mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-01 19:06:01 +00:00
25ea6970a6
* implement goal continuations * fix(goal): address review findings for goal continuations - goal: key the no-progress breaker on a signature of the latest visible assistant evidence instead of the evaluator's volatile free-text, so it actually fires on stalled turns; thread the signature through every worker persist / no-progress call site - goal: align _stand_down_reason default caps with should_continue_goal (8 / 2) so the two gate functions agree on goals missing the fields - runtime: offload the synchronous checkpointer fallback via asyncio.to_thread (goal.py + worker.py) to keep blocking IO off the loop - frontend: i18n the GoalStatus "Goal" label (goalLabel in en/zh/types) - frontend: extract pure composer helpers into input-box-helpers.ts with unit tests (parseGoalCommand, readGoalResponseError, skill suggestions) - tests: cover the evidence-based no-progress and default-cap behavior - docs: align backend/AGENTS.md goal paragraph with actual behavior - e2e: prettier-format chat.spec.ts (fixes the lint-frontend CI failure) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(frontend): hide goal continuation counter until the agent continues The goal status bar rendered a raw "0/8" before any auto-continuation, which read as a mysterious score. Now the counter is hidden until continuation_count > 0, then shows "Continuing N/M" with a tooltip explaining the auto-continuation cap. - Extract getGoalContinuationDisplay into a pure helper (hides at 0) + unit tests - Add goalContinuing / goalContinuationTooltip i18n keys (en/zh/types) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix(goal): address review findings for goal continuations Frontend correctness - Fix the optimistic /goal result permanently shadowing server goal state: the streamed continuation counter never surfaced for a goal set in-session. Extract a shared useActiveGoal hook (used by both chat pages) that reconciles the optimistic copy with server state via a goalReconciliationKey, de-duping the copy-pasted goal block across the two pages. - Stop /goal status|clear failures from escaping handleSubmit as unhandled rejections (handleGoalCommand now returns success; the run only starts when a goal was actually saved). - Use a function replacer for the goal-status toast so an objective containing $&/$1 isn't treated as a replacement pattern. Backend cleanliness / correctness - De-duplicate four byte-identical helpers (_call_checkpointer_method, _message_type, _additional_kwargs, _is_visible_message) by importing them from runtime.goal instead of re-defining them in the run worker. - Remove the dead `checkpoint_tuple.tasks` durability guard (CheckpointTuple has no tasks field) and document that pending_writes is the durability signal. - Decompose the 176-line _prepare_goal_continuation_input: extract _reread_goal_and_checkpoint and a _persist closure so the thread-unchanged guard and stand-down persistence aren't open-coded three times. Document the last-writer-wins write-window limitation as a follow-up. - Add a shared parse_goal_command helper and use it from the TUI and IM-channel /goal handlers (one place for the status/clear/set semantics). Tests - Restore the 11 command-registry tests dropped by the previous goal change (filter_commands ranking/description, build_registry builtins/skills, resolve cases) alongside the new goal tests. - Add coverage for the IM-channel _handle_goal_command, the TUI _handle_goal handler, parse_goal_command, and goalReconciliationKey. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix goal review feedback * fix goal continuation checkpoint races * prioritize goal commands while streaming Route composer submits through a shared helper so /goal commands can be handled before the streaming stop shortcut, while ordinary streaming submits still stop the active run. Testing: cd frontend && pnpm exec rstest run tests/unit/components/workspace/input-box-helpers.test.ts tests/unit/components/workspace/goal-status-helpers.test.ts; cd frontend && pnpm check * preserve goal status during clarification Keep omitted stream goal fields distinct from explicit null clears so clarification interrupts do not hide an active thread goal that is still present in the checkpoint. Testing: pnpm exec rstest run tests/unit/components/workspace/use-active-goal.test.ts tests/unit/components/workspace/input-box-helpers.test.ts tests/unit/components/workspace/goal-status-helpers.test.ts; pnpm check; git diff --check * style: format active goal hook Run Prettier on use-active-goal.ts to satisfy the frontend lint workflow formatting gate. Testing: pnpm format; pnpm exec rstest run tests/unit/components/workspace/use-active-goal.test.ts tests/unit/components/workspace/input-box-helpers.test.ts tests/unit/components/workspace/goal-status-helpers.test.ts; pnpm check; git diff --check * fix goal review followups --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
1290 lines
56 KiB
Python
1290 lines
56 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
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from functools import lru_cache
|
|
from typing import Any, Literal, cast
|
|
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
|
|
from deerflow.agents.goal_state import GoalEvaluation, GoalState
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.runtime.goal import (
|
|
DEFAULT_MAX_GOAL_CONTINUATIONS,
|
|
DEFAULT_MAX_NO_PROGRESS_CONTINUATIONS,
|
|
GoalWriteConflict,
|
|
_call_checkpointer_method,
|
|
_is_visible_message,
|
|
_message_type,
|
|
attach_goal_evaluation,
|
|
compute_no_progress_count,
|
|
create_goal_evaluator_model,
|
|
evaluate_goal_completion,
|
|
goal_thread_lock,
|
|
latest_visible_assistant_signature,
|
|
make_goal_continuation_message,
|
|
read_thread_goal,
|
|
should_continue_goal,
|
|
visible_conversation_signature,
|
|
write_thread_goal,
|
|
)
|
|
from deerflow.runtime.serialization import serialize
|
|
from deerflow.runtime.stream_bridge import StreamBridge
|
|
from deerflow.runtime.user_context import get_effective_user_id
|
|
from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, get_current_trace_id, normalize_trace_id
|
|
from deerflow.tracing import inject_langfuse_metadata
|
|
from deerflow.utils.messages import message_to_text
|
|
|
|
from .manager import RunManager, RunRecord
|
|
from .naming import resolve_root_run_name
|
|
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"}
|
|
|
|
|
|
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 DEERFLOW_TRACE_METADATA_KEY in runtime_context:
|
|
existing_context.setdefault(DEERFLOW_TRACE_METADATA_KEY, runtime_context[DEERFLOW_TRACE_METADATA_KEY])
|
|
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)
|
|
|
|
|
|
class _SubagentEventBuffer:
|
|
"""Buffer subagent ``task_*`` step events and flush them in one locked batch (#3779).
|
|
|
|
The live SSE bridge already forwards these events for real-time display; this
|
|
additionally writes them so the subtask card's step history survives a reload.
|
|
|
|
``RunEventStore.put`` is documented as a low-frequency path — on Postgres each
|
|
call opens its own transaction and takes a per-thread advisory lock. A deep
|
|
subagent (``general-purpose`` runs up to ``max_turns=150``) emits hundreds of
|
|
``task_running`` steps on the hot stream loop, so persisting each with
|
|
``put()`` would serialize against the run's own message-batch writer. This
|
|
accumulates recognized subagent events and writes them with ``put_batch``,
|
|
which acquires the lock once per batch, honoring the store's contract.
|
|
|
|
Best-effort: a missing store (run_events not configured) or an unrecognized
|
|
chunk is a no-op, flush failures are logged but never propagate into the
|
|
stream loop, and terminal ``subagent.end`` events flush eagerly so a completed
|
|
subagent's step history is durable promptly rather than only at run end.
|
|
"""
|
|
|
|
#: Flush once this many events are buffered, bounding memory and reload lag on
|
|
#: a single deep subagent without paying a per-step lock.
|
|
FLUSH_THRESHOLD = 25
|
|
|
|
def __init__(self, event_store: Any | None, thread_id: str, run_id: str) -> None:
|
|
self._event_store = event_store
|
|
self._thread_id = thread_id
|
|
self._run_id = run_id
|
|
self._pending: list[dict[str, Any]] = []
|
|
|
|
async def add(self, chunk: Any) -> None:
|
|
"""Buffer one custom stream chunk; flush on a terminal event or threshold."""
|
|
if self._event_store is None:
|
|
return
|
|
# Lazy import: importing deerflow.subagents at module load triggers its
|
|
# package __init__ (executor → agents → tools → task_tool), which imports
|
|
# back from deerflow.subagents and deadlocks at gateway startup. Deferring
|
|
# it to call time (after all modules are loaded) breaks that cycle.
|
|
from deerflow.subagents.step_events import subagent_run_event
|
|
|
|
record = subagent_run_event(chunk)
|
|
if record is None:
|
|
return
|
|
self._pending.append({"thread_id": self._thread_id, "run_id": self._run_id, **record})
|
|
if record["event_type"] == "subagent.end" or len(self._pending) >= self.FLUSH_THRESHOLD:
|
|
await self.flush()
|
|
|
|
async def flush(self) -> None:
|
|
"""Persist buffered events in one ``put_batch`` call; swallow store errors."""
|
|
if self._event_store is None or not self._pending:
|
|
return
|
|
batch = self._pending
|
|
self._pending = []
|
|
try:
|
|
await self._event_store.put_batch(batch)
|
|
except Exception:
|
|
logger.warning("Run %s: failed to persist %d subagent step event(s)", self._run_id, len(batch), exc_info=True)
|
|
|
|
|
|
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,
|
|
) -> None:
|
|
"""Execute an agent in the background, publishing events to *bridge*."""
|
|
|
|
# Unpack infrastructure dependencies from RunContext.
|
|
checkpointer = ctx.checkpointer
|
|
store = ctx.store
|
|
event_store = ctx.event_store
|
|
run_events_config = ctx.run_events_config
|
|
thread_store = ctx.thread_store
|
|
|
|
run_id = record.run_id
|
|
thread_id = record.thread_id
|
|
requested_modes: set[str] = set(stream_modes or ["values"])
|
|
pre_run_checkpoint_id: str | None = None
|
|
pre_run_snapshot: dict[str, Any] | None = None
|
|
snapshot_capture_failed = False
|
|
llm_error_fallback_message: str | None = None
|
|
|
|
journal = None
|
|
# Buffers subagent step events for batched persistence (#3779); assigned once
|
|
# streaming starts and flushed in the finally block. Pre-bound to None so the
|
|
# finally is safe even if an exception fires before streaming begins.
|
|
subagent_events: _SubagentEventBuffer | None = 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:
|
|
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
|
|
# 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),
|
|
progress_reporter=lambda snapshot: run_manager.update_run_progress(run_id, **snapshot),
|
|
)
|
|
|
|
# 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
|
|
await bridge.publish(
|
|
run_id,
|
|
"metadata",
|
|
{
|
|
"run_id": run_id,
|
|
"thread_id": thread_id,
|
|
},
|
|
)
|
|
|
|
# 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)
|
|
incoming_metadata = config.get("metadata") if isinstance(config.get("metadata"), dict) else {}
|
|
deerflow_trace_id = normalize_trace_id(incoming_metadata.get(DEERFLOW_TRACE_METADATA_KEY)) or get_current_trace_id()
|
|
if deerflow_trace_id:
|
|
runtime_ctx[DEERFLOW_TRACE_METADATA_KEY] = deerflow_trace_id
|
|
# Expose the run-scoped journal under a sentinel key so middleware can
|
|
# write audit events (e.g. SafetyFinishReasonMiddleware recording
|
|
# suppressed tool calls). Double-underscore prefix marks it as a
|
|
# runtime-internal channel; user code must not depend on the key name.
|
|
if journal is not None:
|
|
runtime_ctx["__run_journal"] = journal
|
|
_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)
|
|
|
|
# Inject Langfuse trace-attribute metadata so the langchain CallbackHandler
|
|
# can lift session_id / user_id / trace_name / tags onto the root trace.
|
|
# Shared helper with ``DeerFlowClient.stream`` so both entry points stay
|
|
# in sync; caller-provided metadata wins via setdefault inside the helper.
|
|
inject_langfuse_metadata(
|
|
config,
|
|
thread_id=thread_id,
|
|
user_id=get_effective_user_id(),
|
|
assistant_id=record.assistant_id,
|
|
model_name=record.model_name,
|
|
environment=os.environ.get("DEER_FLOW_ENV") or os.environ.get("ENVIRONMENT"),
|
|
deerflow_trace_id=deerflow_trace_id,
|
|
)
|
|
|
|
# Resolve after runtime context installation so context/configurable reflect
|
|
# the agent name that this run will actually execute.
|
|
config.setdefault("run_name", resolve_root_run_name(config, record.assistant_id))
|
|
initial_runnable_config = RunnableConfig(**config)
|
|
|
|
def _continuation_runnable_config() -> RunnableConfig:
|
|
continuation_config = dict(config)
|
|
configurable = dict(continuation_config.get("configurable", {}) or {})
|
|
configurable["checkpoint_ns"] = ""
|
|
configurable.pop("checkpoint_id", None)
|
|
configurable.pop("checkpoint_map", None)
|
|
continuation_config["configurable"] = configurable
|
|
return RunnableConfig(**continuation_config)
|
|
|
|
if ctx.app_config is not None and _agent_factory_supports_app_config(agent_factory):
|
|
agent = agent_factory(config=initial_runnable_config, app_config=ctx.app_config)
|
|
else:
|
|
agent = agent_factory(config=initial_runnable_config)
|
|
|
|
# Capture the effective (resolved) model name from the agent's metadata.
|
|
# _resolve_model_name in agent.py may return the default model if the
|
|
# requested name is not in the allowlist — this update ensures the
|
|
# persisted model_name reflects the actual model used.
|
|
if record.model_name is not None:
|
|
resolved = getattr(agent, "metadata", {}) or {}
|
|
if isinstance(resolved, dict):
|
|
effective = resolved.get("model_name")
|
|
if effective and effective != record.model_name:
|
|
await run_manager.update_model_name(record.run_id, effective)
|
|
|
|
# 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)
|
|
|
|
# Buffer subagent step events and persist them in batches (#3779) instead
|
|
# of one low-frequency put() per step on the hot stream loop. Flushed in
|
|
# the finally block so buffered steps survive abort/exception paths too.
|
|
subagent_events = _SubagentEventBuffer(event_store, thread_id, run_id)
|
|
|
|
goal_evaluator_model: Any | None = None
|
|
|
|
def _get_goal_evaluator_model() -> Any:
|
|
nonlocal goal_evaluator_model
|
|
if goal_evaluator_model is None:
|
|
goal_evaluator_model = create_goal_evaluator_model(
|
|
model_name=record.model_name,
|
|
app_config=ctx.app_config,
|
|
)
|
|
return goal_evaluator_model
|
|
|
|
async def _stream_once(input_payload: Any, stream_config: RunnableConfig) -> None:
|
|
nonlocal llm_error_fallback_message
|
|
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(input_payload, config=stream_config, stream_mode=single_mode):
|
|
if record.abort_event.is_set():
|
|
logger.info("Run %s abort requested — stopping", run_id)
|
|
break
|
|
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk)
|
|
sse_event = _lg_mode_to_sse_event(single_mode)
|
|
await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode))
|
|
if single_mode == "custom":
|
|
await subagent_events.add(chunk)
|
|
return
|
|
|
|
# Multiple modes or subgraphs: astream yields tuples
|
|
async for item in agent.astream(
|
|
input_payload,
|
|
config=stream_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
|
|
|
|
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk)
|
|
sse_event = _lg_mode_to_sse_event(mode)
|
|
await bridge.publish(run_id, sse_event, serialize(chunk, mode=mode))
|
|
if mode == "custom":
|
|
await subagent_events.add(chunk)
|
|
|
|
# 7. Stream the requested turn, then optionally continue hidden goal turns.
|
|
await _stream_once(graph_input, initial_runnable_config)
|
|
while not record.abort_event.is_set() and not llm_error_fallback_message and (journal is None or not journal.had_llm_error_fallback):
|
|
continuation_input = await _prepare_goal_continuation_input(
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
model_name=record.model_name,
|
|
app_config=ctx.app_config,
|
|
evaluator_model_factory=_get_goal_evaluator_model,
|
|
abort_event=record.abort_event,
|
|
)
|
|
if continuation_input is None or record.abort_event.is_set():
|
|
break
|
|
await _stream_once(continuation_input, _continuation_runnable_config())
|
|
|
|
# 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")
|
|
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 run_manager.set_status(run_id, RunStatus.interrupted)
|
|
elif llm_error_fallback_message or (journal is not None and journal.had_llm_error_fallback):
|
|
error_msg = llm_error_fallback_message
|
|
if error_msg is None and journal is not None:
|
|
error_msg = journal.llm_error_fallback_message
|
|
error_msg = error_msg or "LLM provider failed after retries"
|
|
await run_manager.set_status(run_id, RunStatus.error, error=error_msg)
|
|
else:
|
|
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")
|
|
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 run_manager.set_status(run_id, RunStatus.interrupted)
|
|
logger.info("Run %s was cancelled", run_id)
|
|
|
|
except Exception as exc:
|
|
error_msg = f"{exc}"
|
|
logger.exception("Run %s failed: %s", run_id, error_msg)
|
|
await run_manager.set_status(run_id, RunStatus.error, error=error_msg)
|
|
await bridge.publish(
|
|
run_id,
|
|
"error",
|
|
{
|
|
"message": error_msg,
|
|
"name": type(exc).__name__,
|
|
},
|
|
)
|
|
|
|
finally:
|
|
# Persist any subagent step events still buffered (#3779) — including on
|
|
# abort/exception paths, where the stream loop broke before its own flush.
|
|
if subagent_events is not None:
|
|
await subagent_events.flush()
|
|
|
|
# Flush any buffered journal events and persist completion data
|
|
if journal is not None:
|
|
try:
|
|
await journal.flush()
|
|
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)
|
|
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:
|
|
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
|
ckpt_tuple = await checkpointer.aget_tuple(ckpt_config)
|
|
if ckpt_tuple is not None:
|
|
ckpt = getattr(ckpt_tuple, "checkpoint", {}) or {}
|
|
title = ckpt.get("channel_values", {}).get("title")
|
|
if title:
|
|
await thread_store.update_display_name(thread_id, title)
|
|
except Exception:
|
|
logger.debug("Failed to sync title for thread %s (non-fatal)", thread_id)
|
|
|
|
# Update threads_meta status based on run outcome
|
|
if thread_store is not None:
|
|
try:
|
|
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 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))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _checkpoint_id(checkpoint_tuple: Any) -> str | None:
|
|
config = getattr(checkpoint_tuple, "config", {}) or {}
|
|
configurable = config.get("configurable", {}) if isinstance(config, dict) else {}
|
|
checkpoint_id = configurable.get("checkpoint_id") if isinstance(configurable, dict) else None
|
|
if isinstance(checkpoint_id, str):
|
|
return checkpoint_id
|
|
checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {}
|
|
if isinstance(checkpoint, dict) and isinstance(checkpoint.get("id"), str):
|
|
return checkpoint["id"]
|
|
return None
|
|
|
|
|
|
def _goal_instance_matches(left: GoalState | None, right: GoalState | None) -> bool:
|
|
if not left or not right:
|
|
return False
|
|
same_status = left.get("status") == right.get("status") == "active"
|
|
same_objective = left.get("objective") == right.get("objective")
|
|
same_created_at = left.get("created_at") == right.get("created_at")
|
|
return same_status and same_objective and same_created_at
|
|
|
|
|
|
def _read_checkpoint_messages(checkpoint_tuple: Any) -> list[Any]:
|
|
checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {}
|
|
channel_values = checkpoint.get("channel_values", {}) if isinstance(checkpoint, dict) else {}
|
|
messages = channel_values.get("messages", []) if isinstance(channel_values, dict) else []
|
|
return messages if isinstance(messages, list) else []
|
|
|
|
|
|
def _read_checkpoint_goal(checkpoint_tuple: Any) -> GoalState | None:
|
|
checkpoint = getattr(checkpoint_tuple, "checkpoint", {}) or {}
|
|
channel_values = checkpoint.get("channel_values", {}) if isinstance(checkpoint, dict) else {}
|
|
raw_goal = channel_values.get("goal") if isinstance(channel_values, dict) else None
|
|
return copy.deepcopy(raw_goal) if isinstance(raw_goal, dict) else None
|
|
|
|
|
|
def _has_durable_goal_turn_receipt(checkpoint_tuple: Any, messages: list[Any]) -> bool:
|
|
"""Return true when a completed visible assistant turn is safely checkpointed.
|
|
|
|
``pending_writes`` is the durability signal: a ``CheckpointTuple`` carries no
|
|
``tasks`` field (those live on a ``StateSnapshot``), so the presence of any
|
|
queued writes is what tells us the turn is still in flight.
|
|
"""
|
|
if _checkpoint_id(checkpoint_tuple) is None:
|
|
return False
|
|
if getattr(checkpoint_tuple, "pending_writes", None):
|
|
return False
|
|
visible_messages = []
|
|
for message in messages:
|
|
if _is_visible_message(message) and message_to_text(message).strip():
|
|
visible_messages.append(message)
|
|
if not visible_messages:
|
|
return False
|
|
return _message_type(visible_messages[-1]) == "ai"
|
|
|
|
|
|
def _stand_down_reason(goal: GoalState, evaluation: GoalEvaluation, no_progress_count: int) -> str | None:
|
|
if evaluation["satisfied"]:
|
|
return None
|
|
if evaluation["blocker"] != "goal_not_met_yet":
|
|
return f"blocked:{evaluation['blocker']}"
|
|
# Default caps mirror should_continue_goal so the two gate functions agree on
|
|
# a goal dict that is missing these fields.
|
|
if int(goal.get("continuation_count", 0)) >= int(goal.get("max_continuations", DEFAULT_MAX_GOAL_CONTINUATIONS)):
|
|
return "max_continuations_reached"
|
|
if no_progress_count >= int(goal.get("max_no_progress_continuations", DEFAULT_MAX_NO_PROGRESS_CONTINUATIONS)):
|
|
return "no_progress_detected"
|
|
return None
|
|
|
|
|
|
async def _persist_goal_evaluation(
|
|
*,
|
|
bridge: StreamBridge,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
run_id: str,
|
|
goal: GoalState,
|
|
evaluation: GoalEvaluation,
|
|
no_progress_count: int,
|
|
continuation_count: int | None = None,
|
|
stand_down_reason: str | None = None,
|
|
evidence_signature: str = "",
|
|
) -> GoalState | None:
|
|
try:
|
|
async with goal_thread_lock(thread_id):
|
|
checkpoint_tuple = await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aget_tuple",
|
|
"get_tuple",
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
)
|
|
if checkpoint_tuple is None:
|
|
return None
|
|
current_goal = _read_checkpoint_goal(checkpoint_tuple)
|
|
if current_goal is None or not _goal_instance_matches(goal, current_goal):
|
|
return None
|
|
expected_checkpoint_id = _checkpoint_id(checkpoint_tuple)
|
|
updated_goal = attach_goal_evaluation(
|
|
current_goal,
|
|
evaluation,
|
|
run_id=run_id,
|
|
continuation_count=continuation_count,
|
|
no_progress_count=no_progress_count,
|
|
stand_down_reason=stand_down_reason,
|
|
evidence_signature=evidence_signature,
|
|
)
|
|
values = await write_thread_goal(
|
|
checkpointer,
|
|
thread_id,
|
|
updated_goal,
|
|
as_node="goal_evaluator",
|
|
expected_checkpoint_id=expected_checkpoint_id,
|
|
)
|
|
await bridge.publish(run_id, "values", serialize(values, mode="values"))
|
|
return updated_goal
|
|
except GoalWriteConflict:
|
|
return None
|
|
except Exception:
|
|
logger.warning("Could not persist goal evaluation for thread %s", thread_id, exc_info=True)
|
|
return None
|
|
|
|
|
|
async def _reread_goal_and_checkpoint(checkpointer: Any, thread_id: str) -> tuple[GoalState | None, Any]:
|
|
"""Re-read the goal and latest checkpoint together for a concurrency re-check."""
|
|
goal = await read_thread_goal(checkpointer, thread_id)
|
|
checkpoint_tuple = await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aget_tuple",
|
|
"get_tuple",
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
)
|
|
return goal, checkpoint_tuple
|
|
|
|
|
|
async def _prepare_goal_continuation_input(
|
|
*,
|
|
bridge: StreamBridge,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
run_id: str,
|
|
model_name: str | None,
|
|
app_config: AppConfig | None,
|
|
evaluator_model_factory: Any | None = None,
|
|
abort_event: asyncio.Event | None = None,
|
|
) -> dict[str, Any] | None:
|
|
"""Evaluate the active goal and return a hidden continuation input if needed.
|
|
|
|
NOTE: The re-reads below catch a racing user message or ``/goal clear``
|
|
before we queue a continuation. Goal writes then serialize per thread and
|
|
pass the checkpoint id they read from, so stale evaluator writes stand down
|
|
instead of clobbering a newer goal change.
|
|
"""
|
|
if checkpointer is None:
|
|
return None
|
|
if abort_event is not None and abort_event.is_set():
|
|
return None
|
|
|
|
try:
|
|
goal = await read_thread_goal(checkpointer, thread_id)
|
|
except Exception:
|
|
logger.warning("Could not read goal for thread %s after run %s", thread_id, run_id, exc_info=True)
|
|
return None
|
|
if not goal or goal.get("status") != "active":
|
|
return None
|
|
|
|
async def _persist(
|
|
goal: GoalState,
|
|
evaluation: GoalEvaluation,
|
|
no_progress_count: int,
|
|
*,
|
|
stand_down_reason: str | None = None,
|
|
continuation_count: int | None = None,
|
|
) -> GoalState | None:
|
|
"""Record the evaluation against the still-current goal instance."""
|
|
return await _persist_goal_evaluation(
|
|
bridge=bridge,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
goal=goal,
|
|
evaluation=evaluation,
|
|
no_progress_count=no_progress_count,
|
|
continuation_count=continuation_count,
|
|
stand_down_reason=stand_down_reason,
|
|
evidence_signature=evidence_signature,
|
|
)
|
|
|
|
try:
|
|
checkpoint_tuple = await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aget_tuple",
|
|
"get_tuple",
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
)
|
|
if checkpoint_tuple is None:
|
|
return None
|
|
checkpoint_id_before = _checkpoint_id(checkpoint_tuple)
|
|
messages = _read_checkpoint_messages(checkpoint_tuple)
|
|
conversation_signature_before = visible_conversation_signature(messages)
|
|
evidence_signature = latest_visible_assistant_signature(messages)
|
|
|
|
if not _has_durable_goal_turn_receipt(checkpoint_tuple, messages):
|
|
evaluation = GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="run_failed",
|
|
reason="No durable assistant end-of-turn receipt was available.",
|
|
evidence_summary="",
|
|
)
|
|
no_progress_count = compute_no_progress_count(goal, evaluation, evidence_signature=evidence_signature)
|
|
await _persist(goal, evaluation, no_progress_count, stand_down_reason="no_durable_end_of_turn")
|
|
return None
|
|
|
|
if abort_event is not None and abort_event.is_set():
|
|
return None
|
|
evaluator_model = evaluator_model_factory() if evaluator_model_factory is not None else None
|
|
evaluation = await evaluate_goal_completion(
|
|
goal,
|
|
messages,
|
|
model=evaluator_model,
|
|
model_name=model_name,
|
|
app_config=app_config,
|
|
)
|
|
if abort_event is not None and abort_event.is_set():
|
|
return None
|
|
except Exception:
|
|
logger.warning("Goal evaluator failed for thread %s after run %s", thread_id, run_id, exc_info=True)
|
|
return None
|
|
|
|
no_progress_count = compute_no_progress_count(goal, evaluation, evidence_signature=evidence_signature)
|
|
|
|
# Re-check that neither the goal nor the visible conversation changed while the
|
|
# evaluator ran — a user message or /goal clear racing the evaluation must win.
|
|
try:
|
|
current_goal, current_checkpoint_tuple = await _reread_goal_and_checkpoint(checkpointer, thread_id)
|
|
except Exception:
|
|
logger.warning("Could not re-check goal state for thread %s after evaluation", thread_id, exc_info=True)
|
|
return None
|
|
|
|
if not _goal_instance_matches(goal, current_goal) or current_checkpoint_tuple is None:
|
|
return None
|
|
|
|
checkpoint_changed = _checkpoint_id(current_checkpoint_tuple) != checkpoint_id_before
|
|
messages_changed = visible_conversation_signature(_read_checkpoint_messages(current_checkpoint_tuple)) != conversation_signature_before
|
|
if checkpoint_changed or messages_changed:
|
|
await _persist(current_goal, evaluation, no_progress_count, stand_down_reason="thread_changed_after_evaluation")
|
|
return None
|
|
|
|
if evaluation["satisfied"]:
|
|
try:
|
|
async with goal_thread_lock(thread_id):
|
|
latest_checkpoint_tuple = await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aget_tuple",
|
|
"get_tuple",
|
|
{"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}},
|
|
)
|
|
if latest_checkpoint_tuple is None:
|
|
return None
|
|
latest_goal = _read_checkpoint_goal(latest_checkpoint_tuple)
|
|
if latest_goal is None or not _goal_instance_matches(goal, latest_goal):
|
|
return None
|
|
values = await write_thread_goal(
|
|
checkpointer,
|
|
thread_id,
|
|
None,
|
|
as_node="goal_evaluator",
|
|
expected_checkpoint_id=_checkpoint_id(latest_checkpoint_tuple),
|
|
)
|
|
await bridge.publish(run_id, "values", serialize(values, mode="values"))
|
|
except GoalWriteConflict:
|
|
return None
|
|
except Exception:
|
|
logger.warning("Could not clear satisfied goal for thread %s", thread_id, exc_info=True)
|
|
return None
|
|
|
|
stand_down_reason = _stand_down_reason(goal, evaluation, no_progress_count)
|
|
if stand_down_reason is not None or not should_continue_goal(goal, evaluation, no_progress_count=no_progress_count):
|
|
await _persist(goal, evaluation, no_progress_count, stand_down_reason=stand_down_reason)
|
|
return None
|
|
|
|
next_count = int(goal.get("continuation_count", 0)) + 1
|
|
updated_goal = await _persist(goal, evaluation, no_progress_count, continuation_count=next_count)
|
|
if updated_goal is None:
|
|
return None
|
|
|
|
# Final guard: the persist above bumped the checkpoint id, so only the visible
|
|
# conversation signature is meaningful for detecting a racing user turn here.
|
|
try:
|
|
latest_goal, latest_checkpoint_tuple = await _reread_goal_and_checkpoint(checkpointer, thread_id)
|
|
except Exception:
|
|
logger.warning("Could not verify queued goal continuation for thread %s", thread_id, exc_info=True)
|
|
return None
|
|
if not _goal_instance_matches(updated_goal, latest_goal) or latest_checkpoint_tuple is None:
|
|
return None
|
|
if visible_conversation_signature(_read_checkpoint_messages(latest_checkpoint_tuple)) != conversation_signature_before:
|
|
await _persist(
|
|
latest_goal,
|
|
evaluation,
|
|
no_progress_count,
|
|
continuation_count=next_count,
|
|
stand_down_reason="thread_changed_before_continuation",
|
|
)
|
|
return None
|
|
|
|
logger.info(
|
|
"Run %s continuing thread %s for active goal (%d/%d)",
|
|
run_id,
|
|
thread_id,
|
|
updated_goal.get("continuation_count", next_count),
|
|
updated_goal.get("max_continuations", 0),
|
|
)
|
|
return {"messages": [make_goal_continuation_message(updated_goal, evaluation)]}
|
|
|
|
|
|
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 _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[<channel>]``, 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.
|
|
|
|
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 _error_fallback_message_from_metadata(metadata: dict[str, Any], content: Any) -> str:
|
|
detail = metadata.get("error_detail")
|
|
if isinstance(detail, str) and detail.strip():
|
|
return detail.strip()
|
|
reason = metadata.get("error_reason")
|
|
if isinstance(reason, str) and reason.strip():
|
|
return reason.strip()
|
|
if isinstance(content, str) and content.strip():
|
|
return content.strip()[:2000]
|
|
return "LLM provider failed after retries"
|
|
|
|
|
|
def _try_extract_from_message(obj: Any) -> str | None:
|
|
"""Try to extract fallback marker from a single message object or dict."""
|
|
additional_kwargs = getattr(obj, "additional_kwargs", None)
|
|
if isinstance(additional_kwargs, dict) and additional_kwargs.get("deerflow_error_fallback"):
|
|
return _error_fallback_message_from_metadata(additional_kwargs, getattr(obj, "content", None))
|
|
|
|
if isinstance(obj, dict):
|
|
nested_kwargs = obj.get("additional_kwargs")
|
|
if isinstance(nested_kwargs, dict) and nested_kwargs.get("deerflow_error_fallback"):
|
|
return _error_fallback_message_from_metadata(nested_kwargs, obj.get("content"))
|
|
return None
|
|
|
|
|
|
def _extract_llm_error_fallback_message(value: Any) -> str | None:
|
|
"""Find LLM fallback markers in streamed LangGraph chunks.
|
|
|
|
Error fallback messages returned by model-call middleware are not guaranteed
|
|
to pass through LLM end callbacks, but they do appear in graph state chunks.
|
|
"""
|
|
# Fast path: large state chunks produced by stream_mode="values" have a
|
|
# top-level "messages" list. Scanning only that list avoids expensive deep
|
|
# recursion into large state dicts.
|
|
if isinstance(value, dict):
|
|
messages = value.get("messages")
|
|
if isinstance(messages, (list, tuple)):
|
|
for msg in messages:
|
|
result = _try_extract_from_message(msg)
|
|
if result is not None:
|
|
return result
|
|
# Fallback marker is attached to an AI message in the messages
|
|
# channel; it will never appear elsewhere in a values chunk.
|
|
return None
|
|
# No top-level "messages" — this is likely an "updates" chunk (small
|
|
# dict keyed by node name). Fall through to deep walk, which is cheap
|
|
# for these payloads.
|
|
|
|
# Deep walk for updates / messages / tuple / list modes. Payloads are
|
|
# small, so full recursion is acceptable here.
|
|
seen: set[int] = set()
|
|
|
|
def walk(obj: Any) -> str | None:
|
|
oid = id(obj)
|
|
if oid in seen:
|
|
return None
|
|
seen.add(oid)
|
|
|
|
result = _try_extract_from_message(obj)
|
|
if result is not None:
|
|
return result
|
|
|
|
if isinstance(obj, dict):
|
|
for item in obj.values():
|
|
result = walk(item)
|
|
if result is not None:
|
|
return result
|
|
return None
|
|
|
|
if isinstance(obj, (list, tuple, set)):
|
|
for item in obj:
|
|
result = walk(item)
|
|
if result is not None:
|
|
return result
|
|
return None
|
|
|
|
return walk(value)
|
|
|
|
|
|
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
|