diff --git a/README.md b/README.md index d02c1d9c5..d3f9c5395 100644 --- a/README.md +++ b/README.md @@ -835,17 +835,21 @@ Advanced deployments can enable pluggable authorization with `authorization.enab Advanced deployments can also extend the agent runtime itself by declaring zero-argument `AgentMiddleware` classes under `extensions.middlewares` in `config.yaml` or `extensions_config.json`. DeerFlow loads the same configured class list into the lead-agent and subagent pipelines after their built-in runtime middlewares and loop/token guards, but before the terminal-response/safety/clarification tail, so enterprise forks can add domain guardrails, tool-call governance, or observability hooks without patching the built-in middleware builders. Missing packages, invalid classes, and broken modules fail loudly at agent creation. Treat `config.yaml` and `extensions_config.json` as trusted operator-controlled files: middleware paths are code execution, just like custom tool, model, sandbox, guardrail, MCP server, and MCP interceptor declarations. Gateway skill/MCP toggle endpoints preserve this field but do not expose an API write path for `extensions.middlewares`. Per-context parameterization and separate lead-only/subagent-only middleware lists are not supported yet. -For packaged and configurable middleware integrations, use the top-level `plugins:` list -in `config.yaml`. A plugin exposes `module.path:install`, depends only on the standalone -`deerflow-extension-api` contract package, and can contribute isolated middleware to -semantic lead/subagent model or tool positions without patching DeerFlow's builders. -Plugin order is deterministic, per-plugin configuration is passed to `install()`, and -`required: true` makes load failure abort startup; otherwise failures are reported and -skipped. Plugins load once when the Gateway app is constructed, so changes require a -restart. Because this imports Python code, `plugins:` is intentionally unavailable through -the API-writable `extensions_config.json`. In Docker deployments, install the plugin in the -Gateway image rather than only in the host environment. See `config.example.yaml` for -configuration. +For packaged and configurable runtime integrations, use the top-level `plugins:` list in +`config.yaml`. A plugin exposes `module.path:install`, depends only on the standalone +`deerflow-extension-api` contract package, and can register exactly three contribution +kinds: isolated middleware at semantic lead/subagent model or tool positions, lead and +subagent task-lifecycle hooks, and observers for DeerFlow-owned system model calls such as +goal evaluation, memory extraction, title generation, and summarization. DeerFlow allocates +a task-scoped extension store only when one of those contribution kinds is registered and +uses the Gateway's canonical notification loop for lifecycle and system-model callbacks, +including subagents that execute on isolated loops. Plugin order is deterministic, +per-plugin configuration is passed to `install()`, and `required: true` makes load failure +abort startup; otherwise failures are reported and skipped. Plugins load once when the +Gateway app is constructed, so changes require a restart. Because this imports Python code, +`plugins:` is intentionally unavailable through the API-writable +`extensions_config.json`. In Docker deployments, install the plugin in the Gateway image +rather than only in the host environment. See `config.example.yaml` for configuration. Gateway-generated follow-up suggestions now normalize both plain-string model output and block/list-style rich content before parsing the JSON array response, so provider-specific content wrappers do not silently drop suggestions. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 744a504cc..3738046da 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -438,7 +438,7 @@ Before changing a later authorization phase, read the [authorization RFC](../doc 34. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first 35. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Payloads are versioned: legacy modes (`free_text` / `choice_with_other`) keep `version: 1` unchanged, while the v2 `form` mode (from `fields`) carries `version: 2` so older frontends reject the payload and degrade to the plain-text fallback. Field normalization is deterministic and lives in the middleware, not the tool schema — the middleware short-circuits before tool execution, so tool-arg typing alone provides no runtime validation. Validation is atomic: any structurally broken entry (non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member like `__proto__`/`constructor`, exceeding the caps of 16 fields / 24 options per field / 200 chars per text, or the whole normalized definition exceeding `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8 — the per-item caps alone admit forms whose IM text fallback would blow channel delivery limits and truncate away trailing fields) degrades the whole form to the legacy option/free-text modes, so a card can never render "complete" while silently missing a business field; benign issues keep local degradation (unknown types — including unhashable JSON like `type: []`, which must never raise from the membership probe — and option-less selects become `text`), and options are trimmed/deduped with blanks dropped (both form-level and top-level) because the frontend parser rejects blank option labels. Model-produced XML-to-dict option payloads are recursively flattened from dict/list containers in source order, scalar string/number leaves are retained, and residual XML tags are removed before the same trimming and deduplication. Checkbox fields are booleans that default to an explicit "no"; `required` on a checkbox means must-agree/consent semantics. The response protocol is deliberately unchanged (v1 `text`/`option` only): form cards submit a readable text summary as `response_kind: "text"`, so journal persistence and answered-card recovery need no new allowlist entries. Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context. -### Python Extension System (Middleware Slice) +### Python Extension System (Runtime Slice) Third-party Python packages can expose an `install(registry, config)` function and be loaded, in deterministic order, from the startup-only top-level `plugins:` list in @@ -447,9 +447,10 @@ through Gateway APIs, while importing Python entry points is an operator-control execution boundary. A plugin marked `required: true` fails Gateway construction when it cannot load; optional plugins fail open with attributed diagnostics. -The public package is `packages/extension-api/` and must never import `deerflow`. In this -slice its registry contract exposes middleware contribution only. Each contribution -declares lead/subagent scope, stable order, and a semantic placement (`MODEL_LOGICAL`, +The public package is `packages/extension-api/` and must never import `deerflow`. Its +registry contract exposes exactly three contribution kinds: middleware contributors, +task-lifecycle contributors, and system-model-call observers. Middleware contributions +declare lead/subagent scope, stable order, and a semantic placement (`MODEL_LOGICAL`, `MODEL_PHYSICAL`, `TOOL_VISIBLE`, `TOOL_RAW`, or `STANDARD`) rather than a fragile list index. `extensions/stack.py` is the single final composition point; do not inject inside the shared base builder because the lead builder appends more middleware afterward. @@ -469,27 +470,76 @@ wrapper mirrors lifecycle hooks, tools, transformers, and state schema implement the inner middleware. LangChain treats each sync/async model or tool wrapper pair as one capability, so a single-sided wrapper receives a pass-through counterpart; implement both sides when the extension must observe both synchronous and asynchronous execution -paths. Lead runs and -subagents allocate an `ExtensionData` task store only when middleware contributors are -present and expose it through `EXTENSION_TASK_STORE_KEY`; extensions retrieve it with -`task_store_from_runtime()`. Each run resolves the immutable loaded-extension snapshot -once and binds that same object through task-store allocation and synchronous agent -construction, so a concurrent singleton replacement cannot mix two extension -generations without changing the LangGraph graph-factory ABI. The graph-build binding is -a ContextVar scoped to synchronous construction, so it has already exited by the time the -lead agent delegates; the run worker therefore also publishes the snapshot on runtime -context under the host-internal `EXTENSION_SNAPSHOT_CONTEXT_KEY`, `task_tool` reads it -back through `resolve_run_extensions()` (type-checked — runtime context is -caller-mergeable), and `SubagentExecutor` binds it at construction. That key is written -after the caller merge and popped when the run has none, so a caller-supplied value is -never authoritative. Absent the key — embedded `DeerFlowClient`, standalone LangGraph +paths. + +Lead runs and subagents allocate an `ExtensionData` task store only when at least one of +the three contribution kinds is registered. Middleware and system-call sites recover the +live store through `EXTENSION_TASK_STORE_KEY` / `task_store_from_runtime()`; lifecycle +contributors receive that same store directly. Each task resolves the immutable +loaded-extension snapshot once and binds that same object through task-store allocation, +hooks, and synchronous agent construction, so a concurrent singleton replacement cannot +mix two extension generations without changing the LangGraph graph-factory ABI. The +graph-build binding is a ContextVar scoped to synchronous construction, so it has already +exited by the time the lead agent delegates; the run worker therefore also publishes the +snapshot on runtime context under the host-internal `EXTENSION_SNAPSHOT_CONTEXT_KEY`, +`task_tool` reads it back through `resolve_run_extensions()` (type-checked — runtime +context is caller-mergeable), and `SubagentExecutor` binds it at construction. That key is +written after the caller merge and popped when the run has none, so a caller-supplied value +is never authoritative. Absent the key — embedded `DeerFlowClient`, standalone LangGraph Server — the executor keeps its `get_loaded_extensions()` fallback. +The lead worker awaits `on_task_start` after the run has started and awaits `on_task_stop` +after completion persistence/hooks but before clearing any active finalizing barrier or +publishing the stream end. A subagent with a parent `run_id` wraps its execution with the +same start/stop pair. Outcomes are conservative (`completed`, `aborted`, or `failed`), +contributors run in registration order within one bounded budget, and notification failures +are logged and fail open. + +Fail-open is decided by the *origin* of a failure, not by its base class, because +`CancelledError` reaches a contributor's `except` for two unrelated reasons. Only a genuine +cancellation of the host task increments `asyncio.Task.cancelling()`, so `_notify_each` +propagates on that and contains everything else: a contributor that lets a `CancelledError` +escape — an extension implementing an internal timeout with cancellation, say — must not +skip its successors, and must not reach the worker's deferred-interrupt path, which would +end an otherwise successful run as cancelled. `KeyboardInterrupt` / `SystemExit` still +propagate. + +System-model-call observers cover DeerFlow-owned model invocations that do not pass +through middleware model-call wrappers: goal evaluation, memory extraction, title +generation, and summarization. They receive a request/result snapshot, duration, and the +active task store when one exists; detached system work receives an isolated store. All +three terminal paths are reported without changing the exception the host observes: +success and failure are awaited inline, while cancellation — routine, since +interrupt/rollback admission and shutdown both cancel the run task, with the provider +tokens already spent — is submitted to the notify loop instead of awaited, because a +repeated cancel would interrupt that await before any observer ran. A deployment with no +registered notify loop drops the cancellation observation, exactly as the synchronous +memory bridge does. `SystemModelRequest.messages` normalizes to a tuple at construction: goal and +memory pass a message list while title and summarization pass one prompt string, and a +bare `str` is already a `Sequence`, so without normalization an observer iterating it +would walk characters. Normalizing also copies a live list, which is what makes the frozen +snapshot immutable in fact rather than only by declaration. Gateway registers one canonical extension-notification loop. Awaited lifecycle +hooks and async system observations are dispatched to that loop even when the caller is a +subagent's isolated loop, while synchronous system callbacks submit fire-and-forget work +there. Shutdown stops accepting detached observations before the memory shutdown flush and +resets the loop only after in-flight run/subagent drain ordering is complete. + +The memory kind reaches those observers through a different shape, and the difference is +deliberate rather than an oversight to be "aligned" away. DeerMem must stay vendorable and +cannot import the extension API, so it reports through the `MemoryCallbacks.on_memory_llm_result` +host hook, which the DeerFlow-side callbacks translate into an observation and submit +without awaiting. It also guards its provider call with `BaseException` rather than +`Exception`, which is safe precisely because that whole path runs on a worker thread — the +debounce timer, or the executor `update_memory` offloads to — where cancelling the awaiting +side never interrupts the running thread, so `CancelledError` cannot arrive there at all. +The host hook wrapper around the callback stays at `Exception`: only the hook's own failures +are non-fatal, and an observability path must not swallow `SystemExit` / `KeyboardInterrupt`. + Gateway `create_app()` loads plugins once, stores the immutable registry on `app.state` and in the process-wide singleton, and installs one canonical live diagnostics list. -Changing `plugins` requires a restart. Later extension contribution points must be added -to the public contract and host runtime in the same slice; never accept a registration -method that the current host silently ignores. +Changing `plugins` requires a restart. Any future contribution kind must be added to the +public contract and host runtime in the same slice; never accept a registration method +that the current host silently ignores. ### Configuration System diff --git a/backend/app/gateway/app.py b/backend/app/gateway/app.py index b82a8363a..20472649d 100644 --- a/backend/app/gateway/app.py +++ b/backend/app/gateway/app.py @@ -438,6 +438,17 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: logger.warning("Memory retrieval index rebuild is still running; leaving its connection open during shutdown") manager = None + try: + # Memory shutdown runs on a worker thread and can trigger detached + # system-model callbacks. Stop accepting those callbacks before + # flushing, while keeping the registered loop alive for awaited + # task hooks until langgraph_runtime drains runs and subagents. + from deerflow.extensions.notify import suspend_extension_system_observations + + suspend_extension_system_observations() + except Exception: + logger.debug("Failed to suspend extension system observations (non-fatal)", exc_info=True) + try: app_cfg = get_app_config() if app_cfg.memory.enabled: diff --git a/backend/app/gateway/deps.py b/backend/app/gateway/deps.py index b6856d622..844e6db3a 100644 --- a/backend/app/gateway/deps.py +++ b/backend/app/gateway/deps.py @@ -385,6 +385,33 @@ async def langgraph_runtime(app: FastAPI, startup_config: AppConfig) -> AsyncGen _validate_agent_storage(startup_config) async with AsyncExitStack() as stack: + # Lifecycle and system-model hooks can originate on isolated subagent + # loops. Bind them to the Gateway's serving loop before any runtime + # dependency starts, then reset the binding last through the exit + # stack. Registering the callback synchronously here also covers every + # startup-failure and cancellation path below. + try: + from deerflow.extensions.notify import ( + reset_extension_notify_loop, + set_extension_notify_loop, + ) + + set_extension_notify_loop(asyncio.get_running_loop()) + except Exception: + logger.exception("Failed to register the extension notify loop; sync observations will be dropped") + else: + + def reset_notify_loop_safely() -> None: + try: + reset_extension_notify_loop() + except Exception: + logger.debug( + "Failed to reset the extension notify loop (non-fatal)", + exc_info=True, + ) + + stack.callback(reset_notify_loop_safely) + config = startup_config app.state.checkpoint_channel_mode = freeze_checkpoint_channel_mode(config.database.checkpoint_channel_mode) app.state.checkpoint_snapshot_frequency = freeze_checkpoint_snapshot_frequency(config.database.checkpoint_delta.snapshot_frequency) diff --git a/backend/extension_test_fixtures/demo_extensions.py b/backend/extension_test_fixtures/demo_extensions.py index 79cdb315e..67b81bea4 100644 --- a/backend/extension_test_fixtures/demo_extensions.py +++ b/backend/extension_test_fixtures/demo_extensions.py @@ -23,7 +23,7 @@ def install_ok(registry: ExtensionRegistry, config: Mapping[str, Any]) -> None: @extension(api="0.1", name="stamped") def install_stamped(registry: ExtensionRegistry, config: Mapping[str, Any]) -> None: INSTALLED.append("stamped") - registry.middlewares(_Contributor("stamped")) + registry.task_lifecycle(_Contributor("stamped")) @extension(api="99.0", name="future") @@ -42,8 +42,8 @@ def install_newer_minor_api(registry: ExtensionRegistry, config: Mapping[str, An def install_partial_then_raise(registry: ExtensionRegistry, config: Mapping[str, Any]) -> None: """Registers two contributors, then fails — exercises rollback.""" - registry.middlewares(_Contributor("partial-a")) - registry.middlewares(_Contributor("partial-b")) + registry.middlewares(_Contributor("partial")) + registry.task_lifecycle(_Contributor("partial")) raise ValueError("boom") diff --git a/backend/packages/extension-api/deerflow_extension_api/__init__.py b/backend/packages/extension-api/deerflow_extension_api/__init__.py index ec2778e0a..5b9a65f1b 100644 --- a/backend/packages/extension-api/deerflow_extension_api/__init__.py +++ b/backend/packages/extension-api/deerflow_extension_api/__init__.py @@ -12,6 +12,13 @@ from deerflow_extension_api.contracts import ( ExtensionRegistry, HostPolicySnapshot, MiddlewareContributor, + SystemModelCallObserver, + SystemModelRequest, + SystemModelResult, + SystemOperationKind, + TaskInfo, + TaskLifecycleContributor, + TaskOutcome, extension, ) from deerflow_extension_api.placement import ( @@ -30,7 +37,7 @@ from deerflow_extension_api.state import ExtensionData #: (contributors and observers), so minors may break and only patches promise #: to be additive. From 1.0 on, bump the major on any breaking change; see the #: spec's evolution rules for what counts as additive. -API_VERSION = "0.1.0" +API_VERSION = "0.1.1" __all__ = [ "API_VERSION", @@ -44,6 +51,13 @@ __all__ = [ "MiddlewareContributor", "MiddlewarePlacement", "Placement", + "SystemModelCallObserver", + "SystemModelRequest", + "SystemModelResult", + "SystemOperationKind", + "TaskInfo", + "TaskLifecycleContributor", + "TaskOutcome", "extension", "task_store_from_runtime", ] diff --git a/backend/packages/extension-api/deerflow_extension_api/contracts.py b/backend/packages/extension-api/deerflow_extension_api/contracts.py index d92cd6e50..da66894f9 100644 --- a/backend/packages/extension-api/deerflow_extension_api/contracts.py +++ b/backend/packages/extension-api/deerflow_extension_api/contracts.py @@ -11,7 +11,8 @@ from __future__ import annotations from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Protocol, TypeVar, runtime_checkable +from enum import StrEnum +from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar, runtime_checkable from deerflow_extension_api.state import ExtensionData @@ -42,6 +43,103 @@ class HostPolicySnapshot: max_subagents_per_run: int | None = None +# --- Task lifecycle --------------------------------------------------------- + + +class TaskOutcome(StrEnum): + COMPLETED = "completed" + ABORTED = "aborted" + FAILED = "failed" + + +@dataclass(frozen=True) +class TaskInfo: + """Identity of one lead-agent or subagent execution.""" + + task_id: str + run_id: str + thread_id: str + kind: Literal["lead", "subagent"] + parent_task_id: str | None = None + agent_name: str | None = None + resumed: bool = False + + +class TaskLifecycleContributor(Protocol): + async def on_task_start( + self, + app_store: ExtensionData, + task_store: ExtensionData, + info: TaskInfo, + ) -> None: + return None + + async def on_task_stop( + self, + app_store: ExtensionData, + task_store: ExtensionData, + info: TaskInfo, + outcome: TaskOutcome, + ) -> None: + return None + + +# --- System model calls not wrapped by middleware model-call hooks ---------- + + +class SystemOperationKind(StrEnum): + GOAL = "goal" + MEMORY = "memory" + TITLE = "title" + SUMMARIZATION = "summarization" + + +@dataclass(frozen=True) +class SystemModelRequest: + """Read-only snapshot taken before a system-owned model call.""" + + messages: Sequence[Any] = () + model_name: str | None = None + invoke_config: Mapping[str, Any] | None = None + + def __post_init__(self) -> None: + """Normalize ``messages`` to a tuple so the snapshot is what it claims to be. + + Call sites differ: goal evaluation and memory extraction pass a message list, + while title generation and summarization pass one prompt string. A bare ``str`` + already satisfies ``Sequence``, so without this an observer iterating + ``request.messages`` would silently walk characters. Copying a list also makes + the frozen snapshot immutable in fact, not only by dataclass declaration — the + caller keeps its own list and observations may run after the call returns. + """ + messages = self.messages + if isinstance(messages, tuple): + return + normalized = tuple(messages) if isinstance(messages, Sequence) and not isinstance(messages, str | bytes) else (messages,) + object.__setattr__(self, "messages", normalized) + + +@dataclass(frozen=True) +class SystemModelResult: + """Success or failure snapshot taken after a system-owned model call.""" + + response: Any | None = None + error: BaseException | None = None + duration_ms: float | None = None + + +class SystemModelCallObserver(Protocol): + async def on_system_model_call( + self, + app_store: ExtensionData, + task_store: ExtensionData, + kind: SystemOperationKind, + request: SystemModelRequest, + result: SystemModelResult, + ) -> None: + return None + + # --- Middleware ------------------------------------------------------------- @@ -71,6 +169,12 @@ class ExtensionRegistry(Protocol): def middlewares(self, contributor: MiddlewareContributor) -> None: return None + def task_lifecycle(self, contributor: TaskLifecycleContributor) -> None: + return None + + def system_model_observer(self, observer: SystemModelCallObserver) -> None: + return None + #: The install() entry point signature every extension exposes. ExtensionInstall = Callable[[ExtensionRegistry, Mapping[str, Any]], None] diff --git a/backend/packages/extension-api/pyproject.toml b/backend/packages/extension-api/pyproject.toml index 70a84a157..f245f3fe7 100644 --- a/backend/packages/extension-api/pyproject.toml +++ b/backend/packages/extension-api/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "deerflow-extension-api" -version = "0.1.0" +version = "0.1.1" description = "Public contracts for DeerFlow extensions" requires-python = ">=3.12" # Keep the contract package import-light and independent from the host. Public diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index 688a52f40..70bd072eb 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -235,14 +235,23 @@ def _authorize_model_name( return model_name -def _create_summarization_middleware(*, app_config: AppConfig | None = None, run_model_name: str | None = None) -> DeerFlowSummarizationMiddleware | None: +def _create_summarization_middleware( + *, + app_config: AppConfig | None = None, + run_model_name: str | None = None, + extensions=None, +) -> DeerFlowSummarizationMiddleware | None: """Create and configure the summarization middleware from config. ``run_model_name`` is the resolved run model; it is the source of truth for ``model_name: null`` summarization and the explicit-summary-model fallback, so a custom agent's model is used instead of ``config.models[0]``. """ - return create_summarization_middleware(app_config=app_config, run_model_name=run_model_name) + return create_summarization_middleware( + app_config=app_config, + run_model_name=run_model_name, + extensions=extensions, + ) def _create_todo_list_middleware(is_plan_mode: bool) -> TodoMiddleware | None: @@ -412,6 +421,9 @@ def build_middlewares( List of middleware instances. """ resolved_app_config = app_config or get_app_config() + from deerflow.extensions import get_agent_build_extensions + + resolved_extensions = extensions if extensions is not None else get_agent_build_extensions() runtime_middleware_kwargs = { "app_config": resolved_app_config, "lazy_init": True, @@ -469,7 +481,11 @@ def build_middlewares( ) # Add summarization middleware if enabled - summarization_middleware = _create_summarization_middleware(app_config=resolved_app_config, run_model_name=model_name) + summarization_middleware = _create_summarization_middleware( + app_config=resolved_app_config, + run_model_name=model_name, + extensions=resolved_extensions, + ) if summarization_middleware is not None: middlewares.append(summarization_middleware) @@ -485,7 +501,12 @@ def build_middlewares( middlewares.append(TokenUsageMiddleware()) # Add TitleMiddleware - middlewares.append(TitleMiddleware(app_config=resolved_app_config)) + middlewares.append( + TitleMiddleware( + app_config=resolved_app_config, + extensions=resolved_extensions, + ) + ) # Add MemoryMiddleware after TitleMiddleware. Tool mode normally skips it; # conversation-extraction backends may explicitly retain passive writes. @@ -587,10 +608,8 @@ def build_middlewares( # above, changing what "the final request" means for observers. from deerflow_extension_api import AgentScope - from deerflow.extensions import get_agent_build_extensions from deerflow.extensions.stack import compose_with_extensions - resolved_extensions = extensions if extensions is not None else get_agent_build_extensions() if not resolved_extensions.has_middleware_contributors: return compose_with_extensions(middlewares, AgentScope.LEAD, None, resolved_extensions) diff --git a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py index 78522e04f..a017cf3f9 100644 --- a/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py +++ b/backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py @@ -9,6 +9,7 @@ import json import logging import math import re +import time import uuid from collections import OrderedDict from datetime import UTC, datetime, timedelta @@ -1494,7 +1495,34 @@ class MemoryUpdater: ) logger.info("Invoking memory-update LLM (thread=%s trace_id=%s)", thread_id, trace_id) attempted = True - response = model.invoke(prompt, config=invoke_config) + started = time.monotonic() + try: + response = model.invoke(prompt, config=invoke_config) + # Deliberately broader than `Exception` so no terminal path of the + # provider call goes unobserved. This is NOT about asyncio + # cancellation: this whole method runs on a worker thread (the + # debounce timer, or the executor `update_memory` offloads to), and + # cancelling the awaiting side never interrupts a running thread, so + # `CancelledError` cannot arrive here. Reporting costs nothing on + # this path either way — the hook below is a non-blocking submit. + except BaseException as exc: + self._notify_llm_result( + invoke_config, + prompt=prompt, + response=None, + error=exc, + started=started, + model_name=model_name, + ) + raise + self._notify_llm_result( + invoke_config, + prompt=prompt, + response=response, + error=None, + started=started, + model_name=model_name, + ) success = self._finalize_update( current_memory=current_memory, response_content=response.content, @@ -1533,6 +1561,37 @@ class MemoryUpdater: success=success, ) + def _notify_llm_result( + self, + invoke_config: dict[str, Any], + *, + prompt: Any, + response: Any, + error: BaseException | None, + started: float, + model_name: str | None, + ) -> None: + """Fire the optional host result hook without affecting the update.""" + if self._callbacks is None: + return + hook = getattr(self._callbacks, "on_memory_llm_result", None) + if hook is None: + return + try: + hook( + invoke_config, + prompt=prompt, + response=response, + error=error, + duration_ms=(time.monotonic() - started) * 1000, + model_name=model_name, + ) + except Exception: + # Only the hook's own failures are non-fatal. `SystemExit` / + # `KeyboardInterrupt` mean the process is going down and must not be + # swallowed by an observability path. + logger.warning("Memory LLM result hook failed (non-fatal)", exc_info=True) + def update_memory( self, messages: list[Any], diff --git a/backend/packages/harness/deerflow/agents/memory/manager.py b/backend/packages/harness/deerflow/agents/memory/manager.py index 0335d6dbc..7edee32d6 100644 --- a/backend/packages/harness/deerflow/agents/memory/manager.py +++ b/backend/packages/harness/deerflow/agents/memory/manager.py @@ -63,6 +63,24 @@ class MemoryCallbacks: """Pre-LLM-call: mutate ``invoke_config`` (e.g. merge trace metadata) before the backend invokes the model. Default: no-op.""" + def on_memory_llm_result( + self, + invoke_config: dict[str, Any], + *, + prompt: Any, + response: Any, + error: BaseException | None, + duration_ms: float, + model_name: str | None, + ) -> None: + """Post-LLM-call hook for host-owned observation. Default: no-op. + + This callback keeps the vendorable DeerMem backend independent from + DeerFlow's extension API. It is invoked for both provider success and + failure, and backend callers isolate exceptions raised by an + implementation. + """ + class MemoryManagerError(RuntimeError): """Backend-neutral base error exposed at the MemoryManager boundary.""" @@ -623,6 +641,13 @@ class LangfuseMemoryCallbacks(MemoryCallbacks): langfuse is not an enabled tracing provider. """ + def __init__(self, *, extensions=None) -> None: + if extensions is None: + from deerflow.extensions import get_loaded_extensions + + extensions = get_loaded_extensions() + self._extensions = extensions + def on_memory_llm_call( self, invoke_config: dict[str, Any], @@ -644,6 +669,60 @@ class LangfuseMemoryCallbacks(MemoryCallbacks): deerflow_trace_id=trace_id, ) + def on_memory_llm_result( + self, + invoke_config: dict[str, Any], + *, + prompt: Any, + response: Any, + error: BaseException | None, + duration_ms: float, + model_name: str | None, + ) -> None: + """Forward a DeerMem provider result using the captured snapshot.""" + extensions = self._extensions + if not extensions.has_system_model_observers: + return + try: + from deerflow_extension_api import ( + SystemModelRequest, + SystemModelResult, + SystemOperationKind, + ) + + from deerflow.extensions.notify import ( + dispatch_system_model_observation, + notify_system_model_call, + task_store_for_system_call, + ) + + dispatch_system_model_observation( + notify_system_model_call( + extensions, + task_store_for_system_call(invoke_config), + SystemOperationKind.MEMORY, + SystemModelRequest( + messages=prompt, + model_name=model_name, + invoke_config=invoke_config, + ), + SystemModelResult( + response=response, + error=error, + duration_ms=duration_ms, + ), + ), + SystemOperationKind.MEMORY.value, + ) + except Exception: + # Only the bridge's own failures are non-fatal. A teardown signal + # must propagate, matching the boundary the DeerMem-side call + # site documents and tests. + logger.warning( + "Extension observation of the memory model call failed (non-fatal)", + exc_info=True, + ) + def _host_default_should_keep_hidden_message(additional_kwargs: Any) -> bool: """deer-flow default for DeerMem's ``should_keep_hidden_message`` slot. diff --git a/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py index 44e6cf4ac..3b87db6af 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py @@ -105,6 +105,7 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): configured_model_name: str | None = None, run_model_name: str | None = None, anchor_model_name: str | None = _UNSET, # type: ignore[assignment] + extensions=None, **kwargs, ) -> None: super().__init__(*args, **kwargs) @@ -139,6 +140,11 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): self._anchor_model_name = configured_model_name or self._default_model_name() else: self._anchor_model_name = anchor_model_name + if extensions is None: + from deerflow.extensions import get_agent_build_extensions + + extensions = get_agent_build_extensions() + self._extensions = extensions # Nostream generation models built lazily by name and cached (None = a build # that failed, so a broken candidate config is not retried every turn and does # not escape the fail-open boundary). @@ -270,14 +276,26 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): return text return None - async def _asummarize_with(self, messages_to_summarize: list[AnyMessage], previous_summary: str | None = None) -> str | None: + async def _asummarize_with( + self, + messages_to_summarize: list[AnyMessage], + previous_summary: str | None = None, + *, + task_store=None, + ) -> str | None: """Async counterpart of :meth:`_summarize_with` using the nostream model.""" prompt = self._prepare_summary_prompt(messages_to_summarize, previous_summary) if prompt is None or prompt in _CANNED_SUMMARIES: return prompt names = self._generation_candidate_names() for index, name in enumerate(names): - text = await self._ainvoke_summary(self._model_for(name), prompt, last=index == len(names) - 1) + text = await self._ainvoke_summary( + self._model_for(name), + prompt, + last=index == len(names) - 1, + model_name=name, + task_store=task_store, + ) if text is not None: return text return None @@ -289,6 +307,13 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): response's ``.text`` is part of consuming the provider result, so a failing accessor must convert to a candidate failure (fall through) rather than escape the fail-open boundary. + + Deliberately unobserved by system-model-call extensions, unlike + :meth:`_ainvoke_summary`. Both this method and its only host caller + (``compact_state``) are the sync half of an async-only runtime: the agent runs + through ``abefore_model``, and ``runtime/context_compaction.py`` calls + ``acompact_state``. Notifying from here would have to block the calling thread + on the extension loop for a call site the host never reaches. """ if model is None: return None @@ -299,12 +324,37 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): self._log_summary_error(last) return None - async def _ainvoke_summary(self, model: Any | None, prompt: str, *, last: bool = False) -> str | None: + async def _ainvoke_summary( + self, + model: Any | None, + prompt: str, + *, + last: bool = False, + model_name: str | None = None, + task_store=None, + ) -> str | None: """Async counterpart of :meth:`_invoke_summary`.""" if model is None: return None try: - response = await model.ainvoke(prompt, config={"metadata": {"lc_source": "summarization"}}) + invoke_config = {"metadata": {"lc_source": "summarization"}} + extensions = getattr(self, "_extensions", None) + if extensions is None: + response = await model.ainvoke(prompt, config=invoke_config) + else: + from deerflow_extension_api import SystemOperationKind + + from deerflow.extensions.notify import observe_system_model_call + + response = await observe_system_model_call( + extensions, + SystemOperationKind.SUMMARIZATION, + messages=prompt, + model_name=model_name, + invoke_config=invoke_config, + invoke=lambda: model.ainvoke(prompt, config=invoke_config), + task_store=task_store, + ) return self._checked_summary(response, last) except Exception: self._log_summary_error(last) @@ -530,7 +580,13 @@ class DeerFlowSummarizationMiddleware(SummarizationMiddleware): if prepared is None: return None messages_to_summarize, preserved_messages, previous_summary, total_tokens = prepared - summary = await self._asummarize_with(messages_to_summarize, previous_summary=previous_summary) + from deerflow_extension_api import task_store_from_runtime + + summary = await self._asummarize_with( + messages_to_summarize, + previous_summary=previous_summary, + task_store=task_store_from_runtime(runtime), + ) if summary is None: if raise_on_failure: raise SummaryGenerationError("summary generation failed") @@ -678,6 +734,7 @@ def create_summarization_middleware( keep: tuple[str, int | float] | None = None, skip_memory_flush: bool = False, run_model_name: str | None = None, + extensions=None, ) -> DeerFlowSummarizationMiddleware | None: """Create the configured summarization middleware. @@ -753,4 +810,5 @@ def create_summarization_middleware( configured_model_name=config.model_name, run_model_name=run_model_name, anchor_model_name=anchor_name, + extensions=extensions, ) diff --git a/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py index b2ac92c27..137aeb89f 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/title_middleware.py @@ -32,10 +32,21 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): state_schema = TitleMiddlewareState - def __init__(self, *, app_config: "AppConfig | None" = None, title_config: "TitleConfig | None" = None): + def __init__( + self, + *, + app_config: "AppConfig | None" = None, + title_config: "TitleConfig | None" = None, + extensions=None, + ): super().__init__() self._app_config = app_config self._title_config = title_config + if extensions is None: + from deerflow.extensions import get_agent_build_extensions + + extensions = get_agent_build_extensions() + self._extensions = extensions def _get_title_config(self): if self._title_config is not None: @@ -196,7 +207,12 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): user_msg = self._get_title_user_message(state) return {"title": self._fallback_title(user_msg)} - async def _agenerate_title_result(self, state: TitleMiddlewareState) -> dict | None: + async def _agenerate_title_result( + self, + state: TitleMiddlewareState, + *, + task_store=None, + ) -> dict | None: """Generate a configured LLM title asynchronously and fall back locally.""" if not self._should_generate_title(state): return None @@ -218,7 +234,21 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): if self._app_config is not None: model_kwargs["app_config"] = self._app_config model = create_chat_model(name=config.model_name, **model_kwargs) - response = await model.ainvoke(prompt, config=self._get_runnable_config()) + invoke_config = self._get_runnable_config() + + from deerflow_extension_api import SystemOperationKind + + from deerflow.extensions.notify import observe_system_model_call + + response = await observe_system_model_call( + self._extensions, + SystemOperationKind.TITLE, + messages=prompt, + model_name=config.model_name, + invoke_config=invoke_config, + invoke=lambda: model.ainvoke(prompt, config=invoke_config), + task_store=task_store, + ) title = self._parse_title(response.content) if title: return {"title": title} @@ -232,4 +262,9 @@ class TitleMiddleware(AgentMiddleware[TitleMiddlewareState]): @override async def aafter_model(self, state: TitleMiddlewareState, runtime: Runtime) -> dict | None: - return await self._agenerate_title_result(state) + from deerflow_extension_api import task_store_from_runtime + + return await self._agenerate_title_result( + state, + task_store=task_store_from_runtime(runtime), + ) diff --git a/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py index 92674698f..fb19d058b 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py @@ -323,6 +323,10 @@ def build_subagent_runtime_middlewares( app_config = get_app_config() + from deerflow.extensions import get_agent_build_extensions + + resolved_extensions = extensions if extensions is not None else get_agent_build_extensions() + middlewares = _build_runtime_middlewares( app_config=app_config, include_uploads=False, @@ -502,6 +506,7 @@ def build_subagent_runtime_middlewares( # model (it inherits the parent's), so passing it directly is what makes a # distinct-model subagent summarize with its own model, not the parent's. run_model_name=model_name, + extensions=resolved_extensions, ) if summarization_middleware is not None: middlewares.append(summarization_middleware) @@ -524,10 +529,8 @@ def build_subagent_runtime_middlewares( from deerflow_extension_api import AgentScope - from deerflow.extensions import get_agent_build_extensions from deerflow.extensions.stack import compose_with_extensions - resolved_extensions = extensions if extensions is not None else get_agent_build_extensions() if not resolved_extensions.has_middleware_contributors: return compose_with_extensions(middlewares, AgentScope.SUBAGENT, None, resolved_extensions) diff --git a/backend/packages/harness/deerflow/extensions/notify.py b/backend/packages/harness/deerflow/extensions/notify.py new file mode 100644 index 000000000..b4933bb8d --- /dev/null +++ b/backend/packages/harness/deerflow/extensions/notify.py @@ -0,0 +1,418 @@ +"""Fail-open notification helpers for extension runtime hooks.""" + +from __future__ import annotations + +import asyncio +import logging +import time +from collections.abc import Awaitable, Callable, Coroutine, Mapping +from typing import Any + +from deerflow_extension_api import ( + EXTENSION_TASK_STORE_KEY, + ExtensionData, + SystemModelRequest, + SystemModelResult, + SystemOperationKind, + TaskInfo, + TaskOutcome, +) + +from deerflow.extensions.registry import LoadedExtensions + +logger = logging.getLogger(__name__) + + +def lead_task_id(run_id: str) -> str: + """Return the stable task id for a lead run, including continuations.""" + return run_id + + +def lead_task_outcome(*, aborted: bool, succeeded: bool) -> TaskOutcome: + """Classify a lead run conservatively from its terminal state.""" + if aborted: + return TaskOutcome.ABORTED + if succeeded: + return TaskOutcome.COMPLETED + return TaskOutcome.FAILED + + +def subagent_task_outcome(*, cancelled: bool, succeeded: bool) -> TaskOutcome: + """Classify a subagent execution conservatively from its terminal state.""" + if cancelled: + return TaskOutcome.ABORTED + if succeeded: + return TaskOutcome.COMPLETED + return TaskOutcome.FAILED + + +def _host_is_cancelling() -> bool: + """Whether the host task itself is being cancelled. + + Fail-open has to be decided by the *origin* of a failure, not by its base + class. ``CancelledError`` reaches a contributor's ``except`` for two very + different reasons: the host task was cancelled (must propagate), or the + contributor raised it on its own — an extension implementing an internal + timeout with cancellation, for instance (must stay contained). Only the + first increments the task's cancellation counter, so it is what tells the + two apart. + """ + task = asyncio.current_task() + return task is not None and task.cancelling() > 0 + + +async def _notify_each( + contributors: tuple[tuple[str, Any], ...], + hook: str, + invoke: Callable[[Any], Any], + task_id: str, + timeout: float | None, +) -> None: + """Invoke contributors in order, fail-open, within one shared budget.""" + loop = asyncio.get_running_loop() + deadline = None if timeout is None else loop.time() + timeout + for source, contributor in contributors: + try: + call = invoke(contributor) + if deadline is None: + await call + continue + + remaining = deadline - loop.time() + if remaining <= 0: + close = getattr(call, "close", None) + if callable(close): + close() + logger.warning( + "Extension %s: %s skipped for task %s; the %.1fs notification budget was spent", + source, + hook, + task_id, + timeout, + ) + continue + await asyncio.wait_for(call, remaining) + except TimeoutError: + if deadline is not None and loop.time() >= deadline: + # Budget exhaustion mid-hook is the same expected operational + # condition as the skip above, so it stays a warning rather + # than a hook failure with an asyncio-internal traceback. + logger.warning( + "Extension %s: %s timed out for task %s; the %.1fs notification budget was spent", + source, + hook, + task_id, + timeout, + ) + else: + # A TimeoutError the contributor raised on its own is a hook + # failure like any other. + logger.exception( + "Extension %s: %s failed for task %s", + source, + hook, + task_id, + ) + except asyncio.CancelledError: + if _host_is_cancelling(): + raise + # The contributor raised it, so containing it keeps one broken + # extension from skipping its successors — and, at the task-stop + # site, from turning a run's cleanup into a deferred interrupt. + logger.exception( + "Extension %s: %s raised CancelledError for task %s", + source, + hook, + task_id, + ) + except Exception: + logger.exception( + "Extension %s: %s failed for task %s", + source, + hook, + task_id, + ) + + +# Gateway registers its serving loop here. Subagents can run on isolated event +# loops, but extension resources must always be touched on the loop where they +# were started. +_notify_loop: asyncio.AbstractEventLoop | None = None +_pending_dispatches: set[asyncio.Future[Any]] = set() +_warned_no_loop = False +_system_observations_enabled = True + + +def set_extension_notify_loop(loop: asyncio.AbstractEventLoop | None) -> None: + """Bind extension notifications to the loop that owns extension resources.""" + global _notify_loop, _system_observations_enabled, _warned_no_loop + _notify_loop = loop + _system_observations_enabled = True + _warned_no_loop = False + + +def reset_extension_notify_loop() -> None: + """Remove the process-wide loop binding during host shutdown or tests.""" + global _notify_loop, _system_observations_enabled, _warned_no_loop + _notify_loop = None + _system_observations_enabled = True + _warned_no_loop = False + _pending_dispatches.clear() + + +def suspend_extension_system_observations() -> None: + """Drop new fire-and-forget observations while awaited hooks still drain.""" + global _system_observations_enabled + if _notify_loop is not None: + _system_observations_enabled = False + + +async def _notify_each_on_extension_loop( + contributors: tuple[tuple[str, Any], ...], + hook: str, + invoke: Callable[[Any], Any], + task_id: str, + timeout: float | None, +) -> None: + loop = _notify_loop + current_loop = asyncio.get_running_loop() + if loop is None or loop is current_loop: + await _notify_each(contributors, hook, invoke, task_id, timeout) + return + if not loop.is_running(): + logger.warning( + "No running loop registered for awaited extension hook; %s for %s was dropped", + hook, + task_id, + ) + return + + notification = _notify_each(contributors, hook, invoke, task_id, timeout) + try: + future = asyncio.run_coroutine_threadsafe(notification, loop) + except Exception: + notification.close() + logger.exception( + "Could not dispatch extension %s for task %s to the registered loop", + hook, + task_id, + ) + return + + try: + wrapped = asyncio.wrap_future(future) + if timeout is None: + await wrapped + else: + await asyncio.wait_for(wrapped, timeout) + except TimeoutError: + future.cancel() + logger.warning( + "Extension %s dispatch timed out for task %s after %.1fs", + hook, + task_id, + timeout, + ) + except asyncio.CancelledError: + future.cancel() + raise + except Exception: + logger.exception( + "Extension %s dispatch failed for task %s", + hook, + task_id, + ) + + +async def notify_task_start( + extensions: LoadedExtensions, + task_store: ExtensionData, + info: TaskInfo, + *, + timeout: float | None = None, +) -> None: + await _notify_each_on_extension_loop( + extensions.task_lifecycle, + "on_task_start", + lambda contributor: contributor.on_task_start( + extensions.app_store, + task_store, + info, + ), + info.task_id, + timeout, + ) + + +async def notify_task_stop( + extensions: LoadedExtensions, + task_store: ExtensionData, + info: TaskInfo, + outcome: TaskOutcome, + *, + timeout: float | None = None, +) -> None: + await _notify_each_on_extension_loop( + extensions.task_lifecycle, + "on_task_stop", + lambda contributor: contributor.on_task_stop( + extensions.app_store, + task_store, + info, + outcome, + ), + info.task_id, + timeout, + ) + + +async def notify_system_model_call( + extensions: LoadedExtensions, + task_store: ExtensionData | None, + kind: SystemOperationKind, + request: SystemModelRequest, + result: SystemModelResult, + *, + timeout: float | None = None, +) -> None: + """Notify the observers from one immutable extension snapshot.""" + if not extensions.system_model_observers: + return + store = task_store if task_store is not None else ExtensionData("detached") + await _notify_each_on_extension_loop( + extensions.system_model_observers, + "on_system_model_call", + lambda observer: observer.on_system_model_call( + extensions.app_store, + store, + kind, + request, + result, + ), + f"{store.scope_id} ({kind.value})", + timeout, + ) + + +def task_store_for_system_call(invoke_config: object) -> ExtensionData | None: + """Recover the live task store from a legacy top-level runtime context.""" + if not isinstance(invoke_config, Mapping): + return None + context = invoke_config.get("context") + if not isinstance(context, Mapping): + return None + store = context.get(EXTENSION_TASK_STORE_KEY) + return store if isinstance(store, ExtensionData) else None + + +async def observe_system_model_call( + extensions: LoadedExtensions, + kind: SystemOperationKind, + *, + messages: Any, + model_name: str | None, + invoke_config: Any, + invoke: Callable[[], Awaitable[Any]], + task_store: ExtensionData | None = None, + timeout: float | None = None, +) -> Any: + """Invoke a system-owned model call and report either terminal path.""" + if not extensions.has_system_model_observers: + return await invoke() + + store = task_store if task_store is not None else task_store_for_system_call(invoke_config) + request = SystemModelRequest( + messages=messages, + model_name=model_name, + invoke_config=(invoke_config if isinstance(invoke_config, Mapping) else None), + ) + started = time.monotonic() + try: + response = await invoke() + except asyncio.CancelledError as exc: + # Cancellation is a terminal path as well: interrupt/rollback admission + # and shutdown both cancel the run task, so a user sending a follow-up + # mid-run routinely ends a goal or summarization call here, with the + # provider tokens already spent. Awaiting observers would be unreliable + # — a repeated cancel interrupts that await before any of them runs — so + # this reports through the same non-blocking submission the synchronous + # memory bridge uses, then propagates the cancellation untouched. A + # deployment with no registered notify loop drops it, exactly as that + # bridge does. + dispatch_system_model_observation( + notify_system_model_call( + extensions, + store, + kind, + request, + SystemModelResult( + error=exc, + duration_ms=(time.monotonic() - started) * 1000, + ), + ), + kind.value, + ) + raise + except Exception as exc: + await notify_system_model_call( + extensions, + store, + kind, + request, + SystemModelResult( + error=exc, + duration_ms=(time.monotonic() - started) * 1000, + ), + timeout=timeout, + ) + raise + await notify_system_model_call( + extensions, + store, + kind, + request, + SystemModelResult( + response=response, + duration_ms=(time.monotonic() - started) * 1000, + ), + timeout=timeout, + ) + return response + + +def dispatch_system_model_observation( + coro: Coroutine[Any, Any, None], + what: str, +) -> bool: + """Submit a synchronous call site's observation to the registered loop.""" + global _warned_no_loop + + loop = _notify_loop + submitted = False + try: + if not _system_observations_enabled: + return False + if loop is None or not loop.is_running(): + if not _warned_no_loop: + _warned_no_loop = True + logger.warning( + "No running loop registered for extension observations; %s and later ones are dropped", + what, + ) + return False + try: + future = asyncio.run_coroutine_threadsafe(coro, loop) + except Exception: + logger.debug( + "Could not dispatch %s to the extension notify loop", + what, + exc_info=True, + ) + return False + _pending_dispatches.add(future) + future.add_done_callback(_pending_dispatches.discard) + submitted = True + return True + finally: + if not submitted: + coro.close() diff --git a/backend/packages/harness/deerflow/extensions/registry.py b/backend/packages/harness/deerflow/extensions/registry.py index af7179cd5..7dcbfb4f0 100644 --- a/backend/packages/harness/deerflow/extensions/registry.py +++ b/backend/packages/harness/deerflow/extensions/registry.py @@ -12,7 +12,12 @@ from contextlib import contextmanager from dataclasses import dataclass from typing import Any -from deerflow_extension_api import ExtensionData, MiddlewareContributor +from deerflow_extension_api import ( + ExtensionData, + MiddlewareContributor, + SystemModelCallObserver, + TaskLifecycleContributor, +) from deerflow_extension_api import ExtensionRegistry as ExtensionRegistryContract _Entry = tuple[str, Any] @@ -28,10 +33,14 @@ class LoadedExtensions: app_store: ExtensionData middleware_contributors: tuple[tuple[str, MiddlewareContributor], ...] = () + task_lifecycle: tuple[tuple[str, TaskLifecycleContributor], ...] = () + system_model_observers: tuple[tuple[str, SystemModelCallObserver], ...] = () # Precomputed attributes, not methods: hook sites read one attribute to # short-circuit, so the zero-extension path constructs nothing. has_middleware_contributors: bool = False + has_task_lifecycle: bool = False + has_system_model_observers: bool = False needs_task_store: bool = False @@ -46,6 +55,8 @@ class ExtensionRegistry(ExtensionRegistryContract): def __init__(self) -> None: self._middlewares: list[_Entry] = [] + self._task_lifecycle: list[_Entry] = [] + self._system_model_observers: list[_Entry] = [] self._current_source: str | None = None @contextmanager @@ -66,6 +77,12 @@ class ExtensionRegistry(ExtensionRegistryContract): def middlewares(self, contributor: MiddlewareContributor) -> None: self._middlewares.append((self._source(), contributor)) + def task_lifecycle(self, contributor: TaskLifecycleContributor) -> None: + self._task_lifecycle.append((self._source(), contributor)) + + def system_model_observer(self, observer: SystemModelCallObserver) -> None: + self._system_model_observers.append((self._source(), observer)) + def discard(self, source: str) -> None: """Remove every entry registered by ``source``. @@ -79,27 +96,49 @@ class ExtensionRegistry(ExtensionRegistryContract): that process one install() at a time should prefer ``mark()``/``rollback_to()`` instead. """ - self._middlewares[:] = [entry for entry in self._middlewares if entry[0] != source] + for bucket in ( + self._middlewares, + self._task_lifecycle, + self._system_model_observers, + ): + bucket[:] = [entry for entry in bucket if entry[0] != source] - def mark(self) -> int: + def mark(self) -> tuple[int, int, int]: """Snapshot bucket lengths so one install() can be undone positionally.""" - return len(self._middlewares) + return ( + len(self._middlewares), + len(self._task_lifecycle), + len(self._system_model_observers), + ) - def rollback_to(self, mark: int) -> None: + def rollback_to(self, mark: tuple[int, int, int]) -> None: """Undo every registration made since ``mark``. Positional rather than source-keyed: two specs may legitimately share a ``use`` string with different config, and deleting by source would take the other instance's successful registrations with it. """ - del self._middlewares[mark:] + for bucket, size in zip( + ( + self._middlewares, + self._task_lifecycle, + self._system_model_observers, + ), + mark, + strict=True, + ): + del bucket[size:] def build(self) -> LoadedExtensions: return LoadedExtensions( app_store=ExtensionData("app"), middleware_contributors=tuple(self._middlewares), + task_lifecycle=tuple(self._task_lifecycle), + system_model_observers=tuple(self._system_model_observers), has_middleware_contributors=bool(self._middlewares), - needs_task_store=bool(self._middlewares), + has_task_lifecycle=bool(self._task_lifecycle), + has_system_model_observers=bool(self._system_model_observers), + needs_task_store=bool(self._middlewares or self._task_lifecycle or self._system_model_observers), ) diff --git a/backend/packages/harness/deerflow/runtime/goal.py b/backend/packages/harness/deerflow/runtime/goal.py index ff0be9d4a..3f603a42b 100644 --- a/backend/packages/harness/deerflow/runtime/goal.py +++ b/backend/packages/harness/deerflow/runtime/goal.py @@ -277,6 +277,8 @@ async def evaluate_goal_completion( thread_id: str | None = None, user_id: str | None = None, deerflow_trace_id: str | None = None, + task_store: Any | None = None, + extensions: Any | None = None, ) -> GoalEvaluation: """Ask a small non-thinking model whether the active goal is satisfied. @@ -320,10 +322,26 @@ async def evaluate_goal_completion( environment=_resolve_environment(), deerflow_trace_id=deerflow_trace_id, ) - response = await model.ainvoke( - [SystemMessage(content=system_instruction), HumanMessage(content=user_content)], - config=invoke_config, - ) + prompt_messages = [ + SystemMessage(content=system_instruction), + HumanMessage(content=user_content), + ] + if extensions is None: + response = await model.ainvoke(prompt_messages, config=invoke_config) + else: + from deerflow_extension_api import SystemOperationKind + + from deerflow.extensions.notify import observe_system_model_call + + response = await observe_system_model_call( + extensions, + SystemOperationKind.GOAL, + messages=prompt_messages, + model_name=model_name, + invoke_config=invoke_config, + invoke=lambda: model.ainvoke(prompt_messages, config=invoke_config), + task_store=task_store, + ) return parse_goal_evaluation_response(_extract_response_text(response.content)) diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index b307c393c..52b73d8ac 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -111,6 +111,7 @@ async def _checkpoint_thread_lock(thread_id: str) -> AsyncIterator[None]: _DELIVERY_RECEIPT_RETRY_DELAYS_SECONDS = (0.1, 0.5) +_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS = 3.0 async def _persist_delivery_receipt( @@ -558,12 +559,20 @@ async def run_agent( run_id = record.run_id thread_id = record.thread_id - from deerflow_extension_api import ExtensionData + 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_stop_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 @@ -677,8 +686,25 @@ async def run_agent( return started = True + task_id = lead_task_id(run_id) if extensions.needs_task_store: - task_store = ExtensionData(run_id) + 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: @@ -990,6 +1016,8 @@ async def run_agent( abort_event=record.abort_event, user_id=resolve_runtime_user_id(runtime), deerflow_trace_id=deerflow_trace_id, + task_store=task_store, + extensions=extensions, ) if continuation_input is None or record.abort_event.is_set(): break @@ -1240,12 +1268,44 @@ async def run_agent( await ctx.on_run_completed(record) except Exception: logger.warning("Run completion hook failed for %s (non-fatal)", run_id, exc_info=True) + + 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. + try: + await 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, + ) + 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_stop_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) asyncio.create_task(bridge.cleanup(run_id, delay=60)) + if deferred_stop_interrupt is not None: + raise deferred_stop_interrupt + # --------------------------------------------------------------------------- # Helpers @@ -1410,6 +1470,8 @@ async def _prepare_goal_continuation_input( 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, ) -> dict[str, Any] | None: """Evaluate the active goal and return a hidden continuation input if needed. @@ -1490,6 +1552,8 @@ async def _prepare_goal_continuation_input( thread_id=thread_id, user_id=user_id, deerflow_trace_id=deerflow_trace_id, + task_store=task_store, + extensions=extensions, ) if abort_event is not None and abort_event.is_set(): return None diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index 0e4b70cc2..86e1b671b 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -46,6 +46,8 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS = 3.0 + _previous_shutdown_isolated_subagent_loop = globals().get("_shutdown_isolated_subagent_loop") if callable(_previous_shutdown_isolated_subagent_loop): @@ -807,14 +809,37 @@ class SubagentExecutor: status=SubagentStatus.RUNNING, started_at=datetime.now(), ) + from deerflow_extension_api import ExtensionData, TaskInfo + from deerflow.extensions import get_loaded_extensions + from deerflow.extensions.notify import ( + lead_task_id, + notify_task_start, + notify_task_stop, + subagent_task_outcome, + ) loaded_extensions = self.extensions if self.extensions is not None else get_loaded_extensions() - task_store = None + task_store: ExtensionData | None = None + task_info: TaskInfo | None = None if loaded_extensions.needs_task_store: - from deerflow_extension_api import ExtensionData - task_store = ExtensionData(result.task_id) + if loaded_extensions.has_task_lifecycle and self.run_id: + task_info = TaskInfo( + task_id=result.task_id, + run_id=self.run_id, + thread_id=self.thread_id or "", + kind="subagent", + parent_task_id=lead_task_id(self.run_id), + agent_name=self.config.name, + ) + assert task_store is not None + elif loaded_extensions.has_task_lifecycle: + logger.debug( + "[trace=%s] Subagent %s has no run_id; skipping extension task lifecycle", + self.trace_id, + self.config.name, + ) ai_messages = result.ai_messages if ai_messages is None: ai_messages = [] @@ -831,6 +856,14 @@ class SubagentExecutor: collector: SubagentTokenCollector | None = None try: + if task_info is not None and task_store is not None: + await notify_task_start( + loaded_extensions, + task_store, + task_info, + timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS, + ) + state, final_tools, deferred_setup = await self._build_initial_state(task) agent = self._create_agent( final_tools, @@ -1050,6 +1083,27 @@ class SubagentExecutor: token_usage_records=collector.snapshot_records() if collector is not None else None, ) + finally: + if task_info is not None and task_store is not None: + try: + await notify_task_stop( + loaded_extensions, + task_store, + task_info, + subagent_task_outcome( + cancelled=result.status is SubagentStatus.CANCELLED, + succeeded=result.status is SubagentStatus.COMPLETED, + ), + timeout=_EXTENSION_TASK_NOTIFY_TIMEOUT_SECONDS, + ) + except Exception: + logger.warning( + "[trace=%s] Extension task-stop notification failed for subagent %s (non-fatal)", + self.trace_id, + self.config.name, + exc_info=True, + ) + return result def _execute_in_isolated_loop(self, task: str, result_holder: SubagentResult | None = None) -> SubagentResult: diff --git a/backend/packages/harness/pyproject.toml b/backend/packages/harness/pyproject.toml index 0685e2231..e6e9db2a9 100644 --- a/backend/packages/harness/pyproject.toml +++ b/backend/packages/harness/pyproject.toml @@ -11,7 +11,7 @@ dependencies = [ # the contract version it implements, extensions declare ranges. A range # here would let pip resolve a newer contract package than this harness # implements, making newer extensions look supported at runtime. - "deerflow-extension-api==0.1.0", + "deerflow-extension-api==0.1.1", "dotenv>=0.9.9", "exa-py>=1.0.0", "httpx>=0.28.0", diff --git a/backend/tests/test_extension_api_contracts.py b/backend/tests/test_extension_api_contracts.py index d0681a15c..37a4c2bcb 100644 --- a/backend/tests/test_extension_api_contracts.py +++ b/backend/tests/test_extension_api_contracts.py @@ -8,6 +8,7 @@ are asserted here. from __future__ import annotations +import asyncio import dataclasses import importlib.resources import inspect @@ -24,6 +25,13 @@ from deerflow_extension_api import ( MiddlewareContributor, MiddlewarePlacement, Placement, + SystemModelCallObserver, + SystemModelRequest, + SystemModelResult, + SystemOperationKind, + TaskInfo, + TaskLifecycleContributor, + TaskOutcome, extension, ) from deerflow_extension_api.runtime_bridge import ( @@ -56,6 +64,9 @@ def test_middleware_placement_defaults(): [ HostPolicySnapshot, AgentBuildContext, + TaskInfo, + SystemModelRequest, + SystemModelResult, MiddlewarePlacement, ], ) @@ -66,16 +77,23 @@ def test_every_dataclass_is_frozen(cls): @pytest.mark.parametrize( "cls", - [HostPolicySnapshot], + [HostPolicySnapshot, TaskInfo, SystemModelRequest, SystemModelResult], ) def test_additive_dataclasses_are_constructible_with_required_fields_only(cls): """Fields added later must carry defaults, or old extensions break on upgrade. - HostPolicySnapshot is host-constructed and fully optional. - AgentBuildContext gets its own dedicated test below because its scope is - legitimately required. + HostPolicySnapshot and the two system-call snapshots are host-constructed + and fully optional. TaskInfo has a required identity core and optional + remainder. AgentBuildContext gets its own dedicated test below because its + scope is legitimately required. """ - assert cls() is not None + if cls is TaskInfo: + info = cls(task_id="t", run_id="r", thread_id="th", kind="lead") + assert info.parent_task_id is None + assert info.agent_name is None + assert info.resumed is False + else: + assert cls() is not None def test_agent_build_context_optional_fields_keep_their_defaults(): @@ -97,6 +115,8 @@ def test_agent_build_context_optional_fields_keep_their_defaults(): [ ExtensionRegistry, MiddlewareContributor, + TaskLifecycleContributor, + SystemModelCallObserver, ], ) def test_every_protocol_method_has_a_default_implementation(protocol): @@ -123,6 +143,58 @@ def test_contributor_defaults_return_empty(): assert MiddlewareContributor.contribute_middlewares(bare, ExtensionData("app"), AgentBuildContext(scope=AgentScope.LEAD)) == () +def test_task_lifecycle_contract_is_public_and_defaults_to_noop(): + class _Bare: + pass + + app_store = ExtensionData("app") + task_store = ExtensionData("task-1") + info = TaskInfo( + task_id="task-1", + run_id="run-1", + thread_id="thread-1", + kind="lead", + ) + + assert TaskOutcome.COMPLETED.value == "completed" + assert asyncio.run(TaskLifecycleContributor.on_task_start(_Bare(), app_store, task_store, info)) is None + assert asyncio.run(TaskLifecycleContributor.on_task_stop(_Bare(), app_store, task_store, info, TaskOutcome.COMPLETED)) is None + + +def test_system_model_observer_contract_reports_success_and_failure_shapes(): + class _Bare: + pass + + app_store = ExtensionData("app") + task_store = ExtensionData("task-1") + request = SystemModelRequest(messages=("prompt",), model_name="system-model") + success = SystemModelResult(response="answer", duration_ms=1.5) + failure = SystemModelResult(error=RuntimeError("provider failed"), duration_ms=2.0) + + assert SystemOperationKind.GOAL.value == "goal" + assert asyncio.run(SystemModelCallObserver.on_system_model_call(_Bare(), app_store, task_store, SystemOperationKind.GOAL, request, success)) is None + assert asyncio.run(SystemModelCallObserver.on_system_model_call(_Bare(), app_store, task_store, SystemOperationKind.GOAL, request, failure)) is None + + +def test_system_model_request_normalizes_messages_into_an_immutable_sequence(): + """``messages`` is a snapshot of a message sequence, never a per-character view. + + Title and summarization pass a single prompt string, so a bare ``str`` must not + reach observers as a ``Sequence`` whose items are characters. A live ``list`` from + a call site must also be copied: the snapshot is documented as read-only, and the + caller keeps mutating its own list after the observation is dispatched. + """ + assert SystemModelRequest(messages="one prompt").messages == ("one prompt",) + + live: list[str] = ["first"] + request = SystemModelRequest(messages=live) + live.append("second") + assert request.messages == ("first",) + + assert SystemModelRequest().messages == () + assert SystemModelRequest(messages=("already", "a", "tuple")).messages == ("already", "a", "tuple") + + def test_future_contribution_points_are_not_advertised_before_the_host_supports_them(): """A merged slice must not silently accept registrations it cannot run.""" import deerflow_extension_api @@ -130,8 +202,6 @@ def test_future_contribution_points_are_not_advertised_before_the_host_supports_ for name in ( "ExtensionRuntimeDeps", "ExtensionService", - "SystemModelCallObserver", - "TaskLifecycleContributor", ): assert name not in deerflow_extension_api.__all__ assert not hasattr(deerflow_extension_api, name) @@ -166,6 +236,14 @@ def test_extension_decorator_stamps_api_requirement(): assert install.__deerflow_name__ == "demo" +def test_task_outcome_members(): + assert {outcome.value for outcome in TaskOutcome} == {"completed", "aborted", "failed"} + + +def test_system_operation_kind_members(): + assert {kind.value for kind in SystemOperationKind} == {"goal", "memory", "title", "summarization"} + + def test_registry_and_install_alias_are_part_of_the_public_surface(): """Independent extensions annotate install(registry, config) against the contract package alone — importing the host's concrete registry would pin @@ -211,4 +289,5 @@ def test_runtime_api_version_matches_the_installed_contract_package(): """Every additive contract slice bumps both gates together.""" from importlib.metadata import version + assert API_VERSION == "0.1.1" assert API_VERSION == version("deerflow-extension-api") diff --git a/backend/tests/test_extension_config.py b/backend/tests/test_extension_config.py index 236daa8e5..a8c2f65ff 100644 --- a/backend/tests/test_extension_config.py +++ b/backend/tests/test_extension_config.py @@ -94,6 +94,8 @@ def test_plugins_is_registered_as_startup_only(): def test_singleton_defaults_to_empty(): loaded = get_loaded_extensions() assert loaded.has_middleware_contributors is False + assert loaded.has_task_lifecycle is False + assert loaded.has_system_model_observers is False assert loaded.needs_task_store is False @@ -113,6 +115,8 @@ def test_reset_gives_a_fresh_instance(): assert after is not populated assert after is not EMPTY_EXTENSIONS assert after.has_middleware_contributors is False + assert after.has_task_lifecycle is False + assert after.has_system_model_observers is False def test_reset_does_not_leak_app_store_writes(): diff --git a/backend/tests/test_extension_loader.py b/backend/tests/test_extension_loader.py index b5a564ae2..b14b47146 100644 --- a/backend/tests/test_extension_loader.py +++ b/backend/tests/test_extension_loader.py @@ -87,6 +87,7 @@ def test_install_failure_rolls_back_partial_registration(): sources = {source for source, _ in loaded.middleware_contributors} assert sources == {f"{_FIXTURE}:install_ok"} assert len(loaded.middleware_contributors) == 1, "rollback must clear every partial registration" + assert loaded.task_lifecycle == (), "rollback must clear partial lifecycle registrations too" def test_rollback_does_not_remove_a_different_specs_registrations_sharing_the_same_use(): @@ -294,6 +295,8 @@ def test_compatible_declared_api_loads(): loaded, diagnostics = load_extensions([spec]) assert diagnostics == [] assert demo_extensions.INSTALLED == ["stamped"] + assert loaded.has_task_lifecycle is True + assert loaded.task_lifecycle[0][0] == f"{_FIXTURE}:install_stamped" def test_compatible_follows_semver_windows(): diff --git a/backend/tests/test_extension_registry.py b/backend/tests/test_extension_registry.py index 088420e22..296d0ba5c 100644 --- a/backend/tests/test_extension_registry.py +++ b/backend/tests/test_extension_registry.py @@ -15,18 +15,31 @@ class _Contributor: def test_empty_registry_builds_with_all_flags_false(): loaded = ExtensionRegistry().build() assert loaded.has_middleware_contributors is False + assert loaded.has_task_lifecycle is False + assert loaded.has_system_model_observers is False assert loaded.needs_task_store is False def test_empty_singleton_matches_an_empty_build(): assert EMPTY_EXTENSIONS.has_middleware_contributors is False + assert EMPTY_EXTENSIONS.has_task_lifecycle is False + assert EMPTY_EXTENSIONS.has_system_model_observers is False assert EMPTY_EXTENSIONS.needs_task_store is False -def test_middleware_contributions_require_a_task_store(): +@pytest.mark.parametrize( + "register", + [ + lambda registry, contributor: registry.middlewares(contributor), + lambda registry, contributor: registry.task_lifecycle(contributor), + lambda registry, contributor: registry.system_model_observer(contributor), + ], + ids=["middleware", "task-lifecycle", "system-model-observer"], +) +def test_task_scoped_contributions_require_a_task_store(register): registry = ExtensionRegistry() with registry.attributed_to("demo:install"): - registry.middlewares(_Contributor()) + register(registry, _Contributor()) assert registry.build().needs_task_store is True @@ -41,6 +54,50 @@ def test_entries_carry_their_source(): assert loaded.has_middleware_contributors is True +def test_task_lifecycle_entries_are_attributed_and_require_task_storage(): + registry = ExtensionRegistry() + contributor = _Contributor("lifecycle") + with registry.attributed_to("lifecycle_ext:install"): + registry.task_lifecycle(contributor) + + loaded = registry.build() + + assert loaded.task_lifecycle == (("lifecycle_ext:install", contributor),) + assert loaded.has_task_lifecycle is True + assert loaded.needs_task_store is True + + +def test_system_model_observers_are_attributed_and_require_task_storage(): + registry = ExtensionRegistry() + observer = _Contributor("system-model") + with registry.attributed_to("observer_ext:install"): + registry.system_model_observer(observer) + + loaded = registry.build() + + assert loaded.system_model_observers == (("observer_ext:install", observer),) + assert loaded.has_system_model_observers is True + assert loaded.needs_task_store is True + + +def test_rollback_restores_all_registration_buckets_positionally(): + registry = ExtensionRegistry() + with registry.attributed_to("keep:install"): + registry.middlewares(_Contributor("keep")) + mark = registry.mark() + with registry.attributed_to("drop:install"): + registry.middlewares(_Contributor("drop-middleware")) + registry.task_lifecycle(_Contributor("drop-lifecycle")) + registry.system_model_observer(_Contributor("drop-observer")) + + registry.rollback_to(mark) + loaded = registry.build() + + assert [contributor.tag for _, contributor in loaded.middleware_contributors] == ["keep"] + assert loaded.task_lifecycle == () + assert loaded.system_model_observers == () + + def test_registration_order_is_preserved(): registry = ExtensionRegistry() first, second = _Contributor("a"), _Contributor("b") @@ -59,11 +116,15 @@ def test_discard_removes_every_entry_of_one_source(): keep, drop = _Contributor("keep"), _Contributor("drop") with registry.attributed_to("good:install"): registry.middlewares(keep) + registry.task_lifecycle(keep) with registry.attributed_to("bad:install"): registry.middlewares(drop) + registry.system_model_observer(drop) registry.discard("bad:install") loaded = registry.build() assert loaded.middleware_contributors == (("good:install", keep),) + assert loaded.task_lifecycle == (("good:install", keep),) + assert loaded.system_model_observers == () def test_registering_outside_attributed_to_raises(): diff --git a/backend/tests/test_extension_stack_wiring.py b/backend/tests/test_extension_stack_wiring.py index 4f7cdc772..a909ca21d 100644 --- a/backend/tests/test_extension_stack_wiring.py +++ b/backend/tests/test_extension_stack_wiring.py @@ -120,6 +120,20 @@ def test_bound_build_snapshot_is_used_by_lead_and_subagent_fallbacks(): assert subagent_capture.ctx is not None +def test_lead_system_middlewares_capture_the_explicit_build_snapshot(): + from deerflow.agents.middlewares.title_middleware import TitleMiddleware + + registry = ExtensionRegistry() + with registry.attributed_to("observer:install"): + registry.system_model_observer(object()) + extensions = registry.build() + + stack = _lead_stack(extensions) + title = next(item for item in stack if isinstance(item, TitleMiddleware)) + + assert title._extensions is extensions + + def test_tool_visible_lands_at_the_outermost_position(): stack = _lead_stack(_extensions(MiddlewarePlacement(_Probe("visible"), Placement.TOOL_VISIBLE))) assert _tags(stack)[0] == "visible" diff --git a/backend/tests/test_extension_subagent_lifecycle.py b/backend/tests/test_extension_subagent_lifecycle.py new file mode 100644 index 000000000..c8115ef91 --- /dev/null +++ b/backend/tests/test_extension_subagent_lifecycle.py @@ -0,0 +1,259 @@ +"""Subagents expose the same task lifecycle contract as lead runs.""" + +from __future__ import annotations + +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from deerflow_extension_api import EXTENSION_TASK_STORE_KEY, ExtensionData, TaskInfo, TaskOutcome +from langchain_core.messages import AIMessage + +from deerflow.extensions import reset_loaded_extensions, set_loaded_extensions +from deerflow.extensions.registry import ExtensionRegistry + +_MOCKED_MODULE_NAMES = ( + "deerflow.agents", + "deerflow.agents.thread_state", + "deerflow.agents.middlewares", + "deerflow.agents.middlewares.thread_data_middleware", + "deerflow.sandbox", + "deerflow.sandbox.middleware", + "deerflow.sandbox.security", + "deerflow.models", + "deerflow.skills.storage", +) + + +@pytest.fixture +def env(): + """Import the real executor behind conftest's cycle-breaking mock.""" + reset_loaded_extensions() + original_modules = {name: sys.modules.get(name) for name in _MOCKED_MODULE_NAMES} + original_executor = sys.modules.get("deerflow.subagents.executor") + subagents_pkg = sys.modules.get("deerflow.subagents") + missing = object() + original_executor_attr = getattr(subagents_pkg, "executor", missing) if subagents_pkg is not None else missing + + sys.modules.pop("deerflow.subagents.executor", None) + if subagents_pkg is not None and hasattr(subagents_pkg, "executor"): + delattr(subagents_pkg, "executor") + + try: + for name in _MOCKED_MODULE_NAMES: + sys.modules[name] = MagicMock() + storage_module = ModuleType("deerflow.skills.storage") + storage_module.get_or_new_skill_storage = lambda **kwargs: SimpleNamespace(load_skills=lambda *, enabled_only: []) + storage_module.get_or_new_user_skill_storage = lambda user_id, **kwargs: SimpleNamespace(load_skills=lambda *, enabled_only: []) + sys.modules["deerflow.skills.storage"] = storage_module + + from deerflow.subagents.config import SubagentConfig + from deerflow.subagents.executor import ( + SubagentExecutor, + SubagentResult, + SubagentStatus, + ) + + sys.modules["deerflow.subagents.executor"].get_app_config = lambda: SimpleNamespace( + tool_search=SimpleNamespace(enabled=False), + authorization=SimpleNamespace(enabled=False), + ) + yield SimpleNamespace( + SubagentConfig=SubagentConfig, + SubagentExecutor=SubagentExecutor, + SubagentResult=SubagentResult, + SubagentStatus=SubagentStatus, + ) + finally: + reset_loaded_extensions() + for name, original in original_modules.items(): + if original is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = original + if original_executor is None: + sys.modules.pop("deerflow.subagents.executor", None) + else: + sys.modules["deerflow.subagents.executor"] = original_executor + subagents_pkg = sys.modules.get("deerflow.subagents") + if subagents_pkg is not None: + if original_executor_attr is missing: + if hasattr(subagents_pkg, "executor"): + delattr(subagents_pkg, "executor") + else: + setattr(subagents_pkg, "executor", original_executor_attr) + + +class _Recorder: + def __init__(self) -> None: + self.starts: list[TaskInfo] = [] + self.stops: list[tuple[TaskInfo, TaskOutcome]] = [] + self.stores: list[ExtensionData] = [] + + async def on_task_start(self, app_store, task_store, info): + self.starts.append(info) + self.stores.append(task_store) + + async def on_task_stop(self, app_store, task_store, info, outcome): + self.stops.append((info, outcome)) + self.stores.append(task_store) + + +def _loaded(recorder): + registry = ExtensionRegistry() + with registry.attributed_to("demo:install"): + registry.task_lifecycle(recorder) + return registry.build() + + +def _executor(env, **overrides): + config = env.SubagentConfig( + name="researcher", + description="d", + system_prompt="p", + tools=[], + ) + kwargs = {"run_id": "run-1", "thread_id": "thread-1"} + kwargs.update(overrides) + return env.SubagentExecutor(config=config, tools=[], **kwargs) + + +class _CompletingAgent: + def __init__(self, seen: dict | None = None) -> None: + self.seen = seen + + async def astream(self, *args, **kwargs): + if self.seen is not None: + self.seen["context"] = kwargs.get("context") + yield {"messages": [AIMessage(content="done")]} + + +async def _noop_initial_state(self, task): + return ({}, [], None) + + +@pytest.mark.asyncio +async def test_subagent_success_emits_shaped_start_and_completed_stop(monkeypatch, env): + recorder = _Recorder() + set_loaded_extensions(_loaded(recorder)) + executor = _executor(env) + seen: dict = {} + monkeypatch.setattr(env.SubagentExecutor, "_build_initial_state", _noop_initial_state) + monkeypatch.setattr( + env.SubagentExecutor, + "_create_agent", + lambda self, tools, **kwargs: _CompletingAgent(seen), + ) + + result = await executor._aexecute("do the thing") + + assert result.status is env.SubagentStatus.COMPLETED + [info] = recorder.starts + assert info == TaskInfo( + task_id=result.task_id, + run_id="run-1", + thread_id="thread-1", + kind="subagent", + parent_task_id="run-1", + agent_name="researcher", + ) + assert recorder.stops == [(info, TaskOutcome.COMPLETED)] + assert recorder.stores[0] is recorder.stores[1] + assert seen["context"][EXTENSION_TASK_STORE_KEY] is recorder.stores[0] + + +@pytest.mark.asyncio +async def test_subagent_failure_and_cancellation_map_to_distinct_outcomes(monkeypatch, env): + recorder = _Recorder() + set_loaded_extensions(_loaded(recorder)) + executor = _executor(env) + + async def _fail_before_agent(self, task): + raise RuntimeError("build failed") + + monkeypatch.setattr(env.SubagentExecutor, "_build_initial_state", _fail_before_agent) + failed = await executor._aexecute("fail") + assert failed.status is env.SubagentStatus.FAILED + assert recorder.stops[-1][1] is TaskOutcome.FAILED + + monkeypatch.setattr(env.SubagentExecutor, "_build_initial_state", _noop_initial_state) + monkeypatch.setattr( + env.SubagentExecutor, + "_create_agent", + lambda self, tools, **kwargs: _CompletingAgent(), + ) + holder = env.SubagentResult( + task_id="cancel-me", + trace_id="trace", + status=env.SubagentStatus.RUNNING, + ) + holder.cancel_event.set() + cancelled = await executor._aexecute("cancel", holder) + assert cancelled.status is env.SubagentStatus.CANCELLED + assert recorder.stops[-1][1] is TaskOutcome.ABORTED + assert recorder.stops[-1][0].task_id == "cancel-me" + + +@pytest.mark.asyncio +async def test_subagent_base_exception_still_emits_failed_stop(monkeypatch, env): + recorder = _Recorder() + set_loaded_extensions(_loaded(recorder)) + executor = _executor(env) + + async def _hard_stop(self, task): + raise KeyboardInterrupt("host shutdown") + + monkeypatch.setattr(env.SubagentExecutor, "_build_initial_state", _hard_stop) + with pytest.raises(KeyboardInterrupt): + await executor._aexecute("stop") + + assert recorder.stops[0][1] is TaskOutcome.FAILED + + +@pytest.mark.asyncio +async def test_subagent_without_parent_run_skips_lifecycle_but_keeps_task_store(monkeypatch, env): + recorder = _Recorder() + set_loaded_extensions(_loaded(recorder)) + executor = _executor(env, run_id=None) + seen: dict = {} + monkeypatch.setattr(env.SubagentExecutor, "_build_initial_state", _noop_initial_state) + monkeypatch.setattr( + env.SubagentExecutor, + "_create_agent", + lambda self, tools, **kwargs: _CompletingAgent(seen), + ) + + result = await executor._aexecute("direct") + + assert recorder.starts == [] + assert recorder.stops == [] + assert seen["context"][EXTENSION_TASK_STORE_KEY].scope_id == result.task_id + + +@pytest.mark.asyncio +async def test_subagent_keeps_one_snapshot_across_build_context_and_hooks(monkeypatch, env): + first = _Recorder() + second = _Recorder() + snapshot = _loaded(first) + set_loaded_extensions(snapshot) + executor = _executor(env) + seen: dict = {} + + async def _switch_singleton(self, task): + set_loaded_extensions(_loaded(second)) + return ({}, [], None) + + def _capture_agent(self, tools, *, deferred_setup=None, extensions=None): + seen["extensions"] = extensions + return _CompletingAgent(seen) + + monkeypatch.setattr(env.SubagentExecutor, "_build_initial_state", _switch_singleton) + monkeypatch.setattr(env.SubagentExecutor, "_create_agent", _capture_agent) + + await executor._aexecute("snapshot") + + assert seen["extensions"] is snapshot + assert len(first.starts) == len(first.stops) == 1 + assert second.starts == second.stops == [] + assert seen["context"][EXTENSION_TASK_STORE_KEY] is first.stores[0] diff --git a/backend/tests/test_extension_system_model_calls.py b/backend/tests/test_extension_system_model_calls.py new file mode 100644 index 000000000..5203127f4 --- /dev/null +++ b/backend/tests/test_extension_system_model_calls.py @@ -0,0 +1,727 @@ +"""Observation helpers for model calls made outside the agent graph.""" + +from __future__ import annotations + +import asyncio +import logging +import threading +import time +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from deerflow_extension_api import ( + EXTENSION_TASK_STORE_KEY, + ExtensionData, + SystemModelRequest, + SystemModelResult, + SystemOperationKind, +) +from langchain_core.messages import AIMessage, HumanMessage +from langgraph.runtime import Runtime + +from deerflow.extensions.notify import ( + dispatch_system_model_observation, + notify_system_model_call, + observe_system_model_call, + reset_extension_notify_loop, + set_extension_notify_loop, + suspend_extension_system_observations, + task_store_for_system_call, +) +from deerflow.extensions.registry import ExtensionRegistry + + +class _Observer: + def __init__(self) -> None: + self.calls: list[tuple[SystemOperationKind, SystemModelRequest, SystemModelResult]] = [] + self.stores: list[ExtensionData] = [] + + async def on_system_model_call(self, app_store, task_store, kind, request, result): + self.calls.append((kind, request, result)) + self.stores.append(task_store) + + +def _extensions(*observers): + registry = ExtensionRegistry() + for index, observer in enumerate(observers): + with registry.attributed_to(f"ext{index}:install"): + registry.system_model_observer(observer) + return registry.build() + + +@pytest.mark.asyncio +async def test_system_model_notification_reports_success_failure_and_detached_store(): + observer = _Observer() + extensions = _extensions(observer) + request = SystemModelRequest(messages=("prompt",), model_name="system-model") + + await notify_system_model_call( + extensions, + None, + SystemOperationKind.TITLE, + request, + SystemModelResult(response="ok"), + ) + error = RuntimeError("provider down") + live_store = ExtensionData("task-1") + await notify_system_model_call( + extensions, + live_store, + SystemOperationKind.MEMORY, + request, + SystemModelResult(error=error), + ) + + assert [call[0] for call in observer.calls] == [ + SystemOperationKind.TITLE, + SystemOperationKind.MEMORY, + ] + assert observer.calls[0][2].response == "ok" + assert observer.calls[1][2].error is error + assert observer.stores[0].scope_id == "detached" + assert observer.stores[1] is live_store + + +@pytest.mark.asyncio +async def test_bad_observer_is_fail_open_and_does_not_hide_later_observers(): + class _Boom: + async def on_system_model_call(self, app_store, task_store, kind, request, result): + raise RuntimeError("observer exploded") + + survivor = _Observer() + await notify_system_model_call( + _extensions(_Boom(), survivor), + ExtensionData("task"), + SystemOperationKind.SUMMARIZATION, + SystemModelRequest(), + SystemModelResult(response="summary"), + ) + + assert [call[0] for call in survivor.calls] == [SystemOperationKind.SUMMARIZATION] + + +@pytest.mark.asyncio +async def test_observer_raising_cancellederror_is_fail_open_like_any_other_failure( + _notification_loop_state, +): + # An observer that implements its own timeout with cancellation can let a + # CancelledError escape. Fail-open is about the origin of the failure, not + # its base class: a contributor must never skip its successors or reach the + # host, and CancelledError does not derive from Exception. + class _Rogue: + async def on_system_model_call(self, app_store, task_store, kind, request, result): + raise asyncio.CancelledError() + + survivor = _Observer() + await notify_system_model_call( + _extensions(_Rogue(), survivor), + ExtensionData("task"), + SystemOperationKind.GOAL, + SystemModelRequest(), + SystemModelResult(response="ok"), + ) + + assert [call[0] for call in survivor.calls] == [SystemOperationKind.GOAL] + + +@pytest.mark.asyncio +async def test_genuine_host_cancellation_during_notification_still_propagates( + _notification_loop_state, +): + # The guard above must not swallow a real cancellation of the host task. + entered = asyncio.Event() + survivor = _Observer() + + class _Slow: + async def on_system_model_call(self, app_store, task_store, kind, request, result): + entered.set() + await asyncio.sleep(10) + + async def _body(): + await notify_system_model_call( + _extensions(_Slow(), survivor), + ExtensionData("task"), + SystemOperationKind.GOAL, + SystemModelRequest(), + SystemModelResult(response="ok"), + ) + + task = asyncio.create_task(_body()) + await entered.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + assert survivor.calls == [] + + +@pytest.mark.asyncio +async def test_observe_uses_explicit_snapshot_and_store_on_both_paths(): + observer = _Observer() + extensions = _extensions(observer) + store = ExtensionData("live-task") + + async def _success(): + return "answer" + + response = await observe_system_model_call( + extensions, + SystemOperationKind.GOAL, + messages=("prompt",), + model_name="goal-model", + invoke_config={"run_name": "goal"}, + invoke=_success, + task_store=store, + ) + assert response == "answer" + assert observer.stores == [store] + assert observer.calls[0][2].response == "answer" + assert observer.calls[0][2].duration_ms is not None + + async def _failure(): + raise ValueError("provider down") + + with pytest.raises(ValueError, match="provider down"): + await observe_system_model_call( + extensions, + SystemOperationKind.GOAL, + messages=(), + model_name=None, + invoke_config=None, + invoke=_failure, + task_store=store, + ) + assert isinstance(observer.calls[1][2].error, ValueError) + + +@pytest.mark.asyncio +async def test_observe_reports_cancellation_without_awaiting_inside_the_cancelled_task( + _notification_loop_state, +): + # Interrupt/rollback admission cancels the in-flight run task, so a system + # model call being cancelled is routine, not exotic. Awaiting observers here + # is unreliable (a repeated cancel interrupts that await too), so the + # cancellation terminal path is submitted to the notify loop instead. + observer = _Observer() + store = ExtensionData("live-task") + set_extension_notify_loop(asyncio.get_running_loop()) + entered = asyncio.Event() + + async def _never_returns(): + entered.set() + await asyncio.sleep(10) + + async def _body(): + await observe_system_model_call( + _extensions(observer), + SystemOperationKind.SUMMARIZATION, + messages="prompt", + model_name="sum-model", + invoke_config=None, + invoke=_never_returns, + task_store=store, + ) + + task = asyncio.create_task(_body()) + await entered.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + deadline = time.monotonic() + 2 + while not observer.calls and time.monotonic() < deadline: + await asyncio.sleep(0.01) + + assert [call[0] for call in observer.calls] == [SystemOperationKind.SUMMARIZATION] + assert isinstance(observer.calls[0][2].error, asyncio.CancelledError) + assert observer.calls[0][2].response is None + assert observer.calls[0][2].duration_ms is not None + assert observer.stores == [store] + + +@pytest.mark.asyncio +async def test_zero_observer_path_only_invokes_the_original_call(): + invoked: list[str] = [] + + async def _call(): + invoked.append("call") + return "ok" + + result = await observe_system_model_call( + ExtensionRegistry().build(), + SystemOperationKind.GOAL, + messages=(), + model_name=None, + invoke_config=None, + invoke=_call, + ) + + assert result == "ok" + assert invoked == ["call"] + + +def test_task_store_fallback_reads_only_the_host_runtime_key(): + store = ExtensionData("task-1") + assert task_store_for_system_call({"context": {EXTENSION_TASK_STORE_KEY: store}}) is store + for value in (None, {}, {"context": None}, {"context": {}}, "bad"): + assert task_store_for_system_call(value) is None + + +class _GoalModel: + def __init__(self, error: Exception | None = None) -> None: + self.error = error + + async def ainvoke(self, messages, config=None): + if self.error is not None: + raise self.error + return AIMessage(content=('{"satisfied": true, "blocker": "none", "reason": "done", "evidence_summary": "shipped"}')) + + +@pytest.mark.asyncio +async def test_goal_evaluator_observes_success_and_failure_with_explicit_snapshot(): + from deerflow.runtime.goal import evaluate_goal_completion + + observer = _Observer() + extensions = _extensions(observer) + store = ExtensionData("goal-task") + evidence = [ + HumanMessage(content="Ship it"), + AIMessage(content="It is shipped"), + ] + + await evaluate_goal_completion( + {"objective": "ship it"}, + evidence, + model=_GoalModel(), + model_name="goal-model", + task_store=store, + extensions=extensions, + ) + with pytest.raises(ValueError, match="provider down"): + await evaluate_goal_completion( + {"objective": "ship it"}, + evidence, + model=_GoalModel(ValueError("provider down")), + model_name="goal-model", + task_store=store, + extensions=extensions, + ) + + assert [call[2].error is None for call in observer.calls] == [True, False] + assert observer.stores == [store, store] + + +@pytest.mark.asyncio +async def test_title_middleware_uses_build_bound_snapshot_and_live_task_store(monkeypatch): + from deerflow.agents.middlewares import title_middleware as title_module + from deerflow.config.title_config import TitleConfig + + observer = _Observer() + extensions = _extensions(observer) + store = ExtensionData("title-task") + + class _TitleModel: + async def ainvoke(self, prompt, config=None): + return AIMessage(content="A Good Title") + + monkeypatch.setattr( + title_module, + "create_chat_model", + lambda **kwargs: _TitleModel(), + ) + middleware = title_module.TitleMiddleware( + title_config=TitleConfig(model_name="title-model"), + extensions=extensions, + ) + state = { + "messages": [ + HumanMessage(content="Question"), + AIMessage(content="Answer"), + ] + } + + result = await middleware.aafter_model( + state, + Runtime(context={EXTENSION_TASK_STORE_KEY: store}), + ) + + assert result == {"title": "A Good Title"} + assert [call[0] for call in observer.calls] == [SystemOperationKind.TITLE] + assert observer.stores == [store] + # The title call sends one prompt string, so observers must see it whole + # rather than as a character-by-character sequence. + (prompt,) = observer.calls[0][1].messages + assert isinstance(prompt, str) + assert "Question" in prompt + + +@pytest.mark.asyncio +async def test_async_summarization_observes_each_provider_attempt_and_live_store(): + from deerflow.agents.middlewares.summarization_middleware import ( + DeerFlowSummarizationMiddleware, + ) + + observer = _Observer() + extensions = _extensions(observer) + store = ExtensionData("summary-task") + + class _Failing: + async def ainvoke(self, prompt, config=None): + raise RuntimeError("first provider down") + + class _Working: + async def ainvoke(self, prompt, config=None): + return SimpleNamespace(text=" compact summary ") + + middleware = DeerFlowSummarizationMiddleware.__new__(DeerFlowSummarizationMiddleware) + middleware._extensions = extensions + middleware._prepare_summary_prompt = lambda messages, previous_summary=None: "prompt" + middleware._generation_candidate_names = lambda: ["first", "second"] + models = {"first": _Failing(), "second": _Working()} + middleware._model_for = lambda name: models[name] + + result = await middleware._asummarize_with(["message"], task_store=store) + + assert result == "compact summary" + assert [call[2].error is None for call in observer.calls] == [False, True] + assert [call[1].model_name for call in observer.calls] == ["first", "second"] + assert [call[1].messages for call in observer.calls] == [("prompt",), ("prompt",)] + assert observer.stores == [store, store] + + +@pytest.mark.asyncio +async def test_summarization_public_hook_propagates_the_live_task_store(): + from deerflow.agents.middlewares.summarization_middleware import ( + DeerFlowSummarizationMiddleware, + ) + + observer = _Observer() + extensions = _extensions(observer) + store = ExtensionData("summary-live-task") + model = MagicMock() + model.with_config.return_value = model + model.ainvoke = AsyncMock(return_value=SimpleNamespace(text="compressed")) + middleware = DeerFlowSummarizationMiddleware( + model=model, + trigger=("messages", 4), + keep=("messages", 2), + token_counter=len, + extensions=extensions, + ) + state = { + "messages": [ + HumanMessage(content="user-1"), + AIMessage(content="assistant-1"), + HumanMessage(content="user-2"), + AIMessage(content="assistant-2"), + ] + } + + result = await middleware.abefore_model( + state, + Runtime(context={EXTENSION_TASK_STORE_KEY: store}), + ) + + assert result is not None + assert observer.stores == [store] + + +@pytest.mark.asyncio +async def test_memory_callback_dispatches_the_captured_snapshot_and_live_store(): + from deerflow.agents.memory.manager import LangfuseMemoryCallbacks + from deerflow.extensions import reset_loaded_extensions, set_loaded_extensions + + first = _Observer() + second = _Observer() + captured = _extensions(first) + replacement = _extensions(second) + store = ExtensionData("memory-task") + loop = asyncio.get_running_loop() + set_extension_notify_loop(loop) + try: + callback = LangfuseMemoryCallbacks(extensions=captured) + set_loaded_extensions(replacement) + callback.on_memory_llm_result( + {"context": {EXTENSION_TASK_STORE_KEY: store}}, + prompt=("memory prompt",), + response="memory response", + error=None, + duration_ms=12.5, + model_name="memory-model", + ) + for _ in range(20): + if first.calls: + break + await asyncio.sleep(0) + finally: + reset_extension_notify_loop() + reset_loaded_extensions() + + assert [call[0] for call in first.calls] == [SystemOperationKind.MEMORY] + assert first.calls[0][1].messages == ("memory prompt",) + assert first.calls[0][2].response == "memory response" + assert first.calls[0][2].duration_ms == 12.5 + assert first.stores == [store] + assert second.calls == [] + + +def test_memory_callback_does_not_swallow_interpreter_shutdown(monkeypatch): + # Fail-open covers the bridge's own failures, not a process teardown + # signal — the same boundary the DeerMem-side call site pins in + # `test_memory_updater.py`. + from deerflow.agents.memory.manager import LangfuseMemoryCallbacks + from deerflow.extensions import notify as notify_module + + def _teardown(coro, what): + coro.close() + raise SystemExit("interpreter is going down") + + monkeypatch.setattr(notify_module, "dispatch_system_model_observation", _teardown) + callback = LangfuseMemoryCallbacks(extensions=_extensions(_Observer())) + + with pytest.raises(SystemExit): + callback.on_memory_llm_result( + {}, + prompt=("memory prompt",), + response=None, + error=None, + duration_ms=1.0, + model_name="memory-model", + ) + + +def test_memory_callback_contains_bridge_failures(monkeypatch): + from deerflow.agents.memory.manager import LangfuseMemoryCallbacks + from deerflow.extensions import notify as notify_module + + def _broken(coro, what): + coro.close() + raise RuntimeError("loop is gone") + + monkeypatch.setattr(notify_module, "dispatch_system_model_observation", _broken) + callback = LangfuseMemoryCallbacks(extensions=_extensions(_Observer())) + + callback.on_memory_llm_result( + {}, + prompt=("memory prompt",), + response="memory response", + error=None, + duration_ms=1.0, + model_name="memory-model", + ) + + +@pytest.mark.asyncio +async def test_system_model_failure_logs_identify_the_task_scope(caplog): + class _Broken: + async def on_system_model_call(self, app_store, task_store, kind, request, result): + raise RuntimeError("observer broke") + + with caplog.at_level(logging.WARNING, logger="deerflow.extensions.notify"): + await notify_system_model_call( + _extensions(_Broken()), + ExtensionData("summary-live-task"), + SystemOperationKind.GOAL, + SystemModelRequest(messages=("prompt",), model_name="system-model"), + SystemModelResult(response="ok"), + ) + + messages = [record.getMessage() for record in caplog.records if record.name == "deerflow.extensions.notify"] + assert any("summary-live-task" in message and "goal" in message for message in messages) + + +@pytest.fixture +def _notification_loop_state(): + reset_extension_notify_loop() + yield + reset_extension_notify_loop() + + +class _RunningLoop: + def __init__(self) -> None: + self.loop = asyncio.new_event_loop() + self._ready = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + assert self._ready.wait(2) + + def _run(self) -> None: + asyncio.set_event_loop(self.loop) + self.loop.call_soon(self._ready.set) + self.loop.run_forever() + + def stop(self) -> None: + if self.loop.is_running(): + self.loop.call_soon_threadsafe(self.loop.stop) + self._thread.join(2) + + def close(self) -> None: + self.stop() + if not self.loop.is_closed(): + self.loop.close() + + +def _wait_for_calls(observer: _Observer, expected: int = 1) -> None: + deadline = time.monotonic() + 2 + while len(observer.calls) < expected and time.monotonic() < deadline: + threading.Event().wait(0.01) + + +def test_detached_observation_dispatches_to_the_registered_loop( + _notification_loop_state, +): + observed_loops: list[asyncio.AbstractEventLoop] = [] + + class _LoopObserver(_Observer): + async def on_system_model_call(self, app_store, task_store, kind, request, result): + observed_loops.append(asyncio.get_running_loop()) + await super().on_system_model_call(app_store, task_store, kind, request, result) + + observer = _LoopObserver() + host = _RunningLoop() + set_extension_notify_loop(host.loop) + try: + submitted = dispatch_system_model_observation( + notify_system_model_call( + _extensions(observer), + None, + SystemOperationKind.MEMORY, + SystemModelRequest(), + SystemModelResult(response="ok"), + ), + "memory", + ) + _wait_for_calls(observer) + finally: + host.close() + + assert submitted is True + assert observed_loops == [host.loop] + + +def test_awaited_observation_from_an_isolated_loop_uses_registered_loop( + _notification_loop_state, +): + observed_loops: list[asyncio.AbstractEventLoop] = [] + + class _LoopObserver: + async def on_system_model_call(self, app_store, task_store, kind, request, result): + observed_loops.append(asyncio.get_running_loop()) + + host = _RunningLoop() + set_extension_notify_loop(host.loop) + try: + asyncio.run( + notify_system_model_call( + _extensions(_LoopObserver()), + None, + SystemOperationKind.SUMMARIZATION, + SystemModelRequest(), + SystemModelResult(response="ok"), + ) + ) + finally: + host.close() + + assert observed_loops == [host.loop] + + +def test_detached_observation_drops_when_loop_is_missing_stopped_or_suspended( + _notification_loop_state, +): + observer = _Observer() + extensions = _extensions(observer) + + assert ( + dispatch_system_model_observation( + notify_system_model_call( + extensions, + None, + SystemOperationKind.MEMORY, + SystemModelRequest(), + SystemModelResult(response="missing"), + ), + "missing-loop", + ) + is False + ) + + host = _RunningLoop() + set_extension_notify_loop(host.loop) + host.stop() + assert not host.loop.is_closed() + assert ( + dispatch_system_model_observation( + notify_system_model_call( + extensions, + None, + SystemOperationKind.MEMORY, + SystemModelRequest(), + SystemModelResult(response="stopped"), + ), + "stopped-loop", + ) + is False + ) + host.loop.close() + + active = _RunningLoop() + set_extension_notify_loop(active.loop) + suspend_extension_system_observations() + try: + assert ( + dispatch_system_model_observation( + notify_system_model_call( + extensions, + None, + SystemOperationKind.MEMORY, + SystemModelRequest(), + SystemModelResult(response="suspended"), + ), + "suspended-loop", + ) + is False + ) + finally: + active.close() + + assert observer.calls == [] + + +def test_detached_observation_ignores_the_callers_other_running_loop( + _notification_loop_state, +): + observed_loops: list[asyncio.AbstractEventLoop] = [] + + class _LoopObserver: + async def on_system_model_call(self, app_store, task_store, kind, request, result): + observed_loops.append(asyncio.get_running_loop()) + + registered = _RunningLoop() + other = _RunningLoop() + set_extension_notify_loop(registered.loop) + + async def _dispatch_from_other() -> None: + assert dispatch_system_model_observation( + notify_system_model_call( + _extensions(_LoopObserver()), + None, + SystemOperationKind.MEMORY, + SystemModelRequest(), + SystemModelResult(response="ok"), + ), + "memory-from-other-loop", + ) + + try: + asyncio.run_coroutine_threadsafe(_dispatch_from_other(), other.loop).result(2) + deadline = time.monotonic() + 2 + while not observed_loops and time.monotonic() < deadline: + threading.Event().wait(0.01) + finally: + other.close() + registered.close() + + assert observed_loops == [registered.loop] diff --git a/backend/tests/test_extension_task_lifecycle.py b/backend/tests/test_extension_task_lifecycle.py new file mode 100644 index 000000000..72aa29006 --- /dev/null +++ b/backend/tests/test_extension_task_lifecycle.py @@ -0,0 +1,438 @@ +"""Task lifecycle extension notifications and outcome classification.""" + +from __future__ import annotations + +import asyncio +import logging +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from deerflow_extension_api import ( + EXTENSION_TASK_STORE_KEY, + ExtensionData, + TaskInfo, + TaskOutcome, +) +from langgraph.checkpoint.memory import InMemorySaver + +from deerflow.extensions.notify import ( + lead_task_id, + lead_task_outcome, + notify_task_start, + notify_task_stop, + subagent_task_outcome, +) +from deerflow.extensions.registry import ExtensionRegistry +from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.schemas import RunStatus +from deerflow.runtime.runs.worker import RunContext, run_agent + + +class _Recorder: + def __init__(self) -> None: + self.events: list[tuple[str, str, str]] = [] + + async def on_task_start(self, app_store, task_store, info): + self.events.append(("start", info.task_id, info.kind)) + + async def on_task_stop(self, app_store, task_store, info, outcome): + self.events.append(("stop", info.task_id, outcome.value)) + + +def _extensions(*contributors): + registry = ExtensionRegistry() + for index, contributor in enumerate(contributors): + with registry.attributed_to(f"ext{index}:install"): + registry.task_lifecycle(contributor) + return registry.build() + + +def _info(task_id: str = "task-1", kind: str = "lead") -> TaskInfo: + return TaskInfo(task_id=task_id, run_id="run-1", thread_id="thread-1", kind=kind) + + +def test_task_identity_and_outcome_classification_are_explicit(): + assert lead_task_id("run-abc") == "run-abc" + assert lead_task_outcome(aborted=True, succeeded=True) is TaskOutcome.ABORTED + assert lead_task_outcome(aborted=False, succeeded=True) is TaskOutcome.COMPLETED + assert lead_task_outcome(aborted=False, succeeded=False) is TaskOutcome.FAILED + assert subagent_task_outcome(cancelled=True, succeeded=True) is TaskOutcome.ABORTED + assert subagent_task_outcome(cancelled=False, succeeded=True) is TaskOutcome.COMPLETED + assert subagent_task_outcome(cancelled=False, succeeded=False) is TaskOutcome.FAILED + + +@pytest.mark.asyncio +async def test_start_and_stop_reach_contributors_in_order(): + first = _Recorder() + second = _Recorder() + extensions = _extensions(first, second) + store = ExtensionData("task-1") + + await notify_task_start(extensions, store, _info()) + await notify_task_stop(extensions, store, _info(), TaskOutcome.COMPLETED) + + assert first.events == [("start", "task-1", "lead"), ("stop", "task-1", "completed")] + assert second.events == first.events + + +@pytest.mark.asyncio +async def test_one_malformed_or_failing_contributor_does_not_stop_the_rest(): + class _WrongShape: + def on_task_start(self, app_store, task_store, info): + raise RuntimeError("sync boom") + + survivor = _Recorder() + await notify_task_start( + _extensions(_WrongShape(), survivor), + ExtensionData("task-1"), + _info(), + ) + + assert survivor.events == [("start", "task-1", "lead")] + + +@pytest.mark.asyncio +async def test_notification_timeout_is_one_shared_budget(): + reached: list[str] = [] + + class _Hang: + async def on_task_stop(self, app_store, task_store, info, outcome): + reached.append("hang") + await asyncio.sleep(10) + + class _Starved: + async def on_task_stop(self, app_store, task_store, info, outcome): + reached.append("starved") + + loop = asyncio.get_running_loop() + started = loop.time() + await notify_task_stop( + _extensions(_Hang(), _Starved()), + ExtensionData("task-1"), + _info(), + TaskOutcome.COMPLETED, + timeout=0.02, + ) + + assert reached == ["hang"] + assert loop.time() - started < 1 + + +@pytest.mark.asyncio +async def test_budget_exhaustion_mid_hook_logs_a_warning_not_a_traceback(caplog): + # Spending the shared budget mid-hook is the same expected operational + # condition as the pre-hook skip, not a hook failure. + class _Hang: + async def on_task_stop(self, app_store, task_store, info, outcome): + await asyncio.sleep(10) + + with caplog.at_level(logging.WARNING, logger="deerflow.extensions.notify"): + await notify_task_stop( + _extensions(_Hang()), + ExtensionData("task-1"), + _info(), + TaskOutcome.COMPLETED, + timeout=0.02, + ) + + records = [record for record in caplog.records if record.name == "deerflow.extensions.notify"] + assert [record.levelno for record in records] == [logging.WARNING] + assert "timed out" in records[0].getMessage() + assert "task-1" in records[0].getMessage() + assert records[0].exc_info is None + + +@pytest.mark.asyncio +async def test_contributor_timeout_error_before_budget_is_a_hook_failure(caplog): + # A TimeoutError the contributor raises on its own is not budget + # exhaustion; it stays classified as a hook failure. + class _TimedOut: + async def on_task_stop(self, app_store, task_store, info, outcome): + raise TimeoutError("the contributor's own downstream call timed out") + + with caplog.at_level(logging.WARNING, logger="deerflow.extensions.notify"): + await notify_task_stop( + _extensions(_TimedOut()), + ExtensionData("task-1"), + _info(), + TaskOutcome.COMPLETED, + timeout=30, + ) + # And with no notification budget at all. + await notify_task_stop( + _extensions(_TimedOut()), + ExtensionData("task-1"), + _info(), + TaskOutcome.COMPLETED, + ) + + records = [record for record in caplog.records if record.name == "deerflow.extensions.notify"] + assert [record.levelno for record in records] == [logging.ERROR, logging.ERROR] + assert all("failed" in record.getMessage() for record in records) + + +class _RunRecorder(_Recorder): + def __init__(self) -> None: + super().__init__() + self.start_infos: list[TaskInfo] = [] + self.start_stores: list[ExtensionData] = [] + + async def on_task_start(self, app_store, task_store, info): + self.start_infos.append(info) + self.start_stores.append(task_store) + await super().on_task_start(app_store, task_store, info) + + +class _OkAgent: + def __init__(self) -> None: + self.runtime_context = None + + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + self.runtime_context = (config or {}).get("context") + yield {"messages": []} + + +class _BoomAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + raise RuntimeError("agent exploded") + yield # pragma: no cover + + +def _bridge(): + return SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + + +@pytest.mark.asyncio +async def test_run_agent_uses_the_run_bound_snapshot_for_lifecycle_and_task_store(): + recorder = _RunRecorder() + extensions = _extensions(recorder) + manager = RunManager() + record = await manager.create("thread-ext", assistant_id="custom-agent") + agent = _OkAgent() + + await run_agent( + _bridge(), + manager, + record, + ctx=RunContext(checkpointer=InMemorySaver(), extensions=extensions), + agent_factory=lambda *, config: agent, + graph_input={}, + config={}, + ) + + assert record.status is RunStatus.success + assert recorder.events == [ + ("start", record.run_id, "lead"), + ("stop", record.run_id, "completed"), + ] + assert recorder.start_infos[0].agent_name == "custom-agent" + assert agent.runtime_context[EXTENSION_TASK_STORE_KEY] is recorder.start_stores[0] + + +@pytest.mark.asyncio +async def test_run_agent_reports_failed_and_skips_runs_that_never_started(): + recorder = _RunRecorder() + extensions = _extensions(recorder) + manager = RunManager() + + failed = await manager.create("thread-failed") + await run_agent( + _bridge(), + manager, + failed, + ctx=RunContext(checkpointer=InMemorySaver(), extensions=extensions), + agent_factory=lambda *, config: _BoomAgent(), + graph_input={}, + config={}, + ) + assert failed.status is RunStatus.error + assert recorder.events[-1] == ("stop", failed.run_id, "failed") + + skipped = await manager.create("thread-skipped") + await manager.cancel(skipped.run_id) + before = list(recorder.events) + await run_agent( + _bridge(), + manager, + skipped, + ctx=RunContext(checkpointer=InMemorySaver(), extensions=extensions), + agent_factory=lambda *, config: _OkAgent(), + graph_input={}, + config={}, + ) + assert recorder.events == before + + +@pytest.mark.asyncio +async def test_lead_stop_runs_after_completion_hook_and_before_stream_end(): + events: list[str] = [] + manager = RunManager() + record = await manager.create("thread-order") + bridge = _bridge() + bridge.publish_end.side_effect = lambda run_id: events.append("stream-end") + + class _OrderingRecorder(_RunRecorder): + async def on_task_stop(self, app_store, task_store, info, outcome): + assert record.finalizing is True + assert bridge.publish_end.await_count == 0 + events.append("task-stop") + await super().on_task_stop( + app_store, + task_store, + info, + outcome, + ) + + async def _on_completed(_record): + events.append("run-completed") + + class _CancelledAgent: + async def astream( + self, + graph_input, + config=None, + stream_mode=None, + subgraphs=False, + ): + raise asyncio.CancelledError() + yield # pragma: no cover + + await run_agent( + bridge, + manager, + record, + ctx=RunContext( + checkpointer=InMemorySaver(), + extensions=_extensions(_OrderingRecorder()), + on_run_completed=_on_completed, + ), + agent_factory=lambda *, config: _CancelledAgent(), + graph_input={}, + config={}, + ) + + assert events == ["run-completed", "task-stop", "stream-end"] + assert record.finalizing is False + + +@pytest.mark.asyncio +async def test_lead_stop_interrupt_is_deferred_until_final_cleanup(): + class _StopInterrupt(BaseException): + pass + + class _InterruptingRecorder(_RunRecorder): + async def on_task_stop(self, app_store, task_store, info, outcome): + await super().on_task_stop( + app_store, + task_store, + info, + outcome, + ) + raise _StopInterrupt("shutdown") + + manager = RunManager() + record = await manager.create("thread-interrupt") + bridge = _bridge() + + with pytest.raises(_StopInterrupt, match="shutdown"): + await run_agent( + bridge, + manager, + record, + ctx=RunContext( + checkpointer=InMemorySaver(), + extensions=_extensions(_InterruptingRecorder()), + ), + agent_factory=lambda *, config: _OkAgent(), + graph_input={}, + config={}, + ) + + assert record.finalizing is False + bridge.publish_end.assert_awaited_once_with(record.run_id) + + +@pytest.mark.asyncio +async def test_contributor_raising_cancellederror_cannot_interrupt_run_cleanup(): + # Fail-open is decided by origin, not base class: a contributor that lets a + # CancelledError escape must not skip its successors, and must not reach the + # worker's deferred-interrupt path, which would end an otherwise successful + # run as cancelled. + class _Rogue: + async def on_task_stop(self, app_store, task_store, info, outcome): + raise asyncio.CancelledError() + + survivor = _RunRecorder() + manager = RunManager() + record = await manager.create("thread-rogue-stop") + bridge = _bridge() + + await run_agent( + bridge, + manager, + record, + ctx=RunContext( + checkpointer=InMemorySaver(), + extensions=_extensions(_Rogue(), survivor), + ), + agent_factory=lambda *, config: _OkAgent(), + graph_input={}, + config={}, + ) + + assert record.status is RunStatus.success + assert survivor.events[-1] == ("stop", record.run_id, "completed") + assert record.finalizing is False + bridge.publish_end.assert_awaited_once_with(record.run_id) + + +@pytest.mark.asyncio +async def test_lead_stop_cancellation_is_deferred_rather_than_swallowed(): + # CancelledError derives from BaseException, so the non-fatal `except + # Exception` guard around the stop notification must not absorb it. The + # stimulus has to be a genuine cancellation of the run task — a contributor + # raising CancelledError is contained as an extension failure instead. + entered = asyncio.Event() + + class _SlowRecorder(_RunRecorder): + async def on_task_stop(self, app_store, task_store, info, outcome): + await super().on_task_stop( + app_store, + task_store, + info, + outcome, + ) + entered.set() + await asyncio.sleep(10) + + manager = RunManager() + record = await manager.create("thread-cancel-stop") + bridge = _bridge() + + task = asyncio.create_task( + run_agent( + bridge, + manager, + record, + ctx=RunContext( + checkpointer=InMemorySaver(), + extensions=_extensions(_SlowRecorder()), + ), + agent_factory=lambda *, config: _OkAgent(), + graph_input={}, + config={}, + ) + ) + await entered.wait() + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert record.finalizing is False + bridge.publish_end.assert_awaited_once_with(record.run_id) diff --git a/backend/tests/test_gateway_lifespan_shutdown.py b/backend/tests/test_gateway_lifespan_shutdown.py index 049d84d14..9800c614b 100644 --- a/backend/tests/test_gateway_lifespan_shutdown.py +++ b/backend/tests/test_gateway_lifespan_shutdown.py @@ -119,7 +119,12 @@ def test_lifespan_sweeps_upload_staging_files_on_startup(): stop_channel_service.assert_awaited_once() -async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool | Exception) -> MagicMock: +async def _run_lifespan_with_memory_flush( + *, + enabled: bool, + flush_return: bool | Exception, + shutdown_events: list[str] | None = None, +) -> MagicMock: """Drive lifespan with a spied memory manager.shutdown_flush. Returns the manager mock so the caller can assert the shutdown flush was @@ -151,9 +156,20 @@ async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool | manager = MagicMock() if isinstance(flush_return, Exception): manager.shutdown_flush.side_effect = flush_return + elif shutdown_events is not None: + + def record_memory_flush(_timeout: float) -> bool: + shutdown_events.append("memory_flush_started") + return flush_return + + manager.shutdown_flush.side_effect = record_memory_flush else: manager.shutdown_flush.return_value = flush_return + suspend_system_observations = MagicMock() + if shutdown_events is not None: + suspend_system_observations.side_effect = lambda: shutdown_events.append("system_observations_suspended") + with ( patch("app.gateway.app.get_app_config", return_value=startup_config), patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)), @@ -163,6 +179,7 @@ async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool | patch("app.channels.service.start_channel_service", side_effect=fake_start), patch("app.channels.service.stop_channel_service", stop_channel_service), patch("deerflow.agents.memory.get_memory_manager", return_value=manager), + patch("deerflow.extensions.notify.suspend_extension_system_observations", suspend_system_observations), ): async with lifespan(app): pass @@ -180,6 +197,21 @@ def test_lifespan_drains_memory_on_shutdown_with_configured_timeout(caplog) -> N assert any(r.levelno == logging.INFO and "flush completed" in r.message for r in caplog.records) +def test_lifespan_suspends_system_observations_before_memory_flush() -> None: + """Shutdown-flushed memory calls cannot enqueue observations onto a dying loop.""" + shutdown_events: list[str] = [] + + asyncio.run( + _run_lifespan_with_memory_flush( + enabled=True, + flush_return=True, + shutdown_events=shutdown_events, + ) + ) + + assert shutdown_events == ["system_observations_suspended", "memory_flush_started"] + + def test_lifespan_warns_when_memory_flush_does_not_finish(caplog) -> None: """A False return (timeout/failure) is the path operators actually see when K8s SIGKILLs the drain; the host must log a WARNING (not 'completed'), so diff --git a/backend/tests/test_gateway_run_drain_shutdown.py b/backend/tests/test_gateway_run_drain_shutdown.py index 528997069..7e7bdf704 100644 --- a/backend/tests/test_gateway_run_drain_shutdown.py +++ b/backend/tests/test_gateway_run_drain_shutdown.py @@ -185,6 +185,13 @@ async def test_langgraph_runtime_drains_runs_before_closing_checkpointer(monkeyp async def spy_shutdown(self, *, timeout): # noqa: ANN001 events.append("runs_drained") + def spy_set_extension_notify_loop(loop): # noqa: ANN001 + assert loop is asyncio.get_running_loop() + events.append("extension_loop_set") + + def spy_reset_extension_notify_loop(): + events.append("extension_loop_reset") + monkeypatch.setattr("deerflow.runtime.checkpointer.async_provider.make_checkpointer", probe_checkpointer) monkeypatch.setattr("deerflow.runtime.make_stream_bridge", fake_stream_bridge) monkeypatch.setattr("deerflow.runtime.make_store", fake_store) @@ -194,6 +201,8 @@ async def test_langgraph_runtime_drains_runs_before_closing_checkpointer(monkeyp monkeypatch.setattr("deerflow.runtime.events.store.make_run_event_store", lambda _cfg: object()) monkeypatch.setattr("deerflow.persistence.thread_meta.make_thread_store", lambda _sf, _store: object()) monkeypatch.setattr(RunManager, "shutdown", spy_shutdown, raising=False) + monkeypatch.setattr("deerflow.extensions.notify.set_extension_notify_loop", spy_set_extension_notify_loop) + monkeypatch.setattr("deerflow.extensions.notify.reset_extension_notify_loop", spy_reset_extension_notify_loop) app = FastAPI() startup_config = SimpleNamespace(database=SimpleNamespace(backend="memory", checkpoint_channel_mode="full", checkpoint_delta=SimpleNamespace(snapshot_frequency=10)), run_events=None) @@ -204,6 +213,50 @@ async def test_langgraph_runtime_drains_runs_before_closing_checkpointer(monkeyp assert "runs_drained" in events, "langgraph_runtime never drained in-flight runs on shutdown" assert "checkpointer_closed" in events assert events.index("runs_drained") < events.index("checkpointer_closed"), f"runs must be drained before the checkpointer pool is closed; got order {events}" + assert events[0] == "extension_loop_set" + assert events.index("checkpointer_closed") < events.index("extension_loop_reset"), f"extension loop reset must be the final runtime teardown; got order {events}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("startup_error", [RuntimeError("startup failed"), asyncio.CancelledError()]) +async def test_langgraph_runtime_resets_extension_loop_when_startup_exits_early(monkeypatch, startup_error): + """A partial startup must not leave a stale process-wide loop binding.""" + from fastapi import FastAPI + + from app.gateway.deps import langgraph_runtime + + events: list[str] = [] + + @asynccontextmanager + async def failing_stream_bridge(_config): + raise startup_error + yield # pragma: no cover - makes this an async context manager + + def spy_set_extension_notify_loop(loop): # noqa: ANN001 + assert loop is asyncio.get_running_loop() + events.append("extension_loop_set") + + def spy_reset_extension_notify_loop(): + events.append("extension_loop_reset") + + monkeypatch.setattr("deerflow.runtime.make_stream_bridge", failing_stream_bridge) + monkeypatch.setattr("deerflow.extensions.notify.set_extension_notify_loop", spy_set_extension_notify_loop) + monkeypatch.setattr("deerflow.extensions.notify.reset_extension_notify_loop", spy_reset_extension_notify_loop) + + app = FastAPI() + startup_config = SimpleNamespace( + database=SimpleNamespace( + backend="memory", + checkpoint_channel_mode="full", + checkpoint_delta=SimpleNamespace(snapshot_frequency=10), + ), + ) + + with pytest.raises(type(startup_error)): + async with langgraph_runtime(app, startup_config): + pass + + assert events == ["extension_loop_set", "extension_loop_reset"] @pytest.mark.asyncio diff --git a/backend/tests/test_goal_worker.py b/backend/tests/test_goal_worker.py index 3630f7dc5..eed0938fd 100644 --- a/backend/tests/test_goal_worker.py +++ b/backend/tests/test_goal_worker.py @@ -2,10 +2,12 @@ import asyncio import copy import pytest +from deerflow_extension_api import ExtensionData from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.base import empty_checkpoint, uuid6 from langgraph.checkpoint.memory import InMemorySaver +from deerflow.extensions.registry import ExtensionRegistry from deerflow.runtime.checkpoint_state import CheckpointStateAccessor, build_state_mutation_graph from deerflow.runtime.goal import GoalEvaluation, attach_goal_evaluation, build_goal_state, latest_visible_assistant_signature, read_thread_goal, write_thread_goal from deerflow.runtime.runs import worker @@ -150,10 +152,14 @@ async def test_goal_worker_returns_hidden_continuation_when_goal_is_unmet(monkey thread_id = "goal-thread" await _seed_goal_thread(checkpointer, thread_id=thread_id, goal_text="Finish all tests") bridge = _CollectingBridge() + task_store = ExtensionData("run-1") + extensions = ExtensionRegistry().build() - async def fake_evaluate_goal_completion(goal, messages, **_kwargs): + async def fake_evaluate_goal_completion(goal, messages, **kwargs): assert goal["objective"] == "Finish all tests" assert [message.content for message in messages][-1] == "I made a start, but I am not done." + assert kwargs["task_store"] is task_store + assert kwargs["extensions"] is extensions return GoalEvaluation( satisfied=False, blocker="goal_not_met_yet", @@ -171,6 +177,8 @@ async def test_goal_worker_returns_hidden_continuation_when_goal_is_unmet(monkey run_id="run-1", model_name="test-model", app_config=None, + task_store=task_store, + extensions=extensions, ) assert continuation is not None diff --git a/backend/tests/test_lead_agent_model_resolution.py b/backend/tests/test_lead_agent_model_resolution.py index 804dca9ad..609a4695f 100644 --- a/backend/tests/test_lead_agent_model_resolution.py +++ b/backend/tests/test_lead_agent_model_resolution.py @@ -643,7 +643,7 @@ def test_build_middlewares_passes_explicit_app_config_to_shared_factory(monkeypa monkeypatch.setattr( lead_agent_module, "TitleMiddleware", - lambda *, app_config: captured.setdefault("title_app_config", app_config) or "title-middleware", + lambda *, app_config, extensions: captured.setdefault("title_app_config", app_config) or "title-middleware", ) monkeypatch.setattr( lead_agent_module, @@ -688,7 +688,7 @@ def test_build_middlewares_passes_run_model_name_to_summarization(monkeypatch): lambda **kwargs: captured.update(kwargs) or None, ) monkeypatch.setattr(lead_agent_module, "_create_todo_list_middleware", lambda is_plan_mode: None) - monkeypatch.setattr(lead_agent_module, "TitleMiddleware", lambda *, app_config: "title-middleware") + monkeypatch.setattr(lead_agent_module, "TitleMiddleware", lambda *, app_config, extensions: "title-middleware") monkeypatch.setattr(lead_agent_module, "MemoryMiddleware", lambda agent_name=None, *, memory_config: "memory-middleware") lead_agent_module.build_middlewares( diff --git a/backend/tests/test_memory_manager_interface.py b/backend/tests/test_memory_manager_interface.py index f70cbd141..2ce60e82c 100644 --- a/backend/tests/test_memory_manager_interface.py +++ b/backend/tests/test_memory_manager_interface.py @@ -199,6 +199,14 @@ def test_callbacks_field_optional_and_noop_default(): noop = MemoryCallbacks() # no-op: mutates nothing, raises nothing noop.on_memory_llm_call({}, thread_id="t", user_id="u", trace_id="tr", model_name="m") + noop.on_memory_llm_result( + {}, + prompt="prompt", + response="response", + error=None, + duration_ms=1.0, + model_name="m", + ) manager = _MinimalBackend(backend_config={}, callbacks=noop) assert manager.callbacks is noop diff --git a/backend/tests/test_memory_updater.py b/backend/tests/test_memory_updater.py index 7746fd7a6..c5a99ca70 100644 --- a/backend/tests/test_memory_updater.py +++ b/backend/tests/test_memory_updater.py @@ -3,6 +3,8 @@ import copy import threading from unittest.mock import AsyncMock, MagicMock, patch +import pytest + from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig from deerflow.agents.memory.backends.deermem.deermem.core.prompt import format_conversation_for_update from deerflow.agents.memory.backends.deermem.deermem.core.storage import ( @@ -927,6 +929,81 @@ class TestUpdateMemoryStructuredResponse: assert result is True model.invoke.assert_called_once() + def test_result_callback_observes_successful_provider_call(self): + calls: list[dict[str, object]] = [] + + class _Callbacks: + def on_memory_llm_call(self, invoke_config, **kwargs): + return None + + def on_memory_llm_result(self, invoke_config, **kwargs): + calls.append({"invoke_config": invoke_config, **kwargs}) + + valid_json = '{"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}' + model = self._make_mock_model(valid_json) + updater = _make_updater(llm=model, callbacks=_Callbacks()) + msg = MagicMock(type="human", content="Hello") + ai_msg = MagicMock(type="ai", content="Hi", tool_calls=[]) + + assert updater.update_memory([msg, ai_msg]) is True + assert len(calls) == 1 + assert calls[0]["response"] is model.invoke.return_value + assert calls[0]["error"] is None + assert calls[0]["model_name"] is None + assert calls[0]["duration_ms"] >= 0 + + def test_result_callback_observes_provider_failure_and_is_fail_open(self): + calls: list[dict[str, object]] = [] + + class _Callbacks: + def on_memory_llm_call(self, invoke_config, **kwargs): + return None + + def on_memory_llm_result(self, invoke_config, **kwargs): + calls.append({"invoke_config": invoke_config, **kwargs}) + + provider_error = RuntimeError("provider down") + model = MagicMock() + model.invoke.side_effect = provider_error + updater = _make_updater(llm=model, callbacks=_Callbacks()) + msg = MagicMock(type="human", content="Hello") + ai_msg = MagicMock(type="ai", content="Hi", tool_calls=[]) + + assert updater.update_memory([msg, ai_msg]) is False + assert len(calls) == 1 + assert calls[0]["response"] is None + assert calls[0]["error"] is provider_error + + class _BrokenCallbacks(_Callbacks): + def on_memory_llm_result(self, invoke_config, **kwargs): + raise RuntimeError("observer callback broke") + + working_model = self._make_mock_model('{"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}') + fail_open_updater = _make_updater( + llm=working_model, + callbacks=_BrokenCallbacks(), + ) + assert fail_open_updater.update_memory([msg, ai_msg]) is True + + def test_result_callback_does_not_swallow_interpreter_shutdown(self): + # Fail-open covers the hook's own failures, not a process teardown + # signal: swallowing SystemExit here would let an observability path + # keep a shutting-down interpreter alive. + class _ExitingCallbacks: + def on_memory_llm_call(self, invoke_config, **kwargs): + return None + + def on_memory_llm_result(self, invoke_config, **kwargs): + raise SystemExit("interpreter is going down") + + model = self._make_mock_model('{"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}') + updater = _make_updater(llm=model, callbacks=_ExitingCallbacks()) + msg = MagicMock(type="human", content="Hello") + ai_msg = MagicMock(type="ai", content="Hi", tool_calls=[]) + + with pytest.raises(SystemExit): + updater.update_memory([msg, ai_msg]) + def test_list_content_response_parses(self): """LLM response as list-of-blocks should be extracted, not repr'd.""" valid_json = '{"user": {}, "history": {}, "newFacts": [], "factsToRemove": []}' diff --git a/backend/tests/test_tool_error_handling_middleware.py b/backend/tests/test_tool_error_handling_middleware.py index 03d493b6e..20f65f701 100644 --- a/backend/tests/test_tool_error_handling_middleware.py +++ b/backend/tests/test_tool_error_handling_middleware.py @@ -697,11 +697,19 @@ def test_subagent_runtime_middlewares_attach_durable_context_before_summarizatio sentinel = object() captured: dict[str, object] = {} - def fake_create_summarization_middleware(*, app_config=None, keep=None, skip_memory_flush=False, run_model_name=None): + def fake_create_summarization_middleware( + *, + app_config=None, + keep=None, + skip_memory_flush=False, + run_model_name=None, + extensions=None, + ): captured["app_config"] = app_config captured["keep"] = keep captured["skip_memory_flush"] = skip_memory_flush captured["run_model_name"] = run_model_name + captured["extensions"] = extensions return sentinel # summarization is enabled by default False; flip it on so the factory path @@ -724,6 +732,7 @@ def test_subagent_runtime_middlewares_attach_durable_context_before_summarizatio # so a distinct-model subagent summarizes with its model, not the parent's — the # subagent context/configurable never carries the child model. assert captured["run_model_name"] == "test-model" + assert captured["extensions"] is not None durable = [middleware for middleware in middlewares if isinstance(middleware, DurableContextMiddleware)] assert len(durable) == 1 # ``_skills_root`` is ``posixpath.normpath(container_path)``, so compare against diff --git a/backend/uv.lock b/backend/uv.lock index ff06b556c..f386e8612 100644 --- a/backend/uv.lock +++ b/backend/uv.lock @@ -907,7 +907,7 @@ dev = [ [[package]] name = "deerflow-extension-api" -version = "0.1.0" +version = "0.1.1" source = { editable = "packages/extension-api" } [[package]]