From 7389331e6593c7f39cdacd9b078cf946e4e0b22d Mon Sep 17 00:00:00 2001 From: Nan Gao Date: Tue, 11 Aug 2026 16:33:22 +0800 Subject: [PATCH] feat(extensions): observe task lifecycle and system model calls (#4684) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(extensions): observe task lifecycle and system model calls PR 1 (#4636) gave extensions a middleware chain, and a middleware only sees what passes through the agent graph. Two runtime surfaces stay invisible to it: when a lead run or a subagent begins and ends, and the DeerFlow-owned model calls made outside the graph. This slice adds both, with no new Gateway surface -- routers, services, and the reference extension stay in PR 3. Contract (deerflow-extension-api 0.1.1) --------------------------------------- Two contribution kinds join `middlewares` on the registry: `task_lifecycle` (`on_task_start` / `on_task_stop`, receiving a `TaskInfo` and a conservative `TaskOutcome` of completed / aborted / failed) and `system_model_observer` (`on_system_model_call`, receiving a `SystemOperationKind`, a `SystemModelRequest` snapshot, and a `SystemModelResult` carrying either the response or the provider exception plus a duration). `SystemModelRequest.messages` normalizes to a tuple at construction. Goal evaluation and memory extraction pass a message list while title generation and summarization pass one prompt string, and a bare `str` already satisfies `Sequence` -- without normalization an observer iterating `request.messages` would silently walk characters. Copying also makes the frozen snapshot immutable in fact rather than only by declaration, since observations may run after the call site returns and keeps mutating its own list. Registry marks and rollbacks become per-bucket and positional, so an `install()` that fails after registering two different kinds cannot leave one of them behind. `needs_task_store` now covers all three kinds: a deployment that registers only lifecycle hooks still gets a task store. Task lifecycle -------------- The lead worker notifies start after the run has started and stop after completion persistence and the completion hook, but before clearing the finalizing barrier and publishing the stream end -- holding the barrier across stop is what keeps a same-thread replacement run from overlapping this task's lifecycle. Cancellation raised out of the stop notification is deferred, not propagated in place, so a cancelled run still clears the barrier and emits its end frame. A subagent with a parent `run_id` wraps its execution in the same pair inside `finally`, reporting `parent_task_id` so a delegation tree is reconstructable; a subagent without a `run_id` (embedded client, standalone LangGraph Server) logs and skips rather than inventing a parent. Contributors run in registration order inside one shared 3s budget and every failure is logged and failed open. System model calls ------------------ Four kinds cover the model calls the middleware chain cannot see: goal evaluation, memory extraction, title generation, and summarization. Each site reports both terminal paths without changing the provider exception the host observes, short-circuits on `has_system_model_observers`, and passes the live task store when the runtime has one (detached work gets an isolated store). The sync summarization half stays unobserved on purpose -- it and its only host caller are the sync side of an async-only runtime, so notifying there would block a thread on a call site the host never reaches; the reason is recorded at the call site. The DeerMem backend must stay vendorable and cannot import the extension API, so it reports through a new `MemoryCallbacks.on_memory_llm_result` host hook that the DeerFlow-side callbacks translate into an observation. Notification loop ----------------- Extension resources must be touched on the loop that created them, but subagents can execute on isolated loops and DeerMem runs on a worker thread. The Gateway registers its serving loop before any runtime dependency starts and resets it last through the exit stack, so every startup-failure and cancellation path is covered. Awaited hooks raised on another loop are dispatched across with `run_coroutine_threadsafe` and awaited under the same budget; synchronous sites submit fire-and-forget work. Shutdown stops accepting detached observations before the memory flush -- that flush runs on a worker thread and can emit memory observations -- while keeping the loop alive for awaited task hooks until run and subagent drain completes. Tests ----- `test_extension_task_lifecycle.py`, `test_extension_subagent_lifecycle.py`, and `test_extension_system_model_calls.py` cover ordering, fail-open, budget exhaustion, snapshot binding under a concurrent singleton replacement, the loop-dispatch and shutdown-suspension paths, and both terminal paths at every call site. `test_gateway_run_drain_shutdown.py` pins the stop-before-barrier and drain ordering. * fix(extensions): decide notification fail-open by origin, observe cancellation `_notify_each` only guarded `Exception`, so a contributor letting a `CancelledError` escape — an extension implementing an internal timeout with cancellation, say — skipped its successors and reached the worker's deferred-interrupt path, ending an otherwise successful run as cancelled. Fail-open is about where a failure came from, not its base class: only a genuine cancellation of the host task increments `Task.cancelling()`, so propagate on that and contain everything else. `KeyboardInterrupt` / `SystemExit` still propagate. `observe_system_model_call` skipped observers on cancellation for the same base-class reason, leaving goal / title / summarization silent on a terminal path that is routine — interrupt/rollback admission and shutdown both cancel the run task, with the provider tokens already spent. Awaiting observers there is unreliable (a repeated cancel interrupts that await before any of them runs), so report through the same non-blocking submission the synchronous memory bridge uses, then propagate the cancellation untouched. DeerMem keeps `BaseException` around its provider call, now with the reason recorded: that path runs on a worker thread, where cancelling the awaiting side never interrupts the running thread, so `CancelledError` cannot arrive at all. Its host-hook wrapper narrows to `Exception` — only the hook's own failures are non-fatal, and an observability path must not swallow a process teardown signal. * fix(extensions): warn on budget exhaustion, scope observer logs by task, propagate teardown Review response on #4684: - The memory observation bridge caught BaseException, which would swallow a teardown signal raised while dispatching; it now catches Exception, matching the boundary the DeerMem-side call site documents and tests. - A notification-budget timeout raised mid-hook fell into the generic hook-failure path and logged an asyncio-internal traceback; it now logs a warning like the pre-hook budget skip, while a TimeoutError a contributor raises on its own stays classified as a hook failure. - System model observer logs passed the operation kind as the task id, so log lines said "task goal/title/..."; they now carry the task scope id alongside the kind. --- README.md | 26 +- backend/AGENTS.md | 92 ++- backend/app/gateway/app.py | 11 + backend/app/gateway/deps.py | 27 + .../demo_extensions.py | 6 +- .../deerflow_extension_api/__init__.py | 16 +- .../deerflow_extension_api/contracts.py | 106 ++- backend/packages/extension-api/pyproject.toml | 2 +- .../deerflow/agents/lead_agent/agent.py | 31 +- .../backends/deermem/deermem/core/updater.py | 61 +- .../harness/deerflow/agents/memory/manager.py | 79 ++ .../middlewares/summarization_middleware.py | 68 +- .../agents/middlewares/title_middleware.py | 43 +- .../tool_error_handling_middleware.py | 7 +- .../harness/deerflow/extensions/notify.py | 418 ++++++++++ .../harness/deerflow/extensions/registry.py | 53 +- .../packages/harness/deerflow/runtime/goal.py | 26 +- .../harness/deerflow/runtime/runs/worker.py | 68 +- .../harness/deerflow/subagents/executor.py | 60 +- backend/packages/harness/pyproject.toml | 2 +- backend/tests/test_extension_api_contracts.py | 93 ++- backend/tests/test_extension_config.py | 4 + backend/tests/test_extension_loader.py | 3 + backend/tests/test_extension_registry.py | 65 +- backend/tests/test_extension_stack_wiring.py | 14 + .../test_extension_subagent_lifecycle.py | 259 +++++++ .../test_extension_system_model_calls.py | 727 ++++++++++++++++++ .../tests/test_extension_task_lifecycle.py | 438 +++++++++++ .../tests/test_gateway_lifespan_shutdown.py | 34 +- .../tests/test_gateway_run_drain_shutdown.py | 53 ++ backend/tests/test_goal_worker.py | 10 +- .../tests/test_lead_agent_model_resolution.py | 4 +- .../tests/test_memory_manager_interface.py | 8 + backend/tests/test_memory_updater.py | 77 ++ .../test_tool_error_handling_middleware.py | 11 +- backend/uv.lock | 2 +- 36 files changed, 2916 insertions(+), 88 deletions(-) create mode 100644 backend/packages/harness/deerflow/extensions/notify.py create mode 100644 backend/tests/test_extension_subagent_lifecycle.py create mode 100644 backend/tests/test_extension_system_model_calls.py create mode 100644 backend/tests/test_extension_task_lifecycle.py 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]]