github/deer-flow
1
0
Fork 0
mirror of https://github.com/bytedance/deer-flow.git synced 2026-09-09 21:49:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(extensions): observe task lifecycle and system model calls (#4684)

Browse Source 
* 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.
This commit is contained in:
Nan Gao 2026-08-11 16:33:22 +08:00 committed by GitHub
parent a263af2845
commit 7389331e65
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
36 changed files with 2916 additions and 88 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
README.md
View File

@ -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.

92
backend/AGENTS.md
View File

@ -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

11
backend/app/gateway/app.py
View File

@ -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:

27
backend/app/gateway/deps.py
View File

@ -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)

6
backend/extension_test_fixtures/demo_extensions.py
View File

@ -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")

16
backend/packages/extension-api/deerflow_extension_api/__init__.py
View File

@ -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",
]

106
backend/packages/extension-api/deerflow_extension_api/contracts.py
View File

@ -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]

2
backend/packages/extension-api/pyproject.toml
View File

@ -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

31
backend/packages/harness/deerflow/agents/lead_agent/agent.py
View File

@ -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)

61
backend/packages/harness/deerflow/agents/memory/backends/deermem/deermem/core/updater.py
View File

@ -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],

79
backend/packages/harness/deerflow/agents/memory/manager.py
View File

@ -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.

68
backend/packages/harness/deerflow/agents/middlewares/summarization_middleware.py
View File

@ -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,
)

43
backend/packages/harness/deerflow/agents/middlewares/title_middleware.py
View File

@ -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),
)

7
backend/packages/harness/deerflow/agents/middlewares/tool_error_handling_middleware.py
View File

@ -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)

418
backend/packages/harness/deerflow/extensions/notify.py Normal file
View File

@ -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()

53
backend/packages/harness/deerflow/extensions/registry.py
View File

@ -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),
)

26
backend/packages/harness/deerflow/runtime/goal.py
View File

@ -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))

68
backend/packages/harness/deerflow/runtime/runs/worker.py
View File

@ -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

60
backend/packages/harness/deerflow/subagents/executor.py
View File

@ -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:

2
backend/packages/harness/pyproject.toml
View File

@ -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",

93
backend/tests/test_extension_api_contracts.py
View File

@ -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")

4
backend/tests/test_extension_config.py
View File

@ -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():

3
backend/tests/test_extension_loader.py
View File

@ -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():

65
backend/tests/test_extension_registry.py
View File

@ -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():

14
backend/tests/test_extension_stack_wiring.py
View File

@ -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"

259
backend/tests/test_extension_subagent_lifecycle.py Normal file
View File

@ -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]

727
backend/tests/test_extension_system_model_calls.py Normal file
View File

@ -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]

438
backend/tests/test_extension_task_lifecycle.py Normal file
View File

@ -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)

34
backend/tests/test_gateway_lifespan_shutdown.py
View File

@ -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

53
backend/tests/test_gateway_run_drain_shutdown.py
View File

@ -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

10
backend/tests/test_goal_worker.py
View File

@ -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

4
backend/tests/test_lead_agent_model_resolution.py
View File

@ -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(

8
backend/tests/test_memory_manager_interface.py
View File

@ -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

77
backend/tests/test_memory_updater.py
View File

@ -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": []}'

11
backend/tests/test_tool_error_handling_middleware.py
View File

@ -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

2
backend/uv.lock generated
View File

@ -907,7 +907,7 @@ dev = [
[[package]]
name = "deerflow-extension-api"
version = "0.1.0"
version = "0.1.1"
source = { editable = "packages/extension-api" }
[[package]]