"""Background agent execution. Runs an agent graph inside an ``asyncio.Task``, publishing events to a :class:`StreamBridge` as they are produced. Uses ``graph.astream(stream_mode=[...])`` which gives correct full-state snapshots for ``values`` mode, proper ``{node: writes}`` for ``updates``, and ``(chunk, metadata)`` tuples for ``messages`` mode. Note: ``events`` mode is not supported through the gateway — it requires ``graph.astream_events()`` which cannot simultaneously produce ``values`` snapshots. The JS open-source LangGraph API server works around this via internal checkpoint callbacks that are not exposed in the Python public API. """ from __future__ import annotations import asyncio import copy import inspect import logging from collections.abc import Awaitable from dataclasses import dataclass, field from functools import lru_cache from time import perf_counter from typing import TYPE_CHECKING, Any, Literal, cast from langgraph.checkpoint.base import empty_checkpoint from langgraph.errors import GraphRecursionError if TYPE_CHECKING: from langchain_core.messages import HumanMessage from deerflow.agents.middlewares.title_middleware import ( finalize_thread_title, log_title_debug, ) from deerflow.config.app_config import AppConfig from deerflow.runtime.serialization import serialize from deerflow.runtime.stream_bridge import StreamBridge from deerflow.utils.background_init import run_background_init from .manager import RunManager, RunRecord from .schemas import RunStatus logger = logging.getLogger(__name__) # Valid stream_mode values for LangGraph's graph.astream() _VALID_LG_MODES = {"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"} # This budget applies only after graph execution has returned. Long-running # tools, research and subagents retain their own execution/cancellation policy. _RUN_FINALIZATION_TIMEOUT_SECONDS = 10.0 _PENDING_COMPLETION_TASKS: set[asyncio.Future[Any]] = set() _BACKGROUND_TITLE_TASKS: set[asyncio.Task[Any]] = set() def _log_background_title_finalize_failure(thread_id: str): def _done(task: asyncio.Task[Any]) -> None: _BACKGROUND_TITLE_TASKS.discard(task) try: task.result() except Exception: logger.debug("Failed to finalize background title for thread %s", thread_id, exc_info=True) return _done async def _finalize_title_and_clear_provisional_marker( *, checkpointer: Any, thread_store: Any, thread_id: str, source_state: dict[str, Any], app_config: AppConfig | None ) -> bool: """Finalize the title and always clear the cheap metadata polling marker.""" try: return await finalize_thread_title( checkpointer=checkpointer, thread_store=thread_store, thread_id=thread_id, source_state=source_state, app_config=app_config, ) finally: try: await thread_store.update_metadata( thread_id, {"_deerflow_title_provisional": False}, user_id=None ) except Exception: logger.debug("Failed to clear provisional-title marker for thread %s", thread_id, exc_info=True) async def _sync_thread_title_after_run( *, checkpointer: Any, thread_store: Any, thread_id: str, run_manager: RunManager, record: RunRecord, app_config: AppConfig | None, ) -> None: """Read title state and generate a final title outside run completion.""" try: if not await _completion_is_current(run_manager, record): return ckpt_tuple = await checkpointer.aget_tuple( {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ) if ckpt_tuple is None or not await _completion_is_current(run_manager, record): return checkpoint = getattr(ckpt_tuple, "checkpoint", {}) or {} channel_values = (checkpoint.get("channel_values", {}) or {}) if isinstance(checkpoint, dict) else {} title = channel_values.get("title") if title: await thread_store.update_display_name(thread_id, title) provisional = channel_values.get("title_provisional") is True if provisional: await thread_store.update_metadata( thread_id, {"_deerflow_title_provisional": record.status == RunStatus.success}, ) if record.status == RunStatus.success and provisional: log_title_debug( "Run %s: scheduling background title generation for thread %s", record.run_id, thread_id, ) await _finalize_title_and_clear_provisional_marker( checkpointer=checkpointer, thread_store=thread_store, thread_id=thread_id, source_state=copy.deepcopy(channel_values), app_config=app_config, ) elif record.status == RunStatus.success: log_title_debug( "Run %s: title generation not needed for thread %s (title_provisional=%r title=%r)", record.run_id, thread_id, channel_values.get("title_provisional"), title, ) except asyncio.CancelledError: raise except Exception: logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id, exc_info=True) def _build_runtime_context( thread_id: str, run_id: str, caller_context: Any | None, app_config: AppConfig | None = None, ) -> dict[str, Any]: """Build the dict that becomes ``ToolRuntime.context`` for the run. Always includes ``thread_id`` and ``run_id``. Additional keys from the caller's ``config['context']`` (e.g. ``agent_name`` for the bootstrap flow — issue #2677) are merged in but never override ``thread_id``/``run_id``. The resolved ``AppConfig`` is added by the worker so tools can consume it without ambient global lookups. langgraph 1.1+ surfaces this as ``runtime.context`` via the parent runtime stored under ``config['configurable']['__pregel_runtime']`` — see ``langgraph.pregel.main`` where ``parent_runtime.merge(...)`` is invoked. """ runtime_ctx: dict[str, Any] = {"thread_id": thread_id, "run_id": run_id} if isinstance(caller_context, dict): for key, value in caller_context.items(): runtime_ctx.setdefault(key, value) if app_config is not None: runtime_ctx["app_config"] = app_config return runtime_ctx @dataclass(frozen=True) class RunContext: """Infrastructure dependencies for a single agent run. Groups checkpointer, store, and persistence-related singletons so that ``run_agent`` (and any future callers) receive one object instead of a growing list of keyword arguments. """ checkpointer: Any store: Any | None = field(default=None) event_store: Any | None = field(default=None) run_events_config: Any | None = field(default=None) thread_store: Any | None = field(default=None) app_config: AppConfig | None = field(default=None) def _install_runtime_context(config: dict, runtime_context: dict[str, Any]) -> None: existing_context = config.get("context") if isinstance(existing_context, dict): existing_context.setdefault("thread_id", runtime_context["thread_id"]) existing_context.setdefault("run_id", runtime_context["run_id"]) if "app_config" in runtime_context: existing_context["app_config"] = runtime_context["app_config"] return config["context"] = dict(runtime_context) def _compute_agent_factory_supports_app_config(agent_factory: Any) -> bool: try: return "app_config" in inspect.signature(agent_factory).parameters except (TypeError, ValueError): return False @lru_cache(maxsize=128) def _cached_agent_factory_supports_app_config(agent_factory: Any) -> bool: return _compute_agent_factory_supports_app_config(agent_factory) def _agent_factory_supports_app_config(agent_factory: Any) -> bool: try: return _cached_agent_factory_supports_app_config(agent_factory) except TypeError: # Some callable instances are unhashable; fall back to a direct check. return _compute_agent_factory_supports_app_config(agent_factory) def _build_agent_with_timing(agent_factory: Any, config: Any, app_config: AppConfig | None, run_id: str) -> Any: # Keep the synchronous contract for callers whose factories need their # current event loop. The async wrapper only offloads opted-in factories. started = perf_counter() outcome = "error" try: if app_config is not None and _agent_factory_supports_app_config(agent_factory): agent = agent_factory(config=config, app_config=app_config) else: agent = agent_factory(config=config) outcome = "ok" return agent finally: elapsed_ms = (perf_counter() - started) * 1000 log = logger.warning if elapsed_ms >= 500 else logger.debug try: log("AGENT_BUILD_TIMING run_id=%s outcome=%s elapsed_ms=%.2f", run_id, outcome, elapsed_ms) except Exception: # Timing is diagnostic only; preserve the factory result/exception # even if a deployment's custom log handler is unavailable. pass async def build_agent_with_timing(agent_factory: Any, config: Any, app_config: AppConfig | None, run_id: str) -> Any: """Build agents from async entry points with shared admission and context. All async callers of an opted-in factory must use this path: a synchronous build on the owner loop can block MCP discovery scheduled by another build. """ if getattr(agent_factory, "_deerflow_background_build", False) is not True: return _build_agent_with_timing(agent_factory, config, app_config, run_id) return await run_background_init(_build_agent_with_timing, agent_factory, config, app_config, run_id) async def run_agent( bridge: StreamBridge, run_manager: RunManager, record: RunRecord, *, ctx: RunContext, agent_factory: Any, graph_input: dict, config: dict, stream_modes: list[str] | None = None, stream_subgraphs: bool = False, interrupt_before: list[str] | Literal["*"] | None = None, interrupt_after: list[str] | Literal["*"] | None = None, command: dict[str, Any] | None = None, ) -> None: """Execute an agent in the background, publishing events to *bridge*. ``command`` 来自 LangGraph SDK 的 ``Command`` 协议（``{resume, update, goto}``）。 存在时把它翻译成 ``langgraph.types.Command`` 并作为 ``astream`` 的入参 —— Pregel 检查到 Command 输入会从最近 checkpoint 处恢复（典型场景：interrupt 之后 的 ``Command(resume=user_input)``）；不存在时仍按 ``graph_input`` dict 起新 run。 AI 写作的 4 个干预卡就靠这个把 user_input 传回 ``pause_*_confirm`` 节点。 """ # Unpack infrastructure dependencies from RunContext. checkpointer = ctx.checkpointer store = ctx.store event_store = ctx.event_store run_events_config = ctx.run_events_config run_id = record.run_id thread_id = record.thread_id requested_modes: set[str] = set(stream_modes or ["values"]) internal_stream_call_ids: set[str] = set() pre_run_checkpoint_id: str | None = None pre_run_snapshot: dict[str, Any] | None = None snapshot_capture_failed = False journal = None journal = None # Track whether "events" was requested but skipped if "events" in requested_modes: logger.info( "Run %s: 'events' stream_mode not supported in gateway (requires astream_events + checkpoint callbacks). Skipping.", run_id, ) try: # 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 # path that publishes an "end" event to the SSE bridge — # otherwise a failure here would leave the stream hanging # with no terminator. if event_store is not None: from deerflow.runtime.journal import RunJournal journal = RunJournal( run_id=run_id, thread_id=thread_id, event_store=event_store, track_token_usage=getattr(run_events_config, "track_token_usage", True), ) # 1. Mark running await run_manager.set_status(run_id, RunStatus.running) # Snapshot the latest pre-run checkpoint so rollback can restore it. if checkpointer is not None: try: config_for_check = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ckpt_tuple = await checkpointer.aget_tuple(config_for_check) if ckpt_tuple is not None: ckpt_config = getattr(ckpt_tuple, "config", {}).get("configurable", {}) pre_run_checkpoint_id = ckpt_config.get("checkpoint_id") pre_run_snapshot = { "checkpoint_ns": ckpt_config.get("checkpoint_ns", ""), "checkpoint": copy.deepcopy(getattr(ckpt_tuple, "checkpoint", {})), "metadata": copy.deepcopy(getattr(ckpt_tuple, "metadata", {})), "pending_writes": copy.deepcopy(getattr(ckpt_tuple, "pending_writes", []) or []), } except Exception: snapshot_capture_failed = True logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True) # 2. Publish metadata — useStream needs both run_id AND thread_id metadata_event: dict[str, Any] = { "run_id": run_id, "thread_id": thread_id, } retrieval_trace = (record.metadata or {}).get("knowledge_retrieval_trace") if isinstance(retrieval_trace, dict): metadata_event["knowledge_retrieval_trace"] = retrieval_trace metadata_event["knowledge_retrieval_url"] = ( f"/api/threads/{thread_id}/reference-batches" ) await bridge.publish(run_id, "metadata", metadata_event) # 3. Build the agent from langchain_core.runnables import RunnableConfig from langgraph.runtime import Runtime # Inject runtime context so middlewares and tools (via ToolRuntime.context) can # access thread-level data. langgraph-cli does this automatically; we must do it # manually here because we drive the graph through ``agent.astream(config=...)`` # without passing the official ``context=`` parameter. runtime_ctx = _build_runtime_context(thread_id, run_id, config.get("context"), ctx.app_config) _install_runtime_context(config, runtime_ctx) runtime = Runtime(context=cast(Any, runtime_ctx), store=store) config.setdefault("configurable", {})["__pregel_runtime"] = runtime # Inject RunJournal as a LangChain callback handler. # on_llm_end captures token usage; on_chain_start/end captures lifecycle. if journal is not None: config.setdefault("callbacks", []).append(journal) runnable_config = RunnableConfig(**config) agent = await build_agent_with_timing(agent_factory, runnable_config, ctx.app_config, run_id) if record.abort_event.is_set(): raise asyncio.CancelledError # 4. Attach checkpointer and store if checkpointer is not None: agent.checkpointer = checkpointer if store is not None: agent.store = store # 5. Set interrupt nodes if interrupt_before: agent.interrupt_before_nodes = interrupt_before if interrupt_after: agent.interrupt_after_nodes = interrupt_after # 6. Build LangGraph stream_mode list # "events" is NOT a valid astream mode — skip it # "messages-tuple" maps to LangGraph's "messages" mode lg_modes: list[str] = [] for m in requested_modes: if m == "messages-tuple": lg_modes.append("messages") elif m == "events": # Skipped — see log above continue elif m in _VALID_LG_MODES: lg_modes.append(m) if not lg_modes: lg_modes = ["values"] # Deduplicate while preserving order seen: set[str] = set() deduped: list[str] = [] for m in lg_modes: if m not in seen: seen.add(m) deduped.append(m) lg_modes = deduped logger.info("Run %s: streaming with modes %s (requested: %s)", run_id, lg_modes, requested_modes) # 7. Stream using graph.astream # 如果客户端传了 Command（典型：interrupt 后续提的 ``Command(resume=...)``）， # 把它翻译成 langgraph.types.Command 作为 astream 输入；否则按 graph_input dict # 起新一轮。Pregel 检查到 Command 会从最近 checkpoint 处恢复，而不是重头跑。 stream_input: Any = graph_input if command is not None: # 有 resume/update/goto（如 AI 写作 4 个干预卡的 Command(resume=user_input)）→ # 构造 langgraph.types.Command，Pregel 从 interrupt 处恢复。 # **空 command（无任何指令）= "从最近 checkpoint 继续执行挂起的任务"**：AI 写作 # "暂停后继续" —— 暂停时硬取消了活跃节点（非 interrupt），continue 要让该阶段节点 # 从 checkpoint 续跑。此时必须用 **None 输入**让 Pregel resume pending tasks； # 既不能用 graph_input={}（空 dict 被当新一轮输入 → 从图入口 intent 重跑），也不能 # 构造 Command()（没有挂起的 interrupt 可消费，LangGraph 内部会抛 # "cannot access local variable 'resume_is_map'"）。 has_directive = any(command.get(k) is not None for k in ("resume", "update", "goto")) if has_directive: from langgraph.types import Command as _LGCommand stream_input = _LGCommand( resume=command.get("resume"), update=command.get("update"), goto=command.get("goto"), ) else: stream_input = None if len(lg_modes) == 1 and not stream_subgraphs: # Single mode, no subgraphs: astream yields raw chunks single_mode = lg_modes[0] async for chunk in agent.astream(stream_input, config=runnable_config, stream_mode=single_mode): if record.abort_event.is_set(): logger.info("Run %s abort requested — stopping", run_id) break sse_event = _lg_mode_to_sse_event(single_mode) serialized = serialize( chunk, mode=single_mode, known_internal_tool_call_ids=internal_stream_call_ids, ) if serialized is not None: await bridge.publish(run_id, sse_event, serialized) else: # Multiple modes or subgraphs: astream yields tuples async for item in agent.astream( stream_input, config=runnable_config, stream_mode=lg_modes, subgraphs=stream_subgraphs, ): if record.abort_event.is_set(): logger.info("Run %s abort requested — stopping", run_id) break mode, chunk = _unpack_stream_item(item, lg_modes, stream_subgraphs) if mode is None: continue sse_event = _lg_mode_to_sse_event(mode) serialized = serialize( chunk, mode=mode, known_internal_tool_call_ids=internal_stream_call_ids, ) if serialized is not None: await bridge.publish(run_id, sse_event, serialized) # 8. Final status if record.abort_event.is_set(): action = record.abort_action if action == "rollback": await _set_terminal_status(run_manager, record, RunStatus.error, error="Rolled back by user") try: await _rollback_to_pre_run_checkpoint( checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, pre_run_checkpoint_id=pre_run_checkpoint_id, pre_run_snapshot=pre_run_snapshot, snapshot_capture_failed=snapshot_capture_failed, ) logger.info("Run %s rolled back to pre-run checkpoint %s", run_id, pre_run_checkpoint_id) except Exception: logger.warning("Failed to rollback checkpoint for run %s", run_id, exc_info=True) else: await _set_terminal_status(run_manager, record, RunStatus.interrupted) else: # LLMErrorHandlingMiddleware intentionally converts an exhausted # provider failure into a user-safe AIMessage instead of raising. # That keeps a malformed upstream response out of the SSE transport, # but previously made this worker record the run as ``success``. # The UI consequently had no error event to show and looked as if # generation had silently stopped halfway through. Inspect the # terminal checkpoint before declaring success so the persisted run # state and the user-visible SSE outcome stay consistent. terminal_llm_error = await _await_completion_io( _terminal_llm_error_reason(checkpointer=checkpointer, thread_id=thread_id), run_id=run_id, phase="terminal_checkpoint", ) if terminal_llm_error is not None: error_message = "本轮模型生成未能完成，已保留当前对话内容，请稍后重试。" logger.warning( "Run %s ended with LLM fallback message (reason=%s)", run_id, terminal_llm_error, ) await _set_terminal_status( run_manager, record, RunStatus.error, error=error_message, ) await bridge.publish( run_id, "error", { "message": error_message, "name": "ModelGenerationError", }, ) else: await _set_terminal_status(run_manager, record, RunStatus.success) except asyncio.CancelledError: action = record.abort_action if action == "rollback": await _set_terminal_status(run_manager, record, RunStatus.error, error="Rolled back by user") try: await _rollback_to_pre_run_checkpoint( checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, pre_run_checkpoint_id=pre_run_checkpoint_id, pre_run_snapshot=pre_run_snapshot, snapshot_capture_failed=snapshot_capture_failed, ) logger.info("Run %s was cancelled and rolled back", run_id) except Exception: logger.warning("Run %s cancellation rollback failed", run_id, exc_info=True) else: await _set_terminal_status(run_manager, record, RunStatus.interrupted) logger.info("Run %s was cancelled", run_id) except GraphRecursionError: logger.info("Run %s reached its step budget", run_id) try: await _await_completion_io( _finish_step_limited_run(agent, runnable_config, bridge, run_id), run_id=run_id, phase="step_limit_notice", ) await _set_terminal_status(run_manager, record, RunStatus.interrupted) except Exception as exc: logger.exception("Could not save step-limit notice for run %s", run_id) await _set_terminal_status(run_manager, record, RunStatus.error, error=str(exc)) await bridge.publish(run_id, "error", { "message": "本次问答已达到步骤上限，但保存暂停提示失败，请刷新对话后重试。", "name": "StepLimitSaveError", }) except Exception as exc: error_msg = f"{exc}" logger.exception("Run %s failed: %s", run_id, error_msg) await _set_terminal_status(run_manager, record, RunStatus.error, error=error_msg) await bridge.publish( run_id, "error", { "message": error_msg, "name": type(exc).__name__, }, ) finally: try: await _await_completion_io( _finalize_run( journal=journal, run_manager=run_manager, record=record, ctx=ctx, ), run_id=run_id, phase="metadata_cleanup", ) except TimeoutError: logger.warning("Run %s completion metadata exceeded its time budget; ending SSE", run_id) except asyncio.CancelledError: # A second cancellation can arrive while handling the original # cancellation or during normal cleanup. Never lose the terminator. logger.info("Run %s completion metadata was cancelled; ending SSE", run_id) except Exception: logger.warning("Failed to finalize run %s (non-fatal)", run_id, exc_info=True) finally: await bridge.publish_end(run_id) asyncio.create_task(bridge.cleanup(run_id, delay=60)) async def _finish_step_limited_run(agent: Any, config: dict, bridge: StreamBridge, run_id: str) -> None: """Save a normal chat notice and close outstanding tool calls for the next turn. Never rerun pending tools during recovery: they may have side effects. New user input starts a fresh budget while keeping the existing transcript. """ from langchain_core.messages import AIMessage, ToolMessage config = {**config, "configurable": dict(config.get("configurable", {}))} # Recover from the run's latest checkpoint, even when its input selected # an older checkpoint explicitly. config["configurable"].pop("checkpoint_id", None) snapshot = await agent.aget_state(config) messages = snapshot.values.get("messages", []) pending: dict[str, str] = {} for message in messages: if isinstance(message, AIMessage): for call in message.tool_calls: pending[call["id"]] = call["name"] elif isinstance(message, ToolMessage): pending.pop(message.tool_call_id, None) additions: list = [ ToolMessage( content="本次问答达到步骤上限，工具执行已暂停；此调用没有可用的完成结果。", tool_call_id=call_id, name=name, status="error", ) for call_id, name in pending.items() ] notice = AIMessage( id=f"step-limit-{run_id}", content=( "本次问答已达到步骤上限，已暂停处理。已有对话内容已保留。\n\n" "如需接着完成，请发送“继续”；也可以补充要求，指定接下来优先处理的内容。" ), additional_kwargs={"step_limit_reached": True}, ) additions.append(notice) # Explicitly identify the model writer to avoid ambiguous last-writer # inference when parallel tool nodes reached the budget together. await agent.aupdate_state(config, {"messages": additions}, as_node="model" if "model" in agent.nodes else None) await agent.aupdate_state(config, None, as_node="__end__") final_state = await agent.aget_state(config) await bridge.publish(run_id, "messages", serialize((notice, {"langgraph_node": "model"}), mode="messages")) await bridge.publish(run_id, "values", serialize(final_state.values, mode="values")) async def _await_completion_io[T]( operation: Awaitable[T], *, run_id: str, phase: str, ) -> T: """Bound completion-only I/O even if a backend delays cancellation. Unlike wait_for, wait does not wait indefinitely for the cancelled operation to acknowledge cancellation. Do not use this for executing the graph: detaching a live tool would release run concurrency prematurely. """ task = asyncio.ensure_future(operation) _PENDING_COMPLETION_TASKS.add(task) def consume_result(completed: asyncio.Future) -> None: _PENDING_COMPLETION_TASKS.discard(completed) if not completed.cancelled(): completed.exception() task.add_done_callback(consume_result) try: done, _ = await asyncio.wait({task}, timeout=_RUN_FINALIZATION_TIMEOUT_SECONDS) if not done: logger.warning("RUN_COMPLETION_TIMEOUT run_id=%s phase=%s", run_id, phase) raise TimeoutError(f"Run completion timed out during {phase}") return task.result() finally: if not task.done(): task.cancel() async def _set_terminal_status(run_manager: RunManager, record: RunRecord, status: RunStatus, *, error: str | None = None) -> None: try: await _await_completion_io( run_manager.set_status(record.run_id, status, error=error), run_id=record.run_id, phase="terminal_status", ) except TimeoutError: # set_status updates the registry before awaiting its best-effort store. # Keep that local outcome authoritative if persistence never returns. record.status = status if error is not None: record.error = error logger.warning("Run %s terminal status persistence timed out", record.run_id) async def _finalize_run(*, journal: Any, run_manager: RunManager, record: RunRecord, ctx: RunContext) -> None: """Best-effort metadata writes after graph execution, under one budget.""" run_id = record.run_id thread_id = record.thread_id checkpointer = ctx.checkpointer thread_store = ctx.thread_store # Flush any buffered journal events and persist completion data if journal is not None: try: await journal.flush() _check_completion_cancellation() except Exception: logger.warning("Failed to flush journal for run %s", run_id, exc_info=True) 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) _check_completion_cancellation() except Exception: logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True) # Mark the thread idle before loading checkpoint/title data. Title refresh # must never keep the composer locked after the answer stream has ended. if thread_store is not None: try: if not await _completion_is_current(run_manager, record): return final_status = "idle" if record.status == RunStatus.success else record.status.value await thread_store.update_status(thread_id, final_status) except Exception: logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id) if checkpointer is not None and thread_store is not None: title_task = asyncio.create_task( _sync_thread_title_after_run( checkpointer=checkpointer, thread_store=thread_store, thread_id=thread_id, run_manager=run_manager, record=record, app_config=ctx.app_config, ), name=f"title-sync:{thread_id}", ) _BACKGROUND_TITLE_TASKS.add(title_task) title_task.add_done_callback(_log_background_title_finalize_failure(thread_id)) def _check_completion_cancellation() -> None: task = asyncio.current_task() if task is not None and task.cancelling(): # Some database adapters delay or suppress cancellation. They must not # resume subsequent title/status writes after the cleanup budget ends. raise asyncio.CancelledError async def _completion_is_current(run_manager: RunManager, record: RunRecord) -> bool: _check_completion_cancellation() if hasattr(run_manager, "is_latest"): return run_manager.is_latest(record) records = await run_manager.list_by_thread(record.thread_id) _check_completion_cancellation() return not records or records[-1].run_id == record.run_id # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- async def _call_checkpointer_method(checkpointer: Any, async_name: str, sync_name: str, *args: Any, **kwargs: Any) -> Any: """Call a checkpointer method, supporting async and sync variants.""" method = getattr(checkpointer, async_name, None) or getattr(checkpointer, sync_name, None) if method is None: raise AttributeError(f"Missing checkpointer method: {async_name}/{sync_name}") result = method(*args, **kwargs) if inspect.isawaitable(result): return await result return result async def _terminal_llm_error_reason( *, checkpointer: Any | None, thread_id: str, ) -> str | None: """Return the terminal LLM fallback reason, if this run ended in one. ``LLMErrorHandlingMiddleware`` persists its fallback as an AI message with an internal marker. A normal agent graph treats that message as a valid result and completes without an exception, so the worker must check the saved conversation rather than relying solely on ``astream`` raising. """ if checkpointer is None: return None try: checkpoint_config = { "configurable": { "thread_id": thread_id, "checkpoint_ns": "", } } checkpoint_tuple = await _call_checkpointer_method( checkpointer, "aget_tuple", "get_tuple", checkpoint_config, ) if checkpoint_tuple is None: return None checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {} channel_values = checkpoint.get("channel_values", {}) messages = channel_values.get("messages", []) if isinstance(channel_values, dict) else [] if not isinstance(messages, list): return None # Walk backwards to the latest assistant message. Tool/human messages # may be appended around it by middleware, but an older fallback should # never turn a later successful response into an error. for message in reversed(messages): message_type = message.get("type") if isinstance(message, dict) else getattr(message, "type", None) if message_type not in {"ai", "assistant"}: continue from deerflow.agents.middlewares.llm_error_handling_middleware import ( is_llm_error_message, ) return is_llm_error_message(message) except Exception: # A status check must never turn a completed answer into a failed run # merely because a custom checkpointer cannot expose its latest tuple. logger.debug( "Could not inspect terminal LLM state for thread %s", thread_id, exc_info=True, ) return None async def _rollback_to_pre_run_checkpoint( *, checkpointer: Any, thread_id: str, run_id: str, pre_run_checkpoint_id: str | None, pre_run_snapshot: dict[str, Any] | None, snapshot_capture_failed: bool, ) -> None: """Restore thread state to the checkpoint snapshot captured before run start.""" if checkpointer is None: logger.info("Run %s rollback requested but no checkpointer is configured", run_id) return if snapshot_capture_failed: logger.warning("Run %s rollback skipped: pre-run checkpoint snapshot capture failed", run_id) return if pre_run_snapshot is None: await _call_checkpointer_method(checkpointer, "adelete_thread", "delete_thread", thread_id) logger.info("Run %s rollback reset thread %s to empty state", run_id, thread_id) return checkpoint_to_restore = None metadata_to_restore: dict[str, Any] = {} checkpoint_ns = "" checkpoint = pre_run_snapshot.get("checkpoint") if not isinstance(checkpoint, dict): logger.warning("Run %s rollback skipped: invalid pre-run checkpoint snapshot", run_id) return checkpoint_to_restore = checkpoint if checkpoint_to_restore.get("id") is None and pre_run_checkpoint_id is not None: checkpoint_to_restore = {**checkpoint_to_restore, "id": pre_run_checkpoint_id} if checkpoint_to_restore.get("id") is None: logger.warning("Run %s rollback skipped: pre-run checkpoint has no checkpoint id", run_id) return restore_marker = _new_checkpoint_marker() checkpoint_to_restore = { **checkpoint_to_restore, "id": restore_marker["id"], "ts": restore_marker["ts"], } metadata = pre_run_snapshot.get("metadata", {}) metadata_to_restore = metadata if isinstance(metadata, dict) else {} raw_checkpoint_ns = pre_run_snapshot.get("checkpoint_ns") checkpoint_ns = raw_checkpoint_ns if isinstance(raw_checkpoint_ns, str) else "" channel_versions = checkpoint_to_restore.get("channel_versions") new_versions = dict(channel_versions) if isinstance(channel_versions, dict) else {} restore_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}} restored_config = await _call_checkpointer_method( checkpointer, "aput", "put", restore_config, checkpoint_to_restore, metadata_to_restore if isinstance(metadata_to_restore, dict) else {}, new_versions, ) if not isinstance(restored_config, dict): raise RuntimeError(f"Run {run_id} rollback restore returned invalid config: expected dict") restored_configurable = restored_config.get("configurable", {}) if not isinstance(restored_configurable, dict): raise RuntimeError(f"Run {run_id} rollback restore returned invalid config payload") restored_checkpoint_id = restored_configurable.get("checkpoint_id") if not restored_checkpoint_id: raise RuntimeError(f"Run {run_id} rollback restore did not return checkpoint_id") pending_writes = pre_run_snapshot.get("pending_writes", []) if not pending_writes: return writes_by_task: dict[str, list[tuple[str, Any]]] = {} for item in pending_writes: if not isinstance(item, (tuple, list)) or len(item) != 3: raise RuntimeError(f"Run {run_id} rollback failed: pending_write is not a 3-tuple: {item!r}") task_id, channel, value = item if not isinstance(channel, str): raise RuntimeError(f"Run {run_id} rollback failed: pending_write has non-string channel: task_id={task_id!r}, channel={channel!r}") writes_by_task.setdefault(str(task_id), []).append((channel, value)) for task_id, writes in writes_by_task.items(): await _call_checkpointer_method( checkpointer, "aput_writes", "put_writes", restored_config, writes, task_id=task_id, ) def _new_checkpoint_marker() -> dict[str, str]: marker = empty_checkpoint() return {"id": marker["id"], "ts": marker["ts"]} def _lg_mode_to_sse_event(mode: str) -> str: """Map LangGraph internal stream_mode name to SSE event name. LangGraph's ``astream(stream_mode="messages")`` produces message tuples. The SSE protocol calls this ``messages-tuple`` when the client explicitly requests it, but the default SSE event name used by LangGraph Platform is simply ``"messages"``. """ # All LG modes map 1:1 to SSE event names — "messages" stays "messages" return mode def _extract_human_message(graph_input: dict) -> HumanMessage | None: """Extract or construct a HumanMessage from graph_input for event recording. Returns a LangChain HumanMessage so callers can use .model_dump() to get the checkpoint-aligned serialization format. """ from langchain_core.messages import HumanMessage messages = graph_input.get("messages") if not messages: return None last = messages[-1] if isinstance(messages, list) else messages if isinstance(last, HumanMessage): return last if isinstance(last, str): return HumanMessage(content=last) if last else None if hasattr(last, "content"): content = last.content return HumanMessage(content=content) if isinstance(last, dict): content = last.get("content", "") return HumanMessage(content=content) if content else None return None def _unpack_stream_item( item: Any, lg_modes: list[str], stream_subgraphs: bool, ) -> tuple[str | None, Any]: """Unpack a multi-mode or subgraph stream item into (mode, chunk). Returns ``(None, None)`` if the item cannot be parsed. """ if stream_subgraphs: if isinstance(item, tuple) and len(item) == 3: _ns, mode, chunk = item return str(mode), chunk if isinstance(item, tuple) and len(item) == 2: mode, chunk = item return str(mode), chunk return None, None if isinstance(item, tuple) and len(item) == 2: mode, chunk = item return str(mode), chunk # Fallback: single-element output from first mode return lg_modes[0] if lg_modes else None, item