mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-21 12:06:18 +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>
451 lines
19 KiB
Python
451 lines
19 KiB
Python
"""JSONL file-backed RunEventStore implementation.
|
|
|
|
Each run's events are stored in a single file:
|
|
``.deer-flow/threads/{thread_id}/runs/{run_id}.jsonl``
|
|
|
|
All categories (message, trace, lifecycle) are in the same file.
|
|
This backend is suitable for lightweight single-node deployments.
|
|
|
|
**Single-process guarantee**: the in-memory seq counter is process-local.
|
|
Multi-process deployments sharing the same directory will produce duplicate
|
|
or non-monotonic seq values. Use ``DbRunEventStore`` for multi-process or
|
|
high-concurrency deployments.
|
|
|
|
File I/O is offloaded to a thread pool via ``asyncio.to_thread`` so the
|
|
event loop is never blocked. Per-thread ``asyncio.Lock`` objects serialise
|
|
writes within a single process to prevent interleaved JSONL lines.
|
|
|
|
Known trade-off: ``list_messages()`` must scan all run files for a
|
|
thread since messages from multiple runs need unified seq ordering.
|
|
``list_events()`` reads only one file -- the fast path.
|
|
|
|
Read records using physical newline boundaries, not ``str.splitlines()``:
|
|
Unicode line separators are valid JSON string content and must stay inside
|
|
their record. ``read_text`` normalizes CRLF before the LF split.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import re
|
|
import weakref
|
|
from collections.abc import Callable, Coroutine
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from deerflow.runtime.events.message_identity import message_identity
|
|
from deerflow.runtime.events.store.base import RunEventStore, match_ai_message_run_id, normalize_message_ids
|
|
from deerflow.runtime.user_context import AUTO, _AutoSentinel
|
|
from deerflow.utils.thread_id import validate_thread_id
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SAFE_ID_PATTERN = re.compile(r"^[A-Za-z0-9_\-]+$")
|
|
|
|
|
|
class JsonlRunEventStore(RunEventStore):
|
|
def __init__(self, base_dir: str | Path | None = None):
|
|
self._base_dir = Path(base_dir) if base_dir else Path(".deer-flow")
|
|
self._seq_counters: dict[str, int] = {} # thread_id -> current max seq
|
|
# Weak ownership avoids leaking one lock per historical thread without
|
|
# splitting a live lock generation while a holder/waiter still owns it.
|
|
self._write_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = weakref.WeakValueDictionary()
|
|
|
|
def _get_write_lock(self, thread_id: str) -> asyncio.Lock:
|
|
lock = self._write_locks.get(thread_id)
|
|
if lock is None:
|
|
lock = asyncio.Lock()
|
|
self._write_locks[thread_id] = lock
|
|
return lock
|
|
|
|
async def _run_mutation[T](self, thread_id: str, operation: Callable[[], Coroutine[Any, Any, T]]) -> T:
|
|
"""Drain an admitted mutation before propagating caller cancellation.
|
|
|
|
Cancelling ``to_thread`` only stops its awaiter, not the filesystem
|
|
worker. Keep the thread lock through I/O, rollback and bookkeeping,
|
|
even if the caller is cancelled repeatedly. Queued callers can still
|
|
cancel before acquiring the lock, without starting a mutation.
|
|
"""
|
|
async with self._get_write_lock(thread_id):
|
|
task = asyncio.create_task(operation(), name=f"jsonl-mutation:{thread_id}")
|
|
cancellation: asyncio.CancelledError | None = None
|
|
while not task.done():
|
|
try:
|
|
await asyncio.shield(task)
|
|
except asyncio.CancelledError as exc:
|
|
if cancellation is None:
|
|
cancellation = exc
|
|
except Exception:
|
|
# Retrieve the failure below, after preserving any earlier
|
|
# cancellation. The operation has already finished rollback.
|
|
break
|
|
if cancellation is not None:
|
|
try:
|
|
task.result()
|
|
except Exception as exc:
|
|
raise cancellation from exc
|
|
raise cancellation
|
|
return task.result()
|
|
|
|
@staticmethod
|
|
def _validate_id(value: str, label: str) -> str:
|
|
"""Validate that an ID is safe for use in filesystem paths."""
|
|
if not value or not _SAFE_ID_PATTERN.match(value):
|
|
raise ValueError(f"Invalid {label}: must be alphanumeric/dash/underscore, got {value!r}")
|
|
return value
|
|
|
|
def _thread_dir(self, thread_id: str) -> Path:
|
|
validate_thread_id(thread_id)
|
|
return self._base_dir / "threads" / thread_id / "runs"
|
|
|
|
def _run_file(self, thread_id: str, run_id: str) -> Path:
|
|
self._validate_id(run_id, "run_id")
|
|
return self._thread_dir(thread_id) / f"{run_id}.jsonl"
|
|
|
|
def _next_seq(self, thread_id: str) -> int:
|
|
self._seq_counters[thread_id] = self._seq_counters.get(thread_id, 0) + 1
|
|
return self._seq_counters[thread_id]
|
|
|
|
def _compute_max_seq(self, thread_id: str) -> int:
|
|
"""Scan all run files for a thread and return the current max seq (blocking I/O)."""
|
|
max_seq = 0
|
|
thread_dir = self._thread_dir(thread_id)
|
|
if thread_dir.exists():
|
|
for f in thread_dir.glob("*.jsonl"):
|
|
for line in f.read_text(encoding="utf-8").strip().split("\n"):
|
|
try:
|
|
record = json.loads(line)
|
|
max_seq = max(max_seq, record.get("seq", 0))
|
|
except json.JSONDecodeError:
|
|
logger.debug("Skipping malformed JSONL line in %s", f)
|
|
return max_seq
|
|
|
|
async def _ensure_seq_loaded(self, thread_id: str) -> None:
|
|
"""Load max seq from existing files into the in-memory counter (non-blocking)."""
|
|
if thread_id in self._seq_counters:
|
|
return
|
|
max_seq = await asyncio.to_thread(self._compute_max_seq, thread_id)
|
|
self._seq_counters[thread_id] = max_seq
|
|
|
|
def _write_record(self, record: dict) -> None:
|
|
path = self._run_file(record["thread_id"], record["run_id"])
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with open(path, "a", encoding="utf-8") as f:
|
|
f.write(json.dumps(record, default=str, ensure_ascii=False) + "\n")
|
|
|
|
def _read_thread_events(self, thread_id: str) -> list[dict]:
|
|
"""Read all events for a thread, sorted by seq (blocking I/O)."""
|
|
events = []
|
|
thread_dir = self._thread_dir(thread_id)
|
|
if not thread_dir.exists():
|
|
return events
|
|
for f in sorted(thread_dir.glob("*.jsonl")):
|
|
for line in f.read_text(encoding="utf-8").strip().split("\n"):
|
|
if not line:
|
|
continue
|
|
try:
|
|
events.append(json.loads(line))
|
|
except json.JSONDecodeError:
|
|
logger.debug("Skipping malformed JSONL line in %s", f)
|
|
events.sort(key=lambda e: e.get("seq", 0))
|
|
return events
|
|
|
|
def _read_run_events(self, thread_id: str, run_id: str) -> list[dict]:
|
|
"""Read events for a specific run file (blocking I/O)."""
|
|
path = self._run_file(thread_id, run_id)
|
|
if not path.exists():
|
|
return []
|
|
events = []
|
|
for line in path.read_text(encoding="utf-8").strip().split("\n"):
|
|
if not line:
|
|
continue
|
|
try:
|
|
events.append(json.loads(line))
|
|
except json.JSONDecodeError:
|
|
logger.debug("Skipping malformed JSONL line in %s", path)
|
|
events.sort(key=lambda e: e.get("seq", 0))
|
|
return events
|
|
|
|
def _delete_thread_files(self, thread_id: str) -> None:
|
|
thread_dir = self._thread_dir(thread_id)
|
|
if thread_dir.exists():
|
|
for f in thread_dir.glob("*.jsonl"):
|
|
f.unlink()
|
|
|
|
def _delete_run_file(self, thread_id: str, run_id: str) -> None:
|
|
path = self._run_file(thread_id, run_id)
|
|
if path.exists():
|
|
path.unlink()
|
|
|
|
async def put(self, *, thread_id, run_id, event_type, category, content="", metadata=None, created_at=None):
|
|
async def mutate():
|
|
await self._ensure_seq_loaded(thread_id)
|
|
seq = self._next_seq(thread_id)
|
|
record = {
|
|
"thread_id": thread_id,
|
|
"run_id": run_id,
|
|
"event_type": event_type,
|
|
"category": category,
|
|
"content": content,
|
|
"metadata": metadata or {},
|
|
"seq": seq,
|
|
"created_at": created_at or datetime.now(UTC).isoformat(),
|
|
}
|
|
await asyncio.to_thread(self._write_record, record)
|
|
return record
|
|
|
|
return await self._run_mutation(thread_id, mutate)
|
|
|
|
async def put_batch(self, events):
|
|
"""Persist a batch of events under a per-thread write lock.
|
|
|
|
All seq numbers for the batch are reserved under a single per-thread
|
|
write lock. Records are grouped by run_id and appended to their own
|
|
run files while that lock is held. If a write fails and rollback
|
|
succeeds, already-appended groups for the current thread are restored
|
|
so callers (e.g. worker.py's flush-retry path) may safely re-buffer
|
|
that thread's batch. When a batch contains multiple thread IDs, thread
|
|
groups are processed sequentially, so a later failure does not roll
|
|
back earlier thread groups. Cancellation drains the current thread group
|
|
before propagating, without starting subsequent groups. This rollback
|
|
does not make a multi-file batch crash-atomic.
|
|
"""
|
|
if not events:
|
|
return []
|
|
|
|
# Group by thread_id; each thread has its own write lock and seq counter.
|
|
by_thread: dict[str, list[dict[str, Any]]] = {}
|
|
for ev in events:
|
|
by_thread.setdefault(ev["thread_id"], []).append(ev)
|
|
|
|
results: list[dict[str, Any]] = []
|
|
for thread_id, batch in by_thread.items():
|
|
records = await self._write_batch_async(thread_id, batch)
|
|
results.extend(records)
|
|
return results
|
|
|
|
async def put_if_absent(
|
|
self,
|
|
*,
|
|
thread_id,
|
|
run_id,
|
|
event_type,
|
|
category,
|
|
content="",
|
|
metadata=None,
|
|
created_at=None,
|
|
):
|
|
async def mutate():
|
|
existing = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
|
|
for event in existing:
|
|
if event.get("event_type") == event_type:
|
|
return event, False
|
|
await self._ensure_seq_loaded(thread_id)
|
|
record = {
|
|
"thread_id": thread_id,
|
|
"run_id": run_id,
|
|
"event_type": event_type,
|
|
"category": category,
|
|
"content": content,
|
|
"metadata": metadata or {},
|
|
"seq": self._next_seq(thread_id),
|
|
"created_at": created_at or datetime.now(UTC).isoformat(),
|
|
}
|
|
await asyncio.to_thread(self._write_record, record)
|
|
return record, True
|
|
|
|
return await self._run_mutation(thread_id, mutate)
|
|
|
|
async def _write_batch_async(self, thread_id: str, batch: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
async def mutate():
|
|
await self._ensure_seq_loaded(thread_id)
|
|
records: list[dict[str, Any]] = []
|
|
for ev in batch:
|
|
seq = self._next_seq(thread_id)
|
|
record = {
|
|
"thread_id": thread_id,
|
|
"run_id": ev["run_id"],
|
|
"event_type": ev["event_type"],
|
|
"category": ev["category"],
|
|
"content": ev.get("content", ""),
|
|
"metadata": ev.get("metadata") or {},
|
|
"seq": seq,
|
|
"created_at": ev.get("created_at") or datetime.now(UTC).isoformat(),
|
|
}
|
|
records.append(record)
|
|
records_by_run: dict[str, list[dict[str, Any]]] = {}
|
|
for record in records:
|
|
records_by_run.setdefault(record["run_id"], []).append(record)
|
|
run_batches = [(self._run_file(thread_id, run_id), run_records) for run_id, run_records in records_by_run.items()]
|
|
await asyncio.to_thread(self._append_record_groups, run_batches)
|
|
return records
|
|
|
|
return await self._run_mutation(thread_id, mutate)
|
|
|
|
def _append_records(self, path: Path, records: list[dict[str, Any]]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
lines = "".join(json.dumps(r, default=str, ensure_ascii=False) + "\n" for r in records)
|
|
with open(path, "a", encoding="utf-8") as f:
|
|
f.write(lines)
|
|
|
|
def _append_record_groups(self, groups: list[tuple[Path, list[dict[str, Any]]]]) -> None:
|
|
"""Append run groups and restore their original sizes if one fails."""
|
|
original_sizes: dict[Path, int | None] = {}
|
|
try:
|
|
for path, records in groups:
|
|
original_sizes[path] = path.stat().st_size if path.exists() else None
|
|
self._append_records(path, records)
|
|
except Exception:
|
|
for path, original_size in original_sizes.items():
|
|
try:
|
|
if original_size is None:
|
|
if path.exists():
|
|
path.unlink()
|
|
else:
|
|
with open(path, "r+b") as f:
|
|
f.truncate(original_size)
|
|
except OSError:
|
|
logger.error(
|
|
"Failed to roll back JSONL batch append for %s; retrying the batch may create duplicate records",
|
|
path,
|
|
exc_info=True,
|
|
)
|
|
raise
|
|
|
|
async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None, user_id: str | None | _AutoSentinel = AUTO):
|
|
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
|
|
messages = [e for e in all_events if e.get("category") == "message"]
|
|
|
|
if before_seq is not None:
|
|
messages = [e for e in messages if e["seq"] < before_seq]
|
|
return messages[-limit:]
|
|
elif after_seq is not None:
|
|
messages = [e for e in messages if e["seq"] > after_seq]
|
|
return messages[:limit]
|
|
else:
|
|
return messages[-limit:]
|
|
|
|
async def find_latest_ai_message_run_ids(
|
|
self,
|
|
thread_id: str,
|
|
message_ids: set[str],
|
|
*,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
) -> dict[str, str]:
|
|
pending = normalize_message_ids(message_ids)
|
|
if not pending:
|
|
return {}
|
|
|
|
# Keep the one-pass view stable against this backend's supported
|
|
# single-process writers. Without the write lock, reading run files one
|
|
# by one can mix events from opposite sides of a concurrent append.
|
|
async with self._get_write_lock(thread_id):
|
|
events = await asyncio.to_thread(self._read_thread_events, thread_id)
|
|
result: dict[str, str] = {}
|
|
for event in reversed(events):
|
|
match = match_ai_message_run_id(event, pending)
|
|
if match is None:
|
|
continue
|
|
message_id, run_id = match
|
|
result[message_id] = run_id
|
|
pending.remove(message_id)
|
|
if not pending:
|
|
break
|
|
return result
|
|
|
|
async def list_events(
|
|
self,
|
|
thread_id,
|
|
run_id,
|
|
*,
|
|
event_types=None,
|
|
task_id=None,
|
|
limit=500,
|
|
after_seq=None,
|
|
user_id: str | None | _AutoSentinel = AUTO,
|
|
):
|
|
events = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
|
|
if event_types is not None:
|
|
events = [e for e in events if e.get("event_type") in event_types]
|
|
if task_id is not None:
|
|
events = [e for e in events if (e.get("metadata") or {}).get("task_id") == task_id]
|
|
if after_seq is not None:
|
|
events = [e for e in events if e.get("seq", 0) > after_seq]
|
|
return events[:limit]
|
|
|
|
async def list_messages_by_run(self, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None):
|
|
events = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
|
|
filtered = [e for e in events if e.get("category") == "message"]
|
|
if before_seq is not None:
|
|
filtered = [e for e in filtered if e.get("seq", 0) < before_seq]
|
|
if after_seq is not None:
|
|
filtered = [e for e in filtered if e.get("seq", 0) > after_seq]
|
|
if after_seq is not None:
|
|
return filtered[:limit]
|
|
else:
|
|
return filtered[-limit:] if len(filtered) > limit else filtered
|
|
|
|
async def get_last_visible_ai_seq_by_run(self, thread_id, run_ids, *, user_id: str | None | _AutoSentinel = AUTO):
|
|
def _scan() -> dict[str, int]:
|
|
result: dict[str, int] = {}
|
|
for run_id in run_ids:
|
|
for event in reversed(self._read_run_events(thread_id, run_id)):
|
|
caller = str((event.get("metadata") or {}).get("caller", ""))
|
|
if event.get("category") == "message" and event.get("event_type") in {"llm.ai.response", "ai_message"} and not caller.startswith("middleware:"):
|
|
result[run_id] = event["seq"]
|
|
break
|
|
return result
|
|
|
|
return await asyncio.to_thread(_scan)
|
|
|
|
async def count_messages(self, thread_id):
|
|
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
|
|
return sum(1 for e in all_events if e.get("category") == "message")
|
|
|
|
async def get_message_seqs(self, thread_id, identities, *, user_id: str | None | _AutoSentinel = AUTO):
|
|
wanted = set(identities)
|
|
if not wanted:
|
|
return {}
|
|
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
|
|
found: dict[str, int] = {}
|
|
for event in all_events:
|
|
if event.get("category") != "message":
|
|
continue
|
|
content = event.get("content")
|
|
if not isinstance(content, dict):
|
|
continue
|
|
identity = message_identity(content)
|
|
# Earliest seq wins: a message re-persisted later keeps the position
|
|
# it first occupied in the feed.
|
|
if identity in wanted and identity not in found:
|
|
found[identity] = event["seq"]
|
|
# Later events can only be re-persisted copies that already lose
|
|
# that tiebreak, so the scan ends with the last wanted seq.
|
|
if len(found) == len(wanted):
|
|
break
|
|
return found
|
|
|
|
async def delete_by_thread(self, thread_id):
|
|
async def mutate():
|
|
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
|
|
count = len(all_events)
|
|
await asyncio.to_thread(self._delete_thread_files, thread_id)
|
|
self._seq_counters.pop(thread_id, None)
|
|
# Mutations already queued on this lock resume after deletion; with
|
|
# files and the counter cleared, they recreate the thread at seq 1.
|
|
return count
|
|
|
|
return await self._run_mutation(thread_id, mutate)
|
|
|
|
async def delete_by_run(self, thread_id, run_id):
|
|
async def mutate():
|
|
events = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
|
|
count = len(events)
|
|
await asyncio.to_thread(self._delete_run_file, thread_id, run_id)
|
|
return count
|
|
|
|
return await self._run_mutation(thread_id, mutate)
|