From b7f96c8241d20b080e914ca1e07e9ae7ce798822 Mon Sep 17 00:00:00 2001 From: Ali Khokhar <20476625+Alishahryar1@users.noreply.github.com> Date: Sun, 27 Sep 2026 02:16:48 -0700 Subject: [PATCH] patch: Migrate messaging history to fcc.db (#1919) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Why Messaging stores conversations and tracked platform messages in a JSON file with delayed whole-file writes and separate file-write coordination. Move that state into `fcc.db` so related changes commit atomically and database constraints protect conversation links. ## How - Add numbered migration 4 for messaging trees, nodes, prompt/status references, tracked messages, and a one-time import receipt, with scoped keys, foreign keys, constraints and indexes. - Give the application ownership of the shared database and process lock. Code and messaging release storage only after their workers drain. - Commit admissions and successor claims before launching harness work. Publish workflow admission counters immediately after successful stop/clear commits, before notifications and cleanup. A failed write preserves those counters and the healthy tree workers. - Permit stop and partial clear while interrupted workers are cleaning up. Apply the cleanup readiness check only to new admissions, and validate delayed failure feedback against surviving runnable status references. - Accepted storage-failure behavior: interrupt the affected conversation and queued replies, retire worker ownership, and report the failure. A later admission reconciles interrupted history in the same transaction as its new change. Completed harness work and interrupted queued prompts are never replayed. No retry scheduler is added. - Keep voice handoff draining outside the workflow lock and platform calls outside SQL. Already cancelled transcriptions remain cancelled if a subsequent tree mutation fails. - Drain workers during shutdown even if their interrupted state cannot be saved. Existing startup normalization handles those saved runnable rows on restart. Publish terminal success only after its save commits. - Import supported legacy formats and both platforms at startup, including when messaging is disabled. Delete the old JSON and writer temporary files after a complete committed import. - Isolate each legacy tree and managed-message write with a savepoint. Reject constraint-invalid or unencodable records without retaining partial trees or references. Reserve reference and message IDs only after successful insertion, so rejected records cannot suppress later valid data. - Commit the actual recovered counts and one-time receipt with the recovered rows. Database failures, receipt failures, and commit failures roll back the entire import and preserve its source for retry. - Accepted import behavior: recover valid records or start fresh, preserve incomplete originals, and show a one-time nonblocking warning. Later startups intentionally do not repeat that notice or reimport the retained source. Cleanup failures continue to warn and retry cleanup. No manual repair or confirmation is required. - Preserve native continuation IDs, literal prompt-versus-status reply links, and managed-message retention. Remove the JSON writer, debounce timers, persistence cache and flush machinery, and update the smoke inventory. - Correct two existing browser-test synchronization defects: wait for the captured OpenAI login request and keep simulated Claude Desktop startup state consistent across its two status endpoints. - Updates require FCC to be stopped. Externally modified SQLite files remain outside the supported contract.

RetriggerConfidence Score: 5/5

No blocking findings remain. The outstanding import-warning issue is non-blocking.

Fix All in CodexFindings

1. P2 **Import warning disappears** โ–ถ
Summary The PR moves messaging trees and managed-message history into the application-owned SQLite database. It adds a one-time legacy JSON import that rejects invalid records while preserving the source after a partial import, and updates messaging and browser-test coverage.
Reviews (4) ยท Last reviewed commit: ["Recover malformed legacy messaging recor..."](https://github.com/alishahryar1/free-claude-code/commit/3dfb1e4d1cb641699df8bfc403e8bdbe5f3e9409) --- e2e/code_support.py | 4 +- e2e/conftest.py | 1 + e2e/test_claude_desktop_integration.py | 17 +- e2e/test_connected_accounts.py | 13 +- e2e/test_messaging_storage_recovery.py | 27 + smoke/capabilities.py | 10 +- smoke/features.py | 4 +- smoke/lib/e2e.py | 22 +- smoke/product/test_messaging_product_live.py | 100 +- .../api/admin_static/admin.js | 6 +- .../messaging/command_context.py | 4 +- src/free_claude_code/messaging/commands.py | 8 +- .../messaging/node_event_pipeline.py | 4 +- src/free_claude_code/messaging/node_runner.py | 26 +- .../messaging/session/__init__.py | 5 - .../messaging/session/managed_message_log.py | 135 --- .../messaging/session/persistence.py | 161 --- .../messaging/session/store.py | 162 --- .../messaging/trees/__init__.py | 5 + src/free_claude_code/messaging/trees/graph.py | 14 + .../messaging/trees/manager.py | 288 +++-- src/free_claude_code/messaging/trees/ports.py | 35 + .../messaging/trees/processor.py | 17 +- .../messaging/trees/repository.py | 11 + .../messaging/trees/runtime.py | 38 + .../messaging/trees/transaction.py | 51 + .../messaging/trees/transitions.py | 1 + src/free_claude_code/messaging/turn_intake.py | 11 +- src/free_claude_code/messaging/workflow.py | 117 +-- src/free_claude_code/runtime/application.py | 73 +- src/free_claude_code/runtime/bootstrap.py | 11 +- .../runtime/code_sessions_sqlite.py | 73 +- .../runtime/messaging_import.py | 206 ++++ .../runtime/messaging_sqlite.py | 183 ++++ .../runtime/sqlite_database.py | 118 ++- .../runtime/sqlite_migrations/__init__.py | 8 +- .../v0004_messaging_schema.py | 108 ++ tests/api/test_code_sessions_routes.py | 8 +- tests/application/test_code_sessions.py | 90 +- tests/code_sessions_support.py | 8 + tests/conftest.py | 67 +- tests/messaging/test_handler.py | 262 ++--- .../test_handler_markdown_and_status_edges.py | 62 +- tests/messaging/test_messaging.py | 124 +-- tests/messaging/test_platform_voice_flow.py | 2 +- tests/messaging/test_restart_reply_restore.py | 68 +- .../test_session_store_edge_cases.py | 680 ------------ tests/messaging/test_storage_controls.py | 386 +++++++ tests/messaging/test_tree_concurrency.py | 14 +- .../test_tree_customer_regressions.py | 7 +- .../test_tree_ownership_concurrency.py | 31 +- tests/messaging/test_tree_processor.py | 13 +- tests/messaging/test_tree_repository.py | 11 +- tests/runtime/test_application_runtime.py | 93 +- tests/runtime/test_code_sessions_sqlite.py | 101 +- tests/runtime/test_codex_app_server.py | 15 +- tests/runtime/test_messaging_sqlite.py | 985 ++++++++++++++++++ tests/runtime/test_sqlite_database.py | 55 +- 58 files changed, 3165 insertions(+), 1994 deletions(-) create mode 100644 e2e/test_messaging_storage_recovery.py delete mode 100644 src/free_claude_code/messaging/session/__init__.py delete mode 100644 src/free_claude_code/messaging/session/managed_message_log.py delete mode 100644 src/free_claude_code/messaging/session/persistence.py delete mode 100644 src/free_claude_code/messaging/session/store.py create mode 100644 src/free_claude_code/messaging/trees/ports.py create mode 100644 src/free_claude_code/messaging/trees/transaction.py create mode 100644 src/free_claude_code/runtime/messaging_import.py create mode 100644 src/free_claude_code/runtime/messaging_sqlite.py create mode 100644 src/free_claude_code/runtime/sqlite_migrations/v0004_messaging_schema.py delete mode 100644 tests/messaging/test_session_store_edge_cases.py create mode 100644 tests/messaging/test_storage_controls.py create mode 100644 tests/runtime/test_messaging_sqlite.py diff --git a/e2e/code_support.py b/e2e/code_support.py index 4411d92299..1a98c79281 100644 --- a/e2e/code_support.py +++ b/e2e/code_support.py @@ -6,6 +6,7 @@ from free_claude_code.application.code_sessions import CodeService from free_claude_code.runtime.code_sessions_sqlite import SQLiteCodeStore +from free_claude_code.runtime.sqlite_database import SQLiteDatabase from tests.code_sessions_support import FakeHarness @@ -13,8 +14,9 @@ class CodeControl: def __init__(self, directory): self.harness = FakeHarness() self.folder_picker = FolderPickerControl() + self.database = SQLiteDatabase(directory / "fcc.db", directory / "code.lock") self.service = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), + SQLiteCodeStore(self.database), self.harness, ) self.loop: asyncio.AbstractEventLoop | None = None diff --git a/e2e/conftest.py b/e2e/conftest.py index 76903e3cfb..c9eaa51ece 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -268,6 +268,7 @@ async def fixture_provider(provider_id: str, _settings: Settings) -> BaseProvide configuration=ConfigurationService(ManagedConfigStore()), transcriber=None, code_service=code_control.service, + database=code_control.database, ) monkeypatch.setattr( NativeFolderPicker, diff --git a/e2e/test_claude_desktop_integration.py b/e2e/test_claude_desktop_integration.py index 76fd1655fc..d4ddf3ca68 100644 --- a/e2e/test_claude_desktop_integration.py +++ b/e2e/test_claude_desktop_integration.py @@ -38,19 +38,26 @@ def fail(path, *args, **kwargs): ), ) elif status_failure == "startup": + marker = tmp_path / ".fcc/claude-desktop-disconnect.json" + failed = {"state": "failed", "message": "Startup check failed"} def failed_update(route): result = route.fetch() body = result.json() - if body["disconnect_pending"]: + if marker.exists(): body["connected"] = None - body["update"] = { - "state": "failed", - "message": "Startup check failed", - } + body["update"] = failed + route.fulfill(response=result, json=body) + + def failed_startup(route): + result = route.fetch() + body = result.json() + if marker.exists(): + body["startup"]["integrations"]["claude-desktop"] = failed route.fulfill(response=result, json=body) page.route(endpoint, failed_update) + page.route("**/admin/api/status", failed_startup) dialog.get_by_role("button", name="Disconnect", exact=True).click() expect( dialog.get_by_role( diff --git a/e2e/test_connected_accounts.py b/e2e/test_connected_accounts.py index a5fdcef1ea..3a4b8925de 100644 --- a/e2e/test_connected_accounts.py +++ b/e2e/test_connected_accounts.py @@ -66,6 +66,14 @@ def take_pending_status(self, page: Page) -> Route: page.wait_for_timeout(10) return self.pending_status.pop(0) + def take_pending_login(self, page: Page) -> Route: + deadline = time.monotonic() + 5 + while not self.pending_login: + if time.monotonic() >= deadline: + pytest.fail("Account login request did not reach the route handler") + page.wait_for_timeout(10) + return self.pending_login.pop(0) + def config(self, route: Route) -> None: response = route.fetch() config = response.json() @@ -247,9 +255,10 @@ def test_openai_connect_uses_browser_login( try: assert popup.url == "about:blank" assert popup.evaluate("window.opener === null") is True + login = accounts.take_pending_login(page) assert accounts.login_requests == [("openai", "browser")] - assert len(accounts.pending_login) == 1 - accounts.pending_login.pop().fulfill(json=accounts.statuses["openai"]) + assert not accounts.pending_login + login.fulfill(json=accounts.statuses["openai"]) popup.wait_for_url(f"{admin_base_url}/account-test-sign-in") finally: popup.close() diff --git a/e2e/test_messaging_storage_recovery.py b/e2e/test_messaging_storage_recovery.py new file mode 100644 index 0000000000..2b98be8935 --- /dev/null +++ b/e2e/test_messaging_storage_recovery.py @@ -0,0 +1,27 @@ +import pytest +from playwright.sync_api import expect + + +@pytest.fixture +def admin_client_files(tmp_path): + source = tmp_path / ".fcc/agent_workspace/sessions.json" + source.parent.mkdir(parents=True) + source.write_text('{"conversation":') + return source + + +def test_damaged_legacy_history_warns_without_blocking_settings( + page, admin_base_url, admin_client_files +): + page.goto(f"{admin_base_url}/admin") + page.get_by_role("button", name="Messaging", exact=True).click() + notice = page.locator("#startupMessage") + expect(notice).to_be_visible() + expect(notice).to_contain_text("Messaging can still be used") + expect(page.locator("#field-TELEGRAM_BOT_TOKEN")).to_be_enabled() + assert admin_client_files.read_text() == '{"conversation":' + # A preserved source is not a second writer, and the server stays usable. + response = page.request.get(f"{admin_base_url}/health") + assert response.ok + page.get_by_role("button", name="Providers", exact=True).click() + expect(notice).to_be_hidden() diff --git a/smoke/capabilities.py b/smoke/capabilities.py index 4b47f1ee00..957ce7a6fe 100644 --- a/smoke/capabilities.py +++ b/smoke/capabilities.py @@ -491,7 +491,7 @@ class CapabilityContract: "persistence", "restart_restore", "restart_restore", - "free_claude_code.messaging.session.SessionStore", + "free_claude_code.runtime.messaging_sqlite.SQLiteMessagingStore", "stored tree JSON", "restored scoped reply lookup", "stale pending work marked lost", @@ -501,11 +501,11 @@ class CapabilityContract: "persistence", "session_store", "session_persistence", - "free_claude_code.messaging.session.SessionStore", + "free_claude_code.runtime.messaging_sqlite.SQLiteMessagingStore", "scoped tree data and message log", - "JSON persistence compatible with existing files", - "best-effort flush error logging", - ("tests/messaging/test_session_store_edge_cases.py",), + "transactional SQLite persistence with one-time legacy JSON import", + "storage failures propagate to the owning messaging operation", + ("tests/runtime/test_messaging_sqlite.py",), ), CapabilityContract( "voice", diff --git a/smoke/features.py b/smoke/features.py index 78c8f7c231..f1919eeb04 100644 --- a/smoke/features.py +++ b/smoke/features.py @@ -638,8 +638,8 @@ def has_pytest_coverage(self) -> bool: ), FeatureCoverage( "session_persistence", - "Session JSON preserves scoped trees and message logs", - ("tests/messaging/test_session_store_edge_cases.py",), + "SQLite persistence preserves scoped trees and message logs", + ("tests/runtime/test_messaging_sqlite.py",), (), ("test_restart_restore_and_session_persistence_e2e",), ("messaging",), diff --git a/smoke/lib/e2e.py b/smoke/lib/e2e.py index 92233cb1a4..903b14a6d4 100644 --- a/smoke/lib/e2e.py +++ b/smoke/lib/e2e.py @@ -27,9 +27,10 @@ ) from free_claude_code.harnesses.claude import build_claude_proxy_env from free_claude_code.messaging.models import IncomingMessage, MessageScope -from free_claude_code.messaging.session import SessionStore from free_claude_code.messaging.voice import VoiceCancellationResult from free_claude_code.messaging.workflow import MessagingWorkflow +from free_claude_code.runtime.messaging_sqlite import SQLiteMessagingStore +from free_claude_code.runtime.sqlite_database import SQLiteDatabase from smoke.lib.child_process import run_captured_text from smoke.lib.config import ProviderModel, SmokeConfig, auth_headers from smoke.lib.http import conversation_headers @@ -568,15 +569,18 @@ class FakePlatformDriver: event_batches: list[list[dict[str, Any]]] | None = None platform: FakePlatform = field(init=False) cli_manager: FakeCLIManager = field(init=False) - session_store: SessionStore = field(init=False) + session_store: SQLiteMessagingStore = field(init=False) workflow: MessagingWorkflow = field(init=False) + database: SQLiteDatabase = field(init=False) def __post_init__(self) -> None: self.platform = FakePlatform(self.platform_name) self.cli_manager = FakeCLIManager(self.event_batches) - self.session_store = SessionStore( - storage_path=str(self.tmp_path / f"{self.platform_name}-sessions.json") + self.database = SQLiteDatabase( + self.tmp_path / f"{self.platform_name}-fcc.db", + self.tmp_path / f"{self.platform_name}-fcc.lock", ) + self.session_store = SQLiteMessagingStore(self.database) self.workflow = MessagingWorkflow( self.platform, self.cli_manager, @@ -586,6 +590,14 @@ def __post_init__(self) -> None: ) self.platform.on_message(self.workflow.handle_message) + async def start(self) -> None: + await self.database.start() + await self.workflow.restore() + + async def close(self) -> None: + await self.workflow.close() + await self.database.close() + async def send( self, text: str, @@ -628,7 +640,7 @@ async def wait_for_idle(self, *, timeout_s: float = 5.0) -> None: while time.monotonic() < deadline: pending = [task for task in self.platform._tasks if not task.done()] if not pending and await self._all_tree_nodes_terminal(): - self.session_store.flush_pending_save() + await self.workflow.tree_queue.wait_idle() return await asyncio.sleep(0.02) raise AssertionError("fake platform did not become idle") diff --git a/smoke/product/test_messaging_product_live.py b/smoke/product/test_messaging_product_live.py index eb4ffd0ff5..971024216a 100644 --- a/smoke/product/test_messaging_product_live.py +++ b/smoke/product/test_messaging_product_live.py @@ -1,7 +1,7 @@ import asyncio -import json import pytest +import pytest_asyncio from free_claude_code.messaging.models import MessageScope from free_claude_code.messaging.platforms.ports import MessagingStartupNotice @@ -13,8 +13,10 @@ @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) -async def test_messaging_fake_full_flow_e2e(platform_name: str, tmp_path) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) +async def test_messaging_fake_full_flow_e2e( + driver_factory, platform_name: str, tmp_path +) -> None: + driver = await driver_factory(platform_name, tmp_path) incoming = await driver.send("Please inspect README.", message_id="root_1") @@ -33,7 +35,9 @@ async def test_messaging_fake_full_flow_e2e(platform_name: str, tmp_path) -> Non @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) -async def test_messaging_subagent_control_e2e(platform_name: str, tmp_path) -> None: +async def test_messaging_subagent_control_e2e( + driver_factory, platform_name: str, tmp_path +) -> None: task_events = [ {"type": "session_info", "session_id": "sess_task"}, { @@ -53,7 +57,7 @@ async def test_messaging_subagent_control_e2e(platform_name: str, tmp_path) -> N }, {"type": "exit", "code": 0, "stderr": None}, ] - driver = FakePlatformDriver(platform_name, tmp_path, event_batches=[task_events]) + driver = await driver_factory(platform_name, tmp_path, event_batches=[task_events]) await driver.send("Delegate this safely.", message_id="root_task") @@ -65,9 +69,9 @@ async def test_messaging_subagent_control_e2e(platform_name: str, tmp_path) -> N @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) async def test_messaging_commands_stop_clear_stats_e2e( - platform_name: str, tmp_path + driver_factory, platform_name: str, tmp_path ) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) + driver = await driver_factory(platform_name, tmp_path) root = await driver.send("start work", message_id="root_1") root_status_id = driver.platform.sent[-1]["message_id"] @@ -90,16 +94,17 @@ async def test_messaging_commands_stop_clear_stats_e2e( "clear_1", "clear_all", } <= deleted - assert driver.session_store.load_conversation_snapshot().trees == {} + assert (await driver.session_store.load_conversation_snapshot()).trees == {} @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) async def test_reply_clear_uses_literal_platform_subtree_e2e( + driver_factory, platform_name: str, tmp_path, ) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) + driver = await driver_factory(platform_name, tmp_path) root = await driver.send("root", message_id="root") root_status = driver.platform.sent[-1]["message_id"] sibling = await driver.send( @@ -158,8 +163,10 @@ async def test_reply_clear_uses_literal_platform_subtree_e2e( @pytest.mark.asyncio -async def test_messaging_startup_notice_is_clearable_e2e(tmp_path) -> None: - first_driver = FakePlatformDriver("telegram", tmp_path) +async def test_messaging_startup_notice_is_clearable_e2e( + driver_factory, tmp_path +) -> None: + first_driver = await driver_factory("telegram", tmp_path) scope = MessageScope(platform="telegram", chat_id="chat_1") await first_driver.workflow.publish_startup_notice( @@ -169,9 +176,9 @@ async def test_messaging_startup_notice_is_clearable_e2e(tmp_path) -> None: ) ) first_startup_id = first_driver.platform.sent[-1]["message_id"] - first_driver.session_store.flush_pending_save() - driver = FakePlatformDriver("telegram", tmp_path) + await first_driver.close() + driver = await driver_factory("telegram", tmp_path) await driver.platform.send_message("untracked", "advance fake message ID") await driver.workflow.publish_startup_notice( MessagingStartupNotice( @@ -185,7 +192,7 @@ async def test_messaging_startup_notice_is_clearable_e2e(tmp_path) -> None: "๐Ÿš€ *Claude Code Proxy is online\\!* \\(Bot API\\)" ) assert second_startup["parse_mode"] == "MarkdownV2" - assert driver.session_store.get_tracked_message_ids_for_chat( + assert await driver.session_store.get_tracked_message_ids_for_chat( scope.platform, scope.chat_id ) == [first_startup_id, second_startup_id] @@ -194,7 +201,7 @@ async def test_messaging_startup_notice_is_clearable_e2e(tmp_path) -> None: deleted = {entry["message_id"] for entry in driver.platform.deletes} assert {first_startup_id, second_startup_id, "clear_startup"} <= deleted assert ( - driver.session_store.get_tracked_message_ids_for_chat( + await driver.session_store.get_tracked_message_ids_for_chat( scope.platform, scope.chat_id ) == [] @@ -204,11 +211,12 @@ async def test_messaging_startup_notice_is_clearable_e2e(tmp_path) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) async def test_messaging_active_stop_uses_status_only_e2e( + driver_factory, platform_name: str, tmp_path, monkeypatch, ) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) + driver = await driver_factory(platform_name, tmp_path) started = asyncio.Event() class GatedActiveSession(FakeCLISession): @@ -272,9 +280,9 @@ async def controlled_session(session_id: str | None = None): @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) -async def test_tree_threading_e2e(platform_name: str, tmp_path) -> None: +async def test_tree_threading_e2e(driver_factory, platform_name: str, tmp_path) -> None: batches = [default_cli_events("sess_root"), default_cli_events("sess_branch")] - driver = FakePlatformDriver(platform_name, tmp_path, event_batches=batches) + driver = await driver_factory(platform_name, tmp_path, event_batches=batches) root = await driver.send("root prompt", message_id="root_1") branch = await driver.send( @@ -294,11 +302,12 @@ async def test_tree_threading_e2e(platform_name: str, tmp_path) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) async def test_messaging_queued_scoped_cancel_e2e( + driver_factory, platform_name: str, tmp_path, monkeypatch, ) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) + driver = await driver_factory(platform_name, tmp_path) root_started = asyncio.Event() release_root = asyncio.Event() @@ -392,27 +401,29 @@ async def controlled_session(session_id: str | None = None): edit["message_id"] == cancelled_status_id and "Stopped" in edit["text"] for edit in driver.platform.edits ) - driver.session_store.flush_pending_save() - persisted = driver.session_store.load_conversation_snapshot() + persisted = await driver.session_store.load_conversation_snapshot() root_identity = TreeIdentity(scope=root.scope, root_id=root.message_id) assert persisted.trees[root_identity].nodes["survivor"]["state"] == "completed" @pytest.mark.asyncio -async def test_restart_restore_and_session_persistence_e2e(tmp_path) -> None: - first = FakePlatformDriver("telegram", tmp_path) +async def test_restart_restore_and_session_persistence_e2e( + driver_factory, tmp_path +) -> None: + first = await driver_factory("telegram", tmp_path) root = await first.send("persist me", message_id="root_1") - first.session_store.flush_pending_save() - session_file = tmp_path / "telegram-sessions.json" - payload = json.loads(session_file.read_text(encoding="utf-8")) - assert payload["conversation"]["trees"] - assert payload["managed_messages"] + saved_before = await first.session_store.load_conversation_snapshot() + assert saved_before.trees + assert await first.session_store.get_tracked_message_ids_for_chat( + "telegram", "chat_1" + ) + await first.close() - restored = FakePlatformDriver("telegram", tmp_path) + restored = await driver_factory("telegram", tmp_path) restored.platform.continue_message_sequence_after(first.platform) - restored.workflow.restore() - saved = restored.session_store.load_conversation_snapshot() + await restored.workflow.restore() + saved = await restored.session_store.load_conversation_snapshot() assert saved.trees identity = TreeIdentity(scope=root.scope, root_id=root.message_id) assert saved.get_tree(identity) is not None @@ -436,10 +447,11 @@ async def test_restart_restore_and_session_persistence_e2e(tmp_path) -> None: @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) async def test_same_message_ids_are_isolated_by_chat_e2e( + driver_factory, platform_name: str, tmp_path, ) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) + driver = await driver_factory(platform_name, tmp_path) first = await driver.send("first chat", chat_id="chat_a", message_id="42") second = await driver.send("second chat", chat_id="chat_b", message_id="42") @@ -462,8 +474,10 @@ async def test_same_message_ids_are_isolated_by_chat_e2e( @pytest.mark.asyncio @pytest.mark.parametrize("platform_name", ["discord", "telegram"]) -async def test_voice_platform_fake_e2e(platform_name: str, tmp_path) -> None: - driver = FakePlatformDriver(platform_name, tmp_path) +async def test_voice_platform_fake_e2e( + driver_factory, platform_name: str, tmp_path +) -> None: + driver = await driver_factory(platform_name, tmp_path) driver.platform.seed_pending_voice("chat_1", "voice_msg_1", "voice_status_1") await driver.send("/clear", message_id="clear_voice", reply_to="voice_msg_1") @@ -482,3 +496,21 @@ async def test_voice_platform_fake_e2e(platform_name: str, tmp_path) -> None: edit["message_id"] == "voice_status_2" and "Stopped" in edit["text"] for edit in driver.platform.edits ) + + +@pytest_asyncio.fixture +async def driver_factory(): + drivers = [] + + async def create(*args, **kwargs): + driver = FakePlatformDriver(*args, **kwargs) + await driver.start() + drivers.append(driver) + return driver + + try: + yield create + finally: + for driver in reversed(drivers): + if driver.database._started: + await driver.close() diff --git a/src/free_claude_code/api/admin_static/admin.js b/src/free_claude_code/api/admin_static/admin.js index e7281f087a..74dc9a1d5c 100644 --- a/src/free_claude_code/api/admin_static/admin.js +++ b/src/free_claude_code/api/admin_static/admin.js @@ -123,12 +123,12 @@ function renderStartup() { const message = byId("startupMessage"); if (message) { const messaging = startup?.messaging; - const visible = state.activeView === "messaging" && ["starting", "failed"].includes(messaging?.state); + const visible = state.activeView === "messaging" && (["starting", "failed"].includes(messaging?.state) || Boolean(messaging?.warning)); message.hidden = !visible; message.classList.toggle("startup-spinner", visible && messaging.state === "starting"); message.classList.toggle("error", messaging?.state === "failed"); - message.textContent = messaging?.state === "failed" ? messaging.message || "Messaging could not start." : ""; - message.setAttribute("aria-label", "Messaging is starting"); + message.textContent = messaging?.state === "failed" ? messaging.message || "Messaging could not start." : messaging?.warning || ""; + message.setAttribute("aria-label", messaging?.state === "starting" ? "Messaging is starting" : "Messaging status"); } if (!startup) return; document.querySelectorAll("[data-startup-provider]").forEach((button) => { diff --git a/src/free_claude_code/messaging/command_context.py b/src/free_claude_code/messaging/command_context.py index 5a46f52191..8ad0c4a8a2 100644 --- a/src/free_claude_code/messaging/command_context.py +++ b/src/free_claude_code/messaging/command_context.py @@ -76,7 +76,7 @@ async def clear_chat(self, platform: str, chat_id: str) -> frozenset[str]: """Clear one chat and return every tracked platform message ID.""" ... - def forget_tracked_message_ids( + async def forget_tracked_message_ids( self, platform: str, chat_id: str, @@ -85,7 +85,7 @@ def forget_tracked_message_ids( """Forget platform message IDs removed from the managed conversation.""" ... - def record_outgoing_message( + async def record_outgoing_message( self, platform: str, chat_id: str, diff --git a/src/free_claude_code/messaging/commands.py b/src/free_claude_code/messaging/commands.py index fd042ca36c..72abeef836 100644 --- a/src/free_claude_code/messaging/commands.py +++ b/src/free_claude_code/messaging/commands.py @@ -21,7 +21,7 @@ async def _send_stop_feedback( fire_and_forget=False, message_thread_id=incoming.message_thread_id, ) - handler.record_outgoing_message( + await handler.record_outgoing_message( incoming.platform, incoming.chat_id, msg_id, "command" ) @@ -85,7 +85,7 @@ async def handle_stats_command( fire_and_forget=False, message_thread_id=incoming.message_thread_id, ) - handler.record_outgoing_message( + await handler.record_outgoing_message( incoming.platform, incoming.chat_id, msg_id, "command" ) @@ -157,7 +157,7 @@ async def handle_clear_command( fire_and_forget=False, message_thread_id=incoming.message_thread_id, ) - handler.record_outgoing_message( + await handler.record_outgoing_message( incoming.platform, incoming.chat_id, msg_id, "command" ) return @@ -166,7 +166,7 @@ async def handle_clear_command( if incoming.message_id is not None: delete_message_ids.add(str(incoming.message_id)) await _delete_message_ids(handler, incoming.chat_id, delete_message_ids) - handler.forget_tracked_message_ids( + await handler.forget_tracked_message_ids( incoming.platform, incoming.chat_id, delete_message_ids, diff --git a/src/free_claude_code/messaging/node_event_pipeline.py b/src/free_claude_code/messaging/node_event_pipeline.py index dc7d4d5ce1..e8f13b1e11 100644 --- a/src/free_claude_code/messaging/node_event_pipeline.py +++ b/src/free_claude_code/messaging/node_event_pipeline.py @@ -94,8 +94,8 @@ async def process_parsed_cli_event( node_id=claim.node.node_id, claude_session_id=captured_session_id, ) - await update_ui(format_status("โœ…", "Complete"), force=True) await complete_claim(captured_session_id) + await update_ui(format_status("โœ…", "Complete"), force=True) elif ptype == "error": error_msg = parsed.get("message", "Unknown error") em = error_msg if isinstance(error_msg, str) else str(error_msg) @@ -114,7 +114,7 @@ async def process_parsed_cli_event( "HANDLER: Error event received: message_chars={}", text_len_hint(em), ) - await update_ui(format_status("โŒ", "Error"), force=True) await fail_claim(em, "Parent task failed") + await update_ui(format_status("โŒ", "Error"), force=True) return last_status, had_transcript_events diff --git a/src/free_claude_code/messaging/node_runner.py b/src/free_claude_code/messaging/node_runner.py index ad4bfa3f3f..ee288e53fb 100644 --- a/src/free_claude_code/messaging/node_runner.py +++ b/src/free_claude_code/messaging/node_runner.py @@ -16,9 +16,13 @@ from .node_event_pipeline import handle_session_info_event, process_parsed_cli_event from .platforms.ports import OutboundMessenger from .safe_diagnostics import format_exception_for_log -from .session import SessionStore from .transcript import RenderCtx, TranscriptBuffer -from .trees import CancellationReason, NodeClaim, TreeQueueManager, TreeSnapshot +from .trees import ( + CancellationReason, + MessagingStorageError, + NodeClaim, + TreeQueueManager, +) from .ui_updates import ThrottledTranscriptEditor @@ -31,7 +35,6 @@ def __init__( platform_name: str, outbound: OutboundMessenger, cli_manager: ManagedClaudeSessionManagerProtocol, - session_store: SessionStore, get_tree_queue: Callable[[], TreeQueueManager], format_status: Callable[[str, str, str | None], str], get_parse_mode: Callable[[], str | None], @@ -45,7 +48,6 @@ def __init__( self.platform_name = platform_name self.outbound = outbound self.cli_manager = cli_manager - self.session_store = session_store self._get_tree_queue = get_tree_queue self._format_status = format_status self._get_parse_mode = get_parse_mode @@ -66,23 +68,15 @@ def _create_transcript_and_render_ctx( ) return transcript, self._get_render_ctx() - def _save_snapshot(self, snapshot: TreeSnapshot | None) -> None: - """Persist a snapshot returned by the active aggregate manager.""" - if snapshot is None: - return - self.session_store.save_tree_snapshot(snapshot) - async def _record_session(self, claim: NodeClaim, session_id: str) -> None: - snapshot = await self._get_tree_queue().record_session(claim, session_id) - self._save_snapshot(snapshot) + await self._get_tree_queue().record_session(claim, session_id) async def _complete_claim( self, claim: NodeClaim, session_id: str | None, ) -> None: - snapshot = await self._get_tree_queue().complete_claim(claim, session_id) - self._save_snapshot(snapshot) + await self._get_tree_queue().complete_claim(claim, session_id) async def _fail_claim( self, @@ -95,7 +89,6 @@ async def _fail_claim( claim, propagate=propagate, ) - self._save_snapshot(result.snapshot) if child_status_text is None: return for child in result.affected: @@ -358,6 +351,9 @@ async def fail_parsed_event( claim, propagate=False, ) + except MessagingStorageError: + # The processor owns interruption reporting, including queued turns. + raise except Exception as e: trace_event( stage="claude_cli", diff --git a/src/free_claude_code/messaging/session/__init__.py b/src/free_claude_code/messaging/session/__init__.py deleted file mode 100644 index cdf1f36759..0000000000 --- a/src/free_claude_code/messaging/session/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Public messaging session persistence API.""" - -from .store import SessionStore - -__all__ = ["SessionStore"] diff --git a/src/free_claude_code/messaging/session/managed_message_log.py b/src/free_claude_code/messaging/session/managed_message_log.py deleted file mode 100644 index c62e2adc07..0000000000 --- a/src/free_claude_code/messaging/session/managed_message_log.py +++ /dev/null @@ -1,135 +0,0 @@ -"""Persist platform messages belonging to FCC-managed conversations.""" - -from datetime import UTC, datetime -from typing import Any - - -class ManagedMessageLog: - """Track managed inbound and outbound messages in insertion order.""" - - def __init__(self, *, cap: int | None = None) -> None: - self._items: dict[str, list[dict[str, Any]]] = {} - self._ids: dict[str, set[str]] = {} - self._cap = cap - - @property - def cap(self) -> int | None: - return self._cap - - @classmethod - def from_json(cls, raw_log: Any, *, cap: int | None = None) -> ManagedMessageLog: - """Load current and legacy message-log entries.""" - log = cls(cap=cap) - if not isinstance(raw_log, dict): - return log - for chat_key, items in raw_log.items(): - if not isinstance(chat_key, str) or not isinstance(items, list): - continue - for item in items: - if not isinstance(item, dict): - continue - message_id = item.get("message_id") - direction = str(item.get("direction") or "") - kind = str(item.get("kind") or "") - if message_id is None or direction not in {"in", "out"} or not kind: - continue - log._append( - chat_key, - str(message_id), - ts=str(item.get("ts") or ""), - direction=direction, - kind=kind, - ) - return log - - def to_json(self) -> dict[str, list[dict[str, Any]]]: - return {chat_key: list(items) for chat_key, items in self._items.items()} - - def record( - self, - *, - platform: str, - chat_id: str, - message_id: str, - direction: str, - kind: str, - ) -> bool: - """Record one managed platform message.""" - if direction not in {"in", "out"}: - raise ValueError("Managed message direction must be 'in' or 'out'") - if not kind: - raise ValueError("Managed message kind cannot be empty") - return self._append( - make_chat_key(platform, chat_id), - str(message_id), - ts=datetime.now(UTC).isoformat(), - direction=direction, - kind=str(kind), - ) - - def ids_for_chat(self, platform: str, chat_id: str) -> list[str]: - chat_key = make_chat_key(platform, chat_id) - return [str(item["message_id"]) for item in self._items.get(chat_key, [])] - - def remove_ids(self, platform: str, chat_id: str, message_ids: set[str]) -> bool: - chat_key = make_chat_key(platform, chat_id) - if not message_ids or chat_key not in self._items: - return False - - before_count = len(self._items[chat_key]) - retained = [ - item - for item in self._items[chat_key] - if str(item["message_id"]) not in message_ids - ] - if retained: - self._items[chat_key] = retained - self._ids[chat_key] = {str(item["message_id"]) for item in retained} - else: - self._items.pop(chat_key) - self._ids.pop(chat_key, None) - return len(retained) != before_count - - def clear_chat(self, platform: str, chat_id: str) -> bool: - chat_key = make_chat_key(platform, chat_id) - removed = self._items.pop(chat_key, None) is not None - self._ids.pop(chat_key, None) - return removed - - def _append( - self, - chat_key: str, - message_id: str, - *, - ts: str, - direction: str, - kind: str, - ) -> bool: - seen = self._ids.setdefault(chat_key, set()) - if message_id in seen: - return False - self._items.setdefault(chat_key, []).append( - { - "message_id": message_id, - "ts": ts, - "direction": direction, - "kind": kind, - } - ) - seen.add(message_id) - self._trim(chat_key) - return True - - def _trim(self, chat_key: str) -> None: - if self._cap is None or self._cap <= 0: - return - items = self._items.get(chat_key, []) - if len(items) <= self._cap: - return - retained = items[-self._cap :] - self._items[chat_key] = retained - self._ids[chat_key] = {str(item["message_id"]) for item in retained} - - -def make_chat_key(platform: str, chat_id: str) -> str: - return f"{platform}:{chat_id}" diff --git a/src/free_claude_code/messaging/session/persistence.py b/src/free_claude_code/messaging/session/persistence.py deleted file mode 100644 index 7849e6259f..0000000000 --- a/src/free_claude_code/messaging/session/persistence.py +++ /dev/null @@ -1,161 +0,0 @@ -"""Atomic JSON persistence for messaging session state.""" - -import contextlib -import json -import os -import tempfile -import threading -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -from loguru import logger - - -@dataclass(frozen=True) -class _PendingWrite: - generation: int - snapshot: dict[str, Any] - - -class DebouncedJsonPersistence: - """Thread-safe debounced JSON writer with atomic replace semantics.""" - - def __init__( - self, - storage_path: str, - *, - snapshot: Callable[[], dict[str, Any]], - on_dirty: Callable[[bool], None], - debounce_secs: float = 0.5, - ) -> None: - self.storage_path = storage_path - self._snapshot = snapshot - self._on_dirty = on_dirty - self._debounce_secs = debounce_secs - self._save_timer: threading.Timer | None = None - self._timer_lock = threading.Lock() - self._writer_lock = threading.Lock() - self._save_generation = 0 - - def load_json(self) -> dict[str, Any]: - if not os.path.exists(self.storage_path): - return {} - with open(self.storage_path, encoding="utf-8") as file: - data = json.load(file) - return data if isinstance(data, dict) else {} - - def schedule_save(self) -> None: - self._on_dirty(True) - with self._timer_lock: - if self._save_timer is not None: - self._save_timer.cancel() - self._save_generation += 1 - generation = self._save_generation - timer = threading.Timer( - self._debounce_secs, - self._save_from_timer, - args=(generation,), - ) - timer.daemon = True - self._save_timer = timer - timer.start() - - def flush(self) -> None: - self._on_dirty(True) - pending = self._snapshot_for_write() - if pending is None: - return - self._write_pending(pending) - - def _save_from_timer(self, generation: int) -> None: - try: - pending = self._snapshot_for_write(expected_generation=generation) - if pending is None: - return - self._write_pending(pending) - except Exception as e: - self._on_dirty(True) - logger.error( - "Failed to save sessions: exc_type={}", - type(e).__name__, - ) - - def _write_pending(self, pending: _PendingWrite) -> None: - try: - written = self._write_if_current(pending) - except Exception: - self._on_dirty(True) - raise - if written: - self._mark_clean_if_current(pending.generation) - - def _write_if_current(self, pending: _PendingWrite) -> bool: - """Serialize writers and reject a snapshot superseded before replace.""" - with self._writer_lock: - with self._timer_lock: - if pending.generation != self._save_generation: - return False - self._write_file(pending.snapshot) - return True - - def _snapshot_for_write( - self, *, expected_generation: int | None = None - ) -> _PendingWrite | None: - generation = self._claim_timer(expected_generation) - if generation is None: - return None - snapshot = self._snapshot() - return _PendingWrite(generation=generation, snapshot=snapshot) - - def _claim_timer(self, expected_generation: int | None) -> int | None: - with self._timer_lock: - if expected_generation is not None and ( - expected_generation != self._save_generation or self._save_timer is None - ): - return None - if self._save_timer is not None: - self._save_timer.cancel() - self._save_timer = None - return self._save_generation - - def _mark_clean_if_current(self, generation: int) -> None: - with self._timer_lock: - is_current = ( - self._save_timer is None and generation == self._save_generation - ) - if is_current: - self._on_dirty(False) - - def write_data(self, data: dict[str, Any]) -> None: - """Write authoritative state after invalidating older pending snapshots.""" - self._on_dirty(True) - with self._timer_lock: - if self._save_timer is not None: - self._save_timer.cancel() - self._save_timer = None - self._save_generation += 1 - pending = _PendingWrite( - generation=self._save_generation, - snapshot=data, - ) - self._write_pending(pending) - - def _write_file(self, data: dict[str, Any]) -> None: - abs_target = os.path.abspath(self.storage_path) - dir_name = os.path.dirname(abs_target) or "." - fd, tmp_path = tempfile.mkstemp( - dir=dir_name, - prefix=".sessions.", - suffix=".tmp.json", - ) - try: - with os.fdopen(fd, "w", encoding="utf-8") as file: - json.dump(data, file, indent=2) - file.flush() - os.fsync(file.fileno()) - os.replace(tmp_path, abs_target) - except BaseException: - with contextlib.suppress(OSError): - os.unlink(tmp_path) - raise diff --git a/src/free_claude_code/messaging/session/store.py b/src/free_claude_code/messaging/session/store.py deleted file mode 100644 index 674aa720b5..0000000000 --- a/src/free_claude_code/messaging/session/store.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Persistent messaging conversation state store.""" - -import threading -from copy import deepcopy - -from loguru import logger - -from free_claude_code.messaging.models import MessageScope -from free_claude_code.messaging.trees import ( - ConversationSnapshot, - TreeIdentity, - TreeSnapshot, -) - -from .managed_message_log import ManagedMessageLog -from .persistence import DebouncedJsonPersistence - - -class SessionStore: - """ - Persistent storage for conversation snapshots and managed platform messages. - - The store reads both the old raw ``trees``/``node_to_tree`` shape and the - current typed ``conversation`` snapshot shape. Runtime callers deal in typed - snapshots only. - """ - - def __init__( - self, - storage_path: str = "sessions.json", - *, - managed_message_cap: int | None = None, - ) -> None: - self.storage_path = storage_path - self._lock = threading.RLock() - self._conversation = ConversationSnapshot() - self._managed_messages = ManagedMessageLog(cap=managed_message_cap) - self._dirty = False - self._persistence = DebouncedJsonPersistence( - storage_path, - snapshot=self._snapshot_for_persistence, - on_dirty=self._set_dirty, - ) - self._load() - - @property - def dirty(self) -> bool: - return self._dirty - - def _set_dirty(self, dirty: bool) -> None: - with self._lock: - self._dirty = dirty - - def _load(self) -> None: - try: - data = self._persistence.load_json() - except Exception as e: - logger.error("Failed to load sessions: {}", e) - return - - conversation_data = data.get("conversation") if isinstance(data, dict) else None - if not isinstance(conversation_data, dict): - conversation_data = data - - with self._lock: - self._conversation = ConversationSnapshot.from_json(conversation_data) - raw_messages = {} - if isinstance(data, dict): - raw_messages = data.get("managed_messages", data.get("message_log", {})) - self._managed_messages = ManagedMessageLog.from_json( - raw_messages, - cap=self._managed_messages.cap, - ) - message_count = sum( - len(items) for items in self._managed_messages.to_json().values() - ) - logger.info( - "Loaded {} trees and {} managed message IDs from {}", - len(self._conversation.trees), - message_count, - self.storage_path, - ) - - def _snapshot_for_persistence(self) -> dict: - with self._lock: - return { - "conversation": self._conversation.to_json(), - "managed_messages": self._managed_messages.to_json(), - } - - def load_conversation_snapshot(self) -> ConversationSnapshot: - with self._lock: - return deepcopy(self._conversation) - - def save_conversation_snapshot(self, snapshot: ConversationSnapshot) -> None: - with self._lock: - self._conversation = deepcopy(snapshot) - self._persistence.schedule_save() - - def save_tree_snapshot(self, snapshot: TreeSnapshot) -> None: - with self._lock: - self._conversation = self._conversation.with_tree(deepcopy(snapshot)) - self._persistence.schedule_save() - logger.debug("Saved tree {}", snapshot.root_id) - - def remove_tree_snapshot(self, identity: TreeIdentity) -> None: - with self._lock: - self._conversation = self._conversation.without_tree(identity) - self._persistence.schedule_save() - - def flush_pending_save(self) -> None: - self._persistence.flush() - - def record_message_id( - self, - platform: str, - chat_id: str, - message_id: str, - direction: str, - kind: str, - ) -> None: - if message_id is None: - return - with self._lock: - recorded = self._managed_messages.record( - platform=str(platform), - chat_id=str(chat_id), - message_id=str(message_id), - direction=str(direction), - kind=str(kind), - ) - if recorded: - self._persistence.schedule_save() - - def get_tracked_message_ids_for_chat( - self, platform: str, chat_id: str - ) -> list[str]: - with self._lock: - return self._managed_messages.ids_for_chat(str(platform), str(chat_id)) - - def forget_tracked_message_ids( - self, platform: str, chat_id: str, message_ids: set[str] - ) -> None: - with self._lock: - removed = self._managed_messages.remove_ids( - str(platform), - str(chat_id), - {str(message_id) for message_id in message_ids}, - ) - if removed: - self._persistence.schedule_save() - - def clear_scope(self, scope: MessageScope) -> None: - """Authoritatively clear one platform chat while preserving others.""" - with self._lock: - self._conversation = self._conversation.without_scope(scope) - self._managed_messages.clear_chat(scope.platform, scope.chat_id) - self._write_current_state() - - def _write_current_state(self) -> None: - self._set_dirty(True) - self._persistence.write_data(self._snapshot_for_persistence()) diff --git a/src/free_claude_code/messaging/trees/__init__.py b/src/free_claude_code/messaging/trees/__init__.py index 318776f40a..0cd13146f2 100644 --- a/src/free_claude_code/messaging/trees/__init__.py +++ b/src/free_claude_code/messaging/trees/__init__.py @@ -1,8 +1,10 @@ """Internal messaging tree package facade.""" +from .graph import normalize_tree_snapshot from .identity import TreeIdentity from .manager import TreeQueueManager from .node import MessageReferenceKind, MessageState +from .ports import MessagingStorageError, MessagingStore from .snapshot import ConversationSnapshot, TreeSnapshot from .transitions import ( AdmissionRejection, @@ -29,6 +31,8 @@ "MessageReferenceKind", "MessageState", "MessageSubtreeRemovalResult", + "MessagingStorageError", + "MessagingStore", "NodeClaim", "NodeUiTarget", "NodeView", @@ -38,4 +42,5 @@ "TreeIdentity", "TreeQueueManager", "TreeSnapshot", + "normalize_tree_snapshot", ] diff --git a/src/free_claude_code/messaging/trees/graph.py b/src/free_claude_code/messaging/trees/graph.py index d120df08f2..a4b7349799 100644 --- a/src/free_claude_code/messaging/trees/graph.py +++ b/src/free_claude_code/messaging/trees/graph.py @@ -248,3 +248,17 @@ def from_snapshot(cls, snapshot: TreeSnapshot) -> MessageTreeGraph: if set(graph.get_descendants(graph.root_id)) != set(graph._nodes): raise ValueError("Tree snapshot contains a disconnected branch") return graph + + +def normalize_tree_snapshot(snapshot: TreeSnapshot) -> TreeSnapshot: + """Validate domain links and return detached records in parent-first order.""" + graph = MessageTreeGraph.from_snapshot(snapshot) + normalized = graph.snapshot() + return TreeSnapshot( + normalized.scope, + normalized.root_id, + { + node_id: normalized.nodes[node_id] + for node_id in graph.get_descendants(graph.root_id) + }, + ) diff --git a/src/free_claude_code/messaging/trees/manager.py b/src/free_claude_code/messaging/trees/manager.py index 472055aed1..d5789a4abb 100644 --- a/src/free_claude_code/messaging/trees/manager.py +++ b/src/free_claude_code/messaging/trees/manager.py @@ -1,14 +1,19 @@ """Public facade for atomic messaging tree aggregates.""" import asyncio -from collections.abc import Callable, Coroutine -from typing import Any +from collections.abc import AsyncIterator, Callable, Coroutine +from contextlib import asynccontextmanager +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .ports import MessagingStore from loguru import logger from ..models import IncomingMessage, MessageScope from .identity import TreeIdentity from .node import MessageNode, MessageState +from .ports import MessagingStorageError from .processor import ( CancelledTask, NodeProcessor, @@ -19,6 +24,7 @@ from .repository import TreeRepository from .runtime import MessageTree from .snapshot import ConversationSnapshot, TreeSnapshot +from .transaction import TreeTransaction from .transitions import ( AdmissionRejection, CancellationEffect, @@ -84,6 +90,7 @@ def __init__( self, node_processor: NodeProcessor, *, + store: MessagingStore, queue_update_callback: QueueUpdateCallback | None = None, node_started_callback: NodeStartedCallback | None = None, unexpected_failure_callback: Callable[[FailureResult], None] | None = None, @@ -92,8 +99,11 @@ def __init__( _restored_snapshot: ConversationSnapshot | None = None, _restored_stale_targets: tuple[NodeUiTarget, ...] = (), ) -> None: + self._store = store self._repository = _repository or TreeRepository() self._lock = asyncio.Lock() + self._interrupted: dict[MessageTree, tuple[NodeUiTarget, ...]] = {} + self._closing = False self._processor = TreeQueueProcessor( node_processor, claim_failure_callback=self._handle_processor_failure, @@ -107,6 +117,43 @@ def __init__( self._unexpected_failure_callback = unexpected_failure_callback logger.info("TreeQueueManager initialized") + @asynccontextmanager + async def _transaction( + self, *, worker: NodeClaim | None = None, allow_closing: bool = False + ) -> AsyncIterator[TreeTransaction]: + async with self._lock: + if self._closing and worker is None and not allow_closing: + raise RuntimeError("Messaging is shutting down") + transaction = TreeTransaction(self._repository, self._store) + try: + yield transaction + await transaction.commit() + except BaseException as exc: + await transaction.rollback() + if isinstance(exc, MessagingStorageError) and worker is not None: + tree = self._repository.get_tree(worker.identity) + if tree is not None and await tree.owns_claim(worker.claim_id): + self._interrupted.setdefault( + tree, await tree.runnable_targets() + ) + raise + else: + for tree in transaction.release_interruptions: + self._interrupted.pop(tree, None) + + async def _prepare_admission( + self, transaction: TreeTransaction, tree: MessageTree + ) -> None: + transaction.watch(tree) + if tree not in self._interrupted: + return + if self._processor.owns_tree(tree): + raise MessagingStorageError( + "This conversation is cleaning up after a storage failure. Try again shortly." + ) + await tree.cancel_all() + transaction.release_interruptions.add(tree) + @property def restored_snapshot(self) -> ConversationSnapshot | None: return self._restored_snapshot @@ -126,7 +173,7 @@ async def admit( node_id = str(incoming.message_id) scope = incoming.scope prompt = incoming.text or "" - async with self._lock: + async with self._transaction() as transaction: duplicate_tree = self._repository.get_tree_for_reference(scope, node_id) if duplicate_tree is None: duplicate_tree = self._repository.get_tree_for_reference( @@ -134,6 +181,7 @@ async def admit( status_message_id, ) if duplicate_tree is not None: + await self._prepare_admission(transaction, duplicate_tree) return await duplicate_tree.enqueue_or_claim(node_id) tree: MessageTree | None = None @@ -154,6 +202,7 @@ async def admit( ) if tree is not None and resolved_parent is not None: + await self._prepare_admission(transaction, tree) decision = await tree.add_and_enqueue( node_id, scope, @@ -177,6 +226,7 @@ async def admit( state=MessageState.PENDING, ) tree = MessageTree(root) + transaction.inserted(tree.identity) decision = await tree.enqueue_or_claim(node_id) self._repository.add_tree( tree, @@ -225,8 +275,11 @@ async def record_session( claim: NodeClaim, session_id: str, ) -> TreeSnapshot | None: - async with self._lock: + async with self._transaction(worker=claim) as transaction: tree = self._repository.get_tree(claim.identity) + if self._closing or tree in self._interrupted: + return None + transaction.watch(tree) return ( await tree.record_session(claim.claim_id, session_id) if tree is not None @@ -238,8 +291,11 @@ async def complete_claim( claim: NodeClaim, session_id: str | None, ) -> TreeSnapshot | None: - async with self._lock: + async with self._transaction(worker=claim) as transaction: tree = self._repository.get_tree(claim.identity) + if self._closing or tree in self._interrupted: + return None + transaction.watch(tree) return ( await tree.complete_claim(claim.claim_id, session_id) if tree is not None @@ -260,8 +316,11 @@ async def _fail_claim( *, propagate: bool, ) -> FailureResult: - async with self._lock: + async with self._transaction(worker=claim) as transaction: tree = self._repository.get_tree(claim.identity) + if self._closing or tree in self._interrupted: + return FailureResult(affected=(), queue_update=None, snapshot=None) + transaction.watch(tree) result = ( await tree.fail_claim( claim.claim_id, @@ -280,6 +339,9 @@ async def _handle_processor_failure( ) -> None: """Route an escaped runner failure through manager-owned effects.""" result = await self._fail_claim(claim, propagate=True) + self._notify_failure(result) + + def _notify_failure(self, result: FailureResult) -> None: if self._unexpected_failure_callback is None: return try: @@ -292,16 +354,58 @@ async def _handle_processor_failure( async def _finish_claim(self, tree: MessageTree, claim: NodeClaim) -> None: """Serialize successor task publication with aggregate detachment.""" - async with self._lock: - tree_is_published = self._repository.get_tree(claim.identity) is tree - completion = await tree.finish_and_claim_next(claim.claim_id) - if tree_is_published and completion.next_claim is not None: - self._processor.launch( - tree, - completion.next_claim, - announce_started=True, - queue=completion.queue, + try: + async with self._transaction(worker=claim) as transaction: + if self._closing or tree in self._interrupted: + targets = await self._retire_interrupted(tree, claim) + completion = None + else: + tree_is_published = ( + self._repository.get_tree(claim.identity) is tree + ) + if tree_is_published: + transaction.watch(tree) + completion = await tree.finish_and_claim_next(claim.claim_id) + targets = () + except MessagingStorageError: + async with self._lock: + targets = await self._retire_interrupted(tree, claim) + completion = None + + if targets: + self._notify_failure( + FailureResult( + affected=targets, + queue_update=(), + snapshot=None, + message="Messaging history could not be saved. This conversation was interrupted.", ) + ) + + if ( + completion is not None + and tree_is_published + and completion.next_claim is not None + ): + self._processor.launch( + tree, + completion.next_claim, + announce_started=True, + queue=completion.queue, + ) + + async def _retire_interrupted( + self, tree: MessageTree, claim: NodeClaim + ) -> tuple[NodeUiTarget, ...]: + """Called under the manager lock, after any intervening stop or clear.""" + await tree.retire_claim(claim.claim_id) + current = set(await tree.runnable_targets()) + targets = tuple( + target for target in self._interrupted.get(tree, ()) if target in current + ) + if tree in self._interrupted: + self._interrupted[tree] = () + return targets @staticmethod def _external_effects( @@ -353,19 +457,20 @@ async def _cancel_node( node_id: str, reason: CancellationReason | None, ) -> CancellationResult: - async with self._lock: + async with self._transaction() as transaction: tree = self._repository.get_tree_for_reference(scope, node_id) if tree is None: return CancellationResult() target = await tree.resolve_reply(node_id) if target is None: return CancellationResult() + transaction.watch(tree) transition = await tree.cancel_node(target.node_id) - cancelled_task = ( - self._processor.cancel(transition.active_claim, reason) - if transition.active_claim is not None - else None - ) + cancelled_task = ( + self._processor.cancel(transition.active_claim, reason) + if transition.active_claim is not None + else None + ) if transition.queue_update is not None: await self._processor.notify_queue_updated(transition.queue_update) @@ -381,85 +486,72 @@ async def cancel_all( self, *, reason: CancellationReason | None = None, + on_committed: Callable[[], None] | None = None, ) -> CancellationResult: - return await _finish_transition(self._cancel_all(reason)) + """Publish the synchronous commit receipt before cancellation effects.""" + return await _finish_transition(self._cancel_all(reason, on_committed)) async def _cancel_all( - self, - reason: CancellationReason | None, + self, reason: CancellationReason | None, on_committed: Callable[[], None] | None ) -> CancellationResult: - transitions: list[ - tuple[MessageTree, TreeCancellation, CancelledTask | None] - ] = [] - async with self._lock: + transitions: list[TreeCancellation] = [] + async with self._transaction() as transaction: for tree in self._repository.trees(): - transition = await tree.cancel_all() - cancelled_task = ( - self._processor.cancel(transition.active_claim, reason) - if transition.active_claim is not None - else None - ) - transitions.append((tree, transition, cancelled_task)) - - for _tree, transition, _task in transitions: + transaction.watch(tree) + transitions.append(await tree.cancel_all()) + if on_committed is not None: + on_committed() + return await self._cancel_committed(transitions, reason) + + async def _cancel_committed( + self, transitions: list[TreeCancellation], reason: CancellationReason | None + ) -> CancellationResult: + effects: list[CancellationEffect] = [] + tasks: list[asyncio.Task[None]] = [] + for transition in transitions: + cancelled = ( + self._processor.cancel(transition.active_claim, reason) + if transition.active_claim is not None + else None + ) + effects.extend(self._external_effects(transition, cancelled)) + if cancelled is not None: + tasks.append(cancelled.task) + for transition in transitions: if transition.queue_update is not None: await self._processor.notify_queue_updated(transition.queue_update) - await _drain_cancelled_tasks( - [ - cancelled.task - for _tree, _transition, cancelled in transitions - if cancelled is not None - ] - ) - - effects: list[CancellationEffect] = [] - snapshots: list[TreeSnapshot] = [] - for tree, transition, cancelled_task in transitions: - effects.extend(self._external_effects(transition, cancelled_task)) - snapshot = await self._current_snapshot(tree.identity, tree) - if snapshot is not None: - snapshots.append(snapshot) - return CancellationResult(effects=tuple(effects), snapshots=tuple(snapshots)) + await _drain_cancelled_tasks(tasks) + return CancellationResult(effects=tuple(effects)) async def clear_scope( self, scope: MessageScope, *, reason: CancellationReason | None = None, + on_committed: Callable[[], None] | None = None, ) -> CancellationResult: - """Atomically detach one chat's trees, then drain their active tasks.""" - return await _finish_transition(self._clear_scope(scope, reason)) + """Commit detachment and its receipt before draining active tasks.""" + return await _finish_transition(self._clear_scope(scope, reason, on_committed)) async def _clear_scope( self, scope: MessageScope, reason: CancellationReason | None, + on_committed: Callable[[], None] | None, ) -> CancellationResult: - transitions: list[tuple[TreeCancellation, CancelledTask | None]] = [] - async with self._lock: + transitions: list[TreeCancellation] = [] + async with self._transaction() as transaction: + transaction.clear_scope = scope for tree in self._repository.trees(): if tree.identity.scope != scope: continue - transition = await tree.cancel_all() - cancelled_task = ( - self._processor.cancel(transition.active_claim, reason) - if transition.active_claim is not None - else None - ) - transitions.append((transition, cancelled_task)) + transaction.watch(tree) + transitions.append(await tree.cancel_all()) self._repository.remove_tree(tree.identity) - - await _drain_cancelled_tasks( - [ - cancelled.task - for _transition, cancelled in transitions - if cancelled is not None - ] - ) - effects: list[CancellationEffect] = [] - for transition, cancelled_task in transitions: - effects.extend(self._external_effects(transition, cancelled_task)) - return CancellationResult(effects=tuple(effects)) + transaction.release_interruptions.add(tree) + if on_committed is not None: + on_committed() + return await self._cancel_committed(transitions, reason) async def remove_message_subtree( self, @@ -479,7 +571,7 @@ async def _remove_message_subtree( reference_id: str, reason: CancellationReason | None, ) -> MessageSubtreeRemovalResult: - async with self._lock: + async with self._transaction() as transaction: tree = self._repository.get_tree_for_reference(scope, reference_id) if tree is None: return MessageSubtreeRemovalResult( @@ -497,17 +589,19 @@ async def _remove_message_subtree( tree_matched=False, ) identity = tree.identity + transaction.watch(tree) transition = await tree.remove_message_subtree(target.reference_id) lookup_ids = set(transition.removed_message_ids) if transition.removed_entire_tree: self._repository.remove_tree(identity) + transaction.release_interruptions.add(tree) else: self._repository.unregister_references(identity, lookup_ids) - cancelled_task = ( - self._processor.cancel(transition.cancellation.active_claim, reason) - if transition.cancellation.active_claim is not None - else None - ) + cancelled_task = ( + self._processor.cancel(transition.cancellation.active_claim, reason) + if transition.cancellation.active_claim is not None + else None + ) if transition.cancellation.queue_update is not None: await self._processor.notify_queue_updated( @@ -542,6 +636,34 @@ def get_tree_count(self) -> int: def task_count(self) -> int: return self._processor.task_count() + async def shutdown(self) -> CancellationResult: + """Stop execution even when saving its interruption fails.""" + transitions: list[TreeCancellation] = [] + try: + async with self._transaction(allow_closing=True) as transaction: + if self._closing: + return CancellationResult() + self._closing = True + for tree in self._repository.trees(): + transaction.watch(tree) + transitions.append(await tree.cancel_all()) + except MessagingStorageError: + logger.warning("Messaging shutdown could not save interrupted state") + finally: + cancelled = self._processor.cancel_all(CancellationReason.STOP) + return CancellationResult( + effects=tuple( + effect + for transition in transitions + for effect in self._external_effects( + transition, + cancelled.get(transition.active_claim.claim_id) + if transition.active_claim is not None + else None, + ) + ) + ) + async def wait_idle(self) -> None: """Wait until every processor-owned claim has finished cleanup.""" await self._processor.wait_idle() @@ -566,6 +688,7 @@ def from_snapshot( snapshot: ConversationSnapshot, node_processor: NodeProcessor, *, + store: MessagingStore, queue_update_callback: QueueUpdateCallback | None = None, node_started_callback: NodeStartedCallback | None = None, unexpected_failure_callback: Callable[[FailureResult], None] | None = None, @@ -574,6 +697,7 @@ def from_snapshot( repository, normalized, stale_targets = TreeRepository.from_snapshot(snapshot) return cls( node_processor, + store=store, queue_update_callback=queue_update_callback, node_started_callback=node_started_callback, unexpected_failure_callback=unexpected_failure_callback, diff --git a/src/free_claude_code/messaging/trees/ports.py b/src/free_claude_code/messaging/trees/ports.py new file mode 100644 index 0000000000..6a538f9ba9 --- /dev/null +++ b/src/free_claude_code/messaging/trees/ports.py @@ -0,0 +1,35 @@ +"""Durable messaging operations implemented by the runtime storage adapter.""" + +from typing import Protocol + +from ..models import MessageScope +from .identity import TreeIdentity +from .snapshot import ConversationSnapshot, TreeSnapshot + + +class MessagingStorageError(RuntimeError): + """A messaging transition could not be durably stored.""" + + +class MessagingStore(Protocol): + async def load_conversation_snapshot(self) -> ConversationSnapshot: ... + + async def commit_trees( + self, + snapshots: tuple[TreeSnapshot, ...], + *, + removed: tuple[TreeIdentity, ...] = (), + clear_scope: MessageScope | None = None, + ) -> None: ... + + async def record_message_id( + self, platform: str, chat_id: str, message_id: str, direction: str, kind: str + ) -> None: ... + + async def get_tracked_message_ids_for_chat( + self, platform: str, chat_id: str + ) -> list[str]: ... + + async def forget_tracked_message_ids( + self, platform: str, chat_id: str, message_ids: set[str] + ) -> None: ... diff --git a/src/free_claude_code/messaging/trees/processor.py b/src/free_claude_code/messaging/trees/processor.py index 11039adac6..629e31eabc 100644 --- a/src/free_claude_code/messaging/trees/processor.py +++ b/src/free_claude_code/messaging/trees/processor.py @@ -8,6 +8,7 @@ from loguru import logger from ..safe_diagnostics import format_exception_for_log +from .ports import MessagingStorageError from .runtime import MessageTree from .transitions import CancellationReason, NodeClaim, QueueEntry @@ -186,7 +187,12 @@ async def _run_claim( log_full_message=self._log_messaging_error_details, ), ) - await self._claim_failure_callback(claim) + if not isinstance(exc, MessagingStorageError): + try: + await self._claim_failure_callback(claim) + except MessagingStorageError: + # The manager has gated the tree. Finalization retires it. + logger.warning("Messaging failure state could not be saved") finally: if not slot.transitioned: await self._finish_and_continue(slot) @@ -257,6 +263,15 @@ def task_count(self) -> int: """Return the number of attached claims for observability.""" return len(self._tasks) + def owns_tree(self, tree: MessageTree) -> bool: + return any(slot.tree is tree for slot in self._tasks.values()) + + def cancel_all(self, reason: CancellationReason) -> dict[str, CancelledTask | None]: + return { + slot.claim.claim_id: self.cancel(slot.claim, reason) + for slot in tuple(self._tasks.values()) + } + async def wait_idle(self) -> None: """Wait for every task and hand completion failures to the caller once.""" await self._idle.wait() diff --git a/src/free_claude_code/messaging/trees/repository.py b/src/free_claude_code/messaging/trees/repository.py index c7e60ec4d0..e7d52072f5 100644 --- a/src/free_claude_code/messaging/trees/repository.py +++ b/src/free_claude_code/messaging/trees/repository.py @@ -90,6 +90,17 @@ def remove_tree( def trees(self) -> tuple[MessageTree, ...]: return tuple(self._trees.values()) + async def restore_tree( + self, identity: TreeIdentity, tree: MessageTree | None + ) -> None: + """Restore publication and indexes after a failed durable transition.""" + self.remove_tree(identity) + if tree is not None: + self._trees[identity] = tree + snapshot = await tree.snapshot() + for reference in snapshot.lookup_ids(): + self._reference_to_tree[(identity.scope, reference)] = identity + def tree_count(self) -> int: return len(self._trees) diff --git a/src/free_claude_code/messaging/trees/runtime.py b/src/free_claude_code/messaging/trees/runtime.py index c3d2b32af0..f460b57ea3 100644 --- a/src/free_claude_code/messaging/trees/runtime.py +++ b/src/free_claude_code/messaging/trees/runtime.py @@ -1,6 +1,7 @@ """Atomic runtime aggregate for one messaging conversation tree.""" import asyncio +from copy import deepcopy from dataclasses import dataclass from uuid import uuid4 @@ -35,6 +36,13 @@ class _ActiveClaim: cancellation_requested: bool = False +@dataclass +class TreeCheckpoint: + graph: MessageTreeGraph + queue: MessageNodeQueue + active: _ActiveClaim | None + + class MessageTree: """Own graph, queue, claim identity, and every concurrency invariant.""" @@ -56,6 +64,36 @@ def __init__( def root_id(self) -> str: return self._graph.root_id + def checkpoint(self) -> TreeCheckpoint: + """Copy only runtime data while the manager owns the transition.""" + return deepcopy(TreeCheckpoint(self._graph, self._queue, self._active)) + + def rollback(self, checkpoint: TreeCheckpoint) -> None: + self._graph, self._queue, self._active = ( + checkpoint.graph, + checkpoint.queue, + checkpoint.active, + ) + + async def retire_claim(self, claim_id: str) -> None: + """Retire stopped execution without pretending to persist an outcome.""" + async with self._lock: + if self._active is not None and self._active.claim.claim_id == claim_id: + self._active = None + self._queue.drain() + + async def runnable_targets(self) -> tuple[NodeUiTarget, ...]: + async with self._lock: + return tuple( + self._ui_target(node) + for node in self._graph.all_nodes() + if node.state in (MessageState.PENDING, MessageState.IN_PROGRESS) + ) + + async def owns_claim(self, claim_id: str) -> bool: + async with self._lock: + return self._active is not None and self._active.claim.claim_id == claim_id + @property def identity(self) -> TreeIdentity: return self._graph.identity diff --git a/src/free_claude_code/messaging/trees/transaction.py b/src/free_claude_code/messaging/trees/transaction.py new file mode 100644 index 0000000000..fa942a1e62 --- /dev/null +++ b/src/free_claude_code/messaging/trees/transaction.py @@ -0,0 +1,51 @@ +"""Rollback state for a manager-owned durable messaging transition.""" + +from typing import TYPE_CHECKING + +from ..models import MessageScope +from .identity import TreeIdentity +from .repository import TreeRepository +from .runtime import MessageTree, TreeCheckpoint + +if TYPE_CHECKING: + from .ports import MessagingStore + + +class TreeTransaction: + def __init__(self, repository: TreeRepository, store: MessagingStore) -> None: + self.repository = repository + self.store = store + self.before: dict[TreeIdentity, tuple[MessageTree, TreeCheckpoint] | None] = {} + self.clear_scope: MessageScope | None = None + self.release_interruptions: set[MessageTree] = set() + + def watch(self, tree: MessageTree | None) -> None: + if tree is not None and tree.identity not in self.before: + self.before[tree.identity] = (tree, tree.checkpoint()) + + def inserted(self, identity: TreeIdentity) -> None: + self.before[identity] = None + + async def commit(self) -> None: + changed, removed = [], [] + for identity, previous in self.before.items(): + current = self.repository.get_tree(identity) + if current is None: + removed.append(identity) + else: + snapshot = await current.snapshot() + if previous is None or snapshot != previous[1].graph.snapshot(): + changed.append(snapshot) + if changed or removed or self.clear_scope is not None: + await self.store.commit_trees( + tuple(changed), removed=tuple(removed), clear_scope=self.clear_scope + ) + + async def rollback(self) -> None: + for identity, previous in self.before.items(): + if previous is None: + await self.repository.restore_tree(identity, None) + else: + tree, checkpoint = previous + tree.rollback(checkpoint) + await self.repository.restore_tree(identity, tree) diff --git a/src/free_claude_code/messaging/trees/transitions.py b/src/free_claude_code/messaging/trees/transitions.py index 38b99016e4..ef9e99d0ce 100644 --- a/src/free_claude_code/messaging/trees/transitions.py +++ b/src/free_claude_code/messaging/trees/transitions.py @@ -152,6 +152,7 @@ class FailureResult: affected: tuple[NodeUiTarget, ...] queue_update: tuple[QueueEntry, ...] | None snapshot: TreeSnapshot | None + message: str | None = None __all__ = [ diff --git a/src/free_claude_code/messaging/turn_intake.py b/src/free_claude_code/messaging/turn_intake.py index a16b2a23dd..3f7940c65d 100644 --- a/src/free_claude_code/messaging/turn_intake.py +++ b/src/free_claude_code/messaging/turn_intake.py @@ -14,9 +14,9 @@ ) from .models import AdmissionToken, IncomingMessage, MessageScope from .platforms.ports import OutboundMessenger -from .session import SessionStore from .trees import ( AdmissionRejection, + MessagingStore, NodeClaim, QueueDecision, QueueEntry, @@ -32,7 +32,7 @@ def __init__( *, platform_name: str, outbound: OutboundMessenger, - session_store: SessionStore, + session_store: MessagingStore, command_context: MessagingCommandContext, resolve_reply: Callable[[MessageScope, str], Awaitable[ReplyTarget | None]], admit_turn: Callable[ @@ -41,7 +41,6 @@ def __init__( ], format_status: Callable[[str, str, str | None], str], get_parse_mode: Callable[[], str | None], - record_outgoing_message: Callable[[str, str, str | None, str], bool], ) -> None: self.platform_name = platform_name self.outbound = outbound @@ -51,7 +50,6 @@ def __init__( self._admit_turn = admit_turn self._format_status = format_status self._get_parse_mode = get_parse_mode - self._record_outgoing_message = record_outgoing_message async def handle_message( self, @@ -100,9 +98,6 @@ async def handle_message( fire_and_forget=False, message_thread_id=incoming.message_thread_id, ) - self._record_outgoing_message( - incoming.platform, incoming.chat_id, status_msg_id, "status" - ) if status_msg_id is None: return @@ -183,7 +178,7 @@ async def _discard_rejected_messages( type(exc).__name__, ) try: - self.session_store.forget_tracked_message_ids( + await self.session_store.forget_tracked_message_ids( incoming.platform, incoming.chat_id, message_ids, diff --git a/src/free_claude_code/messaging/workflow.py b/src/free_claude_code/messaging/workflow.py index f0200365d0..411635228f 100644 --- a/src/free_claude_code/messaging/workflow.py +++ b/src/free_claude_code/messaging/workflow.py @@ -20,7 +20,6 @@ ) from .rendering.profiles import build_rendering_profile from .safe_diagnostics import format_exception_for_log -from .session import SessionStore from .transcript import RenderCtx from .trees import ( CancellationReason, @@ -28,6 +27,7 @@ CancellationUiOwner, ConversationSnapshot, FailureResult, + MessagingStore, NodeUiTarget, QueueDecision, ReplyTarget, @@ -92,7 +92,7 @@ def __init__( self, outbound: OutboundMessenger, cli_manager: ManagedClaudeSessionManagerProtocol, - session_store: SessionStore, + session_store: MessagingStore, *, platform_name: str | None = None, voice_cancellation: VoiceCancellation | None = None, @@ -118,7 +118,6 @@ def __init__( platform_name=self.platform_name, outbound=outbound, cli_manager=cli_manager, - session_store=session_store, get_tree_queue=lambda: self._tree_queue, format_status=self.format_status, get_parse_mode=self._parse_mode, @@ -138,7 +137,6 @@ def __init__( admit_turn=self._admit_turn_if_current, format_status=self.format_status, get_parse_mode=self._parse_mode, - record_outgoing_message=self.record_outgoing_message, ) self._tree_queue = self._build_tree_queue() @@ -148,6 +146,7 @@ def _build_tree_queue( if snapshot is None: return TreeQueueManager( self.node_runner.process_node, + store=self.session_store, queue_update_callback=self.turn_intake.update_queue_positions, node_started_callback=self.turn_intake.mark_node_processing, unexpected_failure_callback=self._apply_unexpected_failure, @@ -156,6 +155,7 @@ def _build_tree_queue( return TreeQueueManager.from_snapshot( snapshot, self.node_runner.process_node, + store=self.session_store, queue_update_callback=self.turn_intake.update_queue_positions, node_started_callback=self.turn_intake.mark_node_processing, unexpected_failure_callback=self._apply_unexpected_failure, @@ -179,16 +179,24 @@ def tree_queue(self) -> TreeQueueManager: """Expose the manager facade for diagnostics and smoke tests.""" return self._tree_queue - def restore(self) -> None: + async def restore(self) -> None: """Restore and reconcile persisted conversations before platform start.""" - snapshot = self.session_store.load_conversation_snapshot() + snapshot = await self.session_store.load_conversation_snapshot() if snapshot.is_empty: return logger.info("Restoring {} conversation trees...", len(snapshot.trees)) self._tree_queue = self._build_tree_queue(snapshot) normalized = self._tree_queue.restored_snapshot if normalized is not None and normalized != snapshot: - self.session_store.save_conversation_snapshot(normalized) + await self.session_store.commit_trees( + tuple( + tree + for tree in normalized.trees.values() + if self.platform_name == "messaging" + or tree.scope.platform == self.platform_name + ), + removed=tuple(snapshot.trees.keys() - normalized.trees.keys()), + ) self._pending_restored_status_targets = self._tree_queue.restored_stale_targets async def repair_restored_statuses(self) -> None: @@ -271,7 +279,7 @@ async def _finalize_startup_notice( MessageScope(platform=self.platform_name, chat_id=notice.chat_id), 0 ) if not must_discard: - must_discard = not self.record_outgoing_message( + must_discard = not await self.record_outgoing_message( self.platform_name, notice.chat_id, message_id, @@ -302,7 +310,7 @@ async def _discard_startup_notice( ), ) async with self._state_lock: - tracked = self.record_outgoing_message( + tracked = await self.record_outgoing_message( self.platform_name, chat_id, message_id, @@ -315,7 +323,7 @@ async def _discard_startup_notice( return async with self._state_lock: - self.forget_tracked_message_ids( + await self.forget_tracked_message_ids( self.platform_name, chat_id, {message_id}, @@ -323,9 +331,13 @@ async def _discard_startup_notice( async def close(self) -> None: """Finish every owned task and durable write before releasing delivery.""" - await self.stop_all_tasks() + await self._cancel_all_pending_voices() + async with self._state_lock: + self._stop_generation += 1 + result = await self._tree_queue.shutdown() + self._apply_cancellation_result(result) + await self.cli_manager.stop_all() await self._tree_queue.wait_idle() - self.session_store.flush_pending_save() async def handle_message(self, incoming: IncomingMessage) -> None: """Handle one platform message.""" @@ -353,7 +365,7 @@ async def handle_message(self, incoming: IncomingMessage) -> None: clear_generation=self._clear_generations.get(incoming.scope, 0), ) if not is_standalone_clear: - self._record_incoming_message(incoming) + await self._record_incoming_message(incoming) try: await self.turn_intake.handle_message( incoming, @@ -362,7 +374,7 @@ async def handle_message(self, incoming: IncomingMessage) -> None: except BaseException: if is_standalone_clear: async with self._state_lock: - self._record_incoming_message(incoming) + await self._record_incoming_message(incoming) raise async def resolve_reply( @@ -404,13 +416,14 @@ async def _admit_if_current( ) if admission_token != current_token: return None + await self.record_outgoing_message( + incoming.platform, incoming.chat_id, status_message_id, "status" + ) decision = await self._tree_queue.admit( incoming, status_message_id, parent_reference_id=parent_reference_id, ) - if decision.snapshot is not None: - self.session_store.save_tree_snapshot(decision.snapshot) return decision def get_tree_count(self) -> int: @@ -477,9 +490,6 @@ async def _clear_reply( ) if not subtree.tree_matched and voice_result is None: return None - self._save_cancellation_snapshots(subtree.cancellation) - if subtree.removed_tree_identity is not None: - self.session_store.remove_tree_snapshot(subtree.removed_tree_identity) delete_message_ids = set(subtree.delete_message_ids) if voice_result is not None: @@ -501,9 +511,14 @@ async def _stop_all_tasks(self) -> StopOutcome: for voice in voice_results: self.render_voice_stopped(voice) async with self._state_lock: - self._stop_generation += 1 + + def publish_stop() -> None: + self._stop_generation += 1 + logger.info("Cancelling tree queue tasks...") - result = await self._tree_queue.cancel_all(reason=CancellationReason.STOP) + result = await self._tree_queue.cancel_all( + reason=CancellationReason.STOP, on_committed=publish_stop + ) logger.info("Cancelled {} nodes", len(result.effects)) self._apply_cancellation_result(result) logger.info("Stopping all CLI sessions...") @@ -530,37 +545,23 @@ async def _clear_chat( for voice in voice_results: delete_message_ids.update(voice.delete_message_ids) delete_message_ids.update( - self.session_store.get_tracked_message_ids_for_chat(platform, chat_id) + await self.session_store.get_tracked_message_ids_for_chat( + platform, chat_id + ) ) delete_message_ids.update( await self._tree_queue.get_message_ids_for_chat(platform, chat_id) ) - # All fallible/cancellable reads precede the commit boundary. Once - # the scope generation advances, state removal is one-way work. - self._clear_generations[clear_scope] = ( - self._clear_generations.get(clear_scope, 0) + 1 - ) - failures: list[Exception] = [] - try: - await self._tree_queue.clear_scope( - clear_scope, - reason=CancellationReason.CLEAR, - ) - except Exception as exc: - failures.append(exc) - try: - self.session_store.clear_scope(clear_scope) - except Exception as exc: - failures.append(exc) - logger.warning( - "Failed to persist final cleared session state: {}", - type(exc).__name__, + + def publish_clear() -> None: + self._clear_generations[clear_scope] = ( + self._clear_generations.get(clear_scope, 0) + 1 ) - if len(failures) == 1: - raise failures[0] - if failures: - raise ExceptionGroup("Chat clear failed", failures) + + await self._tree_queue.clear_scope( + clear_scope, reason=CancellationReason.CLEAR, on_committed=publish_clear + ) return frozenset(delete_message_ids) async def _cancel_all_pending_voices( @@ -603,14 +604,14 @@ def render_voice_stopped(self, result: VoiceCancellationResult) -> None: ) ) - def forget_tracked_message_ids( + async def forget_tracked_message_ids( self, platform: str, chat_id: str, message_ids: set[str], ) -> None: try: - self.session_store.forget_tracked_message_ids( + await self.session_store.forget_tracked_message_ids( platform, chat_id, message_ids ) except Exception as exc: @@ -619,7 +620,7 @@ def forget_tracked_message_ids( type(exc).__name__, ) - def record_outgoing_message( + async def record_outgoing_message( self, platform: str, chat_id: str, @@ -630,7 +631,7 @@ def record_outgoing_message( if not msg_id: return False try: - self.session_store.record_message_id( + await self.session_store.record_message_id( platform, chat_id, str(msg_id), @@ -648,7 +649,7 @@ def record_outgoing_message( return False return True - def _record_incoming_message(self, incoming: IncomingMessage) -> bool: + async def _record_incoming_message(self, incoming: IncomingMessage) -> bool: """Record an inbound prompt, voice note, or command for standalone clear.""" command = parse_command_base(incoming.text) kind = ( @@ -659,7 +660,7 @@ def _record_incoming_message(self, incoming: IncomingMessage) -> bool: else "prompt" ) try: - self.session_store.record_message_id( + await self.session_store.record_message_id( incoming.platform, incoming.chat_id, str(incoming.message_id), @@ -677,11 +678,6 @@ def _record_incoming_message(self, incoming: IncomingMessage) -> bool: return False return True - def _save_cancellation_snapshots(self, result: CancellationResult) -> None: - """Persist transition snapshots without publishing cancellation UI.""" - for snapshot in result.snapshots: - self.session_store.save_tree_snapshot(snapshot) - def _apply_cancellation_result(self, result: CancellationResult) -> None: """Apply detached UI and persistence effects from one transition.""" for effect in result.effects: @@ -694,18 +690,15 @@ def _apply_cancellation_result(self, result: CancellationResult) -> None: parse_mode=self._parse_mode(), ) ) - self._save_cancellation_snapshots(result) def _apply_unexpected_failure(self, result: FailureResult) -> None: - """Persist and render a failure that escaped the total node runner.""" - if result.snapshot is not None: - self.session_store.save_tree_snapshot(result.snapshot) + """Render a failure reported by the execution owner.""" for target in result.affected: self.outbound.fire_and_forget( self.outbound.queue_edit_message( target.scope.chat_id, target.status_message_id, - self.format_status("๐Ÿ’ฅ", "Task Failed"), + self.format_status("๐Ÿ’ฅ", "Task Failed", result.message), parse_mode=self._parse_mode(), ) ) diff --git a/src/free_claude_code/runtime/application.py b/src/free_claude_code/runtime/application.py index 3302c98004..c887cf56c6 100644 --- a/src/free_claude_code/runtime/application.py +++ b/src/free_claude_code/runtime/application.py @@ -10,6 +10,7 @@ from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from functools import partial +from pathlib import Path from typing import TYPE_CHECKING, Literal from loguru import logger @@ -67,6 +68,8 @@ import free_claude_code.cli.managed as cli_managed import free_claude_code.messaging.workflow as messaging_workflow_module + from .messaging_sqlite import SQLiteMessagingStore + from free_claude_code.application.readiness import InitializationWait from free_claude_code.core.async_tasks import run_sync_owned @@ -74,6 +77,7 @@ from .folder_picker import NativeFolderPicker from .provider_manager import ProviderRuntimeManager from .retired_chat import remove_retired_chat_history +from .sqlite_database import SQLiteDatabase RestartCallback = Callable[[], None] IntegrationAction = Literal["status", "connect", "disconnect", "refresh"] @@ -184,6 +188,7 @@ def __init__( configuration: ConfigurationService, transcriber: Transcriber | None, code_service: CodeService | None = None, + database: SQLiteDatabase | None = None, transcriber_factory: Callable[[Settings], Awaitable[Transcriber | None]] | None = None, restart_callback: RestartCallback | None = None, @@ -192,6 +197,11 @@ def __init__( self.provider_manager = provider_manager self._configuration = configuration self._code_service = code_service + self._database = database + self._messaging_store: SQLiteMessagingStore | None = None + self._messaging_import_task: asyncio.Task[None] | None = None + self._messaging_warning: str | None = None + self._messaging_storage_error: Exception | None = None self._folder_picker = NativeFolderPicker() self._transcriber = transcriber self._transcriber_factory = transcriber_factory @@ -269,6 +279,13 @@ async def start(self) -> None: name="fcc-retired-chat-cleanup", ) ) + if self._database is not None: + self._messaging_state = "starting" + self._messaging_import_task = asyncio.create_task( + self._initialize_messaging_storage(), + name="fcc-messaging-import", + ) + self._startup_tasks.append(self._messaging_import_task) if self._code_service is not None: self._startup_tasks.append( asyncio.create_task( @@ -864,6 +881,7 @@ async def admin_status(self) -> JsonObject: "messaging": { "state": self._messaging_state, "message": self._messaging_error, + "warning": self._messaging_warning, }, "integrations": { "vscode-chat": self._vscode_update.snapshot(), @@ -1006,12 +1024,15 @@ def _restart_metadata( return result async def _start_messaging_if_configured(self) -> None: + if self._messaging_import_task is not None: + await asyncio.shield(self._messaging_import_task) if self.settings.messaging_platform == "none": + self._messaging_state = "disabled" return try: def load_modules() -> None: - for name in ("cli.managed", "messaging.session", "messaging.workflow"): + for name in ("cli.managed", "messaging.workflow"): importlib.import_module(f"free_claude_code.{name}") importlib.import_module( f"free_claude_code.messaging.platforms.{self.settings.messaging_platform}" @@ -1077,12 +1098,36 @@ def _messaging_options(self) -> MessagingPlatformOptions: log_api_error_tracebacks=settings.log_api_error_tracebacks, ) + async def _initialize_messaging_storage(self) -> None: + if self._database is None: + return + try: + await self._database.start() + from .messaging_import import import_legacy + from .messaging_sqlite import SQLiteMessagingStore + + self._messaging_store = SQLiteMessagingStore( + self._database, + managed_message_cap=self.settings.max_message_log_entries_per_chat, + ) + self._messaging_warning = await import_legacy( + self._database, Path(messaging_state_dir_path()) / "sessions.json" + ) + await self._messaging_store.trim() + except Exception as exc: + self._messaging_storage_error = exc + logger.error( + "Messaging storage initialization failed: {}", type(exc).__name__ + ) + return + if self._messaging_warning: + logger.warning("{}", self._messaging_warning) + async def _start_messaging_workflow( self, components: MessagingPlatformComponents, ) -> None: import free_claude_code.cli.managed as cli_managed - import free_claude_code.messaging.session as messaging_session import free_claude_code.messaging.workflow as messaging_workflow_module settings = self.settings @@ -1093,8 +1138,6 @@ async def _start_messaging_workflow( else os.getcwd() ) await run_sync_owned(partial(os.makedirs, workspace, exist_ok=True)) - data_path = os.path.abspath(messaging_state_dir_path()) - await run_sync_owned(partial(os.makedirs, data_path, exist_ok=True)) allowed_dirs = [workspace] if settings.allowed_dir else [] self._cli_manager = cli_managed.ManagedClaudeSessionManager( @@ -1105,13 +1148,15 @@ async def _start_messaging_workflow( log_raw_cli_diagnostics=settings.log_raw_cli_diagnostics, log_messaging_error_details=settings.log_messaging_error_details, ) - session_store = await run_sync_owned( - partial( - messaging_session.SessionStore, - storage_path=os.path.join(data_path, "sessions.json"), - managed_message_cap=settings.max_message_log_entries_per_chat, - ) - ) + if self._messaging_import_task is not None: + await asyncio.shield(self._messaging_import_task) + if self._messaging_storage_error is not None: + raise ApplicationUnavailableError( + "Messaging storage is unavailable." + ) from self._messaging_storage_error + if self._messaging_store is None: + raise ApplicationUnavailableError("Messaging storage is unavailable.") + session_store = self._messaging_store workflow = messaging_workflow_module.MessagingWorkflow( platform_name=components.name, outbound=components.outbound, @@ -1124,7 +1169,7 @@ async def _start_messaging_workflow( log_messaging_error_details=settings.log_messaging_error_details, ) self._messaging_workflow = workflow - workflow.restore() + await workflow.restore() components.runtime.on_message(workflow.handle_message) await self._http_ready.wait() if self._draining: @@ -1147,6 +1192,10 @@ async def _close_owned_resources(self) -> bool: log_verbose_errors=verbose, ): return False + if self._database is not None and not await best_effort( + "database.close", self._database.close() + ): + return False if not await self._cleanup_transcriber(): return False if not self._provider_manager_closed: diff --git a/src/free_claude_code/runtime/bootstrap.py b/src/free_claude_code/runtime/bootstrap.py index b8c661116e..abdf7d6bfa 100644 --- a/src/free_claude_code/runtime/bootstrap.py +++ b/src/free_claude_code/runtime/bootstrap.py @@ -35,6 +35,7 @@ from .codex_catalog import CodexModelCatalogPublisher from .configuration import ConfigurationService from .provider_manager import ProviderRuntimeManager +from .sqlite_database import SQLiteDatabase from .web_tools.client import HTTPWebToolsClient @@ -73,18 +74,18 @@ def build_asgi_app( ), model_catalog_publisher=CodexModelCatalogPublisher(), ) + database = SQLiteDatabase( + fcc_database_path(), code_lock_path(), legacy_path=legacy_code_database_path() + ) code_service = CodeService( - SQLiteCodeStore( - fcc_database_path(), - code_lock_path(), - legacy_database_path=legacy_code_database_path(), - ), + SQLiteCodeStore(database), CodexHarnessFactory(provider_manager), ) runtime = ApplicationRuntime( provider_manager, configuration=ConfigurationService(ManagedConfigStore()), code_service=code_service, + database=database, transcriber=None, transcriber_factory=_create_transcriber, restart_callback=restart_callback, diff --git a/src/free_claude_code/runtime/code_sessions_sqlite.py b/src/free_claude_code/runtime/code_sessions_sqlite.py index eddb9f926e..b3d5227917 100644 --- a/src/free_claude_code/runtime/code_sessions_sqlite.py +++ b/src/free_claude_code/runtime/code_sessions_sqlite.py @@ -2,13 +2,9 @@ import asyncio import json -import os import sqlite3 from collections.abc import Callable, Sequence from contextlib import closing -from pathlib import Path - -import anyio.to_thread from free_claude_code.application.code_sessions.models import ( CodeConflictError, @@ -23,9 +19,8 @@ Record, now_ms, ) -from free_claude_code.core.interprocess_lock import InterprocessFileLock -from .sqlite_database import initialize_database +from .sqlite_database import SQLiteDatabase _JSON_FIELDS = frozenset( {"raw", "form", "error_details", "request_id", "native_permission_defaults"} @@ -302,16 +297,9 @@ def _write_prompt(connection: sqlite3.Connection, prompt: CodePrompt) -> None: class SQLiteCodeStore: - def __init__( - self, - database_path: Path, - lock_path: Path, - *, - legacy_database_path: Path | None = None, - ) -> None: - self._path = database_path - self._legacy_path = legacy_database_path - self._lock = InterprocessFileLock(lock_path) + def __init__(self, database: SQLiteDatabase) -> None: + self.database = database + self._path = database.path self._started = False self._lifecycle = asyncio.Lock() @@ -319,53 +307,22 @@ async def start(self) -> None: async with self._lifecycle: if self._started: return - initialization = asyncio.create_task( - anyio.to_thread.run_sync(self._initialize) - ) try: - await asyncio.shield(initialization) - except BaseException: - await asyncio.gather(initialization, return_exceptions=True) - self._lock.release() - raise + await self.database.start() + await self.database.work(lambda: self._execute(self._recover)) + except sqlite3.IntegrityError as exc: + raise CodeConflictError( + "Saved Code history conflicts with its schema." + ) from exc + except (sqlite3.Error, OSError) as exc: + raise CodeUnavailableError( + f"Code storage initialization failed: {exc}" + ) from exc self._started = True - def _initialize(self) -> None: - if not self._lock.acquire(): - raise CodeUnavailableError( - "Code sessions is already owned by another FCC server." - ) - try: - self._path.parent.mkdir(parents=True, exist_ok=True) - if os.name != "nt": - self._path.parent.chmod(0o700) - initialize_database(self._path, self._legacy_path) - except sqlite3.IntegrityError as exc: - raise CodeConflictError( - "Saved Code history conflicts with its schema." - ) from exc - except sqlite3.Error as exc: - raise CodeUnavailableError( - f"Code storage initialization failed: {exc}" - ) from exc - except OSError as exc: - raise CodeUnavailableError( - "Code database files could not be initialized or moved." - ) from exc - self._execute(self._recover) - if os.name != "nt": - for path in ( - self._path, - Path(f"{self._path}-wal"), - Path(f"{self._path}-shm"), - ): - if path.exists(): - path.chmod(0o600) - async def close(self) -> None: async with self._lifecycle: self._started = False - await anyio.to_thread.run_sync(self._lock.release) def _execute[T](self, operation: Callable[[sqlite3.Connection], T]) -> T: try: @@ -384,7 +341,7 @@ def _execute[T](self, operation: Callable[[sqlite3.Connection], T]) -> T: async def _run[T](self, operation: Callable[[sqlite3.Connection], T]) -> T: if not self._started: raise CodeUnavailableError("Code session storage is closed.") - return await anyio.to_thread.run_sync(self._execute, operation) + return await self.database.work(lambda: self._execute(operation)) def _recover(self, connection: sqlite3.Connection) -> None: connection.execute( diff --git a/src/free_claude_code/runtime/messaging_import.py b/src/free_claude_code/runtime/messaging_import.py new file mode 100644 index 0000000000..82d50b4324 --- /dev/null +++ b/src/free_claude_code/runtime/messaging_import.py @@ -0,0 +1,206 @@ +"""One-time, best-effort import of FCC's retired messaging JSON file.""" + +import json +import sqlite3 +from collections.abc import Callable +from dataclasses import dataclass, field +from functools import partial +from pathlib import Path +from typing import Any + +from free_claude_code.messaging.trees import ( + TreeIdentity, + TreeSnapshot, + normalize_tree_snapshot, +) + +from .messaging_sqlite import write_tree +from .sqlite_database import SQLiteDatabase + + +@dataclass +class LegacyData: + outcome: str = "complete" + trees: dict[TreeIdentity, TreeSnapshot] = field(default_factory=dict) + messages: list[tuple[str, str, str, str, str, str]] = field(default_factory=list) + skipped: int = 0 + + def reject(self) -> None: + self.outcome = "partial" + self.skipped += 1 + + +def read_legacy(path: Path) -> LegacyData: + result = LegacyData() + try: + if path.is_symlink(): + return LegacyData(outcome="unreadable", skipped=1) + raw = json.loads(path.read_text(encoding="utf-8")) + except FileNotFoundError: + return LegacyData(outcome="absent") + except OSError, UnicodeError, ValueError: + return LegacyData(outcome="unreadable", skipped=1) + if not isinstance(raw, dict): + return LegacyData(outcome="unreadable", skipped=1) + conversation = raw.get("conversation", raw) + if not isinstance(conversation, dict): + result.reject() + conversation = raw + candidates = conversation.get("trees", []) + if isinstance(candidates, dict): + candidates = list(candidates.values()) + if not isinstance(candidates, list): + result.reject() + candidates = [] + # Historical JSON loading used the last duplicate tree identity. + decoded: dict[TreeIdentity, TreeSnapshot] = {} + for candidate in candidates: + try: + snapshot = TreeSnapshot.from_json(candidate) + if snapshot is None: + raise ValueError("Missing tree scope") + if snapshot.identity in decoded: + result.reject() + decoded[snapshot.identity] = snapshot + except KeyError, TypeError, ValueError: + result.reject() + for snapshot in decoded.values(): + try: + normalized = normalize_tree_snapshot(snapshot) + result.trees[normalized.identity] = normalized + except KeyError, TypeError, ValueError: + result.reject() + messages = raw.get("managed_messages", raw.get("message_log", {})) + if not isinstance(messages, dict): + result.reject() + messages = {} + for chat_key, items in messages.items(): + if ( + not isinstance(chat_key, str) + or ":" not in chat_key + or not isinstance(items, list) + ): + result.reject() + continue + platform, chat_id = chat_key.split(":", 1) + for item in items: + record = _message(platform, chat_id, item) + if record is None: + result.reject() + else: + result.messages.append(record) + return result + + +def _message( + platform: str, chat_id: str, item: Any +) -> tuple[str, str, str, str, str, str] | None: + if not isinstance(item, dict) or item.get("message_id") is None: + return None + direction, kind = str(item.get("direction") or ""), str(item.get("kind") or "") + if direction not in {"in", "out"} or not kind: + return None + return ( + platform, + chat_id, + str(item["message_id"]), + str(item.get("ts") or ""), + direction, + kind, + ) + + +async def import_legacy(database: SQLiteDatabase, path: Path) -> str | None: + return await database.work(lambda: _import_legacy(database, path)) + + +def _try_import(connection: sqlite3.Connection, write: Callable[[], object]) -> bool: + """Reject one damaged unit without publishing part of a tree or its references.""" + connection.execute("SAVEPOINT legacy_record") + try: + write() + except sqlite3.IntegrityError, UnicodeEncodeError: + connection.execute("ROLLBACK TO legacy_record") + connection.execute("RELEASE legacy_record") + return False + connection.execute("RELEASE legacy_record") + return True + + +def _import_legacy(database: SQLiteDatabase, path: Path) -> str | None: + receipt = database.execute( + lambda connection: connection.execute( + "SELECT * FROM messaging_legacy_import" + ).fetchone(), + write=False, + ) + warning = None + if receipt is None: + data = read_legacy(path) + + def commit(connection: sqlite3.Connection) -> bool: + references: set[tuple[str, str, str]] = set() + trees = 0 + for snapshot in data.trees.values(): + scoped = { + (snapshot.scope.platform, snapshot.scope.chat_id, reference) + for reference in snapshot.lookup_ids() + } + if scoped & references or not _try_import( + connection, partial(write_tree, connection, snapshot) + ): + data.reject() + continue + references.update(scoped) + trees += 1 + seen: set[tuple[str, str, str]] = set() + for record in data.messages: + key = record[:3] + if key in seen: + continue + if not _try_import( + connection, + partial( + connection.execute, + "INSERT INTO messaging_managed_messages(platform,chat_id,message_id,ts,direction,kind) VALUES (?,?,?,?,?,?)", + record, + ), + ): + data.reject() + continue + seen.add(key) + complete = data.outcome in {"complete", "absent"} + connection.execute( + "INSERT INTO messaging_legacy_import VALUES ('sessions.json',?,?,?,?,?)", + ( + data.outcome, + trees, + len(seen), + data.skipped, + int(complete), + ), + ) + return complete + + complete = database.execute(commit) + if not complete: + return f"Some previous messaging history could not be restored. Messaging can still be used. The original file was kept at {path}." + elif not receipt["cleanup_pending"]: + # Incomplete-import notices are one-time. The receipt still prevents replay. + return None + try: + if path.is_symlink(): + raise OSError("Legacy source is redirected") + path.unlink(missing_ok=True) + for temporary in path.parent.glob(".sessions.*.tmp.json"): + if temporary.is_file() and not temporary.is_symlink(): + temporary.unlink() + except OSError: + warning = f"Messaging history was migrated. Old files at {path.parent} could not be removed yet. FCC will retry cleanup at startup." + else: + database.execute( + lambda connection: connection.execute( + "UPDATE messaging_legacy_import SET cleanup_pending=0" + ) + ) + return warning diff --git a/src/free_claude_code/runtime/messaging_sqlite.py b/src/free_claude_code/runtime/messaging_sqlite.py new file mode 100644 index 0000000000..3152b638ed --- /dev/null +++ b/src/free_claude_code/runtime/messaging_sqlite.py @@ -0,0 +1,183 @@ +"""Transactional messaging records in the application-owned FCC database.""" + +import sqlite3 +from datetime import UTC, datetime + +from free_claude_code.messaging.models import MessageScope +from free_claude_code.messaging.trees import ( + ConversationSnapshot, + MessagingStorageError, + TreeIdentity, + TreeSnapshot, + normalize_tree_snapshot, +) + +from .sqlite_database import SQLiteDatabase + +_FIELDS = ( + "node_id", + "status_message_id", + "state", + "parent_id", + "parent_reference_id", + "session_id", +) + + +def write_tree(connection: sqlite3.Connection, snapshot: TreeSnapshot) -> None: + """Update only changed rows, inserting new parents before their children.""" + normalized = normalize_tree_snapshot(snapshot) + scope = (snapshot.scope.platform, snapshot.scope.chat_id) + identity = (*scope, snapshot.root_id) + connection.execute( + "INSERT INTO messaging_trees VALUES (?,?,?) ON CONFLICT DO NOTHING", identity + ) + previous = { + row["node_id"]: dict(row) + for row in connection.execute( + "SELECT * FROM messaging_nodes WHERE platform=? AND chat_id=? AND root_id=?", + identity, + ) + } + for node_id in previous.keys() - normalized.nodes.keys(): + connection.execute( + "DELETE FROM messaging_nodes WHERE platform=? AND chat_id=? AND node_id=?", + (*scope, node_id), + ) + for node_id in normalized.nodes: + node = normalized.nodes[node_id] + before = previous.get(node_id) + if before is not None and all( + before[field] == node.get(field) for field in _FIELDS + ): + continue + connection.execute( + "INSERT INTO messaging_nodes (platform,chat_id,root_id,node_id,status_message_id,state,parent_id,parent_reference_id,session_id) " + "VALUES (?,?,?,?,?,?,?,?,?) ON CONFLICT(platform,chat_id,node_id) DO UPDATE SET " + "root_id=excluded.root_id,status_message_id=excluded.status_message_id,state=excluded.state," + "parent_id=excluded.parent_id,parent_reference_id=excluded.parent_reference_id,session_id=excluded.session_id", + (*identity, *(node.get(field) for field in _FIELDS)), + ) + + +class SQLiteMessagingStore: + def __init__( + self, database: SQLiteDatabase, *, managed_message_cap: int | None = None + ) -> None: + self.database = database + self._cap = managed_message_cap + + async def load_conversation_snapshot(self) -> ConversationSnapshot: + def load(connection: sqlite3.Connection) -> ConversationSnapshot: + trees: dict[TreeIdentity, TreeSnapshot] = {} + for row in connection.execute( + "SELECT * FROM messaging_nodes ORDER BY rowid" + ): + scope = MessageScope(platform=row["platform"], chat_id=row["chat_id"]) + identity = TreeIdentity(scope=scope, root_id=row["root_id"]) + snapshot = trees.setdefault( + identity, TreeSnapshot(scope, identity.root_id, {}) + ) + snapshot.nodes[row["node_id"]] = { + field: row[field] for field in _FIELDS + } + return ConversationSnapshot(trees) + + return await self.database.run(load, write=False) + + async def commit_trees( + self, + snapshots: tuple[TreeSnapshot, ...], + *, + removed: tuple[TreeIdentity, ...] = (), + clear_scope: MessageScope | None = None, + ) -> None: + def commit(connection: sqlite3.Connection) -> None: + if clear_scope is not None: + scope = (clear_scope.platform, clear_scope.chat_id) + connection.execute( + "DELETE FROM messaging_trees WHERE platform=? AND chat_id=?", scope + ) + connection.execute( + "DELETE FROM messaging_managed_messages WHERE platform=? AND chat_id=?", + scope, + ) + for identity in removed: + connection.execute( + "DELETE FROM messaging_trees WHERE platform=? AND chat_id=? AND root_id=?", + (identity.scope.platform, identity.scope.chat_id, identity.root_id), + ) + for snapshot in snapshots: + write_tree(connection, snapshot) + + try: + await self.database.run(commit) + except sqlite3.Error as exc: + raise MessagingStorageError( + "Messaging history could not be saved." + ) from exc + + def _trim( + self, connection: sqlite3.Connection, platform: str, chat_id: str + ) -> None: + if self._cap is not None and self._cap > 0: + connection.execute( + "DELETE FROM messaging_managed_messages WHERE sequence IN (" + "SELECT sequence FROM messaging_managed_messages WHERE platform=? AND chat_id=? " + "ORDER BY sequence DESC LIMIT -1 OFFSET ?)", + (platform, chat_id, self._cap), + ) + + async def trim(self) -> None: + def trim(connection: sqlite3.Connection) -> None: + for row in connection.execute( + "SELECT DISTINCT platform,chat_id FROM messaging_managed_messages" + ).fetchall(): + self._trim(connection, row[0], row[1]) + + await self.database.run(trim) + + async def record_message_id( + self, platform: str, chat_id: str, message_id: str, direction: str, kind: str + ) -> None: + def record(connection: sqlite3.Connection) -> None: + connection.execute( + "INSERT INTO messaging_managed_messages(platform,chat_id,message_id,ts,direction,kind) " + "VALUES (?,?,?,?,?,?) ON CONFLICT(platform,chat_id,message_id) DO NOTHING", + ( + str(platform), + str(chat_id), + str(message_id), + datetime.now(UTC).isoformat(), + direction, + kind, + ), + ) + self._trim(connection, platform, chat_id) + + await self.database.run(record) + + async def get_tracked_message_ids_for_chat( + self, platform: str, chat_id: str + ) -> list[str]: + return await self.database.run( + lambda connection: [ + row[0] + for row in connection.execute( + "SELECT message_id FROM messaging_managed_messages WHERE platform=? AND chat_id=? ORDER BY sequence", + (platform, chat_id), + ) + ], + write=False, + ) + + async def forget_tracked_message_ids( + self, platform: str, chat_id: str, message_ids: set[str] + ) -> None: + def forget(connection: sqlite3.Connection) -> None: + connection.executemany( + "DELETE FROM messaging_managed_messages WHERE platform=? AND chat_id=? AND message_id=?", + [(platform, chat_id, item) for item in message_ids], + ) + + await self.database.run(forget) diff --git a/src/free_claude_code/runtime/sqlite_database.py b/src/free_claude_code/runtime/sqlite_database.py index 1997171daa..2c74da3169 100644 --- a/src/free_claude_code/runtime/sqlite_database.py +++ b/src/free_claude_code/runtime/sqlite_database.py @@ -1,13 +1,19 @@ """Offline relocation and transactional schema initialization of FCC's database. -The caller must hold the Code owner lock until the store closes. Connections +The application holds the shared owner lock until both features close. Connections opened here never escape initialization, including when migration fails. """ +import asyncio +import os import sqlite3 +from collections.abc import Callable from contextlib import closing from pathlib import Path +from free_claude_code.core.async_tasks import run_sync_owned +from free_claude_code.core.interprocess_lock import InterprocessFileLock + from .sqlite_migrations import MIGRATIONS # Historical columns identify supported files without depending on live models. @@ -21,6 +27,13 @@ "code_deleted": "id", } _SIDECARS = ("-wal", "-shm", "-journal") +_MESSAGING_COLUMNS = { + "messaging_trees": "platform chat_id root_id", + "messaging_nodes": "platform chat_id root_id node_id parent_id parent_reference_id status_message_id session_id state", + "messaging_references": "platform chat_id reference_id node_id kind", + "messaging_managed_messages": "sequence platform chat_id message_id ts direction kind", + "messaging_legacy_import": "source outcome trees messages skipped cleanup_pending", +} def _connect(path: Path, *, existing: bool = False) -> sqlite3.Connection: @@ -47,11 +60,12 @@ def _schema_version(connection: sqlite3.Connection) -> int: } if version == 0 and not tables: return version - if tables != _BASE_COLUMNS.keys(): + columns = _BASE_COLUMNS | (_MESSAGING_COLUMNS if version >= 4 else {}) + if tables != columns.keys(): raise sqlite3.DatabaseError( "Unrecognized FCC database schema. Saved data was preserved." ) - for table, baseline in _BASE_COLUMNS.items(): + for table, baseline in columns.items(): expected = set(baseline.split()) if version >= 2 and table in ("code_sessions", "code_runs"): expected.add("mode") @@ -161,3 +175,101 @@ def initialize_database(path: Path, legacy_path: Path | None = None) -> None: if connection.in_transaction: connection.execute("ROLLBACK") raise + + +class SQLiteDatabase: + """Application-owned initialization and lifetime for both durable features.""" + + def __init__( + self, path: Path, lock_path: Path, *, legacy_path: Path | None = None + ) -> None: + self.path = path + self._legacy_path = legacy_path + self._owner = InterprocessFileLock(lock_path) + self._lifecycle = asyncio.Lock() + self._started = False + self._closing = False + self._startup_error: Exception | None = None + self._operations: set[asyncio.Task] = set() + + async def start(self) -> None: + async with self._lifecycle: + if self._started: + return + if self._closing: + raise sqlite3.OperationalError("FCC database is closing") + if self._startup_error is not None: + raise self._startup_error + try: + await run_sync_owned(self._initialize) + except Exception as exc: + self._startup_error = exc + self._owner.release() + raise + except BaseException: + self._owner.release() + raise + self._started = True + + def _initialize(self) -> None: + if not self._owner.acquire(): + raise sqlite3.OperationalError( + "FCC storage is already owned by another FCC server" + ) + self.path.parent.mkdir(parents=True, exist_ok=True) + if os.name != "nt": + self.path.parent.chmod(0o700) + initialize_database(self.path, self._legacy_path) + if os.name != "nt": + for path in (self.path, Path(f"{self.path}-wal"), Path(f"{self.path}-shm")): + if path.exists(): + path.chmod(0o600) + + def execute[T]( + self, operation: Callable[[sqlite3.Connection], T], *, write: bool = True + ) -> T: + with closing(_connect(self.path)) as connection: + connection.execute("BEGIN IMMEDIATE" if write else "BEGIN") + try: + result = operation(connection) + connection.execute("COMMIT") + return result + except BaseException: + if connection.in_transaction: + connection.execute("ROLLBACK") + raise + + async def run[T]( + self, operation: Callable[[sqlite3.Connection], T], *, write: bool = True + ) -> T: + return await self.work(lambda: self.execute(operation, write=write)) + + async def work[T](self, operation: Callable[[], T]) -> T: + if not self._started or self._closing: + raise sqlite3.OperationalError("FCC storage is closed") + task = asyncio.create_task(run_sync_owned(operation)) + self._operations.add(task) + try: + # Deliver the result before cancellation can separate SQL commit from + # the caller's in-memory publication/rollback. + while True: + try: + return await asyncio.shield(task) + except asyncio.CancelledError: + if task.done(): + return task.result() + finally: + self._operations.discard(task) + + async def close(self) -> None: + async with self._lifecycle: + self._closing = True + if self._operations: + drain = asyncio.gather(*self._operations, return_exceptions=True) + while not drain.done(): + try: + await asyncio.shield(drain) + except asyncio.CancelledError: + continue + self._started = False + await run_sync_owned(self._owner.release) diff --git a/src/free_claude_code/runtime/sqlite_migrations/__init__.py b/src/free_claude_code/runtime/sqlite_migrations/__init__.py index 186b3ac517..0dcebe48d5 100644 --- a/src/free_claude_code/runtime/sqlite_migrations/__init__.py +++ b/src/free_claude_code/runtime/sqlite_migrations/__init__.py @@ -8,10 +8,16 @@ import sqlite3 from collections.abc import Callable -from . import v0001_code_schema, v0002_code_permissions, v0003_code_context_usage +from . import ( + v0001_code_schema, + v0002_code_permissions, + v0003_code_context_usage, + v0004_messaging_schema, +) MIGRATIONS: tuple[tuple[int, Callable[[sqlite3.Connection], None]], ...] = ( (1, v0001_code_schema.upgrade), (2, v0002_code_permissions.upgrade), (3, v0003_code_context_usage.upgrade), + (4, v0004_messaging_schema.upgrade), ) diff --git a/src/free_claude_code/runtime/sqlite_migrations/v0004_messaging_schema.py b/src/free_claude_code/runtime/sqlite_migrations/v0004_messaging_schema.py new file mode 100644 index 0000000000..76db9f7356 --- /dev/null +++ b/src/free_claude_code/runtime/sqlite_migrations/v0004_messaging_schema.py @@ -0,0 +1,108 @@ +"""Relational messaging history and the retired JSON import receipt.""" + +import sqlite3 + + +def upgrade(connection: sqlite3.Connection) -> None: + connection.execute(""" + CREATE TABLE messaging_trees ( + platform TEXT NOT NULL, chat_id TEXT NOT NULL, root_id TEXT NOT NULL, + PRIMARY KEY (platform, chat_id, root_id), + FOREIGN KEY (platform, chat_id, root_id) + REFERENCES messaging_nodes(platform, chat_id, node_id) + DEFERRABLE INITIALLY DEFERRED + ) + """) + connection.execute(""" + CREATE TABLE messaging_nodes ( + platform TEXT NOT NULL, chat_id TEXT NOT NULL, node_id TEXT NOT NULL, + root_id TEXT NOT NULL, parent_id TEXT, parent_reference_id TEXT, + status_message_id TEXT, session_id TEXT, + state TEXT NOT NULL CHECK (state IN ('pending','in_progress','completed','error')), + PRIMARY KEY (platform, chat_id, node_id), + UNIQUE (platform, chat_id, root_id, node_id), + CHECK (length(node_id) > 0), + CHECK (status_message_id IS NULL OR (length(status_message_id) > 0 AND status_message_id <> node_id)), + CHECK (state NOT IN ('pending','in_progress') OR status_message_id IS NOT NULL), + CHECK ((node_id = root_id AND parent_id IS NULL AND parent_reference_id IS NULL) + OR (node_id <> root_id AND parent_id IS NOT NULL AND parent_reference_id IS NOT NULL AND parent_id <> node_id)), + FOREIGN KEY (platform, chat_id, root_id) + REFERENCES messaging_trees(platform, chat_id, root_id) ON DELETE CASCADE, + FOREIGN KEY (platform, chat_id, root_id, parent_id) + REFERENCES messaging_nodes(platform, chat_id, root_id, node_id) + DEFERRABLE INITIALLY DEFERRED, + FOREIGN KEY (platform, chat_id, parent_reference_id, parent_id) + REFERENCES messaging_references(platform, chat_id, reference_id, node_id) + DEFERRABLE INITIALLY DEFERRED + ) + """) + connection.execute(""" + CREATE TABLE messaging_references ( + platform TEXT NOT NULL, chat_id TEXT NOT NULL, reference_id TEXT NOT NULL, + node_id TEXT NOT NULL, kind TEXT NOT NULL CHECK (kind IN ('prompt','status')), + PRIMARY KEY (platform, chat_id, reference_id), + UNIQUE (platform, chat_id, node_id, kind), + UNIQUE (platform, chat_id, reference_id, node_id), + FOREIGN KEY (platform, chat_id, node_id) + REFERENCES messaging_nodes(platform, chat_id, node_id) ON DELETE CASCADE + ) + """) + connection.execute(""" + CREATE TRIGGER messaging_node_parent BEFORE INSERT ON messaging_nodes + WHEN NEW.parent_id IS NOT NULL AND NOT EXISTS ( + SELECT 1 FROM messaging_nodes WHERE platform = NEW.platform AND chat_id = NEW.chat_id + AND root_id = NEW.root_id AND node_id = NEW.parent_id) + BEGIN SELECT RAISE(ABORT, 'Messaging parent must exist before its child'); END + """) + connection.execute(""" + CREATE TRIGGER messaging_node_identity BEFORE UPDATE ON messaging_nodes + WHEN NEW.platform IS NOT OLD.platform OR NEW.chat_id IS NOT OLD.chat_id + OR NEW.node_id IS NOT OLD.node_id OR NEW.root_id IS NOT OLD.root_id + OR NEW.parent_id IS NOT OLD.parent_id OR NEW.parent_reference_id IS NOT OLD.parent_reference_id + BEGIN SELECT RAISE(ABORT, 'Messaging topology is immutable'); END + """) + connection.execute(""" + CREATE TRIGGER messaging_node_references AFTER INSERT ON messaging_nodes + BEGIN + INSERT INTO messaging_references VALUES (NEW.platform, NEW.chat_id, NEW.node_id, NEW.node_id, 'prompt'); + INSERT INTO messaging_references SELECT NEW.platform, NEW.chat_id, NEW.status_message_id, NEW.node_id, 'status' + WHERE NEW.status_message_id IS NOT NULL; + END + """) + connection.execute(""" + CREATE TRIGGER messaging_status_reference AFTER UPDATE OF status_message_id ON messaging_nodes + WHEN NEW.status_message_id IS NOT OLD.status_message_id + BEGIN + DELETE FROM messaging_references WHERE platform = OLD.platform AND chat_id = OLD.chat_id + AND node_id = OLD.node_id AND kind = 'status'; + INSERT INTO messaging_references SELECT NEW.platform, NEW.chat_id, NEW.status_message_id, NEW.node_id, 'status' + WHERE NEW.status_message_id IS NOT NULL; + END + """) + connection.execute( + "CREATE INDEX messaging_parent ON messaging_nodes(platform, chat_id, parent_id)" + ) + connection.execute( + "CREATE INDEX messaging_parent_reference ON messaging_nodes(platform, chat_id, parent_reference_id, parent_id)" + ) + connection.execute(""" + CREATE TABLE messaging_managed_messages ( + sequence INTEGER PRIMARY KEY, platform TEXT NOT NULL, chat_id TEXT NOT NULL, + message_id TEXT NOT NULL, ts TEXT NOT NULL, + direction TEXT NOT NULL CHECK (direction IN ('in','out')), + kind TEXT NOT NULL CHECK (length(kind) > 0), + UNIQUE (platform, chat_id, message_id) + ) + """) + connection.execute( + "CREATE INDEX messaging_message_order ON messaging_managed_messages(platform, chat_id, sequence)" + ) + connection.execute(""" + CREATE TABLE messaging_legacy_import ( + source TEXT PRIMARY KEY NOT NULL CHECK (source = 'sessions.json'), + outcome TEXT NOT NULL CHECK (outcome IN ('complete','partial','unreadable','absent')), + trees INTEGER NOT NULL CHECK (trees >= 0), messages INTEGER NOT NULL CHECK (messages >= 0), + skipped INTEGER NOT NULL CHECK (skipped >= 0), + cleanup_pending INTEGER NOT NULL CHECK (cleanup_pending IN (0,1)) + ) + """) diff --git a/tests/api/test_code_sessions_routes.py b/tests/api/test_code_sessions_routes.py index b7d3e8a827..0f2b63bde5 100644 --- a/tests/api/test_code_sessions_routes.py +++ b/tests/api/test_code_sessions_routes.py @@ -9,14 +9,15 @@ from free_claude_code.application.errors import ApplicationUnavailableError from free_claude_code.runtime.code_sessions_sqlite import SQLiteCodeStore from tests.api.support import create_test_app -from tests.code_sessions_support import CodexPackets, FakeHarness +from tests.code_sessions_support import CodexPackets, FakeHarness, close_code_database @pytest_asyncio.fixture -async def code_api(tmp_path): +async def code_api(database_factory, tmp_path): harness = FakeHarness() code = CodeService( - SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock"), harness + SQLiteCodeStore(database_factory(tmp_path / "code.db", tmp_path / "code.lock")), + harness, ) await code.start() app = create_test_app(code=code) @@ -27,6 +28,7 @@ async def code_api(tmp_path): yield client, code, harness, tmp_path, app finally: await code.close() + await close_code_database(code) await app.state.services.admin.close() diff --git a/tests/application/test_code_sessions.py b/tests/application/test_code_sessions.py index db4962a6cd..ae001fb83d 100644 --- a/tests/application/test_code_sessions.py +++ b/tests/application/test_code_sessions.py @@ -17,7 +17,12 @@ ) from free_claude_code.runtime.code_sessions_sqlite import SQLiteCodeStore from free_claude_code.runtime.codex_protocol import CodexProtocol -from tests.code_sessions_support import CodexPackets, FakeConnection, FakeHarness +from tests.code_sessions_support import ( + CodexPackets, + FakeConnection, + FakeHarness, + close_code_database, +) def new_id(): @@ -25,15 +30,18 @@ def new_id(): @pytest_asyncio.fixture -async def code(tmp_path): +async def code(database_factory, tmp_path): harness = FakeHarness() - store = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) service = CodeService(store, harness) await service.start() try: yield service, harness, tmp_path finally: await service.close() + await close_code_database(service) async def session_for(code): @@ -328,7 +336,7 @@ async def test_child_review_stays_with_first_entry_after_next_turn_begins(code): @pytest.mark.asyncio @pytest.mark.parametrize("closure", ["native", "service"]) async def test_pending_child_review_becomes_unavailable_when_idle_connection_ends( - code, closure + database_factory, code, closure ): service, harness, directory = code session = await session_for(code) @@ -360,8 +368,12 @@ async def test_pending_child_review_becomes_unavailable_when_idle_connection_end assert detail.active_review_ids == () assert detail.run.status == "completed" await service.close() + await close_code_database(service) restarted = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), harness + SQLiteCodeStore( + database_factory(directory / "code.db", directory / "code.lock") + ), + harness, ) await restarted.start() try: @@ -566,7 +578,9 @@ async def reject_defaults(session, *args, **kwargs): @pytest.mark.asyncio -async def test_native_reviews_and_notices_keep_identity_order_and_survive_restart(code): +async def test_native_reviews_and_notices_keep_identity_order_and_survive_restart( + database_factory, code +): service, harness, directory = code session = await session_for(code) await service.send( @@ -625,8 +639,12 @@ async def test_native_reviews_and_notices_keep_identity_order_and_survive_restar assert detail.run.status == "completed" assert detail.items[-1].text == "Native warning" await service.close() + await close_code_database(service) restarted = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), harness + SQLiteCodeStore( + database_factory(directory / "code.db", directory / "code.lock") + ), + harness, ) await restarted.start() try: @@ -751,7 +769,9 @@ async def test_failed_outcome_is_durable_once_after_partial_output_and_older_pag @pytest.mark.asyncio -async def test_retry_notice_does_not_finish_run_and_stale_error_is_ignored(code): +async def test_retry_notice_does_not_finish_run_and_stale_error_is_ignored( + database_factory, code +): service, harness, directory = code session = await session_for(code) await service.send( @@ -801,8 +821,12 @@ async def test_retry_notice_does_not_finish_run_and_stale_error_is_ignored(code) assert (await service.get_detail(session.id)).run.error is None detail = await service.get_detail(session.id) await service.close() + await close_code_database(service) restarted = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), harness + SQLiteCodeStore( + database_factory(directory / "code.db", directory / "code.lock") + ), + harness, ) await restarted.start() try: @@ -1022,19 +1046,21 @@ async def test_poorer_native_history_keeps_captured_source_and_reasoning(code): @pytest.mark.asyncio -async def test_storage_start_failure_is_isolated_to_code(tmp_path): +async def test_storage_start_failure_is_isolated_to_code(database_factory, tmp_path): class BrokenStore(SQLiteCodeStore): async def start(self): raise CodeUnavailableError("Code disk unavailable") service = CodeService( - BrokenStore(tmp_path / "code.db", tmp_path / "code.lock"), FakeHarness() + BrokenStore(database_factory(tmp_path / "code.db", tmp_path / "code.lock")), + FakeHarness(), ) try: await service.start() assert service.availability() == (False, "Code disk unavailable") finally: await service.close() + await close_code_database(service) @pytest.mark.asyncio @@ -1446,10 +1472,12 @@ async def test_idle_delete_removes_both_histories_and_blocks_late_create(code): @pytest.mark.asyncio -async def test_restart_preserves_receipt_and_does_not_replay_input(tmp_path): +async def test_restart_preserves_receipt_and_does_not_replay_input( + database_factory, tmp_path +): harness = FakeHarness() database, lock = tmp_path / "code.db", tmp_path / "code.lock" - first = CodeService(SQLiteCodeStore(database, lock), harness) + first = CodeService(SQLiteCodeStore(database_factory(database, lock)), harness) await first.start() session = await first.create_session(new_id(), str(tmp_path)) run = await first.send( @@ -1457,7 +1485,8 @@ async def test_restart_preserves_receipt_and_does_not_replay_input(tmp_path): ) await harness.started.wait() await first.close() - second = CodeService(SQLiteCodeStore(database, lock), harness) + await close_code_database(first) + second = CodeService(SQLiteCodeStore(database_factory(database, lock)), harness) await second.start() try: repeat = await second.send( @@ -1469,6 +1498,7 @@ async def test_restart_preserves_receipt_and_does_not_replay_input(tmp_path): assert harness.connections[0].inputs == [(run.id, "once", "provider/model")] finally: await second.close() + await close_code_database(second) @pytest.mark.asyncio @@ -1498,6 +1528,7 @@ async def test_two_answers_claim_one_native_prompt(code): @pytest.mark.asyncio async def test_prompt_arrival_has_one_durable_position_despite_repeated_native_item( + database_factory, code, ): service, harness, directory = code @@ -1546,8 +1577,12 @@ async def test_prompt_arrival_has_one_durable_position_despite_repeated_native_i await connection.text("turn-1", "after", "After approvals", complete=True) await connection.finish("turn-1") await service.close() + await close_code_database(service) restored = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), harness + SQLiteCodeStore( + database_factory(directory / "code.db", directory / "code.lock") + ), + harness, ) await restored.start() try: @@ -1838,7 +1873,9 @@ async def test_catalog_replacement_is_limited_to_selected_entry_and_session(code @pytest.mark.asyncio -async def test_cancelled_http_admission_still_commits_and_executes_once(tmp_path): +async def test_cancelled_http_admission_still_commits_and_executes_once( + database_factory, tmp_path +): committed = asyncio.Event() release = asyncio.Event() @@ -1851,7 +1888,8 @@ async def admit_run(self, session, run, item, expected_revision): harness = FakeHarness() service = CodeService( - GatedStore(tmp_path / "code.db", tmp_path / "code.lock"), harness + GatedStore(database_factory(tmp_path / "code.db", tmp_path / "code.lock")), + harness, ) await service.start() try: @@ -1884,6 +1922,7 @@ async def admit_run(self, session, run, item, expected_revision): finally: release.set() await service.close() + await close_code_database(service) @pytest.mark.asyncio @@ -2081,7 +2120,7 @@ def observe(event, data): @pytest.mark.asyncio @pytest.mark.parametrize("native_exists", [False, True]) async def test_startup_reconciles_persisted_deletion_without_another_native_delete( - code, native_exists + database_factory, code, native_exists ): service, harness, directory = code session = await idle_native(code) @@ -2092,10 +2131,14 @@ async def test_startup_reconciles_persisted_deletion_without_another_native_dele session.revision, ) await service.close() + await close_code_database(service) if not native_exists: harness.histories.pop(session.native_thread_id) restarted = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), harness + SQLiteCodeStore( + database_factory(directory / "code.db", directory / "code.lock") + ), + harness, ) try: await restarted.start() @@ -2114,7 +2157,7 @@ async def test_startup_reconciles_persisted_deletion_without_another_native_dele @pytest.mark.asyncio @pytest.mark.parametrize("may_have_input", [False, True]) async def test_missing_native_history_is_recreated_only_before_possible_input( - code, may_have_input + database_factory, code, may_have_input ): service, harness, directory = code session = await session_for(code) @@ -2128,8 +2171,12 @@ async def test_missing_native_history_is_recreated_only_before_possible_input( session.revision, ) await service.close() + await close_code_database(service) restarted = CodeService( - SQLiteCodeStore(directory / "code.db", directory / "code.lock"), harness + SQLiteCodeStore( + database_factory(directory / "code.db", directory / "code.lock") + ), + harness, ) try: await restarted.start() @@ -2186,3 +2233,4 @@ async def gated_close(): finally: release.set() await service.close() + await close_code_database(service) diff --git a/tests/code_sessions_support.py b/tests/code_sessions_support.py index d6a5be4210..5c8576dbf9 100644 --- a/tests/code_sessions_support.py +++ b/tests/code_sessions_support.py @@ -418,3 +418,11 @@ async def close(self): self.harness.interrupt_gate.set() self.harness.delete_gate.set() await self.sink(HarnessEvent(self.generation, self.thread_id, "closed")) + + +async def close_code_database(service): + """Tests composing a service directly also own its injected database.""" + from free_claude_code.runtime.code_sessions_sqlite import SQLiteCodeStore + + assert isinstance(service._store, SQLiteCodeStore) + await service._store.database.close() diff --git a/tests/conftest.py b/tests/conftest.py index 745ddca16a..3da3ad974b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,10 +1,12 @@ import asyncio import contextlib import os +from pathlib import Path from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest +import pytest_asyncio from free_claude_code.config import env_migrations, paths from free_claude_code.config.loader import clear_settings_cache @@ -212,16 +214,11 @@ def _fire_and_forget(task): @pytest.fixture def mock_session_store(): - from free_claude_code.messaging.session import SessionStore - - store = MagicMock(spec=SessionStore) - store.save_tree = MagicMock() - store.get_tree = MagicMock(return_value=None) - store.register_node = MagicMock() - store.record_message_id = MagicMock() - store.get_tracked_message_ids_for_chat = MagicMock(return_value=[]) - store.forget_tracked_message_ids = MagicMock() - store.clear_scope = MagicMock() + from free_claude_code.messaging.trees import ConversationSnapshot, MessagingStore + + store = AsyncMock(spec=MessagingStore) + store.get_tracked_message_ids_for_chat.return_value = [] + store.load_conversation_snapshot.return_value = ConversationSnapshot() return store @@ -279,3 +276,53 @@ def _propagate_loguru_to_caplog(caplog): loguru_logger.remove( handler_id ) # Handler already removed (e.g. by test_logging_config) + + +@pytest_asyncio.fixture +async def messaging_store_factory(tmp_path): + """Real SQLite stores with fixture-owned lifetime and optional old JSON input.""" + from free_claude_code.runtime.messaging_import import import_legacy + from free_claude_code.runtime.messaging_sqlite import SQLiteMessagingStore + from free_claude_code.runtime.sqlite_database import SQLiteDatabase + + databases = {} + + async def create(*, storage_path=None, managed_message_cap=None): + path = Path(storage_path) if storage_path is not None else tmp_path / "fcc.db" + database_path = path.with_suffix(".db") + database = databases.get(database_path) + first = database is None + if first: + database = SQLiteDatabase(database_path, database_path.with_suffix(".lock")) + await database.start() + databases[database_path] = database + store = SQLiteMessagingStore(database, managed_message_cap=managed_message_cap) + if first and path.suffix == ".json": + await import_legacy(database, path) + await store.trim() + return store + + try: + yield create + finally: + for database in databases.values(): + await database.close() + + +@pytest_asyncio.fixture +async def database_factory(): + """Own database resources for tests that compose individual services.""" + from free_claude_code.runtime.sqlite_database import SQLiteDatabase + + databases = [] + + def create(*args, **kwargs): + database = SQLiteDatabase(*args, **kwargs) + databases.append(database) + return database + + try: + yield create + finally: + for database in reversed(databases): + await database.close() diff --git a/tests/messaging/test_handler.py b/tests/messaging/test_handler.py index 4dac0957f0..09127c5a6c 100644 --- a/tests/messaging/test_handler.py +++ b/tests/messaging/test_handler.py @@ -1,12 +1,12 @@ import asyncio -from unittest.mock import AsyncMock, MagicMock, call, patch +from collections.abc import Callable +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch import pytest from free_claude_code.messaging.command_context import ReplyClearResult, StopOutcome from free_claude_code.messaging.models import MessageScope from free_claude_code.messaging.platforms.ports import MessagingStartupNotice -from free_claude_code.messaging.session import SessionStore from free_claude_code.messaging.trees import ( CancellationReason, CancellationResult, @@ -80,12 +80,7 @@ def _session(events) -> MagicMock: async def _wait_for_idle(workflow: MessagingWorkflow) -> None: - for _ in range(200): - if workflow.tree_queue.task_count() == 0: - await asyncio.sleep(0) - return - await asyncio.sleep(0.01) - raise AssertionError("messaging workflow did not become idle") + await asyncio.wait_for(workflow.tree_queue.wait_idle(), timeout=5) @pytest.fixture @@ -602,7 +597,6 @@ async def test_new_turn_uses_public_admission_and_persists_exact_snapshot( await _wait_for_idle(handler) assert "Launching" in mock_platform.queue_send_message.call_args.args[1] - assert mock_session_store.save_tree_snapshot.call_count >= 2 view = await handler.tree_queue.get_node(incoming.scope, "node_1") assert view is not None assert view.state is MessageState.COMPLETED @@ -763,20 +757,26 @@ async def test_stop_all_applies_immutable_ui_ownership_and_snapshots( effects=(workflow_owned, runner_owned), snapshots=(snapshot,), ) + + async def commit_cancel(*, reason, on_committed): + on_committed() + return result + with patch.object( handler.tree_queue, "cancel_all", - AsyncMock(return_value=result), + AsyncMock(side_effect=commit_cancel), ) as cancel_all: outcome = await handler.stop_all_tasks() await asyncio.sleep(0) assert outcome == _stop_outcome(2) - cancel_all.assert_awaited_once_with(reason=CancellationReason.STOP) + cancel_all.assert_awaited_once_with( + reason=CancellationReason.STOP, on_committed=ANY + ) mock_cli_manager.stop_all.assert_awaited_once() assert mock_platform.fire_and_forget.call_count == 1 assert mock_platform.queue_edit_message.call_args.args[1] == "status_queued" - mock_session_store.save_tree_snapshot.assert_called_once_with(snapshot) @pytest.mark.asyncio @@ -810,10 +810,13 @@ async def cancel_voices() -> tuple[VoiceCancellationResult, ...]: events.append("voices") return (voice,) - async def cancel_trees(*, reason: CancellationReason) -> CancellationResult: + async def cancel_trees( + *, reason: CancellationReason, on_committed: Callable[[], None] + ) -> CancellationResult: assert reason is CancellationReason.STOP assert handler._state_lock.locked() events.append("trees") + on_committed() return tree_result mock_platform.cancel_all_pending_voices.side_effect = cancel_voices @@ -865,16 +868,19 @@ async def block_shutdown() -> None: shutdown_started.set() await release_shutdown.wait() + async def commit_cancel(*, reason, on_committed): + on_committed() + return result + mock_cli_manager.stop_all.side_effect = block_shutdown with patch.object( handler.tree_queue, "cancel_all", - AsyncMock(return_value=result), + AsyncMock(side_effect=commit_cancel), ): stop_task = asyncio.create_task(handler.stop_all_tasks()) await shutdown_started.wait() - mock_session_store.save_tree_snapshot.assert_called_once_with(snapshot) stop_task.cancel() stop_task.cancel() await asyncio.sleep(0) @@ -922,7 +928,9 @@ async def stop_cli() -> None: mock_session_store, platform_name="telegram", ) - workflow._tree_queue = TreeQueueManager(cancellation_delayed_runner) + workflow._tree_queue = TreeQueueManager( + cancellation_delayed_runner, store=AsyncMock() + ) await workflow.handle_message( incoming_message_factory(text="work", message_id="work_1") ) @@ -936,13 +944,11 @@ async def stop_cli() -> None: mock_cli_manager.stop_all.assert_awaited_once() assert close_task.done() is False assert workflow.tree_queue.task_count() == 1 - mock_session_store.flush_pending_save.assert_not_called() release_cleanup.set() await asyncio.wait_for(close_task, timeout=1) assert workflow.tree_queue.task_count() == 0 - mock_session_store.flush_pending_save.assert_called_once() finally: release_cleanup.set() if not close_task.done(): @@ -982,7 +988,6 @@ async def test_node_runner_success_uses_claim_and_semantic_completion( await handler.node_runner.process_node(claim) complete_claim.assert_awaited_once_with(claim, "session_1") - mock_session_store.save_tree_snapshot.assert_called_once_with(snapshot) rendered = mock_platform.queue_edit_message.call_args_list[-1].args[2] assert "โœ… *Complete*" in rendered assert "Hello world" in rendered @@ -1054,10 +1059,6 @@ async def test_session_info_records_real_session_through_manager( ) record_session.assert_awaited_once_with(claim, "real_session") complete_claim.assert_awaited_once_with(claim, "real_session") - assert mock_session_store.save_tree_snapshot.call_args_list == [ - ((record_snapshot,), {}), - ((complete_snapshot,), {}), - ] @pytest.mark.asyncio @@ -1148,10 +1149,6 @@ async def test_non_exit_error_defers_child_failure_until_stream_ends( call(claim, propagate=False), call(claim, propagate=True), ] - assert mock_session_store.save_tree_snapshot.call_args_list == [ - call(snapshot), - call(snapshot), - ] rendered = "\n".join( call.args[2] for call in mock_platform.queue_edit_message.call_args_list ) @@ -1602,11 +1599,12 @@ async def finalize_notice(*args, **kwargs) -> None: @pytest.mark.asyncio async def test_concurrent_global_clear_does_not_wait_for_startup_delivery( + messaging_store_factory, mock_platform, mock_cli_manager, tmp_path, ) -> None: - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) + store = await messaging_store_factory(storage_path=str(tmp_path / "sessions.json")) workflow = MessagingWorkflow( mock_platform, mock_cli_manager, @@ -1638,7 +1636,9 @@ async def send_notice(*args, **kwargs) -> str: release_send.set() await asyncio.gather(publish_task, clear_task, return_exceptions=True) - assert store.get_tracked_message_ids_for_chat(_SCOPE.platform, _SCOPE.chat_id) == [] + assert ( + await store.get_tracked_message_ids_for_chat(_SCOPE.platform, _SCOPE.chat_id) + ) == [] mock_platform.queue_delete_messages.assert_awaited_once_with( _SCOPE.chat_id, ["startup_1"], @@ -1686,11 +1686,12 @@ async def send_notice(*args, **kwargs) -> str: @pytest.mark.asyncio async def test_global_clear_precedes_concurrent_startup_notice_publication( + messaging_store_factory, mock_platform, mock_cli_manager, tmp_path, ) -> None: - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) + store = await messaging_store_factory(storage_path=str(tmp_path / "sessions.json")) workflow = MessagingWorkflow( mock_platform, mock_cli_manager, @@ -1702,11 +1703,15 @@ async def test_global_clear_precedes_concurrent_startup_notice_publication( release_clear = asyncio.Event() async def clear_trees( - scope: MessageScope, *, reason: CancellationReason + scope: MessageScope, + *, + reason: CancellationReason, + on_committed: Callable[[], None], ) -> CancellationResult: assert reason is CancellationReason.CLEAR clear_started.set() await release_clear.wait() + on_committed() return CancellationResult() with patch.object(workflow.tree_queue, "clear_scope", side_effect=clear_trees): @@ -1725,20 +1730,21 @@ async def clear_trees( assert await clear_task == frozenset() await publish_task - assert store.get_tracked_message_ids_for_chat(_SCOPE.platform, _SCOPE.chat_id) == [ - "msg_123" - ] + assert ( + await store.get_tracked_message_ids_for_chat(_SCOPE.platform, _SCOPE.chat_id) + ) == ["msg_123"] mock_platform.queue_delete_messages.assert_not_awaited() @pytest.mark.asyncio async def test_clear_command_cannot_evict_startup_notice_at_managed_message_cap( + messaging_store_factory, mock_platform, mock_cli_manager, incoming_message_factory, tmp_path, ) -> None: - store = SessionStore( + store = await messaging_store_factory( storage_path=str(tmp_path / "sessions.json"), managed_message_cap=1, ) @@ -1749,12 +1755,14 @@ async def test_clear_command_cannot_evict_startup_notice_at_managed_message_cap( platform_name="telegram", voice_cancellation=mock_platform, ) - store.record_message_id( - _SCOPE.platform, - _SCOPE.chat_id, - "100", - "out", - "startup", + ( + await store.record_message_id( + _SCOPE.platform, + _SCOPE.chat_id, + "100", + "out", + "startup", + ) ) incoming = incoming_message_factory( text="/clear", @@ -1799,7 +1807,6 @@ async def test_clear_chat_is_fully_scoped_to_the_invoking_chat( assert handler.get_tree_count() == 1 assert await handler.tree_queue.get_node(root_2.scope, "200") is not None mock_cli_manager.stop_all.assert_not_awaited() - mock_session_store.clear_scope.assert_called_once_with(_SCOPE) @pytest.mark.asyncio @@ -1823,11 +1830,15 @@ async def cancel_voices(scope: MessageScope) -> tuple[VoiceCancellationResult, . return (current,) async def clear_trees( - scope: MessageScope, *, reason: CancellationReason + scope: MessageScope, + *, + reason: CancellationReason, + on_committed: Callable[[], None], ) -> CancellationResult: assert reason is CancellationReason.CLEAR assert handler._state_lock.locked() events.append("trees") + on_committed() return CancellationResult() mock_platform.cancel_pending_voices_in_scope.side_effect = cancel_voices @@ -1844,136 +1855,11 @@ async def clear_trees( @pytest.mark.asyncio -async def test_global_clear_persists_after_tree_detach( - handler, - mock_cli_manager, - mock_session_store, -) -> None: - events: list[str] = [] - - async def clear_trees( - scope: MessageScope, *, reason: CancellationReason - ) -> CancellationResult: - assert scope == _SCOPE - assert reason is CancellationReason.CLEAR - events.append("trees.clear_scope") - return CancellationResult() - - mock_session_store.clear_scope.side_effect = lambda scope: events.append( - f"store.clear_scope:{scope.chat_id}" - ) - - with patch.object(handler.tree_queue, "clear_scope", side_effect=clear_trees): - result = await handler.clear_chat("telegram", "chat_1") - - assert result == frozenset() - assert events == ["trees.clear_scope", "store.clear_scope:chat_1"] - mock_cli_manager.stop_all.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_global_clear_finishes_cleanup_before_final_persistence_error_escapes( - handler, - mock_cli_manager, - mock_session_store, -) -> None: - events: list[str] = [] - - async def clear_trees( - scope: MessageScope, *, reason: CancellationReason - ) -> CancellationResult: - assert scope == _SCOPE - assert reason is CancellationReason.CLEAR - events.append("trees.clear_scope") - return CancellationResult() - - def fail_final_clear(scope: MessageScope) -> None: - assert scope == _SCOPE - events.append("store.clear_scope") - raise OSError("final clear failure") - - mock_session_store.clear_scope.side_effect = fail_final_clear - - with ( - patch.object(handler.tree_queue, "clear_scope", side_effect=clear_trees), - pytest.raises(OSError, match="final clear failure"), - ): - await handler.clear_chat("telegram", "chat_1") - - assert events == ["trees.clear_scope", "store.clear_scope"] - mock_cli_manager.stop_all.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_global_clear_preserves_tree_and_persistence_failures( - handler, - mock_cli_manager, - mock_session_store, -) -> None: - events: list[str] = [] - persistence_error = OSError("final clear failure") - tree_error = RuntimeError("tree clear failure") - - async def fail_tree_clear( - scope: MessageScope, *, reason: CancellationReason - ) -> CancellationResult: - assert scope == _SCOPE - assert reason is CancellationReason.CLEAR - events.append("trees.clear_scope") - raise tree_error - - def fail_final_clear(scope: MessageScope) -> None: - assert scope == _SCOPE - events.append("store.clear_scope") - raise persistence_error - - mock_session_store.clear_scope.side_effect = fail_final_clear - - with ( - patch.object(handler.tree_queue, "clear_scope", side_effect=fail_tree_clear), - pytest.raises(ExceptionGroup) as raised, - ): - await handler.clear_chat("telegram", "chat_1") - - assert raised.value.exceptions == (tree_error, persistence_error) - assert events == ["trees.clear_scope", "store.clear_scope"] - mock_cli_manager.stop_all.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_committed_global_clear_attempts_remaining_steps_after_tree_failure( - handler, - mock_cli_manager, - mock_session_store, -) -> None: - events: list[str] = [] - tree_error = RuntimeError("tree clear failure") - - async def fail_tree_clear( - scope: MessageScope, *, reason: CancellationReason - ) -> CancellationResult: - assert scope == _SCOPE - assert reason is CancellationReason.CLEAR - events.append("trees.clear_scope") - raise tree_error - - mock_session_store.clear_scope.side_effect = lambda scope: events.append( - f"store.clear_scope:{scope.chat_id}" - ) - - with ( - patch.object(handler.tree_queue, "clear_scope", side_effect=fail_tree_clear), - pytest.raises(RuntimeError, match="tree clear failure") as raised, - ): +async def test_global_clear_reports_transaction_failure(handler, mock_session_store): + mock_session_store.commit_trees.side_effect = OSError("write failed") + with pytest.raises(OSError, match="write failed"): await handler.clear_chat("telegram", "chat_1") - assert raised.value is tree_error - assert events == [ - "trees.clear_scope", - "store.clear_scope:chat_1", - ] - mock_cli_manager.stop_all.assert_not_awaited() - @pytest.mark.asyncio async def test_cancelled_global_clear_finishes_owned_transaction_before_propagating( @@ -2010,7 +1896,6 @@ async def block_id_read(platform: str, chat_id: str) -> set[str]: assert handler._clear_generations[_SCOPE] == 1 assert handler.get_tree_count() == 0 - mock_session_store.clear_scope.assert_called_once_with(_SCOPE) mock_cli_manager.stop_all.assert_not_awaited() @@ -2093,7 +1978,6 @@ async def capture_delete(chat_id, message_ids, fire_and_forget=True): assert cleared_prompt.state is MessageState.ERROR assert cleared_prompt.session_id is None assert await handler.tree_queue.get_node(root.scope, "100") is not None - mock_session_store.save_tree_snapshot.assert_called_once() mock_session_store.record_message_id.assert_called_once_with( "telegram", "chat_1", @@ -2128,7 +2012,6 @@ async def test_reply_clear_unknown_reports_nothing_to_clear( "in", "command", ) - mock_session_store.clear_scope.assert_not_called() @pytest.mark.asyncio @@ -2157,9 +2040,6 @@ async def capture_delete(chat_id, message_ids, fire_and_forget=True): ) assert set(deleted_ids) == {"100", "101", "150"} - mock_session_store.remove_tree_snapshot.assert_called_once_with( - TreeIdentity(scope=root.scope, root_id="100") - ) assert handler.get_tree_count() == 0 @@ -2195,14 +2075,13 @@ async def blocking_start(*args, **kwargs): await handler.clear_chat("telegram", "chat_1") - mock_session_store.save_tree_snapshot.assert_not_called() - mock_session_store.clear_scope.assert_called_once_with(_SCOPE) mock_platform.queue_edit_message.assert_not_awaited() assert handler.get_tree_count() == 0 @pytest.mark.asyncio async def test_global_clear_removes_snapshot_saved_during_detach_window( + messaging_store_factory, tmp_path, mock_platform, mock_cli_manager, @@ -2227,7 +2106,7 @@ async def finish_after_release(*args, **kwargs): ) mock_platform.queue_send_message.return_value = "status-new" store_path = tmp_path / "sessions.json" - store = SessionStore(storage_path=str(store_path)) + store = await messaging_store_factory(storage_path=str(store_path)) workflow = MessagingWorkflow( mock_platform, mock_cli_manager, @@ -2258,22 +2137,21 @@ async def block_id_read(platform: str, chat_id: str) -> set[str]: release_runner.set() await _wait_for_idle(workflow) - assert not store.load_conversation_snapshot().is_empty - store.flush_pending_save() - assert ( - not SessionStore(storage_path=str(store_path)) - .load_conversation_snapshot() - .is_empty - ) + assert not (await store.load_conversation_snapshot()).is_empty + assert not ( + await ( + await messaging_store_factory(storage_path=str(store_path)) + ).load_conversation_snapshot() + ).is_empty release_id_read.set() await clear_task - assert store.load_conversation_snapshot().is_empty + assert (await store.load_conversation_snapshot()).is_empty assert ( - SessionStore(storage_path=str(store_path)) - .load_conversation_snapshot() - .is_empty - ) + await ( + await messaging_store_factory(storage_path=str(store_path)) + ).load_conversation_snapshot() + ).is_empty finally: release_runner.set() release_id_read.set() diff --git a/tests/messaging/test_handler_markdown_and_status_edges.py b/tests/messaging/test_handler_markdown_and_status_edges.py index 614dd24af5..4d5509579b 100644 --- a/tests/messaging/test_handler_markdown_and_status_edges.py +++ b/tests/messaging/test_handler_markdown_and_status_edges.py @@ -1,5 +1,5 @@ import asyncio -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest @@ -243,7 +243,6 @@ async def test_node_runner_process_node_session_limit_marks_error_and_updates_ui claim, propagate=False, ) - session_store.save_tree_snapshot.assert_called_once_with(snapshot) @pytest.mark.asyncio @@ -286,7 +285,6 @@ async def _cancelled_start_task(*args, **kwargs): claim, propagate=False, ) - session_store.save_tree_snapshot.assert_called_once_with(snapshot) @pytest.mark.asyncio @@ -314,7 +312,12 @@ async def test_stop_all_tasks_saves_tree_for_cancelled_nodes(): ), snapshots=(snapshot,), ) - cancel_all = AsyncMock(return_value=result) + + async def commit_cancel(*, reason, on_committed): + on_committed() + return result + + cancel_all = AsyncMock(side_effect=commit_cancel) with patch.object( handler.tree_queue, "cancel_all", @@ -326,9 +329,10 @@ async def test_stop_all_tasks_saves_tree_for_cancelled_nodes(): status_feedback_scopes=frozenset({_SCOPE}), fallback_required=False, ) - cancel_all.assert_awaited_once_with(reason=CancellationReason.STOP) + cancel_all.assert_awaited_once_with( + reason=CancellationReason.STOP, on_committed=ANY + ) cli_manager.stop_all.assert_awaited_once() - session_store.save_tree_snapshot.assert_called_once_with(snapshot) @pytest.mark.asyncio @@ -523,6 +527,52 @@ async def test_process_parsed_event_failed_complete_does_not_mark_success(): fail_claim.assert_not_awaited() +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event", + [{"type": "complete", "status": "success"}, {"type": "error", "message": "failed"}], +) +@pytest.mark.parametrize("save_fails", [False, True]) +async def test_terminal_status_waits_for_its_durable_outcome(event, save_fails): + from free_claude_code.messaging.trees import MessagingStorageError + + entered, release = asyncio.Event(), asyncio.Event() + ui = AsyncMock() + + async def save(*_args): + entered.set() + await release.wait() + if save_fails: + raise MessagingStorageError("disk full") + + operation = asyncio.create_task( + process_parsed_cli_event( + parsed=event, + transcript=MagicMock(), + update_ui=ui, + last_status=None, + had_transcript_events=True, + claim=_claim(), + captured_session_id="session", + format_status=lambda *args: " ".join(args), + complete_claim=save, + fail_claim=save, + ) + ) + try: + await asyncio.wait_for(entered.wait(), 1) + ui.assert_not_awaited() + finally: + release.set() + if save_fails: + with pytest.raises(MessagingStorageError): + await operation + ui.assert_not_awaited() + else: + await operation + ui.assert_awaited_once() + + @pytest.mark.asyncio async def test_handler_update_ui_edit_failure_does_not_crash(): """When queue_edit_message raises during streaming, node_runner.process_node continues and completes.""" diff --git a/tests/messaging/test_messaging.py b/tests/messaging/test_messaging.py index b97e8046e0..7fd9443b4d 100644 --- a/tests/messaging/test_messaging.py +++ b/tests/messaging/test_messaging.py @@ -1,8 +1,7 @@ """Tests for messaging/ module.""" import asyncio -import json -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock import pytest @@ -74,121 +73,6 @@ def test_components_bundle_runtime_and_outbound(self): assert components.outbound is outbound -class TestSessionStore: - """Test SessionStore.""" - - def test_session_store_init(self, tmp_path): - """Test SessionStore initialization.""" - from free_claude_code.messaging.session import SessionStore - - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - assert store.load_conversation_snapshot().is_empty - - # --- Tree Tests --- - - def test_save_and_get_tree(self, tmp_path): - """Test saving and retrieving trees.""" - from free_claude_code.messaging.models import MessageScope - from free_claude_code.messaging.session import SessionStore - from free_claude_code.messaging.trees import TreeIdentity, TreeSnapshot - - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - scope = MessageScope(platform="telegram", chat_id="chat") - - tree_data = { - "scope": {"platform": scope.platform, "chat_id": scope.chat_id}, - "root_id": "r1", - "nodes": { - "r1": {"node_id": "r1", "status_message_id": "s1"}, - "n1": {"node_id": "n1", "status_message_id": "s2"}, - }, - } - snapshot = TreeSnapshot.from_json(tree_data) - assert snapshot is not None - store.save_tree_snapshot(snapshot) - - identity = TreeIdentity(scope=scope, root_id="r1") - loaded = store.load_conversation_snapshot().get_tree(identity) - assert loaded is not None - assert loaded == snapshot - assert loaded.lookup_ids() == {"r1", "s1", "n1", "s2"} - - # --- Persistence & Edge Cases --- - - def test_load_existing_file_with_trees(self, tmp_path): - """Test loading file with trees (legacy sessions ignored).""" - from free_claude_code.messaging.models import MessageScope - from free_claude_code.messaging.session import SessionStore - from free_claude_code.messaging.trees import TreeIdentity - - data = { - "sessions": {}, - "trees": { - "r1": { - "root_id": "r1", - "nodes": { - "r1": { - "node_id": "r1", - "incoming": { - "platform": "telegram", - "chat_id": "chat", - }, - } - }, - } - }, - "node_to_tree": {"r1": "r1"}, - "message_log": {}, - } - - p = tmp_path / "sessions.json" - with open(p, "w") as f: - json.dump(data, f) - - store = SessionStore(storage_path=str(p)) - identity = TreeIdentity( - scope=MessageScope(platform="telegram", chat_id="chat"), - root_id="r1", - ) - assert store.load_conversation_snapshot().get_tree(identity) is not None - - def test_load_corrupt_file(self, tmp_path): - """Test loading corrupt/invalid json file.""" - p = tmp_path / "sessions.json" - with open(p, "w") as f: - f.write("{invalid json") - - from free_claude_code.messaging.session import SessionStore - - # Should log error and start empty, avoiding crash - store = SessionStore(storage_path=str(p)) - assert store.load_conversation_snapshot().is_empty - - def test_save_error_handling(self, tmp_path): - """Test error during save.""" - from free_claude_code.messaging.models import MessageScope - from free_claude_code.messaging.session import SessionStore - from free_claude_code.messaging.trees import TreeIdentity, TreeSnapshot - - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - scope = MessageScope(platform="telegram", chat_id="chat") - snapshot = TreeSnapshot(scope=scope, root_id="r1", nodes={"r1": {}}) - store.save_tree_snapshot(snapshot) - - with ( - patch( - "free_claude_code.messaging.session.persistence.os.replace", - side_effect=OSError("Disk full"), - ), - pytest.raises(OSError, match="Disk full"), - ): - store.flush_pending_save() - - assert store.dirty is True - identity = TreeIdentity(scope=scope, root_id="r1") - assert store.load_conversation_snapshot().get_tree(identity) is not None - - class TestTreeQueueManager: """Test TreeQueueManager.""" @@ -198,7 +82,7 @@ def test_tree_queue_manager_init(self): async def process(_claim): return None - mgr = TreeQueueManager(process) + mgr = TreeQueueManager(process, store=AsyncMock()) assert mgr.get_tree_count() == 0 @pytest.mark.asyncio @@ -220,7 +104,7 @@ async def processor(claim): platform="test", ) - mgr = TreeQueueManager(processor) + mgr = TreeQueueManager(processor, store=AsyncMock()) decision = await mgr.admit(incoming, "status_1") assert decision.accepted is True @@ -235,7 +119,7 @@ async def test_cancel_unknown_node_is_empty(self): async def process(_claim): return None - mgr = TreeQueueManager(process) + mgr = TreeQueueManager(process, store=AsyncMock()) scope = MessageScope(platform="test", chat_id="1") cancelled = await mgr.cancel_node(scope, "nonexistent") assert cancelled.effects == () diff --git a/tests/messaging/test_platform_voice_flow.py b/tests/messaging/test_platform_voice_flow.py index 40ee4946ae..826ba5ba15 100644 --- a/tests/messaging/test_platform_voice_flow.py +++ b/tests/messaging/test_platform_voice_flow.py @@ -511,7 +511,7 @@ async def mutate_then_block(tree: MessageTree, node_id: str): assert voice_result is True await asyncio.sleep(0) - mock_session_store.save_tree_snapshot.assert_called() + mock_session_store.commit_trees.assert_awaited() if command == "/clear": assert workflow.get_tree_count() == 0 assert await workflow.tree_queue.resolve_node_id(VOICE_SCOPE, "voice") is None diff --git a/tests/messaging/test_restart_reply_restore.py b/tests/messaging/test_restart_reply_restore.py index d42d0165ac..ddafcca310 100644 --- a/tests/messaging/test_restart_reply_restore.py +++ b/tests/messaging/test_restart_reply_restore.py @@ -5,8 +5,7 @@ import pytest from free_claude_code.messaging.models import IncomingMessage, MessageScope -from free_claude_code.messaging.session import SessionStore -from free_claude_code.messaging.trees import TreeIdentity +from free_claude_code.messaging.trees import MessagingStore, TreeIdentity from free_claude_code.messaging.trees.node import MessageNode, MessageState from free_claude_code.messaging.trees.runtime import MessageTree from free_claude_code.messaging.workflow import MessagingWorkflow @@ -15,14 +14,10 @@ async def _wait_for_idle(workflow: MessagingWorkflow) -> None: - for _ in range(100): - if workflow.tree_queue.task_count() == 0: - return - await asyncio.sleep(0) - raise AssertionError("messaging claims did not finish") + await asyncio.wait_for(workflow.tree_queue.wait_idle(), timeout=5) -async def _write_completed_root(store: SessionStore) -> None: +async def _write_completed_root(store: MessagingStore) -> None: tree = MessageTree( MessageNode( node_id="A", @@ -33,11 +28,10 @@ async def _write_completed_root(store: SessionStore) -> None: session_id="sess_A", ) ) - store.save_tree_snapshot(await tree.snapshot()) - store.flush_pending_save() + await store.commit_trees((await tree.snapshot(),)) -async def _write_interrupted_root(store: SessionStore) -> None: +async def _write_interrupted_root(store: MessagingStore) -> None: tree = MessageTree( MessageNode( node_id="A", @@ -47,8 +41,7 @@ async def _write_interrupted_root(store: SessionStore) -> None: state=MessageState.IN_PROGRESS, ) ) - store.save_tree_snapshot(await tree.snapshot()) - store.flush_pending_save() + await store.commit_trees((await tree.snapshot(),)) def _successful_session(session_id: str): @@ -64,16 +57,19 @@ async def events(*_args, **_kwargs): @pytest.mark.asyncio async def test_reply_to_old_status_message_after_restore_routes_to_parent( + messaging_store_factory, tmp_path, mock_platform, mock_cli_manager, ) -> None: store_path = tmp_path / "sessions.json" - await _write_completed_root(SessionStore(storage_path=str(store_path))) + await _write_completed_root( + await messaging_store_factory(storage_path=str(store_path)) + ) - restored_store = SessionStore(storage_path=str(store_path)) + restored_store = await messaging_store_factory(storage_path=str(store_path)) workflow = MessagingWorkflow(mock_platform, mock_cli_manager, restored_store) - workflow.restore() + (await workflow.restore()) mock_platform.queue_send_message = AsyncMock(return_value="status_reply") mock_cli_manager.get_or_create_session.return_value = ( _successful_session("sess_R1"), @@ -102,6 +98,7 @@ async def test_reply_to_old_status_message_after_restore_routes_to_parent( @pytest.mark.asyncio @pytest.mark.parametrize("wrapped", [False, True]) async def test_legacy_session_json_restores_through_workflow_and_routes_reply( + messaging_store_factory, wrapped: bool, tmp_path, mock_platform, @@ -145,9 +142,9 @@ async def test_legacy_session_json_restores_through_workflow_and_routes_reply( workflow = MessagingWorkflow( mock_platform, mock_cli_manager, - SessionStore(storage_path=str(store_path)), + (await messaging_store_factory(storage_path=str(store_path))), ) - workflow.restore() + (await workflow.restore()) mock_platform.queue_send_message = AsyncMock(return_value="status_reply") mock_cli_manager.get_or_create_session.return_value = ( _successful_session("sess_R1"), @@ -174,30 +171,34 @@ async def test_legacy_session_json_restores_through_workflow_and_routes_reply( @pytest.mark.asyncio async def test_save_tree_snapshot_restores_status_lookup_without_manual_index( + messaging_store_factory, tmp_path, mock_platform, mock_cli_manager, ) -> None: store_path = tmp_path / "sessions.json" - await _write_completed_root(SessionStore(storage_path=str(store_path))) + await _write_completed_root( + await messaging_store_factory(storage_path=str(store_path)) + ) workflow = MessagingWorkflow( mock_platform, mock_cli_manager, - SessionStore(storage_path=str(store_path)), + (await messaging_store_factory(storage_path=str(store_path))), ) - workflow.restore() + (await workflow.restore()) assert await workflow.tree_queue.resolve_node_id(TELEGRAM_CHAT_1, "status_A") == "A" @pytest.mark.asyncio async def test_reply_clear_purges_removed_status_mapping_from_persisted_store( + messaging_store_factory, tmp_path, mock_platform, mock_cli_manager, ) -> None: store_path = tmp_path / "sessions.json" - store = SessionStore(storage_path=str(store_path)) + store = await messaging_store_factory(storage_path=str(store_path)) workflow = MessagingWorkflow(mock_platform, mock_cli_manager, store) mock_platform.queue_send_message = AsyncMock( side_effect=["root_status", "child_status"] @@ -238,10 +239,11 @@ async def test_reply_clear_purges_removed_status_mapping_from_persisted_store( reply_to_message_id="child", ) ) - store.flush_pending_save() identity = TreeIdentity(scope=TELEGRAM_CHAT_1, root_id="root") - persisted = SessionStore(storage_path=str(store_path)).load_conversation_snapshot() + persisted = await ( + await messaging_store_factory(storage_path=str(store_path)) + ).load_conversation_snapshot() tree = persisted.get_tree(identity) assert tree is not None assert tree.lookup_ids() == {"root", "root_status"} @@ -249,19 +251,22 @@ async def test_reply_clear_purges_removed_status_mapping_from_persisted_store( @pytest.mark.asyncio async def test_restore_repairs_interrupted_status_after_delivery_starts( + messaging_store_factory, tmp_path, mock_platform, mock_cli_manager, ) -> None: store_path = tmp_path / "sessions.json" - await _write_interrupted_root(SessionStore(storage_path=str(store_path))) + await _write_interrupted_root( + await messaging_store_factory(storage_path=str(store_path)) + ) workflow = MessagingWorkflow( mock_platform, mock_cli_manager, - SessionStore(storage_path=str(store_path)), + (await messaging_store_factory(storage_path=str(store_path))), platform_name="telegram", ) - workflow.restore() + (await workflow.restore()) await workflow.repair_restored_statuses() await workflow.repair_restored_statuses() @@ -276,7 +281,7 @@ async def test_restore_repairs_interrupted_status_after_delivery_starts( @pytest.mark.asyncio -async def test_workflow_close_waits_for_claim_cleanup_before_flushing( +async def test_workflow_close_stops_cli_and_waits_for_claim_cleanup( mock_platform, mock_cli_manager, mock_session_store, @@ -295,18 +300,15 @@ async def wait_idle() -> None: wait_started.set() await cleanup_release.wait() - workflow.stop_all_tasks = AsyncMock(side_effect=stop_all) + mock_cli_manager.stop_all = AsyncMock(side_effect=stop_all) workflow.tree_queue.wait_idle = AsyncMock(side_effect=wait_idle) - mock_session_store.flush_pending_save.side_effect = lambda: events.append("flush") close_task = asyncio.create_task(workflow.close()) await wait_started.wait() assert events == ["stop", "wait"] - mock_session_store.flush_pending_save.assert_not_called() cleanup_release.set() await close_task - assert events == ["stop", "wait", "flush"] - mock_session_store.flush_pending_save.assert_called_once() + assert events == ["stop", "wait"] diff --git a/tests/messaging/test_session_store_edge_cases.py b/tests/messaging/test_session_store_edge_cases.py deleted file mode 100644 index 077060d17b..0000000000 --- a/tests/messaging/test_session_store_edge_cases.py +++ /dev/null @@ -1,680 +0,0 @@ -"""Edge case tests for the messaging session store.""" - -import json -import threading -from collections.abc import Callable -from typing import Any, ClassVar -from unittest.mock import patch - -import pytest - -import free_claude_code.messaging.session.persistence as persistence_module -from free_claude_code.messaging.models import MessageScope -from free_claude_code.messaging.session import SessionStore -from free_claude_code.messaging.session.persistence import DebouncedJsonPersistence -from free_claude_code.messaging.trees import TreeIdentity, TreeSnapshot - -TELEGRAM_C1 = MessageScope(platform="telegram", chat_id="c1") -TELEGRAM_C2 = MessageScope(platform="telegram", chat_id="c2") - - -def _identity(root_id: str, scope: MessageScope = TELEGRAM_C1) -> TreeIdentity: - return TreeIdentity(scope=scope, root_id=root_id) - - -@pytest.fixture -def tmp_store(tmp_path): - """Create a SessionStore using a temp file.""" - path = str(tmp_path / "sessions.json") - return SessionStore(storage_path=path) - - -def _tree_node(node_id: str, status_message_id: str) -> dict: - return { - "node_id": node_id, - "status_message_id": status_message_id, - } - - -class FakeTimer: - instances: ClassVar[list[FakeTimer]] = [] - - def __init__( - self, - interval: float, - function: Callable[..., None], - args: tuple[Any, ...] | None = None, - kwargs: dict[str, Any] | None = None, - ) -> None: - self.interval = interval - self.function = function - self.args = args or () - self.kwargs = kwargs or {} - self.daemon = False - self.canceled = False - self.started = False - self.instances.append(self) - - def cancel(self) -> None: - self.canceled = True - - def start(self) -> None: - self.started = True - - def fire(self, *, force: bool = False) -> None: - if self.canceled and not force: - return - self.function(*self.args, **self.kwargs) - - -class RecordingPersistence(DebouncedJsonPersistence): - def __init__( - self, - storage_path: str, - *, - snapshot: Callable[[], dict[str, Any]], - on_dirty: Callable[[bool], None], - ) -> None: - self.writes: list[dict[str, Any]] = [] - super().__init__(storage_path, snapshot=snapshot, on_dirty=on_dirty) - - def _write_file(self, data: dict[str, Any]) -> None: - self.writes.append(data) - - -class TestSessionStoreLoadEdgeCases: - """Tests for loading corrupted/malformed data.""" - - def test_load_corrupted_json(self, tmp_path): - """Corrupted JSON file is handled gracefully (logs error, starts empty).""" - path = str(tmp_path / "sessions.json") - with open(path, "w") as f: - f.write("{invalid json") - - store = SessionStore(storage_path=path) - assert store.load_conversation_snapshot().is_empty - - def test_load_truncated_json(self, tmp_path): - """Truncated JSON file is handled gracefully.""" - path = str(tmp_path / "sessions.json") - with open(path, "w") as f: - f.write('{"sessions": {"s1": {"session_id": "s1"') - - store = SessionStore(storage_path=path) - assert store.load_conversation_snapshot().is_empty - - def test_load_empty_file(self, tmp_path): - """Empty file is handled gracefully.""" - path = str(tmp_path / "sessions.json") - with open(path, "w") as f: - f.write("") - - store = SessionStore(storage_path=path) - assert store.load_conversation_snapshot().is_empty - - def test_load_nonexistent_file(self, tmp_path): - """Non-existent file starts with empty state.""" - path = str(tmp_path / "nonexistent.json") - store = SessionStore(storage_path=path) - assert store.load_conversation_snapshot().is_empty - - def test_load_legacy_sessions_ignored(self, tmp_path): - """Legacy sessions in file are ignored; trees and message_log load.""" - path = str(tmp_path / "sessions.json") - data = { - "sessions": { - "s1": { - "session_id": "s1", - "chat_id": 12345, - "initial_msg_id": 100, - "last_msg_id": 200, - "platform": "telegram", - "created_at": "2025-01-01T00:00:00+00:00", - "updated_at": "2025-01-01T00:00:00+00:00", - } - }, - "trees": { - "r1": { - "root_id": "r1", - "nodes": { - "r1": { - "node_id": "r1", - "incoming": { - "platform": TELEGRAM_C1.platform, - "chat_id": TELEGRAM_C1.chat_id, - }, - } - }, - } - }, - "node_to_tree": {"r1": "r1"}, - "message_log": {}, - } - with open(path, "w") as f: - json.dump(data, f) - - store = SessionStore(storage_path=path) - assert store.load_conversation_snapshot().get_tree(_identity("r1")) is not None - - -class TestSessionStoreSaveEdgeCases: - """Tests for save failure handling.""" - - def test_explicit_flush_reports_io_error_and_keeps_store_dirty(self, tmp_store): - """A durability request must not hide that the snapshot was not saved.""" - tmp_store.save_tree_snapshot( - TreeSnapshot(scope=TELEGRAM_C1, root_id="r1", nodes={"r1": {}}) - ) - with ( - patch( - "free_claude_code.messaging.session.persistence.os.replace", - side_effect=OSError("disk full"), - ), - pytest.raises(OSError, match="disk full"), - ): - tmp_store.flush_pending_save() - assert tmp_store.dirty is True - - def test_explicit_flush_snapshot_failure_is_dirty_and_retryable( - self, - tmp_path, - ) -> None: - dirty_states: list[bool] = [] - should_fail = True - - def snapshot() -> dict[str, Any]: - if should_fail: - raise RuntimeError("snapshot failed") - return {"saved": True} - - persistence = DebouncedJsonPersistence( - str(tmp_path / "sessions.json"), - snapshot=snapshot, - on_dirty=dirty_states.append, - ) - - with pytest.raises(RuntimeError, match="snapshot failed"): - persistence.flush() - - assert dirty_states[-1] is True - - should_fail = False - persistence.flush() - - assert dirty_states[-1] is False - - def test_timer_save_io_error_is_best_effort_and_keeps_store_dirty( - self, - tmp_path, - monkeypatch, - ) -> None: - """A background timer may report failure without crashing its thread.""" - FakeTimer.instances = [] - monkeypatch.setattr(persistence_module.threading, "Timer", FakeTimer) - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - store.save_tree_snapshot( - TreeSnapshot(scope=TELEGRAM_C1, root_id="r1", nodes={"r1": {}}) - ) - - with patch( - "free_claude_code.messaging.session.persistence.os.replace", - side_effect=OSError("timer disk full"), - ): - FakeTimer.instances[-1].fire() - - assert store.dirty is True - - store.flush_pending_save() - - assert store.dirty is False - assert ( - SessionStore(storage_path=store.storage_path) - .load_conversation_snapshot() - .get_tree(_identity("r1")) - is not None - ) - - def test_timer_snapshot_failure_is_best_effort_type_only_and_stays_dirty( - self, - tmp_path, - monkeypatch, - ) -> None: - FakeTimer.instances = [] - monkeypatch.setattr(persistence_module.threading, "Timer", FakeTimer) - dirty_states: list[bool] = [] - - def fail_snapshot() -> dict[str, Any]: - raise RuntimeError("SECRET_SNAPSHOT_DETAIL") - - persistence = DebouncedJsonPersistence( - str(tmp_path / "sessions.json"), - snapshot=fail_snapshot, - on_dirty=dirty_states.append, - ) - persistence.schedule_save() - - with patch.object(persistence_module.logger, "error") as error: - FakeTimer.instances[-1].fire() - - assert dirty_states[-1] is True - error.assert_called_once_with( - "Failed to save sessions: exc_type={}", - "RuntimeError", - ) - assert "SECRET_SNAPSHOT_DETAIL" not in str(error.call_args) - - def test_stale_timer_callback_cannot_clear_newer_timer(self, tmp_path, monkeypatch): - """An already-running old timer cannot consume the newest save.""" - FakeTimer.instances = [] - monkeypatch.setattr(persistence_module.threading, "Timer", FakeTimer) - - dirty_states: list[bool] = [] - snapshot_count = 0 - - def snapshot() -> dict[str, Any]: - nonlocal snapshot_count - snapshot_count += 1 - return {"snapshot": snapshot_count} - - persistence = RecordingPersistence( - str(tmp_path / "sessions.json"), - snapshot=snapshot, - on_dirty=dirty_states.append, - ) - - persistence.schedule_save() - first_timer = FakeTimer.instances[0] - persistence.schedule_save() - second_timer = FakeTimer.instances[1] - - first_timer.fire(force=True) - assert persistence.writes == [] - assert dirty_states[-1] is True - assert second_timer.canceled is False - - second_timer.fire() - assert persistence.writes == [{"snapshot": 1}] - assert dirty_states[-1] is False - - def test_running_old_write_finishes_before_newer_flush(self, tmp_path, monkeypatch): - """A claimed old snapshot cannot land after a newer flushed snapshot.""" - FakeTimer.instances = [] - monkeypatch.setattr(persistence_module.threading, "Timer", FakeTimer) - - state = {"version": "old"} - dirty_states: list[bool] = [] - old_write_started = threading.Event() - release_old_write = threading.Event() - - class BlockingPersistence(RecordingPersistence): - def _write_file(self, data: dict[str, Any]) -> None: - if data == {"version": "old"}: - old_write_started.set() - release_old_write.wait(timeout=2) - super()._write_file(data) - - persistence = BlockingPersistence( - str(tmp_path / "sessions.json"), - snapshot=lambda: dict(state), - on_dirty=dirty_states.append, - ) - persistence.schedule_save() - old_writer = threading.Thread(target=FakeTimer.instances[0].fire) - old_writer.start() - assert old_write_started.wait(timeout=2) - - state["version"] = "new" - persistence.schedule_save() - new_writer = threading.Thread(target=persistence.flush) - new_writer.start() - assert new_writer.is_alive() - - release_old_write.set() - old_writer.join(timeout=2) - new_writer.join(timeout=2) - - assert not old_writer.is_alive() - assert not new_writer.is_alive() - assert persistence.writes == [{"version": "old"}, {"version": "new"}] - assert dirty_states[-1] is False - - -class TestSessionStoreTreeSnapshots: - def test_unscoped_tree_without_legacy_ingress_is_reported_and_skipped( - self, tmp_path - ): - path = tmp_path / "sessions.json" - path.write_text( - json.dumps( - { - "conversation": { - "trees": [ - { - "root_id": "root", - "nodes": { - "root": { - "node_id": "root", - "status_message_id": "status", - "state": "completed", - } - }, - } - ] - } - } - ), - encoding="utf-8", - ) - - with patch( - "free_claude_code.messaging.trees.snapshot.logger.warning" - ) as warning: - store = SessionStore(storage_path=str(path)) - - assert store.load_conversation_snapshot().is_empty - warning.assert_called_once_with( - "Skipping messaging tree snapshot without recoverable scope: root_id={}", - "root", - ) - - def test_snapshot_ingress_and_egress_are_deeply_detached(self, tmp_path): - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - snapshot = TreeSnapshot( - scope=TELEGRAM_C1, - root_id="root", - nodes={"root": {"node_id": "root", "state": "completed"}}, - ) - store.save_tree_snapshot(snapshot) - snapshot.nodes["root"]["state"] = "mutated-after-save" - - loaded = store.load_conversation_snapshot() - loaded_tree = loaded.get_tree(_identity("root")) - assert loaded_tree is not None - assert loaded_tree.nodes["root"]["state"] == "completed" - loaded_tree.nodes["root"]["state"] = "mutated-after-load" - - reloaded = store.load_conversation_snapshot().get_tree(_identity("root")) - assert reloaded is not None - assert reloaded.nodes["root"]["state"] == "completed" - - def test_save_tree_replaces_snapshot_for_scoped_identity(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - - store.save_tree_snapshot( - TreeSnapshot( - scope=TELEGRAM_C1, - root_id="root", - nodes={ - "root": _tree_node("root", "root_status"), - "child": _tree_node("child", "child_status"), - }, - ) - ) - - saved = store.load_conversation_snapshot().get_tree(_identity("root")) - assert saved is not None - assert saved.lookup_ids() == { - "root", - "root_status", - "child", - "child_status", - } - - store.save_tree_snapshot( - TreeSnapshot( - scope=TELEGRAM_C1, - root_id="root", - nodes={ - "root": _tree_node("root", "root_status"), - }, - ) - ) - - replaced = store.load_conversation_snapshot().get_tree(_identity("root")) - assert replaced is not None - assert replaced.lookup_ids() == {"root", "root_status"} - - def test_remove_tree_removes_only_scoped_identity(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - store.save_tree_snapshot( - TreeSnapshot( - scope=TELEGRAM_C1, - root_id="root", - nodes={ - "root": _tree_node("root", "root_status"), - "child": _tree_node("child", "child_status"), - }, - ) - ) - - store.remove_tree_snapshot(_identity("root")) - - assert store.load_conversation_snapshot().get_tree(_identity("root")) is None - - -class TestSessionStoreAtomicWrites: - """Atomic persistence: failed replace must not truncate the prior file.""" - - def test_failed_replace_keeps_prior_bytes_and_marks_dirty(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - store.save_tree_snapshot( - TreeSnapshot(scope=TELEGRAM_C1, root_id="r1", nodes={"r1": {}}) - ) - store.flush_pending_save() - with open(path, encoding="utf-8") as f: - disk_after_first = f.read() - - store.save_tree_snapshot( - TreeSnapshot(scope=TELEGRAM_C1, root_id="r2", nodes={"r2": {}}) - ) - - with ( - patch( - "free_claude_code.messaging.session.persistence.os.replace", - side_effect=OSError("replace failed"), - ), - pytest.raises(OSError, match="replace failed"), - ): - store.flush_pending_save() - - with open(path, encoding="utf-8") as f: - disk_after_failed = f.read() - assert disk_after_failed == disk_after_first - assert store.dirty is True - assert store.load_conversation_snapshot().get_tree(_identity("r2")) is not None - - def test_failed_authoritative_clear_is_visible_and_retryable(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - store.save_tree_snapshot( - TreeSnapshot(scope=TELEGRAM_C1, root_id="r1", nodes={"r1": {}}) - ) - store.flush_pending_save() - - with ( - patch( - "free_claude_code.messaging.session.persistence.os.replace", - side_effect=OSError("clear replace failed"), - ), - pytest.raises(OSError, match="clear replace failed"), - ): - store.clear_scope(TELEGRAM_C1) - - assert store.load_conversation_snapshot().is_empty - assert store.dirty is True - assert not SessionStore(storage_path=path).load_conversation_snapshot().is_empty - - store.flush_pending_save() - - assert store.dirty is False - assert SessionStore(storage_path=path).load_conversation_snapshot().is_empty - - def test_authoritative_snapshot_failure_marks_store_dirty_before_retry( - self, - tmp_path, - ) -> None: - store = SessionStore(storage_path=str(tmp_path / "sessions.json")) - - with ( - patch.object( - store, - "_snapshot_for_persistence", - side_effect=RuntimeError("authoritative snapshot failed"), - ), - pytest.raises(RuntimeError, match="authoritative snapshot failed"), - ): - store.clear_scope(TELEGRAM_C1) - - assert store.dirty is True - - store.clear_scope(TELEGRAM_C1) - - assert store.dirty is False - - -class TestSessionStoreClearScope: - def test_clear_scope_wipes_matching_state_and_persists(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - - store.save_tree_snapshot( - TreeSnapshot( - scope=TELEGRAM_C1, - root_id="root1", - nodes={ - "root1": { - "node_id": "root1", - "incoming": { - "text": "hello", - "chat_id": "c1", - "user_id": "u1", - "message_id": "m1", - "platform": "telegram", - "reply_to_message_id": None, - }, - "status_message_id": "status1", - "state": "pending", - "parent_id": None, - "session_id": None, - "children_ids": [], - "created_at": "2025-01-01T00:00:00+00:00", - "completed_at": None, - "error_message": None, - } - }, - ) - ) - - store.clear_scope(TELEGRAM_C1) - - assert store.load_conversation_snapshot().is_empty - - with open(path, encoding="utf-8") as f: - data = json.load(f) - assert data["conversation"]["trees"] == [] - assert data["managed_messages"] == {} - - store2 = SessionStore(storage_path=path) - assert store2.load_conversation_snapshot().is_empty - - def test_clear_scope_preserves_other_chat_trees_and_messages(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - other_tree = TreeSnapshot( - scope=TELEGRAM_C2, - root_id="root2", - nodes={ - "root2": { - "node_id": "root2", - "status_message_id": "status2", - "state": "completed", - "parent_id": None, - "parent_reference_id": None, - "session_id": "session2", - } - }, - ) - store.save_tree_snapshot(other_tree) - store.record_message_id("telegram", "c1", "message1", "in", "prompt") - store.record_message_id("telegram", "c2", "message2", "in", "prompt") - - store.clear_scope(TELEGRAM_C1) - - assert ( - store.load_conversation_snapshot().get_tree(other_tree.identity) is not None - ) - assert store.get_tracked_message_ids_for_chat("telegram", "c1") == [] - assert store.get_tracked_message_ids_for_chat("telegram", "c2") == ["message2"] - restored = SessionStore(storage_path=path) - assert ( - restored.load_conversation_snapshot().get_tree(other_tree.identity) - is not None - ) - assert restored.get_tracked_message_ids_for_chat("telegram", "c2") == [ - "message2" - ] - - def test_managed_message_log_persists_and_dedups(self, tmp_path): - path = str(tmp_path / "sessions.json") - store = SessionStore(storage_path=path) - - store.record_message_id("telegram", "c1", "2", "out", "command") - store.record_message_id("telegram", "c1", "2", "out", "command") - store.record_message_id("telegram", "c1", "3", "in", "command") - - ids = store.get_tracked_message_ids_for_chat("telegram", "c1") - assert ids == ["2", "3"] - - store.flush_pending_save() - store2 = SessionStore(storage_path=path) - assert store2.get_tracked_message_ids_for_chat("telegram", "c1") == [ - "2", - "3", - ] - - def test_load_preserves_legacy_managed_messages(self, tmp_path): - path = tmp_path / "sessions.json" - path.write_text( - json.dumps( - { - "conversation": {"trees": []}, - "message_log": { - "telegram:c1": [ - { - "message_id": "prompt", - "direction": "in", - "kind": "content", - }, - { - "message_id": "old-command", - "direction": "in", - "kind": "command", - }, - { - "message_id": "status", - "direction": "out", - "kind": "status", - }, - { - "message_id": "clear-command", - "direction": "in", - "kind": "clear_command", - }, - ] - }, - } - ), - encoding="utf-8", - ) - - store = SessionStore(storage_path=str(path)) - - assert store.get_tracked_message_ids_for_chat("telegram", "c1") == [ - "prompt", - "old-command", - "status", - "clear-command", - ] diff --git a/tests/messaging/test_storage_controls.py b/tests/messaging/test_storage_controls.py new file mode 100644 index 0000000000..19cfe08424 --- /dev/null +++ b/tests/messaging/test_storage_controls.py @@ -0,0 +1,386 @@ +"""Workflow publication and control effects at real SQLite commit boundaries.""" + +import asyncio +import sqlite3 +from unittest.mock import patch + +import pytest +import pytest_asyncio + +from free_claude_code.messaging.models import IncomingMessage +from free_claude_code.messaging.platforms.ports import MessagingStartupNotice +from free_claude_code.messaging.trees import MessagingStorageError, TreeQueueManager +from free_claude_code.messaging.workflow import MessagingWorkflow + +pytestmark = pytest.mark.asyncio + + +def message(node): + return IncomingMessage( + platform="telegram", chat_id="chat", user_id="user", message_id=node, text=node + ) + + +async def control(workflow, operation): + if operation == "stop": + return await workflow.stop_all_tasks() + return await workflow.clear_chat("telegram", "chat") + + +@pytest.mark.parametrize("operation", ["stop", "clear"]) +@pytest.mark.parametrize("write_fails", [False, True]) +async def test_delayed_prompt_is_invalidated_only_by_committed_control( + messaging_store_factory, mock_platform, mock_cli_manager, operation, write_fails +): + storage = await messaging_store_factory() + workflow = MessagingWorkflow(mock_platform, mock_cli_manager, storage) + release, status_started, status_release = (asyncio.Event() for _ in range(3)) + processed = [] + + async def process(claim): + processed.append(claim.node.node_id) + await release.wait() + await manager.complete_claim(claim, "native") + + manager = TreeQueueManager(process, store=storage) + workflow._tree_queue = manager + await manager.admit(message("existing"), "status-existing") + + async def send(*args, **kwargs): + status_started.set() + await status_release.wait() + return "status-new" + + mock_platform.queue_send_message.side_effect = send + pending = asyncio.create_task(workflow.handle_message(message("new"))) + await status_started.wait() + before = await storage.load_conversation_snapshot() + execute = storage.database.execute + + def reject(operation, *, write=True): + if write: + raise sqlite3.OperationalError("test disk failure") + return execute(operation, write=False) + + try: + if write_fails: + with ( + patch.object(storage.database, "execute", reject), + pytest.raises(MessagingStorageError), + ): + await control(workflow, operation) + assert await storage.load_conversation_snapshot() == before + assert await manager.snapshot() == before + else: + await control(workflow, operation) + status_release.set() + await pending + release.set() + await asyncio.wait_for(manager.wait_idle(), 3) + assert processed.count("new") == (1 if write_fails else 0) + if write_fails: + mock_platform.queue_delete_messages.assert_not_awaited() + else: + mock_platform.queue_delete_messages.assert_awaited_once_with( + "chat", ["status-new"], fire_and_forget=False + ) + finally: + status_release.set() + release.set() + await pending + await workflow.close() + + +async def test_failed_clear_preserves_delayed_startup_notice( + messaging_store_factory, mock_platform, mock_cli_manager +): + storage = await messaging_store_factory() + workflow = MessagingWorkflow( + mock_platform, mock_cli_manager, storage, platform_name="telegram" + ) + started, release = asyncio.Event(), asyncio.Event() + + async def send(*args, **kwargs): + started.set() + await release.wait() + return "notice" + + mock_platform.queue_send_message.side_effect = send + notice = asyncio.create_task( + workflow.publish_startup_notice( + MessagingStartupNotice(chat_id="chat", transport_label="test") + ) + ) + await started.wait() + execute = storage.database.execute + + def reject(operation, *, write=True): + if write: + raise sqlite3.OperationalError("test disk failure") + return execute(operation, write=False) + + try: + with ( + patch.object(storage.database, "execute", reject), + pytest.raises(MessagingStorageError), + ): + await workflow.clear_chat("telegram", "chat") + finally: + release.set() + await notice + await workflow.close() + assert await storage.get_tracked_message_ids_for_chat("telegram", "chat") == [ + "notice" + ] + mock_platform.queue_delete_messages.assert_not_awaited() + + +@pytest_asyncio.fixture +async def interrupted_workflow( + messaging_store_factory, mock_platform, mock_cli_manager, monkeypatch +): + monkeypatch.setattr( + "free_claude_code.messaging.trees.manager.CANCEL_TASK_DRAIN_TIMEOUT_S", 0.01 + ) + storage = await messaging_store_factory() + workflow = MessagingWorkflow(mock_platform, mock_cli_manager, storage) + fail_now, faulted, release = (asyncio.Event() for _ in range(3)) + cancelled, notifications = set(), [] + + async def process(claim): + node = claim.node.node_id + if node == "root": + await fail_now.wait() + try: + await manager.record_session(claim, "rejected-session") + except MessagingStorageError: + faulted.set() + while not release.is_set(): + try: + await release.wait() + except asyncio.CancelledError: + cancelled.add(node) + raise + else: + try: + await release.wait() + except asyncio.CancelledError: + cancelled.add(node) + raise + + manager = TreeQueueManager( + process, store=storage, unexpected_failure_callback=notifications.append + ) + workflow._tree_queue = manager + await manager.admit(message("root"), "status-root") + await manager.admit( + message("child"), "status-child", parent_reference_id="status-root" + ) + await manager.admit( + message("sibling"), "status-sibling", parent_reference_id="root" + ) + await manager.admit(message("healthy"), "status-healthy") + await storage.database.run( + lambda connection: connection.execute( + "CREATE TRIGGER reject_session BEFORE UPDATE ON messaging_nodes " + "WHEN NEW.session_id='rejected-session' " + "BEGIN SELECT RAISE(ABORT, 'test storage failure'); END" + ) + ) + fail_now.set() + await asyncio.wait_for(faulted.wait(), 3) + await storage.database.run( + lambda connection: connection.execute("DROP TRIGGER reject_session") + ) + try: + yield workflow, storage, release, notifications, cancelled + finally: + release.set() + await workflow.close() + + +@pytest.mark.parametrize( + ("operation", "remaining_feedback"), + [ + ("stop-all", set()), + ("stop-root", {"child", "sibling"}), + ("stop-child", {"root", "sibling"}), + ("clear-child", {"root", "sibling"}), + ("clear-status", {"sibling"}), + ], +) +async def test_controls_during_interrupted_cleanup_preserve_ownership_and_feedback( + interrupted_workflow, mock_cli_manager, operation, remaining_feedback +): + workflow, storage, release, notifications, cancelled = interrupted_workflow + manager = workflow.tree_queue + scope = message("root").scope + if operation == "stop-all": + await workflow.stop_all_tasks() + mock_cli_manager.stop_all.assert_awaited_once() + assert cancelled == {"root", "healthy"} + elif operation.startswith("stop-"): + await workflow.stop_reply(scope, operation.removeprefix("stop-")) + else: + result = await workflow.clear_reply( + scope, "child" if operation == "clear-child" else "status-root" + ) + assert result is not None + assert "status-child" in result.delete_message_ids + # Controls do not reopen admission while the interrupted owner is alive. + with pytest.raises(MessagingStorageError, match="cleaning up"): + await manager.admit( + message("later"), "status-later", parent_reference_id="root" + ) + await manager.cancel_node(scope, "healthy") + release.set() + await asyncio.wait_for(manager.wait_idle(), 3) + assert { + target.node_id for result in notifications for target in result.affected + } == remaining_feedback + assert await manager.snapshot() == await storage.load_conversation_snapshot() + + +@pytest.mark.parametrize( + "operation", ["stop-all", "stop-root", "clear-child", "clear-status", "clear-all"] +) +async def test_failed_control_during_cleanup_preserves_other_workers_and_gate( + interrupted_workflow, mock_cli_manager, operation +): + workflow, storage, _release, _notifications, cancelled = interrupted_workflow + manager = workflow.tree_queue + before = await storage.load_conversation_snapshot() + execute = storage.database.execute + + def reject(operation, *, write=True): + if write: + raise sqlite3.OperationalError("test disk failure") + return execute(operation, write=False) + + with ( + patch.object(storage.database, "execute", reject), + pytest.raises(MessagingStorageError), + ): + if operation == "stop-all": + await workflow.stop_all_tasks() + elif operation == "stop-root": + await workflow.stop_reply(message("root").scope, "root") + elif operation == "clear-all": + await workflow.clear_chat("telegram", "chat") + else: + await workflow.clear_reply( + message("root").scope, + "child" if operation == "clear-child" else "status-root", + ) + assert await storage.load_conversation_snapshot() == before + assert await manager.snapshot() == before + assert not cancelled + mock_cli_manager.stop_all.assert_not_awaited() + with pytest.raises(MessagingStorageError, match="cleaning up"): + await manager.admit( + message("later"), "status-later", parent_reference_id="root" + ) + + +@pytest.mark.parametrize( + ("operation", "failure"), + [("stop", "callback"), ("clear", "callback"), ("stop", "cli")], +) +async def test_committed_barrier_survives_postcommit_failure( + messaging_store_factory, mock_platform, mock_cli_manager, operation, failure +): + storage = await messaging_store_factory() + workflow = MessagingWorkflow(mock_platform, mock_cli_manager, storage) + status_started, status_release, effect_started, effect_release = ( + asyncio.Event() for _ in range(4) + ) + processed = [] + + async def process(claim): + processed.append(claim.node.node_id) + await asyncio.Event().wait() + + async def fail_effect(*_args): + effect_started.set() + await effect_release.wait() + if failure == "callback": + raise asyncio.CancelledError("delivery cancelled") + raise RuntimeError("CLI cleanup failed") + + manager = TreeQueueManager( + process, + store=storage, + queue_update_callback=fail_effect if failure == "callback" else None, + ) + workflow._tree_queue = manager + root = await manager.admit(message("root"), "status-root") + assert root.claim is not None + await manager.admit(message("child"), "status-child", parent_reference_id="root") + + async def send(*args, **kwargs): + status_started.set() + await status_release.wait() + return "status-new" + + mock_platform.queue_send_message.side_effect = send + if failure == "cli": + mock_cli_manager.stop_all.side_effect = fail_effect + pending = asyncio.create_task(workflow.handle_message(message("new"))) + await status_started.wait() + command = asyncio.create_task(control(workflow, operation)) + try: + await asyncio.wait_for(effect_started.wait(), 3) + saved = await storage.load_conversation_snapshot() + if operation == "clear": + assert saved.is_empty + else: + assert all( + node["state"] == "error" + for node in saved.trees[root.claim.identity].nodes.values() + ) + effect_release.set() + with pytest.raises( + asyncio.CancelledError if failure == "callback" else RuntimeError + ): + await command + status_release.set() + await pending + assert processed == ["root"] + mock_platform.queue_delete_messages.assert_awaited_once_with( + "chat", ["status-new"], fire_and_forget=False + ) + finally: + effect_release.set() + status_release.set() + await asyncio.gather(pending, command, return_exceptions=True) + mock_cli_manager.stop_all.side_effect = None + await workflow.close() + + +@pytest.mark.parametrize("operation", ["stop", "clear"]) +async def test_successful_empty_control_still_invalidates_delayed_admission( + messaging_store_factory, mock_platform, mock_cli_manager, operation +): + storage = await messaging_store_factory() + workflow = MessagingWorkflow(mock_platform, mock_cli_manager, storage) + started, release = asyncio.Event(), asyncio.Event() + + async def send(*args, **kwargs): + started.set() + await release.wait() + return "status-new" + + mock_platform.queue_send_message.side_effect = send + prompt = asyncio.create_task(workflow.handle_message(message("new"))) + await started.wait() + try: + await control(workflow, operation) + finally: + release.set() + await prompt + await workflow.close() + assert (await storage.load_conversation_snapshot()).is_empty + mock_cli_manager.get_or_create_session.assert_not_awaited() + mock_platform.queue_delete_messages.assert_awaited_once_with( + "chat", ["status-new"], fire_and_forget=False + ) diff --git a/tests/messaging/test_tree_concurrency.py b/tests/messaging/test_tree_concurrency.py index bd0e95181e..aefae20be3 100644 --- a/tests/messaging/test_tree_concurrency.py +++ b/tests/messaging/test_tree_concurrency.py @@ -1,6 +1,7 @@ """Deterministic manager concurrency contracts.""" import asyncio +from unittest.mock import AsyncMock import pytest @@ -67,6 +68,7 @@ async def capture_queue(queue: tuple[QueueEntry, ...]) -> None: process, queue_update_callback=capture_queue, node_started_callback=capture_started, + store=AsyncMock(), ) root = await manager.admit(_incoming("root"), "status-root") assert root.claim is not None @@ -131,7 +133,7 @@ async def process(claim: NodeClaim) -> None: active -= 1 completed[node_id].set() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await asyncio.gather( manager.admit(_incoming("one"), "status-one"), manager.admit(_incoming("two"), "status-two"), @@ -160,7 +162,7 @@ async def process(claim: NodeClaim) -> None: active_started[node_id].set() await asyncio.Event().wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("one"), "status-one") await manager.admit(_incoming("two"), "status-two") await asyncio.gather(*(event.wait() for event in active_started.values())) @@ -184,10 +186,10 @@ async def process(claim: NodeClaim) -> None: "two": CancellationUiOwner.RUNNER, "two-child": CancellationUiOwner.WORKFLOW, } - assert len(result.snapshots) == 2 + assert len((await manager.snapshot()).trees) == 2 assert { node["state"] - for snapshot in result.snapshots + for snapshot in (await manager.snapshot()).trees.values() for node in snapshot.nodes.values() } == {"error"} assert set(processed) == {"one", "two"} @@ -215,7 +217,7 @@ async def process(claim: NodeClaim) -> None: else: unexpected.append(node_id) - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await root_started.wait() await manager.admit( @@ -273,7 +275,7 @@ async def process(claim: NodeClaim) -> None: root_started.set() await asyncio.Event().wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await root_started.wait() await manager.admit( diff --git a/tests/messaging/test_tree_customer_regressions.py b/tests/messaging/test_tree_customer_regressions.py index 45d7e6619b..f77827a193 100644 --- a/tests/messaging/test_tree_customer_regressions.py +++ b/tests/messaging/test_tree_customer_regressions.py @@ -1,6 +1,7 @@ """Customer-visible regressions at the tree manager ownership boundary.""" import asyncio +from unittest.mock import AsyncMock import pytest @@ -43,7 +44,7 @@ async def test_same_telegram_ids_in_different_chats_remain_independent_after_res async def process(_claim: NodeClaim) -> None: return - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) first = await manager.admit(_incoming("chat-a"), "99") second = await manager.admit(_incoming("chat-b"), "99") @@ -61,7 +62,7 @@ async def process(_claim: NodeClaim) -> None: assert snapshot.get_tree(TreeIdentity(scope=chat_a, root_id="42")) is not None assert snapshot.get_tree(TreeIdentity(scope=chat_b, root_id="42")) is not None - restored = TreeQueueManager.from_snapshot(snapshot, process) + restored = TreeQueueManager.from_snapshot(snapshot, process, store=AsyncMock()) assert restored.get_tree_count() == 2 assert _snapshot_chat_ids(await restored.snapshot()) == {"chat-a", "chat-b"} @@ -82,7 +83,7 @@ async def process(_claim: NodeClaim) -> None: started.set() await release.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) decision = await manager.admit(_incoming("chat"), "99") assert decision.claim is not None await started.wait() diff --git a/tests/messaging/test_tree_ownership_concurrency.py b/tests/messaging/test_tree_ownership_concurrency.py index 3b116eda99..710a4f8ea1 100644 --- a/tests/messaging/test_tree_ownership_concurrency.py +++ b/tests/messaging/test_tree_ownership_concurrency.py @@ -1,6 +1,7 @@ import asyncio import contextlib import sys +from unittest.mock import AsyncMock import pytest @@ -67,7 +68,7 @@ async def process(claim: NodeClaim) -> None: finally: active -= 1 - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await root_started.wait() @@ -113,6 +114,7 @@ async def process(claim: NodeClaim) -> None: manager = TreeQueueManager( process, unexpected_failure_callback=unexpected_failures.append, + store=AsyncMock(), ) await manager.admit(_incoming("root"), "status-root") await root_started.wait() @@ -163,7 +165,7 @@ async def process(claim: NodeClaim) -> None: finally: child_exited.set() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) root = await manager.admit(_incoming("root"), "status-root") assert root.claim is not None await root_started.wait() @@ -240,7 +242,7 @@ async def process(claim: NodeClaim) -> None: started.set() await release.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) first = await manager.admit(_incoming("root"), "status-root") duplicate = await manager.admit(_incoming("root"), "status-duplicate") await started.wait() @@ -260,7 +262,7 @@ async def test_node_and_status_references_are_published_together() -> None: async def process(_claim: NodeClaim) -> None: await release.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") assert await manager.resolve_node_id(_SCOPE, "root") == "root" @@ -292,6 +294,7 @@ async def broken_started_callback(_claim) -> None: process, queue_update_callback=broken_queue_callback, node_started_callback=broken_started_callback, + store=AsyncMock(), ) await manager.admit(_incoming("root"), "status-root") await manager.admit( @@ -315,7 +318,7 @@ async def process(claim: NodeClaim) -> None: root_started.set() await release_root.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await root_started.wait() await manager.admit( @@ -355,7 +358,7 @@ async def process(claim: NodeClaim) -> None: finally: cleanup_finished.set() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await terminal.wait() @@ -387,7 +390,7 @@ async def process(_claim: NodeClaim) -> None: cleanup_finished.set() raise - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") checkpoint = asyncio.Event() asyncio.get_running_loop().call_soon(checkpoint.set) @@ -429,7 +432,7 @@ async def process(_claim: NodeClaim) -> None: started.set() await release.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) loop = asyncio.get_running_loop() original_factory = loop.get_task_factory() try: @@ -467,7 +470,7 @@ def portable_create_task(coro, *, name=None, context=None): return original_create_task(coro, name=name, context=context) monkeypatch.setattr(loop, "create_task", portable_create_task) - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) decision = await manager.admit(_incoming("root"), "status-root") await asyncio.wait_for(manager.wait_idle(), timeout=1) @@ -489,7 +492,7 @@ async def scenario() -> None: async def process(_claim: NodeClaim) -> None: processed.set() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) decision = await manager.admit(_incoming("root"), "status-root") await asyncio.wait_for(manager.wait_idle(), timeout=1) @@ -523,7 +526,9 @@ async def announce_started(claim: NodeClaim) -> None: except asyncio.CancelledError: callback_absorbed_cancellation.set() - manager = TreeQueueManager(process, node_started_callback=announce_started) + manager = TreeQueueManager( + process, node_started_callback=announce_started, store=AsyncMock() + ) await manager.admit(_incoming("root"), "status-root") await root_started.wait() child = await manager.admit( @@ -576,7 +581,7 @@ async def process(claim: NodeClaim) -> None: new_started.set() await release_new.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) old = await manager.admit(_incoming("root"), "status-root") await old_started.wait() @@ -628,7 +633,7 @@ async def process(_claim: NodeClaim) -> None: second_cancellation.set() raise - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await started.wait() diff --git a/tests/messaging/test_tree_processor.py b/tests/messaging/test_tree_processor.py index 277688a79d..cab7123ef6 100644 --- a/tests/messaging/test_tree_processor.py +++ b/tests/messaging/test_tree_processor.py @@ -2,6 +2,7 @@ import asyncio import logging +from unittest.mock import AsyncMock import pytest @@ -57,7 +58,7 @@ async def process(_claim: NodeClaim) -> None: started.set() await asyncio.Event().wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await started.wait() @@ -96,7 +97,9 @@ async def capture_queue(queue: tuple[QueueEntry, ...]) -> None: tuple((entry.node.node_id, entry.position) for entry in queue) ) - manager = TreeQueueManager(process, queue_update_callback=capture_queue) + manager = TreeQueueManager( + process, queue_update_callback=capture_queue, store=AsyncMock() + ) await manager.admit(_incoming("root"), "status-root") await root_started.wait() decision = await manager.admit( @@ -142,7 +145,7 @@ async def process(_claim: NodeClaim) -> None: await release_cleanup.wait() raise - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await manager.admit(_incoming("root"), "status-root") await started.wait() @@ -186,6 +189,7 @@ async def capture_queue(queue: tuple[QueueEntry, ...]) -> None: process, queue_update_callback=capture_queue, unexpected_failure_callback=failures.append, + store=AsyncMock(), ) await manager.admit(_incoming("root"), "status-root") await started.wait() @@ -222,7 +226,7 @@ async def process(claim: NodeClaim) -> None: child_started.set() await release_child.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) await asyncio.wait_for(manager.wait_idle(), timeout=0.1) await manager.admit(_incoming("root"), "status-root") await asyncio.wait_for(root_started.wait(), timeout=1) @@ -426,6 +430,7 @@ async def process(_claim: NodeClaim) -> None: manager = TreeQueueManager( process, log_messaging_error_details=log_messaging_error_details, + store=AsyncMock(), ) with caplog.at_level(logging.ERROR): await manager.admit(_incoming("root"), "status-root") diff --git a/tests/messaging/test_tree_repository.py b/tests/messaging/test_tree_repository.py index 18e817c445..ff4f134b3a 100644 --- a/tests/messaging/test_tree_repository.py +++ b/tests/messaging/test_tree_repository.py @@ -1,6 +1,7 @@ """Persistence and restored-manager contracts for messaging trees.""" import asyncio +from unittest.mock import AsyncMock import pytest @@ -95,7 +96,9 @@ async def test_restore_reconciles_interrupted_nodes_before_manager_exposure() -> async def process(claim: NodeClaim) -> None: processed.append(claim.node.node_id) - manager = TreeQueueManager.from_snapshot(_interrupted_conversation(), process) + manager = TreeQueueManager.from_snapshot( + _interrupted_conversation(), process, store=AsyncMock() + ) assert len(manager.restored_stale_targets) == 2 assert manager.restored_snapshot is not None @@ -144,7 +147,9 @@ def test_restore_rejects_tree_that_violates_scoped_reference_invariants( else: tree.nodes["failed"]["status_message_id"] = "status-pending" - manager = TreeQueueManager.from_snapshot(snapshot, lambda _claim: asyncio.sleep(0)) + manager = TreeQueueManager.from_snapshot( + snapshot, lambda _claim: asyncio.sleep(0), store=AsyncMock() + ) assert manager.get_tree_count() == 0 assert manager.restored_snapshot == ConversationSnapshot() @@ -159,7 +164,7 @@ async def process(_claim: NodeClaim) -> None: started.set() await release.wait() - manager = TreeQueueManager(process) + manager = TreeQueueManager(process, store=AsyncMock()) decision = await manager.admit(_incoming("root"), "status-root") assert decision.claim is not None await started.wait() diff --git a/tests/runtime/test_application_runtime.py b/tests/runtime/test_application_runtime.py index 31b1d051cb..1feb66d1af 100644 --- a/tests/runtime/test_application_runtime.py +++ b/tests/runtime/test_application_runtime.py @@ -7,7 +7,6 @@ import pytest -import free_claude_code.messaging.session.persistence as persistence_module from free_claude_code.application.connected_accounts import ( ConnectedAccountState, ConnectedAccountStatus, @@ -25,8 +24,6 @@ MessagingPlatformComponents, MessagingStartupNotice, ) -from free_claude_code.messaging.session import SessionStore -from free_claude_code.messaging.workflow import MessagingWorkflow from free_claude_code.providers.base import BaseProvider from free_claude_code.providers.credential_validation import ( CredentialCheck, @@ -854,92 +851,6 @@ async def close_workflow() -> None: assert runtime._closed is True -@pytest.mark.asyncio -async def test_close_retries_real_workflow_persistence_without_losing_latest_state( - tmp_path: Path, -) -> None: - events: list[str] = [] - manager = ProviderRuntimeManager(_settings("nvidia_nim/model")) - runtime = ApplicationRuntime( - manager, configuration=AsyncMock(spec=ConfigurationService), transcriber=None - ) - messaging = TrackingMessagingRuntime(events) - cli_manager = MagicMock() - cli_manager.stop_all = AsyncMock() - outbound = MagicMock() - store_path = tmp_path / "sessions.json" - real_replace = persistence_module.os.replace - replace_calls = 0 - - def fail_first_replace(source: str, target: str) -> None: - nonlocal replace_calls - replace_calls += 1 - if replace_calls == 1: - raise OSError("replace failed once") - real_replace(source, target) - - with patch.object(persistence_module.threading, "Timer"): - store = SessionStore(storage_path=str(store_path)) - store.record_message_id( - "telegram", - "chat_1", - "old", - "out", - "status", - ) - store.flush_pending_save() - store.record_message_id( - "telegram", - "chat_1", - "latest", - "out", - "status", - ) - workflow = MessagingWorkflow(outbound, cli_manager, store) - runtime._messaging_runtime = messaging - runtime._messaging_workflow = workflow - runtime._cli_manager = cli_manager - - with patch.object( - persistence_module.os, - "replace", - side_effect=fail_first_replace, - ): - assert await runtime.close() is False - - assert runtime._messaging_runtime is messaging - assert runtime._messaging_workflow is workflow - assert runtime._cli_manager is cli_manager - assert runtime.is_closed is False - assert store.dirty is True - assert store.get_tracked_message_ids_for_chat("telegram", "chat_1") == [ - "old", - "latest", - ] - assert SessionStore( - storage_path=str(store_path) - ).get_tracked_message_ids_for_chat("telegram", "chat_1") == ["old"] - - assert await runtime.close() is True - - assert runtime._messaging_runtime is None - assert runtime._messaging_workflow is None - assert runtime._cli_manager is None - assert runtime.is_closed is True - assert store.dirty is False - assert SessionStore( - storage_path=str(store_path) - ).get_tracked_message_ids_for_chat("telegram", "chat_1") == [ - "old", - "latest", - ] - assert events == [ - "messaging.quiesce", - "messaging.quiesce", - "messaging.close", - ] - - @pytest.mark.asyncio async def test_cancelled_transcriber_close_retains_ownership() -> None: events: list[str] = [] @@ -1279,7 +1190,7 @@ async def test_composition_publishes_startup_notice_after_runtime_and_repair() - ) workflow = MagicMock() workflow.handle_message = AsyncMock() - workflow.restore.side_effect = lambda: events.append("workflow.restore") + workflow.restore = AsyncMock(side_effect=lambda: events.append("workflow.restore")) workflow.repair_restored_statuses = AsyncMock( side_effect=lambda: events.append("workflow.repair") ) @@ -1294,7 +1205,7 @@ async def test_composition_publishes_startup_notice_after_runtime_and_repair() - "free_claude_code.cli.managed.ManagedClaudeSessionManager", return_value=cli_manager, ) as manager_constructor, - patch("free_claude_code.messaging.session.SessionStore"), + patch.object(runtime, "_messaging_store", AsyncMock()), patch( "free_claude_code.messaging.workflow.MessagingWorkflow", return_value=workflow, diff --git a/tests/runtime/test_code_sessions_sqlite.py b/tests/runtime/test_code_sessions_sqlite.py index d610015078..49ae7da4f2 100644 --- a/tests/runtime/test_code_sessions_sqlite.py +++ b/tests/runtime/test_code_sessions_sqlite.py @@ -18,24 +18,26 @@ CodeUnavailableError, ) from free_claude_code.runtime.code_sessions_sqlite import SQLiteCodeStore +from free_claude_code.runtime.sqlite_database import SQLiteDatabase @pytest.mark.asyncio async def test_cancelled_initialization_drains_thread_before_releasing_owner_lock( + database_factory, tmp_path, ): entered = threading.Event() release = threading.Event() initialized = threading.Event() - class GatedStore(SQLiteCodeStore): + class GatedDatabase(SQLiteDatabase): def _initialize(self): entered.set() assert release.wait(5) super()._initialize() initialized.set() - store = GatedStore(tmp_path / "code.db", tmp_path / "code.lock") + store = SQLiteCodeStore(GatedDatabase(tmp_path / "code.db", tmp_path / "code.lock")) starting = asyncio.create_task(store.start()) await asyncio.to_thread(entered.wait, 3) starting.cancel() @@ -43,7 +45,9 @@ def _initialize(self): with pytest.raises(asyncio.CancelledError): await starting await asyncio.to_thread(initialized.wait, 3) - second = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") + second = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) try: await second.start() session = await second.create( @@ -52,26 +56,43 @@ def _initialize(self): assert (await second.get_session(session.id)).id == session.id finally: await store.close() + await store.database.close() await second.close() + await second.database.close() @pytest.mark.asyncio -async def test_store_has_one_process_owner_and_can_reopen_after_close(tmp_path): - first = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") - second = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") +async def test_store_has_one_process_owner_and_can_reopen_after_close( + database_factory, tmp_path +): + first = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) + second = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await first.start() try: with pytest.raises(CodeUnavailableError, match="another FCC"): await second.start() finally: await first.close() + await first.database.close() + second = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await second.start() await second.close() + await second.database.close() @pytest.mark.asyncio -async def test_reused_item_id_cannot_overwrite_another_session(tmp_path): - store = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") +async def test_reused_item_id_cannot_overwrite_another_session( + database_factory, tmp_path +): + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await store.start() try: first, second = [ @@ -99,6 +120,7 @@ async def test_reused_item_id_cannot_overwrite_another_session(tmp_path): assert (await store.items(second.id, None, None))[0].text == "second" finally: await store.close() + await store.database.close() async def _admit(store, session, *, run_id=None, text="message", sequence=1): @@ -122,11 +144,14 @@ async def _admit(store, session, *, run_id=None, text="message", sequence=1): @pytest_asyncio.fixture -async def store(tmp_path): - result = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") +async def store(database_factory, tmp_path): + result = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await result.start() yield result await result.close() + await result.database.close() async def _session(store): @@ -212,10 +237,13 @@ async def test_mode_is_guarded_by_sqlite_busy_and_run_immutability(store): @pytest.mark.asyncio -async def test_version_one_database_gains_mode_without_losing_history(store, tmp_path): +async def test_version_one_database_gains_mode_without_losing_history( + database_factory, store, tmp_path +): session, run = await _admit(store, await _session(store)) items = await store.items(session.id, None, None) await store.close() + await store.database.close() with closing(sqlite3.connect(tmp_path / "code.db")) as connection, connection: connection.execute("ALTER TABLE code_sessions DROP COLUMN mode") connection.execute( @@ -223,16 +251,21 @@ async def test_version_one_database_gains_mode_without_losing_history(store, tmp ) connection.execute("ALTER TABLE code_sessions DROP COLUMN context_used_tokens") connection.execute("ALTER TABLE code_runs DROP COLUMN mode") + drop_messaging_schema(connection) connection.execute("PRAGMA user_version = 1") for _ in range(2): + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await store.start() saved = await store.get_session(session.id) assert saved.mode == "config" assert saved.native_permission_defaults is None - assert (await store.get_run(session.id, run.id)).mode == "config" + saved_run = await store.get_run(session.id, run.id) + assert saved_run is not None and saved_run.mode == "config" assert await store.items(session.id, None, None) == items with closing(sqlite3.connect(tmp_path / "code.db")) as connection: - assert connection.execute("PRAGMA user_version").fetchone()[0] == 3 + assert connection.execute("PRAGMA user_version").fetchone()[0] == 4 with pytest.raises(sqlite3.IntegrityError): connection.execute("UPDATE code_sessions SET mode = 'unknown'") with pytest.raises(sqlite3.IntegrityError): @@ -240,20 +273,28 @@ async def test_version_one_database_gains_mode_without_losing_history(store, tmp "UPDATE code_sessions SET native_permission_defaults = '[]'" ) await store.close() + await store.database.close() @pytest.mark.asyncio -async def test_version_two_database_gains_nullable_context_usage(store, tmp_path): +async def test_version_two_database_gains_nullable_context_usage( + database_factory, store, tmp_path +): session = await _session(store) await store.close() + await store.database.close() with closing(sqlite3.connect(tmp_path / "code.db")) as connection, connection: connection.execute("ALTER TABLE code_sessions DROP COLUMN context_used_tokens") + drop_messaging_schema(connection) connection.execute("PRAGMA user_version = 2") for _ in range(2): + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await store.start() assert (await store.get_session(session.id)).context_used_tokens is None with closing(sqlite3.connect(tmp_path / "code.db")) as connection: - assert connection.execute("PRAGMA user_version").fetchone()[0] == 3 + assert connection.execute("PRAGMA user_version").fetchone()[0] == 4 column = next( row for row in connection.execute("PRAGMA table_info(code_sessions)") @@ -266,6 +307,7 @@ async def test_version_two_database_gains_nullable_context_usage(store, tmp_path (session.id,), ) await store.close() + await store.database.close() @pytest.mark.asyncio @@ -522,6 +564,7 @@ async def _legacy_prompt_database(store, path): ) original = await store.items(session.id, None, None) await store.close() + await store.database.close() cases: list[ tuple[str | None, str | None, Literal["resolved", "expired", "pending"]] ] = [ @@ -556,6 +599,7 @@ async def _legacy_prompt_database(store, path): ) connection.execute("ALTER TABLE code_sessions DROP COLUMN context_used_tokens") connection.execute("ALTER TABLE code_runs DROP COLUMN mode") + drop_messaging_schema(connection) connection.execute("PRAGMA user_version = 0") for prompt in prompts: connection.execute( @@ -580,11 +624,14 @@ async def _legacy_prompt_database(store, path): @pytest.mark.asyncio async def test_legacy_prompts_migrate_once_to_run_ends_without_resequencing( - store, tmp_path + database_factory, store, tmp_path ): path = tmp_path / "code.db" session, old, latest, original, prompts = await _legacy_prompt_database(store, path) for _ in range(2): + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await store.start() items = await store.items(session.id, None, None) assert [item.id for item in items] == [ @@ -614,18 +661,19 @@ async def test_legacy_prompts_migrate_once_to_run_ends_without_resequencing( ) assert saved[prompt.id] == expected with closing(sqlite3.connect(path)) as connection: - assert connection.execute("PRAGMA user_version").fetchone()[0] == 3 + assert connection.execute("PRAGMA user_version").fetchone()[0] == 4 assert connection.execute("PRAGMA foreign_key_check").fetchall() == [] assert any( row[2] == "code_items" for row in connection.execute("PRAGMA foreign_key_list(code_prompts)") ) await store.close() + await store.database.close() @pytest.mark.asyncio async def test_failed_prompt_migration_rolls_back_items_schema_and_version( - store, tmp_path + database_factory, store, tmp_path ): path = tmp_path / "code.db" _, _, _, original, prompts = await _legacy_prompt_database(store, path) @@ -637,6 +685,9 @@ async def test_failed_prompt_migration_rolls_back_items_schema_and_version( + "' BEGIN SELECT RAISE(ABORT, 'blocked'); END" ) with pytest.raises(CodeConflictError): + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await store.start() with closing(sqlite3.connect(path)) as connection, connection: assert connection.execute("PRAGMA user_version").fetchone()[0] == 0 @@ -651,7 +702,21 @@ async def test_failed_prompt_migration_rolls_back_items_schema_and_version( for row in connection.execute("PRAGMA foreign_key_list(code_prompts)") ) connection.execute("DROP TRIGGER refuse_second_prompt") + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) await store.start() assert len(await store.items(prompts[0].session_id, None, None)) == len( original ) + len(prompts) + + +def drop_messaging_schema(connection): + for table in ( + "messaging_references", + "messaging_nodes", + "messaging_trees", + "messaging_managed_messages", + "messaging_legacy_import", + ): + connection.execute(f"DROP TABLE {table}") diff --git a/tests/runtime/test_codex_app_server.py b/tests/runtime/test_codex_app_server.py index f25ea756a4..e398225ae0 100644 --- a/tests/runtime/test_codex_app_server.py +++ b/tests/runtime/test_codex_app_server.py @@ -25,7 +25,7 @@ CodexAppServer, CodexHarnessFactory, ) -from tests.code_sessions_support import FakeHarness +from tests.code_sessions_support import FakeHarness, close_code_database @dataclass @@ -200,7 +200,7 @@ async def test_history_open_needs_no_fcc_setup_or_inventory(tmp_path, factory_pe @pytest.mark.asyncio async def test_child_warning_during_shutdown_allows_connection_replacement( - tmp_path, monkeypatch + database_factory, tmp_path, monkeypatch ): harness = FakeHarness() prepare = harness.prepare @@ -242,7 +242,8 @@ async def receive(event): monkeypatch.setattr(harness, "prepare", selection_for) service = CodeService( - SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock"), harness + SQLiteCodeStore(database_factory(tmp_path / "code.db", tmp_path / "code.lock")), + harness, ) await service.start() observing = asyncio.create_task(warning.wait()) @@ -286,6 +287,7 @@ async def receive(event): observing.cancel() await asyncio.gather(observing, return_exceptions=True) await service.close() + await close_code_database(service) @pytest.mark.asyncio @@ -468,7 +470,7 @@ async def test_jsonl_large_unicode_events_can_precede_rpc_ack(tmp_path): @pytest.mark.asyncio async def test_terminal_storage_failure_drains_the_native_dispatcher( - tmp_path, monkeypatch + database_factory, tmp_path, monkeypatch ): harness = FakeHarness() prepare = harness.prepare @@ -498,7 +500,9 @@ async def open_native(cwd, sink): return selection monkeypatch.setattr(harness, "prepare", select) - store = SQLiteCodeStore(tmp_path / "code.db", tmp_path / "code.lock") + store = SQLiteCodeStore( + database_factory(tmp_path / "code.db", tmp_path / "code.lock") + ) save = store.save_progress rejected = [] @@ -532,6 +536,7 @@ async def reject_terminal(session, revision, **values): await service.get_detail(session.id) finally: await asyncio.wait_for(service.close(), 5) + await close_code_database(service) @pytest.mark.asyncio diff --git a/tests/runtime/test_messaging_sqlite.py b/tests/runtime/test_messaging_sqlite.py new file mode 100644 index 0000000000..e331726d90 --- /dev/null +++ b/tests/runtime/test_messaging_sqlite.py @@ -0,0 +1,985 @@ +import asyncio +import json +import sqlite3 +import threading +from contextlib import closing, suppress +from unittest.mock import AsyncMock, Mock, patch + +import pytest +import pytest_asyncio + +from free_claude_code.messaging.models import IncomingMessage, MessageScope +from free_claude_code.messaging.trees import MessagingStorageError, TreeQueueManager +from free_claude_code.messaging.trees.snapshot import TreeSnapshot +from free_claude_code.messaging.workflow import MessagingWorkflow +from free_claude_code.runtime.code_sessions_sqlite import SQLiteCodeStore +from free_claude_code.runtime.messaging_import import import_legacy +from free_claude_code.runtime.messaging_sqlite import SQLiteMessagingStore +from free_claude_code.runtime.sqlite_database import SQLiteDatabase + +pytestmark = pytest.mark.asyncio + + +def tree(platform="telegram", chat="chat", root="root"): + return TreeSnapshot( + scope=MessageScope(platform=platform, chat_id=chat), + root_id=root, + nodes={ + root: { + "node_id": root, + "status_message_id": f"status-{root}", + "state": "completed", + "parent_id": None, + "parent_reference_id": None, + "session_id": "native-session", + } + }, + ) + + +@pytest_asyncio.fixture +async def storage(tmp_path): + database = SQLiteDatabase(tmp_path / "fcc.db", tmp_path / "code.lock") + await database.start() + try: + yield SQLiteMessagingStore(database, managed_message_cap=2) + finally: + await database.close() + + +async def test_tree_and_managed_messages_are_durable_and_scope_isolated(storage): + first, second = tree(), tree("discord") + await storage.commit_trees((first, second)) + for message in ("one", "two", "two", "three"): + await storage.record_message_id("telegram", "chat", message, "in", "prompt") + assert await storage.get_tracked_message_ids_for_chat("telegram", "chat") == [ + "two", + "three", + ] + await storage.commit_trees((), clear_scope=first.scope) + loaded = await storage.load_conversation_snapshot() + assert list(loaded.trees) == [second.identity] + assert loaded.trees[second.identity].nodes["root"]["session_id"] == "native-session" + assert not await storage.get_tracked_message_ids_for_chat("telegram", "chat") + + +async def test_reference_collision_rolls_back_whole_write(storage): + original = tree() + await storage.commit_trees((original,)) + collision = tree(root="status-root") + with pytest.raises(MessagingStorageError) as failure: + await storage.commit_trees((collision,)) + assert isinstance(failure.value.__cause__, sqlite3.IntegrityError) + assert (await storage.load_conversation_snapshot()).trees == { + original.identity: original + } + + +@pytest.mark.parametrize("damaged", [False, True]) +async def test_import_keeps_valid_data_and_never_replays_after_clear( + storage, tmp_path, damaged +): + source = tmp_path / "sessions.json" + valid = tree() + source.write_text( + json.dumps( + { + "conversation": { + "trees": [valid.to_json(), *([{"broken": True}] if damaged else [])] + } + } + ) + ) + warning = await import_legacy(storage.database, source) + assert bool(warning) is damaged + assert source.exists() is damaged + assert (await storage.load_conversation_snapshot()).trees == {valid.identity: valid} + await storage.commit_trees((), clear_scope=valid.scope) + assert await import_legacy(storage.database, source) is None + assert (await storage.load_conversation_snapshot()).is_empty + + +async def test_unreadable_json_does_not_block_new_conversations(storage, tmp_path): + source = tmp_path / "sessions.json" + source.write_text('{"conversation":') + assert await import_legacy(storage.database, source) + assert source.read_text() == '{"conversation":' + await storage.commit_trees((tree(),)) + assert not (await storage.load_conversation_snapshot()).is_empty + assert await import_legacy(storage.database, source) is None + + +async def test_pending_import_keeps_restart_repair_information(storage, tmp_path): + pending = tree() + pending.nodes["root"]["state"] = "in_progress" + source = tmp_path / "sessions.json" + source.write_text(json.dumps(pending_to_legacy(pending))) + await import_legacy(storage.database, source) + assert (await storage.load_conversation_snapshot()).trees[pending.identity].nodes[ + "root" + ]["state"] == "in_progress" + + +def pending_to_legacy(snapshot): + return {"trees": {snapshot.root_id: snapshot.to_json()}} + + +async def test_new_schema_preserves_foreign_keys(storage): + await storage.commit_trees((tree(),)) + with closing(sqlite3.connect(storage.database.path)) as connection: + assert connection.execute("PRAGMA foreign_key_check").fetchall() == [] + + +async def test_sqlite_prevents_deleting_a_root_without_its_tree(storage): + original = tree() + await storage.commit_trees((original,)) + with pytest.raises(sqlite3.IntegrityError): + await storage.database.run( + lambda connection: connection.execute("DELETE FROM messaging_nodes") + ) + assert (await storage.load_conversation_snapshot()).trees == { + original.identity: original + } + + +@pytest.mark.parametrize("wrapped", [False, True]) +@pytest.mark.parametrize("dictionary", [False, True]) +async def test_historical_formats_and_message_log_preserve_resume_and_order( + storage, tmp_path, wrapped, dictionary +): + original = tree() + legacy = original.to_json() + del legacy["scope"] + legacy["nodes"]["root"]["incoming"] = {"platform": "telegram", "chat_id": "chat"} + trees = {"old": legacy} if dictionary else [legacy] + conversation = {"trees": trees} + payload = {"conversation": conversation} if wrapped else conversation + payload["message_log"] = { + "telegram:chat": [ + {"message_id": value, "direction": "out", "kind": "notice"} + for value in (1, 2, 2, 3) + ] + } + source = tmp_path / "sessions.json" + source.write_text(json.dumps(payload)) + assert await import_legacy(storage.database, source) is None + assert not source.exists() + loaded = (await storage.load_conversation_snapshot()).trees[original.identity] + assert loaded.nodes["root"]["session_id"] == "native-session" + await storage.trim() + assert await storage.get_tracked_message_ids_for_chat("telegram", "chat") == [ + "2", + "3", + ] + + +async def test_cleanup_failure_cannot_resurrect_cleared_history(storage, tmp_path): + source = tmp_path / "sessions.json" + original = tree() + source.write_text(json.dumps(pending_to_legacy(original))) + with patch.object(type(source), "unlink", side_effect=PermissionError("busy")): + assert await import_legacy(storage.database, source) + assert source.exists() + await storage.commit_trees((), clear_scope=original.scope) + assert await import_legacy(storage.database, source) is None + assert not source.exists() + assert (await storage.load_conversation_snapshot()).is_empty + + +async def test_invalid_graph_and_invalid_log_entry_do_not_discard_good_records( + storage, tmp_path +): + original, invalid = tree(), tree(root="bad") + invalid.nodes["bad"]["parent_id"] = "missing" + source = tmp_path / "sessions.json" + source.write_text( + json.dumps( + { + "trees": [original.to_json(), invalid.to_json()], + "managed_messages": { + "discord:other": [ + {"message_id": "good", "direction": "in", "kind": "prompt"}, + {"message_id": "bad"}, + ] + }, + } + ) + ) + assert await import_legacy(storage.database, source) + assert source.exists() + assert list((await storage.load_conversation_snapshot()).trees) == [ + original.identity + ] + assert await storage.get_tracked_message_ids_for_chat("discord", "other") == [ + "good" + ] + + +async def test_failed_import_transaction_keeps_source_and_allows_retry( + storage, tmp_path +): + from free_claude_code.runtime.messaging_sqlite import write_tree + + source = tmp_path / "sessions.json" + source.write_text( + json.dumps({"trees": [tree().to_json(), tree(root="second").to_json()]}) + ) + + def fail_second_tree(connection, snapshot): + write_tree(connection, snapshot) + if snapshot.root_id == "second": + raise sqlite3.OperationalError("disk failure") + + with ( + patch( + "free_claude_code.runtime.messaging_import.write_tree", + side_effect=fail_second_tree, + ), + pytest.raises(sqlite3.OperationalError), + ): + await import_legacy(storage.database, source) + assert source.exists() + assert (await storage.load_conversation_snapshot()).is_empty + assert ( + storage.database.execute( + lambda connection: connection.execute( + "SELECT count(*) FROM messaging_legacy_import" + ).fetchone()[0], + write=False, + ) + == 0 + ) + assert await import_legacy(storage.database, source) is None + assert not source.exists() + + +def incoming(node, *, text="prompt"): + return IncomingMessage( + platform="telegram", chat_id="chat", message_id=node, user_id="user", text=text + ) + + +def import_receipt(database): + return database.execute( + lambda connection: dict( + connection.execute("SELECT * FROM messaging_legacy_import").fetchone() + ), + write=False, + ) + + +@pytest.mark.parametrize( + "field,value", + [ + ("platform", "\ud800"), + ("chat_id", "\ud800"), + ("root_id", "\ud800"), + ("root_id", "\0root"), + ("node_id", "\ud800"), + ("node_id", "\0child"), + ("status_message_id", "\ud800"), + ("status_message_id", "\0status"), + ("session_id", "\ud800"), + ], +) +async def test_import_rejects_storage_invalid_tree_without_reserving_references( + storage, tmp_path, field, value +): + before, after = tree(root="before"), tree(root="after") + invalid = tree(root=value if field == "root_id" else "bad").to_json() + root = invalid["root_id"] + child = value if field == "node_id" else "child" + invalid["nodes"][child] = { + "node_id": child, + "parent_id": root, + "parent_reference_id": invalid["nodes"][root]["status_message_id"], + "status_message_id": "shared-reference", + "state": "completed", + "session_id": "native-child", + } + if field in {"platform", "chat_id"}: + invalid["scope"][field] = value + elif field not in {"root_id", "node_id"}: + invalid["nodes"][child][field] = value + # A rejected tree must not claim a reference needed by a later valid tree. + after.nodes["after"]["status_message_id"] = "shared-reference" + source = tmp_path / "sessions.json" + original = json.dumps({"trees": [before.to_json(), invalid, after.to_json()]}) + source.write_text(original) + + assert await import_legacy(storage.database, source) + assert (await storage.load_conversation_snapshot()).trees == { + before.identity: before, + after.identity: after, + } + assert source.read_text() == original + assert import_receipt(storage.database) == { + "source": "sessions.json", + "outcome": "partial", + "trees": 2, + "messages": 0, + "skipped": 1, + "cleanup_pending": 0, + } + with closing(sqlite3.connect(storage.database.path)) as connection: + assert connection.execute("PRAGMA foreign_key_check").fetchall() == [] + assert ( + connection.execute("SELECT count(*) FROM messaging_references").fetchone()[ + 0 + ] + == 4 + ) + assert ( + connection.execute("SELECT count(*) FROM messaging_trees").fetchone()[0] + == 2 + ) + + +@pytest.mark.parametrize("field", ["platform", "chat_id", "message_id", "ts", "kind"]) +@pytest.mark.parametrize("value", ["\ud800", "\0kind"]) +async def test_import_message_storage_validation_and_successful_duplicate_ownership( + storage, tmp_path, field, value +): + bad = {"message_id": "same", "direction": "in", "kind": "prompt", "ts": "old"} + key = "telegram:chat" + if field == "platform": + key = f"{value}:chat" + elif field == "chat_id": + key = f"telegram:{value}" + else: + bad[field] = value + good = {"message_id": "same", "direction": "out", "kind": "notice", "ts": "new"} + messages = {key: [bad]} + messages.setdefault("telegram:chat", []).extend([good, {**good, "ts": "duplicate"}]) + source = tmp_path / "sessions.json" + source.write_text(json.dumps({"managed_messages": messages})) + + rejected = value == "\ud800" or field == "kind" + assert bool(await import_legacy(storage.database, source)) is rejected + assert source.exists() is rejected + rows = storage.database.execute( + lambda connection: [ + tuple(row) + for row in connection.execute( + "SELECT platform,chat_id,message_id,ts,direction,kind FROM messaging_managed_messages ORDER BY sequence" + ) + ], + write=False, + ) + if rejected: + assert rows == [("telegram", "chat", "same", "new", "out", "notice")] + else: + # NUL is legal in fields without a length CHECK. Preserve accepted text. + assert rows[0] == ( + *key.split(":", 1), + bad["message_id"], + bad["ts"], + "in", + bad["kind"], + ) + assert len(rows) == (1 if field == "ts" else 2) + receipt = import_receipt(storage.database) + assert receipt["messages"] == len(rows) + assert receipt["skipped"] == int(rejected) + + +@pytest.mark.parametrize("last_invalid", [False, True]) +async def test_import_last_tree_identity_wins_even_when_unrecoverable( + storage, tmp_path, last_invalid +): + first, last = tree(), tree() + last.nodes["root"]["session_id"] = "\ud800" if last_invalid else "last-session" + source = tmp_path / "sessions.json" + source.write_text(json.dumps({"trees": [first.to_json(), last.to_json()]})) + assert await import_legacy(storage.database, source) + assert (await storage.load_conversation_snapshot()).trees == ( + {} if last_invalid else {last.identity: last} + ) + assert import_receipt(storage.database)["skipped"] == 1 + int(last_invalid) + await storage.commit_trees((tree(root="fresh"),)) + + +async def test_partial_import_survives_reopen_and_clear_without_replay( + storage, tmp_path +): + source = tmp_path / "sessions.json" + source.write_text( + json.dumps( + { + "trees": [tree().to_json()], + "managed_messages": { + "telegram:chat": [ + {"message_id": "bad", "direction": "in", "kind": "\0"} + ] + }, + } + ) + ) + assert await import_legacy(storage.database, source) + receipt = import_receipt(storage.database) + await storage.database.close() + for restart in range(2): + database = SQLiteDatabase(storage.database.path, tmp_path / "code.lock") + await database.start() + try: + reopened = SQLiteMessagingStore(database) + assert await import_legacy(database, source) is None + assert import_receipt(database) == receipt + assert (await reopened.load_conversation_snapshot()).is_empty is bool( + restart + ) + await reopened.commit_trees((), clear_scope=tree().scope) + finally: + await database.close() + assert source.exists() + + +@pytest.mark.parametrize("failure_stage", ["receipt", "commit"]) +async def test_import_outer_failure_rolls_back_released_units( + storage, tmp_path, failure_stage +): + from free_claude_code.runtime import sqlite_database + + source = tmp_path / "sessions.json" + source.write_text( + json.dumps( + { + "trees": [tree().to_json()], + "managed_messages": { + "telegram:chat": [ + {"message_id": "one", "direction": "in", "kind": "prompt"} + ] + }, + } + ) + ) + connect = sqlite_database._connect + inserted = [] + + def faulting_connect(path, *, existing=False): + connection = connect(path, existing=existing) + + def authorize(action, arg1, arg2, database_name, trigger): + if action == sqlite3.SQLITE_INSERT: + inserted.append(arg1) + if failure_stage == "receipt" and arg1 == "messaging_legacy_import": + return sqlite3.SQLITE_DENY + if ( + failure_stage == "commit" + and action == sqlite3.SQLITE_TRANSACTION + and arg1 == "COMMIT" + and "messaging_managed_messages" in inserted + ): + return sqlite3.SQLITE_DENY + return sqlite3.SQLITE_OK + + connection.set_authorizer(authorize) + return connection + + with ( + patch.object(sqlite_database, "_connect", faulting_connect), + pytest.raises(sqlite3.DatabaseError), + ): + await import_legacy(storage.database, source) + assert "messaging_nodes" in inserted and "messaging_managed_messages" in inserted + assert source.exists() + assert (await storage.load_conversation_snapshot()).is_empty + assert not await storage.get_tracked_message_ids_for_chat("telegram", "chat") + assert ( + storage.database.execute( + lambda connection: connection.execute( + "SELECT count(*) FROM messaging_legacy_import" + ).fetchone()[0], + write=False, + ) + == 0 + ) + assert await import_legacy(storage.database, source) is None + assert import_receipt(storage.database)["trees"] == 1 + assert import_receipt(storage.database)["messages"] == 1 + + +@pytest.mark.parametrize("platform", ["none", "telegram"]) +async def test_startup_partial_import_preserves_links_scopes_and_code( + storage, tmp_path, platform +): + from free_claude_code.application.code_sessions.models import CodeSession + from free_claude_code.config.settings import Settings + from free_claude_code.runtime.application import ApplicationRuntime + from free_claude_code.runtime.configuration import ConfigurationService + from free_claude_code.runtime.provider_manager import ProviderRuntimeManager + + code = SQLiteCodeStore(storage.database) + await code.start() + session = await code.create( + CodeSession(id="saved-code", cwd=str(tmp_path), model="provider/model") + ) + valid = tree(root="root-\U0001f600\0suffix") + parent = valid.root_id + for node_id in ("child", "grandchild"): + valid.nodes[node_id] = { + "node_id": node_id, + "status_message_id": f"status-{node_id}", + "parent_id": parent, + "parent_reference_id": valid.nodes[parent]["status_message_id"], + "state": "in_progress", + "session_id": "native-\U0001f600", + } + parent = node_id + inactive = tree("discord", root=valid.root_id) + conflict = tree(root="status-child") + source = tmp_path / "sessions.json" + source.write_text( + json.dumps( + { + "trees": [valid.to_json(), conflict.to_json(), inactive.to_json()], + "managed_messages": { + "telegram:chat": [ + { + "message_id": str(index), + "direction": "in", + "kind": "\0" if index == 0 else "prompt", + } + for index in range(4) + ] + }, + } + ) + ) + manager = ProviderRuntimeManager( + Settings().model_copy( + update={ + "messaging_platform": platform, + "max_message_log_entries_per_chat": 2, + } + ) + ) + runtime = ApplicationRuntime( + manager, + configuration=AsyncMock(spec=ConfigurationService), + transcriber=None, + database=storage.database, + ) + try: + with patch( + "free_claude_code.runtime.application.messaging_state_dir_path", + return_value=str(tmp_path), + ): + await runtime._initialize_messaging_storage() + assert runtime._messaging_storage_error is None + assert runtime._messaging_warning + assert (await storage.load_conversation_snapshot()).trees == { + valid.identity: valid, + inactive.identity: inactive, + } + assert await storage.get_tracked_message_ids_for_chat("telegram", "chat") == [ + "2", + "3", + ] + assert import_receipt(storage.database)["skipped"] == 2 + assert import_receipt(storage.database)["messages"] == 3 + assert await code.get_session(session.id) == session + assert ( + storage.database.execute( + lambda connection: connection.execute( + "PRAGMA foreign_key_check" + ).fetchall(), + write=False, + ) + == [] + ) + finally: + await manager.close() + await code.close() + + +async def test_failed_admission_never_launches_harness_or_publishes_memory(storage): + processor = AsyncMock() + manager = TreeQueueManager(processor, store=storage) + with ( + patch.object( + storage, "commit_trees", side_effect=sqlite3.OperationalError("disk full") + ), + pytest.raises(sqlite3.OperationalError), + ): + await manager.admit(incoming("new"), "status-new") + assert manager.get_tree_count() == 0 + assert manager.task_count() == 0 + processor.assert_not_awaited() + assert (await storage.load_conversation_snapshot()).is_empty + + +async def test_failed_clear_restores_runtime_prompts_queue_and_claim(storage): + started, release = asyncio.Event(), asyncio.Event() + prompts = [] + + async def process(claim): + prompts.append(claim.prompt) + started.set() + await release.wait() + await manager.complete_claim(claim, "resumable") + + manager = TreeQueueManager(process, store=storage) + root = await manager.admit(incoming("root", text="first"), "status-root") + assert root.claim is not None + await started.wait() + await manager.admit( + incoming("child", text="second"), + "status-child", + parent_reference_id="status-root", + ) + try: + with ( + patch.object( + storage, + "commit_trees", + side_effect=sqlite3.OperationalError("disk full"), + ), + pytest.raises(sqlite3.OperationalError), + ): + await manager.clear_scope(root.claim.identity.scope) + assert manager.task_count() == 1 + assert ( + len( + (await storage.load_conversation_snapshot()) + .trees[root.claim.identity] + .nodes + ) + == 2 + ) + finally: + release.set() + await asyncio.wait_for(manager.wait_idle(), 5) + assert prompts == ["first", "second"] + child = await manager.get_node(root.claim.identity.scope, "child") + assert child is not None and child.session_id == "resumable" + + +async def test_status_subtree_delete_preserves_prompt_sibling_branch(storage): + original = tree() + for node, reference in (("status-reply", "status-root"), ("prompt-reply", "root")): + original.nodes[node] = { + "node_id": node, + "status_message_id": f"status-{node}", + "state": "completed", + "parent_id": "root", + "parent_reference_id": reference, + "session_id": node, + } + await storage.commit_trees((original,)) + manager = TreeQueueManager.from_snapshot( + await storage.load_conversation_snapshot(), AsyncMock(), store=storage + ) + await manager.remove_message_subtree(original.scope, "status-root") + persisted = (await storage.load_conversation_snapshot()).trees[original.identity] + assert set(persisted.nodes) == {"root", "prompt-reply"} + assert persisted.nodes["root"]["status_message_id"] is None + assert persisted.nodes["root"]["session_id"] is None + assert persisted.nodes["prompt-reply"]["parent_reference_id"] == "root" + + +async def test_failure_after_scope_delete_rolls_back_both_trees_and_message_log( + storage, +): + original = tree() + await storage.commit_trees((original,)) + await storage.record_message_id("telegram", "chat", "tracked", "out", "notice") + invalid = tree(root="invalid") + invalid.nodes["invalid"]["state"] = "not-a-state" + with pytest.raises(ValueError): + await storage.commit_trees((invalid,), clear_scope=original.scope) + assert (await storage.load_conversation_snapshot()).trees == { + original.identity: original + } + assert await storage.get_tracked_message_ids_for_chat("telegram", "chat") == [ + "tracked" + ] + + +async def test_code_close_does_not_release_shared_database(storage, tmp_path): + code = SQLiteCodeStore(storage.database) + await code.start() + await code.close() + other = SQLiteDatabase(tmp_path / "fcc.db", tmp_path / "code.lock") + try: + with pytest.raises(sqlite3.OperationalError, match="another FCC"): + await other.start() + await storage.commit_trees((tree(),)) + assert not (await storage.load_conversation_snapshot()).is_empty + finally: + await other.close() + + +async def test_cancellation_during_commit_preserves_memory_and_database_agreement( + storage, +): + entered, release = threading.Event(), threading.Event() + execute = storage.database.execute + + def blocked(operation, *, write=True): + entered.set() + if not release.wait(5): + raise TimeoutError("test commit barrier") + return execute(operation, write=write) + + manager = TreeQueueManager(AsyncMock(), store=storage) + with patch.object(storage.database, "execute", blocked): + admission = asyncio.create_task(manager.admit(incoming("new"), "status-new")) + try: + assert await asyncio.to_thread(entered.wait, 3) + admission.cancel() + finally: + release.set() + decision = await admission + assert decision.accepted + await asyncio.wait_for(manager.wait_idle(), 5) + assert (await manager.snapshot()).trees == ( + await storage.load_conversation_snapshot() + ).trees + + +async def test_restore_does_not_consume_inactive_platform_status_repairs(storage): + telegram, discord = tree(), tree("discord") + telegram.nodes["root"]["state"] = "in_progress" + discord.nodes["root"]["state"] = "pending" + await storage.commit_trees((telegram, discord)) + workflow = MessagingWorkflow( + AsyncMock(), AsyncMock(), storage, platform_name="telegram" + ) + await workflow.restore() + await workflow.close() + loaded = await storage.load_conversation_snapshot() + assert loaded.trees[telegram.identity].nodes["root"]["state"] == "error" + assert loaded.trees[discord.identity].nodes["root"]["state"] == "pending" + next_workflow = MessagingWorkflow( + AsyncMock(), AsyncMock(), storage, platform_name="discord" + ) + await next_workflow.restore() + assert any( + target.scope.platform == "discord" + for target in next_workflow.tree_queue.restored_stale_targets + ) + + +@pytest.mark.parametrize("boundary", ["session", "terminal", "failure", "successor"]) +async def test_worker_write_failure_interrupts_queue_and_new_reply_recovers( + storage, boundary +): + started, release = asyncio.Event(), asyncio.Event() + processed, interrupted = [], [] + + async def process(claim): + processed.append(claim.node.node_id) + if claim.node.node_id == "root": + started.set() + await release.wait() + if boundary == "session": + await manager.record_session(claim, "native") + elif boundary == "failure": + await manager.fail_claim(claim) + await manager.complete_claim(claim, "native") + + manager = TreeQueueManager( + process, store=storage, unexpected_failure_callback=interrupted.append + ) + root = await manager.admit(incoming("root"), "status-root") + assert root.claim is not None + await started.wait() + await manager.admit(incoming("child"), "status-child", parent_reference_id="root") + condition = { + "session": "NEW.node_id='root' AND NEW.session_id='native'", + "terminal": "NEW.node_id='root' AND NEW.state='completed'", + "failure": "NEW.node_id='root' AND NEW.state='error'", + "successor": "NEW.node_id='child' AND NEW.state='in_progress'", + }[boundary] + await storage.database.run( + lambda connection: connection.execute( + f"CREATE TRIGGER reject_transition BEFORE UPDATE ON messaging_nodes WHEN {condition} " + "BEGIN SELECT RAISE(ABORT, 'test storage failure'); END" + ) + ) + release.set() + await asyncio.wait_for(manager.wait_idle(), 5) + assert processed == ["root"] + assert manager.task_count() == 0 + assert interrupted + assert "child" in { + node.node_id for result in interrupted for node in result.affected + } + saved = await storage.load_conversation_snapshot() + assert (await manager.snapshot()).trees == saved.trees + await storage.database.run( + lambda connection: connection.execute("DROP TRIGGER reject_transition") + ) + await manager.admit(incoming("later"), "status-later", parent_reference_id="root") + await asyncio.wait_for(manager.wait_idle(), 5) + assert processed == ["root", "later"] + nodes = ( + (await storage.load_conversation_snapshot()).trees[root.claim.identity].nodes + ) + assert nodes["root"]["state"] == ( + "completed" if boundary == "successor" else "error" + ) + assert nodes["child"]["state"] == "error" + assert nodes["later"]["state"] == "completed" + + +async def test_shutdown_drains_workers_even_when_interruption_cannot_be_saved( + storage, mock_platform +): + started, stopped = asyncio.Event(), asyncio.Event() + cli = AsyncMock() + workflow = MessagingWorkflow(mock_platform, cli, storage, platform_name="telegram") + + async def process(claim): + started.set() + try: + await asyncio.Event().wait() + finally: + await manager.fail_claim(claim) + stopped.set() + + manager = TreeQueueManager(process, store=storage) + workflow._tree_queue = manager + await manager.admit(incoming("root"), "status-root") + await started.wait() + await manager.admit(incoming("child"), "status-child", parent_reference_id="root") + execute = storage.database.execute + writes = 0 + + def broken(operation, *, write=True): + nonlocal writes + if write: + writes += 1 + raise sqlite3.OperationalError("disk full") + return execute(operation, write=write) + + try: + with patch.object(storage.database, "execute", broken): + await asyncio.wait_for(workflow.close(), 5) + assert writes == 1 + assert stopped.is_set() + assert manager.task_count() == 0 + cli.stop_all.assert_awaited() + assert any( + call.args[1] == "status-child" + for call in mock_platform.queue_edit_message.call_args_list + ) + finally: + await manager.shutdown() + await manager.wait_idle() + + +async def test_persistent_failure_has_no_retries_and_failed_recovery_keeps_history( + storage, +): + started, release = asyncio.Event(), asyncio.Event() + processed = [] + # An outbound failure must not prevent task retirement. + notify = Mock(side_effect=RuntimeError("platform unavailable")) + + async def process(claim): + processed.append(claim.node.node_id) + started.set() + await release.wait() + await manager.complete_claim(claim, "native") + + manager = TreeQueueManager( + process, store=storage, unexpected_failure_callback=notify + ) + root = await manager.admit(incoming("root"), "status-root") + await started.wait() + await manager.admit(incoming("child"), "status-child", parent_reference_id="root") + original = await storage.load_conversation_snapshot() + writes = 0 + execute = storage.database.execute + + def broken(operation, *, write=True): + nonlocal writes + if write: + writes += 1 + raise sqlite3.OperationalError("disk full") + return execute(operation, write=False) + + with patch.object(storage.database, "execute", broken): + release.set() + await asyncio.wait_for(manager.wait_idle(), 5) + assert writes == 1 + assert processed == ["root"] + assert manager.task_count() == 0 + with pytest.raises(MessagingStorageError): + await manager.admit( + incoming("new"), "status-new", parent_reference_id="root" + ) + assert writes == 2 + assert (await manager.snapshot()) == original + assert (await storage.load_conversation_snapshot()) == original + notify.assert_called_once() + await manager.admit(incoming("new"), "status-new", parent_reference_id="root") + await asyncio.wait_for(manager.wait_idle(), 5) + assert processed == ["root", "new"] + nodes = ( + (await storage.load_conversation_snapshot()).trees[root.claim.identity].nodes + ) + assert nodes["root"]["state"] == nodes["child"]["state"] == "error" + assert nodes["new"]["state"] == "completed" + + +async def test_gated_tree_can_be_cleared_and_replaced_before_old_cleanup_finishes( + storage, monkeypatch +): + monkeypatch.setattr( + "free_claude_code.messaging.trees.manager.CANCEL_TASK_DRAIN_TIMEOUT_S", 0.01 + ) + faulted, cleanup_release = asyncio.Event(), asyncio.Event() + new_started, new_release = asyncio.Event(), asyncio.Event() + claims = [] + notifications = [] + + async def process(claim): + claims.append(claim) + if len(claims) == 1: + try: + with patch.object( + storage.database, + "execute", + side_effect=sqlite3.OperationalError("disk full"), + ): + await manager.record_session(claim, "old-session") + except MessagingStorageError: + faulted.set() + while not cleanup_release.is_set(): + with suppress(asyncio.CancelledError): + await cleanup_release.wait() + # This late update must not affect the same-ID replacement. + await manager.complete_claim(claim, "old-session") + raise + new_started.set() + await new_release.wait() + await manager.complete_claim(claim, "new-session") + + manager = TreeQueueManager( + process, store=storage, unexpected_failure_callback=notifications.append + ) + old = await manager.admit(incoming("root"), "status-root") + assert old.claim is not None + await faulted.wait() + try: + with pytest.raises(MessagingStorageError, match="cleaning up"): + await manager.admit( + incoming("child"), "status-child", parent_reference_id="root" + ) + await manager.clear_scope(old.claim.identity.scope) + replacement = await manager.admit(incoming("root"), "status-root") + assert replacement.claim is not None + await new_started.wait() + assert replacement.claim.claim_id != old.claim.claim_id + assert manager.task_count() == 2 + finally: + cleanup_release.set() + new_release.set() + await asyncio.wait_for(manager.wait_idle(), 5) + nodes = (await storage.load_conversation_snapshot()).trees[old.claim.identity].nodes + assert list(nodes) == ["root"] + assert nodes["root"]["session_id"] == "new-session" + assert not notifications diff --git a/tests/runtime/test_sqlite_database.py b/tests/runtime/test_sqlite_database.py index 6f10791fa3..10a992597d 100644 --- a/tests/runtime/test_sqlite_database.py +++ b/tests/runtime/test_sqlite_database.py @@ -116,7 +116,7 @@ def test_historical_database_moves_and_upgrades_without_losing_records( assert len(after["code_items"]) == 2 assert after["code_items"][-1]["kind"] == "prompt" with closing(sqlite3.connect(target)) as connection: - assert connection.execute("PRAGMA user_version").fetchone() == (3,) + assert connection.execute("PRAGMA user_version").fetchone() == (4,) assert connection.execute("PRAGMA journal_mode").fetchone() == ("wal",) assert connection.execute("PRAGMA foreign_key_check").fetchall() == [] @@ -153,7 +153,7 @@ def test_orphan_journals_are_not_treated_as_a_fresh_database( assert not (tmp_path / "fcc.db").exists() -@pytest.mark.parametrize("version", [-1, 4]) +@pytest.mark.parametrize("version", [-1, 5]) def test_unsupported_version_is_not_moved_or_mutated(tmp_path, version): old, target = tmp_path / "code/code.db", tmp_path / "fcc.db" historical_database(old, 3) @@ -289,6 +289,7 @@ def fail_upgrade(connection): def test_current_schema_does_not_execute_historical_migrations(tmp_path, monkeypatch): target = tmp_path / "fcc.db" historical_database(target, 3) + initialize_database(target) def already_applied(connection): pytest.fail("An applied migration was executed again") @@ -296,11 +297,37 @@ def already_applied(connection): monkeypatch.setattr( sqlite_database, "MIGRATIONS", - tuple((version, already_applied) for version in (1, 2, 3)), + tuple((version, already_applied) for version in (1, 2, 3, 4)), ) initialize_database(target) +def test_messaging_schema_failure_rolls_back_without_changing_code_history( + tmp_path, monkeypatch +): + target = tmp_path / "fcc.db" + historical_database(target, 3) + before = snapshot(target) + migrations = sqlite_database.MIGRATIONS + + def fail(connection): + migrations[-1][1](connection) + raise sqlite3.DatabaseError("interrupted messaging schema") + + with monkeypatch.context() as patch: + patch.setattr(sqlite_database, "MIGRATIONS", (*migrations[:-1], (4, fail))) + with pytest.raises(sqlite3.DatabaseError, match="interrupted"): + initialize_database(target) + assert snapshot(target) == before + with closing(sqlite3.connect(target)) as connection: + assert connection.execute("PRAGMA user_version").fetchone() == (3,) + assert not connection.execute( + "SELECT name FROM sqlite_schema WHERE name LIKE 'messaging_%'" + ).fetchall() + initialize_database(target) + assert snapshot(target) == before + + def test_busy_wal_is_not_renamed_and_retries_after_reader_closes(tmp_path, monkeypatch): old, target = tmp_path / "code/code.db", tmp_path / "fcc.db" historical_database(old, 3) @@ -349,7 +376,9 @@ def test_redirected_legacy_storage_is_not_moved(tmp_path, redirect): @pytest.mark.asyncio -async def test_restart_recovery_runs_after_current_schema_relocation(tmp_path): +async def test_restart_recovery_runs_after_current_schema_relocation( + database_factory, tmp_path +): old, target, lock = ( tmp_path / "code/code.db", tmp_path / "fcc.db", @@ -363,7 +392,7 @@ async def test_restart_recovery_runs_after_current_schema_relocation(tmp_path): "UPDATE code_runs SET status='running', finished_at=NULL" ) connection.execute("UPDATE code_prompts SET status='answering'") - store = SQLiteCodeStore(target, lock, legacy_database_path=old) + store = SQLiteCodeStore(database_factory(target, lock, legacy_path=old)) try: await store.start() run = await store.get_run("s", "r") @@ -372,18 +401,21 @@ async def test_restart_recovery_runs_after_current_schema_relocation(tmp_path): assert (await store.prompts("s"))[0].status == "expired" finally: await store.close() + await store.database.close() @pytest.mark.asyncio -async def test_old_owner_blocks_move_and_new_owner_uses_same_lock(tmp_path): +async def test_old_owner_blocks_move_and_new_owner_uses_same_lock( + database_factory, tmp_path +): old, target, lock = ( tmp_path / "code/code.db", tmp_path / "fcc.db", tmp_path / "code/code.lock", ) historical_database(old, 3) - first = SQLiteCodeStore(old, lock) - second = SQLiteCodeStore(target, lock, legacy_database_path=old) + first = SQLiteCodeStore(database_factory(old, lock)) + second = SQLiteCodeStore(database_factory(target, lock, legacy_path=old)) await first.start() try: with pytest.raises(CodeUnavailableError, match="another FCC"): @@ -391,12 +423,15 @@ async def test_old_owner_blocks_move_and_new_owner_uses_same_lock(tmp_path): assert old.exists() and not target.exists() finally: await first.close() + await first.database.close() + second = SQLiteCodeStore(database_factory(target, lock, legacy_path=old)) await second.start() try: assert (await second.get_session("s")).title == "Saved" assert lock.exists() and not old.exists() finally: await second.close() + await second.database.close() @pytest.mark.skipif(os.name == "nt", reason="POSIX permission contract") @@ -416,7 +451,7 @@ async def test_production_startup_keeps_http_available_and_chat_cleanup_independ from free_claude_code.config import paths from free_claude_code.config.settings import Settings from free_claude_code.core.interprocess_lock import InterprocessFileLock - from free_claude_code.runtime import bootstrap, code_sessions_sqlite + from free_claude_code.runtime import bootstrap from free_claude_code.runtime.provider_manager import ProviderRuntimeManager from free_claude_code.runtime.retired_chat import remove_retired_chat_history @@ -436,7 +471,7 @@ def gated_migration(*args): raise sqlite3.DatabaseError("injected migration failure") initialize_database(*args) - monkeypatch.setattr(code_sessions_sqlite, "initialize_database", gated_migration) + monkeypatch.setattr(sqlite_database, "initialize_database", gated_migration) monkeypatch.setattr(bootstrap, "configure_logging", lambda *args, **kwargs: None) monkeypatch.setattr( ProviderRuntimeManager, "start_model_list_refresh", lambda _: None