diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 43f09580bc6..ff1e45e746b 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -66,7 +66,7 @@ ) 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) - - # 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) + 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) + + 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") 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 - - stamp_turn_duration_on_last_ai(serialized_msgs, 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. + 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 + + 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 + + # 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 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"] diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index 2c27ba96740..3ebb7ddd8c3 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -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` diff --git a/backend/packages/harness/deerflow/runtime/events/store/base.py b/backend/packages/harness/deerflow/runtime/events/store/base.py index df552e40d61..8cf4ee8bb69 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/base.py +++ b/backend/packages/harness/deerflow/runtime/events/store/base.py @@ -16,6 +16,31 @@ 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 @@ async def list_messages( 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, diff --git a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py index 440d361bc2f..9185adbfdd2 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py +++ b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py @@ -30,7 +30,7 @@ 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 @@ async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq 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: diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 9827afd118e..38ab9fd6644 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -549,6 +549,26 @@ async def update_run_progress(self, run_id: str, **kwargs) -> None: 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, diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index d21cf495832..a690ea31ff9 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -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 _stream_once(input_payload: Any, stream_config: RunnableConfig) -> Non 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 _stream_once(input_payload: Any, stream_config: RunnableConfig) -> Non 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 _stream_once(input_payload: Any, stream_config: RunnableConfig) -> Non 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, diff --git a/backend/tests/test_run_duration_checkpoint.py b/backend/tests/test_run_duration_checkpoint.py index 7119de203c1..0ad1f6a5727 100644 --- a/backend/tests/test_run_duration_checkpoint.py +++ b/backend/tests/test_run_duration_checkpoint.py @@ -9,9 +9,12 @@ 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 @@ def factory(*, config): 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.""" diff --git a/backend/tests/test_run_event_store.py b/backend/tests/test_run_event_store.py index 00b543cb4d7..76ba803b8a3 100644 --- a/backend/tests/test_run_event_store.py +++ b/backend/tests/test_run_event_store.py @@ -4,6 +4,8 @@ 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 @@ async def test_pagination_with_interleaved_trace_events(self, store): 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 @@ async def test_basic_crud(self, tmp_path): 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 @@ async def test_basic_crud(self, tmp_path): 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 diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index 740b08665ca..567e1d74666 100644 --- a/backend/tests/test_threads_router.py +++ b/backend/tests/test_threads_router.py @@ -1355,6 +1355,7 @@ async def _seed() -> 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 @@ async def list_by_thread(_: str) -> list[SimpleNamespace]: ), ] - 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 @@ async def list_messages(thread: str, *, limit: int) -> list[dict]: 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 @@ async def list_by_thread(_: str) -> list[SimpleNamespace]: ), ] - 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 @@ async def list_messages(_: str, *, limit: int) -> list[dict]: 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 _run(run_id: str, seconds: int) -> RunRecord: 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