diff --git a/examples/deerflow-extension-example/README.md b/examples/deerflow-extension-example/README.md index 5890fa7b2b..e6ce490a64 100644 --- a/examples/deerflow-extension-example/README.md +++ b/examples/deerflow-extension-example/README.md @@ -12,7 +12,7 @@ depends on FastAPI, LangChain, and LangGraph in `pyproject.toml`. | Contribution | Example behavior | | --- | --- | -| Middleware | Counts tool calls through one `TOOL_VISIBLE` middleware for lead agents and subagents | +| Middleware | Counts synchronous and asynchronous tool calls, outcomes, and elapsed time at `TOOL_VISIBLE` for lead agents and subagents | | Task lifecycle | Creates task-scoped stats on start and folds them into app scope on stop | | System-model observer | Counts DeerFlow-owned model calls, including failures | | Service | Binds `ExtensionRuntimeDeps` only while the Gateway is running | @@ -24,6 +24,48 @@ the same `ExampleService` object: its FastAPI dependency returns `503` before `start()`, after `stop()`, or when no app store was bound. This keeps the route topology stable while runtime capabilities arrive later. +### Tool statistics + +The existing `tool_calls` field counts calls entering the wrapper. The stats +response also includes these fields (this is illustrative output, not a benchmark): + +```json +{ + "tool_calls": 3, + "tool_outcomes": {"returned": 1, "raised": 1, "cancelled": 1}, + "tool_duration_samples": 3, + "tool_duration_total_ms": 60.0 +} +``` + +- `returned` means the downstream handler returned normally, including an error + `ToolMessage`. It does **not** measure business success. A tool exception + converted into a normal return downstream of `TOOL_VISIBLE` counts here too. +- `cancelled` counts `asyncio.CancelledError` escaping the handler, including + actual task cancellation. `raised` counts other escaping exceptions, including + LangGraph `GraphBubbleUp` control-flow signals such as graph interrupts. + Results and exceptions propagate unchanged; each wrapper calls its handler once. + Resuming an interrupted tool node re-executes it and re-enters the wrapper: + the interrupted attempt adds a `raised` sample, and the resumed attempt adds + its own `tool_calls` increment, duration sample, and outcome. These statistics + count observed attempts, not unique logical tool calls across interrupt/resume. +- Each completed observation adds one duration sample, including zero-duration + calls. Elapsed time uses a monotonic clock and includes downstream middleware + processing. Durations of overlapping lead/subagent calls are summed, so this + total is **not** elapsed run time or exclusive tool execution time. + +Outcome counters, duration samples, and elapsed totals are updated together under +a lock. Task stats are copied and merged into app scope on task stop, after the +task's tool calls have settled. The first stop consumes the task stats; duplicate +stops (or a stop without a matching start) add neither tasks nor tool statistics. +If a host stops a task while calls are still running, only the snapshot at that +stop is merged; later completions are not published. Active tasks are not included +in the app's tool totals. The system-model observer retains its existing behavior. + +The new fields contain only fixed counters and timing totals, with no tool names, +arguments, outputs, exception text, or identifying labels. Stats are process-local +and reset on restart; they are not persisted or shared across workers. + ## Run the package tests `deerflow-extension-api` is currently sourced from this checkout. Install it @@ -118,7 +160,7 @@ After one or more runs, request the extension route: curl -s http://localhost:2026/api/extension-example/stats ``` -The response contains aggregated task outcomes, tool-call counts, system-model +The response contains aggregated task outcomes, tool-call counts, outcomes and timing, system-model call counts, the app scope id, and a small projection of the host policy. The route passes through the Gateway's normal authentication middleware; use an authenticated browser session when authentication is enabled. @@ -145,5 +187,6 @@ deerflow_extension_example/ └── plugin.py # state plus all five small contribution implementations tests/ ├── test_entry_point.py -└── test_plugin.py +├── test_plugin.py +└── test_tool_stats.py ``` diff --git a/examples/deerflow-extension-example/deerflow_extension_example/plugin.py b/examples/deerflow-extension-example/deerflow_extension_example/plugin.py index fb301c6d7b..c99efc6164 100644 --- a/examples/deerflow-extension-example/deerflow_extension_example/plugin.py +++ b/examples/deerflow-extension-example/deerflow_extension_example/plugin.py @@ -2,9 +2,12 @@ from __future__ import annotations -from collections.abc import Awaitable, Callable, Sequence +from asyncio import CancelledError +from collections.abc import Awaitable, Callable, Iterator, Sequence +from contextlib import contextmanager from dataclasses import dataclass, field from threading import Lock +from time import monotonic_ns from typing import Any from deerflow_extension_api import ( @@ -31,6 +34,9 @@ class ExampleStats: """Small extension-owned value used in both app and task stores.""" tool_calls: int = 0 + tool_outcomes: dict[str, int] = field(default_factory=lambda: {"returned": 0, "raised": 0, "cancelled": 0}) + tool_duration_samples: int = 0 + tool_duration_total_ns: int = 0 tasks: dict[str, int] = field(default_factory=dict) system_model_calls: dict[str, dict[str, int]] = field(default_factory=dict) _lock: Lock = field(default_factory=Lock, repr=False, compare=False) @@ -39,13 +45,25 @@ class ExampleStats: with self._lock: self.tool_calls += 1 - def task_tool_calls(self) -> int: + def note_tool_outcome(self, outcome: str, duration_ns: int) -> None: with self._lock: - return self.tool_calls + self.tool_outcomes[outcome] += 1 + self.tool_duration_samples += 1 + self.tool_duration_total_ns += duration_ns - def absorb_task(self, tool_calls: int, outcome: TaskOutcome) -> None: + def absorb_task(self, task_stats: ExampleStats, outcome: TaskOutcome) -> None: + # Copy the task under its lock, then merge under the app lock. + with task_stats._lock: + tool_calls = task_stats.tool_calls + outcomes = dict(task_stats.tool_outcomes) + samples = task_stats.tool_duration_samples + duration_ns = task_stats.tool_duration_total_ns with self._lock: self.tool_calls += tool_calls + for name, count in outcomes.items(): + self.tool_outcomes[name] += count + self.tool_duration_samples += samples + self.tool_duration_total_ns += duration_ns key = outcome.value self.tasks[key] = self.tasks.get(key, 0) + 1 @@ -64,6 +82,9 @@ class ExampleStats: return { "tasks": dict(self.tasks), "tool_calls": self.tool_calls, + "tool_outcomes": dict(self.tool_outcomes), + "tool_duration_samples": self.tool_duration_samples, + "tool_duration_total_ms": self.tool_duration_total_ns / 1_000_000, "system_model_calls": {kind: dict(counts) for kind, counts in self.system_model_calls.items()}, } @@ -73,16 +94,43 @@ def _stats(store: ExtensionData) -> ExampleStats: class ExampleMiddleware(AgentMiddleware): + @contextmanager + def _observe(self, request: ToolCallRequest) -> Iterator[None]: + task_store = task_store_from_runtime(getattr(request, "runtime", None)) + task_stats = task_store.get(ExampleStats) if task_store is not None else None + if task_stats is None: + yield + return + task_stats.note_tool_call() + started = monotonic_ns() + outcome = "returned" + try: + yield + except CancelledError: + outcome = "cancelled" + raise + except BaseException: + # Includes graph interrupts; observation must not consume control flow. + outcome = "raised" + raise + finally: + task_stats.note_tool_outcome(outcome, monotonic_ns() - started) + + def wrap_tool_call( + self, + request: ToolCallRequest, + handler: Callable[[ToolCallRequest], Any], + ) -> Any: + with self._observe(request): + return handler(request) + async def awrap_tool_call( self, request: ToolCallRequest, handler: Callable[[ToolCallRequest], Awaitable[Any]], ) -> Any: - task_store = task_store_from_runtime(getattr(request, "runtime", None)) - task_stats = task_store.get(ExampleStats) if task_store is not None else None - if task_stats is not None: - task_stats.note_tool_call() - return await handler(request) + with self._observe(request): + return await handler(request) class ExampleMiddlewareContributor: @@ -117,10 +165,8 @@ class ExampleTaskLifecycle: outcome: TaskOutcome, ) -> None: task_stats = task_store.remove(ExampleStats) - _stats(app_store).absorb_task( - task_stats.task_tool_calls() if task_stats is not None else 0, - outcome, - ) + if task_stats is not None: + _stats(app_store).absorb_task(task_stats, outcome) class ExampleSystemObserver: diff --git a/examples/deerflow-extension-example/tests/test_plugin.py b/examples/deerflow-extension-example/tests/test_plugin.py index 7707d75ef6..95b2c8f142 100644 --- a/examples/deerflow-extension-example/tests/test_plugin.py +++ b/examples/deerflow-extension-example/tests/test_plugin.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio from dataclasses import dataclass, field from typing import Any +from unittest.mock import patch from deerflow_extension_api import ( EXTENSION_TASK_STORE_KEY, @@ -25,6 +26,9 @@ from deerflow_extension_example import install class FakeRegistry: + def plugin(self, contribution: Any) -> bool: + return False + def __init__(self) -> None: self.middleware_contributors: list[Any] = [] self.task_lifecycle_contributors: list[Any] = [] @@ -158,7 +162,8 @@ def test_registered_contributions_publish_one_shared_stats_snapshot() -> None: after_stop.status_code, ) - before_start, status_code, body, after_stop = asyncio.run(exercise_contributions()) + with patch("deerflow_extension_example.plugin.monotonic_ns", side_effect=[0, 10_000_000]): + before_start, status_code, body, after_stop = asyncio.run(exercise_contributions()) assert before_start == 503 assert status_code == 200 @@ -169,5 +174,8 @@ def test_registered_contributions_publish_one_shared_stats_snapshot() -> None: "host_policy": {"max_subagents_per_run": 6}, "tasks": {"completed": 1}, "tool_calls": 1, + "tool_outcomes": {"returned": 1, "raised": 0, "cancelled": 0}, + "tool_duration_samples": 1, + "tool_duration_total_ms": 10.0, "system_model_calls": {"title": {"calls": 1, "errors": 1}}, } diff --git a/examples/deerflow-extension-example/tests/test_tool_stats.py b/examples/deerflow-extension-example/tests/test_tool_stats.py new file mode 100644 index 0000000000..f80cdfc06f --- /dev/null +++ b/examples/deerflow-extension-example/tests/test_tool_stats.py @@ -0,0 +1,221 @@ +from __future__ import annotations + +import asyncio +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from deerflow_extension_api import EXTENSION_TASK_STORE_KEY, ExtensionData, TaskInfo, TaskOutcome +from langchain_core.messages import ToolMessage +from langgraph.errors import GraphInterrupt + +from deerflow_extension_example.plugin import ExampleMiddleware, ExampleStats, ExampleTaskLifecycle + + +def task_context(): + app = ExtensionData("app") + task = ExtensionData("task") + info = TaskInfo(task_id="task", run_id="run", thread_id="thread", kind="lead") + lifecycle = ExampleTaskLifecycle() + asyncio.run(lifecycle.on_task_start(app, task, info)) + request = SimpleNamespace(runtime=SimpleNamespace(context={EXTENSION_TASK_STORE_KEY: task})) + return app, task, info, lifecycle, request + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("outcome", ["returned", "raised", "interrupt"]) +def test_preserves_result_exception_and_single_invocation(asynchronous, outcome): + app, task, info, lifecycle, request = task_context() + result = ToolMessage(content="tool reported an error", tool_call_id="call", status="error") + error = GraphInterrupt(()) if outcome == "interrupt" else ValueError("tool failed") + calls = [] + + def handler(received): + calls.append(received) + if outcome != "returned": + raise error + return result + + async def async_handler(received): + return handler(received) + + def invoke(): + middleware = ExampleMiddleware() + if asynchronous: + return asyncio.run(middleware.awrap_tool_call(request, async_handler)) + return middleware.wrap_tool_call(request, handler) + + with patch("deerflow_extension_example.plugin.monotonic_ns", side_effect=[0, 10_000_000]): + if outcome == "returned": + assert invoke() is result + else: + with pytest.raises(type(error)) as caught: + invoke() + assert caught.value is error + assert len(calls) == 1 and calls[0] is request + asyncio.run(lifecycle.on_task_stop(app, task, info, TaskOutcome.COMPLETED)) + snapshot = app.get(ExampleStats).snapshot() + assert snapshot["tool_calls"] == 1 + assert snapshot["tool_outcomes"] == {"returned": int(outcome == "returned"), "raised": int(outcome != "returned"), "cancelled": 0} + assert snapshot["tool_duration_samples"] == 1 + assert snapshot["tool_duration_total_ms"] == 10.0 + + +def test_actual_task_cancellation_propagates_and_is_counted(): + app, task, info, lifecycle, request = task_context() + errors = [] + calls = [] + + async def exercise(): + entered = asyncio.Event() + + async def handler(received): + calls.append(received) + entered.set() + try: + await asyncio.Future() + except asyncio.CancelledError as error: + errors.append(error) + raise + + running = asyncio.create_task(ExampleMiddleware().awrap_tool_call(request, handler)) + await entered.wait() + running.cancel() + with pytest.raises(asyncio.CancelledError) as caught: + await running + assert caught.value is errors[0] + assert running.cancelled() + await lifecycle.on_task_stop(app, task, info, TaskOutcome.ABORTED) + + with patch("deerflow_extension_example.plugin.monotonic_ns", side_effect=[0, 20_000_000]): + asyncio.run(exercise()) + assert len(calls) == 1 and calls[0] is request + snapshot = app.get(ExampleStats).snapshot() + assert snapshot["tool_outcomes"] == {"returned": 0, "raised": 0, "cancelled": 1} + assert snapshot["tool_duration_samples"] == 1 + assert snapshot["tool_duration_total_ms"] == 20.0 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("scope", ["no-runtime", "no-store", "not-started"]) +def test_missing_task_scope_passes_through(asynchronous, scope): + request = SimpleNamespace() + if scope != "no-runtime": + request.runtime = SimpleNamespace(context={}) + if scope == "not-started": + request.runtime.context[EXTENSION_TASK_STORE_KEY] = ExtensionData("unused") + calls = [] + result = object() + + def handler(received): + calls.append(received) + return result + + async def async_handler(received): + return handler(received) + + middleware = ExampleMiddleware() + if asynchronous: + assert asyncio.run(middleware.awrap_tool_call(request, async_handler)) is result + else: + assert middleware.wrap_tool_call(request, handler) is result + assert len(calls) == 1 and calls[0] is request + + +def test_concurrent_sync_calls_and_duplicate_finalization(): + app, task, info, lifecycle, request = task_context() + barrier = Barrier(8) + + def handler(received): + assert received is request + barrier.wait(timeout=10) + return received + + def invoke(_): + assert ExampleMiddleware().wrap_tool_call(request, handler) is request + + # A constant clock also verifies that zero durations still produce samples. + with patch("deerflow_extension_example.plugin.monotonic_ns", return_value=42): + with ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(invoke, range(8))) + assert app.get(ExampleStats) is None + + def finalize(_): + asyncio.run(lifecycle.on_task_stop(app, task, info, TaskOutcome.COMPLETED)) + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(finalize, range(4))) + snapshot = app.get(ExampleStats).snapshot() + assert snapshot["tasks"] == {"completed": 1} + assert snapshot["tool_calls"] == 8 + assert snapshot["tool_outcomes"] == {"returned": 8, "raised": 0, "cancelled": 0} + assert snapshot["tool_duration_samples"] == 8 + assert snapshot["tool_duration_total_ms"] == 0.0 + snapshot["tool_outcomes"]["returned"] = -1 + assert app.get(ExampleStats).snapshot()["tool_outcomes"]["returned"] == 8 + + +def test_overlapping_async_calls_merge_outcomes_and_elapsed_time(): + app, task, info, lifecycle, request = task_context() + + async def exercise(): + ready = asyncio.Event() + entered = 0 + + async def handler(_): + nonlocal entered + entered += 1 + number = entered + if entered == 3: + ready.set() + await ready.wait() + if number == 2: + raise ValueError("failed") + return number + + results = await asyncio.gather(*(ExampleMiddleware().awrap_tool_call(request, handler) for _ in range(3)), return_exceptions=True) + assert sum(isinstance(result, ValueError) for result in results) == 1 + await lifecycle.on_task_stop(app, task, info, TaskOutcome.COMPLETED) + + # Three starts at zero, then completions after 10, 20 and 30 milliseconds. + with patch("deerflow_extension_example.plugin.monotonic_ns", side_effect=[0, 0, 0, 10_000_000, 20_000_000, 30_000_000]): + asyncio.run(exercise()) + snapshot = app.get(ExampleStats).snapshot() + assert snapshot["tool_calls"] == 3 + assert snapshot["tool_outcomes"] == {"returned": 2, "raised": 1, "cancelled": 0} + assert snapshot["tool_duration_samples"] == 3 + assert snapshot["tool_duration_total_ms"] == 60.0 + + +def test_multiple_task_stores_merge_once_and_unstarted_stop_is_ignored(): + app, task, info, lifecycle, request = task_context() + second = ExtensionData("second") + unstarted = ExtensionData("unstarted") + + async def exercise(): + await lifecycle.on_task_stop(app, unstarted, info, TaskOutcome.FAILED) + assert app.get(ExampleStats) is None + await lifecycle.on_task_start(app, second, info) + + async def handler(_): + return None + + await ExampleMiddleware().awrap_tool_call(request, handler) + await lifecycle.on_task_stop(app, task, info, TaskOutcome.COMPLETED) + first_snapshot = app.get(ExampleStats).snapshot() + await lifecycle.on_task_stop(app, task, info, TaskOutcome.FAILED) + assert app.get(ExampleStats).snapshot() == first_snapshot + second_request = SimpleNamespace(runtime=SimpleNamespace(context={EXTENSION_TASK_STORE_KEY: second})) + await ExampleMiddleware().awrap_tool_call(second_request, handler) + await lifecycle.on_task_stop(app, second, info, TaskOutcome.ABORTED) + + with patch("deerflow_extension_example.plugin.monotonic_ns", side_effect=[0, 10_000_000, 0, 20_000_000]): + asyncio.run(exercise()) + snapshot = app.get(ExampleStats).snapshot() + assert snapshot["tasks"] == {"completed": 1, "aborted": 1} + assert snapshot["tool_calls"] == 2 + assert snapshot["tool_outcomes"] == {"returned": 2, "raised": 0, "cancelled": 0} + assert snapshot["tool_duration_samples"] == 2 + assert snapshot["tool_duration_total_ms"] == 30.0