179 lines
8.8 KiB
Python
179 lines
8.8 KiB
Python
"""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)
|