"""Background agent execution. Runs an agent graph inside an ``asyncio.Task``, publishing events to a :class:`StreamBridge` as they are produced. Uses ``graph.astream(stream_mode=[...])`` which gives correct full-state snapshots for ``values`` mode, proper ``{node: writes}`` for ``updates``, and ``(chunk, metadata)`` tuples for ``messages`` mode. Note: ``events`` mode is not supported through the gateway — it requires ``graph.astream_events()`` which cannot simultaneously produce ``values`` snapshots. The JS open-source LangGraph API server works around this via internal checkpoint callbacks that are not exposed in the Python public API. """ from __future__ import annotations import asyncio import copy import inspect import logging from typing import Any, Literal from deerflow.runtime.serialization import serialize from deerflow.runtime.stream_bridge import StreamBridge from .manager import RunManager, RunRecord from .schemas import RunStatus logger = logging.getLogger(__name__) # Valid stream_mode values for LangGraph's graph.astream() _VALID_LG_MODES = {"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"} async def run_agent( bridge: StreamBridge, run_manager: RunManager, record: RunRecord, *, checkpointer: Any, store: Any | None = None, agent_factory: Any, graph_input: dict, config: dict, stream_modes: list[str] | None = None, stream_subgraphs: bool = False, interrupt_before: list[str] | Literal["*"] | None = None, interrupt_after: list[str] | Literal["*"] | None = None, ) -> None: """Execute an agent in the background, publishing events to *bridge*.""" run_id = record.run_id thread_id = record.thread_id requested_modes: set[str] = set(stream_modes or ["values"]) pre_run_checkpoint_id: str | None = None pre_run_snapshot: dict[str, Any] | None = None snapshot_capture_failed = False # Track whether "events" was requested but skipped if "events" in requested_modes: logger.info( "Run %s: 'events' stream_mode not supported in gateway (requires astream_events + checkpoint callbacks). Skipping.", run_id, ) try: # 1. Mark running await run_manager.set_status(run_id, RunStatus.running) # Snapshot the latest pre-run checkpoint so rollback can restore it. if checkpointer is not None: try: config_for_check = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}} ckpt_tuple = await checkpointer.aget_tuple(config_for_check) if ckpt_tuple is not None: ckpt_config = getattr(ckpt_tuple, "config", {}).get("configurable", {}) pre_run_checkpoint_id = ckpt_config.get("checkpoint_id") pre_run_snapshot = { "checkpoint_ns": ckpt_config.get("checkpoint_ns", ""), "checkpoint": copy.deepcopy(getattr(ckpt_tuple, "checkpoint", {})), "metadata": copy.deepcopy(getattr(ckpt_tuple, "metadata", {})), "pending_writes": copy.deepcopy(getattr(ckpt_tuple, "pending_writes", []) or []), } except Exception: snapshot_capture_failed = True logger.warning("Could not capture pre-run checkpoint snapshot for run %s", run_id, exc_info=True) # 2. Publish metadata — useStream needs both run_id AND thread_id await bridge.publish( run_id, "metadata", { "run_id": run_id, "thread_id": thread_id, }, ) # 3. Build the agent from langchain_core.runnables import RunnableConfig from langgraph.runtime import Runtime # Inject runtime context so middlewares can access thread_id # (langgraph-cli does this automatically; we must do it manually) runtime = Runtime(context={"thread_id": thread_id}, store=store) # If the caller already set a ``context`` key (LangGraph >= 0.6.0 # prefers it over ``configurable`` for thread-level data), make # sure ``thread_id`` is available there too. if "context" in config and isinstance(config["context"], dict): config["context"].setdefault("thread_id", thread_id) config.setdefault("configurable", {})["__pregel_runtime"] = runtime runnable_config = RunnableConfig(**config) agent = agent_factory(config=runnable_config) # 4. Attach checkpointer and store if checkpointer is not None: agent.checkpointer = checkpointer if store is not None: agent.store = store # 5. Set interrupt nodes if interrupt_before: agent.interrupt_before_nodes = interrupt_before if interrupt_after: agent.interrupt_after_nodes = interrupt_after # 6. Build LangGraph stream_mode list # "events" is NOT a valid astream mode — skip it # "messages-tuple" maps to LangGraph's "messages" mode lg_modes: list[str] = [] for m in requested_modes: if m == "messages-tuple": lg_modes.append("messages") elif m == "events": # Skipped — see log above continue elif m in _VALID_LG_MODES: lg_modes.append(m) if not lg_modes: lg_modes = ["values"] # Deduplicate while preserving order seen: set[str] = set() deduped: list[str] = [] for m in lg_modes: if m not in seen: seen.add(m) deduped.append(m) lg_modes = deduped logger.info("Run %s: streaming with modes %s (requested: %s)", run_id, lg_modes, requested_modes) # 7. Stream using graph.astream if len(lg_modes) == 1 and not stream_subgraphs: # Single mode, no subgraphs: astream yields raw chunks single_mode = lg_modes[0] async for chunk in agent.astream(graph_input, config=runnable_config, stream_mode=single_mode): if record.abort_event.is_set(): logger.info("Run %s abort requested — stopping", run_id) break sse_event = _lg_mode_to_sse_event(single_mode) await bridge.publish(run_id, sse_event, serialize(chunk, mode=single_mode)) else: # Multiple modes or subgraphs: astream yields tuples async for item in agent.astream( graph_input, config=runnable_config, stream_mode=lg_modes, subgraphs=stream_subgraphs, ): if record.abort_event.is_set(): logger.info("Run %s abort requested — stopping", run_id) break mode, chunk = _unpack_stream_item(item, lg_modes, stream_subgraphs) if mode is None: continue sse_event = _lg_mode_to_sse_event(mode) await bridge.publish(run_id, sse_event, serialize(chunk, mode=mode)) # 8. Final status if record.abort_event.is_set(): action = record.abort_action if action == "rollback": await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user") try: await _rollback_to_pre_run_checkpoint( checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, pre_run_checkpoint_id=pre_run_checkpoint_id, pre_run_snapshot=pre_run_snapshot, snapshot_capture_failed=snapshot_capture_failed, ) logger.info("Run %s rolled back to pre-run checkpoint %s", run_id, pre_run_checkpoint_id) except Exception: logger.warning("Failed to rollback checkpoint for run %s", run_id, exc_info=True) else: await run_manager.set_status(run_id, RunStatus.interrupted) else: await run_manager.set_status(run_id, RunStatus.success) except asyncio.CancelledError: action = record.abort_action if action == "rollback": await run_manager.set_status(run_id, RunStatus.error, error="Rolled back by user") try: await _rollback_to_pre_run_checkpoint( checkpointer=checkpointer, thread_id=thread_id, run_id=run_id, pre_run_checkpoint_id=pre_run_checkpoint_id, pre_run_snapshot=pre_run_snapshot, snapshot_capture_failed=snapshot_capture_failed, ) logger.info("Run %s was cancelled and rolled back", run_id) except Exception: logger.warning("Run %s cancellation rollback failed", run_id, exc_info=True) else: await run_manager.set_status(run_id, RunStatus.interrupted) logger.info("Run %s was cancelled", run_id) except Exception as exc: error_msg = f"{exc}" logger.exception("Run %s failed: %s", run_id, error_msg) await run_manager.set_status(run_id, RunStatus.error, error=error_msg) await bridge.publish( run_id, "error", { "message": error_msg, "name": type(exc).__name__, }, ) finally: await bridge.publish_end(run_id) asyncio.create_task(bridge.cleanup(run_id, delay=60)) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- async def _call_checkpointer_method(checkpointer: Any, async_name: str, sync_name: str, *args: Any, **kwargs: Any) -> Any: """Call a checkpointer method, supporting async and sync variants.""" method = getattr(checkpointer, async_name, None) or getattr(checkpointer, sync_name, None) if method is None: raise AttributeError(f"Missing checkpointer method: {async_name}/{sync_name}") result = method(*args, **kwargs) if inspect.isawaitable(result): return await result return result async def _rollback_to_pre_run_checkpoint( *, checkpointer: Any, thread_id: str, run_id: str, pre_run_checkpoint_id: str | None, pre_run_snapshot: dict[str, Any] | None, snapshot_capture_failed: bool, ) -> None: """Restore thread state to the checkpoint snapshot captured before run start.""" if checkpointer is None: logger.info("Run %s rollback requested but no checkpointer is configured", run_id) return if snapshot_capture_failed: logger.warning("Run %s rollback skipped: pre-run checkpoint snapshot capture failed", run_id) return if pre_run_snapshot is None: await _call_checkpointer_method(checkpointer, "adelete_thread", "delete_thread", thread_id) logger.info("Run %s rollback reset thread %s to empty state", run_id, thread_id) return checkpoint_to_restore = None metadata_to_restore: dict[str, Any] = {} checkpoint_ns = "" checkpoint = pre_run_snapshot.get("checkpoint") if not isinstance(checkpoint, dict): logger.warning("Run %s rollback skipped: invalid pre-run checkpoint snapshot", run_id) return checkpoint_to_restore = checkpoint if checkpoint_to_restore.get("id") is None and pre_run_checkpoint_id is not None: checkpoint_to_restore = {**checkpoint_to_restore, "id": pre_run_checkpoint_id} if checkpoint_to_restore.get("id") is None: logger.warning("Run %s rollback skipped: pre-run checkpoint has no checkpoint id", run_id) return metadata = pre_run_snapshot.get("metadata", {}) metadata_to_restore = metadata if isinstance(metadata, dict) else {} raw_checkpoint_ns = pre_run_snapshot.get("checkpoint_ns") checkpoint_ns = raw_checkpoint_ns if isinstance(raw_checkpoint_ns, str) else "" channel_versions = checkpoint_to_restore.get("channel_versions") new_versions = dict(channel_versions) if isinstance(channel_versions, dict) else {} restore_config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": checkpoint_ns}} restored_config = await _call_checkpointer_method( checkpointer, "aput", "put", restore_config, checkpoint_to_restore, metadata_to_restore if isinstance(metadata_to_restore, dict) else {}, new_versions, ) if not isinstance(restored_config, dict): raise RuntimeError(f"Run {run_id} rollback restore returned invalid config: expected dict") restored_configurable = restored_config.get("configurable", {}) if not isinstance(restored_configurable, dict): raise RuntimeError(f"Run {run_id} rollback restore returned invalid config payload") restored_checkpoint_id = restored_configurable.get("checkpoint_id") if not restored_checkpoint_id: raise RuntimeError(f"Run {run_id} rollback restore did not return checkpoint_id") pending_writes = pre_run_snapshot.get("pending_writes", []) if not pending_writes: return writes_by_task: dict[str, list[tuple[str, Any]]] = {} for item in pending_writes: if not isinstance(item, (tuple, list)) or len(item) != 3: raise RuntimeError(f"Run {run_id} rollback failed: pending_write is not a 3-tuple: {item!r}") task_id, channel, value = item if not isinstance(channel, str): raise RuntimeError(f"Run {run_id} rollback failed: pending_write has non-string channel: task_id={task_id!r}, channel={channel!r}") writes_by_task.setdefault(str(task_id), []).append((channel, value)) for task_id, writes in writes_by_task.items(): await _call_checkpointer_method( checkpointer, "aput_writes", "put_writes", restored_config, writes, task_id=task_id, ) def _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 _unpack_stream_item( item: Any, lg_modes: list[str], stream_subgraphs: bool, ) -> tuple[str | None, Any]: """Unpack a multi-mode or subgraph stream item into (mode, chunk). Returns ``(None, None)`` if the item cannot be parsed. """ if stream_subgraphs: if isinstance(item, tuple) and len(item) == 3: _ns, mode, chunk = item return str(mode), chunk if isinstance(item, tuple) and len(item) == 2: mode, chunk = item return str(mode), chunk return None, None if isinstance(item, tuple) and len(item) == 2: mode, chunk = item return str(mode), chunk # Fallback: single-element output from first mode return lg_modes[0] if lg_modes else None, item