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