mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-14 16:08:41 +00:00
7389331e65
* 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.
148 lines
5.3 KiB
Python
148 lines
5.3 KiB
Python
"""Tests for the extension registry and its immutable build product."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
|
|
from deerflow.extensions.registry import EMPTY_EXTENSIONS, ExtensionRegistry
|
|
|
|
|
|
class _Contributor:
|
|
def __init__(self, tag: str = "") -> None:
|
|
self.tag = tag
|
|
|
|
|
|
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
|
|
|
|
|
|
@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"):
|
|
register(registry, _Contributor())
|
|
|
|
assert registry.build().needs_task_store is True
|
|
|
|
|
|
def test_entries_carry_their_source():
|
|
registry = ExtensionRegistry()
|
|
contributor = _Contributor("mw")
|
|
with registry.attributed_to("demo_ext:install"):
|
|
registry.middlewares(contributor)
|
|
loaded = registry.build()
|
|
assert loaded.middleware_contributors == (("demo_ext:install", contributor),)
|
|
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")
|
|
with registry.attributed_to("a_ext:install"):
|
|
registry.middlewares(first)
|
|
with registry.attributed_to("b_ext:install"):
|
|
registry.middlewares(second)
|
|
loaded = registry.build()
|
|
assert [source for source, _ in loaded.middleware_contributors] == ["a_ext:install", "b_ext:install"]
|
|
|
|
|
|
def test_discard_removes_every_entry_of_one_source():
|
|
"""A partially-registered extension is worse than an absent one: the data
|
|
it produces looks complete but is not."""
|
|
registry = ExtensionRegistry()
|
|
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():
|
|
registry = ExtensionRegistry()
|
|
with pytest.raises(RuntimeError, match="attributed_to"):
|
|
registry.middlewares(_Contributor())
|
|
|
|
|
|
def test_build_result_is_frozen():
|
|
loaded = ExtensionRegistry().build()
|
|
with pytest.raises(Exception):
|
|
loaded.middleware_contributors = () # type: ignore[misc]
|
|
|
|
|
|
def test_app_store_is_created_at_build_time():
|
|
"""The app store must exist before binding so the registration phase and
|
|
the binding phase see the same object."""
|
|
loaded = ExtensionRegistry().build()
|
|
assert loaded.app_store is not None
|
|
assert loaded.app_store.scope_id == "app"
|