diff --git a/backend/packages/harness/deerflow/runtime/events/store/memory.py b/backend/packages/harness/deerflow/runtime/events/store/memory.py index a2bb54819..c8e732d44 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/memory.py +++ b/backend/packages/harness/deerflow/runtime/events/store/memory.py @@ -19,6 +19,14 @@ class MemoryRunEventStore(RunEventStore): # kept in seq order so message pagination is O(log m + page) via bisect # instead of re-scanning every event on each request. self._messages: dict[str, list[dict]] = {} # thread_id -> seq-sorted message list + # Run-keyed projections of the two lists above (same dict objects, no + # copies), kept in seq order. Per-run reads then cost O(events-in-run) + # instead of O(events-in-thread): without these, ``list_events`` and + # ``list_messages_by_run`` re-scan the whole thread's event log on every + # request even though one run holds only a handful of events. This is + # the per-run analogue of the thread-wide ``_messages`` projection. + self._events_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted events + self._messages_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted messages self._seq_counters: dict[str, int] = {} # thread_id -> last assigned seq def _next_seq(self, thread_id: str) -> int: @@ -50,8 +58,10 @@ class MemoryRunEventStore(RunEventStore): "created_at": created_at or datetime.now(UTC).isoformat(), } self._events.setdefault(thread_id, []).append(record) + self._events_by_run.setdefault(thread_id, {}).setdefault(run_id, []).append(record) if category == "message": self._messages.setdefault(thread_id, []).append(record) + self._messages_by_run.setdefault(thread_id, {}).setdefault(run_id, []).append(record) return record async def put( @@ -100,23 +110,27 @@ class MemoryRunEventStore(RunEventStore): return messages[-limit:] async def list_events(self, thread_id, run_id, *, event_types=None, limit=500): - all_events = self._events.get(thread_id, []) - filtered = [e for e in all_events if e["run_id"] == run_id] + # ``_events_by_run`` is already scoped to this run and seq-ordered, so we + # touch only this run's events instead of scanning the whole thread. + run_events = self._events_by_run.get(thread_id, {}).get(run_id, []) if event_types is not None: - filtered = [e for e in filtered if e["event_type"] in event_types] - return filtered[:limit] + run_events = [e for e in run_events if e["event_type"] in event_types] + return run_events[:limit] async def list_messages_by_run(self, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None): - all_events = self._events.get(thread_id, []) - filtered = [e for e in all_events if e["run_id"] == run_id and e["category"] == "message"] - if before_seq is not None: - filtered = [e for e in filtered if e["seq"] < before_seq] + # Per-run, messages-only, seq-sorted: the seq window is a contiguous + # slice located with bisect (O(log m_run)) over only this run's + # messages, instead of re-scanning the whole thread's event log. + messages = self._messages_by_run.get(thread_id, {}).get(run_id, []) + lo = 0 if after_seq is None else bisect.bisect_right(messages, after_seq, key=lambda e: e["seq"]) + hi = len(messages) if before_seq is None else bisect.bisect_left(messages, before_seq, key=lambda e: e["seq"]) + window = messages[lo:hi] + # An ``after_seq`` cursor pages forward (first ``limit``); otherwise + # return the last ``limit`` (the latest page, or the page ending just + # before ``before_seq``). Matches the prior filter-based semantics. if after_seq is not None: - filtered = [e for e in filtered if e["seq"] > after_seq] - if after_seq is not None: - return filtered[:limit] - else: - return filtered[-limit:] if len(filtered) > limit else filtered + return window[:limit] + return window[-limit:] async def count_messages(self, thread_id): return len(self._messages.get(thread_id, [])) @@ -124,6 +138,8 @@ class MemoryRunEventStore(RunEventStore): async def delete_by_thread(self, thread_id): events = self._events.pop(thread_id, []) self._messages.pop(thread_id, None) + self._events_by_run.pop(thread_id, None) + self._messages_by_run.pop(thread_id, None) self._seq_counters.pop(thread_id, None) return len(events) @@ -136,4 +152,7 @@ class MemoryRunEventStore(RunEventStore): self._events[thread_id] = remaining # Keep the message projection in lockstep (same surviving dict objects). self._messages[thread_id] = [e for e in remaining if e["category"] == "message"] + # Drop the deleted run from the run-keyed projections. + self._events_by_run.get(thread_id, {}).pop(run_id, None) + self._messages_by_run.get(thread_id, {}).pop(run_id, None) return removed diff --git a/backend/tests/test_run_event_store_by_run_index.py b/backend/tests/test_run_event_store_by_run_index.py new file mode 100644 index 000000000..a16dba16f --- /dev/null +++ b/backend/tests/test_run_event_store_by_run_index.py @@ -0,0 +1,130 @@ +"""Regression tests for MemoryRunEventStore's run-keyed event/message index. + +``list_events`` and ``list_messages_by_run`` are served from per-run +projections so a single run's reads cost O(events-in-run) instead of +re-scanning O(events-in-thread). These tests pin the indexed implementation to +the exact semantics of a brute-force full-thread scan -- including interleaved +trace events (non-contiguous message seqs), both cursors supplied at once, and +index upkeep after ``delete_by_run`` -- so the optimization can never silently +drift from the reference behavior. +""" + +import pytest + +from deerflow.runtime.events.store.memory import MemoryRunEventStore + + +def _ref_messages_by_run(records, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None): + """Brute-force reference: the pre-index full-thread scan it replaced.""" + filtered = [e for e in records if e["thread_id"] == thread_id and e["run_id"] == run_id and e["category"] == "message"] + if before_seq is not None: + filtered = [e for e in filtered if e["seq"] < before_seq] + if after_seq is not None: + filtered = [e for e in filtered if e["seq"] > after_seq] + if after_seq is not None: + return filtered[:limit] + return filtered[-limit:] if len(filtered) > limit else filtered + + +def _ref_events(records, thread_id, run_id, *, event_types=None, limit=500): + filtered = [e for e in records if e["thread_id"] == thread_id and e["run_id"] == run_id] + if event_types is not None: + filtered = [e for e in filtered if e["event_type"] in event_types] + return filtered[:limit] + + +async def _seed(store): + """Two runs interleaved within one thread; messages and traces mixed so + each run's message seqs are non-contiguous (the bisect must handle gaps).""" + plan = [ + ("run-a", "message"), + ("run-a", "trace"), + ("run-b", "message"), + ("run-a", "message"), + ("run-b", "trace"), + ("run-b", "message"), + ("run-a", "trace"), + ("run-a", "message"), + ("run-b", "message"), + ("run-a", "message"), + ("run-b", "message"), + ("run-a", "message"), + ] + records = [] + for i, (run_id, category) in enumerate(plan): + rec = await store.put(thread_id="t1", run_id=run_id, event_type=f"e{i}", category=category, content=str(i)) + records.append(rec) + return records + + +@pytest.mark.anyio +async def test_list_messages_by_run_matches_reference_across_cursors(): + store = MemoryRunEventStore() + records = await _seed(store) + seqs = [r["seq"] for r in records] + cursors = [None, 0, *seqs, max(seqs) + 1] + for run_id in ("run-a", "run-b", "run-missing"): + for limit in (1, 2, 3, 50): + for before_seq in cursors: + for after_seq in cursors: + got = await store.list_messages_by_run("t1", run_id, limit=limit, before_seq=before_seq, after_seq=after_seq) + want = _ref_messages_by_run(records, "t1", run_id, limit=limit, before_seq=before_seq, after_seq=after_seq) + assert got == want, (run_id, limit, before_seq, after_seq) + + +@pytest.mark.anyio +async def test_list_events_matches_reference_with_filters(): + store = MemoryRunEventStore() + records = await _seed(store) + all_types = sorted({r["event_type"] for r in records}) + for run_id in ("run-a", "run-b", "run-missing"): + assert await store.list_events("t1", run_id) == _ref_events(records, "t1", run_id) + assert await store.list_events("t1", run_id, limit=2) == _ref_events(records, "t1", run_id, limit=2) + for et in all_types: + assert await store.list_events("t1", run_id, event_types=[et]) == _ref_events(records, "t1", run_id, event_types=[et]) + + +@pytest.mark.anyio +async def test_run_keyed_index_partitions_every_event(): + """Every stored event is filed under exactly its (thread, run), each run's + list is seq-ordered, and the union reconstructs the flat event log.""" + store = MemoryRunEventStore() + records = await _seed(store) + indexed = [e for run_events in store._events_by_run["t1"].values() for e in run_events] + assert sorted(e["seq"] for e in indexed) == sorted(r["seq"] for r in records) + for run_id, run_events in store._events_by_run["t1"].items(): + assert all(e["run_id"] == run_id for e in run_events) + assert [e["seq"] for e in run_events] == sorted(e["seq"] for e in run_events) + for run_id, run_msgs in store._messages_by_run["t1"].items(): + assert all(e["run_id"] == run_id and e["category"] == "message" for e in run_msgs) + + +@pytest.mark.anyio +async def test_run_index_stays_in_lockstep_after_delete_by_run(): + store = MemoryRunEventStore() + await _seed(store) + removed = await store.delete_by_run("t1", "run-a") + assert removed == 7 # run-a: 5 messages + 2 traces + + # The deleted run vanishes from both per-run reads. + assert await store.list_events("t1", "run-a") == [] + assert await store.list_messages_by_run("t1", "run-a") == [] + assert "run-a" not in store._events_by_run.get("t1", {}) + assert "run-a" not in store._messages_by_run.get("t1", {}) + + # The surviving run is untouched, and the thread-wide projection agrees. + msgs_b = await store.list_messages_by_run("t1", "run-b") + assert len(msgs_b) == 4 + assert all(m["run_id"] == "run-b" for m in msgs_b) + assert all(m["run_id"] == "run-b" for m in await store.list_messages("t1")) + + +@pytest.mark.anyio +async def test_delete_by_thread_clears_run_indexes(): + store = MemoryRunEventStore() + await _seed(store) + await store.delete_by_thread("t1") + assert "t1" not in store._events_by_run + assert "t1" not in store._messages_by_run + assert await store.list_events("t1", "run-a") == [] + assert await store.list_messages_by_run("t1", "run-b") == []