mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-28 14:10:37 +00:00
c994826d06
* feat(channels): generalize IM group targets and add task/goal run notifications
- Generalize proactive group notification capability (proactive_group_notifications), binding validation, and scheduled task group selection across Feishu, QQ, and WeCom while keeping WeChat personal-only.
- Add QQ group active notification delivery via POST /v2/groups/{group_openid}/messages (msg_type=2) and WeCom group active notification delivery via WS Client.sendMessage(chatid, body) with bot_id:group:chatid workspace isolation.
- Generalize notification outbox sources via migration 0043_runtime_im_notices (source_kind, source_key, run_id, nullable task_id/task_run_id) with backward-compatible SQLite/PostgreSQL upgrades.
- Add opt-in ordinary task & goal run IM notifications (run_completed, run_failed, human_input_requested) on the shared outbox, delivery worker, renderer, and settings UI.
* fix(notifications): address review feedback and group/runtime notification edge cases
* test(migrations): update CURRENT_HEAD to 0043_runtime_im_notices across migration chain tests
* fix(channels): preserve group binding names and legacy authorization
Carry owner-chosen group names through existing connect state, keep private bindings out of legacy group fallback, and restore completion metadata retries. Cover real binding and worker finalization paths and synchronize notification docs.
---------
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
3513 lines
166 KiB
Python
3513 lines
166 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 rejected by 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 gc
|
|
import inspect
|
|
import logging
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
import weakref
|
|
from collections.abc import Callable, Coroutine, Mapping
|
|
from contextlib import AbstractAsyncContextManager
|
|
from contextvars import Context
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime
|
|
from functools import lru_cache
|
|
from typing import Any, Final, Literal, cast
|
|
|
|
from deerflow_extension_api.agent_runs import AGENT_RUNS_CONTEXT_KEY
|
|
from langgraph.checkpoint.base import empty_checkpoint
|
|
from langgraph.types import Overwrite
|
|
|
|
from deerflow.agents.goal_state import GoalEvaluation, GoalState
|
|
from deerflow.agents.interaction_policy import RunInteractionPolicy, resolve_run_interaction_policy
|
|
from deerflow.agents.middlewares.input_sanitization_middleware import neutralize_untrusted_tags
|
|
from deerflow.config.app_config import AppConfig
|
|
from deerflow.config.database_config import CheckpointChannelMode
|
|
from deerflow.constants import CONVERSATION_READER_CONTEXT_KEY, TOOL_RESULTS_DIRNAME
|
|
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope
|
|
from deerflow.mcp_scope import (
|
|
THREAD_INCARNATION_CONTEXT_KEY,
|
|
THREAD_INCARNATION_METADATA_GUARD_KEY,
|
|
)
|
|
from deerflow.runtime.checkpoint_mode import (
|
|
aensure_checkpoint_mode_compatible,
|
|
inject_checkpoint_mode,
|
|
)
|
|
from deerflow.runtime.checkpoint_state import (
|
|
CheckpointStateAccessor,
|
|
build_state_mutation_graph,
|
|
graph_reducer_channels,
|
|
graph_state_schema,
|
|
graph_writable_channels,
|
|
)
|
|
from deerflow.runtime.context_keys import (
|
|
CHECKPOINT_AGENT_NAME_METADATA_KEY,
|
|
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
|
|
DEFAULT_AGENT_NAME_METADATA_VALUE,
|
|
PROJECT_CONTEXT_KEY,
|
|
checkpoint_agent_binding_metadata,
|
|
)
|
|
from deerflow.runtime.events.message_identity import attach_message_seq, message_identity
|
|
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,
|
|
build_goal_outcome,
|
|
build_goal_state,
|
|
compute_no_progress_count,
|
|
create_goal_evaluator_model,
|
|
evaluate_goal_completion,
|
|
goal_thread_lock,
|
|
latest_visible_assistant_message_id,
|
|
latest_visible_assistant_signature,
|
|
make_goal_continuation_message,
|
|
read_thread_goal,
|
|
should_continue_goal,
|
|
visible_conversation_signature,
|
|
write_thread_goal,
|
|
)
|
|
from deerflow.runtime.keyed_lock import AsyncKeyedLockTable
|
|
from deerflow.runtime.runs.stream_cleanup import close_agent_stream
|
|
from deerflow.runtime.serialization import serialize
|
|
from deerflow.runtime.stream_bridge import StreamBridge
|
|
from deerflow.runtime.stream_modes import normalize_stream_modes, to_langgraph_stream_modes
|
|
from deerflow.runtime.user_context import get_current_user, get_effective_user_id, resolve_runtime_user_id
|
|
from deerflow.sandbox.lease import SANDBOX_SERVER_OWNED_CONTEXT_KEYS
|
|
from deerflow.scheduler.runtime import SCHEDULER_CAPABILITY_CONTEXT_KEY, SchedulerRunCapability
|
|
from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_id
|
|
from deerflow.tracing import inject_langfuse_metadata
|
|
from deerflow.utils.assembly_io import run_assembly
|
|
from deerflow.utils.messages import message_to_text
|
|
from deerflow.workspace_changes import capture_workspace_snapshot, get_changed_output_paths, record_workspace_changes
|
|
from deerflow.workspace_changes.types import WorkspaceSnapshot
|
|
|
|
from .manager import ConflictError, RunManager, RunRecord, RunStartOutcome
|
|
from .naming import resolve_root_run_name
|
|
from .schemas import RunStatus, ThreadOperationKind
|
|
from .store.base import canonical_run_created_at
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_THREAD_INCARNATION_UNSET = object()
|
|
|
|
|
|
def _log_cancelled_stream_close_failure(
|
|
exc: asyncio.CancelledError,
|
|
*,
|
|
run_id: str,
|
|
abort_requested: bool,
|
|
) -> None:
|
|
close_failure = exc.__cause__
|
|
if not isinstance(close_failure, Exception):
|
|
return
|
|
log = logger.warning if abort_requested else logger.debug
|
|
message = "Could not close aborted agent stream for run %s" if abort_requested else "Could not close agent stream for run %s"
|
|
log(
|
|
message,
|
|
run_id,
|
|
exc_info=(type(close_failure), close_failure, close_failure.__traceback__),
|
|
)
|
|
|
|
|
|
_checkpoint_locks = AsyncKeyedLockTable[str]()
|
|
|
|
# Completed LangGraph runs can leave callback Contexts and AsyncPregelLoop
|
|
# instances in unreachable reference cycles. They are collectable, but a busy
|
|
# Gateway may promote those cycles into older GC generations faster than the
|
|
# automatic collector revisits them, producing a rising post-GC heap floor.
|
|
# Coalesce terminal full collections so the cycles have a bounded lifetime
|
|
# without paying for a stop-the-world collection on every run.
|
|
_TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS = 10.0
|
|
_TERMINAL_CYCLE_COLLECTION_INFO_THRESHOLD_SECONDS = 0.1
|
|
_terminal_cycle_collection_guard = threading.Lock()
|
|
_terminal_cycle_collection_last_at = time.monotonic()
|
|
_terminal_cycle_collection_scheduled_loops: weakref.WeakSet[asyncio.AbstractEventLoop] = weakref.WeakSet()
|
|
|
|
|
|
def _create_contextless_task(coro: Coroutine[Any, Any, Any]) -> asyncio.Task[Any]:
|
|
"""Schedule terminal housekeeping without retaining the run's ContextVars."""
|
|
return asyncio.create_task(coro, context=Context())
|
|
|
|
|
|
def _schedule_terminal_cycle_collection() -> None:
|
|
"""Coalesce full cyclic-GC passes after completed LangGraph runs."""
|
|
loop = asyncio.get_running_loop()
|
|
with _terminal_cycle_collection_guard:
|
|
if loop in _terminal_cycle_collection_scheduled_loops:
|
|
return
|
|
elapsed = time.monotonic() - _terminal_cycle_collection_last_at
|
|
delay = max(0.0, _TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS - elapsed)
|
|
_terminal_cycle_collection_scheduled_loops.add(loop)
|
|
|
|
async def _collect() -> None:
|
|
global _terminal_cycle_collection_last_at
|
|
|
|
try:
|
|
with _terminal_cycle_collection_guard:
|
|
now = time.monotonic()
|
|
if now - _terminal_cycle_collection_last_at < _TERMINAL_CYCLE_COLLECTION_INTERVAL_SECONDS:
|
|
return
|
|
_terminal_cycle_collection_last_at = now
|
|
|
|
started_at = time.perf_counter()
|
|
# Do not run a heap walk synchronously from the event-loop timer.
|
|
# CPython's collector can still contend for the GIL, so surface slow
|
|
# passes below rather than claiming this removes every pause.
|
|
collected = await loop.run_in_executor(None, gc.collect)
|
|
duration = time.perf_counter() - started_at
|
|
if duration >= _TERMINAL_CYCLE_COLLECTION_INFO_THRESHOLD_SECONDS:
|
|
logger.info(
|
|
"Terminal cyclic GC collected %d object(s) in %.3f seconds",
|
|
collected,
|
|
duration,
|
|
)
|
|
else:
|
|
logger.debug(
|
|
"Terminal cyclic GC collected %d object(s) in %.3f seconds",
|
|
collected,
|
|
duration,
|
|
)
|
|
finally:
|
|
with _terminal_cycle_collection_guard:
|
|
_terminal_cycle_collection_scheduled_loops.discard(loop)
|
|
|
|
def _start_collection() -> None:
|
|
_create_contextless_task(_collect())
|
|
|
|
# A blank Context prevents the timer itself from retaining the completed
|
|
# run. The loop owns the TimerHandle; the WeakSet never keeps a test or
|
|
# short-lived embedded-client event loop alive.
|
|
loop.call_later(delay, _start_collection, context=Context())
|
|
|
|
|
|
def _remove_callback(config: dict[str, Any], handler: Any) -> None:
|
|
callbacks = config.get("callbacks")
|
|
if isinstance(callbacks, list):
|
|
callbacks[:] = [callback for callback in callbacks if callback is not handler]
|
|
return
|
|
remove_handler = getattr(callbacks, "remove_handler", None)
|
|
if callable(remove_handler):
|
|
try:
|
|
remove_handler(handler)
|
|
except Exception:
|
|
logger.debug("could not detach terminal callback", exc_info=True)
|
|
|
|
|
|
def _release_run_scoped_references(
|
|
configs: list[dict[str, Any]],
|
|
runtime_context: dict[str, Any] | None,
|
|
journal: Any | None,
|
|
) -> None:
|
|
"""Remove worker-owned graph references once durable finalization is done."""
|
|
internal_context_keys = {
|
|
"__run_journal",
|
|
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
|
|
CONVERSATION_READER_CONTEXT_KEY,
|
|
SCHEDULER_CAPABILITY_CONTEXT_KEY,
|
|
AGENT_RUNS_CONTEXT_KEY,
|
|
}
|
|
try:
|
|
from deerflow.extensions import EXTENSION_SNAPSHOT_CONTEXT_KEY
|
|
|
|
internal_context_keys.add(EXTENSION_SNAPSHOT_CONTEXT_KEY)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
from deerflow_extension_api import EXTENSION_TASK_STORE_KEY
|
|
|
|
internal_context_keys.add(EXTENSION_TASK_STORE_KEY)
|
|
except Exception:
|
|
pass
|
|
|
|
seen_configs: set[int] = set()
|
|
handlers = [journal] if journal is not None else []
|
|
for runnable_config in configs:
|
|
if not isinstance(runnable_config, dict) or id(runnable_config) in seen_configs:
|
|
continue
|
|
seen_configs.add(id(runnable_config))
|
|
configurable = runnable_config.get("configurable")
|
|
if isinstance(configurable, dict):
|
|
configurable.pop("__pregel_runtime", None)
|
|
configurable.pop(CONVERSATION_READER_CONTEXT_KEY, None)
|
|
configurable.pop(SCHEDULER_CAPABILITY_CONTEXT_KEY, None)
|
|
configurable.pop(AGENT_RUNS_CONTEXT_KEY, None)
|
|
context = runnable_config.get("context")
|
|
if isinstance(context, dict):
|
|
for key in internal_context_keys:
|
|
context.pop(key, None)
|
|
for handler in handlers:
|
|
_remove_callback(runnable_config, handler)
|
|
|
|
if isinstance(runtime_context, dict):
|
|
for key in internal_context_keys:
|
|
runtime_context.pop(key, None)
|
|
|
|
|
|
def _checkpoint_thread_lock(thread_id: str) -> AbstractAsyncContextManager[None]:
|
|
"""Serialize checkpoint mutations for one thread without blocking goal commands."""
|
|
return _checkpoint_locks.hold(thread_id)
|
|
|
|
|
|
_DELIVERY_RECEIPT_RETRY_DELAYS_SECONDS = (0.1, 0.5)
|
|
_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS = 3.0
|
|
|
|
|
|
def _project_background_tasks(task_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
|
"""Build the bounded model-state projection without trusting display names."""
|
|
return [
|
|
{
|
|
"task_id": row["id"],
|
|
"task_name": neutralize_untrusted_tags(str(row["task_name"])),
|
|
"status": row["status"],
|
|
"updated_at": row["updated_at"],
|
|
}
|
|
for row in task_rows
|
|
]
|
|
|
|
|
|
async def _persist_delivery_receipt(
|
|
event_store: Any,
|
|
*,
|
|
thread_id: str,
|
|
run_id: str,
|
|
content: dict[str, Any],
|
|
) -> bool:
|
|
"""Persist a terminal receipt with short bounded retries.
|
|
|
|
The owning worker still knows the real terminal outcome and renews its
|
|
lease while this coroutine runs. Retrying here handles transient event
|
|
store failures without handing a successful run to orphan recovery, which
|
|
cannot reconstruct either the terminal status or the detailed receipt.
|
|
"""
|
|
attempts = len(_DELIVERY_RECEIPT_RETRY_DELAYS_SECONDS) + 1
|
|
for attempt in range(attempts):
|
|
try:
|
|
await event_store.put_if_absent(
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
event_type="run.delivery",
|
|
category="outputs",
|
|
content=content,
|
|
)
|
|
return True
|
|
except Exception:
|
|
if attempt == attempts - 1:
|
|
logger.warning(
|
|
"Failed to persist delivery receipt for run %s after %d attempts; applying terminal delivery semantics without a receipt",
|
|
run_id,
|
|
attempts,
|
|
exc_info=True,
|
|
)
|
|
return False
|
|
delay = _DELIVERY_RECEIPT_RETRY_DELAYS_SECONDS[attempt]
|
|
logger.warning(
|
|
"Failed to persist delivery receipt for run %s (attempt %d/%d); retrying in %.1fs",
|
|
run_id,
|
|
attempt + 1,
|
|
attempts,
|
|
delay,
|
|
exc_info=True,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
|
|
return False # pragma: no cover - loop always returns
|
|
|
|
|
|
_DELIVERY_INCOMPLETE_ERROR = "Artifact delivery incomplete: no produced output artifact was presented"
|
|
_DELIVERY_RECEIPT_FAILED_ERROR = "Artifact delivery verification failed: terminal delivery receipt could not be persisted"
|
|
|
|
|
|
def _empty_delivery_content() -> dict[str, Any]:
|
|
return {"presented": 0, "paths": [], "by_tool": {}}
|
|
|
|
|
|
def _presented_path_covers_output(presented_path: str, produced_path: str) -> bool:
|
|
presented_path = presented_path.rstrip("/")
|
|
return bool(presented_path) and (produced_path == presented_path or produced_path.startswith(f"{presented_path}/"))
|
|
|
|
|
|
def _delivery_content_with_outputs(
|
|
content: dict[str, Any],
|
|
produced_paths: list[str],
|
|
) -> dict[str, Any]:
|
|
"""Attach a delivery verdict when this run created or modified outputs."""
|
|
if not produced_paths:
|
|
return content
|
|
|
|
presented_paths = content.get("by_tool", {}).get("present_files", [])
|
|
matched_paths = [produced_path for produced_path in produced_paths if any(_presented_path_covers_output(presented_path, produced_path) for presented_path in presented_paths)]
|
|
satisfied = bool(matched_paths)
|
|
return {
|
|
**content,
|
|
"verification": {
|
|
"source": "outputs_changed",
|
|
"requirement": "present_files_matches_produced_output",
|
|
},
|
|
"produced_paths": produced_paths,
|
|
"presented_paths": presented_paths,
|
|
"matched_paths": matched_paths,
|
|
"stage": "presented" if satisfied else ("mismatched" if presented_paths else "not_started"),
|
|
"satisfied": satisfied,
|
|
}
|
|
|
|
|
|
def _delivery_error(content: dict[str, Any]) -> str | None:
|
|
"""Return the terminal error when no changed output was presented."""
|
|
if not content.get("produced_paths") or content.get("satisfied") is True:
|
|
return None
|
|
return _DELIVERY_INCOMPLETE_ERROR
|
|
|
|
|
|
def _workspace_excluded_dir_names(app_config: AppConfig | None) -> frozenset[str]:
|
|
"""Directory names workspace snapshots must skip for this deployment.
|
|
|
|
The tool-output budget middleware externalizes oversized tool outputs into
|
|
a storage subdir under outputs (default ``.tool-results``). Those files are
|
|
process feedback referenced from the budget preview via ``read_file``, not
|
|
deliverables: counting them as produced artifacts would fail run delivery
|
|
verification for any run that externalized a tool output without also
|
|
presenting a real artifact. The default name is excluded by the scanner
|
|
itself; a custom ``tool_output.storage_subdir`` (a single-segment name,
|
|
enforced by ``ToolOutputConfig`` so the scanner's dir-name pruning always
|
|
matches) is threaded through the snapshot capture here so before/after
|
|
diffs stay consistent.
|
|
"""
|
|
storage_subdir = app_config.tool_output.storage_subdir if app_config is not None else TOOL_RESULTS_DIRNAME
|
|
return frozenset({storage_subdir})
|
|
|
|
|
|
async def _produced_output_paths(
|
|
before: WorkspaceSnapshot | None,
|
|
*,
|
|
thread_id: str,
|
|
user_id: str | None,
|
|
extra_excluded_dir_names: frozenset[str] | None = None,
|
|
) -> list[str]:
|
|
"""Detect regular output files created or modified by this run."""
|
|
if before is None:
|
|
return []
|
|
try:
|
|
after = await capture_workspace_snapshot(thread_id, user_id=user_id, include_text=False, extra_excluded_dir_names=extra_excluded_dir_names)
|
|
return get_changed_output_paths(before, after)
|
|
except Exception:
|
|
logger.warning("Could not detect produced output artifacts for run thread %s", thread_id, exc_info=True)
|
|
return []
|
|
|
|
|
|
# Keep this streaming policy separate from middleware write-authorization sets.
|
|
_LARGE_FILE_TOOL_NAMES = frozenset({"str_replace", "write_file"})
|
|
_LARGE_FILE_TOOL_BATCH_SIZE = 32
|
|
|
|
|
|
@dataclass
|
|
class _LargeFileToolChunkBatcher:
|
|
"""Batch file-body argument deltas to avoid quadratic browser parsing.
|
|
|
|
Normal assistant text and non-file tool calls remain token-streamed. Large
|
|
file arguments still update progressively, but in bounded batches instead
|
|
of forcing the browser to reparse the growing JSON on every model token.
|
|
"""
|
|
|
|
batch_size: int = _LARGE_FILE_TOOL_BATCH_SIZE
|
|
tool_names: dict[tuple[str, str, str], str] = field(default_factory=dict)
|
|
pending_identity: tuple[str, str, str] | None = None
|
|
pending_message: Any | None = None
|
|
pending_metadata: dict[str, Any] = field(default_factory=dict)
|
|
pending_count: int = 0
|
|
|
|
def push(self, chunk: Any) -> list[Any]:
|
|
if not isinstance(chunk, tuple) or len(chunk) != 2:
|
|
return [*self.flush(), chunk]
|
|
|
|
message, metadata = chunk
|
|
message_id = getattr(message, "id", None)
|
|
tool_call_chunks = getattr(message, "tool_call_chunks", None)
|
|
if not isinstance(message_id, str) or not message_id or not isinstance(tool_call_chunks, list) or len(tool_call_chunks) != 1:
|
|
return [*self.flush(), chunk]
|
|
|
|
tool_chunk = tool_call_chunks[0]
|
|
if not isinstance(tool_chunk, dict):
|
|
return [*self.flush(), chunk]
|
|
index = tool_chunk.get("index")
|
|
tool_call_id = tool_chunk.get("id")
|
|
if isinstance(index, int):
|
|
discriminator = f"index:{index}"
|
|
elif isinstance(tool_call_id, str) and tool_call_id:
|
|
discriminator = f"id:{tool_call_id}"
|
|
else:
|
|
discriminator = "single"
|
|
raw_namespace = None
|
|
if isinstance(metadata, dict):
|
|
raw_namespace = metadata.get("langgraph_checkpoint_ns") or metadata.get("checkpoint_ns")
|
|
namespace = raw_namespace if isinstance(raw_namespace, str) else ""
|
|
identity = (namespace, message_id, discriminator)
|
|
name_fragment = tool_chunk.get("name")
|
|
tool_name = self.tool_names.get(identity, "")
|
|
if tool_name not in _LARGE_FILE_TOOL_NAMES and isinstance(name_fragment, str) and name_fragment:
|
|
tool_name += name_fragment
|
|
if any(candidate.startswith(tool_name) for candidate in _LARGE_FILE_TOOL_NAMES):
|
|
self.tool_names[identity] = tool_name
|
|
else:
|
|
self.tool_names.pop(identity, None)
|
|
# Batching starts only after the accumulated name matches; split or
|
|
# incomplete name fragments stream per-chunk until then.
|
|
if tool_name not in _LARGE_FILE_TOOL_NAMES:
|
|
return [*self.flush(), chunk]
|
|
|
|
model_copy = getattr(message, "model_copy", None)
|
|
if not callable(model_copy):
|
|
return [*self.flush(), chunk]
|
|
additional_kwargs = getattr(message, "additional_kwargs", None)
|
|
sanitized_additional_kwargs = additional_kwargs
|
|
if isinstance(additional_kwargs, dict) and ("function_call" in additional_kwargs or "tool_calls" in additional_kwargs):
|
|
sanitized_additional_kwargs = {key: value for key, value in additional_kwargs.items() if key not in {"function_call", "tool_calls"}}
|
|
has_non_tool_payload = bool(getattr(message, "content", None) or sanitized_additional_kwargs or getattr(message, "usage_metadata", None) or getattr(message, "response_metadata", None))
|
|
outputs: list[Any] = []
|
|
if self.pending_identity is not None and self.pending_identity != identity:
|
|
outputs.extend(self.flush())
|
|
if has_non_tool_payload:
|
|
visible_message = model_copy(
|
|
update={
|
|
"additional_kwargs": sanitized_additional_kwargs,
|
|
"invalid_tool_calls": [],
|
|
"tool_call_chunks": [],
|
|
"tool_calls": [],
|
|
}
|
|
)
|
|
outputs.append((visible_message, metadata))
|
|
|
|
tool_only_message = model_copy(
|
|
update={
|
|
"additional_kwargs": {},
|
|
"content": "",
|
|
"invalid_tool_calls": [],
|
|
"response_metadata": {},
|
|
"tool_calls": [],
|
|
"usage_metadata": None,
|
|
}
|
|
)
|
|
self.pending_identity = identity
|
|
self.pending_message = tool_only_message if self.pending_message is None else self.pending_message + tool_only_message
|
|
if isinstance(metadata, dict):
|
|
self.pending_metadata.update(metadata)
|
|
self.pending_count += 1
|
|
if self.pending_count >= self.batch_size:
|
|
outputs.extend(self.flush())
|
|
return outputs
|
|
|
|
def flush(self) -> list[Any]:
|
|
if self.pending_message is None:
|
|
return []
|
|
chunk = (self.pending_message, self.pending_metadata)
|
|
self.pending_identity = None
|
|
self.pending_message = None
|
|
self.pending_metadata = {}
|
|
self.pending_count = 0
|
|
return [chunk]
|
|
|
|
def finish(self) -> list[Any]:
|
|
"""Flush and release identities at a values or end-of-stream boundary.
|
|
|
|
A regular batch-size or interleaved-mode flush must retain identities
|
|
because continuation chunks commonly omit the tool name.
|
|
"""
|
|
chunks = self.flush()
|
|
self.tool_names.clear()
|
|
return chunks
|
|
|
|
|
|
# Runtime-context keys the worker owns outright. A same-named key in the
|
|
# caller's ``config['context']`` is dropped rather than merged: the Gateway
|
|
# strips ``__``-prefixed keys in build_run_config, but embedded harness callers
|
|
# have no such filter and ``deerflow_trace_id`` carries no prefix to be caught
|
|
# by it anyway.
|
|
_SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = (
|
|
frozenset(
|
|
{
|
|
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
|
|
DEERFLOW_TRACE_METADATA_KEY,
|
|
CONVERSATION_READER_CONTEXT_KEY,
|
|
SCHEDULER_CAPABILITY_CONTEXT_KEY,
|
|
AGENT_RUNS_CONTEXT_KEY,
|
|
THREAD_INCARNATION_CONTEXT_KEY,
|
|
THREAD_INCARNATION_METADATA_GUARD_KEY,
|
|
"is_subagent",
|
|
"agent_id",
|
|
"__run_loop_detection_recorder",
|
|
"__run_tool_promotion_recorder",
|
|
"__run_tool_progress_recorder",
|
|
# The Gateway pins the run's project snapshot under this key at
|
|
# admission (spec §7.1); a caller-supplied value in
|
|
# ``config['context']`` must never be merged (§12).
|
|
PROJECT_CONTEXT_KEY,
|
|
KNOWLEDGE_SCOPE_RUNTIME_KEY,
|
|
}
|
|
)
|
|
| SANDBOX_SERVER_OWNED_CONTEXT_KEYS
|
|
)
|
|
|
|
|
|
def _build_runtime_context(
|
|
thread_id: str,
|
|
run_id: str,
|
|
caller_context: Any | None,
|
|
app_config: AppConfig | None = None,
|
|
task_store: Any | None = None,
|
|
extensions: Any | None = None,
|
|
conversation_reader: Any | None = None,
|
|
agent_runs: Any | None = None,
|
|
*,
|
|
thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET,
|
|
scheduler_capability: SchedulerRunCapability | 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`` or the server-owned
|
|
keys in ``_SERVER_OWNED_RUNTIME_CONTEXT_KEYS``. 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 thread_incarnation is not _THREAD_INCARNATION_UNSET:
|
|
runtime_ctx[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation
|
|
if isinstance(caller_context, dict):
|
|
for key, value in caller_context.items():
|
|
if key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS:
|
|
continue
|
|
runtime_ctx.setdefault(key, value)
|
|
if app_config is not None:
|
|
runtime_ctx["app_config"] = app_config
|
|
if agent_runs is not None:
|
|
runtime_ctx[AGENT_RUNS_CONTEXT_KEY] = agent_runs
|
|
if conversation_reader is not None:
|
|
runtime_ctx[CONVERSATION_READER_CONTEXT_KEY] = conversation_reader
|
|
if scheduler_capability is not None:
|
|
runtime_ctx[SCHEDULER_CAPABILITY_CONTEXT_KEY] = scheduler_capability
|
|
if task_store is not None:
|
|
from deerflow_extension_api import EXTENSION_TASK_STORE_KEY
|
|
|
|
runtime_ctx[EXTENSION_TASK_STORE_KEY] = task_store
|
|
# Publish the run's extension snapshot so work dispatched during graph
|
|
# execution (task delegation) binds the same generation the lead agent was
|
|
# built with, instead of re-reading a singleton that may have been replaced
|
|
# mid-run. Written after the caller merge and popped when absent, because a
|
|
# caller-supplied value for this host-internal key is never authoritative.
|
|
from deerflow.extensions import EXTENSION_SNAPSHOT_CONTEXT_KEY
|
|
|
|
if extensions is not None:
|
|
runtime_ctx[EXTENSION_SNAPSHOT_CONTEXT_KEY] = extensions
|
|
else:
|
|
runtime_ctx.pop(EXTENSION_SNAPSHOT_CONTEXT_KEY, None)
|
|
return runtime_ctx
|
|
|
|
|
|
def _pin_admission_project_context(config: dict, runtime_context: dict[str, Any]) -> None:
|
|
"""Re-inject the admission-pinned project snapshot into the runtime context.
|
|
|
|
The Gateway resolves and stamps ``config['context'][PROJECT_CONTEXT_KEY]``
|
|
after stripping client-supplied values at admission (``build_run_config``
|
|
drops ``__``-prefixed context keys; the services pop-set covers both
|
|
sections), so a value surviving to this point is server-authoritative.
|
|
``_build_runtime_context`` refuses server-owned keys from the caller
|
|
merge, and ``_install_runtime_context`` treats the runtime context as the
|
|
authoritative view — hoisting here is what lets middleware and tools read
|
|
exactly the snapshot admission pinned.
|
|
"""
|
|
caller_context = config.get("context")
|
|
if isinstance(caller_context, dict) and PROJECT_CONTEXT_KEY in caller_context:
|
|
runtime_context[PROJECT_CONTEXT_KEY] = caller_context[PROJECT_CONTEXT_KEY]
|
|
|
|
|
|
@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)
|
|
mcp_task_repo: Any | None = field(default=None)
|
|
app_config: AppConfig | None = field(default=None)
|
|
extensions: Any | None = field(default=None)
|
|
checkpoint_channel_mode: CheckpointChannelMode = "full"
|
|
# Delta snapshot cadence frozen at startup; ``None`` means "not frozen in
|
|
# this process" (embedded/tests) and resolves to the config default.
|
|
checkpoint_snapshot_frequency: int | None = None
|
|
on_run_completed: Any | None = field(default=None)
|
|
# The host binds this capability to one run's authenticated reader and references.
|
|
conversation_reader: Any | None = field(default=None)
|
|
agent_runs: Any | None = field(default=None)
|
|
scheduler_capability: SchedulerRunCapability | None = field(default=None)
|
|
# Server-selected occurrence identity. Metadata alone never activates a goal.
|
|
scheduled_task_runtime: Mapping[str, Any] | None = field(default=None)
|
|
|
|
|
|
def _install_runtime_context(config: dict, runtime_context: dict[str, Any]) -> None:
|
|
# Configurable participates in lead-agent option merging and checkpoint
|
|
# persistence; the reader capability belongs only to host-owned context.
|
|
configurable = config.get("configurable")
|
|
if isinstance(configurable, dict):
|
|
configurable.pop(CONVERSATION_READER_CONTEXT_KEY, None)
|
|
configurable.pop(SCHEDULER_CAPABILITY_CONTEXT_KEY, None)
|
|
configurable.pop(AGENT_RUNS_CONTEXT_KEY, 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"])
|
|
# Keep both context views authoritative. A server-owned value is
|
|
# assigned from the runtime context when present and removed otherwise,
|
|
# so an embedded caller cannot preserve a forged lifecycle identity in
|
|
# ``config['context']`` after it was rejected by _build_runtime_context.
|
|
for key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS:
|
|
if key in runtime_context:
|
|
existing_context[key] = runtime_context[key]
|
|
else:
|
|
existing_context.pop(key, None)
|
|
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 _agent_graph(agent_result: Any) -> Any:
|
|
"""Unwrap the lead assembly, leaving any other factory result untouched."""
|
|
try:
|
|
from deerflow.agents.lead_agent.agent import unwrap_agent_graph
|
|
except Exception:
|
|
# A custom factory must keep working even if importing the lead
|
|
# assembly type fails.
|
|
return agent_result
|
|
return unwrap_agent_graph(agent_result)
|
|
|
|
|
|
def _assembled_model_name(agent_result: Any) -> str | None:
|
|
"""Return the selected model only for the trusted lead assembly result."""
|
|
try:
|
|
from deerflow.agents.lead_agent.agent import LeadAgentAssembly
|
|
except Exception:
|
|
return None
|
|
if isinstance(agent_result, LeadAgentAssembly):
|
|
return agent_result.effective_model
|
|
return None
|
|
|
|
|
|
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 = []
|
|
write_task = asyncio.create_task(self._event_store.put_batch(batch))
|
|
cancellation: asyncio.CancelledError | None = None
|
|
failure: Exception | None = None
|
|
while True:
|
|
try:
|
|
await asyncio.shield(write_task)
|
|
break
|
|
except asyncio.CancelledError as exc:
|
|
host_task = asyncio.current_task()
|
|
if host_task is None or not host_task.cancelling():
|
|
# The store task itself was cancelled. It did not confirm
|
|
# durability, so retain the batch for a later flush.
|
|
self._pending = batch + self._pending
|
|
raise
|
|
if cancellation is None:
|
|
cancellation = exc
|
|
while host_task.cancelling():
|
|
host_task.uncancel()
|
|
except Exception as exc:
|
|
failure = exc
|
|
break
|
|
|
|
if failure is not None:
|
|
# Re-buffer the failed batch (ahead of any events queued since) so a
|
|
# transient store error does not silently drop subagent step events.
|
|
self._pending = batch + self._pending
|
|
logger.warning(
|
|
"Run %s: failed to persist %d subagent step event(s)",
|
|
self._run_id,
|
|
len(batch),
|
|
exc_info=(type(failure), failure, failure.__traceback__),
|
|
)
|
|
if cancellation is not None:
|
|
raise cancellation
|
|
|
|
|
|
def _bind_trace_id(config: dict[str, Any], runtime_ctx: dict[str, Any]) -> str:
|
|
"""Record the current request trace id on the runtime context and metadata.
|
|
|
|
The ContextVar is the only source. A ``deerflow_trace_id`` the caller sent
|
|
in ``config["metadata"]`` is overwritten rather than read: honouring it
|
|
would let the persisted run disagree with the ``X-Trace-Id`` and the log
|
|
lines the same request already produced, which is the correlation the id
|
|
exists to provide in the first place.
|
|
|
|
The two destinations serve different purposes. ``runtime_ctx`` is the
|
|
carrier across boundaries the ContextVar does not cross -- subagent
|
|
delegation, the memory update running on a Timer/executor thread -- while
|
|
``config["metadata"]`` is persisted with the checkpoint and is what makes a
|
|
finished run traceable after the fact.
|
|
"""
|
|
trace_id = ensure_trace_id()
|
|
runtime_ctx[DEERFLOW_TRACE_METADATA_KEY] = trace_id
|
|
incoming_metadata = config.get("metadata")
|
|
# Replaced rather than mutated through: this mapping can be shared with the
|
|
# caller's request body.
|
|
merged_metadata = dict(incoming_metadata) if isinstance(incoming_metadata, dict) else {}
|
|
merged_metadata[DEERFLOW_TRACE_METADATA_KEY] = trace_id
|
|
config["metadata"] = merged_metadata
|
|
return trace_id
|
|
|
|
|
|
def _defer_finalization_interrupt(
|
|
deferred: BaseException | None,
|
|
interrupt: BaseException,
|
|
) -> BaseException:
|
|
"""Preserve the first interrupt while allowing terminal awaits to finish."""
|
|
if isinstance(interrupt, asyncio.CancelledError):
|
|
task = asyncio.current_task()
|
|
if task is not None:
|
|
while task.cancelling():
|
|
task.uncancel()
|
|
return deferred if deferred is not None else interrupt
|
|
|
|
|
|
async def _await_task_stop_after_host_cancellation(
|
|
task: asyncio.Task[None],
|
|
deferred: BaseException | None,
|
|
) -> BaseException | None:
|
|
"""Wait for one task-stop fan-out despite repeated host cancellation."""
|
|
while True:
|
|
try:
|
|
await asyncio.shield(task)
|
|
return deferred
|
|
except asyncio.CancelledError as exc:
|
|
host = asyncio.current_task()
|
|
if host is None or not host.cancelling():
|
|
# The fan-out task itself was cancelled rather than the host.
|
|
raise
|
|
deferred = _defer_finalization_interrupt(deferred, exc)
|
|
|
|
|
|
async def run_agent(
|
|
bridge: StreamBridge,
|
|
run_manager: RunManager,
|
|
record: RunRecord,
|
|
*,
|
|
ctx: RunContext,
|
|
agent_factory: Any,
|
|
graph_input: dict,
|
|
config: dict,
|
|
thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET,
|
|
stream_modes: list[str] | None = None,
|
|
stream_subgraphs: bool = False,
|
|
interrupt_before: list[str] | Literal["*"] | None = None,
|
|
interrupt_after: list[str] | Literal["*"] | None = None,
|
|
knowledge_scope: dict[str, Any] | 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
|
|
terminal_status_kwargs = {"persist": False} if event_store is not None else {}
|
|
# Staged terminal statuses commit only after receipt and duration writes;
|
|
# keep the still-active durable row's lease renewed until then.
|
|
record.terminal_commit_pending = event_store is not None
|
|
|
|
run_id = record.run_id
|
|
thread_id = record.thread_id
|
|
|
|
from deerflow_extension_api import ExtensionData, TaskInfo
|
|
|
|
from deerflow.extensions import get_loaded_extensions
|
|
from deerflow.extensions.notify import (
|
|
lead_task_id,
|
|
lead_task_outcome,
|
|
notify_task_start,
|
|
notify_task_stop,
|
|
)
|
|
|
|
extensions = ctx.extensions if ctx.extensions is not None else get_loaded_extensions()
|
|
task_store: ExtensionData | None = None
|
|
task_info: TaskInfo | None = None
|
|
deferred_finalization_interrupt: BaseException | None = None
|
|
pre_run_checkpoint_id: str | None = None
|
|
pre_run_workspace_snapshot: WorkspaceSnapshot | None = None
|
|
workspace_changes_user_id: str | None = None
|
|
workspace_excluded_dir_names: frozenset[str] | None = None
|
|
# Fail closed until a pre-run capture completes: a cancellation delivered
|
|
# before or during capture leaves no rollback point, and treating that as
|
|
# "the thread had no pre-run checkpoint" would delete a live thread's
|
|
# checkpoints on a rollback cancel.
|
|
snapshot_capture_failed = True
|
|
llm_error_fallback_message: str | None = None
|
|
checkpoint_rollback_completed = False
|
|
# Message ids checkpointed *before* this run started. The stream loop uses
|
|
# this set to mask out ``deerflow_error_fallback`` markers that belong to
|
|
# earlier runs on the same thread — without it, one stale fallback in
|
|
# history would mark every subsequent run on this thread as ``error``.
|
|
pre_existing_message_ids: set[str] = set()
|
|
|
|
# Bound agent graph accessor + captured pre-run rollback point; assigned
|
|
# inside the try block so the finally rollback path can fork the pre-run
|
|
# checkpoint lineage (see below).
|
|
accessor: CheckpointStateAccessor | None = None
|
|
rollback_point: RollbackPoint | None = None
|
|
journal = None
|
|
runtime_ctx: dict[str, Any] | None = None
|
|
runtime: Any | None = None
|
|
agent: Any | None = None
|
|
runnable_configs: list[dict[str, Any]] = [config]
|
|
goal_evaluator_model: Any | None = None
|
|
delivery_content: dict[str, Any] | None = None
|
|
goal_completion: _GoalCompletionCandidate | None = None
|
|
scheduled_goal: GoalState | None = None
|
|
final_goal_verdict: dict[str, Any] | None = None
|
|
produced_output_paths: list[str] | None = None
|
|
# Journal construction moved ahead of preflight so every terminal run can
|
|
# emit a receipt. Completion persistence keeps its prior boundary: before
|
|
# #4272 the journal did not exist until preflight had succeeded, so early
|
|
# checkpoint failures / cancellation while waiting did not write an empty
|
|
# completion snapshot into RunStore.
|
|
persist_completion = False
|
|
completion_data: dict[str, Any] | None = None
|
|
# Buffers subagent step events for batched persistence (#3779); assigned once
|
|
# streaming starts and flushed in the finally block. Pre-bound to None so the
|
|
# finally is safe even if an exception fires before streaming begins.
|
|
subagent_events: _SubagentEventBuffer | None = None
|
|
started = False
|
|
|
|
async def _finish_cancellation(
|
|
action: str,
|
|
*,
|
|
restore_checkpoint: bool = True,
|
|
) -> None:
|
|
nonlocal checkpoint_rollback_completed
|
|
await run_manager.set_finalizing(run_id, True)
|
|
if action == "rollback":
|
|
await run_manager.set_status(
|
|
run_id,
|
|
RunStatus.error,
|
|
error="Rolled back by user",
|
|
**terminal_status_kwargs,
|
|
)
|
|
if not restore_checkpoint:
|
|
return
|
|
try:
|
|
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
|
|
accessor=accessor,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
rollback_point=rollback_point,
|
|
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(
|
|
"Run %s cancellation rollback failed",
|
|
run_id,
|
|
exc_info=True,
|
|
)
|
|
else:
|
|
await run_manager.set_status(
|
|
run_id,
|
|
RunStatus.interrupted,
|
|
**terminal_status_kwargs,
|
|
)
|
|
logger.info("Run %s was cancelled", run_id)
|
|
|
|
try:
|
|
normalized_stream_modes = normalize_stream_modes(stream_modes)
|
|
requested_modes: set[str] = set(normalized_stream_modes)
|
|
lg_modes = to_langgraph_stream_modes(normalized_stream_modes)
|
|
# Initialize the run-scoped journal before any fallible or cancellable
|
|
# preflight work. Every terminal run with an event store must reach the
|
|
# shared finally block with a journal available for its run.delivery
|
|
# receipt, including checkpoint validation failures and cancellation
|
|
# while waiting for an earlier run to finish finalizing.
|
|
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),
|
|
)
|
|
|
|
# Keep cancellable preflight work under the worker's terminal guard so
|
|
# cancellation cannot strand a pending RunRecord or stream subscriber.
|
|
if ctx.mcp_task_repo is not None and record.user_id is not None:
|
|
try:
|
|
if thread_incarnation is _THREAD_INCARNATION_UNSET:
|
|
raise RuntimeError("MCP task projection requires a server-owned thread incarnation")
|
|
assert thread_incarnation is None or isinstance(thread_incarnation, str)
|
|
task_rows = await ctx.mcp_task_repo.list_by_thread(
|
|
thread_id,
|
|
user_id=record.user_id,
|
|
thread_incarnation=thread_incarnation,
|
|
limit=20,
|
|
)
|
|
graph_input = {
|
|
**graph_input,
|
|
"background_tasks": _project_background_tasks(task_rows),
|
|
}
|
|
except Exception:
|
|
logger.warning("Run %s: failed to project MCP task state", run_id, exc_info=True)
|
|
|
|
await run_manager.wait_for_prior_finalizing(
|
|
thread_id,
|
|
run_id,
|
|
abort_event=record.abort_event,
|
|
)
|
|
|
|
start_outcome = await run_manager.try_start(run_id)
|
|
if start_outcome is not RunStartOutcome.started:
|
|
if record.abort_event.is_set():
|
|
await _finish_cancellation(
|
|
record.abort_action,
|
|
restore_checkpoint=False,
|
|
)
|
|
return
|
|
started = True
|
|
|
|
task_id = lead_task_id(run_id)
|
|
if extensions.needs_task_store:
|
|
task_store = ExtensionData(task_id)
|
|
|
|
if extensions.has_task_lifecycle:
|
|
task_info = TaskInfo(
|
|
task_id=task_id,
|
|
run_id=run_id,
|
|
thread_id=thread_id,
|
|
kind="lead",
|
|
agent_name=record.assistant_id,
|
|
)
|
|
assert task_store is not None
|
|
await notify_task_start(
|
|
extensions,
|
|
task_store,
|
|
task_info,
|
|
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
|
|
)
|
|
|
|
if not record.ownership_lost and thread_store is not None:
|
|
try:
|
|
await thread_store.update_status(thread_id, "running")
|
|
except Exception:
|
|
logger.debug("Failed to update thread_meta status for %s (non-fatal)", thread_id)
|
|
mode = ctx.checkpoint_channel_mode
|
|
inject_checkpoint_mode(config, mode)
|
|
checkpoint_config = {
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": "",
|
|
}
|
|
}
|
|
if checkpointer is not None:
|
|
await aensure_checkpoint_mode_compatible(
|
|
checkpointer,
|
|
checkpoint_config,
|
|
mode,
|
|
)
|
|
configurable = config["configurable"]
|
|
selected_configurable = {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": configurable.get("checkpoint_ns", ""),
|
|
}
|
|
for selector_key in ("checkpoint_id", "checkpoint_map"):
|
|
if selector_key in configurable:
|
|
selected_configurable[selector_key] = configurable[selector_key]
|
|
selected_checkpoint_config = {
|
|
"configurable": selected_configurable,
|
|
}
|
|
if selected_checkpoint_config != checkpoint_config:
|
|
await aensure_checkpoint_mode_compatible(
|
|
checkpointer,
|
|
selected_checkpoint_config,
|
|
mode,
|
|
)
|
|
|
|
persist_completion = True
|
|
|
|
if event_store is not None:
|
|
workspace_changes_user_id = get_effective_user_id()
|
|
# Resolved once per run so the pre-run snapshot, the post-run
|
|
# delivery scan, and the workspace-changes scan all agree on the
|
|
# same exclusion set.
|
|
workspace_excluded_dir_names = _workspace_excluded_dir_names(ctx.app_config)
|
|
try:
|
|
pre_run_workspace_snapshot = await capture_workspace_snapshot(
|
|
thread_id,
|
|
user_id=workspace_changes_user_id,
|
|
extra_excluded_dir_names=workspace_excluded_dir_names,
|
|
)
|
|
except Exception:
|
|
logger.warning("Could not capture pre-run workspace 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,
|
|
task_store,
|
|
extensions,
|
|
ctx.conversation_reader,
|
|
ctx.agent_runs,
|
|
thread_incarnation=thread_incarnation,
|
|
scheduler_capability=ctx.scheduler_capability,
|
|
)
|
|
# Bind every checkpoint produced by this run to the effective agent
|
|
# identity that produced its state. Manual compaction uses only this
|
|
# server-overwritten value for memory policy; request metadata cannot
|
|
# forge it, and an explicit default sentinel distinguishes new default
|
|
# checkpoints from unbound legacy state.
|
|
if "agent_name" in runtime_ctx:
|
|
checkpoint_agent_name = runtime_ctx["agent_name"]
|
|
else:
|
|
configurable = config.get("configurable")
|
|
checkpoint_agent_name = configurable.get("agent_name") if isinstance(configurable, dict) else None
|
|
checkpoint_metadata = config.get("metadata")
|
|
if not isinstance(checkpoint_metadata, dict):
|
|
checkpoint_metadata = {}
|
|
config["metadata"] = checkpoint_metadata
|
|
checkpoint_metadata[CHECKPOINT_AGENT_NAME_METADATA_KEY] = DEFAULT_AGENT_NAME_METADATA_VALUE if checkpoint_agent_name is None else checkpoint_agent_name
|
|
_pin_admission_project_context(config, runtime_ctx)
|
|
if knowledge_scope is not None:
|
|
runtime_ctx[KNOWLEDGE_SCOPE_RUNTIME_KEY] = execution_scope(knowledge_scope)
|
|
deerflow_trace_id = _bind_trace_id(config, runtime_ctx)
|
|
# 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)
|
|
|
|
# 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)
|
|
runnable_configs.append(initial_runnable_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
|
|
continuation = RunnableConfig(**continuation_config)
|
|
runnable_configs.append(continuation)
|
|
return continuation
|
|
|
|
agent_factory_kwargs: dict[str, Any] = {"config": initial_runnable_config}
|
|
if ctx.app_config is not None and _agent_factory_supports_app_config(agent_factory):
|
|
agent_factory_kwargs["app_config"] = ctx.app_config
|
|
from deerflow.extensions import bind_agent_build_extensions
|
|
|
|
with bind_agent_build_extensions(extensions):
|
|
# Assemble off-loop: agent construction re-enters
|
|
# get_available_tools(), which may block on MCP cache
|
|
# initialization — it must not stall the calling event loop
|
|
# (issue #5172).
|
|
agent_result = await run_assembly(agent_factory, **agent_factory_kwargs)
|
|
agent = _agent_graph(agent_result)
|
|
|
|
# Assembly resolves request, agent, and authorization fallbacks. Trace the
|
|
# model that will run, rather than the model name originally requested.
|
|
effective_model = _assembled_model_name(agent_result) or record.model_name
|
|
for trace_config in (config, initial_runnable_config):
|
|
inject_langfuse_metadata(
|
|
trace_config,
|
|
thread_id=thread_id,
|
|
user_id=resolve_runtime_user_id(runtime),
|
|
assistant_id=record.assistant_id,
|
|
model_name=effective_model,
|
|
environment=os.environ.get("DEER_FLOW_ENV") or os.environ.get("ENVIRONMENT"),
|
|
deerflow_trace_id=deerflow_trace_id,
|
|
)
|
|
|
|
accessor = CheckpointStateAccessor.bind(
|
|
agent,
|
|
checkpointer,
|
|
store=store,
|
|
mode=mode,
|
|
)
|
|
|
|
# Capture the pre-run rollback point (materialized state + raw pending
|
|
# writes) before this run mutates the thread. Raw checkpoint blobs
|
|
# cannot reconstruct Delta-channel messages (their checkpoints omit
|
|
# channel_values), so rollback forks the pre-run lineage through the
|
|
# graph and needs the materialized messages up front. Any capture
|
|
# failure disables rollback: restoring an empty or partial message
|
|
# history would silently truncate the thread.
|
|
if checkpointer is not None:
|
|
# A previous successful run may still be persisting duration
|
|
# metadata after its active admission slot is released. Share its
|
|
# checkpoint lock so the rollback snapshot and any resume rewrite
|
|
# are one uninterrupted read/write sequence against the head.
|
|
async with _checkpoint_thread_lock(thread_id):
|
|
try:
|
|
rollback_point = await _capture_rollback_point(accessor, checkpointer, checkpoint_config)
|
|
except Exception:
|
|
logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True)
|
|
else:
|
|
# Only a completed capture proves what the pre-run state was.
|
|
# ``None`` now unambiguously means the thread had no
|
|
# checkpoint, so rollback may safely reset it to empty.
|
|
snapshot_capture_failed = False
|
|
if rollback_point is not None:
|
|
pre_run_checkpoint_id = rollback_point.config.get("configurable", {}).get("checkpoint_id")
|
|
pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(rollback_point.messages)})
|
|
|
|
# Resuming from an older checkpoint is a fork, and a delta fork
|
|
# materializes the abandoned sibling's writes back into state
|
|
# (#4458). Rewrite it as a linear head write *after* the rollback
|
|
# point is captured, so cancel-with-rollback still restores the
|
|
# real pre-run head rather than the rolled-back one.
|
|
resumed_messages = await _linearize_delta_checkpoint_resume(
|
|
accessor=accessor,
|
|
checkpointer=checkpointer,
|
|
config=config,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
)
|
|
if resumed_messages is not None:
|
|
# The graph now starts from the selected state, so the
|
|
# current-run message boundary is that state, not the head we
|
|
# captured for rollback.
|
|
pre_existing_message_ids = _collect_pre_existing_message_ids({"messages": list(resumed_messages)})
|
|
initial_runnable_config = RunnableConfig(**config)
|
|
runnable_configs.append(initial_runnable_config)
|
|
|
|
runtime_ctx[CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY] = frozenset(pre_existing_message_ids)
|
|
_install_runtime_context(config, runtime_ctx)
|
|
|
|
scheduled_objective = _scheduled_goal_objective(record, ctx.scheduled_task_runtime)
|
|
if scheduled_objective is not None:
|
|
# Track the instance before admitting the write: synchronous saver
|
|
# commits are drained even if host cancellation arrives mid-write.
|
|
scheduled_goal = build_goal_state(scheduled_objective, now=canonical_run_created_at(record.created_at))
|
|
record.scheduled_goal_cleanup_pending = True
|
|
terminal_status_kwargs["persist"] = False
|
|
try:
|
|
await _install_scheduled_goal(checkpointer=checkpointer, accessor=accessor, thread_id=thread_id, goal=scheduled_goal)
|
|
except GoalWriteConflict:
|
|
# Rejected before the write was admitted. Even an equal goal
|
|
# snapshot belongs to its original writer, never this run.
|
|
scheduled_goal = None
|
|
record.scheduled_goal_cleanup_pending = False
|
|
if event_store is None:
|
|
terminal_status_kwargs.pop("persist", None)
|
|
raise
|
|
|
|
if scheduled_goal is None:
|
|
await _clear_stale_scheduled_goal(record=record, run_manager=run_manager, bridge=bridge, checkpointer=checkpointer)
|
|
|
|
# Full-mode replay keeps its historical message fork, but the owner's
|
|
# current goal and budget must win. Read after scheduled-goal preflight
|
|
# so a cleaned-up occurrence goal cannot be reintroduced by the input.
|
|
configurable = config.get("configurable", {})
|
|
selected_checkpoint_id = configurable.get("checkpoint_id")
|
|
if checkpointer is not None and mode == "full" and isinstance(graph_input, dict) and isinstance(selected_checkpoint_id, str) and selected_checkpoint_id and not configurable.get("checkpoint_ns"):
|
|
async with _checkpoint_thread_lock(thread_id):
|
|
current = await accessor.aget(checkpoint_config)
|
|
if _checkpoint_id(current) != selected_checkpoint_id:
|
|
current_values = getattr(current, "values", None) or {}
|
|
graph_input = dict(graph_input)
|
|
if graph_input.get("goal") is None:
|
|
# merge_goal ignores plain None, including a goal that
|
|
# the user cleared or completed after the selected turn.
|
|
graph_input["goal"] = Overwrite(copy.deepcopy(current_values.get("goal")))
|
|
graph_input["goal_outcome"] = None
|
|
|
|
# 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
|
|
|
|
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)
|
|
|
|
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
|
|
|
|
# Built once per run, not per _stream_once call: goal continuations
|
|
# re-enter the stream and would otherwise discard the resolved seqs.
|
|
seq_stamper = _build_seq_stamper(event_store, thread_id, journal) if "values" in requested_modes else None
|
|
|
|
async def _stream_once(input_payload: Any, stream_config: RunnableConfig) -> None:
|
|
nonlocal llm_error_fallback_message
|
|
file_tool_chunk_batcher = _LargeFileToolChunkBatcher() if "messages-tuple" in requested_modes else None
|
|
try:
|
|
async with _checkpoint_thread_lock(thread_id):
|
|
if len(lg_modes) == 1 and not stream_subgraphs:
|
|
# Single mode, no subgraphs: astream yields raw chunks
|
|
single_mode = lg_modes[0]
|
|
stream = agent.astream(input_payload, config=stream_config, stream_mode=single_mode)
|
|
broke_on_abort = False
|
|
try:
|
|
async for chunk in stream:
|
|
if record.abort_event.is_set():
|
|
broke_on_abort = True
|
|
logger.info("Run %s abort requested — stopping", run_id)
|
|
break
|
|
if single_mode != "custom":
|
|
# Custom frames carry task_* events whose payload can hold a delegated
|
|
# subagent's messages; see the multi-mode branch below.
|
|
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids)
|
|
sse_event = _lg_mode_to_sse_event(single_mode)
|
|
single_payload = serialize(chunk, mode=single_mode)
|
|
if single_mode == "values" and seq_stamper is not None:
|
|
single_payload = await seq_stamper.stamp(single_payload)
|
|
await bridge.publish(run_id, sse_event, single_payload)
|
|
if single_mode == "custom":
|
|
await subagent_events.add(chunk)
|
|
finally:
|
|
close_error = sys.exception()
|
|
try:
|
|
await close_agent_stream(stream)
|
|
except asyncio.CancelledError as exc:
|
|
_log_cancelled_stream_close_failure(
|
|
exc,
|
|
run_id=run_id,
|
|
abort_requested=broke_on_abort or record.abort_event.is_set(),
|
|
)
|
|
raise
|
|
except Exception:
|
|
abort_requested = broke_on_abort or record.abort_event.is_set()
|
|
if close_error is None and not abort_requested:
|
|
raise
|
|
if abort_requested:
|
|
logger.warning("Could not close aborted agent stream for run %s", run_id, exc_info=True)
|
|
else:
|
|
logger.debug("Could not close agent stream for run %s", run_id, exc_info=True)
|
|
return
|
|
# Multiple modes or subgraphs: astream yields tuples
|
|
stream = agent.astream(
|
|
input_payload,
|
|
config=stream_config,
|
|
stream_mode=lg_modes,
|
|
subgraphs=stream_subgraphs,
|
|
)
|
|
broke_on_abort = False
|
|
try:
|
|
async for item in stream:
|
|
if record.abort_event.is_set():
|
|
broke_on_abort = True
|
|
logger.info("Run %s abort requested — stopping", run_id)
|
|
break
|
|
|
|
mode, chunk, namespace = _unpack_stream_item(item, lg_modes, stream_subgraphs)
|
|
if mode is None:
|
|
continue
|
|
|
|
if not namespace and mode != "custom":
|
|
# Only root-graph frames may decide the parent run's error
|
|
# fallback: a delegated subagent's marked fallback is the
|
|
# executor's to map (task_failed), not this run's. That
|
|
# includes the child messages task_running custom events
|
|
# carry, which are root frames too.
|
|
llm_error_fallback_message = llm_error_fallback_message or _extract_llm_error_fallback_message(chunk, pre_existing_message_ids)
|
|
await _publish_stream_item(
|
|
bridge=bridge,
|
|
run_id=run_id,
|
|
mode=mode,
|
|
chunk=chunk,
|
|
namespace=namespace,
|
|
file_tool_chunk_batcher=file_tool_chunk_batcher,
|
|
subagent_events=subagent_events,
|
|
seq_stamper=seq_stamper,
|
|
)
|
|
finally:
|
|
close_error = sys.exception()
|
|
try:
|
|
await close_agent_stream(stream)
|
|
except asyncio.CancelledError as exc:
|
|
_log_cancelled_stream_close_failure(
|
|
exc,
|
|
run_id=run_id,
|
|
abort_requested=broke_on_abort or record.abort_event.is_set(),
|
|
)
|
|
raise
|
|
except Exception:
|
|
abort_requested = broke_on_abort or record.abort_event.is_set()
|
|
if close_error is None and not abort_requested:
|
|
raise
|
|
if abort_requested:
|
|
logger.warning("Could not close aborted agent stream for run %s", run_id, exc_info=True)
|
|
else:
|
|
logger.debug("Could not close agent stream for run %s", run_id, exc_info=True)
|
|
finally:
|
|
stream_error = sys.exception()
|
|
if file_tool_chunk_batcher is not None:
|
|
try:
|
|
for publish_chunk in file_tool_chunk_batcher.finish():
|
|
await bridge.publish(run_id, "messages", serialize(publish_chunk, mode="messages"))
|
|
except Exception:
|
|
if stream_error is None:
|
|
raise
|
|
logger.debug("Could not flush pending file-tool chunks for run %s", run_id, exc_info=True)
|
|
|
|
# 7. Stream the requested turn, then optionally continue hidden goal turns.
|
|
# Clear any stale stop_reason before the first (user-visible) turn only.
|
|
# Continuation turns preserve a cap reason from the user turn: a run that
|
|
# hits a cap during the user turn IS capped even if hidden goal-evaluator
|
|
# turns complete cleanly afterward (#4176 review).
|
|
if isinstance(runtime.context, dict):
|
|
runtime.context.pop("stop_reason", None)
|
|
|
|
def capture_goal_verdict(evaluated_goal: GoalState, evaluation: GoalEvaluation, stand_down_reason: str | None) -> None:
|
|
nonlocal final_goal_verdict
|
|
if scheduled_goal is not None and not _goal_instance_matches(scheduled_goal, evaluated_goal):
|
|
return
|
|
final_goal_verdict = dict(evaluation)
|
|
final_goal_verdict["relied_on_assumption"] = bool(evaluation.get("relied_on_assumption", False)) if evaluation["satisfied"] else False
|
|
# Hidden goal turns this run took; the task page shows the count.
|
|
final_goal_verdict["continuations"] = int(evaluated_goal.get("continuation_count", 0))
|
|
if stand_down_reason:
|
|
final_goal_verdict["stand_down_reason"] = stand_down_reason
|
|
if stand_down_reason.startswith("thread_changed"):
|
|
final_goal_verdict.update(satisfied=False, blocker="missing_evidence", relied_on_assumption=False)
|
|
|
|
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,
|
|
accessor=accessor,
|
|
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,
|
|
user_id=resolve_runtime_user_id(runtime),
|
|
deerflow_trace_id=deerflow_trace_id,
|
|
task_store=task_store,
|
|
extensions=extensions,
|
|
run_stop_reason=runtime.context.get("stop_reason") if isinstance(runtime.context, dict) else None,
|
|
interaction_policy=resolve_run_interaction_policy(config),
|
|
expected_goal=scheduled_goal,
|
|
verdict_callback=capture_goal_verdict,
|
|
usage_callback=journal.record_external_llm_usage_records if journal is not None else None,
|
|
)
|
|
if isinstance(continuation_input, _GoalCompletionCandidate):
|
|
goal_completion = continuation_input
|
|
break
|
|
if continuation_input is None or record.abort_event.is_set():
|
|
break
|
|
await _stream_once(continuation_input, _continuation_runnable_config())
|
|
|
|
record.goal_verdict = final_goal_verdict
|
|
|
|
# 8. Final status
|
|
if record.abort_event.is_set():
|
|
await _finish_cancellation(record.abort_action)
|
|
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 _ensure_finalizing_before_edit_failure(run_manager, record)
|
|
cancel_action = await run_manager.set_status_if_not_cancelled(
|
|
run_id,
|
|
RunStatus.error,
|
|
error=error_msg,
|
|
**terminal_status_kwargs,
|
|
)
|
|
if cancel_action is not None:
|
|
await _finish_cancellation(cancel_action)
|
|
else:
|
|
runtime_context = runtime.context if isinstance(runtime.context, dict) else None
|
|
# Guard middlewares that hard-stop a run by stripping tool_calls
|
|
# stamp stop_reason into runtime.context so the worker can surface
|
|
# it on the run record:
|
|
# loop_detection -> "loop_capped"
|
|
# token_budget -> "token_capped"
|
|
# safety_finish_reason -> "safety_capped"
|
|
# subagent_limit -> "subagent_limit_capped"
|
|
# model_length_finish_reason -> "model_length_capped"
|
|
#
|
|
# If more guards grow stop_reason semantics, consider a publish/
|
|
# collect pattern (e.g. each guard middleware publishes its cap
|
|
# reason to a dedicated runtime.context channel, and the worker
|
|
# collects the most severe / first / all reasons) instead of each
|
|
# guard writing directly to the same key.
|
|
stop_reason = runtime_context.get("stop_reason") if runtime_context is not None else None
|
|
produced_output_paths = await _produced_output_paths(
|
|
pre_run_workspace_snapshot,
|
|
thread_id=thread_id,
|
|
user_id=workspace_changes_user_id,
|
|
extra_excluded_dir_names=workspace_excluded_dir_names,
|
|
)
|
|
delivery_content = _delivery_content_with_outputs(
|
|
journal.get_delivery_content() if journal is not None else _empty_delivery_content(),
|
|
produced_output_paths,
|
|
)
|
|
delivery_error = _delivery_error(delivery_content)
|
|
cancel_action = await run_manager.set_status_if_not_cancelled(
|
|
run_id,
|
|
RunStatus.error if delivery_error else RunStatus.success,
|
|
error=delivery_error,
|
|
stop_reason=stop_reason,
|
|
goal_verdict=record.goal_verdict,
|
|
**terminal_status_kwargs,
|
|
)
|
|
if cancel_action is not None:
|
|
await _finish_cancellation(cancel_action)
|
|
|
|
except asyncio.CancelledError:
|
|
await _finish_cancellation(record.abort_action)
|
|
|
|
except Exception as exc:
|
|
error_msg = f"{exc}"
|
|
logger.exception("Run %s failed: %s", run_id, error_msg)
|
|
await _ensure_finalizing_before_edit_failure(run_manager, record)
|
|
cancel_action = await run_manager.set_status_if_not_cancelled(
|
|
run_id,
|
|
RunStatus.error,
|
|
error=error_msg,
|
|
**terminal_status_kwargs,
|
|
)
|
|
if cancel_action is not None:
|
|
await _finish_cancellation(cancel_action)
|
|
else:
|
|
await bridge.publish(
|
|
run_id,
|
|
"error",
|
|
{
|
|
"message": error_msg,
|
|
"name": type(exc).__name__,
|
|
},
|
|
)
|
|
|
|
finally:
|
|
try:
|
|
if record.ownership_lost:
|
|
logger.warning(
|
|
"Skipping durable finalization for run %s because this worker no longer owns its lease",
|
|
run_id,
|
|
)
|
|
|
|
if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success:
|
|
if not record.finalizing:
|
|
await run_manager.set_finalizing(run_id, True)
|
|
try:
|
|
if not checkpoint_rollback_completed:
|
|
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
|
|
accessor=accessor,
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
rollback_point=rollback_point,
|
|
snapshot_capture_failed=snapshot_capture_failed,
|
|
)
|
|
if checkpoint_rollback_completed:
|
|
await _publish_restored_checkpoint_values(
|
|
bridge=bridge,
|
|
run_id=run_id,
|
|
accessor=accessor,
|
|
thread_id=thread_id,
|
|
)
|
|
logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
|
|
except Exception:
|
|
logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True)
|
|
|
|
# Persist any subagent step events still buffered (#3779) — including on
|
|
# abort/exception paths, where the stream loop broke before its own flush.
|
|
if not record.ownership_lost and subagent_events is not None:
|
|
await subagent_events.flush()
|
|
|
|
if not record.ownership_lost and event_store is not None and pre_run_workspace_snapshot is not None:
|
|
try:
|
|
await record_workspace_changes(
|
|
event_store,
|
|
thread_id,
|
|
run_id,
|
|
pre_run_workspace_snapshot,
|
|
user_id=workspace_changes_user_id,
|
|
extra_excluded_dir_names=workspace_excluded_dir_names,
|
|
)
|
|
except Exception:
|
|
logger.warning("Failed to record workspace changes for run %s", run_id, exc_info=True)
|
|
|
|
# Flush buffered journal events before the terminal receipt. The
|
|
# receipt uses a run-scoped idempotent write shared with recovery, then
|
|
# the staged terminal status is persisted. This ordering closes the
|
|
# crash window where a terminal run could otherwise outlive its receipt.
|
|
# A fenced worker leaves receipt recovery to the peer that claimed it.
|
|
if not record.ownership_lost and journal is not None:
|
|
try:
|
|
await journal.flush()
|
|
except Exception:
|
|
logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
|
|
|
|
if delivery_content is None:
|
|
if produced_output_paths is None:
|
|
produced_output_paths = await _produced_output_paths(
|
|
pre_run_workspace_snapshot,
|
|
thread_id=thread_id,
|
|
user_id=workspace_changes_user_id,
|
|
extra_excluded_dir_names=workspace_excluded_dir_names,
|
|
)
|
|
delivery_content = _delivery_content_with_outputs(journal.get_delivery_content(), produced_output_paths)
|
|
receipt_persisted = await _persist_delivery_receipt(
|
|
event_store,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
content=delivery_content,
|
|
)
|
|
if produced_output_paths and record.status == RunStatus.success and not receipt_persisted:
|
|
await run_manager.set_status(
|
|
run_id,
|
|
RunStatus.error,
|
|
error=_DELIVERY_RECEIPT_FAILED_ERROR,
|
|
persist=False,
|
|
)
|
|
|
|
if not record.ownership_lost and journal is not None and persist_completion:
|
|
try:
|
|
# Advance the final completion fields and timestamp without
|
|
# terminalizing the durable row. That active row continues to
|
|
# fence peer checkpoint writers through the duration write.
|
|
completion_data = journal.get_completion_data()
|
|
await run_manager.update_finalizing_progress(run_id, **completion_data)
|
|
except Exception:
|
|
logger.warning("Failed to persist finalizing run progress for %s (non-fatal)", run_id, exc_info=True)
|
|
|
|
# Keep the durable run row active through its final duration checkpoint
|
|
# write. A peer Gateway admits history migration from the durable row,
|
|
# not this worker's staged terminal status; terminalizing first would
|
|
# let that migration read an unfinished lifetime and race this write.
|
|
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.success:
|
|
try:
|
|
created = datetime.fromisoformat(record.created_at.replace("Z", "+00:00"))
|
|
updated = datetime.fromisoformat(record.updated_at.replace("Z", "+00:00"))
|
|
# Match legacy history semantics: turn_duration is the whole
|
|
# RunRecord lifetime in integer seconds, including admission
|
|
# delay. Persist zero for sub-second successful turns.
|
|
duration = max(0, int((updated - created).total_seconds()))
|
|
await _persist_run_duration(
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
run_id=run_id,
|
|
duration_seconds=duration,
|
|
)
|
|
except Exception:
|
|
logger.debug("Failed to persist run duration for thread %s run %s (non-fatal)", thread_id, run_id)
|
|
|
|
if scheduled_goal is not None and not record.ownership_lost and checkpointer is not None:
|
|
if record.abort_event.is_set() and record.abort_action == "rollback" and not checkpoint_rollback_completed:
|
|
await _finish_cancellation("rollback")
|
|
# Scheduled goals are occurrence-scoped. Keep the real run slot
|
|
# active through cleanup so a peer's next user turn cannot
|
|
# inherit this goal. Ordinary goals retain terminal-first cleanup.
|
|
cleanup_task = asyncio.create_task(
|
|
_clear_scheduled_goal(goal=scheduled_goal, record=record, run_manager=run_manager, bridge=bridge, checkpointer=checkpointer, owned_admission=True),
|
|
name=f"scheduled-goal-cleanup-{run_id}",
|
|
)
|
|
try:
|
|
deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation(cleanup_task, deferred_finalization_interrupt)
|
|
except Exception:
|
|
logger.warning("Could not remove scheduled goal for run %s", run_id, exc_info=True)
|
|
await run_manager.set_status(run_id, RunStatus.error, error="Scheduled goal cleanup failed", persist=False)
|
|
|
|
if not record.ownership_lost and (event_store is not None or scheduled_goal is not None):
|
|
try:
|
|
# Even after bounded receipt retries are exhausted, persist the
|
|
# real worker outcome. Leaving a successful row inflight would
|
|
# let lease recovery rewrite it as an error with a synthetic
|
|
# zero receipt.
|
|
if record.abort_event.is_set():
|
|
if scheduled_goal is not None:
|
|
# A cancel may arrive while scheduled cleanup was
|
|
# awaiting its saver. The local status may still be
|
|
# staged success; apply the accepted action, including
|
|
# rollback, before committing any terminal outcome.
|
|
await _finish_cancellation(record.abort_action, restore_checkpoint=not checkpoint_rollback_completed)
|
|
await run_manager.persist_current_status(run_id)
|
|
else:
|
|
cancel_action = await run_manager.set_status_if_not_cancelled(
|
|
run_id,
|
|
record.status,
|
|
error=record.error,
|
|
stop_reason=record.stop_reason,
|
|
goal_verdict=record.goal_verdict,
|
|
)
|
|
if cancel_action is not None:
|
|
await _finish_cancellation(cancel_action)
|
|
await run_manager.persist_current_status(run_id)
|
|
except Exception:
|
|
logger.warning("Failed to persist terminal status for run %s after delivery receipt attempts", run_id, exc_info=True)
|
|
# The deferred commit has been attempted. A failed write is left to
|
|
# completion fallback or lease recovery, as before.
|
|
record.terminal_commit_pending = False
|
|
|
|
if not record.ownership_lost and journal is not None and persist_completion:
|
|
try:
|
|
# Persist token usage + convenience fields to RunStore
|
|
completion_data = completion_data or journal.get_completion_data()
|
|
await run_manager.update_run_completion(run_id, status=record.status.value, **completion_data)
|
|
except Exception:
|
|
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
|
|
|
|
if scheduled_goal is not None and not record.ownership_lost:
|
|
if not await run_manager.persist_current_status(run_id):
|
|
await run_manager._mark_ownership_lost(record, reason="Scheduled run terminal status could not be confirmed after goal cleanup.", require_active=False)
|
|
# Late cancellation can still drain a rollback after goal
|
|
# cleanup. Retain admission/heartbeat protection through that
|
|
# rollback and the actual durable terminal acknowledgement.
|
|
record.scheduled_goal_cleanup_pending = False
|
|
|
|
# A satisfied evaluator is only a candidate until artifact delivery,
|
|
# receipt persistence and durable cancellation arbitration have ended.
|
|
# Status writes are best-effort in single-worker mode: confirm the
|
|
# existing outcome before deleting recoverable goal state.
|
|
if scheduled_goal is None and goal_completion is not None and record.status == RunStatus.success and not record.abort_event.is_set() and not record.ownership_lost:
|
|
try:
|
|
if await run_manager.persist_current_status(run_id):
|
|
await _clear_completed_goal(
|
|
candidate=goal_completion,
|
|
record=record,
|
|
run_manager=run_manager,
|
|
bridge=bridge,
|
|
accessor=accessor,
|
|
checkpointer=checkpointer,
|
|
)
|
|
except Exception:
|
|
logger.warning("Could not finalize satisfied goal for thread %s", thread_id, exc_info=True)
|
|
|
|
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record):
|
|
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 started and not record.ownership_lost and 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 started and not record.ownership_lost and 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 started and not record.ownership_lost and accessor is not None and record.status == RunStatus.success and not record.abort_event.is_set() and record.human_input_request is None:
|
|
try:
|
|
from deerflow.agents.human_input import extract_unanswered_human_input_request
|
|
|
|
record.human_input_request = extract_unanswered_human_input_request(await _materialized_checkpoint_messages(accessor, thread_id))
|
|
except Exception:
|
|
logger.debug("Failed to inspect trailing human_input_request for run %s (non-fatal)", run_id, exc_info=True)
|
|
|
|
if not record.ownership_lost and not getattr(record, "terminal_status_persisted", True):
|
|
try:
|
|
await run_manager.persist_current_status(run_id)
|
|
except Exception:
|
|
logger.warning("Failed to confirm terminal status for run %s before completion hook", run_id, exc_info=True)
|
|
|
|
if not record.ownership_lost and getattr(record, "terminal_status_persisted", True) and ctx.on_run_completed is not None:
|
|
try:
|
|
await ctx.on_run_completed(record)
|
|
except Exception:
|
|
logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True)
|
|
except BaseException as exc:
|
|
# A terminal hook must not leave replacement runs blocked or
|
|
# stream consumers waiting indefinitely.
|
|
deferred_finalization_interrupt = _defer_finalization_interrupt(
|
|
deferred_finalization_interrupt,
|
|
exc,
|
|
)
|
|
logger.warning(
|
|
"Run completion hook interrupted for %s; completing finalization first",
|
|
run_id,
|
|
)
|
|
|
|
if task_info is not None and task_store is not None:
|
|
# Keep the finalizing barrier held until stop observers finish, so
|
|
# a same-thread replacement cannot overlap this task's lifecycle.
|
|
task_stop = asyncio.create_task(
|
|
notify_task_stop(
|
|
extensions,
|
|
task_store,
|
|
task_info,
|
|
lead_task_outcome(
|
|
aborted=(record.abort_event.is_set() or record.status == RunStatus.interrupted),
|
|
succeeded=record.status == RunStatus.success,
|
|
),
|
|
timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS,
|
|
),
|
|
name=f"extension-task-stop-{run_id}",
|
|
)
|
|
try:
|
|
deferred_finalization_interrupt = await _await_task_stop_after_host_cancellation(
|
|
task_stop,
|
|
deferred_finalization_interrupt,
|
|
)
|
|
except Exception:
|
|
logger.warning(
|
|
"Extension task-stop notification failed for run %s (non-fatal)",
|
|
run_id,
|
|
exc_info=True,
|
|
)
|
|
except BaseException as exc:
|
|
# Cancellation here must not strand the finalizing barrier or
|
|
# leave stream consumers waiting for the end frame.
|
|
deferred_finalization_interrupt = _defer_finalization_interrupt(
|
|
deferred_finalization_interrupt,
|
|
exc,
|
|
)
|
|
logger.warning(
|
|
"Extension task-stop notification interrupted for run %s; completing cleanup first",
|
|
run_id,
|
|
)
|
|
if record.finalizing:
|
|
await run_manager.set_finalizing(run_id, False)
|
|
|
|
await bridge.publish_end(run_id)
|
|
|
|
if deferred_finalization_interrupt is not None:
|
|
raise deferred_finalization_interrupt
|
|
finally:
|
|
try:
|
|
if journal is not None:
|
|
try:
|
|
await journal.close(flush=not record.ownership_lost)
|
|
except Exception:
|
|
logger.warning("Failed to close journal for run %s", run_id, exc_info=True)
|
|
finally:
|
|
lease_cleanup_interrupt: BaseException | None = None
|
|
try:
|
|
from deerflow.sandbox.lease import release_sandbox_execution_lease_async
|
|
|
|
await release_sandbox_execution_lease_async(runtime_ctx)
|
|
except Exception:
|
|
logger.warning("Failed to release sandbox execution lease for run %s", run_id, exc_info=True)
|
|
except BaseException as exc:
|
|
# release_async completes the underlying cleanup before it
|
|
# re-raises cancellation. Defer that interruption until the
|
|
# worker has dropped all other run-scoped references too.
|
|
lease_cleanup_interrupt = exc
|
|
logger.warning(
|
|
"Sandbox execution lease cleanup was interrupted for run %s; completing local cleanup first",
|
|
run_id,
|
|
)
|
|
finally:
|
|
_release_run_scoped_references(
|
|
runnable_configs,
|
|
runtime_ctx,
|
|
journal,
|
|
)
|
|
# Drop graph and per-run payload references before the terminal
|
|
# worker task itself becomes collectable.
|
|
agent = None
|
|
agent_result = None
|
|
accessor = None
|
|
runtime = None
|
|
runtime_ctx = None
|
|
rollback_point = None
|
|
subagent_events = None
|
|
goal_evaluator_model = None
|
|
goal_completion = None
|
|
task_store = None
|
|
task_info = None
|
|
pre_run_workspace_snapshot = None
|
|
delivery_content = None
|
|
produced_output_paths = None
|
|
graph_input = {}
|
|
|
|
# Durable finalization and terminal publication may depend on
|
|
# external backends, but local housekeeping must always run.
|
|
record.scheduled_goal_cleanup_pending = False
|
|
_create_contextless_task(bridge.cleanup(run_id, delay=60))
|
|
# Preserve the existing five-minute grace period for local
|
|
# join/status paths, then release the terminal record, completed
|
|
# task, and request payload. Durable run history remains available
|
|
# through RunStore.
|
|
_create_contextless_task(run_manager.cleanup(run_id))
|
|
_schedule_terminal_cycle_collection()
|
|
|
|
if lease_cleanup_interrupt is not None:
|
|
raise lease_cleanup_interrupt
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _GoalCompletionCandidate:
|
|
goal: GoalState
|
|
conversation_signature: str
|
|
evaluation: GoalEvaluation
|
|
|
|
|
|
def _scheduled_goal_objective(record: RunRecord, scheduled_runtime: Mapping[str, Any] | None) -> str | None:
|
|
if scheduled_runtime is None or scheduled_runtime.get("goal_objective") is None:
|
|
return None
|
|
metadata = record.metadata or {}
|
|
if not record.user_id or scheduled_runtime.get("user_id") != record.user_id or scheduled_runtime.get("task_id") != metadata.get("scheduled_task_id") or scheduled_runtime.get("occurrence_id") != metadata.get("scheduled_task_run_id"):
|
|
raise ValueError("Scheduled goal authority does not match the admitted occurrence")
|
|
objective = metadata.get("scheduled_goal_objective")
|
|
if not isinstance(objective, str) or objective != scheduled_runtime.get("goal_objective"):
|
|
raise ValueError("Scheduled goal snapshot does not match the admitted objective")
|
|
return objective
|
|
|
|
|
|
async def _install_scheduled_goal(*, checkpointer: Any, accessor: CheckpointStateAccessor, thread_id: str, goal: GoalState) -> None:
|
|
if checkpointer is None:
|
|
raise RuntimeError("Scheduled goals require a checkpointer")
|
|
async with goal_thread_lock(thread_id), _checkpoint_thread_lock(thread_id):
|
|
checkpoint = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
|
|
if checkpoint is not None:
|
|
if _read_checkpoint_goal(checkpoint) is not None or any(_is_visible_message(message) for message in await _materialized_checkpoint_messages(accessor, thread_id)):
|
|
raise GoalWriteConflict("A scheduled goal requires a fresh conversation")
|
|
await write_thread_goal(checkpointer, thread_id, goal, as_node="scheduled_goal", create_if_missing=True, expected_checkpoint_id=_checkpoint_id(checkpoint))
|
|
|
|
|
|
async def _clear_scheduled_goal(*, goal: GoalState, record: RunRecord, run_manager: RunManager, bridge: StreamBridge | None, checkpointer: Any, owned_admission: bool = False) -> None:
|
|
"""Release only this occurrence's goal under a real run/checkpoint slot."""
|
|
try:
|
|
reservation = _checkpoint_thread_lock(record.thread_id) if owned_admission else run_manager.reserve_thread_operation(record.thread_id, kind=ThreadOperationKind.checkpoint_write, user_id=record.user_id)
|
|
async with goal_thread_lock(record.thread_id), reservation:
|
|
if owned_admission and not await run_manager.owns_active_admission(record):
|
|
return
|
|
checkpoint = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", {"configurable": {"thread_id": record.thread_id, "checkpoint_ns": ""}})
|
|
if checkpoint is None or not _goal_instance_matches(goal, _read_checkpoint_goal(checkpoint)) or record.ownership_lost:
|
|
return
|
|
values = await write_thread_goal(checkpointer, record.thread_id, None, as_node="scheduled_goal", expected_checkpoint_id=_checkpoint_id(checkpoint))
|
|
if bridge is not None:
|
|
await bridge.publish(record.run_id, "values", serialize(values, mode="values"))
|
|
except (GoalWriteConflict, ConflictError):
|
|
return
|
|
|
|
|
|
async def clear_recovered_scheduled_goal(record: RunRecord, *, run_manager: RunManager, checkpointer: Any) -> None:
|
|
"""Clean an idle recovered occurrence; an admitted new run wins safely.
|
|
|
|
Verify durable owner/thread/timestamp/objective identity before acquiring
|
|
the existing checkpoint reservation. Busy threads defer to worker preflight.
|
|
No ORM writer transaction spans a checkpointer write.
|
|
"""
|
|
metadata = record.metadata or {}
|
|
objective = metadata.get("scheduled_goal_objective")
|
|
if checkpointer is None or not record.user_id or not isinstance(objective, str):
|
|
return
|
|
try:
|
|
timestamp = canonical_run_created_at(record.created_at)
|
|
expected = build_goal_state(objective, now=timestamp)
|
|
source = await run_manager.scheduled_goal_source(record.thread_id, user_id=record.user_id, created_at=timestamp, objective=expected["objective"])
|
|
if source is None or source.run_id != record.run_id:
|
|
return
|
|
await _clear_scheduled_goal(goal=expected, record=record, run_manager=run_manager, bridge=None, checkpointer=checkpointer)
|
|
except Exception:
|
|
logger.warning("Unable to remove recovered scheduled goal for run %s; preserving it for guarded worker preflight", record.run_id, exc_info=True)
|
|
|
|
|
|
async def _clear_stale_scheduled_goal(*, record: RunRecord, run_manager: RunManager, bridge: StreamBridge, checkpointer: Any) -> None:
|
|
"""Before a new turn, remove a provably terminal predecessor's goal only."""
|
|
if checkpointer is None or not record.user_id:
|
|
return
|
|
async with goal_thread_lock(record.thread_id), _checkpoint_thread_lock(record.thread_id):
|
|
checkpoint = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", {"configurable": {"thread_id": record.thread_id, "checkpoint_ns": ""}})
|
|
current = _read_checkpoint_goal(checkpoint) if checkpoint is not None else None
|
|
if current is None:
|
|
return
|
|
try:
|
|
created_at = canonical_run_created_at(current.get("created_at", ""))
|
|
except (TypeError, ValueError):
|
|
return
|
|
if current.get("created_at") != created_at:
|
|
# Only this canonical representation is emitted by the scheduled
|
|
# installer; equivalent user timestamps are not provenance.
|
|
return
|
|
try:
|
|
source = await run_manager.scheduled_goal_source(record.thread_id, user_id=record.user_id, created_at=created_at, objective=current["objective"])
|
|
except Exception as exc:
|
|
# Idle recovery can defer cleanup, but a new run must not inherit
|
|
# an unverified goal. Keep cancellation and the checkpoint intact.
|
|
raise RuntimeError("Scheduled goal recovery could not verify the existing goal's ownership. Agent execution did not start; the goal was preserved.") from exc
|
|
if source is None or not await run_manager.owns_active_admission(record):
|
|
return
|
|
try:
|
|
values = await write_thread_goal(checkpointer, record.thread_id, None, as_node="scheduled_goal_recovery", expected_checkpoint_id=_checkpoint_id(checkpoint))
|
|
await bridge.publish(record.run_id, "values", serialize(values, mode="values"))
|
|
except GoalWriteConflict:
|
|
return
|
|
|
|
|
|
async def _clear_completed_goal(
|
|
*,
|
|
candidate: _GoalCompletionCandidate,
|
|
record: RunRecord,
|
|
run_manager: RunManager,
|
|
bridge: StreamBridge,
|
|
accessor: CheckpointStateAccessor,
|
|
checkpointer: Any,
|
|
) -> None:
|
|
"""Clear only the evaluated goal after the run has successfully finalized."""
|
|
try:
|
|
# Terminal status has released the run's admission. Use the same durable
|
|
# reservation as other out-of-run checkpoint writers; a peer that already
|
|
# admitted a new run wins, even before it has written any checkpoint.
|
|
async with (
|
|
goal_thread_lock(record.thread_id),
|
|
run_manager.reserve_thread_operation(record.thread_id, kind=ThreadOperationKind.checkpoint_write, user_id=record.user_id),
|
|
):
|
|
checkpoint_tuple = await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aget_tuple",
|
|
"get_tuple",
|
|
{"configurable": {"thread_id": record.thread_id, "checkpoint_ns": ""}},
|
|
)
|
|
if checkpoint_tuple is None:
|
|
logger.debug("Skipping goal completion for run %s: checkpoint is unavailable", record.run_id)
|
|
return
|
|
# Full equality is intentional: even same-instance goal updates must win over clearing.
|
|
if _read_checkpoint_goal(checkpoint_tuple) != candidate.goal:
|
|
logger.debug("Skipping goal completion for run %s: goal snapshot changed", record.run_id)
|
|
return
|
|
messages = await _materialized_checkpoint_messages(accessor, record.thread_id)
|
|
if visible_conversation_signature(messages) != candidate.conversation_signature:
|
|
logger.debug("Skipping goal completion for run %s: visible conversation changed", record.run_id)
|
|
return
|
|
if record.status != RunStatus.success or record.abort_event.is_set() or record.ownership_lost:
|
|
logger.debug(
|
|
"Skipping goal completion for run %s: status=%s, aborted=%s, ownership_lost=%s",
|
|
record.run_id,
|
|
record.status.value,
|
|
record.abort_event.is_set(),
|
|
record.ownership_lost,
|
|
)
|
|
return
|
|
# Duration bookkeeping may advance the checkpoint after evaluation.
|
|
# Compare goal/conversation above, then guard against stale writes.
|
|
# The record lands in the same checkpoint as the clear, or not at all.
|
|
values = await write_thread_goal(
|
|
checkpointer,
|
|
record.thread_id,
|
|
None,
|
|
as_node="goal_evaluator",
|
|
expected_checkpoint_id=_checkpoint_id(checkpoint_tuple),
|
|
outcome=build_goal_outcome(candidate.goal, candidate.evaluation, reply_message_id=latest_visible_assistant_message_id(messages)),
|
|
)
|
|
await bridge.publish(record.run_id, "values", serialize(values, mode="values"))
|
|
except GoalWriteConflict:
|
|
logger.debug("Skipping goal completion for run %s: checkpoint changed before the goal write", record.run_id)
|
|
return
|
|
except ConflictError:
|
|
return
|
|
|
|
|
|
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
|
|
|
|
|
|
async def _materialized_checkpoint_messages(accessor: CheckpointStateAccessor, thread_id: str) -> list[Any]:
|
|
"""Read ``messages`` through the mode-matched accessor.
|
|
|
|
Raw ``channel_values`` reads see a sentinel in delta mode; only a
|
|
materialized read reconstructs the list. Raw checkpoint tuples remain
|
|
valid for tuple-level metadata (checkpoint id, ``pending_writes``).
|
|
"""
|
|
snapshot = await accessor.aget({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
|
|
values = getattr(snapshot, "values", None) or {}
|
|
messages = values.get("messages") if isinstance(values, dict) else None
|
|
return list(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 _ends_on_human_input_request(messages: list[Any]) -> bool:
|
|
"""Return true when the turn ended on a Human Input Card the user has not answered.
|
|
|
|
``ask_clarification`` and the sandbox network prompt put the request in a
|
|
ToolMessage artifact and end the graph there, so it sits in the trailing run of
|
|
tool results. Continuing would tell the agent to keep going while the question is
|
|
still open, so the evaluator is not asked.
|
|
"""
|
|
for message in reversed(messages):
|
|
if _message_type(message) != "tool":
|
|
return False
|
|
artifact = message.get("artifact") if isinstance(message, dict) else getattr(message, "artifact", None)
|
|
human_input = artifact.get("human_input") if isinstance(artifact, Mapping) else None
|
|
if isinstance(human_input, Mapping) and human_input.get("kind") == "human_input_request":
|
|
return True
|
|
return False
|
|
|
|
|
|
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
|
|
# Defensive: compute continuation_count from the fresh current_goal
|
|
# inside the lock. The caller computed it from a possibly-stale goal
|
|
# snapshot; a racing continuation may have already bumped the count.
|
|
if continuation_count is not None:
|
|
current_count = int(current_goal.get("continuation_count", 0))
|
|
continuation_count = max(continuation_count, current_count + 1)
|
|
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,
|
|
accessor: CheckpointStateAccessor,
|
|
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,
|
|
user_id: str | None = None,
|
|
deerflow_trace_id: str | None = None,
|
|
task_store: Any | None = None,
|
|
extensions: Any | None = None,
|
|
run_stop_reason: str | None = None,
|
|
interaction_policy: RunInteractionPolicy | None = None,
|
|
expected_goal: GoalState | None = None,
|
|
verdict_callback: Callable[[GoalState, GoalEvaluation, str | None], None] | None = None,
|
|
usage_callback: Callable[[list[dict[str, int | str | None]]], None] | None = None,
|
|
) -> dict[str, Any] | _GoalCompletionCandidate | None:
|
|
"""Return a continuation input or a completion candidate for finalization.
|
|
|
|
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
|
|
if expected_goal is not None and not _goal_instance_matches(expected_goal, goal):
|
|
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."""
|
|
updated = 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,
|
|
)
|
|
if updated is not None and verdict_callback is not None:
|
|
verdict_callback(updated, evaluation, stand_down_reason)
|
|
return updated
|
|
|
|
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 = await _materialized_checkpoint_messages(accessor, thread_id)
|
|
conversation_signature_before = visible_conversation_signature(messages)
|
|
evidence_signature = latest_visible_assistant_signature(messages)
|
|
|
|
if _ends_on_human_input_request(messages):
|
|
# The agent asked the user something. Continuing would tell it to keep
|
|
# going while the question is still open on screen.
|
|
evaluation = GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="needs_user_input",
|
|
reason="The turn ended on a question to the user that has not been answered.",
|
|
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=_stand_down_reason(goal, evaluation, no_progress_count))
|
|
return None
|
|
|
|
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
|
|
except Exception:
|
|
logger.warning("Could not prepare goal evaluation for thread %s after run %s", thread_id, run_id, exc_info=True)
|
|
return None
|
|
|
|
try:
|
|
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,
|
|
thread_id=thread_id,
|
|
user_id=user_id,
|
|
deerflow_trace_id=deerflow_trace_id,
|
|
task_store=task_store,
|
|
extensions=extensions,
|
|
interaction_policy=interaction_policy,
|
|
usage_callback=usage_callback,
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("Goal evaluator failed for thread %s after run %s", thread_id, run_id, exc_info=True)
|
|
if abort_event is not None and abort_event.is_set():
|
|
return None
|
|
# Record the failure like the other stand-downs; otherwise the goal keeps the
|
|
# previous run's verdict, or none. Only the exception type is stored: a provider's
|
|
# error message can carry request details, and the traceback is logged above.
|
|
evaluation = GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="run_failed",
|
|
reason=f"The goal evaluator did not return a verdict ({type(exc).__name__}).",
|
|
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="evaluator_failed")
|
|
return None
|
|
if abort_event is not None and abort_event.is_set():
|
|
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(await _materialized_checkpoint_messages(accessor, thread_id)) != 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"]:
|
|
if verdict_callback is not None:
|
|
verdict_callback(current_goal, evaluation, None)
|
|
return _GoalCompletionCandidate(copy.deepcopy(current_goal), conversation_signature_before, evaluation)
|
|
|
|
stand_down_reason = _stand_down_reason(goal, evaluation, no_progress_count)
|
|
if stand_down_reason is None and run_stop_reason == "token_capped":
|
|
# The run already used up its token budget, and continuations share that
|
|
# budget, so another hidden turn would spend one more model call only to
|
|
# have its tool calls stripped.
|
|
stand_down_reason = "token_capped"
|
|
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
|
|
|
|
# Built before the continuation is counted: the message is redacted under the run's
|
|
# pii_redaction, and a redaction error fails the check as an evaluator redaction error
|
|
# does, so the raw objective and reason are never sent and the budget is not spent.
|
|
try:
|
|
continuation_message = make_goal_continuation_message(goal, evaluation, pii_redaction=getattr(app_config, "pii_redaction", None))
|
|
except Exception as exc:
|
|
logger.warning("Could not redact the goal continuation for thread %s after run %s", thread_id, run_id, exc_info=True)
|
|
# Only the exception type is stored: the error message could quote the text.
|
|
evaluation = GoalEvaluation(
|
|
satisfied=False,
|
|
blocker="run_failed",
|
|
reason=f"The goal continuation could not be redacted ({type(exc).__name__}).",
|
|
evidence_summary="",
|
|
)
|
|
await _persist(goal, evaluation, compute_no_progress_count(goal, evaluation, evidence_signature=evidence_signature), stand_down_reason="evaluator_failed")
|
|
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(await _materialized_checkpoint_messages(accessor, thread_id)) != conversation_signature_before:
|
|
# Do not pass continuation_count here: the persist above already
|
|
# committed it (as next_count). Re-passing next_count would make
|
|
# _persist_goal_evaluation's race guard (#4088) see that same write as
|
|
# a "current_count" bump and add another +1 on top of it, silently
|
|
# double-counting this single continuation attempt against the
|
|
# continuation budget even though it is being stood down, not
|
|
# delivered. Omitting it leaves the already-committed count untouched,
|
|
# matching every other stand-down call site in this function.
|
|
await _persist(
|
|
latest_goal,
|
|
evaluation,
|
|
no_progress_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": [continuation_message]}
|
|
|
|
|
|
def _is_edit_replay_run(record: RunRecord) -> bool:
|
|
metadata = record.metadata or {}
|
|
return metadata.get("replay_kind") == "edit"
|
|
|
|
|
|
async def _ensure_finalizing_before_edit_failure(run_manager: RunManager, record: RunRecord) -> None:
|
|
if _is_edit_replay_run(record) and not record.finalizing:
|
|
await run_manager.set_finalizing(record.run_id, True)
|
|
|
|
|
|
async def _publish_restored_checkpoint_values(
|
|
*,
|
|
bridge: StreamBridge,
|
|
run_id: str,
|
|
accessor: CheckpointStateAccessor | None,
|
|
thread_id: str,
|
|
) -> None:
|
|
if accessor is None:
|
|
return
|
|
snapshot = await accessor.aget({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
|
|
values = getattr(snapshot, "values", None)
|
|
if isinstance(values, dict):
|
|
await bridge.publish(run_id, "values", serialize(values, mode="values"))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RollbackPoint:
|
|
"""Materialized pre-run state used to restore the thread after cancellation.
|
|
|
|
Raw checkpoint blobs cannot reconstruct Delta-channel messages (their
|
|
checkpoints omit the materialized value), so rollback preserves those
|
|
messages plus delta mode's materialized non-message state in addition to
|
|
the raw pending writes.
|
|
"""
|
|
|
|
config: dict[str, Any]
|
|
state_values: dict[str, Any]
|
|
messages: tuple[Any, ...]
|
|
metadata: dict[str, Any]
|
|
pending_writes: tuple[tuple[str, str, Any], ...]
|
|
|
|
|
|
async def _capture_rollback_point(
|
|
accessor: CheckpointStateAccessor,
|
|
checkpointer: Any,
|
|
read_config: dict[str, Any],
|
|
) -> RollbackPoint | None:
|
|
"""Materialize the pre-run checkpoint state and its raw pending writes.
|
|
|
|
Returns ``None`` when the thread has no checkpoint yet; the caller keeps
|
|
the existing delete/reset rollback contract for that case.
|
|
"""
|
|
snapshot = await accessor.aget(read_config)
|
|
snapshot_config = getattr(snapshot, "config", None) or {}
|
|
configurable = snapshot_config.get("configurable") or {}
|
|
if not configurable.get("checkpoint_id"):
|
|
return None
|
|
checkpoint_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", snapshot_config)
|
|
raw_values = getattr(snapshot, "values", None) or {}
|
|
messages = raw_values.get("messages") if isinstance(raw_values, dict) else None
|
|
state_values = copy.deepcopy({key: value for key, value in raw_values.items() if key != "messages"}) if accessor.mode == "delta" and isinstance(raw_values, dict) else {}
|
|
return RollbackPoint(
|
|
config={
|
|
"configurable": {
|
|
"thread_id": configurable.get("thread_id"),
|
|
"checkpoint_ns": configurable.get("checkpoint_ns") or "",
|
|
"checkpoint_id": configurable.get("checkpoint_id"),
|
|
}
|
|
},
|
|
state_values=state_values,
|
|
messages=tuple(messages or ()),
|
|
metadata=dict(getattr(snapshot, "metadata", None) or {}),
|
|
pending_writes=tuple(getattr(checkpoint_tuple, "pending_writes", ()) or ()),
|
|
)
|
|
|
|
|
|
def _complete_state_replacement_values(
|
|
*,
|
|
mutation_graph: Any,
|
|
selected_values: dict[str, Any],
|
|
current_values: dict[str, Any],
|
|
run_id: str,
|
|
operation: str,
|
|
) -> dict[str, Any]:
|
|
"""Build a whole-state replacement through the graph's effective schema."""
|
|
writable_fields = graph_writable_channels(mutation_graph)
|
|
reducer_fields = graph_reducer_channels(mutation_graph)
|
|
if writable_fields is None or reducer_fields is None:
|
|
raise RuntimeError(f"Run {run_id} could not inspect the state schema for {operation}")
|
|
|
|
replacement_values: dict[str, Any] = {}
|
|
for field_name in writable_fields:
|
|
if field_name in selected_values:
|
|
replacement = copy.deepcopy(selected_values[field_name])
|
|
elif field_name in current_values:
|
|
# LangGraph has no public "unset channel" update. A fresh channel
|
|
# exposes its schema default when one exists (for example [] / {});
|
|
# optional and otherwise-unconstructible channels reset to None.
|
|
channel = mutation_graph.channels.get(field_name)
|
|
replacement = copy.deepcopy(channel.get()) if channel is not None and channel.is_available() else None
|
|
else:
|
|
continue
|
|
replacement_values[field_name] = Overwrite(replacement) if field_name in reducer_fields else replacement
|
|
return replacement_values
|
|
|
|
|
|
async def _linearize_delta_checkpoint_resume(
|
|
*,
|
|
accessor: CheckpointStateAccessor,
|
|
checkpointer: Any,
|
|
config: dict[str, Any],
|
|
thread_id: str,
|
|
run_id: str,
|
|
) -> list[Any] | None:
|
|
"""Replace a delta-mode checkpoint fork with an equivalent linear write.
|
|
|
|
Resuming from an older checkpoint forks the lineage, and in ``delta`` mode
|
|
the fork's state cannot be materialized correctly: the delta history walk
|
|
collects **every** ``pending_writes`` entry stored on each on-path
|
|
ancestor, but a shared parent also carries the writes of the sibling child
|
|
that was abandoned. Those writes are replayed into the fork, so the run
|
|
starts from a message list that still contains the answer it was supposed
|
|
to replace — regenerating in a branched thread surfaced this as the old
|
|
assistant message reappearing beside the new one after a reload (#4458).
|
|
Reproduced on postgres, sqlite, and the in-memory saver; ``full`` mode is
|
|
unaffected because its checkpoints carry complete ``channel_values`` and
|
|
need no replay.
|
|
|
|
The upstream contract (`BaseCheckpointSaver.get_delta_channel_history` and
|
|
the savers overriding it) is where write-to-child ownership belongs, so
|
|
this does not reimplement it. Instead the fork is expressed as what it
|
|
means: materialize the requested checkpoint's state and write it with
|
|
replace semantics on the **current head**, which has no other children,
|
|
then run linearly. Every materialized channel is restored; channels that
|
|
exist only on the newer head are reset to their schema default (or
|
|
``None`` when the channel has no constructible default). The abandoned
|
|
turn stays in checkpoint history as the rewritten head's ancestry.
|
|
|
|
The owner's current goal and continuation budget survive replay; the old
|
|
achievement record cannot certify the new answer. Full mode applies that
|
|
policy through the run input after scheduled-goal preflight instead.
|
|
|
|
Returns the materialized messages when the resume was linearized, or
|
|
``None`` when there was nothing to do (full mode, no checkpoint selector,
|
|
a non-root namespace, or a selector that already names the head). Failures
|
|
propagate: silently falling back to the fork would persist the corrupted
|
|
history this exists to prevent. The worker call site holds
|
|
``_checkpoint_thread_lock`` across rollback capture and this rewrite; do
|
|
not reacquire that non-reentrant lock inside this helper.
|
|
"""
|
|
if checkpointer is None or accessor.mode != "delta":
|
|
return None
|
|
configurable = config.get("configurable")
|
|
if not isinstance(configurable, dict):
|
|
return None
|
|
checkpoint_id = configurable.get("checkpoint_id")
|
|
if not isinstance(checkpoint_id, str) or not checkpoint_id:
|
|
return None
|
|
if configurable.get("checkpoint_ns"):
|
|
# Subgraph namespaces have their own lineage; the Gateway only selects
|
|
# root checkpoints, so leave anything else untouched.
|
|
return None
|
|
|
|
head_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
|
head = await accessor.aget(head_config)
|
|
if _checkpoint_id(head) == checkpoint_id:
|
|
# Selecting the head is already linear — no sibling can exist yet.
|
|
return None
|
|
|
|
source_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": "", "checkpoint_id": checkpoint_id}}
|
|
snapshot = await accessor.aget(source_config)
|
|
values = getattr(snapshot, "values", None) or {}
|
|
messages = values.get("messages") if isinstance(values, dict) else None
|
|
if not isinstance(messages, list):
|
|
raise RuntimeError(f"Run {run_id} could not materialize resume checkpoint {checkpoint_id}")
|
|
head_config["metadata"] = checkpoint_agent_binding_metadata(getattr(snapshot, "metadata", None))
|
|
|
|
# Write through the thread's effective schema so every application and
|
|
# middleware channel can be restored. Reducer channels need Overwrite to
|
|
# replace their already-aggregated value instead of merging it again.
|
|
mutation_graph = build_state_mutation_graph("checkpoint_resume", accessor.mode, graph_state_schema(getattr(accessor, "graph", None)))
|
|
selected_values = dict(values)
|
|
head_values = getattr(head, "values", None) or {}
|
|
head_values = dict(head_values) if isinstance(head_values, dict) else {}
|
|
selected_values["goal"] = head_values.get("goal")
|
|
selected_values["goal_outcome"] = None
|
|
replacement_values = _complete_state_replacement_values(
|
|
mutation_graph=mutation_graph,
|
|
selected_values=selected_values,
|
|
current_values=head_values,
|
|
run_id=run_id,
|
|
operation="checkpoint resume",
|
|
)
|
|
|
|
mutation_accessor = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode=accessor.mode)
|
|
await mutation_accessor.aupdate(head_config, replacement_values, as_node="checkpoint_resume")
|
|
configurable.pop("checkpoint_id", None)
|
|
configurable.pop("checkpoint_map", None)
|
|
logger.info("Run %s linearized a delta-mode resume of checkpoint %s onto thread %s", run_id, checkpoint_id, thread_id)
|
|
return list(messages)
|
|
|
|
|
|
async def _rollback_to_pre_run_checkpoint(
|
|
*,
|
|
accessor: CheckpointStateAccessor | None,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
run_id: str,
|
|
rollback_point: RollbackPoint | None,
|
|
snapshot_capture_failed: bool,
|
|
) -> bool:
|
|
"""Restore the complete pre-run state and report whether it completed.
|
|
|
|
Full mode forks the captured pre-run checkpoint and overwrites messages;
|
|
all other channels inherit from that parent. Delta mode cannot safely fork
|
|
once the cancelled path has attached writes to the same parent, so it
|
|
replaces every captured channel on the current head instead. Both writes
|
|
use a state-only mutation graph whose synthetic ``rollback_restore`` node
|
|
finishes immediately and schedules no agent work.
|
|
"""
|
|
if checkpointer is None:
|
|
logger.info("Run %s rollback requested but no checkpointer is configured", run_id)
|
|
return False
|
|
|
|
if snapshot_capture_failed:
|
|
logger.warning("Run %s rollback skipped: pre-run checkpoint capture failed", run_id)
|
|
return False
|
|
|
|
if rollback_point 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 True
|
|
|
|
configurable = rollback_point.config.get("configurable", {})
|
|
if not configurable.get("checkpoint_id"):
|
|
logger.warning("Run %s rollback skipped: pre-run checkpoint has no checkpoint id", run_id)
|
|
return False
|
|
|
|
if accessor is None:
|
|
# Unreachable in practice: a rollback point can only be captured
|
|
# through the bound accessor. Stay fail-closed.
|
|
logger.warning("Run %s rollback skipped: agent accessor unavailable", run_id)
|
|
return False
|
|
|
|
# Compile with the thread's effective schema so middleware-contributed
|
|
# channels survive (the base ThreadState fallback would silently drop
|
|
# them).
|
|
mutation_graph = build_state_mutation_graph("rollback_restore", accessor.mode, graph_state_schema(getattr(accessor, "graph", None)))
|
|
mutation_accessor = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode=accessor.mode)
|
|
if accessor.mode == "delta":
|
|
# A delta rollback fork has the same write-ownership problem as a
|
|
# checkpoint resume: the captured parent now carries writes from the
|
|
# cancelled sibling. Restore linearly on the current head instead.
|
|
restore_config: dict[str, Any] = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
|
current = await accessor.aget(restore_config)
|
|
raw_current_values = getattr(current, "values", None) or {}
|
|
current_values = dict(raw_current_values) if isinstance(raw_current_values, dict) else {}
|
|
selected_values = copy.deepcopy(rollback_point.state_values)
|
|
selected_values["messages"] = list(rollback_point.messages)
|
|
replacement_values = _complete_state_replacement_values(
|
|
mutation_graph=mutation_graph,
|
|
selected_values=selected_values,
|
|
current_values=current_values,
|
|
run_id=run_id,
|
|
operation="rollback",
|
|
)
|
|
else:
|
|
restore_config = {
|
|
**rollback_point.config,
|
|
"configurable": dict(rollback_point.config.get("configurable", {})),
|
|
}
|
|
replacement_values = {"messages": Overwrite(list(rollback_point.messages))}
|
|
|
|
restore_config["metadata"] = checkpoint_agent_binding_metadata(rollback_point.metadata)
|
|
|
|
restored_config = await mutation_accessor.aupdate(
|
|
restore_config,
|
|
replacement_values,
|
|
as_node="rollback_restore",
|
|
)
|
|
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 = rollback_point.pending_writes
|
|
if not pending_writes:
|
|
return True
|
|
|
|
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,
|
|
)
|
|
return True
|
|
|
|
|
|
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
|
|
|
|
|
|
def valid_duration_entry(run_id: Any, duration_seconds: Any) -> bool:
|
|
"""Check that (run_id, duration_seconds) is a well-formed duration entry."""
|
|
return isinstance(run_id, str) and bool(run_id) and isinstance(duration_seconds, int) and not isinstance(duration_seconds, bool)
|
|
|
|
|
|
RUN_MESSAGE_IDS_METADATA_KEY = "run_message_ids"
|
|
|
|
|
|
def valid_run_message_id_entry(message_id: Any, run_id: Any) -> bool:
|
|
"""Check that a persisted legacy message-to-run attribution is well formed."""
|
|
return isinstance(message_id, str) and bool(message_id) and isinstance(run_id, str) and bool(run_id)
|
|
|
|
|
|
async def persist_run_history_metadata(
|
|
*,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
durations: dict[str, int] | None = None,
|
|
message_run_ids: dict[str, str] | None = None,
|
|
) -> bool:
|
|
"""Merge validated run history indexes into a metadata-only checkpoint.
|
|
|
|
Durations accumulate so the history fast path can serve every known turn
|
|
from the latest checkpoint. Legacy AI-message attributions are persisted
|
|
alongside them for every audited AI ID, including boundary fallbacks whose
|
|
event lookup was exhaustively empty. The full mapping is deliberate: it is
|
|
both the exact-attribution cache and the negative-result coverage proof.
|
|
While the materialized message set at the head remains unchanged, later
|
|
reads query only uncached IDs. This metadata-only merge retains existing
|
|
entries, so compaction timing or historical migration may leave stale IDs;
|
|
reads ignore them because they only consult IDs in the materialized history.
|
|
"""
|
|
duration_updates = {run_id: max(0, duration_seconds) for run_id, duration_seconds in (durations or {}).items() if valid_duration_entry(run_id, duration_seconds)}
|
|
message_run_id_updates = {message_id: run_id for message_id, run_id in (message_run_ids or {}).items() if valid_run_message_id_entry(message_id, run_id)}
|
|
if not duration_updates and not message_run_id_updates:
|
|
return False
|
|
|
|
ckpt_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
|
|
async with _checkpoint_thread_lock(thread_id):
|
|
for _attempt in range(3):
|
|
ckpt_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", ckpt_config)
|
|
if ckpt_tuple is None:
|
|
return False
|
|
|
|
checkpoint = dict(getattr(ckpt_tuple, "checkpoint", {}) or {})
|
|
metadata = dict(getattr(ckpt_tuple, "metadata", {}) or {})
|
|
raw_run_durations = metadata.get("run_durations")
|
|
run_durations = {key: value for key, value in raw_run_durations.items() if valid_duration_entry(key, value)} if isinstance(raw_run_durations, dict) else {}
|
|
raw_message_run_ids = metadata.get(RUN_MESSAGE_IDS_METADATA_KEY)
|
|
run_message_ids = {message_id: run_id for message_id, run_id in raw_message_run_ids.items() if valid_run_message_id_entry(message_id, run_id)} if isinstance(raw_message_run_ids, dict) else {}
|
|
changed_durations = {run_id: duration for run_id, duration in duration_updates.items() if run_durations.get(run_id) != duration}
|
|
changed_message_run_ids = {message_id: run_id for message_id, run_id in message_run_id_updates.items() if run_message_ids.get(message_id) != run_id}
|
|
if not changed_durations and not changed_message_run_ids:
|
|
return False
|
|
|
|
run_durations.update(changed_durations)
|
|
run_message_ids.update(changed_message_run_ids)
|
|
parent_checkpoint_id = _checkpoint_identity(ckpt_tuple, checkpoint)
|
|
latest_tuple = await _call_checkpointer_method(checkpointer, "aget_tuple", "get_tuple", ckpt_config)
|
|
latest_checkpoint = dict(getattr(latest_tuple, "checkpoint", {}) or {}) if latest_tuple is not None else {}
|
|
if _checkpoint_identity(latest_tuple, latest_checkpoint) != parent_checkpoint_id:
|
|
continue
|
|
|
|
checkpoint.update(_new_checkpoint_marker())
|
|
metadata["source"] = "update"
|
|
prev_step = metadata.get("step")
|
|
metadata["step"] = (prev_step + 1) if isinstance(prev_step, int) else 1
|
|
metadata["run_durations"] = run_durations
|
|
if run_message_ids:
|
|
metadata[RUN_MESSAGE_IDS_METADATA_KEY] = run_message_ids
|
|
else:
|
|
metadata.pop(RUN_MESSAGE_IDS_METADATA_KEY, None)
|
|
metadata["writes"] = {
|
|
"runtime_run_duration": {
|
|
"run_ids": sorted(changed_durations),
|
|
"message_ids": sorted(changed_message_run_ids),
|
|
}
|
|
}
|
|
|
|
checkpoint_ns = _checkpoint_namespace(ckpt_tuple)
|
|
write_config = {
|
|
"configurable": {
|
|
"thread_id": thread_id,
|
|
"checkpoint_ns": checkpoint_ns,
|
|
"checkpoint_id": parent_checkpoint_id,
|
|
}
|
|
}
|
|
await _call_checkpointer_method(
|
|
checkpointer,
|
|
"aput",
|
|
"put",
|
|
write_config,
|
|
checkpoint,
|
|
metadata,
|
|
{},
|
|
)
|
|
return True
|
|
return False
|
|
|
|
|
|
async def persist_run_durations(
|
|
*,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
durations: dict[str, int],
|
|
) -> bool:
|
|
"""Merge validated run durations into a metadata-only checkpoint."""
|
|
return await persist_run_history_metadata(
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
durations=durations,
|
|
)
|
|
|
|
|
|
async def _persist_run_duration(
|
|
*,
|
|
checkpointer: Any,
|
|
thread_id: str,
|
|
run_id: str,
|
|
duration_seconds: int,
|
|
) -> None:
|
|
"""Persist one completed run duration in the thread checkpoint metadata."""
|
|
await persist_run_durations(
|
|
checkpointer=checkpointer,
|
|
thread_id=thread_id,
|
|
durations={run_id: duration_seconds},
|
|
)
|
|
|
|
|
|
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)
|
|
# Parent to the checkpoint this write was derived from - a parentless
|
|
# raw write would sever Delta-channel replay ancestry (and truncate
|
|
# full-mode history walks).
|
|
write_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns, "checkpoint_id": latest_identity}}
|
|
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 _message_id(obj: Any) -> str | None:
|
|
"""Best-effort extraction of a stable message id from a message-like object."""
|
|
msg_id = getattr(obj, "id", None)
|
|
if isinstance(msg_id, str) and msg_id:
|
|
return msg_id
|
|
if isinstance(obj, dict):
|
|
raw = obj.get("id")
|
|
if isinstance(raw, str) and raw:
|
|
return raw
|
|
return None
|
|
|
|
|
|
def _try_extract_from_message(obj: Any, pre_existing_ids: set[str] | None = None) -> str | None:
|
|
"""Try to extract fallback marker from a single message object or dict.
|
|
|
|
Messages whose id appears in ``pre_existing_ids`` are skipped — those are
|
|
history checkpointed by a *prior* run on this thread and any fallback
|
|
marker on them was already accounted for when that earlier run finished.
|
|
Without this filter, a single past run that ended with a fallback marker
|
|
would mark every subsequent run on the same thread as ``error``, because
|
|
LangGraph replays the full message history through ``stream_mode="values"``.
|
|
"""
|
|
if pre_existing_ids:
|
|
msg_id = _message_id(obj)
|
|
if msg_id is not None and msg_id in pre_existing_ids:
|
|
return None
|
|
|
|
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, pre_existing_ids: set[str] | None = None) -> 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.
|
|
|
|
Messages whose id appears in ``pre_existing_ids`` are ignored — they are
|
|
history from prior runs on the same thread (LangGraph replays the full
|
|
messages channel in ``stream_mode="values"`` chunks), and any error
|
|
fallback in that history was already resolved when its run finished.
|
|
"""
|
|
# 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, pre_existing_ids)
|
|
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, pre_existing_ids)
|
|
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 _collect_pre_existing_message_ids(values: Any) -> set[str]:
|
|
"""Collect stable message IDs from graph-materialized channel values."""
|
|
if not isinstance(values, dict):
|
|
return set()
|
|
messages = values.get("messages")
|
|
if not isinstance(messages, (list, tuple)):
|
|
return set()
|
|
return {message_id for message in messages if (message_id := _message_id(message)) is not None}
|
|
|
|
|
|
def _unpack_stream_item(
|
|
item: Any,
|
|
lg_modes: list[str],
|
|
stream_subgraphs: bool,
|
|
) -> tuple[str | None, Any, tuple[str, ...]]:
|
|
"""Unpack a multi-mode or subgraph stream item into (mode, chunk, namespace).
|
|
|
|
``namespace`` is the subgraph namespace tuple LangGraph prefixes onto each
|
|
frame when ``subgraphs=True``; it is empty for root-graph frames. Delegated
|
|
subagent graphs inherit the parent's checkpoint namespace (see
|
|
``subagents/executor.py``), so their frames arrive here with a non-empty
|
|
namespace and must not be mistaken for root frames.
|
|
|
|
Returns ``(None, None, ())`` if the item cannot be parsed.
|
|
"""
|
|
if stream_subgraphs:
|
|
if isinstance(item, tuple) and len(item) == 3:
|
|
ns, mode, chunk = item
|
|
namespace = tuple(str(part) for part in ns) if isinstance(ns, (list, tuple)) else (str(ns),)
|
|
return str(mode), chunk, namespace
|
|
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, ()
|
|
|
|
|
|
def _compose_sse_event(sse_event: str, namespace: tuple[str, ...]) -> str:
|
|
"""Namespace-qualified SSE event name, LangGraph Platform style.
|
|
|
|
Root frames keep the bare event name; subgraph frames become
|
|
``mode|ns1|ns2`` so clients can tell them apart. The LangGraph SDK parses
|
|
exactly this shape (``event.split("|").slice(1)``) and routes
|
|
subagent-namespaced values away from the thread view.
|
|
"""
|
|
if not namespace:
|
|
return sse_event
|
|
return "|".join((sse_event, *namespace))
|
|
|
|
|
|
#: Sentinel generation for an identity no lookup has missed at yet. Below every
|
|
#: real generation, so it never reads as "already asked at this generation".
|
|
_NEVER_LOOKED_UP = -1
|
|
|
|
|
|
def _NO_FEED_WRITES() -> int: # noqa: N802 — a callable constant, not a class
|
|
"""Generation source for a path with no journal: the feed never moves."""
|
|
return 0
|
|
|
|
|
|
class _MessageSeqStamper:
|
|
"""Attach each already-persisted message's feed seq to a ``values`` frame.
|
|
|
|
A checkpoint carries no seq of its own and loses messages to summarization,
|
|
so a client merging it with the seq-ordered thread feed cannot place a
|
|
surviving old message once the feed's loaded page window no longer reaches
|
|
back to it (#4666). The seq exists in the event store keyed by message
|
|
identity; this carries it to the client and writes nothing back to the
|
|
checkpoint.
|
|
|
|
Cost is bounded to frames that introduce identities it has not resolved
|
|
yet: a resolved seq is final (the feed's earliest-seq-wins rule), so it is
|
|
never looked up twice, and in a real run the only frame that pays for a
|
|
query is the one where compaction brings older messages back into view.
|
|
|
|
A miss, unlike a hit, is only provisional. A message this run produces
|
|
reaches a frame before ``RunJournal`` flushes it, so it legitimately misses
|
|
and would stay unstamped for the whole run if that answer were kept — the
|
|
long run that then rolls past the history page and compacts is exactly the
|
|
case this stamper exists for. Misses are therefore re-asked, but only once
|
|
the feed has actually gained rows, which *feed_generation* reports; frames
|
|
alone never trigger a retry. A failed lookup is a miss under the same rule,
|
|
so a transient store error costs one generation, not the run.
|
|
|
|
An unstamped message needs no seq while it is still streaming: appending it
|
|
at the tail is already its correct position.
|
|
"""
|
|
|
|
__slots__ = ("_store", "_thread_id", "_user_id", "_seqs", "_missing", "_feed_generation")
|
|
|
|
def __init__(self, event_store: Any, thread_id: str, *, feed_generation: Callable[[], int] | None = None) -> None:
|
|
self._store = event_store
|
|
self._thread_id = thread_id
|
|
# Without a generation source (no journal on this path) nothing can
|
|
# report a feed write, so a miss stays cached rather than being re-asked
|
|
# on every frame: a constant reads as "the feed has not moved".
|
|
self._feed_generation = feed_generation if feed_generation is not None else _NO_FEED_WRITES
|
|
# Soft-resolved once at build time, like the worker's write paths: a
|
|
# launch path that never inherits the auth contextvar (e.g. a
|
|
# null-owner scheduled task) writes rows with no user_id, so the
|
|
# lookup filters by the same id the writes stamped — or not at all —
|
|
# instead of the db store's strict AUTO default raising per frame.
|
|
user = get_current_user()
|
|
self._user_id: str | None = str(user.id) if user is not None else None
|
|
self._seqs: dict[str, int] = {}
|
|
# identity -> the feed generation its last lookup missed at.
|
|
self._missing: dict[str, int] = {}
|
|
|
|
async def stamp(self, payload: Any) -> Any:
|
|
if self._store is None or not isinstance(payload, Mapping):
|
|
return payload
|
|
messages = payload.get("messages")
|
|
if not isinstance(messages, list) or not messages:
|
|
return payload
|
|
|
|
identities = [message_identity(m) if isinstance(m, Mapping) else None for m in messages]
|
|
# Read before the lookup, never after: a write landing mid-query then
|
|
# advances past the generation the miss is recorded under, so the next
|
|
# frame re-asks. The reverse order could bury such a write.
|
|
generation = self._feed_generation()
|
|
unresolved = {i for i in identities if i is not None and i not in self._seqs and self._missing.get(i, _NEVER_LOOKED_UP) != generation}
|
|
if unresolved:
|
|
try:
|
|
found = await self._store.get_message_seqs(self._thread_id, sorted(unresolved), user_id=self._user_id)
|
|
except Exception:
|
|
# Placement is an enhancement: a client without seq falls back
|
|
# to its own ordering rule. Never fail the frame over it.
|
|
logger.warning("Failed to resolve message seqs for thread %s", self._thread_id, exc_info=True)
|
|
found = {}
|
|
self._seqs.update(found)
|
|
self._missing.update(dict.fromkeys(unresolved - found.keys(), generation))
|
|
|
|
stamped = [attach_message_seq(message, seq) if identity is not None and (seq := self._seqs.get(identity)) is not None else message for message, identity in zip(messages, identities, strict=True)]
|
|
return {**payload, "messages": stamped}
|
|
|
|
|
|
def _build_seq_stamper(event_store: Any, thread_id: str, journal: Any) -> _MessageSeqStamper:
|
|
"""Build the run's stamper, reading feed writes from *journal*.
|
|
|
|
The journal owns the writes that turn a lookup miss into a hit, so it is
|
|
also what can tell the stamper that a cached miss is worth re-asking. A run
|
|
without one has no writer to report, and the stamper falls back to keeping
|
|
its misses.
|
|
"""
|
|
return _MessageSeqStamper(
|
|
event_store,
|
|
thread_id,
|
|
feed_generation=(lambda: journal.feed_generation) if journal is not None else None,
|
|
)
|
|
|
|
|
|
async def _publish_stream_item(
|
|
*,
|
|
bridge: Any,
|
|
run_id: str,
|
|
mode: str,
|
|
chunk: Any,
|
|
namespace: tuple[str, ...],
|
|
file_tool_chunk_batcher: Any,
|
|
subagent_events: Any,
|
|
seq_stamper: Any = None,
|
|
) -> None:
|
|
"""Publish one stream frame, preserving the subgraph namespace.
|
|
|
|
A subgraph frame published under a bare event name impersonates the root
|
|
graph: a delegated subagent's ``values`` snapshot then replaces the whole
|
|
thread view in SDK clients and its token chunks flood the parent message
|
|
stream (#4399). Subgraph frames therefore keep their namespace in the event
|
|
name and bypass the root-only consumers (file-tool chunk batcher, subagent
|
|
event persistence — task_* lifecycle events are root frames already).
|
|
"""
|
|
sse_event = _compose_sse_event(_lg_mode_to_sse_event(mode), namespace)
|
|
if namespace:
|
|
await bridge.publish(run_id, sse_event, serialize(chunk, mode=mode))
|
|
return
|
|
if file_tool_chunk_batcher is not None and mode != "messages":
|
|
pending_chunks = file_tool_chunk_batcher.finish() if mode == "values" else file_tool_chunk_batcher.flush()
|
|
for publish_chunk in pending_chunks:
|
|
await bridge.publish(run_id, "messages", serialize(publish_chunk, mode="messages"))
|
|
chunks_to_publish = file_tool_chunk_batcher.push(chunk) if mode == "messages" and file_tool_chunk_batcher is not None else [chunk]
|
|
for publish_chunk in chunks_to_publish:
|
|
payload = serialize(publish_chunk, mode=mode)
|
|
if mode == "values" and seq_stamper is not None:
|
|
# Root frames only: a subagent's snapshot is not part of this
|
|
# thread's feed ordering (the namespaced branch returned above).
|
|
payload = await seq_stamper.stamp(payload)
|
|
await bridge.publish(run_id, sse_event, payload)
|
|
if mode == "custom":
|
|
await subagent_events.add(chunk)
|