From 06c0f5268dda610eeb219ec3573a8d9501fcc879 Mon Sep 17 00:00:00 2001 From: agessaman Date: Tue, 14 Jul 2026 18:35:05 -0700 Subject: [PATCH] fix(companion): persist messages before client connection --- repeater/companion/frame_server.py | 8 +++++++ tests/test_companion_advert_persistence.py | 27 +++++++++++++++++++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/repeater/companion/frame_server.py b/repeater/companion/frame_server.py index 4f4315a..90c90c8 100644 --- a/repeater/companion/frame_server.py +++ b/repeater/companion/frame_server.py @@ -54,6 +54,14 @@ class CompanionFrameServer(_BaseFrameServer): ) self.sqlite_handler = sqlite_handler + async def start(self) -> None: + """Start persistence before accepting companion client connections.""" + if self.sqlite_handler: + self.bridge.on_message_received(self._on_message_received) + self.bridge.on_channel_message_received(self._on_channel_message_received) + self.bridge.on_channel_data_received(self._on_channel_data_received) + await super().start() + # ----------------------------------------------------------------- # Persistence hook overrides # ----------------------------------------------------------------- diff --git a/tests/test_companion_advert_persistence.py b/tests/test_companion_advert_persistence.py index 6de4bb0..1d515d0 100644 --- a/tests/test_companion_advert_persistence.py +++ b/tests/test_companion_advert_persistence.py @@ -2,7 +2,7 @@ import asyncio import sqlite3 -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import pytest @@ -59,6 +59,31 @@ def _bridge(sqlite_handler) -> RepeaterCompanionBridge: ) +@pytest.mark.asyncio +async def test_message_before_first_client_connection_survives_restart(tmp_path): + handler = SQLiteHandler(tmp_path) + bridge = _bridge(handler) + server = CompanionFrameServer(bridge, COMPANION_HASH, port=0, sqlite_handler=handler) + with patch("repeater.companion.frame_server._BaseFrameServer.start", AsyncMock()): + await server.start() + assert server._client_writer is None + await bridge._handle_new_message( + { + "contact_pubkey": "01" * 32, + "message_text": "before-first-connect", + "timestamp": 123, + "txt_type": 0, + } + ) + assert bridge.message_queue.is_empty() + + restarted_handler = SQLiteHandler(tmp_path) + message = restarted_handler.companion_pop_message(COMPANION_HASH) + assert message is not None + assert message["text"] == "before-first-connect" + assert message["timestamp"] == 123 + + def test_bulk_save_load_preserves_raw_advert_blob_and_upsert_replaces_it(tmp_path): handler = SQLiteHandler(tmp_path) assert handler.companion_save_contacts(COMPANION_HASH, [_contact()])