fix(gateway): preserve exact history attribution beyond event page limits (#4953)

* fix(gateway): preserve exact history run attribution

* fix(gateway): make history migration authoritative

* docs(runtime): keep history contract within guidance budget

* fix(runtime): fence final run duration write
This commit is contained in:
Beautyl0ve 2026-08-25 08:22:57 +08:00 committed by GitHub
parent ff0a6768c2
commit e8410cebfc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 1584 additions and 103 deletions

View File

@ -66,7 +66,7 @@ from deerflow.runtime.goal import (
)
from deerflow.runtime.journal import build_branch_history_seed_events
from deerflow.runtime.runs.manager import ConflictError
from deerflow.runtime.runs.worker import valid_duration_entry
from deerflow.runtime.runs.worker import RUN_MESSAGE_IDS_METADATA_KEY, valid_duration_entry, valid_run_message_id_entry
from deerflow.runtime.secret_context import redact_metadata_secrets
from deerflow.runtime.user_context import get_effective_user_id
from deerflow.utils.file_io import run_file_io
@ -1355,6 +1355,96 @@ def _checkpoint_run_durations(metadata: Any) -> dict[str, int]:
return {run_id: duration_seconds for run_id, duration_seconds in raw_durations.items() if valid_duration_entry(run_id, duration_seconds)}
def _checkpoint_run_message_ids(metadata: Any) -> dict[str, str]:
raw_message_run_ids = metadata.get(RUN_MESSAGE_IDS_METADATA_KEY) if isinstance(metadata, dict) else None
if not isinstance(raw_message_run_ids, dict):
return {}
return {message_id: run_id for message_id, run_id in raw_message_run_ids.items() if valid_run_message_id_entry(message_id, run_id)}
async def _load_run_durations(
*,
run_manager: Any,
thread_id: str,
user_id: str | None,
run_ids: set[str],
) -> dict[str, int]:
"""Batch-hydrate the requested runs and compute their latest durations."""
if not run_ids:
return {}
from app.gateway.routers.thread_runs import compute_run_durations
runs = await run_manager.list_by_thread(
thread_id,
user_id=user_id,
limit=max(100, len(run_ids)),
)
known_run_ids = {run.run_id for run in runs}
for run_id in sorted(run_ids - known_run_ids):
run = await run_manager.get(run_id, user_id=user_id)
if run is not None:
runs.append(run)
known_run_ids.add(run_id)
computed_durations = compute_run_durations(runs)
return {run_id: duration for run_id, duration in computed_durations.items() if run_id in run_ids}
async def _persist_run_history_metadata_background(
*,
request: Request,
checkpointer: Any,
thread_id: str,
user_id: str | None,
duration_run_ids: set[str],
message_run_ids: dict[str, str],
audited_message_ids: set[str],
) -> None:
"""Best-effort history migration behind durable checkpoint admission."""
from deerflow.runtime.runs.worker import persist_run_history_metadata
try:
async with reserve_checkpoint_write(request, thread_id, user_id=user_id):
from app.gateway.deps import get_run_event_store, get_run_manager
authoritative_message_run_ids = dict(message_run_ids)
authoritative_duration_run_ids = set(duration_run_ids)
if audited_message_ids:
exact_after_admission = await get_run_event_store(request).find_latest_ai_message_run_ids(
thread_id,
audited_message_ids,
user_id=user_id,
)
for message_id in audited_message_ids:
exact_run_id = exact_after_admission.get(message_id)
if valid_run_message_id_entry(message_id, exact_run_id):
if authoritative_message_run_ids.get(message_id) != exact_run_id:
authoritative_duration_run_ids.add(exact_run_id)
authoritative_message_run_ids[message_id] = exact_run_id
authoritative_durations = await _load_run_durations(
run_manager=get_run_manager(request),
thread_id=thread_id,
user_id=user_id,
run_ids=authoritative_duration_run_ids,
)
await persist_run_history_metadata(
checkpointer=checkpointer,
thread_id=thread_id,
durations=authoritative_durations,
message_run_ids=authoritative_message_run_ids,
)
except ConflictError:
# A live run or another checkpoint writer owns the thread. The mapping
# is a read-through optimization, so the next history request can retry
# instead of racing a user-visible state mutation.
logger.debug("Skipped run-history metadata migration for busy thread %s", sanitize_log_param(thread_id))
except Exception:
logger.warning("Failed to persist run-history metadata for thread %s", sanitize_log_param(thread_id), exc_info=True)
@router.post("/{thread_id}/history", response_model=list[HistoryEntry])
@require_permission("threads", "read", owner_check=True)
async def get_thread_history(
@ -1408,8 +1498,10 @@ async def get_thread_history(
# carry the completed turns' durations in metadata, so the
# messages channel stays unchanged.
checkpoint_run_durations = _checkpoint_run_durations(metadata)
checkpoint_run_message_ids = _checkpoint_run_message_ids(metadata)
current_turn_run_id = None
turn_run_ids: set[str] = set()
legacy_ai_message_ids: set[str] = set()
for msg in serialized_msgs:
if msg.get("type") == "human":
additional_kwargs = msg.get("additional_kwargs")
@ -1419,75 +1511,121 @@ async def get_thread_history(
current_turn_run_id = run_id
continue
if msg.get("type") not in {"ai", "tool"} or not current_turn_run_id:
message_type = msg.get("type")
if message_type not in {"ai", "tool"}:
continue
msg.setdefault("run_id", current_turn_run_id)
if msg.get("type") == "ai":
turn_run_ids.add(current_turn_run_id)
if message_type == "ai":
message_id = msg.get("id")
persisted_run_id = checkpoint_run_message_ids.get(message_id) if isinstance(message_id, str) else None
if persisted_run_id:
msg["run_id"] = persisted_run_id
elif not isinstance(msg.get("run_id"), str) or not msg.get("run_id"):
if current_turn_run_id:
msg["run_id"] = current_turn_run_id
if isinstance(message_id, str) and message_id:
legacy_ai_message_ids.add(message_id)
# Stamp each run's duration on its last AI message only,
# same as the live message endpoints — never every AI
# message in a multi-message turn (#4152).
stamp_turn_duration_on_last_ai(serialized_msgs, checkpoint_run_durations)
run_id = msg.get("run_id")
if isinstance(run_id, str) and run_id:
turn_run_ids.add(run_id)
elif current_turn_run_id:
msg.setdefault("run_id", current_turn_run_id)
# Runs referenced by this checkpoint's AI messages but
# absent from checkpoint metadata are either legacy
# (never migrated) or just completed. Correlate once via
# event-store + run-manager, then upgrade by a
# metadata-only checkpoint write.
# absent from duration metadata are either legacy
# (never migrated) or just completed. Exact attribution
# has its own completeness condition: duration-only
# checkpoints written before #4949 still need their AI
# IDs correlated and persisted. Correlate once via the
# event store, then hydrate only the run rows whose
# durations are actually required.
resolved_run_durations = dict(checkpoint_run_durations)
missing_run_ids = turn_run_ids - set(checkpoint_run_durations)
if missing_run_ids:
if missing_run_ids or legacy_ai_message_ids:
from app.gateway.deps import get_run_event_store, get_run_manager
from app.gateway.routers.thread_runs import compute_run_durations
from deerflow.runtime.runs.worker import persist_run_durations
run_mgr = get_run_manager(request)
event_store = get_run_event_store(request)
runs = await run_mgr.list_by_thread(thread_id)
events = await event_store.list_messages(thread_id, limit=1000)
if runs:
run_durations = compute_run_durations(runs)
msg_to_run = {}
for event in events:
content = event.get("content", {})
run_id = event.get("run_id")
if isinstance(content, dict) and content.get("type") == "ai" and "id" in content and isinstance(run_id, str) and run_id:
msg_to_run[content["id"]] = run_id
current_turn_run_id = None
user_id = get_effective_user_id()
ai_message_ids = set(legacy_ai_message_ids)
try:
msg_to_run = (
await event_store.find_latest_ai_message_run_ids(
thread_id,
ai_message_ids,
user_id=user_id,
)
if ai_message_ids
else {}
)
except Exception:
# A failed exact lookup must not masquerade as a
# successful boundary attribution. Removing the
# synthesized ids leaves the response incomplete
# rather than deterministically wrong. Durations
# backed by persisted mappings remain provable
# and should still survive this degraded path.
for msg in serialized_msgs:
if msg.get("type") == "human":
additional_kwargs = msg.get("additional_kwargs")
if isinstance(additional_kwargs, dict):
run_id = additional_kwargs.get("run_id")
if isinstance(run_id, str) and run_id:
current_turn_run_id = run_id
continue
if msg.get("type") == "ai" and msg.get("id") in ai_message_ids:
msg.pop("run_id", None)
stamp_turn_duration_on_last_ai(
serialized_msgs,
checkpoint_run_durations,
)
raise
if msg.get("type") not in {"ai", "tool"}:
continue
run_id = msg_to_run.get(msg.get("id")) or current_turn_run_id
if run_id:
msg["run_id"] = run_id
for msg in serialized_msgs:
if msg.get("type") != "ai":
continue
exact_run_id = msg_to_run.get(msg.get("id"))
if exact_run_id:
msg["run_id"] = exact_run_id
stamp_turn_duration_on_last_ai(serialized_msgs, run_durations)
# Cache the complete audited attribution, including
# boundary fallbacks for IDs with no event. Without
# those negative-result entries, every history read
# would rescan the same pre-event-store prefix.
message_run_ids_to_persist = {
message_id: run_id
for msg in serialized_msgs
if msg.get("type") == "ai" and isinstance((message_id := msg.get("id")), str) and message_id in ai_message_ids and isinstance((run_id := msg.get("run_id")), str) and run_id
}
required_run_ids = {run_id for msg in serialized_msgs if msg.get("type") == "ai" and isinstance((run_id := msg.get("run_id")), str) and run_id and run_id not in checkpoint_run_durations}
run_durations = await _load_run_durations(
run_manager=run_mgr,
thread_id=thread_id,
user_id=user_id,
run_ids=required_run_ids,
)
resolved_run_durations.update(run_durations)
# Intentional, best-effort write-on-read migration:
# persist legacy metadata after the response so the
# history request never waits on an active stream's
# same-thread checkpoint lock.
# Intentional, best-effort write-on-read migration:
# persist both exact attribution and duration after
# the response so subsequent reads stay exact without
# waiting on an active stream's checkpoint lock.
if required_run_ids or message_run_ids_to_persist:
background_tasks.add_task(
persist_run_durations,
_persist_run_history_metadata_background,
request=request,
checkpointer=checkpointer,
thread_id=thread_id,
durations=run_durations,
user_id=user_id,
duration_run_ids=required_run_ids,
message_run_ids=message_run_ids_to_persist,
audited_message_ids=ai_message_ids,
)
# Stamp only after exact attribution is final. Stamping
# the synthesized boundary first can leave its duration
# attached to a message whose run ID is later corrected.
stamp_turn_duration_on_last_ai(
serialized_msgs,
resolved_run_durations,
)
except Exception:
logger.warning("Failed to inject turn_duration for thread %s", thread_id, exc_info=True)
logger.warning("Failed to inject turn_duration for thread %s", sanitize_log_param(thread_id), exc_info=True)
values["messages"] = serialized_msgs
@ -1496,7 +1634,7 @@ async def get_thread_history(
next_tasks = list(snapshot.next or ())
# Strip LangGraph internal keys from metadata
user_meta = {k: v for k, v in metadata.items() if k not in ("created_at", "updated_at", "step", "source", "writes", "parents", "run_durations")}
user_meta = {k: v for k, v in metadata.items() if k not in ("created_at", "updated_at", "step", "source", "writes", "parents", "run_durations", RUN_MESSAGE_IDS_METADATA_KEY)}
# Keep step for ordering context
if "step" in metadata:
user_meta["step"] = metadata["step"]

View File

@ -20,6 +20,39 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. Edit replay runs (`metadata.replay_kind="edit"`) also restore the pre-run checkpoint on failed, timed-out, or interrupted completion and publish the restored `values` snapshot to the stream before `end`, so clients do not remain on a transient edited branch when the replay did not produce a successful replacement.
**Targeted run-event attribution** (`runtime/events/store/`):
`RunEventStore.find_latest_ai_message_run_ids()` has a complete-or-error
contract. Its default implementation walks `list_messages()` backward in
1000-row pages, preserves the first page's high-watermark through the exclusive
`before_seq` cursor, and raises when a full page has no safe progressing `seq`.
Memory and database stores use that bounded path; the JSONL store overrides it
with one complete thread-log read because each JSONL page would otherwise
rescan every run file. The default and JSONL paths share the public
`normalize_message_ids()` and `match_ai_message_run_id()` helpers from
`events/store/base.py`. Database owner filtering is inherited on every page.
Callers may use a missing key as proof that no valid AI event exists only after
an ordinary return, never after an exception. A caller that crosses a run or
checkpoint-write admission boundary must repeat the complete audit after
admission; a pre-admission exact hit can be superseded by a later event just as
a pre-admission miss can become an exact hit.
Gateway `POST /api/threads/{id}/history` uses that lookup to migrate legacy AI
messages. An exhaustive miss preserves the human-boundary fallback; an
incomplete lookup removes unproven synthesized IDs. Its metadata-only
write-on-read cache stores `run_message_ids` for every audited AI ID (including
exhaustive misses) plus required `run_durations`; duration presence alone does
not prove attribution. Historical `body.before` reads write the audit to the
head, and the merge may retain IDs no longer in materialized history, which
readers ignore. Migration must acquire the durable `checkpoint_write`
reservation, then repeat the whole message audit and batch-reload required run
rows before persisting. Post-admission exact hits replace foreground exact or
boundary mappings, and recomputed final durations replace foreground snapshots.
Successful workers keep their durable run row active through the final duration
checkpoint write, so a peer migration cannot enter during terminalization.
The first `RunManager.list_by_thread()` hydration page uses a 100-row floor or
the number of required IDs, whichever is larger; missing exact runs use targeted
`get()` calls.
**Where things live**:
- `runtime/checkpoint_mode.py` — mode + snapshot-frequency freeze, marker injection, delta detection, compatibility gate, both error types
- `runtime/checkpoint_state.py``CheckpointStateAccessor`, `build_state_mutation_graph`, `RollbackPoint`

View File

@ -16,6 +16,31 @@ import abc
from deerflow.runtime.user_context import AUTO, _AutoSentinel
_AI_MESSAGE_RUN_LOOKUP_PAGE_SIZE = 1000
class IncompleteMessageRunLookupError(RuntimeError):
"""Raised when a store cannot prove that a targeted lookup is complete."""
def normalize_message_ids(message_ids: set[str]) -> set[str]:
"""Return the non-empty string IDs that can participate in a lookup."""
return {message_id for message_id in message_ids if isinstance(message_id, str) and message_id}
def match_ai_message_run_id(event: object, message_ids: set[str]) -> tuple[str, str] | None:
"""Return a target AI message ID and its valid run ID, if present."""
if not isinstance(event, dict) or event.get("category") != "message":
return None
content = event.get("content")
run_id = event.get("run_id")
if not isinstance(content, dict) or content.get("type") != "ai" or not isinstance(run_id, str) or not run_id:
return None
message_id = content.get("id")
if not isinstance(message_id, str) or message_id not in message_ids:
return None
return message_id, run_id
class RunEventStore(abc.ABC):
"""Run event stream storage interface.
@ -27,6 +52,8 @@ class RunEventStore(abc.ABC):
4. list_events() returns all events for the specified run
5. Returned dicts contain the required RunEvent envelope fields; backends
may add documented fields such as DbRunEventStore.user_id
6. find_latest_ai_message_run_ids() returns the newest valid AI message
event for each requested ID and performs no storage work for empty input
"""
@abc.abstractmethod
@ -92,6 +119,70 @@ class RunEventStore(abc.ABC):
user-scoped backends must apply it according to their isolation model.
"""
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]:
"""Map target message IDs to their newest valid AI event's run ID.
Only ``category="message"`` events whose structured content has
``type="ai"`` and whose ``run_id`` is a non-empty string qualify. An
empty target set must return immediately without storage work. The
default implementation pages backward in bounded windows. It raises
:class:`IncompleteMessageRunLookupError` instead of returning a
partial result when a full page lacks a safe, progressing ``seq``
cursor; callers may only treat an ordinary return as an exhaustive
lookup for unresolved IDs.
``user_id`` follows the same explicit-caller semantics as
:meth:`list_messages`.
"""
pending = normalize_message_ids(message_ids)
if not pending:
return {}
result: dict[str, str] = {}
before_seq: int | None = None
while pending:
page = await self.list_messages(
thread_id,
limit=_AI_MESSAGE_RUN_LOOKUP_PAGE_SIZE,
before_seq=before_seq,
user_id=user_id,
)
if not page:
break
for event in reversed(page):
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
if not pending or len(page) < _AI_MESSAGE_RUN_LOOKUP_PAGE_SIZE:
break
seqs: list[int] = []
for event in page:
seq = event.get("seq") if isinstance(event, dict) else None
if not isinstance(seq, int) or isinstance(seq, bool):
raise IncompleteMessageRunLookupError("Run event lookup could not form a safe backward cursor from a full page")
seqs.append(seq)
next_before_seq = min(seqs)
if before_seq is not None and next_before_seq >= before_seq:
raise IncompleteMessageRunLookupError("Run event lookup could not form a safe backward cursor because seq did not progress")
before_seq = next_before_seq
return result
@abc.abstractmethod
async def list_events(
self,

View File

@ -30,7 +30,7 @@ from datetime import UTC, datetime
from pathlib import Path
from typing import Any
from deerflow.runtime.events.store.base import RunEventStore
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
@ -280,6 +280,34 @@ class JsonlRunEventStore(RunEventStore):
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):
events = await asyncio.to_thread(self._read_run_events, thread_id, run_id)
if event_types is not None:

View File

@ -549,6 +549,26 @@ class RunManager:
except Exception:
logger.warning("Failed to persist run progress for %s", run_id, exc_info=True)
async def update_finalizing_progress(self, run_id: str, **kwargs) -> None:
"""Persist final fields while the durable row is deliberately active."""
should_persist = False
async with self._lock:
record = self._runs.get(run_id)
if record is not None and not record.ownership_lost:
should_persist = record.status not in (RunStatus.pending, RunStatus.running)
if should_persist:
for key, value in kwargs.items():
if hasattr(record, key) and value is not None:
setattr(record, key, value)
record.updated_at = _now_iso()
if should_persist and self._store is not None:
try:
# The local status is already staged as terminal, but the store
# row intentionally remains running until checkpoint finalization.
await self._store.update_run_progress(run_id, **kwargs)
except Exception:
logger.warning("Failed to persist finalizing progress for %s", run_id, exc_info=True)
async def create(
self,
thread_id: str,

View File

@ -626,6 +626,7 @@ async def run_agent(
# checkpoint failures / cancellation while waiting did not write an empty
# completion snapshot into RunStore.
persist_completion = False
completion_data: dict[str, Any] | None = None
# Buffers subagent step events for batched persistence (#3779); assigned once
# streaming starts and flushed in the finally block. Pre-bound to None so the
# finally is safe even if an exception fires before streaming begins.
@ -1226,6 +1227,37 @@ async def run_agent(
persist=False,
)
if not record.ownership_lost and journal is not None and persist_completion:
try:
# Advance the final completion fields and timestamp without
# terminalizing the durable row. That active row continues to
# fence peer checkpoint writers through the duration write.
completion_data = journal.get_completion_data()
await run_manager.update_finalizing_progress(run_id, **completion_data)
except Exception:
logger.warning("Failed to persist finalizing run progress for %s (non-fatal)", run_id, exc_info=True)
# Keep the durable run row active through its final duration checkpoint
# write. A peer Gateway admits history migration from the durable row,
# not this worker's staged terminal status; terminalizing first would
# let that migration read an unfinished lifetime and race this write.
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
try:
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
# Match legacy history semantics: turn_duration is the whole
# RunRecord lifetime in integer seconds, including admission
# delay. Persist zero for sub-second successful turns.
duration = max(0, int((updated - created).total_seconds()))
await _persist_run_duration(
checkpointer=checkpointer,
thread_id=thread_id,
run_id=run_id,
duration_seconds=duration,
)
except Exception:
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
if not record.ownership_lost and event_store is not None:
try:
# Even after bounded receipt retries are exhausted, persist the
@ -1250,8 +1282,8 @@ async def run_agent(
if not record.ownership_lost and journal is not None and persist_completion:
try:
# Persist token usage + convenience fields to RunStore
completion = journal.get_completion_data()
await run_manager.update_run_completion(run_id, status=record.status.value, **completion)
completion_data = completion_data or journal.get_completion_data()
await run_manager.update_run_completion(run_id, status=record.status.value, **completion_data)
except Exception:
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
@ -1276,25 +1308,6 @@ async def run_agent(
except Exception:
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
# Persist run duration to checkpoint metadata so history reads
# don't need to correlate runs and events.
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
try:
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
# Match legacy history semantics: turn_duration is the whole
# RunRecord lifetime in integer seconds, including admission
# delay. Persist zero for sub-second successful turns.
duration = max(0, int((updated - created).total_seconds()))
await _persist_run_duration(
checkpointer=checkpointer,
thread_id=thread_id,
run_id=run_id,
duration_seconds=duration,
)
except Exception:
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
# Update threads_meta status based on run outcome
if started and not record.ownership_lost and thread_store is not None:
try:
@ -2085,21 +2098,36 @@ def valid_duration_entry(run_id: Any, duration_seconds: Any) -> bool:
return isinstance(run_id, str) and bool(run_id) and isinstance(duration_seconds, int) and not isinstance(duration_seconds, bool)
async def persist_run_durations(
RUN_MESSAGE_IDS_METADATA_KEY = "run_message_ids"
def valid_run_message_id_entry(message_id: Any, run_id: Any) -> bool:
"""Check that a persisted legacy message-to-run attribution is well formed."""
return isinstance(message_id, str) and bool(message_id) and isinstance(run_id, str) and bool(run_id)
async def persist_run_history_metadata(
*,
checkpointer: Any,
thread_id: str,
durations: dict[str, int],
durations: dict[str, int] | None = None,
message_run_ids: dict[str, str] | None = None,
) -> bool:
"""Merge validated run durations into a metadata-only checkpoint.
"""Merge validated run history indexes into a metadata-only checkpoint.
Durations accumulate so the history fast path can serve every known turn
from the latest checkpoint. Per-entry overhead is negligible (~50 bytes
per run_id) compared to the messages channel blob written on every graph
checkpoint, so no pruning is needed.
from the latest checkpoint. Legacy AI-message attributions are persisted
alongside them for every audited AI ID, including boundary fallbacks whose
event lookup was exhaustively empty. The full mapping is deliberate: it is
both the exact-attribution cache and the negative-result coverage proof.
While the materialized message set at the head remains unchanged, later
reads query only uncached IDs. This metadata-only merge retains existing
entries, so compaction timing or historical migration may leave stale IDs;
reads ignore them because they only consult IDs in the materialized history.
"""
updates = {run_id: max(0, duration_seconds) for run_id, duration_seconds in durations.items() if valid_duration_entry(run_id, duration_seconds)}
if not updates:
duration_updates = {run_id: max(0, duration_seconds) for run_id, duration_seconds in (durations or {}).items() if valid_duration_entry(run_id, duration_seconds)}
message_run_id_updates = {message_id: run_id for message_id, run_id in (message_run_ids or {}).items() if valid_run_message_id_entry(message_id, run_id)}
if not duration_updates and not message_run_id_updates:
return False
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
@ -2113,11 +2141,15 @@ async def persist_run_durations(
metadata = dict(getattr(ckpt_tuple, "metadata", {}) or {})
raw_run_durations = metadata.get("run_durations")
run_durations = {key: value for key, value in raw_run_durations.items() if valid_duration_entry(key, value)} if isinstance(raw_run_durations, dict) else {}
changed_durations = {run_id: duration for run_id, duration in updates.items() if run_durations.get(run_id) != duration}
if not changed_durations:
raw_message_run_ids = metadata.get(RUN_MESSAGE_IDS_METADATA_KEY)
run_message_ids = {message_id: run_id for message_id, run_id in raw_message_run_ids.items() if valid_run_message_id_entry(message_id, run_id)} if isinstance(raw_message_run_ids, dict) else {}
changed_durations = {run_id: duration for run_id, duration in duration_updates.items() if run_durations.get(run_id) != duration}
changed_message_run_ids = {message_id: run_id for message_id, run_id in message_run_id_updates.items() if run_message_ids.get(message_id) != run_id}
if not changed_durations and not changed_message_run_ids:
return False
run_durations.update(changed_durations)
run_message_ids.update(changed_message_run_ids)
parent_checkpoint_id = _checkpoint_identity(ckpt_tuple, checkpoint)
latest_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", ckpt_config)
latest_checkpoint = dict(getattr(latest_tuple, "checkpoint", {}) or {}) if latest_tuple is not None else {}
@ -2129,7 +2161,16 @@ async def persist_run_durations(
prev_step = metadata.get("step")
metadata["step"] = (prev_step + 1) if isinstance(prev_step, int) else 1
metadata["run_durations"] = run_durations
metadata["writes"] = {"runtime_run_duration": {"run_ids": sorted(changed_durations)}}
if run_message_ids:
metadata[RUN_MESSAGE_IDS_METADATA_KEY] = run_message_ids
else:
metadata.pop(RUN_MESSAGE_IDS_METADATA_KEY, None)
metadata["writes"] = {
"runtime_run_duration": {
"run_ids": sorted(changed_durations),
"message_ids": sorted(changed_message_run_ids),
}
}
checkpoint_ns = _checkpoint_namespace(ckpt_tuple)
write_config = {
@ -2152,6 +2193,20 @@ async def persist_run_durations(
return False
async def persist_run_durations(
*,
checkpointer: Any,
thread_id: str,
durations: dict[str, int],
) -> bool:
"""Merge validated run durations into a metadata-only checkpoint."""
return await persist_run_history_metadata(
checkpointer=checkpointer,
thread_id=thread_id,
durations=durations,
)
async def _persist_run_duration(
*,
checkpointer: Any,

View File

@ -9,9 +9,12 @@ from langgraph.checkpoint.base import empty_checkpoint, uuid6
from langgraph.checkpoint.memory import InMemorySaver
import deerflow.runtime.runs.worker as worker
from deerflow.runtime import ConflictError, ThreadOperationKind
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.goal import goal_thread_lock
from deerflow.runtime.runs.manager import RunManager, RunStartOutcome
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.runs.store.memory import MemoryRunStore
from deerflow.runtime.runs.worker import RunContext, _persist_run_duration, run_agent
@ -312,6 +315,75 @@ async def test_agent_stream_serializes_with_duration_checkpoint_write() -> None:
assert finished_during_stream is False
@pytest.mark.anyio
async def test_successful_run_stays_durably_active_through_final_duration_write(monkeypatch: pytest.MonkeyPatch) -> None:
"""A peer migration cannot enter before the terminal duration is stored."""
run_store = MemoryRunStore()
owner = RunManager(store=run_store, worker_id="duration-owner")
peer = RunManager(store=run_store, worker_id="duration-peer")
record = await owner.create_or_reject("duration-finalization-admission")
checkpointer = InMemorySaver()
await _put_checkpoint(
checkpointer,
thread_id=record.thread_id,
checkpoint_id="00000000-0000-6000-8000-000000000001",
messages=[
HumanMessage(
id="human-1",
content="Question",
additional_kwargs={"run_id": record.run_id},
),
AIMessage(id="ai-1", content="Answer"),
],
step=1,
)
observed_store_statuses: list[str] = []
persist_duration = worker._persist_run_duration
async def assert_active_then_persist(**kwargs) -> None:
stored = await run_store.get(record.run_id)
assert stored is not None
observed_store_statuses.append(stored["status"])
with pytest.raises(ConflictError):
async with peer.reserve_thread_operation(
record.thread_id,
kind=ThreadOperationKind.checkpoint_write,
):
pass
await persist_duration(**kwargs)
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
yield {"messages": []}
monkeypatch.setattr(worker, "_persist_run_duration", assert_active_then_persist)
await run_agent(
SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
),
owner,
record,
ctx=RunContext(
checkpointer=checkpointer,
event_store=MemoryRunEventStore(),
),
agent_factory=lambda *, config: DummyAgent(),
graph_input={},
config={},
)
assert observed_store_statuses == [RunStatus.running.value]
stored = await run_store.get(record.run_id)
assert stored is not None
assert stored["status"] == RunStatus.success.value
latest = await checkpointer.aget_tuple({"configurable": {"thread_id": record.thread_id, "checkpoint_ns": ""}})
assert latest is not None
assert record.run_id in latest.metadata["run_durations"]
@pytest.mark.anyio
async def test_agent_stream_allows_graph_goal_state_access() -> None:
"""A graph node may acquire the goal lock while a run is streaming."""

View File

@ -4,6 +4,8 @@ Uses a helper to create the store for each backend type.
Memory tests run directly; DB and JSONL tests create stores inside each test.
"""
from unittest.mock import AsyncMock, patch
import pytest
from deerflow.runtime.events.store.memory import MemoryRunEventStore
@ -14,6 +16,73 @@ def store():
return MemoryRunEventStore()
async def _assert_find_latest_ai_message_run_ids_contract(store, *, allow_empty_run_id: bool) -> None:
assert await store.find_latest_ai_message_run_ids("t1", set()) == {}
await store.put(
thread_id="t1",
run_id="trace-run",
event_type="llm.ai.response",
category="trace",
content={"type": "ai", "id": "target"},
)
await store.put(
thread_id="t1",
run_id="string-content-run",
event_type="llm.ai.response",
category="message",
content="target",
)
await store.put(
thread_id="t1",
run_id="human-run",
event_type="human_message",
category="message",
content={"type": "human", "id": "target"},
)
await store.put(
thread_id="t1",
run_id="old-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
await store.put(
thread_id="t1",
run_id="other-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "other"},
)
await store.put(
thread_id="t1",
run_id="decoy-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "decoy", "note": "target"},
)
await store.put(
thread_id="t1",
run_id="new-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
if allow_empty_run_id:
await store.put(
thread_id="t1",
run_id="",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
assert await store.find_latest_ai_message_run_ids("t1", {"target", "other", "missing"}, user_id=None) == {
"target": "new-run",
"other": "other-run",
}
# -- Basic write and query --
@ -165,6 +234,121 @@ class TestListMessages:
assert [m["seq"] for m in await store.list_messages("t1", after_seq=5, limit=5)] == [7, 9]
class TestFindLatestAiMessageRunIds:
@pytest.mark.anyio
async def test_memory_contract(self, store):
await _assert_find_latest_ai_message_run_ids_contract(store, allow_empty_run_id=True)
@pytest.mark.anyio
async def test_memory_stops_after_all_targets_are_found(self, store):
from deerflow.runtime.events.store import base as event_store_base
await store.put(
thread_id="t1",
run_id="old-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "old"},
)
await store.put(
thread_id="t1",
run_id="new-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
store.list_messages = AsyncMock(wraps=store.list_messages)
with patch.object(event_store_base, "match_ai_message_run_id", wraps=event_store_base.match_ai_message_run_id) as match_event:
assert await store.find_latest_ai_message_run_ids("t1", {"target"}, user_id=None) == {"target": "new-run"}
assert match_event.call_count == 1
store.list_messages.assert_awaited_once_with("t1", limit=1000, before_seq=None, user_id=None)
@pytest.mark.anyio
async def test_memory_pages_in_bounded_windows_and_keeps_initial_high_watermark(self, store):
await store.put(
thread_id="t1",
run_id="old-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
await store.put_batch(
[
{
"thread_id": "t1",
"run_id": "noise-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": f"noise-{index}"},
}
for index in range(1000)
]
)
original_list_messages = store.list_messages
calls: list[dict] = []
async def list_messages(*args, **kwargs):
page = await original_list_messages(*args, **kwargs)
calls.append(kwargs)
if len(calls) == 1:
# This duplicate is newer than the first page's snapshot. A
# backward cursor must not let it replace the older answer
# while resolving the rest of that same lookup.
await store.put(
thread_id="t1",
run_id="concurrent-new-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
return page
store.list_messages = AsyncMock(side_effect=list_messages)
assert await store.find_latest_ai_message_run_ids("t1", {"target"}, user_id=None) == {"target": "old-run"}
assert len(calls) == 2
assert all(call["limit"] == 1000 for call in calls)
assert calls[0].get("before_seq") is None
assert calls[1]["before_seq"] == 2
@pytest.mark.anyio
@pytest.mark.parametrize("malformed_page", ["missing-seq", "non-progressing-seq"])
async def test_default_lookup_raises_instead_of_looping_on_unsafe_cursor(self, store, malformed_page):
from deerflow.runtime.events.store.base import RunEventStore
calls = 0
async def list_messages(*_args, **_kwargs):
nonlocal calls
calls += 1
if malformed_page == "missing-seq":
return [
{
"category": "message",
"content": {"type": "ai", "id": f"noise-{index}"},
"run_id": "noise-run",
}
for index in range(1000)
]
return [
{
"category": "message",
"content": {"type": "ai", "id": f"noise-{index}"},
"run_id": "noise-run",
"seq": index + 1,
}
for index in range(1000)
]
store.list_messages = AsyncMock(side_effect=list_messages)
with pytest.raises(RuntimeError, match="safe backward cursor"):
await RunEventStore.find_latest_ai_message_run_ids(store, "t1", {"missing"}, user_id=None)
assert calls == (1 if malformed_page == "missing-seq" else 2)
# -- list_events --
@ -365,6 +549,154 @@ class TestDbRunEventStore:
await close_engine()
@pytest.mark.anyio
async def test_find_latest_ai_message_run_ids_contract_and_owner_filter(self, tmp_path):
from types import SimpleNamespace
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.runtime.events.store.db import DbRunEventStore
from deerflow.runtime.user_context import reset_current_user, set_current_user
url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}"
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
try:
store = DbRunEventStore(get_session_factory())
await _assert_find_latest_ai_message_run_ids_contract(store, allow_empty_run_id=True)
owner_a_token = set_current_user(SimpleNamespace(id="owner-a"))
try:
await store.put(
thread_id="owned-thread",
run_id="owner-a-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "shared-id"},
)
finally:
reset_current_user(owner_a_token)
owner_b_token = set_current_user(SimpleNamespace(id="owner-b"))
try:
await store.put(
thread_id="owned-thread",
run_id="owner-b-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "shared-id"},
)
finally:
reset_current_user(owner_b_token)
assert await store.find_latest_ai_message_run_ids("owned-thread", {"shared-id"}, user_id="owner-a") == {"shared-id": "owner-a-run"}
assert await store.find_latest_ai_message_run_ids("owned-thread", {"shared-id"}, user_id="owner-b") == {"shared-id": "owner-b-run"}
finally:
await close_engine()
@pytest.mark.anyio
async def test_find_latest_ai_message_run_ids_handles_large_target_sets_and_special_ids(self, tmp_path):
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.runtime.events.store.db import DbRunEventStore
url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}"
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
try:
store = DbRunEventStore(get_session_factory())
target_ids = {f"id-{index:03d}" for index in range(201)}
special_id = 'message-%_/"-雪'
target_ids.add(special_id)
await store.put_batch(
[
{
"thread_id": "t1",
"run_id": "first-chunk-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": "id-000"},
},
{
"thread_id": "t1",
"run_id": "last-chunk-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": "id-200"},
},
{
"thread_id": "t1",
"run_id": "special-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": special_id},
},
]
)
assert await store.find_latest_ai_message_run_ids("t1", target_ids, user_id=None) == {
"id-000": "first-chunk-run",
"id-200": "last-chunk-run",
special_id: "special-run",
}
finally:
await close_engine()
@pytest.mark.anyio
async def test_find_latest_ai_message_run_ids_pages_db_with_owner_scoped_high_watermark(self, tmp_path):
from types import SimpleNamespace
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.runtime.events.store.db import DbRunEventStore
from deerflow.runtime.user_context import reset_current_user, set_current_user
url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}"
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
owner_token = set_current_user(SimpleNamespace(id="owner-a"))
try:
store = DbRunEventStore(get_session_factory())
await store.put(
thread_id="t1",
run_id="old-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
await store.put_batch(
[
{
"thread_id": "t1",
"run_id": "noise-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": f"noise-{index}"},
}
for index in range(1000)
]
)
original_list_messages = store.list_messages
calls: list[dict] = []
async def list_messages(*args, **kwargs):
page = await original_list_messages(*args, **kwargs)
calls.append(kwargs)
if len(calls) == 1:
await store.put(
thread_id="t1",
run_id="concurrent-new-run",
event_type="llm.ai.response",
category="message",
content={"type": "ai", "id": "target"},
)
return page
store.list_messages = AsyncMock(side_effect=list_messages)
assert await store.find_latest_ai_message_run_ids("t1", {"target"}, user_id="owner-a") == {"target": "old-run"}
assert len(calls) == 2
assert all(call["limit"] == 1000 and call["user_id"] == "owner-a" for call in calls)
assert calls[0].get("before_seq") is None
assert calls[1]["before_seq"] == 2
finally:
reset_current_user(owner_token)
await close_engine()
@pytest.mark.anyio
async def test_put_if_absent_is_idempotent(self, tmp_path):
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
@ -749,6 +1081,44 @@ class TestJsonlRunEventStore:
messages = await s.list_messages("t1")
assert len(messages) == 1
@pytest.mark.anyio
async def test_find_latest_ai_message_run_ids_contract(self, tmp_path):
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
store = JsonlRunEventStore(base_dir=tmp_path / "jsonl")
await _assert_find_latest_ai_message_run_ids_contract(store, allow_empty_run_id=False)
@pytest.mark.anyio
async def test_find_latest_ai_message_run_ids_reads_thread_once_and_ignores_empty_run(self, tmp_path):
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore
store = JsonlRunEventStore(base_dir=tmp_path / "jsonl")
events = [
{
"thread_id": "t1",
"run_id": "valid-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": "target"},
"seq": 1,
},
{
"thread_id": "t1",
"run_id": "",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": "target"},
"seq": 2,
},
]
with patch.object(store, "_read_thread_events", return_value=events) as read_thread_events:
assert await store.find_latest_ai_message_run_ids("t1", {"target"}, user_id=None) == {"target": "valid-run"}
read_thread_events.assert_called_once_with("t1")
with patch.object(store, "_read_thread_events", side_effect=AssertionError("empty input must not read")) as read_thread_events:
assert await store.find_latest_ai_message_run_ids("t1", set()) == {}
read_thread_events.assert_not_called()
@pytest.mark.anyio
async def test_put_if_absent_is_idempotent(self, tmp_path):
from deerflow.runtime.events.store.jsonl import JsonlRunEventStore

View File

@ -1355,6 +1355,7 @@ def test_get_thread_history_returns_iso_for_legacy_checkpoint_metadata() -> None
def test_get_thread_history_associates_tool_messages_from_checkpoint_turn() -> None:
app, _store, checkpointer = _build_thread_app()
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=AsyncMock(return_value={}))
thread_id = "history-tool-run"
messages = [
HumanMessage(id="human-1", content="Use a tool", additional_kwargs={"run_id": "run-1"}),
@ -1414,11 +1415,14 @@ def test_get_thread_history_fast_path_skips_runs_already_in_checkpoint_metadata(
"checkpoint-partial",
messages,
step=1,
metadata={"run_durations": {"run-migrated": 4}},
metadata={
"run_durations": {"run-migrated": 4},
"run_message_ids": {"ai-1": "run-migrated"},
},
)
)
async def list_by_thread(_: str) -> list[SimpleNamespace]:
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
return [
SimpleNamespace(
run_id="run-pending",
@ -1427,14 +1431,19 @@ def test_get_thread_history_fast_path_skips_runs_already_in_checkpoint_metadata(
),
]
list_messages_calls: list[str] = []
lookup_calls: list[set[str]] = []
async def list_messages(thread: str, *, limit: int) -> list[dict]:
list_messages_calls.append(thread)
return []
async def find_latest_ai_message_run_ids(thread: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
assert thread == thread_id
assert message_ids == {"ai-2"}
lookup_calls.append(message_ids)
return {}
app.state.run_manager = SimpleNamespace(list_by_thread=list_by_thread)
app.state.run_event_store = SimpleNamespace(list_messages=list_messages)
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
@ -1443,9 +1452,475 @@ def test_get_thread_history_fast_path_skips_runs_already_in_checkpoint_metadata(
history_messages = response.json()[0]["values"]["messages"]
assert history_messages[1]["additional_kwargs"]["turn_duration"] == 4
assert history_messages[3]["additional_kwargs"]["turn_duration"] == 6
# The fallback still runs (run-pending was missing), but it is the only
# reason it ran — proven by it firing exactly once, not skipped entirely.
assert list_messages_calls == [thread_id]
# The missing ID is checked once for the response and once after write
# admission; the already migrated ID is absent from both lookups.
assert lookup_calls == [{"ai-2"}, {"ai-2"}]
def test_get_thread_history_backfills_exact_mapping_when_durations_already_exist() -> None:
"""Duration metadata alone does not prove exact message attribution.
A pre-#4949 checkpoint can already carry every run duration while lacking
``run_message_ids``. The history read must still consult the event index;
otherwise the synthesized human-boundary run becomes permanent.
"""
app, _store, checkpointer = _build_thread_app()
thread_id = "history-duration-without-attribution"
messages = [
HumanMessage(id="human-1", content="Question", additional_kwargs={"run_id": "boundary-run"}),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(
_write_checkpoint(
checkpointer,
thread_id,
"00000000-0000-6000-8000-000000000010",
messages,
step=1,
metadata={"run_durations": {"boundary-run": 3, "exact-run": 7}},
)
)
lookup_calls: list[set[str]] = []
async def find_latest_ai_message_run_ids(_: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
lookup_calls.append(message_ids)
return {"ai-1": "exact-run"}
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
return []
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
ai_message = response.json()[0]["values"]["messages"][1]
assert ai_message["run_id"] == "exact-run"
assert ai_message["additional_kwargs"]["turn_duration"] == 7
assert lookup_calls == [{"ai-1"}, {"ai-1"}]
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_message_ids"] == {"ai-1": "exact-run"}
def test_get_thread_history_preserves_boundary_fallback_after_complete_partial_lookup() -> None:
"""A complete lookup may legitimately find no event for old messages.
Pre-event-store checkpoints still rely on the human turn boundary. A
partial result therefore corrects the IDs it can prove and preserves that
compatibility fallback for IDs that are definitively absent.
"""
app, _store, checkpointer = _build_thread_app()
thread_id = "history-partial-exact-attribution"
messages = [
HumanMessage(id="human-1", content="First", additional_kwargs={"run_id": "boundary-1"}),
AIMessage(id="ai-1", content="First answer"),
HumanMessage(id="human-2", content="Second", additional_kwargs={"run_id": "boundary-2"}),
AIMessage(id="ai-2", content="Second answer"),
]
asyncio.run(_write_checkpoint(checkpointer, thread_id, "00000000-0000-6000-8000-000000000011", messages, step=1))
lookup_calls: list[set[str]] = []
async def find_latest_ai_message_run_ids(_: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
lookup_calls.append(message_ids)
if message_ids == {"ai-1", "ai-2"}:
return {"ai-1": "exact-1"}
assert message_ids == {"ai-2"}
return {}
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
return [
SimpleNamespace(
run_id="exact-1",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:07+00:00",
),
SimpleNamespace(
run_id="boundary-2",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:03+00:00",
),
]
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
first_ai, second_ai = [message for message in response.json()[0]["values"]["messages"] if message["type"] == "ai"]
assert first_ai["run_id"] == "exact-1"
assert first_ai["additional_kwargs"]["turn_duration"] == 7
assert second_ai["run_id"] == "boundary-2"
assert second_ai["additional_kwargs"]["turn_duration"] == 3
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_message_ids"] == {"ai-1": "exact-1", "ai-2": "boundary-2"}
assert latest.metadata["run_durations"] == {"boundary-2": 3, "exact-1": 7}
assert lookup_calls == [{"ai-1", "ai-2"}, {"ai-1", "ai-2"}]
def test_get_thread_history_removes_synthesized_boundary_when_exact_lookup_is_incomplete() -> None:
"""Unsafe pagination removes only attribution it cannot prove."""
from deerflow.runtime.events.store.base import IncompleteMessageRunLookupError
app, _store, checkpointer = _build_thread_app()
thread_id = "history-incomplete-exact-attribution"
messages = [
HumanMessage(id="human-1", content="Proven question", additional_kwargs={"run_id": "proven-run"}),
AIMessage(id="ai-1", content="Proven answer"),
HumanMessage(id="human-2", content="Legacy question", additional_kwargs={"run_id": "boundary-run"}),
AIMessage(id="ai-2", content="Legacy answer"),
]
asyncio.run(
_write_checkpoint(
checkpointer,
thread_id,
"00000000-0000-6000-8000-000000000012",
messages,
step=1,
metadata={
"run_durations": {"proven-run": 4},
"run_message_ids": {"ai-1": "proven-run"},
},
)
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=AsyncMock(side_effect=IncompleteMessageRunLookupError("Run event lookup could not form a safe backward cursor")))
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
first_ai, second_ai = [message for message in response.json()[0]["values"]["messages"] if message["type"] == "ai"]
assert first_ai["run_id"] == "proven-run"
assert first_ai["additional_kwargs"]["turn_duration"] == 4
assert "run_id" not in second_ai
assert "turn_duration" not in (second_ai.get("additional_kwargs") or {})
assert app.state.run_manager.reservations == []
def test_get_thread_history_caches_complete_boundary_attribution() -> None:
"""A complete audit, including a negative event result, is a one-time scan."""
app, _store, checkpointer = _build_thread_app()
thread_id = "history-sparse-exact-attribution"
messages = [
HumanMessage(id="human-1", content="Question", additional_kwargs={"run_id": "boundary-run"}),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(_write_checkpoint(checkpointer, thread_id, "00000000-0000-6000-8000-000000000013", messages, step=1))
lookup_calls: list[set[str]] = []
async def find_latest_ai_message_run_ids(_: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
assert message_ids == {"ai-1"}
lookup_calls.append(message_ids)
return {}
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
return [
SimpleNamespace(
run_id="boundary-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:03+00:00",
)
]
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
second_response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
ai_message = response.json()[0]["values"]["messages"][1]
assert ai_message["run_id"] == "boundary-run"
assert ai_message["additional_kwargs"]["turn_duration"] == 3
assert second_response.status_code == 200, second_response.text
assert lookup_calls == [{"ai-1"}, {"ai-1"}]
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_durations"] == {"boundary-run": 3}
assert latest.metadata["run_message_ids"] == {"ai-1": "boundary-run"}
def test_get_thread_history_revalidates_boundary_fallback_after_reservation() -> None:
"""A run may flush its exact event before the metadata task is admitted.
The foreground lookup can exhaust the event log while the run's journal is
still buffered. If the background task acquires its checkpoint reservation
only after that run flushes and releases the thread, persisting the earlier
human-boundary fallback would make the temporary miss permanent.
"""
app, _store, checkpointer = _build_thread_app()
thread_id = "history-fallback-reservation-race"
messages = [
HumanMessage(
id="human-1",
content="Question",
additional_kwargs={"run_id": "boundary-run"},
),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(
_write_checkpoint(
checkpointer,
thread_id,
"00000000-0000-6000-8000-000000000014",
messages,
step=1,
)
)
event_visible = False
lookup_visibility: list[bool] = []
async def find_latest_ai_message_run_ids(
_: str,
message_ids: set[str],
*,
user_id=None,
) -> dict[str, str]:
assert message_ids == {"ai-1"}
lookup_visibility.append(event_visible)
return {"ai-1": "exact-run"} if event_visible else {}
class RunManager(_ThreadTestRunManager):
async def list_by_thread(
self,
_thread_id: str,
*,
user_id=None,
limit: int = 100,
) -> list[SimpleNamespace]:
runs = [
SimpleNamespace(
run_id="boundary-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:03+00:00",
)
]
if event_visible:
runs.append(
SimpleNamespace(
run_id="exact-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:07+00:00",
)
)
return runs
@asynccontextmanager
async def reserve_thread_operation(self, _thread_id: str, **kwargs):
nonlocal event_visible
self.reservations.append((_thread_id, kwargs))
event_visible = True
yield
app.state.run_manager = RunManager()
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(
f"/api/threads/{thread_id}/history",
json={"limit": 10},
)
assert response.status_code == 200, response.text
response_ai = response.json()[0]["values"]["messages"][1]
assert response_ai["run_id"] == "boundary-run"
assert response_ai["additional_kwargs"]["turn_duration"] == 3
assert lookup_visibility == [False, True]
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_message_ids"] == {"ai-1": "exact-run"}
assert latest.metadata["run_durations"]["exact-run"] == 7
def test_get_thread_history_revalidates_exact_attribution_after_reservation() -> None:
"""A newer exact event must replace the foreground mapping before persistence."""
app, _store, checkpointer = _build_thread_app()
thread_id = "history-exact-reservation-race"
messages = [
HumanMessage(
id="human-1",
content="Question",
additional_kwargs={"run_id": "boundary-run"},
),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(
_write_checkpoint(
checkpointer,
thread_id,
"00000000-0000-6000-8000-000000000015",
messages,
step=1,
)
)
admitted = False
lookup_states: list[bool] = []
async def find_latest_ai_message_run_ids(
_: str,
message_ids: set[str],
*,
user_id=None,
) -> dict[str, str]:
assert message_ids == {"ai-1"}
lookup_states.append(admitted)
return {"ai-1": "new-run" if admitted else "old-run"}
class RunManager(_ThreadTestRunManager):
async def list_by_thread(
self,
_thread_id: str,
*,
user_id=None,
limit: int = 100,
) -> list[SimpleNamespace]:
runs = [
SimpleNamespace(
run_id="old-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:04+00:00",
)
]
if admitted:
runs.append(
SimpleNamespace(
run_id="new-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:08+00:00",
)
)
return runs
@asynccontextmanager
async def reserve_thread_operation(self, _thread_id: str, **kwargs):
nonlocal admitted
self.reservations.append((_thread_id, kwargs))
admitted = True
yield
app.state.run_manager = RunManager()
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(
f"/api/threads/{thread_id}/history",
json={"limit": 10},
)
assert response.status_code == 200, response.text
response_ai = response.json()[0]["values"]["messages"][1]
assert response_ai["run_id"] == "old-run"
assert response_ai["additional_kwargs"]["turn_duration"] == 4
assert lookup_states == [False, True]
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_message_ids"] == {"ai-1": "new-run"}
assert latest.metadata["run_durations"]["new-run"] == 8
def test_get_thread_history_recomputes_duration_after_reservation() -> None:
"""A final run row must replace a stale foreground duration before persistence."""
app, _store, checkpointer = _build_thread_app()
thread_id = "history-duration-reservation-race"
messages = [
HumanMessage(
id="human-1",
content="Question",
additional_kwargs={"run_id": "run-1"},
),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(
_write_checkpoint(
checkpointer,
thread_id,
"00000000-0000-6000-8000-000000000016",
messages,
step=1,
)
)
admitted = False
lookup_states: list[bool] = []
async def find_latest_ai_message_run_ids(
_: str,
message_ids: set[str],
*,
user_id=None,
) -> dict[str, str]:
assert message_ids == {"ai-1"}
lookup_states.append(admitted)
return {"ai-1": "run-1"}
class RunManager(_ThreadTestRunManager):
async def list_by_thread(
self,
_thread_id: str,
*,
user_id=None,
limit: int = 100,
) -> list[SimpleNamespace]:
return [
SimpleNamespace(
run_id="run-1",
created_at="2026-07-05T00:00:00+00:00",
updated_at=("2026-07-05T00:00:09+00:00" if admitted else "2026-07-05T00:00:03+00:00"),
)
]
@asynccontextmanager
async def reserve_thread_operation(self, _thread_id: str, **kwargs):
nonlocal admitted
self.reservations.append((_thread_id, kwargs))
admitted = True
yield
app.state.run_manager = RunManager()
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(
f"/api/threads/{thread_id}/history",
json={"limit": 10},
)
assert response.status_code == 200, response.text
response_ai = response.json()[0]["values"]["messages"][1]
assert response_ai["run_id"] == "run-1"
assert response_ai["additional_kwargs"]["turn_duration"] == 3
assert lookup_states == [False, True]
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_message_ids"] == {"ai-1": "run-1"}
assert latest.metadata["run_durations"]["run-1"] == 9
def test_get_thread_history_backfills_legacy_durations_with_exact_event_run_id() -> None:
@ -1458,7 +1933,7 @@ def test_get_thread_history_backfills_legacy_durations_with_exact_event_run_id()
]
asyncio.run(_write_checkpoint(checkpointer, thread_id, "00000000-0000-6000-8000-000000000001", messages, step=1))
async def list_by_thread(_: str) -> list[SimpleNamespace]:
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
return [
SimpleNamespace(
run_id="boundary-run",
@ -1472,15 +1947,23 @@ def test_get_thread_history_backfills_legacy_durations_with_exact_event_run_id()
),
]
async def list_messages(_: str, *, limit: int) -> list[dict]:
assert limit == 1000
return [{"content": {"type": "ai", "id": "ai-1"}, "run_id": "exact-run"}]
list_messages_calls: list[str] = []
app.state.run_manager = SimpleNamespace(list_by_thread=list_by_thread)
app.state.run_event_store = SimpleNamespace(list_messages=list_messages)
async def find_latest_ai_message_run_ids(thread: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
assert message_ids == {"ai-1"}
list_messages_calls.append(thread)
return {"ai-1": "exact-run"}
reservation_owner = app.state.run_manager
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
reserve_thread_operation=reservation_owner.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
second_response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
entry = response.json()[0]
@ -1489,10 +1972,200 @@ def test_get_thread_history_backfills_legacy_durations_with_exact_event_run_id()
assert history_messages[1]["additional_kwargs"]["turn_duration"] == 7
assert history_messages[2]["run_id"] == "boundary-run"
assert "run_durations" not in entry["metadata"]
assert list_messages_calls == [thread_id, thread_id]
assert len(reservation_owner.reservations) == 1
reserved_thread_id, reservation_kwargs = reservation_owner.reservations[0]
assert reserved_thread_id == thread_id
assert reservation_kwargs["kind"] is ThreadOperationKind.checkpoint_write
assert isinstance(reservation_kwargs["user_id"], str)
assert second_response.status_code == 200, second_response.text
second_history_messages = second_response.json()[0]["values"]["messages"]
assert second_history_messages[1]["run_id"] == "exact-run"
assert second_history_messages[1]["additional_kwargs"]["turn_duration"] == 7
latest = asyncio.run(checkpointer.aget_tuple({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}))
assert latest is not None
assert latest.metadata["run_durations"] == {"boundary-run": 3, "exact-run": 7}
assert latest.metadata["run_durations"] == {"exact-run": 7}
assert latest.metadata["run_message_ids"] == {"ai-1": "exact-run"}
def test_get_thread_history_finds_ai_event_beyond_ten_thousand_newer_events() -> None:
"""#4949: no arbitrary page cap may turn an old exact run into a boundary run."""
from deerflow.runtime.events.store.memory import MemoryRunEventStore
app, _store, checkpointer = _build_thread_app()
thread_id = "legacy-history-run-id-paginated"
messages = [
HumanMessage(id="human-1", content="Question", additional_kwargs={"run_id": "boundary-run"}),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(_write_checkpoint(checkpointer, thread_id, "00000000-0000-6000-8000-000000000002", messages, step=1))
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
return [
SimpleNamespace(
run_id="boundary-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:03+00:00",
),
SimpleNamespace(
run_id="exact-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:07+00:00",
),
]
event_store = MemoryRunEventStore()
events = [
{
"thread_id": thread_id,
"run_id": "exact-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": "ai-1"},
},
*[
{
"thread_id": thread_id,
"run_id": "noise-run",
"event_type": "llm.ai.response",
"category": "message",
"content": {"type": "ai", "id": f"noise-{index}"},
}
for index in range(10_000)
],
]
asyncio.run(event_store.put_batch(events))
event_store.find_latest_ai_message_run_ids = AsyncMock(wraps=event_store.find_latest_ai_message_run_ids)
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = event_store
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
history_messages = response.json()[0]["values"]["messages"]
assert history_messages[1]["run_id"] == "exact-run"
assert history_messages[1]["additional_kwargs"]["turn_duration"] == 7
assert event_store.find_latest_ai_message_run_ids.await_count == 2
def test_get_thread_history_sizes_initial_run_page_to_required_attributions() -> None:
"""A long thread should batch-hydrate its common migration path."""
app, _store, checkpointer = _build_thread_app()
thread_id = "legacy-history-run-page-sizing"
run_count = 101
messages = []
runs = []
message_run_ids: dict[str, str] = {}
for index in range(run_count):
boundary_run_id = f"boundary-{index}"
exact_run_id = f"exact-{index}"
message_id = f"ai-{index}"
messages.extend(
[
HumanMessage(id=f"human-{index}", content=f"Question {index}", additional_kwargs={"run_id": boundary_run_id}),
AIMessage(id=message_id, content=f"Answer {index}"),
]
)
message_run_ids[message_id] = exact_run_id
runs.append(
SimpleNamespace(
run_id=exact_run_id,
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:05+00:00",
)
)
asyncio.run(_write_checkpoint(checkpointer, thread_id, "00000000-0000-6000-8000-000000000020", messages, step=1))
list_limits: list[int] = []
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
list_limits.append(limit)
return runs[:limit]
async def get(run_id: str, *, user_id=None) -> SimpleNamespace | None:
return next((run for run in runs if run.run_id == run_id), None)
get_mock = AsyncMock(side_effect=get)
async def find_latest_ai_message_run_ids(_: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
assert message_ids == set(message_run_ids)
return message_run_ids
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
get=get_mock,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
ai_messages = [message for message in response.json()[0]["values"]["messages"] if message["type"] == "ai"]
assert len(ai_messages) == run_count
assert ai_messages[-1]["additional_kwargs"]["turn_duration"] == 5
assert list_limits == [run_count, run_count]
get_mock.assert_not_awaited()
def test_get_thread_history_fetches_exact_run_older_than_default_run_page() -> None:
"""The event index may resolve a run outside RunManager's newest-100 page."""
app, _store, checkpointer = _build_thread_app()
thread_id = "legacy-history-old-exact-run"
messages = [
HumanMessage(id="human-1", content="Question", additional_kwargs={"run_id": "boundary-run"}),
AIMessage(id="ai-1", content="Answer"),
]
asyncio.run(_write_checkpoint(checkpointer, thread_id, "00000000-0000-6000-8000-000000000003", messages, step=1))
boundary_run = SimpleNamespace(
run_id="boundary-run",
created_at="2026-07-05T00:00:00+00:00",
updated_at="2026-07-05T00:00:03+00:00",
)
exact_run = SimpleNamespace(
run_id="old-exact-run",
created_at="2026-06-01T00:00:00+00:00",
updated_at="2026-06-01T00:00:09+00:00",
)
get_calls: list[str] = []
async def list_by_thread(_: str, *, user_id=None, limit: int = 100) -> list[SimpleNamespace]:
assert limit == 100
return [boundary_run]
async def get(run_id: str, *, user_id=None) -> SimpleNamespace | None:
get_calls.append(run_id)
return exact_run if run_id == exact_run.run_id else None
async def find_latest_ai_message_run_ids(_: str, message_ids: set[str], *, user_id=None) -> dict[str, str]:
assert message_ids == {"ai-1"}
return {"ai-1": exact_run.run_id}
app.state.run_manager = SimpleNamespace(
list_by_thread=list_by_thread,
get=get,
reserve_thread_operation=app.state.run_manager.reserve_thread_operation,
)
app.state.run_event_store = SimpleNamespace(find_latest_ai_message_run_ids=find_latest_ai_message_run_ids)
with TestClient(app) as client:
response = client.post(f"/api/threads/{thread_id}/history", json={"limit": 10})
assert response.status_code == 200, response.text
history_messages = response.json()[0]["values"]["messages"]
assert history_messages[1]["run_id"] == exact_run.run_id
assert history_messages[1]["additional_kwargs"]["turn_duration"] == 9
assert get_calls == [exact_run.run_id, exact_run.run_id]
def test_get_thread_history_injects_turn_duration_once_per_run() -> None:
@ -1529,8 +2202,9 @@ def test_get_thread_history_injects_turn_duration_once_per_run() -> None:
run_manager = AsyncMock()
run_manager.list_by_thread = AsyncMock(return_value=[_run("run-1", 5), _run("run-2", 9)])
run_manager.reserve_thread_operation = _ThreadTestRunManager().reserve_thread_operation
event_store = MagicMock()
event_store.list_messages = AsyncMock(return_value=[])
event_store.find_latest_ai_message_run_ids = AsyncMock(return_value={})
app.state.run_manager = run_manager
app.state.run_event_store = event_store