From 2e15e3fe0dbe8bd3b3d813a6a3445891b4973c71 Mon Sep 17 00:00:00 2001 From: xiawiie <106433113+xiawiie@users.noreply.github.com> Date: Thu, 2 Jul 2026 00:13:01 +0800 Subject: [PATCH] fix: generate fallback title for interrupted first-turn runs (#3874) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: generate title for interrupted first turn * test(title): cover partial-exchange + dict-form messages Harden the interrupted-run fallback path added in 19fc34fd: - TitleMiddleware._should_generate_title now accepts a lone first-turn user message when allow_partial_exchange=True, so the worker can still derive a title if cancellation lands before any AI chunk is checkpointed. - runtime/runs/worker._ensure_interrupted_title computes the next checkpoint step defensively (treat missing/non-int step as 0) and renames a shadowed ckpt_config local for readability. - Add four unit tests in tests/test_title_middleware_core_logic.py: partial-exchange allows user-only, partial-exchange still respects an existing title, dict-form messages are recognized, and the sync fallback path derives a title from dict-form messages — matching what channel_values stores in the checkpoint. Refs #3859. * fix: persist interrupted-title via channel_versions bump Address PR #3874 review feedback: ``_ensure_interrupted_title`` previously called ``aput(..., new_versions={})``. LangGraph's DB-backed savers (``PostgresSaver`` and the v4 ``SqliteSaver`` blob layout) strip inline ``channel_values`` from ``put`` and only persist blobs for channels named in ``new_versions`` — so the fallback ``title`` channel was dropped on read-back and ``threads_meta.display_name`` stayed ``"Untitled"`` after refresh on those backends. The original in-memory e2e passed because ``InMemorySaver`` keeps the inline snapshot verbatim. Fix mirrors ``_rollback_to_pre_run_checkpoint`` in the same file: bump ``channel_versions["title"]`` (via ``checkpointer.get_next_version`` when available, else int/string fallbacks), persist the new version on the checkpoint, and declare it in ``new_versions`` so the DB savers actually write the blob. Regression coverage in ``tests/test_run_worker_rollback.py``: - ``test_ensure_interrupted_title_bumps_channel_version_and_declares_it_in_new_versions`` — exact ``aput`` invariants: ``new_versions == {"title": 1}``, the written checkpoint's ``channel_versions["title"]`` is bumped, and the pre-existing ``messages`` version is preserved. - ``test_ensure_interrupted_title_bumps_existing_string_version`` — string-shaped prior version (some savers use UUID-style versions); bumped value must differ from the prior, no overwrite-in-place. - ``test_ensure_interrupted_title_skips_when_title_already_set`` — title short-circuit; no extra ``aput``. - ``test_ensure_interrupted_title_returns_none_when_no_checkpoint`` — no checkpoint yet; returns ``None`` without writing. - ``test_ensure_interrupted_title_round_trip_with_real_sqlite_checkpointer`` — full round-trip against a real ``AsyncSqliteSaver`` on a disk-backed DB, then closes and re-opens the saver to simulate a fresh connection. The fallback title must still be present on the second ``aget_tuple``. This is the exact scenario the review flagged. Validated locally with the full backend suite: 5195 passed, 18 skipped. Refs #3859. Addresses review on #3874. * test(worker, title): harden interrupted-title fallback for every saver Defensive coverage on top of the channel_versions fix (commit 05253957), addressing edge cases surfaced during a second-pass review of #3874. Worker: - Extract version bump into ``_bump_channel_version(checkpointer, current)`` with explicit fallbacks for int / float / numeric-string / UUID-shaped string / None / bool, AND a wrap-around defense when the saver's ``get_next_version`` raises or returns an unchanged value. The invariant is: returned version MUST differ from the prior. Without this, a saver bug (or a custom backend) could leave ``new_versions={"title": v}`` no-op on DB savers — the very class of bug the original review pointed out. Title middleware: - Coerce ``state.get("messages")`` from ``None`` to ``[]`` on both ``_should_generate_title`` and ``_build_title_prompt``. A partially-initialized checkpoint can carry ``messages=None`` on the channel_values channel (the worker reads raw channel_values, not BaseMessages), and the default kwarg only protects against a missing key. Repro: ``TypeError: 'NoneType' object is not iterable`` from the next() generator — confirmed by reverting the fix and watching ``test_*_handles_none_messages_channel`` go red. Tests (TDD-verified red→green for the new asserts): - ``test_run_worker_rollback.py``: * ``_bump_channel_version`` — 8 tests covering every version type (int, float, numeric string, UUID-style string, None, bool) and every saver-side fault mode (no ``get_next_version`` / raising / stuck on identity). * ``test_ensure_interrupted_title_*`` — 5 additional helper boundary tests: title.enabled=false short-circuit; empty messages list; messages=None; aput-error propagation (helper contract: caller swallows, not the helper); idempotency on a real InMemorySaver across two invocations. * ``test_ensure_interrupted_title_preserves_non_title_channel_versions`` — pins that ``new_versions`` only contains ``"title"`` and that other channels' versions are untouched (regression anchor for a sloppier draft that bumped every channel). * ``test_worker_finally_block_swallows_helper_exceptions`` — pins the integration contract: even if the helper raises, the worker's threads_meta status sync still runs and ``publish_end`` is still awaited so the SSE stream closes cleanly. - ``test_title_middleware_core_logic.py``: * 4 additional tests: ``messages=None`` on both ``_should_generate_title`` and ``_build_title_prompt``; the ``role: user`` / ``role: assistant`` (OpenAI-style) dict normalization; partial-exchange path with a dict-form message. Verification: - ``PYTHONPATH=. uv run pytest tests/ -x --ignore=tests/blocking_io -q`` → 5215 passed, 18 skipped. - ``ruff check`` + ``ruff format --check`` clean on every touched file. - Red/green TDD verification: temporarily reverted the ``new_versions={}`` fix → 4 new tests went red as expected; restored and the suite is green again. Same red/green dance for the ``messages=None`` coercion. Refs #3859. Addresses second-pass review on #3874. * fix(title): ignore dict context reminders in fallback * fix(worker): link interrupted-title checkpoint to its parent The title-bump checkpoint written by ``_ensure_interrupted_title`` was landing without a ``parent_checkpoint_id`` — a real orphan in the LangGraph history graph. Reproduction (disk-backed AsyncSqliteSaver): [seed] checkpoint_id = 1f173dbc... [helper] wrote title = "Why is the sky blue?" [issue 1] new checkpoint = 1f173dbc..., parent = None [issue 1] is new checkpoint orphaned? True Root cause: ``_ensure_interrupted_title`` built ``write_config`` as ``{"thread_id": ..., "checkpoint_ns": ...}`` only. ``BaseCheckpointSaver`` implementations read ``configurable.checkpoint_id`` from that config as the *parent* id when inserting (see ``langgraph/checkpoint/sqlite/aio.py`` ``aput``: ``config["configurable"].get("checkpoint_id")`` becomes the ``parent_checkpoint_id`` column). With no value, the saver writes NULL — the new checkpoint is a tree root. Consequences: - Any future LangGraph ``runs.resume_from`` / time-travel feature has no backward edge to walk past the title-bump. - History-visualization UIs built on ``alist()`` render the title-bump as a sibling of the prior checkpoint, not its descendant. Fix: read ``checkpoint_id`` off the tuple's own config and thread it into ``write_config["configurable"]["checkpoint_id"]`` before calling ``aput``, the same pattern every middleware-driven write uses. Three new regression tests against real ``AsyncSqliteSaver`` (disk-backed, fresh connections so we exercise the on-disk read path): - ``test_ensure_interrupted_title_links_new_checkpoint_to_its_parent`` — asserts ``latest.parent_config["configurable"]["checkpoint_id"]`` equals the seeded checkpoint id. TDD red-green verified: reverting the fix flips this test red with ``AssertionError: title-bump checkpoint must have a parent_config``. - ``test_ensure_interrupted_title_appears_in_history_with_audit_marker`` — pins the audit contract: the title-bump entry in ``alist()`` carries ``metadata.source == "update"`` and ``metadata.writes`` contains ``runtime_interrupt_title``. This is a deliberate design choice — we do NOT hide the entry from history (audit trail belongs in the saver), but its source and writes marker MUST be unambiguous so UIs/tools can identify it. - ``test_ensure_interrupted_title_survives_immediate_next_turn`` — cancel → immediate user follow-up scenario. Simulates the agent's next turn appending a (user, ai) pair without touching the title channel, then opens a fresh saver and verifies the title is still present after the next-turn checkpoint write. Pins the channel-version-blob invariant established by commit 05253957 — without the ``new_versions={"title": v}`` declaration there, the title blob would vanish from the DB and this test would read back ``None``. Verification: - ``PYTHONPATH=. uv run pytest tests/ -x --ignore=tests/blocking_io -q`` → 5222 passed, 15 skipped. - ``ruff check`` + ``ruff format --check`` clean on every touched file. - Reproduction script confirms ``parent_checkpoint_id`` is now non-null and the next-turn read-back preserves the fallback title. Refs #3859. * Revert "fix(worker): link interrupted-title checkpoint to its parent" This reverts commit c763ed9334781db1acdce0f5f33d663d8d5f80ff. * test: trim over-engineered test coverage Reduce review surface area on PR #3874 by dropping defensive tests that don't pin a real invariant. After self-review: - ``_bump_channel_version``: 8 tests → 2 (happy path + saver-error fallback). Dropped float / bool / numeric-string / UUID-string / missing-get-next-version / stuck-get-next-version branches — those are speculative scaffolding for savers we don't ship. - ``_ensure_interrupted_title``: dropped ``returns_none_when_title_disabled``, ``returns_none_with_no_user_message``, ``returns_none_when_no_checkpoint`` — boundary guards already exercised by the e2e test and the ``handles_none_messages_channel`` regression anchor. Net: -107 lines of test code. Remaining coverage still pins every red-green-verified invariant (channel_versions bump, string-version bump, idempotency, sqlite round-trip, non-title channel preservation, aput-error contract, worker finally swallowing, partial-exchange). Verification: 5209 passed, 15 skipped. * fix: harden interrupted title finalization * fix: serialize interrupted title finalization * fix: preserve interrupt semantics during title finalization * fix: preserve delayed interrupted title recovery --------- Co-authored-by: Willem Jiang --- README.md | 2 + backend/AGENTS.md | 2 +- .../docs/TITLE_GENERATION_IMPLEMENTATION.md | 6 + .../agents/middlewares/title_middleware.py | 60 +- .../harness/deerflow/runtime/runs/manager.py | 73 +- .../harness/deerflow/runtime/runs/worker.py | 176 ++++ backend/tests/test_run_worker_rollback.py | 984 +++++++++++++++++- backend/tests/test_runtime_lifecycle_e2e.py | 151 ++- .../tests/test_title_middleware_core_logic.py | 107 ++ .../tests/test_worker_langfuse_metadata.py | 9 + frontend/AGENTS.md | 5 +- frontend/src/core/threads/hooks.ts | 77 +- .../tests/unit/core/threads/infinite.test.ts | 191 +++- 13 files changed, 1806 insertions(+), 37 deletions(-) diff --git a/README.md b/README.md index fba1f3e6a..5b7fffaad 100644 --- a/README.md +++ b/README.md @@ -606,6 +606,8 @@ Tools follow the same philosophy. DeerFlow comes with a core toolset — web sea Gateway-generated follow-up suggestions now normalize both plain-string model output and block/list-style rich content before parsing the JSON array response, so provider-specific content wrappers do not silently drop suggestions. +Interrupted first-turn runs still persist a fallback conversation title, so stopping a streaming response does not leave the thread as "Untitled" after refresh. + ``` # Paths inside the sandbox container /mnt/skills/public diff --git a/backend/AGENTS.md b/backend/AGENTS.md index c9fd31405..e6945d25e 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -224,7 +224,7 @@ Lead-agent middlewares are assembled in strict order across three functions: the 14. **SummarizationMiddleware** - *(optional, if enabled)* Context reduction when approaching token limits 15. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with the `write_todos` tool 16. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is merged back into the dispatching AIMessage by message position -17. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model +17. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint. 18. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses) 19. **ViewImageMiddleware** - *(optional, if the model supports vision)* Injects base64 image data before the LLM call 20. **DeferredToolFilterMiddleware** - *(optional, if `tool_search.enabled`)* Hides deferred (MCP) tool schemas from the bound model until `tool_search` promotes them (reads per-thread promotions from `ThreadState.promoted`, hash-scoped) diff --git a/backend/docs/TITLE_GENERATION_IMPLEMENTATION.md b/backend/docs/TITLE_GENERATION_IMPLEMENTATION.md index 316b355f8..b0a7a0b03 100644 --- a/backend/docs/TITLE_GENERATION_IMPLEMENTATION.md +++ b/backend/docs/TITLE_GENERATION_IMPLEMENTATION.md @@ -186,6 +186,12 @@ pytest 2. LangGraph Platform 会自动持久化 3. 检查数据库确认 checkpointer 工作正常 +### 中断首轮后仍显示默认标题？ + +1. `runtime/runs/worker.py` 会在 interrupted-run cleanup 中保持 run 处于 finalizing 状态，避免同线程新 run 在 fallback title 写入期间覆盖 checkpoint +2. 如果取消发生在可用 checkpoint 写入前，worker 会使用本次 `graph_input` 中的首条用户消息生成本地 fallback title +3. fallback title 写入前会重新读取 latest checkpoint；如果同线程状态已经前进，只对最新 snapshot 做 title-only 更新，避免旧消息重新成为 latest + --- ## 📊 性能影响 diff --git a/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py index f2118b46f..0b75a71f0 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py @@ -63,16 +63,42 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): return "" + @staticmethod + def _message_type(message: object) -> str | None: + message_type = getattr(message, "type", None) + if message_type is None and isinstance(message, dict): + message_type = message.get("type") or message.get("role") + if message_type == "user": + return "human" + if message_type == "assistant": + return "ai" + return message_type if isinstance(message_type, str) else None + + @staticmethod + def _message_content(message: object) -> object: + if isinstance(message, dict): + return message.get("content", "") + return getattr(message, "content", "") + + @staticmethod + def _is_dynamic_context_reminder_message(message: object) -> bool: + if is_dynamic_context_reminder(message): + return True + if isinstance(message, dict): + additional_kwargs = message.get("additional_kwargs") + return isinstance(additional_kwargs, dict) and bool(additional_kwargs.get("dynamic_context_reminder")) + return False + @staticmethod def _is_user_message_for_title(message: object) -> bool: - return getattr(message, "type", None) == "human" and not is_dynamic_context_reminder(message) + return TitleMiddleware._message_type(message) == "human" and not TitleMiddleware._is_dynamic_context_reminder_message(message) def _get_title_user_message(self, state: TitleMiddlewareState) -> str: - messages = state.get("messages", []) - user_msg_content = next((m.content for m in messages if self._is_user_message_for_title(m)), "") + messages = state.get("messages") or [] + user_msg_content = next((self._message_content(m) for m in messages if self._is_user_message_for_title(m)), "") return self._normalize_content(user_msg_content) - def _should_generate_title(self, state: TitleMiddlewareState) -> bool: + def _should_generate_title(self, state: TitleMiddlewareState, *, allow_partial_exchange: bool = False) -> bool: """Check if we should generate a title for this thread.""" config = self._get_title_config() if not config.enabled: @@ -82,17 +108,23 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): if state.get("title"): return False - # Check if this is the first turn (has at least one user message and one assistant response) - messages = state.get("messages", []) - if len(messages) < 2: + # Check if this is the first turn (has at least one user message and one assistant response). + # Defensively coerce a None ``messages`` channel (possible when reading a + # partially-initialized checkpoint) into an empty list so ``len()`` is safe. + messages = state.get("messages") or [] + min_messages = 1 if allow_partial_exchange else 2 + if len(messages) < min_messages: return False # Count user and assistant messages user_messages = [m for m in messages if self._is_user_message_for_title(m)] - assistant_messages = [m for m in messages if m.type == "ai"] + assistant_messages = [m for m in messages if self._message_type(m) == "ai"] - # Generate title after first complete exchange - return len(user_messages) == 1 and len(assistant_messages) >= 1 + # Normal path: title only after first complete exchange. Interrupted path + # (``allow_partial_exchange=True``) accepts a lone first-turn user message + # so a fallback title can still be persisted when the run is cancelled + # before any AI chunk reaches the checkpoint. + return len(user_messages) == 1 and (len(assistant_messages) >= 1 or allow_partial_exchange) def _build_title_prompt(self, state: TitleMiddlewareState) -> tuple[str, str]: """Extract user/assistant messages and build the title prompt. @@ -100,9 +132,9 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): Returns (prompt_string, user_msg) so callers can use user_msg as fallback. """ config = self._get_title_config() - messages = state.get("messages", []) + messages = state.get("messages") or [] - assistant_msg_content = next((m.content for m in messages if m.type == "ai"), "") + assistant_msg_content = next((self._message_content(m) for m in messages if self._message_type(m) == "ai"), "") user_msg = self._get_title_user_message(state) assistant_msg = self._strip_think_tags(self._normalize_content(assistant_msg_content)) @@ -152,9 +184,9 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): ] return config - def _generate_title_result(self, state: TitleMiddlewareState) -> dict | None: + def _generate_title_result(self, state: TitleMiddlewareState, *, allow_partial_exchange: bool = False) -> dict | None: """Generate a local fallback title without blocking on an LLM call.""" - if not self._should_generate_title(state): + if not self._should_generate_title(state, allow_partial_exchange=allow_partial_exchange): return None user_msg = self._get_title_user_message(state) diff --git a/backend/packages/harness/deerflow/runtime/runs/manager.py b/backend/packages/harness/deerflow/runtime/runs/manager.py index 4b52d0aca..60646be7b 100644 --- a/backend/packages/harness/deerflow/runtime/runs/manager.py +++ b/backend/packages/harness/deerflow/runtime/runs/manager.py @@ -104,6 +104,7 @@ class RunRecord: message_count: int = 0 last_ai_message: str | None = None first_human_message: str | None = None + finalizing: bool = False class RunManager: @@ -484,6 +485,64 @@ class RunManager: await self._persist_status(record, status, error=error) logger.info("Run %s -> %s", run_id, status.value) + async def set_finalizing(self, run_id: str, finalizing: bool) -> None: + """Mark whether a run is performing post-cancel cleanup.""" + async with self._lock: + record = self._runs.get(run_id) + if record is None: + logger.warning("set_finalizing called for unknown run %s", run_id) + return + record.finalizing = finalizing + record.updated_at = _now_iso() + + async def wait_for_prior_finalizing( + self, + thread_id: str, + run_id: str, + *, + poll_interval: float = 0.01, + ) -> None: + """Wait until older same-thread runs have finished post-cancel cleanup.""" + while True: + async with self._lock: + found_current = False + prior_finalizing = False + for record in self._thread_records_locked(thread_id): + if record.run_id == run_id: + found_current = True + break + if record.finalizing: + prior_finalizing = True + + if not found_current or not prior_finalizing: + return + + await asyncio.sleep(poll_interval) + + async def has_later_run(self, thread_id: str, run_id: str) -> bool: + """Return whether a newer in-memory run has been admitted for the thread.""" + async with self._lock: + seen_current = False + for record in self._thread_records_locked(thread_id): + if record.run_id == run_id: + seen_current = True + continue + if seen_current: + return True + return False + + async def has_later_started_run(self, thread_id: str, run_id: str) -> bool: + """Return whether a newer same-thread run may have already advanced state.""" + async with self._lock: + seen_current = False + for record in self._thread_records_locked(thread_id): + if record.run_id == run_id: + seen_current = True + continue + if seen_current and (record.status != RunStatus.pending or record.finalizing): + return True + return False + async def _persist_model_name(self, run_id: str, model_name: str | None) -> None: """Best-effort persist model_name update to the backing store.""" if self._store is None: @@ -532,7 +591,9 @@ class RunManager: return False record.abort_action = action record.abort_event.set() - if record.task is not None and not record.task.done(): + task_active = record.task is not None and not record.task.done() + record.finalizing = task_active + if task_active: record.task.cancel() record.status = RunStatus.interrupted record.updated_at = _now_iso() @@ -571,7 +632,7 @@ class RunManager: if multitask_strategy not in _supported_strategies: raise UnsupportedStrategyError(f"Multitask strategy '{multitask_strategy}' is not yet supported. Supported strategies: {', '.join(_supported_strategies)}") - inflight = [r for r in self._thread_records_locked(thread_id) if r.status in (RunStatus.pending, RunStatus.running)] + inflight = [r for r in self._thread_records_locked(thread_id) if r.status in (RunStatus.pending, RunStatus.running) or r.finalizing] if multitask_strategy == "reject" and inflight: raise ConflictError(f"Thread {thread_id} already has an active run") @@ -615,9 +676,13 @@ class RunManager: if multitask_strategy in ("interrupt", "rollback") and inflight: for r in inflight: + if r.finalizing: + continue r.abort_action = multitask_strategy r.abort_event.set() - if r.task is not None and not r.task.done(): + task_active = r.task is not None and not r.task.done() + r.finalizing = task_active + if task_active: r.task.cancel() r.status = RunStatus.interrupted r.updated_at = now @@ -685,7 +750,7 @@ class RunManager: async def has_inflight(self, thread_id: str) -> bool: """Return ``True`` if *thread_id* has a pending or running run.""" async with self._lock: - return any(r.status in (RunStatus.pending, RunStatus.running) for r in self._thread_records_locked(thread_id)) + return any(r.status in (RunStatus.pending, RunStatus.running) or r.finalizing for r in self._thread_records_locked(thread_id)) async def cleanup(self, run_id: str, *, delay: float = 300) -> None: """Remove a run record after an optional delay.""" diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 8948354ae..2e90828bd 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -159,6 +159,8 @@ async def run_agent( ) try: + await run_manager.wait_for_prior_finalizing(thread_id, run_id) + # Initialize RunJournal + write human_message event. # These are inside the try block so any exception (e.g. a DB # error writing the event) flows through the except/finally @@ -335,6 +337,7 @@ async def run_agent( # 8. Final status if record.abort_event.is_set(): + await run_manager.set_finalizing(run_id, True) action = record.abort_action if action == "rollback": await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user") @@ -362,6 +365,7 @@ async def run_agent( await run_manager.set_status(run_id, RunStatus.success) except asyncio.CancelledError: + await run_manager.set_finalizing(run_id, True) action = record.abort_action if action == "rollback": await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user") @@ -409,6 +413,14 @@ async def run_agent( except Exception: logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) + if checkpointer is not None and record.status == RunStatus.interrupted: + try: + await run_manager.wait_for_prior_finalizing(thread_id, run_id) + if not await run_manager.has_later_started_run(thread_id, run_id): + await _ensure_interrupted_title(checkpointer=checkpointer, thread_id=thread_id, app_config=ctx.app_config, graph_input=graph_input) + except Exception: + logger.debug("Failed to generate interrupted title for thread %s (non-fatal)", thread_id) + # Sync title from checkpoint to threads_meta.display_name if checkpointer is not None and thread_store is not None: try: @@ -430,6 +442,9 @@ async def run_agent( except Exception: logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id) + if record.finalizing: + await run_manager.set_finalizing(run_id, False) + await bridge.publish_end(run_id) asyncio.create_task(bridge.cleanup(run_id, delay=60)) @@ -548,6 +563,167 @@ def _new_checkpoint_marker() -> dict[str, str]: return {"id": marker["id"], "ts": marker["ts"]} +def _bump_channel_version(checkpointer: Any, current_version: Any) -> Any: + """Return a strictly-different next version for a checkpoint channel. + + DB-backed LangGraph savers (PostgresSaver / v4 SqliteSaver blob layout) + persist channel blobs keyed by ``channel_versions[]``, so the + new value MUST differ from the prior value. We delegate to the + checkpointer's ``get_next_version`` when available — that is the canonical + versioning scheme each saver picks (int, monotonic float, or + UUID-shaped string). When the checkpointer doesn't expose it (or it + returns ``None``/an unchanged value), fall back to a defensive bump that + still guarantees inequality. + """ + get_next_version = getattr(checkpointer, "get_next_version", None) + if callable(get_next_version): + try: + next_version = get_next_version(current_version, None) + except Exception: + next_version = None + if next_version is not None and next_version != current_version: + return next_version + # fall through to defensive bump + + if isinstance(current_version, bool): + # ``bool`` is a subclass of ``int``; treat True/False as 1/0 instead of + # adding to the boolean itself, which would produce an int anyway but + # via a path that surprises readers. + return int(current_version) + 1 + if isinstance(current_version, int): + return current_version + 1 + if isinstance(current_version, float): + # Match LangGraph's default float versioning (monotonic increment). + return current_version + 1.0 + if isinstance(current_version, str): + try: + return str(int(current_version) + 1) + except ValueError: + return f"{current_version}.1" + return 1 + + +def _checkpoint_identity(ckpt_tuple: Any | None, checkpoint: dict[str, Any]) -> str | None: + tuple_config = getattr(ckpt_tuple, "config", {}) or {} + tuple_configurable = tuple_config.get("configurable", {}) if isinstance(tuple_config, dict) else {} + if isinstance(tuple_configurable, dict): + checkpoint_id = tuple_configurable.get("checkpoint_id") + if isinstance(checkpoint_id, str) and checkpoint_id: + return checkpoint_id + checkpoint_id = checkpoint.get("id") + return checkpoint_id if isinstance(checkpoint_id, str) and checkpoint_id else None + + +def _checkpoint_namespace(ckpt_tuple: Any | None) -> str: + tuple_config = getattr(ckpt_tuple, "config", {}) or {} + tuple_configurable = tuple_config.get("configurable", {}) if isinstance(tuple_config, dict) else {} + checkpoint_ns = tuple_configurable.get("checkpoint_ns", "") if isinstance(tuple_configurable, dict) else "" + return checkpoint_ns if isinstance(checkpoint_ns, str) else "" + + +def _graph_input_messages(graph_input: Any | None) -> list[Any]: + if not isinstance(graph_input, dict): + return [] + messages = graph_input.get("messages") + if isinstance(messages, list): + return messages + if isinstance(messages, tuple): + return list(messages) + return [] + + +def _title_generation_state(channel_values: dict[str, Any], graph_input: Any | None) -> dict[str, Any]: + state = dict(channel_values) + messages = state.get("messages") + if not messages: + fallback_messages = _graph_input_messages(graph_input) + if fallback_messages: + state["messages"] = fallback_messages + return state + + +async def _ensure_interrupted_title(*, checkpointer: Any, thread_id: str, app_config: AppConfig | None, graph_input: Any | None = None) -> str | None: + """Persist a local fallback title for interrupted first-turn runs. + + Returns the title that is now persisted (existing or newly written), or + ``None`` when no checkpoint is available or no title text can be derived. + Idempotent: re-invoking against a checkpoint that already carries a title + short-circuits without writing a new checkpoint. + """ + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + middleware = TitleMiddleware(app_config=app_config) if app_config is not None else TitleMiddleware() + ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} + + for _attempt in range(3): + ckpt_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", ckpt_config) + checkpoint = copy.deepcopy(getattr(ckpt_tuple, "checkpoint", {}) or {}) if ckpt_tuple is not None else empty_checkpoint() + channel_values = dict(checkpoint.get("channel_values", {}) or {}) + existing_title = channel_values.get("title") + if existing_title: + return existing_title + + result = middleware._generate_title_result(_title_generation_state(channel_values, graph_input), allow_partial_exchange=True) + title = result.get("title") if isinstance(result, dict) else None + if not title: + return None + + # ``empty_checkpoint()`` creates a fresh id every time; only real tuples + # carry an identity stable enough for the stale-snapshot comparison. + base_identity = _checkpoint_identity(ckpt_tuple, checkpoint) if ckpt_tuple is not None else None + latest_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", ckpt_config) + latest_checkpoint = copy.deepcopy(getattr(latest_tuple, "checkpoint", {}) or {}) if latest_tuple is not None else empty_checkpoint() + latest_identity = _checkpoint_identity(latest_tuple, latest_checkpoint) if latest_tuple is not None else None + if base_identity is None: + if latest_identity is not None: + continue + elif latest_identity != base_identity: + continue + + checkpoint = latest_checkpoint + channel_values = dict(checkpoint.get("channel_values", {}) or {}) + existing_title = channel_values.get("title") + if existing_title: + return existing_title + + channel_values["title"] = title + marker = _new_checkpoint_marker() + checkpoint.update({"id": marker["id"], "ts": marker["ts"], "channel_values": channel_values}) + + # Bump ``channel_versions["title"]`` and declare the bump in ``new_versions`` + # so DB-backed savers (SqliteSaver v4 / PostgresSaver) actually persist the + # new blob — those savers strip inline ``channel_values`` from ``put`` and + # only write blobs for channels listed in ``new_versions``. The legacy + # single-table sqlite saver ignores ``new_versions`` and inlines the + # snapshot, so this path is correct for both layouts. Mirrors + # ``_rollback_to_pre_run_checkpoint`` in the same file. + channel_versions = dict(checkpoint.get("channel_versions", {}) or {}) + next_title_version = _bump_channel_version(checkpointer, channel_versions.get("title")) + channel_versions["title"] = next_title_version + checkpoint["channel_versions"] = channel_versions + + metadata = dict(getattr(latest_tuple, "metadata", {}) or {}) + metadata["source"] = "update" + prev_step = metadata.get("step") + metadata["step"] = (prev_step + 1) if isinstance(prev_step, int) else 1 + metadata["writes"] = {"runtime_interrupt_title": {"title": title}} + + checkpoint_ns = _checkpoint_namespace(latest_tuple) + write_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}} + await _call_checkpointer_method( + checkpointer, + "aput", + "put", + write_config, + checkpoint, + metadata, + {"title": next_title_version}, + ) + return title + + return None + + def _lg_mode_to_sse_event(mode: str) -> str: """Map LangGraph internal stream_mode name to SSE event name. diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 00844d88d..3e9b10094 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -1,5 +1,8 @@ import asyncio +import copy +from contextlib import suppress from types import SimpleNamespace +from typing import Any from unittest.mock import AsyncMock, call import pytest @@ -7,12 +10,14 @@ from langchain_core.messages import AIMessage from langgraph.checkpoint.base import empty_checkpoint from langgraph.checkpoint.memory import InMemorySaver -from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.manager import ConflictError, RunManager from deerflow.runtime.runs.schemas import RunStatus from deerflow.runtime.runs.worker import ( RunContext, _agent_factory_supports_app_config, _build_runtime_context, + _bump_channel_version, + _ensure_interrupted_title, _extract_llm_error_fallback_message, _install_runtime_context, _rollback_to_pre_run_checkpoint, @@ -672,3 +677,980 @@ def test_extract_llm_error_fallback_message_updates_mode_no_fallback(): ] } assert _extract_llm_error_fallback_message(update_chunk) is None + + +class _FakeCheckpointTuple: + """Minimal stand-in for ``CheckpointTuple`` used by ``_ensure_interrupted_title``.""" + + def __init__(self, *, checkpoint: dict, metadata: dict, config: dict | None = None): + self.checkpoint = checkpoint + self.metadata = metadata + self.config = config or {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} + + +class _TitleCheckpointer: + """Captures ``aput`` arguments and exposes ``get_next_version`` like DB savers.""" + + def __init__(self, *, tuple_value: _FakeCheckpointTuple | None, put_result: dict | None = None): + self.aget_tuple = AsyncMock(return_value=tuple_value) + self.aput = AsyncMock(return_value=put_result or {}) + + def get_next_version(self, current, _channel): + if current is None: + return 1 + if isinstance(current, int): + return current + 1 + if isinstance(current, str): + try: + return str(int(current) + 1) + except ValueError: + return f"{current}.1" + return 1 + + +@pytest.mark.anyio +async def test_interrupted_title_finalization_blocks_new_same_thread_run(monkeypatch): + """A cancelled run must remain active while its title-only checkpoint is finalizing.""" + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Old Prompt"}, + ) + + initial_checkpoint = { + "id": "ckpt-old", + "ts": "2026-06-29T00:00:00Z", + "channel_values": {"messages": [{"type": "human", "content": "old prompt"}]}, + "channel_versions": {"messages": 1}, + } + + class _BlockingTitleCheckpointer: + def __init__(self) -> None: + self.latest_checkpoint = copy.deepcopy(initial_checkpoint) + self.latest_metadata = {"source": "loop", "step": 1} + self.title_write_started = asyncio.Event() + self.release_title_write = asyncio.Event() + + async def aget_tuple(self, config): + del config + return _FakeCheckpointTuple( + checkpoint=copy.deepcopy(self.latest_checkpoint), + metadata=dict(self.latest_metadata), + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": self.latest_checkpoint["id"], + } + }, + ) + + async def aput(self, config, checkpoint, metadata, new_versions): + del config, new_versions + self.title_write_started.set() + await self.release_title_write.wait() + self.latest_checkpoint = copy.deepcopy(checkpoint) + self.latest_metadata = dict(metadata) + return { + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": checkpoint["id"], + } + } + + def get_next_version(self, current, _channel): + return (current or 0) + 1 + + run_manager = RunManager() + record = await run_manager.create("thread-1") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + checkpointer = _BlockingTitleCheckpointer() + + class _AbortingAgent: + metadata = {"model_name": "fake-test-model"} + checkpointer: Any | None = None + store: Any | None = None + interrupt_before_nodes = None + interrupt_after_nodes = None + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + record.abort_event.set() + if False: + yield # pragma: no cover + + def factory(*, config): + del config + return _AbortingAgent() + + task = asyncio.create_task( + run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=checkpointer), + agent_factory=factory, + graph_input={"messages": [{"role": "user", "content": "old prompt"}]}, + config={}, + ) + ) + record.task = task + + try: + await asyncio.wait_for(checkpointer.title_write_started.wait(), timeout=1.0) + with pytest.raises(ConflictError, match="active run"): + await run_manager.create_or_reject("thread-1", multitask_strategy="reject") + finally: + checkpointer.release_title_write.set() + await task + + records = await run_manager.list_by_thread("thread-1") + assert [item.run_id for item in records] == [record.run_id] + assert checkpointer.latest_checkpoint["channel_values"]["messages"] == [{"type": "human", "content": "old prompt"}] + + +@pytest.mark.anyio +async def test_finalizing_run_only_blocks_reject_strategy(): + """A finalizing run must not break interrupt/rollback superseding semantics.""" + + async def _seed_finalizing_run(): + run_manager = RunManager() + record = await run_manager.create("thread-1") + release_cleanup = asyncio.Event() + cleanup_cancelled = asyncio.Event() + + async def _cleanup_task(): + try: + await release_cleanup.wait() + except asyncio.CancelledError: + cleanup_cancelled.set() + raise + + task = asyncio.create_task(_cleanup_task()) + record.task = task + await run_manager.set_status(record.run_id, RunStatus.interrupted) + await run_manager.set_finalizing(record.run_id, True) + return run_manager, record, task, release_cleanup, cleanup_cancelled + + for strategy in ("interrupt", "rollback"): + run_manager, record, task, release_cleanup, cleanup_cancelled = await _seed_finalizing_run() + try: + replacement = await run_manager.create_or_reject("thread-1", multitask_strategy=strategy) + await asyncio.sleep(0) + + assert replacement.run_id != record.run_id + assert record.status == RunStatus.interrupted + assert record.finalizing is True + assert not cleanup_cancelled.is_set() + assert not task.done() + finally: + release_cleanup.set() + with suppress(asyncio.CancelledError): + await task + + run_manager, _record, task, release_cleanup, _cleanup_cancelled = await _seed_finalizing_run() + try: + with pytest.raises(ConflictError, match="active run"): + await run_manager.create_or_reject("thread-1", multitask_strategy="reject") + finally: + release_cleanup.set() + with suppress(asyncio.CancelledError): + await task + + +@pytest.mark.anyio +async def test_admitted_pending_replacement_does_not_steal_interrupted_title_recovery(monkeypatch): + """The old run must still write the fallback title before releasing a serialized replacement.""" + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Old Prompt"}, + ) + + initial_checkpoint = { + "id": "ckpt-old", + "ts": "2026-06-29T00:00:00Z", + "channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]}, + "channel_versions": {"messages": 1}, + } + + run_manager = RunManager() + old_record = await run_manager.create("thread-1") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + checkpointer = _TitleCheckpointer( + tuple_value=_FakeCheckpointTuple( + checkpoint=initial_checkpoint, + metadata={"source": "loop", "step": 1}, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": "ckpt-old", + } + }, + ), + ) + + old_title_gate_entered = asyncio.Event() + release_old_title_gate = asyncio.Event() + original_wait_for_prior_finalizing = run_manager.wait_for_prior_finalizing + + async def _wait_for_prior_finalizing(thread_id, run_id, **kwargs): + if run_id == old_record.run_id and old_record.status == RunStatus.interrupted and old_record.finalizing: + old_title_gate_entered.set() + await release_old_title_gate.wait() + return await original_wait_for_prior_finalizing(thread_id, run_id, **kwargs) + + run_manager.wait_for_prior_finalizing = _wait_for_prior_finalizing # type: ignore[method-assign] + + class _AbortingAgent: + metadata = {"model_name": "fake-test-model"} + checkpointer: Any | None = None + store: Any | None = None + interrupt_before_nodes = None + interrupt_after_nodes = None + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + old_record.abort_event.set() + if False: + yield # pragma: no cover + + def factory(*, config): + del config + return _AbortingAgent() + + old_task = asyncio.create_task( + run_agent( + bridge, + run_manager, + old_record, + ctx=RunContext(checkpointer=checkpointer), + agent_factory=factory, + graph_input={"messages": [{"role": "user", "content": "Old prompt"}]}, + config={}, + ) + ) + old_record.task = old_task + + try: + await asyncio.wait_for(old_title_gate_entered.wait(), timeout=1.0) + replacement_record = await run_manager.create_or_reject("thread-1", multitask_strategy="interrupt") + assert replacement_record.status == RunStatus.pending + + release_old_title_gate.set() + await old_task + finally: + release_old_title_gate.set() + if not old_task.done(): + old_task.cancel() + with suppress(asyncio.CancelledError): + await old_task + + checkpointer.aput.assert_awaited_once() + _, written_checkpoint, _, _ = checkpointer.aput.await_args.args + assert written_checkpoint["channel_values"]["title"] == "Old Prompt" + + +@pytest.mark.anyio +async def test_interrupted_title_does_not_overwrite_checkpoint_from_admitted_replacement(monkeypatch): + """A replacement run admitted by multitask interrupt must not lose its newer checkpoint.""" + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Old prompt"}, + ) + + old_checkpoint = { + "id": "ckpt-old", + "ts": "2026-06-29T00:00:00Z", + "channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]}, + "channel_versions": {"messages": 1}, + } + replacement_messages = [ + {"type": "human", "content": "Old prompt"}, + {"type": "human", "content": "Replacement prompt"}, + ] + replacement_checkpoint = { + "id": "ckpt-replacement", + "ts": "2026-06-29T00:00:01Z", + "channel_values": {"messages": replacement_messages}, + "channel_versions": {"messages": 2}, + } + + class _ReplacementRaceCheckpointer: + def __init__(self) -> None: + self.latest_checkpoint = copy.deepcopy(old_checkpoint) + self.latest_metadata = {"source": "loop", "step": 1} + self.title_write_started = asyncio.Event() + self.replacement_checkpoint_written = asyncio.Event() + + async def aget_tuple(self, config): + del config + return _FakeCheckpointTuple( + checkpoint=copy.deepcopy(self.latest_checkpoint), + metadata=dict(self.latest_metadata), + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": self.latest_checkpoint["id"], + } + }, + ) + + async def aput(self, config, checkpoint, metadata, new_versions): + del config, new_versions + self.title_write_started.set() + await self.replacement_checkpoint_written.wait() + self.latest_checkpoint = copy.deepcopy(checkpoint) + self.latest_metadata = dict(metadata) + return { + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": checkpoint["id"], + } + } + + def get_next_version(self, current, _channel): + return (current or 0) + 1 + + run_manager = RunManager() + old_record = await run_manager.create("thread-1") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + checkpointer = _ReplacementRaceCheckpointer() + old_agent_started = asyncio.Event() + + class _BlockingAgent: + metadata = {"model_name": "fake-test-model"} + checkpointer: Any | None = None + store: Any | None = None + interrupt_before_nodes = None + interrupt_after_nodes = None + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + old_agent_started.set() + while True: + await asyncio.sleep(0.05) + if False: + yield # pragma: no cover + + def factory(*, config): + del config + return _BlockingAgent() + + old_task = asyncio.create_task( + run_agent( + bridge, + run_manager, + old_record, + ctx=RunContext(checkpointer=checkpointer), + agent_factory=factory, + graph_input={"messages": [{"role": "user", "content": "Old prompt"}]}, + config={}, + ) + ) + old_record.task = old_task + + try: + await asyncio.wait_for(old_agent_started.wait(), timeout=1.0) + replacement_record = await run_manager.create_or_reject("thread-1", multitask_strategy="interrupt") + assert replacement_record.run_id != old_record.run_id + await run_manager.set_status(replacement_record.run_id, RunStatus.running) + + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(checkpointer.title_write_started.wait(), timeout=0.25) + + checkpointer.latest_checkpoint = copy.deepcopy(replacement_checkpoint) + checkpointer.latest_metadata = {"source": "loop", "step": 2} + checkpointer.replacement_checkpoint_written.set() + await old_task + finally: + checkpointer.replacement_checkpoint_written.set() + if not old_task.done(): + old_task.cancel() + with suppress(asyncio.CancelledError): + await old_task + + assert checkpointer.latest_checkpoint["channel_values"]["messages"] == replacement_messages + + +@pytest.mark.anyio +async def test_replacement_run_waits_for_prior_finalizing_run(): + """Replacement workers must not enter the graph while an older run is finalizing.""" + run_manager = RunManager() + old_record = await run_manager.create("thread-1") + replacement_record = await run_manager.create("thread-1") + await run_manager.set_finalizing(old_record.run_id, True) + + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + replacement_started = asyncio.Event() + + class _ReplacementAgent: + metadata = {"model_name": "fake-test-model"} + checkpointer: Any | None = None + store: Any | None = None + interrupt_before_nodes = None + interrupt_after_nodes = None + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + replacement_started.set() + if False: + yield # pragma: no cover + + def factory(*, config): + del config + return _ReplacementAgent() + + task = asyncio.create_task( + run_agent( + bridge, + run_manager, + replacement_record, + ctx=RunContext(checkpointer=None), + agent_factory=factory, + graph_input={"messages": [{"role": "user", "content": "Replacement prompt"}]}, + config={}, + ) + ) + replacement_record.task = task + + try: + await asyncio.sleep(0.1) + assert not replacement_started.is_set() + + await run_manager.set_finalizing(old_record.run_id, False) + await asyncio.wait_for(replacement_started.wait(), timeout=1.0) + await task + finally: + await run_manager.set_finalizing(old_record.run_id, False) + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_reloads_latest_checkpoint_before_write(): + """If the checkpoint advances before the title write, preserve the newer messages.""" + from deerflow.config.title_config import TitleConfig + + old_checkpoint = { + "id": "ckpt-old", + "ts": "2026-06-29T00:00:00Z", + "channel_values": {"messages": [{"type": "human", "content": "Old prompt"}]}, + "channel_versions": {"messages": 1}, + } + new_messages = [{"type": "human", "content": "New prompt"}] + new_checkpoint = { + "id": "ckpt-new", + "ts": "2026-06-29T00:00:01Z", + "channel_values": {"messages": new_messages}, + "channel_versions": {"messages": 2}, + } + + class _AdvancingTitleCheckpointer: + def __init__(self) -> None: + self.read_count = 0 + self.aput = AsyncMock(return_value={}) + + async def aget_tuple(self, config): + del config + self.read_count += 1 + checkpoint = old_checkpoint if self.read_count == 1 else new_checkpoint + return _FakeCheckpointTuple( + checkpoint=copy.deepcopy(checkpoint), + metadata={"source": "loop", "step": self.read_count}, + config={ + "configurable": { + "thread_id": "thread-1", + "checkpoint_ns": "", + "checkpoint_id": checkpoint["id"], + } + }, + ) + + def get_next_version(self, current, _channel): + return (current or 0) + 1 + + checkpointer = _AdvancingTitleCheckpointer() + app_config = SimpleNamespace(title=TitleConfig(enabled=True, max_chars=40, max_words=20)) + + title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=app_config) + + assert title == "New prompt" + _, written_checkpoint, _, _ = checkpointer.aput.await_args.args + assert written_checkpoint["channel_values"]["messages"] == new_messages + assert written_checkpoint["channel_values"]["title"] == "New prompt" + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_bumps_channel_version_and_declares_it_in_new_versions(monkeypatch): + """Regression for #3859 review: DB-backed savers (Sqlite/Postgres) strip inline + ``channel_values`` from ``put`` and only persist blobs for channels listed in + ``new_versions``. The helper must therefore bump ``channel_versions["title"]`` + and pass ``{"title": next_version}`` so the fallback title actually survives + a fresh ``aget_tuple`` after the worker's finally hook. + """ + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Generated Title"}, + ) + + initial_checkpoint = { + "id": "ckpt-1", + "ts": "2026-06-29T00:00:00Z", + "channel_values": {"messages": [{"type": "human", "content": "hi"}]}, + "channel_versions": {"messages": 5}, + } + checkpointer = _TitleCheckpointer( + tuple_value=_FakeCheckpointTuple( + checkpoint=initial_checkpoint, + metadata={"source": "loop", "step": 7}, + ), + ) + + title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + + assert title == "Generated Title" + checkpointer.aput.assert_awaited_once() + write_config, written_checkpoint, written_metadata, new_versions = checkpointer.aput.await_args.args + + # The title channel must be declared in new_versions — without this, DB + # savers drop the inline channel_values["title"] from the persisted blob. + assert new_versions == {"title": 1} + # Channel versions on the checkpoint itself must also reflect the bump, + # so a subsequent aget_tuple reconstructs channel_values with the title. + assert written_checkpoint["channel_versions"]["title"] == 1 + # Pre-existing channel versions must be preserved. + assert written_checkpoint["channel_versions"]["messages"] == 5 + # The fallback title rides into channel_values for the (legacy / single-table) + # savers that inline the snapshot. + assert written_checkpoint["channel_values"]["title"] == "Generated Title" + assert written_metadata["source"] == "update" + assert written_metadata["step"] == 8 + assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Generated Title"}} + assert write_config == {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_writes_graph_input_fallback_without_checkpoint(monkeypatch): + """When no checkpoint exists, graph_input should still seed the fallback title write.""" + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + captured_state: dict[str, Any] = {} + + def _generate(self, state, allow_partial_exchange=False): + del self + captured_state.update(state) + assert allow_partial_exchange is True + return {"title": "Graph Input Title"} + + monkeypatch.setattr(TitleMiddleware, "_generate_title_result", _generate) + + checkpointer = _TitleCheckpointer(tuple_value=None) + + title = await _ensure_interrupted_title( + checkpointer=checkpointer, + thread_id="thread-1", + app_config=None, + graph_input={"messages": [{"role": "user", "content": "Please name this thread"}]}, + ) + + assert title == "Graph Input Title" + assert captured_state["messages"] == [{"role": "user", "content": "Please name this thread"}] + checkpointer.aput.assert_awaited_once() + _, written_checkpoint, written_metadata, new_versions = checkpointer.aput.await_args.args + assert written_checkpoint["channel_values"]["title"] == "Graph Input Title" + assert written_metadata["writes"] == {"runtime_interrupt_title": {"title": "Graph Input Title"}} + assert set(new_versions.keys()) == {"title"} + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_bumps_existing_string_version(monkeypatch): + """When the checkpointer lacks ``get_next_version`` and the prior title + version is a string (some savers use UUID-shaped versions), the helper must + still produce a strictly different value rather than overwriting in place. + """ + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "T"}, + ) + + initial_checkpoint = { + "id": "ckpt-1", + "ts": "2026-06-29T00:00:00Z", + "channel_values": {"messages": [{"type": "human", "content": "hi"}]}, + "channel_versions": {"title": "v3"}, + } + + class _NoGetNextVersion: + def __init__(self): + self.aget_tuple = AsyncMock( + return_value=_FakeCheckpointTuple( + checkpoint=initial_checkpoint, + metadata={}, + ), + ) + self.aput = AsyncMock(return_value={}) + + checkpointer = _NoGetNextVersion() + await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + + _, written_checkpoint, _, new_versions = checkpointer.aput.await_args.args + bumped = written_checkpoint["channel_versions"]["title"] + assert bumped != "v3", "title version must change so DB savers persist the update" + assert new_versions == {"title": bumped} + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_skips_when_title_already_set(): + """If the checkpoint already carries a title, no new checkpoint is written.""" + checkpointer = _TitleCheckpointer( + tuple_value=_FakeCheckpointTuple( + checkpoint={ + "id": "ckpt-1", + "channel_values": {"messages": [], "title": "Already there"}, + "channel_versions": {"title": 1}, + }, + metadata={}, + ), + ) + + title = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + + assert title == "Already there" + checkpointer.aput.assert_not_awaited() + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_round_trip_with_real_sqlite_checkpointer(tmp_path): + """Full round-trip against a real ``AsyncSqliteSaver`` on a disk-backed DB. + + Mirrors what Gateway constructs in production via ``make_checkpointer`` when + ``database.backend == "sqlite"``, then closes and re-opens the saver to + simulate a fresh connection. The fallback title must survive that boundary — + this is the scenario the #3874 review flagged as broken before the + ``new_versions={"title": ...}`` fix. + """ + from langchain_core.messages import HumanMessage + from langgraph.checkpoint.base import empty_checkpoint + from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver + + from deerflow.config.title_config import TitleConfig + + db_path = str(tmp_path / "ckpt.db") + thread_cfg = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} + + # 1. Seed a first-turn checkpoint that has a human message and NO title — + # the same shape the agent leaves behind when interrupted mid-stream. + async with AsyncSqliteSaver.from_conn_string(db_path) as writer: + await writer.setup() + ck = empty_checkpoint() + ck["channel_values"] = { + "messages": [HumanMessage(content="Why is the sky blue?").model_dump()], + } + ck["channel_versions"] = {"messages": 1} + await writer.aput(thread_cfg, ck, {"source": "loop", "step": 1, "writes": {}}, {"messages": 1}) + + # 2. Run the worker helper through a *fresh* saver instance — this is what + # the lifespan-owned checkpointer pool does for each request. + title_config = TitleConfig(enabled=True, max_chars=40, max_words=20) + app_config = SimpleNamespace(title=title_config) + async with AsyncSqliteSaver.from_conn_string(db_path) as worker_saver: + title = await _ensure_interrupted_title( + checkpointer=worker_saver, + thread_id="thread-1", + app_config=app_config, + ) + assert title, "fallback title must be generated from the seeded user message" + + # 3. Open ANOTHER fresh saver and confirm the title survives — this is the + # invariant the #3874 review was guarding: ``new_versions={}`` would + # cause DB savers to drop the title blob, so a fresh aget_tuple would + # read back without it. + async with AsyncSqliteSaver.from_conn_string(db_path) as reader: + tup = await reader.aget_tuple(thread_cfg) + assert tup is not None + persisted = tup.checkpoint.get("channel_values", {}).get("title") + assert persisted == title + + +# --------------------------------------------------------------------------- +# _bump_channel_version — invariant: the returned version MUST differ from +# the prior value, no matter the checkpointer's versioning scheme. +# --------------------------------------------------------------------------- + + +class _CheckpointerWithIntVersion: + """A checkpointer whose ``get_next_version`` increments integers (default LangGraph behavior).""" + + @staticmethod + def get_next_version(current, _channel): + return (current or 0) + 1 + + +class _CheckpointerWithBrokenGetNextVersion: + """A checkpointer whose ``get_next_version`` raises — must fall back, not propagate.""" + + @staticmethod + def get_next_version(current, _channel): + raise RuntimeError("simulated saver bug") + + +def test_bump_channel_version_uses_checkpointer_get_next_version_when_available(): + """Happy path — saver's ``get_next_version`` result is preferred over our fallback.""" + assert _bump_channel_version(_CheckpointerWithIntVersion(), 5) == 6 + + +def test_bump_channel_version_falls_back_on_broken_get_next_version(): + """A raising ``get_next_version`` must not propagate; the defensive path bumps from prior. + + Without this, a saver bug would leave ``new_versions={"title": v}`` no-op + on DB savers — the very class of bug the #3874 review flagged. + """ + bumped = _bump_channel_version(_CheckpointerWithBrokenGetNextVersion(), 7) + assert bumped == 8 + + +# --------------------------------------------------------------------------- +# _ensure_interrupted_title — additional defensive boundaries +# --------------------------------------------------------------------------- + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_handles_none_messages_channel(monkeypatch): + """A partially-initialized checkpoint with ``messages=None`` must not crash.""" + from deerflow.config.title_config import TitleConfig + + initial_checkpoint = { + "id": "ckpt-1", + "channel_values": {"messages": None}, + "channel_versions": {}, + } + checkpointer = _TitleCheckpointer( + tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}), + ) + app_config = SimpleNamespace(title=TitleConfig(enabled=True, max_chars=40, max_words=20)) + + assert await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=app_config) is None + checkpointer.aput.assert_not_awaited() + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_propagates_aput_error_to_caller(monkeypatch): + """Exceptions from ``aput`` propagate — the caller (worker.run_agent finally block) is responsible for swallowing them. + + This test pins the contract: the helper itself does NOT silently eat saver errors, + so a structural saver regression remains visible in the logs at the call site. + """ + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Generated"}, + ) + + initial_checkpoint = { + "id": "ckpt-1", + "channel_values": {"messages": [{"type": "human", "content": "hi"}]}, + "channel_versions": {"messages": 1}, + } + checkpointer = _TitleCheckpointer( + tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}), + ) + checkpointer.aput.side_effect = RuntimeError("simulated DB write failure") + + with pytest.raises(RuntimeError, match="simulated DB write failure"): + await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_idempotent_across_repeated_calls(monkeypatch): + """Second invocation against the now-titled checkpoint must not re-write. + + Regression anchor for the case where a brittle helper might re-trigger + on subsequent finally-hook runs (e.g. retries) and rewrite the title. + """ + from langgraph.checkpoint.memory import InMemorySaver + + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "First Title"}, + ) + + checkpointer = InMemorySaver() + cfg = {"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}} + ck = empty_checkpoint() + ck["channel_values"] = {"messages": [{"type": "human", "content": "hi"}]} + ck["channel_versions"] = {"messages": 1} + await checkpointer.aput(cfg, ck, {"source": "loop", "step": 1, "writes": {}}, {"messages": 1}) + + first = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + assert first == "First Title" + + # Second call: title is now present, so the helper short-circuits without + # rewriting — even if the middleware were to suggest a different title. + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Different Title"}, + ) + second = await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + assert second == "First Title" + + tup = await checkpointer.aget_tuple(cfg) + assert tup.checkpoint["channel_values"]["title"] == "First Title" + + +@pytest.mark.anyio +async def test_ensure_interrupted_title_preserves_non_title_channel_versions(monkeypatch): + """Bumping ``channel_versions["title"]`` must not modify other channels' versions. + + Regression anchor: an earlier draft built ``new_versions`` from + ``dict(channel_versions)`` and would have erroneously declared every + channel as "needs new blob" on DB savers. + """ + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + monkeypatch.setattr( + TitleMiddleware, + "_generate_title_result", + lambda self, state, allow_partial_exchange=False: {"title": "Generated"}, + ) + + initial_checkpoint = { + "id": "ckpt-1", + "channel_values": { + "messages": [{"type": "human", "content": "hi"}], + "artifacts": [], + "todos": None, + }, + "channel_versions": {"messages": 5, "artifacts": 3, "todos": 1}, + } + checkpointer = _TitleCheckpointer( + tuple_value=_FakeCheckpointTuple(checkpoint=initial_checkpoint, metadata={}), + ) + + await _ensure_interrupted_title(checkpointer=checkpointer, thread_id="thread-1", app_config=None) + + _, written_checkpoint, _, new_versions = checkpointer.aput.await_args.args + # Only the title channel is declared in new_versions. + assert set(new_versions.keys()) == {"title"} + # Other channel versions are preserved verbatim on the written checkpoint. + assert written_checkpoint["channel_versions"]["messages"] == 5 + assert written_checkpoint["channel_versions"]["artifacts"] == 3 + assert written_checkpoint["channel_versions"]["todos"] == 1 + + +@pytest.mark.anyio +async def test_worker_finally_block_swallows_helper_exceptions(monkeypatch): + """The worker's interrupted-title hook must remain non-fatal — any exception + from the helper (DB saver bug, middleware bug, etc.) must not propagate past + the run boundary or prevent the subsequent threads_meta sync block from + running. This pins the integration of helper + finally try/except, not just + the helper itself. + """ + import deerflow.runtime.runs.worker as worker_module + + async def _boom(*_args, **_kwargs): + raise RuntimeError("forced helper failure") + + monkeypatch.setattr(worker_module, "_ensure_interrupted_title", _boom) + + run_manager = RunManager() + record = await run_manager.create("thread-1") + record.status = RunStatus.interrupted + + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + + class _MinimalCheckpointer: + async def aget_tuple(self, config): + return None + + async def aput(self, *args, **kwargs): + return {} + + captured_status: dict[str, Any] = {} + + class _ThreadStore: + async def update_display_name(self, thread_id, title): + captured_status["display_name"] = (thread_id, title) + + async def update_status(self, thread_id, status): + captured_status["status"] = (thread_id, status) + + class _AbortingAgent: + def __init__(self) -> None: + self.metadata = {"model_name": "fake-test-model"} + self.checkpointer: Any | None = None + self.store: Any | None = None + self.interrupt_before_nodes = None + self.interrupt_after_nodes = None + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + # Abort immediately so the run lands in the interrupted branch. + record.abort_event.set() + if False: + yield # pragma: no cover — make this an async generator + return + + def factory(*, config): + del config + return _AbortingAgent() + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=_MinimalCheckpointer(), thread_store=_ThreadStore()), + agent_factory=factory, + graph_input={"messages": []}, + config={}, + ) + + # The helper raised, but the run still reaches the threads_meta status sync + # and ``publish_end`` — i.e. the SSE stream is closed cleanly and the row + # reflects the run outcome. + assert captured_status.get("status") == ("thread-1", "interrupted") + bridge.publish_end.assert_awaited_once_with(record.run_id) diff --git a/backend/tests/test_runtime_lifecycle_e2e.py b/backend/tests/test_runtime_lifecycle_e2e.py index be3fb506e..508cc69d8 100644 --- a/backend/tests/test_runtime_lifecycle_e2e.py +++ b/backend/tests/test_runtime_lifecycle_e2e.py @@ -80,11 +80,15 @@ class _ScriptedAgent: title: str, answer: str, block_after_first_chunk: bool = False, + block_before_checkpoint: bool = False, + write_title: bool = True, ) -> None: self.controller = controller self.title = title self.answer = answer self.block_after_first_chunk = block_after_first_chunk + self.block_before_checkpoint = block_before_checkpoint + self.write_title = write_title self.checkpointer: Any | None = None self.store: Any | None = None self.metadata = {"model_name": "fake-test-model"} @@ -98,10 +102,15 @@ class _ScriptedAgent: try: thread_id = _thread_id_from_config(config) + if self.block_before_checkpoint: + while not self.controller.release.is_set(): + await asyncio.sleep(0.05) human_text = _last_human_text(graph_input) human = HumanMessage(content=human_text) ai = await self.model.ainvoke([human], config=config) - state = {"messages": [human.model_dump(), ai.model_dump()], "title": self.title} + state = {"messages": [human.model_dump(), ai.model_dump()]} + if self.write_title: + state["title"] = self.title if self.checkpointer is not None: await _write_checkpoint(self.checkpointer, thread_id=thread_id, state=state) @@ -265,6 +274,25 @@ def isolated_app(isolated_deer_flow_home: Path, monkeypatch: pytest.MonkeyPatch) return create_app() +@pytest.fixture +def isolated_app_with_title(isolated_deer_flow_home: Path, monkeypatch: pytest.MonkeyPatch): + config_path = isolated_deer_flow_home.parent / "config-title-enabled.yaml" + config_path.write_text(_MINIMAL_CONFIG_YAML.replace("title:\n enabled: false", "title:\n enabled: true"), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_CONFIG_PATH", str(config_path)) + + _preserve_process_config_singletons(monkeypatch) + _reset_process_singletons(monkeypatch) + + from deerflow.config import app_config as app_config_module + + cfg = app_config_module.get_app_config() + cfg.database.sqlite_dir = str(isolated_deer_flow_home / "db") + + from app.gateway.app import create_app + + return create_app() + + def _register_user(client, *, email: str = "runtime-e2e@example.com") -> str: response = client.post( "/api/v1/auth/register", @@ -379,6 +407,34 @@ def _wait_for_status(client, thread_id: str, run_id: str, status: str, *, timeou raise AssertionError(f"Run {run_id} did not reach {status!r}; last={last!r}") +def _wait_for_thread_title(client, thread_id: str, expected_title: str, *, timeout: float = 5.0) -> dict: + deadline = time.monotonic() + timeout + last: dict | None = None + while time.monotonic() < deadline: + response = client.get(f"/api/threads/{thread_id}") + assert response.status_code == 200, response.text + last = response.json() + if last.get("values", {}).get("title") == expected_title: + return last + time.sleep(0.05) + raise AssertionError(f"Thread {thread_id} did not reach title {expected_title!r}; last={last!r}") + + +def _wait_for_search_title(client, csrf_token: str, thread_id: str, expected_title: str, *, timeout: float = 5.0) -> dict: + deadline = time.monotonic() + timeout + last_match: dict | None = None + while time.monotonic() < deadline: + response = client.post("/api/threads/search", json={"limit": 20}, headers={"X-CSRF-Token": csrf_token}) + assert response.status_code == 200, response.text + matching = [item for item in response.json() if item["thread_id"] == thread_id] + if matching: + last_match = matching[0] + if last_match.get("values", {}).get("title") == expected_title: + return last_match + time.sleep(0.05) + raise AssertionError(f"Search result for {thread_id} did not reach title {expected_title!r}; last={last_match!r}") + + def _thread_id_from_config(config: dict | None) -> str: config = config or {} context = config.get("context") if isinstance(config.get("context"), dict) else {} @@ -581,6 +637,99 @@ def test_cancel_interrupt_stops_running_background_run(isolated_app): assert thread.json()["status"] == "idle" +def test_cancel_interrupt_generates_missing_title_from_checkpoint(isolated_app_with_title): + """Interrupted first-turn runs should still persist an automatic thread title.""" + from starlette.testclient import TestClient + + controller = _RunController() + factory = _make_agent_factory( + controller, + title="", + answer="This run should be interrupted before a title is written.", + block_after_first_chunk=True, + write_title=False, + ) + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=factory), + TestClient(isolated_app_with_title) as client, + ): + csrf_token = _register_user(client, email="interrupt-title-e2e@example.com") + thread_id = _create_thread(client, csrf_token) + + created = client.post( + f"/api/threads/{thread_id}/runs", + json=_run_body(), + headers={"X-CSRF-Token": csrf_token}, + ) + assert created.status_code == 200, created.text + run_id = created.json()["run_id"] + assert controller.checkpoint_written.wait(5), "fake agent never wrote checkpoint" + + cancelled = client.post( + f"/api/threads/{thread_id}/runs/{run_id}/cancel?wait=true&action=interrupt", + headers={"X-CSRF-Token": csrf_token}, + ) + assert cancelled.status_code == 204, cancelled.text + + thread = client.get(f"/api/threads/{thread_id}") + assert thread.status_code == 200, thread.text + assert thread.json()["values"]["title"] == "Run lifecycle E2E prompt" + + search = client.post("/api/threads/search", json={"limit": 20}, headers={"X-CSRF-Token": csrf_token}) + assert search.status_code == 200, search.text + matching = [item for item in search.json() if item["thread_id"] == thread_id] + assert matching[0]["values"]["title"] == "Run lifecycle E2E prompt" + + +def test_cancel_wait_false_generates_title_from_graph_input_before_checkpoint(isolated_app_with_title): + """Fire-and-forget cancel should title early interruptions before checkpoint.""" + from starlette.testclient import TestClient + + controller = _RunController() + factory = _make_agent_factory( + controller, + title="", + answer="This answer should never be checkpointed.", + block_before_checkpoint=True, + write_title=False, + ) + + with ( + patch("app.gateway.services.resolve_agent_factory", return_value=factory), + TestClient(isolated_app_with_title) as client, + ): + csrf_token = _register_user(client, email="interrupt-title-early-e2e@example.com") + thread_id = _create_thread(client, csrf_token) + + created = client.post( + f"/api/threads/{thread_id}/runs", + json=_run_body(), + headers={"X-CSRF-Token": csrf_token}, + ) + assert created.status_code == 200, created.text + run_id = created.json()["run_id"] + assert controller.started.wait(5), "fake agent never started" + assert not controller.checkpoint_written.is_set() + + cancelled = client.post( + f"/api/threads/{thread_id}/runs/{run_id}/cancel?wait=false&action=interrupt", + headers={"X-CSRF-Token": csrf_token}, + ) + assert cancelled.status_code == 202, cancelled.text + assert controller.cancelled.wait(5), "fake agent task was not cancelled" + assert not controller.checkpoint_written.is_set() + + run = _wait_for_status(client, thread_id, run_id, "interrupted") + assert run["status"] == "interrupted" + + thread = _wait_for_thread_title(client, thread_id, "Run lifecycle E2E prompt") + assert thread["values"]["title"] == "Run lifecycle E2E prompt" + + matching = _wait_for_search_title(client, csrf_token, thread_id, "Run lifecycle E2E prompt") + assert matching["values"]["title"] == "Run lifecycle E2E prompt" + + @pytest.mark.anyio async def test_sse_consumer_disconnect_cancels_inflight_run(): """A disconnected SSE request should cancel an in-flight run when configured.""" diff --git a/backend/tests/test_title_middleware_core_logic.py b/backend/tests/test_title_middleware_core_logic.py index bb2bc7083..8f99f2054 100644 --- a/backend/tests/test_title_middleware_core_logic.py +++ b/backend/tests/test_title_middleware_core_logic.py @@ -367,6 +367,113 @@ class TestTitleMiddlewareCoreLogic: assert "" not in prompt assert "User prefers Python" not in prompt + def test_should_generate_title_partial_exchange_allows_user_only(self): + """Interrupted-run path can produce a fallback from a lone human message.""" + _set_test_title_config(enabled=True) + middleware = TitleMiddleware() + state = {"messages": [HumanMessage(content="只有人类消息，AI 还没回复")]} + + assert middleware._should_generate_title(state) is False + assert middleware._should_generate_title(state, allow_partial_exchange=True) is True + + def test_should_generate_title_partial_exchange_skips_when_titled(self): + """Existing title still wins, even on the interrupted-run path.""" + _set_test_title_config(enabled=True) + middleware = TitleMiddleware() + state = { + "messages": [HumanMessage(content="问题")], + "title": "Already set", + } + assert middleware._should_generate_title(state, allow_partial_exchange=True) is False + + def test_should_generate_title_handles_dict_messages(self): + """Checkpoint channel_values store messages as dicts; the middleware must accept them.""" + _set_test_title_config(enabled=True) + middleware = TitleMiddleware() + state = { + "messages": [ + {"type": "human", "content": "问"}, + {"type": "ai", "content": "答"}, + ] + } + assert middleware._should_generate_title(state) is True + + def test_sync_generate_title_from_dict_messages(self): + """Sync fallback path can derive title text from dict-form messages.""" + _set_test_title_config(max_chars=20) + middleware = TitleMiddleware() + state = { + "messages": [ + {"role": "user", "content": "请帮我写测试"}, + {"role": "assistant", "content": "好的"}, + ] + } + assert middleware._generate_title_result(state) == {"title": "请帮我写测试"} + + def test_should_generate_title_handles_none_messages_channel(self): + """A checkpoint with ``messages=None`` (partially-initialized state) must not crash.""" + _set_test_title_config(enabled=True) + middleware = TitleMiddleware() + # ``messages`` key exists but is None — ``state.get("messages", [])`` would + # have returned ``None`` (default only applies on missing key), so this + # exercises the ``or []`` coercion the helper relies on. + state = {"messages": None} + + assert middleware._should_generate_title(state) is False + assert middleware._should_generate_title(state, allow_partial_exchange=True) is False + + def test_build_title_prompt_handles_none_messages_channel(self): + """``_build_title_prompt`` must also tolerate a None messages channel.""" + _set_test_title_config(enabled=True) + middleware = TitleMiddleware() + state = {"messages": None} + + prompt, user_msg = middleware._build_title_prompt(state) + assert user_msg == "" + # Prompt is still well-formed — just empty user/assistant slots. + assert "{user_msg}" not in prompt # the template was formatted, not left raw + + def test_should_generate_title_dict_messages_role_normalization(self): + """Dict-form messages may use either ``type`` or ``role``; both must map correctly.""" + _set_test_title_config(enabled=True) + middleware = TitleMiddleware() + state = { + "messages": [ + # ``role: user`` should be normalized to ``human`` + {"role": "user", "content": "Q"}, + # ``role: assistant`` should be normalized to ``ai`` + {"role": "assistant", "content": "A"}, + ] + } + assert middleware._should_generate_title(state) is True + + def test_partial_exchange_with_dict_human_message(self): + """Interrupted-run path must accept a lone dict-form first-turn user message.""" + _set_test_title_config(enabled=True, max_chars=20) + middleware = TitleMiddleware() + state = {"messages": [{"role": "user", "content": "请帮我写测试"}]} + + result = middleware._generate_title_result(state, allow_partial_exchange=True) + assert result == {"title": "请帮我写测试"} + + def test_partial_exchange_ignores_dict_dynamic_context_reminder(self): + """Checkpoint dicts can include hidden memory reminders that should not count as real user turns.""" + _set_test_title_config(enabled=True, max_chars=20) + middleware = TitleMiddleware() + state = { + "messages": [ + { + "type": "human", + "content": "User prefers concise titles.", + "additional_kwargs": {"hide_from_ui": True, _DYNAMIC_CONTEXT_REMINDER_KEY: True}, + }, + {"type": "human", "content": "请帮我写测试", "additional_kwargs": {}}, + ] + } + + assert middleware._should_generate_title(state, allow_partial_exchange=True) is True + assert middleware._generate_title_result(state, allow_partial_exchange=True) == {"title": "请帮我写测试"} + def test_generate_title_async_strips_think_tags_in_response(self, monkeypatch): """Async title generation strips blocks from the model response.""" _set_test_title_config(max_chars=50, model_name="title-model") diff --git a/backend/tests/test_worker_langfuse_metadata.py b/backend/tests/test_worker_langfuse_metadata.py index 7b7544771..63a2edf95 100644 --- a/backend/tests/test_worker_langfuse_metadata.py +++ b/backend/tests/test_worker_langfuse_metadata.py @@ -36,6 +36,15 @@ class _FakeAgent: class _FakeRunManager: + async def wait_for_prior_finalizing(self, *_args, **_kwargs) -> None: + return None + + async def has_later_run(self, *_args, **_kwargs) -> bool: + return False + + async def has_later_started_run(self, *_args, **_kwargs) -> bool: + return False + async def set_status(self, *_args, **_kwargs) -> None: return None diff --git a/frontend/AGENTS.md b/frontend/AGENTS.md index dfd34e151..85a1d86f8 100644 --- a/frontend/AGENTS.md +++ b/frontend/AGENTS.md @@ -65,8 +65,9 @@ The frontend is a stateful chat application. Users create **threads** (conversat 1. User input → thread hooks (`core/threads/hooks.ts`) → LangGraph SDK streaming 2. Stream events update thread state (messages, artifacts, todos) -3. TanStack Query manages server state; localStorage stores user settings -4. Components subscribe to thread state and render updates +3. Stop actions call the LangGraph SDK stream stop path; `core/threads/hooks.ts` invalidates current-thread, token-usage, and sidebar/search caches immediately and schedules one follow-up refetch because SDK stop may finish via abort + fire-and-forget cancel before backend title finalization commits +4. TanStack Query manages server state; localStorage stores user settings +5. Components subscribe to thread state and render updates ### Key Patterns diff --git a/frontend/src/core/threads/hooks.ts b/frontend/src/core/threads/hooks.ts index 2e9c9df70..2ef614ab8 100644 --- a/frontend/src/core/threads/hooks.ts +++ b/frontend/src/core/threads/hooks.ts @@ -606,6 +606,58 @@ export function upsertThreadInInfiniteCache( ); } +export function invalidateStoppedThreadCaches( + queryClient: QueryClient, + threadId: string | null | undefined, + isMock = false, +) { + void queryClient.invalidateQueries({ queryKey: ["threads", "search"] }); + void queryClient.invalidateQueries({ + queryKey: INFINITE_THREADS_QUERY_KEY_PREFIX, + }); + + if (!threadId || isMock) { + return; + } + + void queryClient.invalidateQueries({ queryKey: ["thread", threadId] }); + void queryClient.invalidateQueries({ + queryKey: ["thread", "metadata", threadId, isMock], + }); + void queryClient.invalidateQueries({ + queryKey: threadTokenUsageQueryKey(threadId), + }); +} + +export const STOP_THREAD_FINALIZATION_REFETCH_DELAY_MS = 1500; + +function scheduleStoppedThreadFinalizationRefetch( + queryClient: QueryClient, + threadId: string | null | undefined, + isMock = false, +) { + if (isMock) { + return; + } + globalThis.setTimeout(() => { + invalidateStoppedThreadCaches(queryClient, threadId, isMock); + }, STOP_THREAD_FINALIZATION_REFETCH_DELAY_MS); +} + +export async function stopThreadAndInvalidateCaches( + queryClient: QueryClient, + stop: () => Promise | void, + threadId: string | null | undefined, + isMock = false, +) { + try { + await stop(); + } finally { + invalidateStoppedThreadCaches(queryClient, threadId, isMock); + scheduleStoppedThreadFinalizationRefetch(queryClient, threadId, isMock); + } +} + function getStreamErrorMessage(error: unknown): string { if (typeof error === "string" && error.trim()) { return error; @@ -958,21 +1010,21 @@ export function useThreadStream({ .map(messageIdentity) .filter((id): id is string => Boolean(id)), ); - void queryClient.invalidateQueries({ queryKey: ["threads", "search"] }); - void queryClient.invalidateQueries({ - queryKey: INFINITE_THREADS_QUERY_KEY_PREFIX, - }); - if (threadIdRef.current && !isMock) { - void queryClient.invalidateQueries({ - queryKey: ["thread", threadIdRef.current], - }); - void queryClient.invalidateQueries({ - queryKey: threadTokenUsageQueryKey(threadIdRef.current), - }); - } + invalidateStoppedThreadCaches(queryClient, threadIdRef.current, isMock); }, }); + const stopThread = useCallback(async () => { + const stoppedThreadId = + threadIdRef.current ?? displayThreadId ?? threadId ?? null; + await stopThreadAndInvalidateCaches( + queryClient, + () => thread.stop(), + stoppedThreadId, + isMock, + ); + }, [displayThreadId, isMock, queryClient, thread, threadId]); + const hasVisibleStreamState = Boolean(threadId) || liveMessagesThreadId === currentViewThreadId; const persistedMessages = useMemo( @@ -1443,6 +1495,7 @@ export function useThreadStream({ // History messages may overlap with thread.messages; thread.messages take precedence const mergedThread = { ...thread, + stop: stopThread, values: hasVisibleStreamState ? thread.values : EMPTY_THREAD_VALUES, messages: mergedMessages, } as typeof thread; diff --git a/frontend/tests/unit/core/threads/infinite.test.ts b/frontend/tests/unit/core/threads/infinite.test.ts index 52bd6224f..0ac0fee3a 100644 --- a/frontend/tests/unit/core/threads/infinite.test.ts +++ b/frontend/tests/unit/core/threads/infinite.test.ts @@ -1,12 +1,19 @@ -import { describe, expect, test } from "@rstest/core"; -import { QueryClient, type InfiniteData } from "@tanstack/react-query"; +import { describe, expect, rs, test } from "@rstest/core"; +import { + QueryClient, + QueryObserver, + type InfiniteData, +} from "@tanstack/react-query"; import { filterInfiniteThreadsCache, getInfiniteThreadsNextPageParam, INFINITE_THREADS_PAGE_SIZE, INFINITE_THREADS_QUERY_KEY_PREFIX, + invalidateStoppedThreadCaches, mapInfiniteThreadsCache, + STOP_THREAD_FINALIZATION_REFETCH_DELAY_MS, + stopThreadAndInvalidateCaches, upsertThreadInInfiniteCache, } from "@/core/threads/hooks"; import type { AgentThread } from "@/core/threads/types"; @@ -226,3 +233,183 @@ describe("upsertThreadInInfiniteCache", () => { expect(cache?.pages[0]?.[0]?.values.title).toBe("Old title"); }); }); + +describe("invalidateStoppedThreadCaches", () => { + function invalidatedQueryKeys(client: QueryClient) { + const invalidate = rs.spyOn(client, "invalidateQueries"); + return { + invalidate, + queryKeys: () => + invalidate.mock.calls.map(([filters]) => filters?.queryKey), + }; + } + + test("refreshes current thread and sidebar caches after fire-and-forget stop", () => { + const client = new QueryClient(); + const { queryKeys } = invalidatedQueryKeys(client); + + invalidateStoppedThreadCaches(client, "thread-1", false); + + expect(queryKeys()).toContainEqual(["threads", "search"]); + expect(queryKeys()).toContainEqual(INFINITE_THREADS_QUERY_KEY_PREFIX); + expect(queryKeys()).toContainEqual(["thread", "thread-1"]); + expect(queryKeys()).toContainEqual([ + "thread", + "metadata", + "thread-1", + false, + ]); + expect(queryKeys()).toContainEqual(["thread-token-usage", "thread-1"]); + }); + + test("does not refresh per-thread API caches for mock threads", () => { + const client = new QueryClient(); + const { queryKeys } = invalidatedQueryKeys(client); + + invalidateStoppedThreadCaches(client, "thread-1", true); + + expect(queryKeys()).toContainEqual(["threads", "search"]); + expect(queryKeys()).toContainEqual(INFINITE_THREADS_QUERY_KEY_PREFIX); + expect(queryKeys()).not.toContainEqual(["thread", "thread-1"]); + expect(queryKeys()).not.toContainEqual([ + "thread", + "metadata", + "thread-1", + true, + ]); + expect(queryKeys()).not.toContainEqual(["thread-token-usage", "thread-1"]); + }); + + test("wraps SDK stop and refreshes caches after it resolves", async () => { + const client = new QueryClient(); + const stop = rs.fn(() => Promise.resolve()); + const { queryKeys } = invalidatedQueryKeys(client); + + await stopThreadAndInvalidateCaches(client, stop, "thread-1", false); + + expect(stop).toHaveBeenCalledTimes(1); + expect(queryKeys()).toContainEqual([ + "thread", + "metadata", + "thread-1", + false, + ]); + }); + + test("still refreshes caches when SDK stop rejects", async () => { + const client = new QueryClient(); + const stop = rs.fn(async () => { + throw new Error("cancel failed"); + }); + const { queryKeys } = invalidatedQueryKeys(client); + + await expect( + stopThreadAndInvalidateCaches(client, stop, "thread-1", false), + ).rejects.toThrow("cancel failed"); + + expect(queryKeys()).toContainEqual(["threads", "search"]); + expect(queryKeys()).toContainEqual([ + "thread", + "metadata", + "thread-1", + false, + ]); + }); + + test("schedules sidebar refetch even if stopped thread id is not known", async () => { + rs.useFakeTimers(); + + const client = new QueryClient(); + const { queryKeys } = invalidatedQueryKeys(client); + + try { + await stopThreadAndInvalidateCaches( + client, + () => Promise.resolve(), + null, + false, + ); + + const countSearchInvalidations = () => + queryKeys().filter( + (queryKey) => + queryKey?.length === 2 && + queryKey[0] === "threads" && + queryKey[1] === "search", + ).length; + + expect(countSearchInvalidations()).toBe(1); + + await rs.advanceTimersByTimeAsync( + STOP_THREAD_FINALIZATION_REFETCH_DELAY_MS, + ); + + expect(countSearchInvalidations()).toBe(2); + expect(queryKeys()).not.toContainEqual(["thread", null]); + } finally { + client.clear(); + rs.useRealTimers(); + } + }); + + test("scheduled refetch lets sidebar receive delayed backend title finalization", async () => { + rs.useFakeTimers(); + + const client = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + let finalized = false; + let fetchCount = 0; + const observer = new QueryObserver(client, { + queryKey: ["threads", "search"], + queryFn: async () => { + fetchCount += 1; + return [ + makeThread( + "thread-1", + finalized ? "Generated Title" : "New Conversation", + ), + ]; + }, + }); + const unsubscribe = observer.subscribe((result) => { + void result.status; + }); + + try { + await observer.refetch(); + expect( + client.getQueryData(["threads", "search"])?.[0]?.values + ?.title, + ).toBe("New Conversation"); + + await stopThreadAndInvalidateCaches( + client, + () => Promise.resolve(), + "thread-1", + false, + ); + await Promise.resolve(); + + expect( + client.getQueryData(["threads", "search"])?.[0]?.values + ?.title, + ).toBe("New Conversation"); + + finalized = true; + await rs.advanceTimersByTimeAsync( + STOP_THREAD_FINALIZATION_REFETCH_DELAY_MS, + ); + + expect( + client.getQueryData(["threads", "search"])?.[0]?.values + ?.title, + ).toBe("Generated Title"); + expect(fetchCount).toBeGreaterThanOrEqual(3); + } finally { + unsubscribe(); + client.clear(); + rs.useRealTimers(); + } + }); +});