diff --git a/repeater/companion/frame_server.py b/repeater/companion/frame_server.py index 0d55eb1..1764694 100644 --- a/repeater/companion/frame_server.py +++ b/repeater/companion/frame_server.py @@ -59,13 +59,23 @@ class CompanionFrameServer(_BaseFrameServer): # ----------------------------------------------------------------- async def _persist_companion_message(self, msg_dict: dict) -> None: - """Persist message to SQLite and pop from bridge queue.""" + """Persist message to SQLite and pop from bridge queue. + + The bridge's ``offline_queue_size`` (``message_queue._max_size``) doubles + as the SQLite retention limit: 0 disables offline storage entirely, so the + message is dropped instead of persisted. + """ if not self.sqlite_handler: return + retention = getattr(self.bridge.message_queue, "_max_size", None) + if retention == 0: + self.bridge.message_queue.pop_last() + return await asyncio.to_thread( self.sqlite_handler.companion_push_message, self.companion_hash, msg_dict, + retention, ) self.bridge.message_queue.pop_last() diff --git a/repeater/companion/utils.py b/repeater/companion/utils.py index 6f2d2c0..ca64092 100644 --- a/repeater/companion/utils.py +++ b/repeater/companion/utils.py @@ -1,7 +1,46 @@ """Shared utilities for Companion (e.g. validation for config sync).""" +from __future__ import annotations + +import logging +from typing import Any, Dict, Optional + +from pymc_core.companion.constants import DEFAULT_MAX_CONTACTS + +logger = logging.getLogger(__name__) + _INVALID_NODE_NAME_CHARS = "\n\r\x00" +# Optional per-companion RepeaterCompanionBridge constructor settings (power-user). +COMPANION_BRIDGE_SETTING_KEYS = frozenset({"max_contacts", "offline_queue_size"}) + +# Settings that must not be applied from config (fixed at pymc_core defaults). +_COMPANION_IGNORED_BRIDGE_KEYS = frozenset({"max_channels", "adv_type"}) + +# Contact flag bit 0 marks a favourite (protected from forced-trim eviction). +_CONTACT_FLAG_FAVOURITE = 0x01 + + +class CompanionContactCapacityError(Exception): + """Persisted companion contacts exceed configured max_contacts.""" + + def __init__( + self, + companion_hash: str, + stored_count: int, + max_contacts: int, + companion_name: Optional[str] = None, + ) -> None: + self.companion_hash = companion_hash + self.stored_count = stored_count + self.max_contacts = max_contacts + self.companion_name = companion_name + label = f"'{companion_name}'" if companion_name else companion_hash + super().__init__( + f"Companion {label}: {stored_count} contacts in storage exceeds " + f"max_contacts={max_contacts}. Increase max_contacts or remove contacts before starting." + ) + def normalize_companion_identity_key(identity_key: str) -> str: """Strip whitespace and remove optional 0x prefix so fromhex() is consistent across installs.""" @@ -23,3 +62,219 @@ def validate_companion_node_name(value: str) -> str: if any(c in s for c in _INVALID_NODE_NAME_CHARS): raise ValueError("node_name contains invalid characters") return s + + +def parse_positive_int(value: Any, field_name: str, *, minimum: int = 1) -> int: + """Parse a positive integer from config or API input.""" + try: + n = int(value) + except (TypeError, ValueError) as e: + raise ValueError(f"{field_name} must be a positive integer") from e + if n < minimum: + raise ValueError(f"{field_name} must be >= {minimum}") + return n + + +def parse_companion_bridge_kwargs(settings: dict) -> Dict[str, int]: + """Extract optional RepeaterCompanionBridge kwargs from companion settings. + + Only ``max_contacts`` and ``offline_queue_size`` are honored. ``max_channels`` and + ``adv_type`` are ignored with a warning if present. + """ + if not settings: + return {} + for key in _COMPANION_IGNORED_BRIDGE_KEYS: + if key in settings: + logger.warning( + "Companion setting %r is not supported and will be ignored (fixed default)", + key, + ) + kwargs: Dict[str, int] = {} + if "max_contacts" in settings: + max_contacts = parse_positive_int(settings["max_contacts"], "max_contacts") + kwargs["max_contacts"] = max_contacts + if "offline_queue_size" in settings: + # 0 is valid and means "off" (no offline message storage). + kwargs["offline_queue_size"] = parse_positive_int( + settings["offline_queue_size"], "offline_queue_size", minimum=0 + ) + return kwargs + + +def effective_max_contacts(bridge_kwargs: Dict[str, int]) -> int: + """Return max_contacts from parsed kwargs or pymc_core default.""" + return bridge_kwargs.get("max_contacts", DEFAULT_MAX_CONTACTS) + + +def merge_companion_settings_update(current_settings: dict, patch: dict) -> Dict[str, Any]: + """Merge a companion settings PATCH into current settings. + + Raises: + ValueError: Unknown setting or invalid bridge setting value. + """ + merged = dict(current_settings or {}) + for key, value in patch.items(): + if key not in COMPANION_SETTINGS_ALLOWLIST: + raise ValueError(f"Unknown companion setting: {key}") + if key in COMPANION_BRIDGE_SETTING_KEYS: + parsed = parse_companion_bridge_kwargs({key: value}) + merged[key] = parsed[key] + else: + merged[key] = value + return merged + + +def validate_companion_config_capacity( + identity: dict, + sqlite_handler: Any, + *, + companion_name: Optional[str] = None, + settings: Optional[dict] = None, +) -> None: + """Raise CompanionContactCapacityError if persisted contacts exceed configured max_contacts.""" + if sqlite_handler is None: + return + identity_key = identity.get("identity_key") + if not identity_key: + return + merged_settings = settings if settings is not None else (identity.get("settings") or {}) + max_contacts = effective_max_contacts(parse_companion_bridge_kwargs(merged_settings)) + companion_hash = companion_hash_str_from_identity_key(identity_key) + check_companion_contact_capacity( + companion_hash, + max_contacts, + sqlite_handler, + companion_name=companion_name, + ) + + +def check_companion_contact_capacity( + companion_hash: str, + max_contacts: int, + sqlite_handler: Any, + *, + companion_name: Optional[str] = None, +) -> None: + """Raise CompanionContactCapacityError if persisted contacts exceed max_contacts.""" + if sqlite_handler is None: + return + stored_count = sqlite_handler.companion_count_contacts(companion_hash) + if stored_count > max_contacts: + raise CompanionContactCapacityError( + companion_hash, stored_count, max_contacts, companion_name=companion_name + ) + + +def select_companion_contacts_to_trim(contacts, max_contacts: int): + """Select which persisted contacts to keep/remove to fit ``max_contacts``. + + Mirrors ``ContactStore.add_or_overwrite`` eviction: the oldest non-favourite + contacts (by ``lastmod``) are removed first; favourites (flags bit 0) are + never evicted. + + Returns: + (keep, removed): lists of contact dicts. + + Raises: + ValueError: favourites alone exceed ``max_contacts`` (cannot trim). + """ + contacts = list(contacts) + if len(contacts) <= max_contacts: + return contacts, [] + favourites = [c for c in contacts if int(c.get("flags", 0)) & _CONTACT_FLAG_FAVOURITE] + if len(favourites) > max_contacts: + raise ValueError( + f"Cannot trim to max_contacts={max_contacts}: " + f"{len(favourites)} favourite contacts cannot be evicted" + ) + non_favourites = [c for c in contacts if not int(c.get("flags", 0)) & _CONTACT_FLAG_FAVOURITE] + # Keep the newest non-favourites by lastmod; evict the oldest. + non_favourites.sort(key=lambda c: int(c.get("lastmod", 0))) + keep_count = max_contacts - len(favourites) + removed = non_favourites[: len(non_favourites) - keep_count] + kept_non_favourites = non_favourites[len(non_favourites) - keep_count :] + return favourites + kept_non_favourites, removed + + +def trim_companion_contacts_to_fit( + sqlite_handler: Any, companion_hash: str, max_contacts: int +) -> int: + """Trim persisted contacts (favourite-aware) down to ``max_contacts``. + + Loads the companion's contacts, evicts the oldest non-favourites per + :func:`select_companion_contacts_to_trim`, persists the kept set, and returns + the number removed (0 if already within the limit). + + Raises: + ValueError: favourites alone exceed ``max_contacts`` (cannot trim). + RuntimeError: persisting the trimmed contact list failed. + """ + if sqlite_handler is None: + return 0 + contacts = sqlite_handler.companion_load_contacts(companion_hash) + keep, removed = select_companion_contacts_to_trim(contacts, max_contacts) + if not removed: + return 0 + if not sqlite_handler.companion_save_contacts(companion_hash, keep): + raise RuntimeError(f"Failed to persist trimmed contacts for {companion_hash}") + return len(removed) + + +def enforce_companion_contact_capacity( + companion_hash: str, + max_contacts: int, + sqlite_handler: Any, + *, + trim: bool = False, + companion_name: Optional[str] = None, +) -> int: + """Ensure persisted contacts fit ``max_contacts`` at load time. + + With ``trim=False`` (default) this is a guard: it raises + :class:`CompanionContactCapacityError` when over capacity. With ``trim=True`` + (the ``trim_contacts_on_overflow`` policy) it trims favourite-aware to fit, + persists, and returns the number of contacts removed. + """ + if not trim: + check_companion_contact_capacity( + companion_hash, max_contacts, sqlite_handler, companion_name=companion_name + ) + return 0 + return trim_companion_contacts_to_fit(sqlite_handler, companion_hash, max_contacts) + + +def format_companion_bridge_limits(bridge_kwargs: Dict[str, int]) -> str: + """Format non-default bridge limits for log lines.""" + if not bridge_kwargs: + return "" + parts = [f"{k}={v}" for k, v in sorted(bridge_kwargs.items())] + return ", " + ", ".join(parts) + + +def companion_hash_str_from_identity_key(identity_key: Any) -> str: + """Derive companion_hash storage key (0xHH) from an identity_key config value.""" + from pymc_core import LocalIdentity + + if isinstance(identity_key, str): + key_bytes = bytes.fromhex(normalize_companion_identity_key(identity_key)) + elif isinstance(identity_key, bytes): + key_bytes = identity_key + else: + raise ValueError("identity_key has unknown type") + pubkey_byte = LocalIdentity(seed=key_bytes).get_public_key()[0] + return f"0x{pubkey_byte:02x}" + + +# All companion settings writable via identity API (tcp + bridge power-user keys). +COMPANION_SETTINGS_ALLOWLIST = frozenset( + { + "node_name", + "tcp_port", + "bind_address", + "tcp_timeout", + # Persistent opt-in: trim oldest non-favourite contacts to fit max_contacts + # at load instead of refusing to start when over capacity. + "trim_contacts_on_overflow", + *COMPANION_BRIDGE_SETTING_KEYS, + } +) diff --git a/repeater/data_acquisition/sqlite_handler.py b/repeater/data_acquisition/sqlite_handler.py index 9ab18cb..8ea4362 100644 --- a/repeater/data_acquisition/sqlite_handler.py +++ b/repeater/data_acquisition/sqlite_handler.py @@ -2331,6 +2331,20 @@ class SQLiteHandler: return 0 # Companion persistence methods + def companion_count_contacts(self, companion_hash: str) -> int: + """Return the number of persisted contacts for a companion.""" + try: + with self._connect() as conn: + cursor = conn.execute( + "SELECT COUNT(*) FROM companion_contacts WHERE companion_hash = ?", + (companion_hash,), + ) + row = cursor.fetchone() + return int(row[0]) if row else 0 + except Exception as e: + logger.error(f"Failed to count companion contacts: {e}") + return 0 + def companion_load_contacts(self, companion_hash: str) -> List[Dict]: """Load contacts for a companion from storage.""" try: @@ -2627,7 +2641,9 @@ class SQLiteHandler: logger.error(f"Failed to load companion messages: {e}") return [] - def companion_push_message(self, companion_hash: str, msg: Dict) -> bool: + def companion_push_message( + self, companion_hash: str, msg: Dict, max_messages: Optional[int] = None + ) -> bool: """Append a message to the companion's queue. Deduplicates by (companion_hash, packet_hash) using INSERT OR IGNORE @@ -2635,6 +2651,9 @@ class SQLiteHandler: previous SELECT + INSERT round-trip (two statements, two SD-card reads) with a single atomic statement. + When ``max_messages`` is set, the oldest rows beyond that retention limit + are trimmed after a successful insert (power-user ``offline_queue_size``). + Returns True if inserted, False if the message was a duplicate (skipped). """ try: @@ -2663,8 +2682,23 @@ class SQLiteHandler: time.time(), ), ) + inserted = cursor.rowcount > 0 + if inserted and max_messages is not None: + # Keep the newest `max_messages` rows; drop older overflow. + conn.execute( + """ + DELETE FROM companion_messages + WHERE companion_hash = ? AND id NOT IN ( + SELECT id FROM companion_messages + WHERE companion_hash = ? + ORDER BY created_at DESC, id DESC + LIMIT ? + ) + """, + (companion_hash, companion_hash, max_messages), + ) conn.commit() - return cursor.rowcount > 0 + return inserted except Exception as e: logger.error(f"Failed to push companion message: {e}") return False diff --git a/repeater/main.py b/repeater/main.py index f6355b6..3297e0b 100644 --- a/repeater/main.py +++ b/repeater/main.py @@ -7,7 +7,15 @@ import socket import sys import time -from repeater.companion.utils import normalize_companion_identity_key, validate_companion_node_name +from repeater.companion.utils import ( + CompanionContactCapacityError, + effective_max_contacts, + enforce_companion_contact_capacity, + format_companion_bridge_limits, + normalize_companion_identity_key, + parse_companion_bridge_kwargs, + validate_companion_node_name, +) from repeater.config import NullRadio, get_radio_for_board, load_config, save_config from repeater.config_manager import ConfigManager from repeater.data_acquisition.glass_handler import GlassHandler @@ -570,6 +578,25 @@ class RepeaterDaemon: return _sync + bridge_kwargs = parse_companion_bridge_kwargs(settings) + max_contacts = effective_max_contacts(bridge_kwargs) + if sqlite_handler: + trimmed = enforce_companion_contact_capacity( + companion_hash_str, + max_contacts, + sqlite_handler, + trim=bool(settings.get("trim_contacts_on_overflow")), + companion_name=name, + ) + if trimmed: + logger.warning( + "Companion '%s': trimmed %d contact(s) to fit " + "max_contacts=%d (trim_contacts_on_overflow)", + name, + trimmed, + max_contacts, + ) + bridge = RepeaterCompanionBridge( identity=identity, # Tag the injector with this companion's hash so inject_packet can @@ -583,6 +610,7 @@ class RepeaterDaemon: sqlite_handler=sqlite_handler, companion_hash=companion_hash_str, on_prefs_saved=_make_sync_node_name_to_config(name), + **bridge_kwargs, ) # Load contacts from SQLite @@ -616,24 +644,29 @@ class RepeaterDaemon: ch = Channel(name=row.get("name", ""), secret=raw) bridge.channels.set(row.get("channel_idx", 0), ch) - # Preload queued messages from SQLite into bridge - for msg_dict in sqlite_handler.companion_load_messages(companion_hash_str): - from pymc_core.companion.models import QueuedMessage + # Preload queued messages from SQLite into bridge, bounded by + # offline_queue_size (0 disables offline storage entirely). + retention = getattr(bridge.message_queue, "_max_size", None) + if retention != 0: + for msg_dict in sqlite_handler.companion_load_messages( + companion_hash_str, limit=retention or 100 + ): + from pymc_core.companion.models import QueuedMessage - sk = msg_dict.get("sender_key", b"") - if isinstance(sk, str): - sk = bytes.fromhex(sk) - bridge.message_queue.push( - QueuedMessage( - sender_key=sk, - txt_type=msg_dict.get("txt_type", 0), - timestamp=msg_dict.get("timestamp", 0), - text=msg_dict.get("text", ""), - is_channel=bool(msg_dict.get("is_channel", False)), - channel_idx=msg_dict.get("channel_idx", 0), - path_len=msg_dict.get("path_len", 0), + sk = msg_dict.get("sender_key", b"") + if isinstance(sk, str): + sk = bytes.fromhex(sk) + bridge.message_queue.push( + QueuedMessage( + sender_key=sk, + txt_type=msg_dict.get("txt_type", 0), + timestamp=msg_dict.get("timestamp", 0), + text=msg_dict.get("text", ""), + is_channel=bool(msg_dict.get("is_channel", False)), + channel_idx=msg_dict.get("channel_idx", 0), + path_len=msg_dict.get("path_len", 0), + ) ) - ) # Ensure public channel (0) exists with default key for new companions from repeater.companion.constants import DEFAULT_PUBLIC_CHANNEL_SECRET @@ -666,11 +699,15 @@ class RepeaterDaemon: identity_type="companion", ) + limits = format_companion_bridge_limits(bridge_kwargs) logger.info( f"Loaded companion '{name}': hash=0x{companion_hash:02x}, " - f"port={tcp_port}, bind={bind_address}, client_idle_timeout_sec={client_idle_timeout_sec}" + f"port={tcp_port}, bind={bind_address}, " + f"client_idle_timeout_sec={client_idle_timeout_sec}{limits}" ) + except CompanionContactCapacityError as e: + logger.error("%s", e) except Exception as e: logger.error(f"Failed to load companion '{name}': {e}", exc_info=True) @@ -736,13 +773,35 @@ class RepeaterDaemon: tcp_timeout_raw = settings.get("tcp_timeout", 120) client_idle_timeout_sec = None if tcp_timeout_raw == 0 else int(tcp_timeout_raw) + bridge_kwargs = parse_companion_bridge_kwargs(settings) + max_contacts = effective_max_contacts(bridge_kwargs) + if sqlite_handler: + trimmed = enforce_companion_contact_capacity( + companion_hash_str, + max_contacts, + sqlite_handler, + trim=bool(settings.get("trim_contacts_on_overflow")), + companion_name=name, + ) + if trimmed: + logger.warning( + "Hot-reload companion '%s': trimmed %d contact(s) to fit " + "max_contacts=%d (trim_contacts_on_overflow)", + name, + trimmed, + max_contacts, + ) + bridge = RepeaterCompanionBridge( identity=identity, - packet_injector=self.router.inject_packet, + packet_injector=functools.partial( + self.router.inject_packet, origin_hash=companion_hash_str + ), node_name=node_name, radio_config=radio_config, sqlite_handler=sqlite_handler, companion_hash=companion_hash_str, + **bridge_kwargs, ) if sqlite_handler: @@ -773,23 +832,27 @@ class RepeaterDaemon: ch = Channel(name=row.get("name", ""), secret=raw) bridge.channels.set(row.get("channel_idx", 0), ch) - for msg_dict in sqlite_handler.companion_load_messages(companion_hash_str): - from pymc_core.companion.models import QueuedMessage + retention = getattr(bridge.message_queue, "_max_size", None) + if retention != 0: + for msg_dict in sqlite_handler.companion_load_messages( + companion_hash_str, limit=retention or 100 + ): + from pymc_core.companion.models import QueuedMessage - sk = msg_dict.get("sender_key", b"") - if isinstance(sk, str): - sk = bytes.fromhex(sk) - bridge.message_queue.push( - QueuedMessage( - sender_key=sk, - txt_type=msg_dict.get("txt_type", 0), - timestamp=msg_dict.get("timestamp", 0), - text=msg_dict.get("text", ""), - is_channel=bool(msg_dict.get("is_channel", False)), - channel_idx=msg_dict.get("channel_idx", 0), - path_len=msg_dict.get("path_len", 0), + sk = msg_dict.get("sender_key", b"") + if isinstance(sk, str): + sk = bytes.fromhex(sk) + bridge.message_queue.push( + QueuedMessage( + sender_key=sk, + txt_type=msg_dict.get("txt_type", 0), + timestamp=msg_dict.get("timestamp", 0), + text=msg_dict.get("text", ""), + is_channel=bool(msg_dict.get("is_channel", False)), + channel_idx=msg_dict.get("channel_idx", 0), + path_len=msg_dict.get("path_len", 0), + ) ) - ) if bridge.get_channel(0) is None: bridge.set_channel(0, "Public", DEFAULT_PUBLIC_CHANNEL_SECRET) @@ -819,9 +882,11 @@ class RepeaterDaemon: identity_type="companion", ) + limits = format_companion_bridge_limits(bridge_kwargs) logger.info( f"Hot-reload: Loaded companion '{name}': hash=0x{companion_hash:02x}, " - f"port={tcp_port}, bind={bind_address}, client_idle_timeout_sec={client_idle_timeout_sec}" + f"port={tcp_port}, bind={bind_address}, " + f"client_idle_timeout_sec={client_idle_timeout_sec}{limits}" ) async def _on_raw_rx_for_companions( diff --git a/repeater/web/api_endpoints.py b/repeater/web/api_endpoints.py index 2bd5887..5dad39a 100644 --- a/repeater/web/api_endpoints.py +++ b/repeater/web/api_endpoints.py @@ -17,6 +17,13 @@ from repeater.companion.identity_resolve import ( find_companion_index, heal_companion_empty_names, ) +from repeater.companion.utils import ( + CompanionContactCapacityError, + merge_companion_settings_update, + parse_companion_bridge_kwargs, + trim_companion_contacts_to_fit, + validate_companion_config_capacity, +) from repeater.config import resolve_storage_dir from repeater.policy_engine import PolicyEngine from repeater.service_utils import get_buildroot_image_info @@ -4263,6 +4270,11 @@ class APIEndpoints: if any(str(c.get("name") or "").strip() == name for c in companions): return self._error(f"Companion with name '{name}' already exists") + try: + bridge_settings = parse_companion_bridge_kwargs(settings) + except ValueError as e: + return self._error(str(e)) + comp_settings = { "node_name": settings.get("node_name") or name, "tcp_port": settings.get("tcp_port", 5000), @@ -4270,12 +4282,37 @@ class APIEndpoints: } if "tcp_timeout" in settings: comp_settings["tcp_timeout"] = settings["tcp_timeout"] + if "trim_contacts_on_overflow" in settings: + comp_settings["trim_contacts_on_overflow"] = bool( + settings["trim_contacts_on_overflow"] + ) + comp_settings.update(bridge_settings) new_identity = { "name": name, "identity_key": identity_key, "type": identity_type, "settings": comp_settings, } + sqlite_handler = None + repeater_handler = ( + getattr(self.daemon_instance, "repeater_handler", None) + if self.daemon_instance + else None + ) + if repeater_handler and getattr(repeater_handler, "storage", None): + sqlite_handler = repeater_handler.storage.sqlite_handler + if sqlite_handler and identity_key: + try: + validate_companion_config_capacity( + new_identity, + sqlite_handler, + companion_name=name, + settings=comp_settings, + ) + except CompanionContactCapacityError as e: + return self._error(str(e)) + except (ValueError, TypeError) as e: + return self._error(str(e)) companions.append(new_identity) self.config["identities"]["companions"] = companions else: @@ -4304,6 +4341,7 @@ class APIEndpoints: # Hot reload - register identity immediately registration_success = False + companion_activation_error = None if identity_type == "room_server" and self.daemon_instance: try: from pymc_core import LocalIdentity @@ -4358,6 +4396,10 @@ class APIEndpoints: future.result(timeout=15) registration_success = True logger.info(f"Hot reload: Companion '{name}' activated immediately") + except CompanionContactCapacityError as cap_error: + # A restart won't fix a capacity overflow; report the real cause. + companion_activation_error = str(cap_error) + logger.warning(f"Hot reload companion '{name}' not activated: {cap_error}") except Exception as comp_error: logger.warning( f"Hot reload companion '{name}' failed: {comp_error}. Restart required to activate.", @@ -4365,11 +4407,17 @@ class APIEndpoints: ) if identity_type == "companion": - message = ( - f"Companion '{name}' created successfully and activated immediately!" - if registration_success - else f"Companion '{name}' created successfully. Restart required to activate." - ) + if registration_success: + message = f"Companion '{name}' created successfully and activated immediately!" + elif companion_activation_error: + message = ( + f"Companion '{name}' created, but not activated: " + f"{companion_activation_error}" + ) + else: + message = ( + f"Companion '{name}' created successfully. Restart required to activate." + ) else: message = ( f"Identity '{name}' created successfully and activated immediately!" @@ -4473,13 +4521,56 @@ class APIEndpoints: except ValueError: pass + trimmed_count = 0 if "settings" in data: - if "settings" not in identity: - identity["settings"] = {} - # Only allow companion settings - for k, v in data["settings"].items(): - if k in ("node_name", "tcp_port", "bind_address", "tcp_timeout"): - identity["settings"][k] = v + try: + merged_settings = merge_companion_settings_update( + identity.get("settings") or {}, + data["settings"], + ) + except ValueError as e: + return self._error(str(e)) + + sqlite_handler = None + repeater_handler = ( + getattr(self.daemon_instance, "repeater_handler", None) + if self.daemon_instance + else None + ) + if repeater_handler and getattr(repeater_handler, "storage", None): + sqlite_handler = repeater_handler.storage.sqlite_handler + if sqlite_handler and identity.get("identity_key"): + try: + validate_companion_config_capacity( + identity, + sqlite_handler, + companion_name=resolved_name, + settings=merged_settings, + ) + except CompanionContactCapacityError as e: + if not data.get("force_trim"): + return self._error(str(e)) + # Power-user opt-in: trim persisted contacts down to the + # new limit (favourite-aware) instead of rejecting. + try: + trimmed_count = trim_companion_contacts_to_fit( + sqlite_handler, e.companion_hash, e.max_contacts + ) + except ValueError as trim_err: + return self._error(str(trim_err)) + except RuntimeError: + return self._error("Failed to persist trimmed contacts") + logger.info( + "Force-trimmed %d contact(s) for companion '%s' " + "to fit max_contacts=%d", + trimmed_count, + resolved_name, + e.max_contacts, + ) + except (ValueError, TypeError) as e: + return self._error(str(e)) + + identity["settings"] = merged_settings companions[identity_index] = identity self.config["identities"]["companions"] = companions @@ -4491,6 +4582,12 @@ class APIEndpoints: f"Companion '{resolved_name}' updated successfully. " "Restart required to apply changes." ) + if trimmed_count: + message = ( + f"Companion '{resolved_name}' updated successfully; " + f"trimmed {trimmed_count} contact(s) to fit the new limit. " + "Restart required to apply changes." + ) return self._success(identity, message=message) # Room server path diff --git a/repeater/web/companion_endpoints.py b/repeater/web/companion_endpoints.py index 0f433d3..0d3c4a8 100644 --- a/repeater/web/companion_endpoints.py +++ b/repeater/web/companion_endpoints.py @@ -15,6 +15,7 @@ import time from typing import Optional import cherrypy +from pymc_core.companion.constants import DEFAULT_OFFLINE_QUEUE_SIZE from repeater.companion.utils import validate_companion_node_name @@ -257,6 +258,10 @@ class CompanionAPIEndpoints: "is_running": b.is_running, "contacts_count": b.contacts.get_count(), "channels_count": b.channels.get_count(), + "max_contacts": b.contacts.max_contacts, + "offline_queue_size": getattr( + b.message_queue, "_max_size", DEFAULT_OFFLINE_QUEUE_SIZE + ), } ) return self._success(items) diff --git a/tests/test_companion_bridge_frame_utils.py b/tests/test_companion_bridge_frame_utils.py index fcd1c09..c0e7c09 100644 --- a/tests/test_companion_bridge_frame_utils.py +++ b/tests/test_companion_bridge_frame_utils.py @@ -147,7 +147,7 @@ async def test_frame_server_persistence_paths_and_stop(): srv._build_message_frame = MagicMock(return_value=b"frame") await srv._persist_companion_message({"text": "x"}) - sqlite.companion_push_message.assert_called_once_with("h", {"text": "x"}) + sqlite.companion_push_message.assert_called_once_with("h", {"text": "x"}, None) bridge.message_queue.pop_last.assert_called_once() msg = srv._sync_next_from_persistence() diff --git a/tests/test_companion_settings.py b/tests/test_companion_settings.py new file mode 100644 index 0000000..92ccdd8 --- /dev/null +++ b/tests/test_companion_settings.py @@ -0,0 +1,270 @@ +"""Tests for per-companion bridge settings parsing and startup guard.""" + +from __future__ import annotations + +import logging +from unittest.mock import MagicMock + +import pytest + +from repeater.companion.utils import ( + COMPANION_SETTINGS_ALLOWLIST, + CompanionContactCapacityError, + check_companion_contact_capacity, + effective_max_contacts, + enforce_companion_contact_capacity, + merge_companion_settings_update, + parse_companion_bridge_kwargs, + parse_positive_int, + select_companion_contacts_to_trim, + trim_companion_contacts_to_fit, + validate_companion_config_capacity, +) + +# pymc_core defaults (CompanionBridge / ContactStore) +_DEFAULT_MAX_CONTACTS = 1000 + + +class TestParsePositiveInt: + def test_valid(self): + assert parse_positive_int("100", "max_contacts") == 100 + + def test_invalid_type(self): + with pytest.raises(ValueError, match="max_contacts"): + parse_positive_int("abc", "max_contacts") + + def test_below_minimum(self): + with pytest.raises(ValueError, match="max_contacts"): + parse_positive_int(0, "max_contacts") + + +class TestParseCompanionBridgeKwargs: + def test_empty_settings(self): + assert parse_companion_bridge_kwargs({}) == {} + + def test_max_contacts_and_offline_queue(self): + assert parse_companion_bridge_kwargs( + {"max_contacts": 2000, "offline_queue_size": 1024} + ) == {"max_contacts": 2000, "offline_queue_size": 1024} + + def test_ignored_keys_warn(self, caplog): + caplog.set_level(logging.WARNING) + result = parse_companion_bridge_kwargs( + {"max_contacts": 500, "max_channels": 64, "adv_type": 2} + ) + assert result == {"max_contacts": 500} + assert any("max_channels" in r.message for r in caplog.records) + assert any("adv_type" in r.message for r in caplog.records) + + def test_invalid_max_contacts(self): + with pytest.raises(ValueError): + parse_companion_bridge_kwargs({"max_contacts": -1}) + + +class TestEffectiveMaxContacts: + def test_default(self): + assert effective_max_contacts({}) == _DEFAULT_MAX_CONTACTS + + def test_override(self): + assert effective_max_contacts({"max_contacts": 500}) == 500 + + +class TestMergeCompanionSettingsUpdate: + def test_merges_bridge_settings(self): + merged = merge_companion_settings_update( + {"node_name": "a"}, + {"max_contacts": 500}, + ) + assert merged == {"node_name": "a", "max_contacts": 500} + + def test_unknown_key_raises(self): + with pytest.raises(ValueError, match="Unknown companion setting"): + merge_companion_settings_update({}, {"max_channels": 64}) + + +class TestValidateCompanionConfigCapacity: + def test_uses_merged_settings_not_stale_identity(self): + identity = { + "identity_key": "aa" * 32, + "settings": {"max_contacts": 1000}, + } + sqlite = MagicMock() + sqlite.companion_count_contacts.return_value = 600 + with pytest.raises(CompanionContactCapacityError): + validate_companion_config_capacity( + identity, + sqlite, + settings={"max_contacts": 500}, + ) + sqlite.companion_count_contacts.assert_called_once() + + +class TestCheckCompanionContactCapacity: + def test_skips_without_sqlite(self): + check_companion_contact_capacity("0x01", 100, None) + + def test_passes_when_under_limit(self): + sqlite = MagicMock() + sqlite.companion_count_contacts.return_value = 100 + check_companion_contact_capacity("0x01", 500, sqlite) + + def test_raises_when_over_limit(self): + sqlite = MagicMock() + sqlite.companion_count_contacts.return_value = 812 + with pytest.raises(CompanionContactCapacityError) as exc: + check_companion_contact_capacity("0xab", 500, sqlite, companion_name="BotCompanion") + assert exc.value.stored_count == 812 + assert exc.value.max_contacts == 500 + assert "BotCompanion" in str(exc.value) + + +class TestOfflineQueueOff: + def test_zero_allowed(self): + assert parse_companion_bridge_kwargs({"offline_queue_size": 0}) == {"offline_queue_size": 0} + + def test_max_contacts_zero_still_rejected(self): + with pytest.raises(ValueError, match="max_contacts"): + parse_companion_bridge_kwargs({"max_contacts": 0}) + + +class TestSelectCompanionContactsToTrim: + @staticmethod + def _c(pk, flags=0, lastmod=0): + return {"pubkey": pk, "flags": flags, "lastmod": lastmod} + + def test_under_limit_keeps_all(self): + contacts = [self._c(b"\x01"), self._c(b"\x02")] + keep, removed = select_companion_contacts_to_trim(contacts, 5) + assert removed == [] + assert keep == contacts + + def test_evicts_oldest_non_favourite_and_protects_favourites(self): + contacts = [ + self._c(b"\x01", lastmod=10), + self._c(b"\x02", lastmod=30), + self._c(b"\x03", flags=1, lastmod=5), # favourite + oldest -> protected + self._c(b"\x04", lastmod=20), + ] + keep, removed = select_companion_contacts_to_trim(contacts, 2) + assert {c["pubkey"] for c in keep} == {b"\x03", b"\x02"} + assert {c["pubkey"] for c in removed} == {b"\x01", b"\x04"} + + def test_refuses_when_favourites_exceed_limit(self): + contacts = [ + self._c(b"\x01", flags=1, lastmod=1), + self._c(b"\x02", flags=1, lastmod=2), + ] + with pytest.raises(ValueError, match="favourite"): + select_companion_contacts_to_trim(contacts, 1) + + +class TestSqliteRetentionTrim: + @staticmethod + def _handler(tmp_path): + from repeater.data_acquisition.sqlite_handler import SQLiteHandler + + return SQLiteHandler(tmp_path) + + @staticmethod + def _push(h, companion_hash, i, max_messages=None): + return h.companion_push_message( + companion_hash, + {"text": f"m{i}", "timestamp": i, "packet_hash": f"{companion_hash}-{i}"}, + max_messages=max_messages, + ) + + def test_trims_to_max_messages(self, tmp_path): + h = self._handler(tmp_path) + for i in range(5): + self._push(h, "0x01", i, max_messages=3) + assert len(h.companion_load_messages("0x01")) == 3 + + def test_none_keeps_all(self, tmp_path): + h = self._handler(tmp_path) + for i in range(5): + self._push(h, "0x01", i, max_messages=None) + assert len(h.companion_load_messages("0x01")) == 5 + + def test_trim_isolated_per_companion(self, tmp_path): + h = self._handler(tmp_path) + for i in range(4): + self._push(h, "0x01", i, max_messages=2) + for i in range(3): + self._push(h, "0x02", i, max_messages=None) + assert len(h.companion_load_messages("0x01")) == 2 + assert len(h.companion_load_messages("0x02")) == 3 + + +class TestTrimContactsOnOverflowPolicy: + @staticmethod + def _contacts(n, favourites=0): + out = [] + for i in range(n): + flags = 1 if i < favourites else 0 + out.append({"pubkey": i.to_bytes(2, "big"), "flags": flags, "lastmod": i}) + return out + + def test_allowlist_includes_policy_key(self): + assert "trim_contacts_on_overflow" in COMPANION_SETTINGS_ALLOWLIST + # And it is accepted by the settings merge. + merged = merge_companion_settings_update({}, {"trim_contacts_on_overflow": True}) + assert merged == {"trim_contacts_on_overflow": True} + + def test_trim_helper_persists_kept_set(self): + sqlite = MagicMock() + sqlite.companion_load_contacts.return_value = self._contacts(5) + sqlite.companion_save_contacts.return_value = True + removed = trim_companion_contacts_to_fit(sqlite, "0x01", 3) + assert removed == 2 + saved_hash, saved_contacts = sqlite.companion_save_contacts.call_args[0] + assert saved_hash == "0x01" + assert len(saved_contacts) == 3 + + def test_trim_helper_noop_when_under_limit(self): + sqlite = MagicMock() + sqlite.companion_load_contacts.return_value = self._contacts(2) + assert trim_companion_contacts_to_fit(sqlite, "0x01", 5) == 0 + sqlite.companion_save_contacts.assert_not_called() + + def test_enforce_guards_by_default(self): + sqlite = MagicMock() + sqlite.companion_count_contacts.return_value = 600 + with pytest.raises(CompanionContactCapacityError): + enforce_companion_contact_capacity("0x01", 500, sqlite) + sqlite.companion_save_contacts.assert_not_called() + + def test_enforce_trims_when_policy_enabled(self): + sqlite = MagicMock() + sqlite.companion_load_contacts.return_value = self._contacts(600) + sqlite.companion_save_contacts.return_value = True + removed = enforce_companion_contact_capacity("0x01", 500, sqlite, trim=True) + assert removed == 100 + + +class TestPersistSkipWhenOff: + @staticmethod + def _frame_server(max_size): + from repeater.companion.frame_server import CompanionFrameServer + + fs = CompanionFrameServer.__new__(CompanionFrameServer) + fs.sqlite_handler = MagicMock() + fs.companion_hash = "0x01" + bridge = MagicMock() + bridge.message_queue._max_size = max_size + fs.bridge = bridge + return fs + + def test_skips_persistence_when_retention_zero(self): + import asyncio + + fs = self._frame_server(0) + asyncio.run(fs._persist_companion_message({"text": "x"})) + fs.sqlite_handler.companion_push_message.assert_not_called() + fs.bridge.message_queue.pop_last.assert_called_once() + + def test_persists_with_retention(self): + import asyncio + + fs = self._frame_server(7) + asyncio.run(fs._persist_companion_message({"text": "x"})) + fs.sqlite_handler.companion_push_message.assert_called_once_with("0x01", {"text": "x"}, 7)