"""SQL-backed retirement preserves shared/private reads and terminal reconnects.""" import asyncio from types import SimpleNamespace from unittest.mock import AsyncMock import pytest from langgraph.checkpoint.base import empty_checkpoint from langgraph.checkpoint.memory import InMemorySaver from test_thread_lazy_checkpoint import client_for from app.gateway.routers import artifacts, thread_runs from deerflow.persistence.run.model import RunRow from deerflow.persistence.run.sql import RunRepository from deerflow.runtime.events.store.spool import SpoolingRunEventStore from deerflow.runtime.runs.manager import RunManager from deerflow.runtime.runs.schemas import RunStatus from deerflow.runtime.runs.store.memory import MemoryRunStore from deerflow.runtime.stream_bridge.memory import MemoryStreamBridge from deerflow.runtime.user_context import reset_current_user, set_current_user @pytest.mark.asyncio @pytest.mark.parametrize("shared", [False, True]) async def test_retired_sql_run_history_and_reconnect_keep_authorization(tmp_path, shared): saver = InMemorySaver() async with client_for(saver) as (client, app): sf = app.state.thread_store._sf async with sf.kw["bind"].begin() as connection: await connection.run_sync(RunRow.__table__.create) manager = RunManager(RunRepository(sf)) app.state.run_manager = manager app.state.stream_bridge = MemoryStreamBridge() app.state.run_event_store = SpoolingRunEventStore(tmp_path / "events.db") app.include_router(thread_runs.router) metadata = {"taskId": "task", "agent_id": "agent", "actionId": "action"} if shared else {} assert (await client.post("/api/threads", json={"thread_id": "t", "metadata": metadata})).status_code == 200 token = set_current_user(SimpleNamespace(id="owner")) try: record = await manager.create("t", metadata={"user_id": "owner"}, kwargs={"input": {"messages": ["original user question"]}}) record.metadata["late_artifact"] = "report.md" await manager.set_status(record.run_id, RunStatus.success) await app.state.run_event_store.put(thread_id="t", run_id=record.run_id, event_type="answer", category="message", content="answer must survive retirement") checkpoint = empty_checkpoint() checkpoint["channel_values"] = {"messages": [{"type": "ai", "content": "full checkpoint answer"}]} checkpoint["channel_versions"] = {"messages": 1} await saver.aput({"configurable": {"thread_id": "t", "checkpoint_ns": ""}}, checkpoint, {}, {"messages": 1}) await manager.reap_completed(retention_seconds=0) assert manager.get(record.run_id) is None persisted = await manager.get_persisted(record.run_id) assert persisted.kwargs == record.kwargs assert persisted.metadata["late_artifact"] == "report.md" finally: reset_current_user(token) for user in ["owner", "other", "anonymous"]: headers = {"X-User": user} expected = 401 if user == "anonymous" else 200 if shared or user == "owner" else 404 for suffix in [f"runs/{record.run_id}", "runs", f"runs/{record.run_id}/stream", f"runs/{record.run_id}/messages"]: response = await asyncio.wait_for(client.get(f"/api/threads/t/{suffix}", headers=headers), 2) assert response.status_code == expected, response.text if response.status_code == 200 and suffix.endswith("/stream"): assert "event: end" in response.text if response.status_code == 200 and suffix == "runs": assert record.run_id in [row["run_id"] for row in response.json()] state = await client.get("/api/threads/t/state", headers=headers) assert state.status_code == expected if expected == 200: assert state.json()["values"]["messages"][0]["content"] == "full checkpoint answer" assert await manager.get_persisted(record.run_id, authorized_thread_id="wrong-thread") is None @pytest.mark.asyncio async def test_retirement_keeps_live_task_and_unpersisted_or_failed_snapshot(): store = MemoryRunStore() manager = RunManager(store) record = await manager.create("t", kwargs={"input": "keep"}) await manager.set_status(record.run_id, RunStatus.interrupted) release = asyncio.Event() record.task = asyncio.create_task(release.wait()) manager.watch_completion(record) await manager.cleanup(record.run_id, delay=0) assert manager.get(record.run_id) is record release.set() await record.task original_get = store.get store.get = AsyncMock(side_effect=OSError("database unavailable")) await manager.cleanup(record.run_id, delay=0) assert manager.get(record.run_id) is record store.get = AsyncMock(return_value=None) await manager.cleanup(record.run_id, delay=0) assert manager.get(record.run_id) is record store.get = original_get await manager.cleanup(record.run_id, delay=0) assert manager.get(record.run_id) is None @pytest.mark.asyncio async def test_retention_preserves_five_minute_grace_and_close_stops_sweeper(): manager = RunManager(MemoryRunStore()) record = await manager.create("t") await manager.set_status(record.run_id, RunStatus.success) await manager.reap_completed() assert manager.get(record.run_id) is record manager.start_retention(retention_seconds=0, interval_seconds=0.01) async with asyncio.timeout(2): while manager.get(record.run_id): await asyncio.sleep(0.005) await manager.aclose() assert manager._retention_task.done() @pytest.mark.asyncio async def test_artifact_directory_scan_does_not_block_other_user_metadata(monkeypatch): import threading entered, release = threading.Event(), threading.Event() def slow_scan(_): entered.set() assert release.wait(3) return [] monkeypatch.setattr(artifacts, "collect_outputs_listing", slow_scan) monkeypatch.setattr(artifacts, "aresolve_path_user_id", AsyncMock(return_value="owner")) async with client_for(InMemorySaver()) as (client, app): app.include_router(artifacts.router) assert (await client.post("/api/threads", json={"thread_id": "scan"})).status_code == 200 scanning = asyncio.create_task(client.get("/api/threads/scan/artifacts")) try: async with asyncio.timeout(2): while not entered.is_set(): await asyncio.sleep(0.005) assert not scanning.done() response = await asyncio.wait_for(client.post("/api/threads", json={"thread_id": "unrelated"}, headers={"X-User": "other"}), 1) assert response.status_code == 200 denied = await client.get("/api/threads/scan/artifacts", headers={"X-User": "other"}) assert denied.status_code == 404 finally: release.set() assert (await scanning).status_code == 200 @pytest.mark.asyncio async def test_finished_run_rejoin_after_buffer_expiry_ends_instead_of_hanging(): from app.gateway.services import sse_consumer manager = RunManager(MemoryRunStore()) record = await manager.create("t") await manager.set_status(record.run_id, RunStatus.success) bridge = MemoryStreamBridge() await bridge.publish(record.run_id, "values", {"messages": []}) await bridge.publish_end(record.run_id) await bridge.cleanup(record.run_id) frames = [frame async for frame in sse_consumer(bridge, record, SimpleNamespace(headers={}), manager)] assert len(frames) == 1 and "event: end" in frames[0] @pytest.mark.asyncio async def test_retirement_preserves_last_event_id_replay_and_fences_old_completion(): from app.gateway.services import sse_consumer from deerflow.runtime.runs.schemas import DisconnectMode manager = RunManager(MemoryRunStore()) older = await manager.create("t", on_disconnect=DisconnectMode.continue_) latest = await manager.create("t", on_disconnect=DisconnectMode.continue_) await manager.set_status(latest.run_id, RunStatus.success) bridge = MemoryStreamBridge() await bridge.publish(latest.run_id, "messages", {"content": "first"}) first_id = bridge._streams[latest.run_id].events[-1].id await bridge.publish(latest.run_id, "messages", {"content": "second"}) await bridge.publish_end(latest.run_id) await manager.cleanup(latest.run_id, delay=0) assert not manager.is_latest(older), "Retiring a newer run must not un-fence an old finalizer" frames = [frame async for frame in sse_consumer(bridge, latest, SimpleNamespace(headers={"Last-Event-ID": first_id}), manager)] assert len(frames) == 2 assert '"second"' in frames[0] and '"first"' not in frames[0] assert "event: end" in frames[1] assert bridge.has_stream(latest.run_id)