deer-flow/backend/tests/test_history_batch_queries.py
Hyeonsang Cho 6f81daefff
fix(runtime): resolve omitted run owners before idempotent reuse (#5401)
* fix(runtime): resolve omitted run owners before idempotent reuse

HTTP run admissions do not pass a user_id. The SQL run store stamps the
request user from the contextvar onto the row, but RunManager kept None
on its process-local RunRecord, so the two disagreed about who owns the
run. A keyed retry that reached a peer worker, or the owning worker after
cleanup() released its local record, hydrated the stamped row, compared
its owner with None, and raised "Run idempotency key resolved to a
different thread or user", which start_run surfaced as a 500.
MemoryRunStore stored None on both sides and never hit the check.

The same mismatch hid HTTP runs from owner-scoped history reads, which
filter local records by the current user, and skipped the worker's MCP
background_tasks projection, which only runs for records with an owner.

create() and _admit_thread_operation() now resolve an omitted owner from
the current user before building the record, and keep None when no user
is in context instead of falling back to the default bucket, so the
local record and every store agree on the owner. HTTP runs now receive
the background_tasks projection, so the replay golden's values frames
gain that key.

* docs(changelog): reference #5401 in the keyed retry owner fix entry

* test(runtime): close the SQL engine when peer-reuse test setup fails

Review follow-up on #5401: the sql case of the two-worker start_run test
initialized the engine above the try whose finally calls close_engine().
init_engine() assigns the module-global engine and session factory before
bootstrapping the schema, so a failure there skipped the teardown and left
a stale engine for later tests in the same process.

Store setup and the RunManager workers now live inside the try, so the
teardown runs whether setup or the test body fails.
2026-09-13 20:28:32 +08:00

442 lines
15 KiB
Python

"""Cross-store contracts used by thread-global history pagination."""
from __future__ import annotations
import pytest
from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.manager import EditReplayVisibility
from deerflow.runtime.runs.store.memory import MemoryRunStore
async def _seed_ai_messages(store):
await store.put(
thread_id="t1",
run_id="r1",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "first"},
metadata={"caller": "lead_agent"},
)
await store.put(
thread_id="t1",
run_id="r1",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "middleware"},
metadata={"caller": "middleware:title"},
)
last = await store.put(
thread_id="t1",
run_id="r1",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "last"},
metadata={"caller": "lead_agent"},
)
other = await store.put(
thread_id="t1",
run_id="r2",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "other"},
metadata={"caller": "lead_agent"},
)
await store.put(
thread_id="t1",
run_id="r_mw",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "middleware only"},
metadata={"caller": "middleware:title"},
)
return {"r1": last["seq"], "r2": other["seq"]}
@pytest.mark.anyio
async def test_memory_event_store_returns_global_last_non_middleware_ai_seq():
store = MemoryRunEventStore()
expected = await _seed_ai_messages(store)
result = await store.get_last_visible_ai_seq_by_run("t1", {"r1", "r2", "r_mw", "missing"})
assert result == expected
assert "r_mw" not in result
@pytest.mark.anyio
async def test_memory_event_store_defensively_rechecks_message_category():
store = MemoryRunEventStore()
expected = await store.put(
thread_id="t1",
run_id="r1",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "visible"},
metadata={"caller": "lead_agent"},
)
mutated = await store.put(
thread_id="t1",
run_id="r1",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "content": "no longer a message"},
metadata={"caller": "lead_agent"},
)
# Memory projections intentionally share their row dictionaries. Recheck
# category at read time so an accidental mutation cannot violate the same
# contract that the DB and JSONL stores enforce explicitly.
mutated["category"] = "trace"
assert await store.get_last_visible_ai_seq_by_run("t1", {"r1"}) == {"r1": expected["seq"]}
@pytest.mark.anyio
async def test_jsonl_event_store_returns_global_last_non_middleware_ai_seq(tmp_path):
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
store = JsonlRunEventStore(base_dir=tmp_path)
expected = await _seed_ai_messages(store)
result = await store.get_last_visible_ai_seq_by_run("t1", {"r1", "r2", "r_mw", "missing"})
assert result == expected
assert "r_mw" not in result
@pytest.mark.anyio
async def test_db_event_store_returns_global_last_non_middleware_ai_seq(tmp_path):
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.runtime.events.store.db import DbRunEventStore
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'events.db'}", sqlite_dir=str(tmp_path))
try:
store = DbRunEventStore(get_session_factory())
expected = await _seed_ai_messages(store)
result = await store.get_last_visible_ai_seq_by_run("t1", {"r1", "r2", "r_mw", "missing"})
assert result == expected
assert "r_mw" not in result
finally:
await close_engine()
@pytest.mark.anyio
async def test_memory_run_store_supersession_is_unbounded_and_owner_scoped():
store = MemoryRunStore()
for index in range(105):
await store.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
await store.put(
"regen-success",
thread_id="t1",
user_id="alice",
status="success",
metadata={"regenerate_from_run_id": "source-success"},
)
await store.put(
"regen-failed",
thread_id="t1",
user_id="alice",
status="error",
metadata={"regenerate_from_run_id": "source-failed"},
)
await store.put(
"regen-bob",
thread_id="t1",
user_id="bob",
status="success",
metadata={"regenerate_from_run_id": "source-bob"},
)
assert await store.list_successful_regenerate_sources("t1", user_id="alice") == {"source-success"}
@pytest.mark.anyio
async def test_run_repository_batch_queries_are_unbounded_and_owner_scoped(tmp_path):
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.persistence.run import RunRepository
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}", sqlite_dir=str(tmp_path))
try:
repo = RunRepository(get_session_factory())
for index in range(105):
await repo.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
await repo.put(
"regen-a",
thread_id="t1",
user_id="alice",
status="success",
metadata={"regenerate_from_run_id": "source-a"},
)
await repo.put(
"regen-b",
thread_id="t1",
user_id="bob",
status="success",
metadata={"regenerate_from_run_id": "source-b"},
)
assert await repo.list_successful_regenerate_sources("t1", user_id="alice") == {"source-a"}
rows = await repo.get_many_by_thread("t1", {"normal-0", "regen-a", "regen-b"}, user_id="alice")
assert set(rows) == {"normal-0", "regen-a"}
finally:
await close_engine()
@pytest.mark.anyio
async def test_run_manager_prefers_latest_in_memory_regenerate_status():
store = MemoryRunStore()
await store.put(
"regen",
thread_id="t1",
status="success",
metadata={"regenerate_from_run_id": "source"},
)
manager = RunManager(store=store)
# Simulate the same logical run being newer in memory than its persisted
# successful snapshot.
persisted = await manager.get("regen")
assert persisted is not None
manager._runs["regen"] = persisted
manager._index_run_locked(persisted)
persisted.status = RunStatus.error
assert await manager.list_successful_regenerate_sources("t1", user_id=None) == set()
@pytest.mark.anyio
async def test_run_manager_uses_latest_attempt_for_shared_regenerate_source():
manager = RunManager()
older = await manager.create(
"t1",
metadata={"regenerate_from_run_id": "source"},
)
older.status = RunStatus.success
newer = await manager.create(
"t1",
metadata={"regenerate_from_run_id": "source"},
)
newer.status = RunStatus.error
assert await manager.list_successful_regenerate_sources("t1", user_id=None) == set()
@pytest.mark.anyio
async def test_run_manager_batch_history_methods_default_to_current_user():
from types import SimpleNamespace
from deerflow.runtime.user_context import reset_current_user, set_current_user
store = MemoryRunStore()
await store.put(
"regen-alice",
thread_id="shared-thread",
user_id="alice",
status="success",
metadata={"regenerate_from_run_id": "source-alice"},
)
await store.put(
"regen-bob",
thread_id="shared-thread",
user_id="bob",
status="success",
metadata={"regenerate_from_run_id": "source-bob"},
)
manager = RunManager(store=store)
token = set_current_user(SimpleNamespace(id="alice"))
try:
sources = await manager.list_successful_regenerate_sources("shared-thread")
records = await manager.get_many_by_thread("shared-thread", {"regen-alice", "regen-bob"})
finally:
reset_current_user(token)
assert sources == {"source-alice"}
assert set(records) == {"regen-alice"}
@pytest.mark.anyio
async def test_run_manager_owner_scoped_history_includes_runs_admitted_without_explicit_user():
"""A run admitted with the ambient user must count as that user's run."""
from types import SimpleNamespace
from deerflow.runtime.user_context import reset_current_user, set_current_user
manager = RunManager(store=MemoryRunStore())
token = set_current_user(SimpleNamespace(id="alice"))
try:
record = await manager.create_or_reject("t1", metadata={"regenerate_from_run_id": "source"})
await manager.set_status(record.run_id, RunStatus.success)
sources = await manager.list_successful_regenerate_sources("t1")
records = await manager.get_many_by_thread("t1", {record.run_id})
finally:
reset_current_user(token)
assert sources == {"source"}
assert set(records) == {record.run_id}
@pytest.mark.anyio
@pytest.mark.no_auto_user
async def test_run_manager_keeps_omitted_owner_unset_without_user_context():
"""No ambient user means no owner, never a fallback bucket such as ``default``."""
manager = RunManager(store=MemoryRunStore())
record = await manager.create_or_reject("t1")
assert record.user_id is None
@pytest.mark.anyio
async def test_run_manager_batch_history_methods_fail_closed_without_user_context():
from deerflow.runtime import user_context
manager = RunManager(store=MemoryRunStore())
token = user_context._current_user.set(None)
try:
with pytest.raises(RuntimeError, match="user_id=AUTO"):
await manager.list_successful_regenerate_sources("t1")
with pytest.raises(RuntimeError, match="user_id=AUTO"):
await manager.get_many_by_thread("t1", {"run-1"})
finally:
user_context._current_user.reset(token)
@pytest.mark.anyio
async def test_run_manager_batch_history_methods_allow_explicit_unscoped_access():
store = MemoryRunStore()
await store.put(
"regen-alice",
thread_id="shared-thread",
user_id="alice",
status="success",
metadata={"regenerate_from_run_id": "source-alice"},
)
await store.put(
"regen-bob",
thread_id="shared-thread",
user_id="bob",
status="success",
metadata={"regenerate_from_run_id": "source-bob"},
)
manager = RunManager(store=store)
sources = await manager.list_successful_regenerate_sources("shared-thread", user_id=None)
records = await manager.get_many_by_thread("shared-thread", {"regen-alice", "regen-bob"}, user_id=None)
assert sources == {"source-alice", "source-bob"}
assert set(records) == {"regen-alice", "regen-bob"}
@pytest.mark.anyio
async def test_memory_run_store_lists_edit_replay_runs_unbounded_and_owner_scoped():
store = MemoryRunStore()
for index in range(105):
await store.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
await store.put(
"edit-success",
thread_id="t1",
user_id="alice",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
)
await store.put(
"edit-error",
thread_id="t1",
user_id="alice",
status="error",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"},
)
await store.put(
"edit-bob",
thread_id="t1",
user_id="bob",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"},
)
rows = await store.list_edit_regenerate_runs("t1", user_id="alice")
assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"]
@pytest.mark.anyio
async def test_run_repository_lists_edit_replay_runs_unbounded_and_owner_scoped(tmp_path):
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.persistence.run import RunRepository
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}", sqlite_dir=str(tmp_path))
try:
repo = RunRepository(get_session_factory())
for index in range(105):
await repo.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
await repo.put(
"edit-success",
thread_id="t1",
user_id="alice",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
)
await repo.put(
"edit-error",
thread_id="t1",
user_id="alice",
status="error",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"},
)
await repo.put(
"edit-bob",
thread_id="t1",
user_id="bob",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"},
)
rows = await repo.list_edit_regenerate_runs("t1", user_id="alice")
assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"]
finally:
await close_engine()
@pytest.mark.anyio
async def test_run_manager_computes_edit_replay_visibility_by_latest_attempt():
manager = RunManager()
running = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-running"},
)
running.status = RunStatus.running
success = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
)
success.status = RunStatus.success
failed = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-failed"},
)
failed.status = RunStatus.error
older_success = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"},
)
older_success.status = RunStatus.success
newer_failed = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"},
)
newer_failed.status = RunStatus.interrupted
older_failed_then_success = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"},
)
older_failed_then_success.status = RunStatus.error
newer_success_after_failure = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"},
)
newer_success_after_failure.status = RunStatus.success
visibility = await manager.list_edit_replay_visibility("t1", user_id=None)
assert visibility == EditReplayVisibility(
hidden_source_run_ids={"source-running", "source-success", "source-retry-success"},
hidden_attempt_run_ids={failed.run_id, newer_failed.run_id, older_failed_then_success.run_id},
)