1009 lines
42 KiB
Python
1009 lines
42 KiB
Python
"""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
|