mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-19 19:16:17 +00:00
f0cb67b223
* feat(extensions): expose incremental run evidence reader * fix(extensions): address run evidence review feedback * docs(extensions): clarify run deletion reconciliation * docs(migrations): align current head documentation * fix(extensions): isolate run evidence event reads * test: avoid pinning run change migration to latest head --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
344 lines
14 KiB
Python
344 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from deerflow_extension_api import InvalidRunEvidenceCursor
|
|
|
|
from deerflow.extensions.run_evidence import StoreRunEvidenceReader
|
|
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
|
from deerflow.runtime.runs.store.memory import MemoryRunStore
|
|
|
|
|
|
async def _put_run(store: MemoryRunStore, run_id: str, thread_id: str, *, user_id: str = "user-1") -> None:
|
|
await store.put(run_id, thread_id=thread_id, user_id=user_id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_changed_run_cursor_pages_without_loss_and_replays_updates():
|
|
runs = MemoryRunStore()
|
|
events = MemoryRunEventStore()
|
|
reader = StoreRunEvidenceReader(runs, events, user_id="user-1")
|
|
await _put_run(runs, "run-a", "thread-a")
|
|
await _put_run(runs, "run-b", "thread-b")
|
|
|
|
first = await reader.list_changed_runs(cursor=None, limit=1)
|
|
assert [item.run_id for item in first.items] == ["run-a"]
|
|
assert first.has_more is True
|
|
assert first.next_cursor
|
|
|
|
await runs.update_status("run-a", "running")
|
|
second = await reader.list_changed_runs(cursor=first.next_cursor, limit=10)
|
|
assert [item.run_id for item in second.items] == ["run-b", "run-a"]
|
|
assert second.has_more is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_changed_run_reader_is_bound_to_one_user():
|
|
runs = MemoryRunStore()
|
|
events = MemoryRunEventStore()
|
|
await _put_run(runs, "mine", "thread-1", user_id="user-1")
|
|
await _put_run(runs, "theirs", "thread-2", user_id="user-2")
|
|
|
|
page = await StoreRunEvidenceReader(runs, events, user_id="user-1").list_changed_runs(cursor=None, limit=10)
|
|
assert [item.run_id for item in page.items] == ["mine"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_event_pages_resume_and_status_is_authoritative():
|
|
runs = MemoryRunStore()
|
|
events = MemoryRunEventStore()
|
|
await _put_run(runs, "run-a", "thread-a")
|
|
await runs.update_status("run-a", "error", error="failed")
|
|
for index in range(3):
|
|
await events.put(
|
|
thread_id="thread-a",
|
|
run_id="run-a",
|
|
event_type="test.event",
|
|
category="trace",
|
|
content={"index": index},
|
|
)
|
|
|
|
reader = StoreRunEvidenceReader(runs, events, user_id="user-1")
|
|
first = await reader.list_run_events(thread_id="thread-a", run_id="run-a", after_seq=None, limit=2)
|
|
assert [item.content for item in first.items] == [{"index": 0}, {"index": 1}]
|
|
assert first.has_more is True
|
|
second = await reader.list_run_events(
|
|
thread_id="thread-a",
|
|
run_id="run-a",
|
|
after_seq=first.next_after_seq,
|
|
limit=2,
|
|
)
|
|
assert [item.content for item in second.items] == [{"index": 2}]
|
|
assert second.has_more is False
|
|
|
|
status = await reader.get_run_status(thread_id="thread-a", run_id="run-a")
|
|
assert status is not None
|
|
assert status.status == "error"
|
|
assert status.error == "failed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("event_backend", ["memory", "jsonl", "db"])
|
|
async def test_event_views_detach_nested_payloads_from_store(tmp_path, event_backend):
|
|
engine = None
|
|
runs = MemoryRunStore()
|
|
if event_backend == "memory":
|
|
events = MemoryRunEventStore()
|
|
elif event_backend == "jsonl":
|
|
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
|
|
|
|
events = JsonlRunEventStore(tmp_path / "events")
|
|
else:
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.runtime.events.store.db import DbRunEventStore
|
|
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'events.db'}")
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
events = DbRunEventStore(async_sessionmaker(engine, expire_on_commit=False))
|
|
|
|
await _put_run(runs, "run-a", "thread-a")
|
|
try:
|
|
await events.put(
|
|
thread_id="thread-a",
|
|
run_id="run-a",
|
|
event_type="test.event",
|
|
category="trace",
|
|
content={"text": ["original"]},
|
|
metadata={"custom": {"value": "original"}},
|
|
)
|
|
|
|
reader = StoreRunEvidenceReader(runs, events, user_id=None)
|
|
page = await reader.list_run_events(thread_id="thread-a", run_id="run-a", after_seq=None, limit=10)
|
|
page.items[0].content["text"][0] = "changed"
|
|
page.items[0].metadata["custom"]["value"] = "changed"
|
|
|
|
reread = await reader.list_run_events(thread_id="thread-a", run_id="run-a", after_seq=None, limit=10)
|
|
stored = await events.list_events("thread-a", "run-a", user_id=None)
|
|
assert reread.items[0].content == {"text": ["original"]}
|
|
assert reread.items[0].metadata["custom"] == {"value": "original"}
|
|
assert stored[0]["content"] == {"text": ["original"]}
|
|
assert stored[0]["metadata"]["custom"] == {"value": "original"}
|
|
finally:
|
|
if engine is not None:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_db_event_reads_use_reader_scope_not_ambient_user(tmp_path):
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.run import RunRepository
|
|
from deerflow.runtime import user_context
|
|
from deerflow.runtime.events.store.db import DbRunEventStore
|
|
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'scoped-events.db'}")
|
|
try:
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
runs = RunRepository(factory)
|
|
events = DbRunEventStore(factory)
|
|
await runs.put("run-a", thread_id="thread-a", user_id="alice")
|
|
|
|
alice_token = user_context.set_current_user(SimpleNamespace(id="alice"))
|
|
try:
|
|
await events.put(
|
|
thread_id="thread-a",
|
|
run_id="run-a",
|
|
event_type="test.event",
|
|
category="trace",
|
|
content={"owner": "alice"},
|
|
)
|
|
finally:
|
|
user_context.reset_current_user(alice_token)
|
|
|
|
for ambient_user in (None, SimpleNamespace(id="bob")):
|
|
ambient_token = user_context._current_user.set(ambient_user)
|
|
try:
|
|
for reader_user in (None, "alice"):
|
|
reader = StoreRunEvidenceReader(runs, events, user_id=reader_user)
|
|
page = await reader.list_run_events(
|
|
thread_id="thread-a",
|
|
run_id="run-a",
|
|
after_seq=None,
|
|
limit=10,
|
|
)
|
|
assert [item.content for item in page.items] == [{"owner": "alice"}]
|
|
finally:
|
|
user_context._current_user.reset(ambient_token)
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_reader_hides_runs_outside_scope_and_rejects_cursor_from_another_scope():
|
|
runs = MemoryRunStore()
|
|
events = MemoryRunEventStore()
|
|
await _put_run(runs, "run-a", "thread-a", user_id="user-1")
|
|
first = await StoreRunEvidenceReader(runs, events, user_id="user-1").list_changed_runs(cursor=None, limit=1)
|
|
|
|
other = StoreRunEvidenceReader(runs, events, user_id="user-2")
|
|
with pytest.raises(InvalidRunEvidenceCursor, match="cursor scope"):
|
|
await other.list_changed_runs(cursor=first.next_cursor, limit=1)
|
|
assert await other.get_run_status(thread_id="thread-a", run_id="run-a") is None
|
|
assert (await other.list_run_events(thread_id="thread-a", run_id="run-a", after_seq=None, limit=10)).items == ()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_memory_progress_updates_do_not_advance_changed_run_cursor():
|
|
runs = MemoryRunStore()
|
|
reader = StoreRunEvidenceReader(runs, MemoryRunEventStore(), user_id="user-1")
|
|
await _put_run(runs, "run-a", "thread-a")
|
|
await runs.start_run("run-a")
|
|
before_progress = await reader.list_changed_runs(cursor=None, limit=10)
|
|
|
|
await runs.update_run_progress("run-a", total_tokens=10, last_ai_message="working")
|
|
|
|
after_progress = await reader.list_changed_runs(cursor=before_progress.next_cursor, limit=10)
|
|
assert after_progress.items == ()
|
|
assert after_progress.next_cursor == before_progress.next_cursor
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_changed_run_cursor_survives_repository_restart(tmp_path):
|
|
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
|
|
from deerflow.persistence.run import RunRepository
|
|
|
|
url = f"sqlite+aiosqlite:///{tmp_path / 'evidence.db'}"
|
|
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
|
try:
|
|
repo = RunRepository(get_session_factory())
|
|
await repo.put("run-a", thread_id="thread-a", user_id="user-1")
|
|
reader = StoreRunEvidenceReader(repo, MemoryRunEventStore(), user_id="user-1")
|
|
first = await reader.list_changed_runs(cursor=None, limit=1)
|
|
await repo.put("run-b", thread_id="thread-b", user_id="user-1")
|
|
cursor = first.next_cursor
|
|
finally:
|
|
await close_engine()
|
|
|
|
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
|
|
try:
|
|
restarted = StoreRunEvidenceReader(
|
|
RunRepository(get_session_factory()),
|
|
MemoryRunEventStore(),
|
|
user_id="user-1",
|
|
)
|
|
page = await restarted.list_changed_runs(cursor=cursor, limit=10)
|
|
assert [item.run_id for item in page.items] == ["run-b"]
|
|
finally:
|
|
await close_engine()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_progress_updates_do_not_advance_changed_run_cursor(tmp_path):
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.run import RunRepository
|
|
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'progress.db'}")
|
|
try:
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
repo = RunRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
await repo.put("run-a", thread_id="thread-a", user_id="user-1")
|
|
await repo.start_run("run-a")
|
|
reader = StoreRunEvidenceReader(repo, MemoryRunEventStore(), user_id="user-1")
|
|
before_progress = await reader.list_changed_runs(cursor=None, limit=10)
|
|
|
|
await repo.update_run_progress("run-a", total_tokens=10, last_ai_message="working")
|
|
|
|
after_progress = await reader.list_changed_runs(cursor=before_progress.next_cursor, limit=10)
|
|
assert after_progress.items == ()
|
|
assert after_progress.next_cursor == before_progress.next_cursor
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_atomic_interrupt_allocates_before_run_lock_and_shares_position(tmp_path):
|
|
from sqlalchemy import event
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.run import RunRepository
|
|
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'atomic.db'}")
|
|
try:
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
repo = RunRepository(async_sessionmaker(engine, expire_on_commit=False))
|
|
await repo.put(
|
|
"run-old",
|
|
thread_id="thread-a",
|
|
user_id="user-1",
|
|
status="running",
|
|
owner_worker_id="worker-1",
|
|
)
|
|
statements: list[str] = []
|
|
|
|
@event.listens_for(engine.sync_engine, "before_cursor_execute")
|
|
def capture_statement(_connection, _cursor, statement, _parameters, _context, _executemany):
|
|
statements.append(" ".join(statement.lower().split()))
|
|
|
|
new_run, claimed = await repo.create_thread_operation_atomic(
|
|
"run-new",
|
|
thread_id="thread-a",
|
|
user_id="user-1",
|
|
owner_worker_id="worker-1",
|
|
lease_expires_at=None,
|
|
multitask_strategy="interrupt",
|
|
)
|
|
|
|
assert [run["run_id"] for run in claimed] == ["run-old"]
|
|
assert claimed[0]["change_seq"] == new_run["change_seq"]
|
|
clock_update = next(index for index, statement in enumerate(statements) if statement.startswith("update run_change_clock"))
|
|
run_lock_query = next(index for index, statement in enumerate(statements) if statement.startswith("select runs."))
|
|
assert clock_update < run_lock_query
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sql_change_positions_are_unique_across_concurrent_threads(tmp_path):
|
|
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
|
|
|
|
from deerflow.persistence.base import Base
|
|
from deerflow.persistence.run import RunRepository
|
|
|
|
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'concurrent.db'}")
|
|
try:
|
|
async with engine.begin() as connection:
|
|
await connection.run_sync(Base.metadata.create_all)
|
|
factory = async_sessionmaker(engine, expire_on_commit=False)
|
|
repositories = [RunRepository(factory) for _ in range(12)]
|
|
await asyncio.gather(
|
|
*(
|
|
repository.put(
|
|
f"run-{index:02d}",
|
|
thread_id=f"thread-{index:02d}",
|
|
user_id="user-1",
|
|
)
|
|
for index, repository in enumerate(repositories)
|
|
)
|
|
)
|
|
|
|
rows = await repositories[0].list_changed(
|
|
after_change_seq=-1,
|
|
after_run_id="",
|
|
user_id="user-1",
|
|
limit=20,
|
|
)
|
|
positions = [row["change_seq"] for row in rows]
|
|
assert len(rows) == 12
|
|
assert positions == sorted(positions)
|
|
assert len(set(positions)) == len(positions)
|
|
finally:
|
|
await engine.dispose()
|