mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-11-08 03:30:34 +00:00
1dd6d05627
* feat(extension-api): add task identity fields to CompactionEvent (0.2.10) CompactionEvent carried the transform, its hashes and its counts, but nothing that names the task, run or thread the compaction happened in. An observer that wants to relate a compaction to the run that produced it had to match hashes across sessions, which is unreliable. Append three optional fields, task_id, run_id and thread_id (str | None, default None). Existing keyword and positional construction is unchanged, so already-released extensions keep working. The host fills them in the next commit. Bump deerflow-extension-api to 0.2.10 and the harness's exact pin with it. * fix(extensions): hand compaction observers the live task store and its identity The summarization middleware already recovers the live task store from the runtime to observe its summary system call, and the registry allocates that store precisely because a compaction observer is registered. notify_context_compacted() then discarded it and handed every observer a detached store, with a docstring claiming there was no live task. - The lead run worker and the subagent executor now seed every task store they allocate with the task's TaskInfo, not only when a lifecycle contributor is registered; the lifecycle notification reuses the seeded object. A subagent without a parent run_id still gets no TaskInfo. - compact_state() and acompact_state() pass the store from task_store_from_runtime() to _record_compaction(). - notify_context_compacted() takes an optional task_store. A live store goes to the observers and fills the event's task_id/run_id/thread_id from its TaskInfo; without one (manual /compact runs outside any run) observers still receive a detached store and the ids stay None. For a background subagent the store keeps its public scope id (the provider tool_call_id) while TaskInfo.task_id stays the execution id, as lifecycle contributors already see it; the event follows TaskInfo. * docs(extensions): document compaction event identity and seeded TaskInfo The extensions guide and the developer manual still said compaction observers always receive a detached store because no task is live at the call site. Describe what the host does now: lead runs and subagents with a parent run seed every task store with TaskInfo, compaction observers get the live store, and CompactionEvent's task_id/run_id/thread_id come from it; only manual /compact falls back to a detached store with None ids. Adds the 0.2.10 row to the contract version table (PR number pending). * docs(changelog): note compaction observers receive the live task identity * docs(extensions): list compaction observers among task-store triggers * docs: reference #7831 for the compaction event identity change * docs(extensions): backfill contract history rows 0.2.5 to 0.2.9 The version table jumped from 0.2.4 to 0.2.10 although 0.2.5-0.2.9 shipped. Add those rows, attributing names merged without a bump to the first release that carried them (#5796 in 0.2.5, #6920 and #5645 in 0.2.8, #7729 in 0.2.10). The page intro now says it documents 0.2.4 plus the 0.2.10 CompactionEvent fields, and points to the history for the rest.
580 lines
20 KiB
Python
580 lines
20 KiB
Python
"""Fail-open notification helpers for extension runtime hooks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import concurrent.futures
|
|
import dataclasses
|
|
import logging
|
|
import threading
|
|
import time
|
|
from collections.abc import Awaitable, Callable, Coroutine, Mapping
|
|
from typing import Any
|
|
|
|
from deerflow_extension_api import (
|
|
EXTENSION_TASK_STORE_KEY,
|
|
CompactionEvent,
|
|
ExtensionData,
|
|
SystemModelRequest,
|
|
SystemModelResult,
|
|
SystemOperationKind,
|
|
TaskInfo,
|
|
TaskOutcome,
|
|
)
|
|
|
|
from deerflow.extensions.registry import LoadedExtensions
|
|
from deerflow.utils.file_io import await_drained
|
|
|
|
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.
|
|
"""
|
|
try:
|
|
task = asyncio.current_task()
|
|
except RuntimeError:
|
|
# Synchronous hook sites (agent assembly) can run with no loop at all,
|
|
# and "no loop" means there is no host task being cancelled.
|
|
return False
|
|
return task is not None and task.cancelling() > 0
|
|
|
|
|
|
def notify_agent_assembled(descriptor: object, extensions: object | None = None) -> None:
|
|
"""Fan a completed assembly out to observers, in registration order.
|
|
|
|
Synchronous: agent construction is synchronous and there is no loop to
|
|
dispatch onto. Failures are contained per observer — a broken observer must
|
|
not prevent an agent from being built.
|
|
"""
|
|
resolved = extensions
|
|
if resolved is None:
|
|
from deerflow.extensions import get_agent_build_extensions
|
|
|
|
resolved = get_agent_build_extensions()
|
|
observers = getattr(resolved, "agent_assembly_observers", ())
|
|
if not observers:
|
|
return
|
|
app_store = getattr(resolved, "app_store", None)
|
|
for source, observer in observers:
|
|
try:
|
|
observer.on_agent_assembled(app_store, descriptor)
|
|
except asyncio.CancelledError:
|
|
if _host_is_cancelling():
|
|
raise
|
|
# Same rule as the awaited hooks: an observer raising it on its own
|
|
# must not skip its successors, and must not turn graph
|
|
# construction into a deferred interrupt.
|
|
logger.exception(
|
|
"Extension %s: on_agent_assembled raised CancelledError for %s",
|
|
source,
|
|
type(descriptor).__name__,
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
"Extension %s: on_agent_assembled failed for %s",
|
|
source,
|
|
type(descriptor).__name__,
|
|
)
|
|
|
|
|
|
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
|
|
# ``budget_spent`` latches once the shared budget runs out, so the skip
|
|
# decision never depends on a second clock read. The loop runs a scheduled
|
|
# timeout up to one ``clock_resolution`` *early* (``BaseEventLoop._run_once``
|
|
# compares each handle's ``_when`` against ``now + clock_resolution``), so on
|
|
# a coarse clock — Windows' ~15.6 ms against budgets as small as 20 ms —
|
|
# ``loop.time()`` can still sit strictly before ``deadline`` at the moment
|
|
# the budget has in fact expired. That raced both the log classification and
|
|
# the skip of later contributors (#6453).
|
|
budget_spent = False
|
|
for source, contributor in contributors:
|
|
timeout_scope: asyncio.Timeout | None = None
|
|
try:
|
|
call = invoke(contributor)
|
|
if deadline is None:
|
|
await call
|
|
continue
|
|
|
|
remaining = 0.0 if budget_spent else deadline - loop.time()
|
|
if remaining <= 0:
|
|
budget_spent = True
|
|
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
|
|
# ``asyncio.timeout`` records whether *it* expired, so the handler
|
|
# below never re-reads the clock to tell a spent budget from a
|
|
# ``TimeoutError`` the contributor raised on its own.
|
|
timeout_scope = asyncio.timeout(remaining)
|
|
async with timeout_scope:
|
|
await call
|
|
except TimeoutError:
|
|
if timeout_scope is not None and timeout_scope.expired():
|
|
budget_spent = True
|
|
# 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[concurrent.futures.Future[None]] = set()
|
|
_warned_no_loop = False
|
|
_system_observations_enabled = True
|
|
# Serializes detached-dispatch admission with shutdown suspension and snapshots.
|
|
_dispatch_state_lock = threading.Lock()
|
|
|
|
|
|
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
|
|
with _dispatch_state_lock:
|
|
_notify_loop = loop
|
|
_system_observations_enabled = True
|
|
_warned_no_loop = False
|
|
|
|
|
|
async def drain_extension_notify_dispatches(*, timeout: float = 5.0) -> None:
|
|
"""Boundedly drain already-submitted fire-and-forget observations.
|
|
|
|
New detached observations are suspended before the snapshot so shutdown
|
|
cannot open a submit-after-snapshot window. The bounded wait keeps a wedged
|
|
extension observer from holding Gateway teardown forever.
|
|
"""
|
|
suspend_extension_system_observations()
|
|
# Once suspension owns the same lock as dispatch admission, no new future
|
|
# can be registered; completed futures may still be removed by callbacks.
|
|
with _dispatch_state_lock:
|
|
pending = tuple(_pending_dispatches)
|
|
if not pending:
|
|
return
|
|
|
|
async def _wait() -> None:
|
|
try:
|
|
async with asyncio.timeout(timeout):
|
|
await asyncio.gather(
|
|
*(asyncio.wrap_future(future) for future in pending),
|
|
return_exceptions=True,
|
|
)
|
|
except TimeoutError:
|
|
logger.warning(
|
|
"Timed out after %.1fs draining %d pending extension observation(s)",
|
|
timeout,
|
|
len(pending),
|
|
)
|
|
|
|
await await_drained(_wait())
|
|
|
|
|
|
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
|
|
with _dispatch_state_lock:
|
|
_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
|
|
with _dispatch_state_lock:
|
|
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 _discard_pending_dispatch(future: concurrent.futures.Future[None]) -> None:
|
|
with _dispatch_state_lock:
|
|
_pending_dispatches.discard(future)
|
|
|
|
|
|
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
|
|
|
|
submitted = False
|
|
try:
|
|
with _dispatch_state_lock:
|
|
if not _system_observations_enabled:
|
|
return False
|
|
loop = _notify_loop
|
|
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
|
|
# Registration is part of admission: shutdown cannot suspend and
|
|
# snapshot between the enabled check and this insertion.
|
|
_pending_dispatches.add(future)
|
|
submitted = True
|
|
# Register outside the lock because add_done_callback may run inline for
|
|
# an already-completed Future; the callback acquires the same lock.
|
|
future.add_done_callback(_discard_pending_dispatch)
|
|
return True
|
|
finally:
|
|
if not submitted:
|
|
coro.close()
|
|
|
|
|
|
def notify_context_compacted(
|
|
event: CompactionEvent,
|
|
extensions: LoadedExtensions | None = None,
|
|
task_store: ExtensionData | None = None,
|
|
) -> None:
|
|
"""Fan a completed compaction out to observers, fire-and-forget.
|
|
|
|
The compaction seam sits in the summarization middleware's ``before_model`` /
|
|
``abefore_model`` hooks; the sync half has no loop to await onto, and the async
|
|
half must not block the model-call turn on observer latency. Both therefore call
|
|
this synchronous entry point, which dispatches to the registered extension-notify
|
|
loop the same non-blocking way a synchronous system-model-call cancellation does,
|
|
reusing the same fail-open cancellation containment inside ``_notify_each``.
|
|
|
|
Inside a lead run or subagent the caller passes the live task store it recovered
|
|
from the runtime context — the same store the summary system call is observed
|
|
with. Observers receive it, and the event's ``task_id``/``run_id``/``thread_id``
|
|
are filled from the ``TaskInfo`` the host seeded into it. A compaction outside any
|
|
task (manual ``/compact``) passes none, so observers receive a detached store —
|
|
the same fallback ``notify_system_model_call`` uses — and the ids stay ``None``.
|
|
"""
|
|
resolved = extensions
|
|
if resolved is None:
|
|
from deerflow.extensions import get_agent_build_extensions
|
|
|
|
resolved = get_agent_build_extensions()
|
|
observers = resolved.context_compaction_observers
|
|
if not observers:
|
|
return
|
|
app_store = resolved.app_store
|
|
if task_store is None:
|
|
task_store = ExtensionData("detached")
|
|
else:
|
|
info = task_store.get(TaskInfo)
|
|
if info is not None:
|
|
event = dataclasses.replace(event, task_id=info.task_id, run_id=info.run_id, thread_id=info.thread_id)
|
|
what = f"compaction ({event.transform_kind})"
|
|
dispatch_system_model_observation(
|
|
_notify_each(
|
|
observers,
|
|
"on_context_compacted",
|
|
lambda observer: observer.on_context_compacted(app_store, task_store, event),
|
|
what,
|
|
None,
|
|
),
|
|
what,
|
|
)
|