From 88252e9b318d34e7e1867155ad2c77993320788e Mon Sep 17 00:00:00 2001 From: ChiHaYa <97534761+ZeroMadLife@users.noreply.github.com> Date: Wed, 12 Aug 2026 09:25:05 +0800 Subject: [PATCH] fix(subagents): isolate background tasks from reused tool call IDs (#4758) * fix(subagents): isolate background execution IDs * fix(subagents): preserve correlation scope and isolate usage * fix(subagents): make usage attribution idempotent --- README.md | 2 +- backend/AGENTS.md | 3 +- .../middlewares/token_usage_middleware.py | 28 +++- .../harness/deerflow/subagents/executor.py | 78 ++++++----- .../deerflow/tools/builtins/task_tool.py | 124 +++++++----------- backend/tests/test_custom_events.py | 1 - .../test_extension_task_store_runtime.py | 40 +++++- backend/tests/test_subagent_executor.py | 75 +++++++++++ backend/tests/test_task_tool_core_logic.py | 93 +++---------- backend/tests/test_token_usage_middleware.py | 100 +++++++++++--- 10 files changed, 324 insertions(+), 220 deletions(-) diff --git a/README.md b/README.md index 9fd72aa69..37c76c85a 100644 --- a/README.md +++ b/README.md @@ -947,7 +947,7 @@ The chat header also shows a context-window gauge when the selected model has a Sub-agents are an optimization, not the default response to a complex request. -The lead agent can spawn sub-agents on the fly — each with its own scoped context, tools, and termination conditions — when delegation has clear net benefit from real parallel latency, specialist capability, or context isolation. It keeps interdependent scopes and overlapping side effects out of parallel dispatch; a bounded sequential chain can still run in one sub-agent when specialist or context-isolation benefit clearly wins. The lead uses the fewest useful sub-agents and re-evaluates later batches instead of fanning out solely because a task is large or multi-step. Sub-agents report back structured results, and the lead agent verifies and synthesizes them into a coherent output. Their configured skills are resolved from the same user-scoped catalog as the lead agent, so user-owned custom skills remain available without exposing another user's version. Their internal AI and tool messages stay scoped to the delegated graph instead of entering the parent chat stream. Reloaded thread history enforces the same boundary: callback-captured sub-agent AI responses remain available in run-event diagnostics but are excluded from the parent transcript, while the parent `task` result remains attached to its subtask card. Long-running sub-agents compact older history when summarization is enabled and re-inject the summary as guarded, hidden durable context before continuing, so recent assistant/tool activity remains grounded in the task. Provider/model request failures are reported as failed sub-agent tasks rather than successful results, so the lead agent and Web UI can react to them correctly. Collapsed sub-agent cards show the effective model and, when the provider returns usage metadata, a cumulative token total that updates after each completed sub-agent LLM call and persists after a reload. When token usage tracking is enabled, completed sub-agent usage is also attributed back to the dispatching step. +The lead agent can spawn sub-agents on the fly — each with its own scoped context, tools, and termination conditions — when delegation has clear net benefit from real parallel latency, specialist capability, or context isolation. It keeps interdependent scopes and overlapping side effects out of parallel dispatch; a bounded sequential chain can still run in one sub-agent when specialist or context-isolation benefit clearly wins. The lead uses the fewest useful sub-agents and re-evaluates later batches instead of fanning out solely because a task is large or multi-step. Sub-agents report back structured results, and the lead agent verifies and synthesizes them into a coherent output. Their configured skills are resolved from the same user-scoped catalog as the lead agent, so user-owned custom skills remain available without exposing another user's version. Their internal AI and tool messages stay scoped to the delegated graph instead of entering the parent chat stream. Reloaded thread history enforces the same boundary: callback-captured sub-agent AI responses remain available in run-event diagnostics but are excluded from the parent transcript, while the parent `task` result remains attached to its subtask card. Long-running sub-agents compact older history when summarization is enabled and re-inject the summary as guarded, hidden durable context before continuing, so recent assistant/tool activity remains grounded in the task. Provider/model request failures are reported as failed sub-agent tasks rather than successful results, so the lead agent and Web UI can react to them correctly. Concurrent parent runs also receive independent server-side sub-agent execution IDs, so a provider that reuses a tool-call ID cannot make one run poll, cancel, or clean up another run's background task. Collapsed sub-agent cards show the effective model and, when the provider returns usage metadata, a cumulative token total that updates after each completed sub-agent LLM call and persists after a reload. When token usage tracking is enabled, completed sub-agent usage is attributed back to the dispatching step from that run's terminal tool-message metadata rather than a process-global provider-ID cache. For example, independent read-only research can run concurrently when the wall-clock savings outweigh duplicated discovery and synthesis cost, while a repository refactor with shared files and sequential test feedback remains with the lead agent. When `max_concurrent_subagents` is `1`, parallel and multi-batch routing guidance is disabled; delegation remains available only for material specialist or context-isolation benefit. diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 4a58e2ffd..d97665749 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -16,6 +16,7 @@ DeerFlow is a LangGraph-based AI super agent system with a full-stack architectu - `make dev`, Docker dev, and production all run the agent runtime in Gateway via `RunManager` + `run_agent()` + `StreamBridge` (`packages/harness/deerflow/runtime/`). Nginx exposes that runtime at `/api/langgraph/*` and rewrites it to Gateway's native `/api/*` routers. - Gateway streams `write_file` and `str_replace` argument deltas in bounded batches when clients also subscribe to `values`; messages-only consumers retain the original per-chunk contract, while `values` preserves the complete tool call. - With `stream_subgraphs`, subgraph frames keep their namespace in the SSE event name (`values|`, LangGraph Platform style) instead of impersonating root frames — a delegated subagent inherits the parent checkpoint namespace, so publishing its `values` snapshot as bare `values` replaces the whole thread view in SDK clients (#4399). Root-only consumers (file-tool chunk batcher, subagent event persistence, LLM error-fallback detection) ignore namespaced frames. The web frontend does not request subgraph streaming; subtask progress rides root-namespace `task_*` custom events. +- Background subagent identity is deliberately split: the provider `tool_call_id` remains the correlation key for `ToolMessage`, `task_*` SSE events, persisted lifecycle events, frontend cards, and the public `ExtensionData.scope_id` contract (stored as `SubagentResult.external_task_id`), while `SubagentExecutor.execute_async()` generates a full server-side `execution_id` for `SubagentResult.task_id`, the process-wide registry, polling, cancellation, timeout handling, and cleanup. Provider IDs are not globally unique across parent runs, so they must never become registry ownership keys; scheduler closures retain their own `SubagentResult` rather than resolving ownership again through the mutable registry. Terminal subagent token usage travels in the current run's `ToolMessage.additional_kwargs` and is attributed from message state, never through a process-global provider-ID cache. - Scheduled-task executions must reuse that same Gateway run lifecycle. The scheduler may decide *when* work runs, but it must dispatch through the existing run path rather than introducing a parallel execution stack. - Long-running MCP work uses a separate durable task runtime rather than keeping remote task IDs or status polling inside the Agent loop. `McpTaskService` claims due rows with leases, resolves a protocol-specific `McpTaskDriver`, and writes normalized snapshots back to `mcp_tasks`; expired leases are the restart-recovery mechanism, and a result returned after expiry must be discarded even when the owner token still matches. The database is the source of truth. `ThreadState` may receive only a bounded projection in later integration work, never the sole recoverable copy. - Scheduled-task dispatch enforces "at most one active run per task when `overlap_policy=skip`" at the DB layer via the partial unique index `uq_scheduled_task_run_active` (`scheduled_task_runs.task_id WHERE status IN ('queued','running')`). `ScheduledTaskService.dispatch_task`'s `has_active_runs` check is a non-atomic fast path (its own session, separated from the `create()` insert by `await` points), so two concurrent dispatches — a manual `POST /scheduled-tasks/{id}/trigger` racing the poller, a double-click, or a client retry — can both pass it; the index is the atomic arbiter, and the losing `create` surfaces as `ActiveScheduledRunConflict` (translated from `IntegrityError` in the repository) and collapses to the same outcome as the fast path (manual → 409 conflict, scheduled → a `"skipped"` tombstone). The scheduled-skip tombstone is created directly as terminal `"skipped"` (not a transient `"queued"`) so it never occupies the active slot the pre-existing run still holds. Sibling of the `runs` table's `uq_runs_thread_active` (PR #4003), which keys on `thread_id` and so does not cover the default `fresh_thread_per_run` context where every dispatch gets a new thread. Index is status-only, not `overlap_policy`-conditional (the policy is fixed to `"skip"` in the MVP). @@ -427,7 +428,7 @@ Before changing a later authorization phase, read the [authorization RFC](../doc 17. **DurableContextMiddleware** - Captures `task` delegations into `ThreadState.delegations` (including in-progress dispatches and terminal result summaries) and loaded skill-file references (name/path/description, parsed in-memory - not the body) into `ThreadState.skill_context` before summarization can compact the paired tool-call/result messages, then projects durable context into each model request. Static authority rules are injected as a `SystemMessage`; untrusted field values (`summary_text`, delegation results, skill descriptions) are injected separately as a hidden `HumanMessage` data block so compressed history, delegated work, and which skills are active stay visible without being stored as `messages` or promoted to system-role instructions. `build_subagent_runtime_middlewares` also attaches this middleware immediately before subagent summarization so a compacted `summary_text` is projected ahead of a preserved assistant/tool tail instead of leaving strict providers with an assistant-first request. 18. **SummarizationMiddleware** - *(optional, if enabled)* Context reduction when approaching token limits 19. **TodoListMiddleware** - *(optional, if `is_plan_mode`)* Task tracking with the `write_todos` tool -20. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is merged back into the dispatching AIMessage by message position +20. **TokenUsageMiddleware** - *(optional, if `token_usage.enabled`)* Records token usage metrics; subagent usage is read from terminal `ToolMessage.additional_kwargs` in the current run and merged back into the dispatching AIMessage by message position. The same state update marks the ToolMessage with `subagent_token_usage_attributed=true`, so checkpoint replay or middleware re-entry cannot add the cumulative snapshot twice; missing/malformed usage or a result with no matching dispatch remains unmarked and retryable. 21. **TitleMiddleware** - Auto-generates the thread title after the first complete exchange and normalizes structured message content before prompting the title model. If a first-turn run is interrupted before this middleware can write a title, `runtime/runs/worker.py` keeps the run in a finalizing state, persists a local fallback title from the latest checkpoint or original run input, and then syncs it to `threads_meta.display_name`. Replacement runs admitted by `multitask_strategy="interrupt"` / `"rollback"` wait for older same-thread finalization before entering the graph; the interrupted run only skips the fallback title write once a later run has started and may have advanced the checkpoint. 22. **MemoryMiddleware** - Queues conversations for async memory update (filters to user + final AI responses); captures the runtime-resolved user so standalone LangGraph Server reads and writes stay in the same bucket 23. **ViewImageMiddleware** - *(optional, if the model supports vision)* Injects a hidden HumanMessage with base64 image data, identified by a reserved ID prefix plus a server-owned metadata marker, before the LLM call. Because `before_model`, `model`, and `after_model` are separate graph nodes, the `before_model` and `model` node checkpoints for that call still contain the payload; `after_model` / `aafter_model` then emits `RemoveMessage`, so subsequent checkpoints do not retain it diff --git a/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py index 2575043c2..946e5903a 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py @@ -12,9 +12,12 @@ from langchain.agents.middleware.todo import Todo from langchain_core.messages import AIMessage, ToolMessage from langgraph.runtime import Runtime +from deerflow.subagents.status_contract import SUBAGENT_TOKEN_USAGE_KEY, normalize_token_usage + logger = logging.getLogger(__name__) TOKEN_USAGE_ATTRIBUTION_KEY = "token_usage_attribution" +SUBAGENT_TOKEN_USAGE_ATTRIBUTED_KEY = "subagent_token_usage_attributed" def _string_arg(value: Any) -> str | None: @@ -228,6 +231,16 @@ def _has_tool_call(message: AIMessage, tool_call_id: str) -> bool: return False +def _subagent_usage_from_tool_message(message: ToolMessage) -> dict[str, int] | None: + """Read validated subagent usage from the current run's message state.""" + additional_kwargs = getattr(message, "additional_kwargs", None) + if not isinstance(additional_kwargs, dict): + return None + if additional_kwargs.get(SUBAGENT_TOKEN_USAGE_ATTRIBUTED_KEY) is True: + return None + return normalize_token_usage(additional_kwargs.get(SUBAGENT_TOKEN_USAGE_KEY)) + + def _build_attribution(message: AIMessage, todos: list[Todo]) -> dict[str, Any]: tool_calls = getattr(message, "tool_calls", None) or [] actions: list[dict[str, Any]] = [] @@ -273,22 +286,20 @@ class TokenUsageMiddleware(AgentMiddleware): return None # Annotate subagent token usage onto the AIMessage that dispatched it. - # When a task tool completes, its usage is cached by tool_call_id. Detect - # the ToolMessage → search backward for the corresponding AIMessage → merge. + # Terminal usage travels with the current run's ToolMessage, so provider + # tool-call IDs reused by another run cannot overwrite attribution state. # Walk backward through consecutive ToolMessages before the new AIMessage # so that multiple concurrent task tool calls all get their subagent tokens # written back to the same dispatch message (merging into one update). - state_updates: dict[int, AIMessage] = {} + state_updates: dict[int, AIMessage | ToolMessage] = {} if len(messages) >= 2: - from deerflow.tools.builtins.task_tool import pop_cached_subagent_usage - idx = len(messages) - 2 while idx >= 0: tool_msg = messages[idx] if not isinstance(tool_msg, ToolMessage) or not tool_msg.tool_call_id: break - subagent_usage = pop_cached_subagent_usage(tool_msg.tool_call_id) + subagent_usage = _subagent_usage_from_tool_message(tool_msg) if subagent_usage: # Search backward from the ToolMessage to find the AIMessage # that dispatched it. A single model response can dispatch @@ -301,7 +312,7 @@ class TokenUsageMiddleware(AgentMiddleware): # AIMessage (multiple task calls in one response), # or merge fresh from the original message. existing_update = state_updates.get(dispatch_idx) - prev = existing_update.usage_metadata if existing_update else (getattr(candidate, "usage_metadata", None) or {}) + prev = existing_update.usage_metadata if isinstance(existing_update, AIMessage) else (getattr(candidate, "usage_metadata", None) or {}) merged = { **prev, "input_tokens": prev.get("input_tokens", 0) + subagent_usage["input_tokens"], @@ -309,6 +320,9 @@ class TokenUsageMiddleware(AgentMiddleware): "total_tokens": prev.get("total_tokens", 0) + subagent_usage["total_tokens"], } state_updates[dispatch_idx] = candidate.model_copy(update={"usage_metadata": merged}) + tool_metadata = dict(tool_msg.additional_kwargs) + tool_metadata[SUBAGENT_TOKEN_USAGE_ATTRIBUTED_KEY] = True + state_updates[idx] = tool_msg.model_copy(update={"additional_kwargs": tool_metadata}) break dispatch_idx -= 1 idx -= 1 diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index 86e1b671b..45b88743e 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -80,7 +80,9 @@ class SubagentResult: """Result of a subagent execution. Attributes: - task_id: Unique identifier for this execution. + task_id: Server-generated identifier that owns this execution. + external_task_id: Optional provider correlation ID. This stays separate + because provider tool-call IDs can repeat across parent runs. trace_id: Trace ID for distributed tracing (links parent and subagent logs). status: Current status of the execution. result: The final result message (if completed). @@ -99,6 +101,7 @@ class SubagentResult: task_id: str trace_id: str status: SubagentStatus + external_task_id: str | None = field(default=None, kw_only=True) result: str | None = None error: str | None = None stop_reason: str | None = None @@ -823,7 +826,7 @@ class SubagentExecutor: task_store: ExtensionData | None = None task_info: TaskInfo | None = None if loaded_extensions.needs_task_store: - task_store = ExtensionData(result.task_id) + task_store = ExtensionData(result.external_task_id or result.task_id) if loaded_extensions.has_task_lifecycle and self.run_id: task_info = TaskInfo( task_id=result.task_id, @@ -1191,42 +1194,49 @@ class SubagentExecutor: Args: task: The task description for the subagent. - task_id: Optional task ID to use. If not provided, a random UUID will be generated. + task_id: Optional external correlation ID for logs. It is never used + as the process-wide background registry key because provider + tool-call IDs can repeat across concurrent parent runs. Returns: - Task ID that can be used to check status later. + Unique execution ID that can be used to check status later. """ - # Use provided task_id or generate a new one - if task_id is None: - task_id = str(uuid.uuid4())[:8] + execution_id = str(uuid.uuid4()) # Create initial pending result result = SubagentResult( - task_id=task_id, + task_id=execution_id, + external_task_id=task_id, trace_id=self.trace_id, status=SubagentStatus.PENDING, ) - logger.info(f"[trace={self.trace_id}] Subagent {self.config.name} starting async execution, task_id={task_id}, timeout={self.config.timeout_seconds}s") + logger.info( + "[trace=%s] Subagent %s starting async execution, execution_id=%s, external_task_id=%s, timeout=%ss", + self.trace_id, + self.config.name, + execution_id, + task_id, + self.config.timeout_seconds, + ) with _background_tasks_lock: - _background_tasks[task_id] = result + _background_tasks[execution_id] = result parent_context = _copy_isolated_subagent_context() # Submit to scheduler pool def run_task(): with _background_tasks_lock: - _background_tasks[task_id].status = SubagentStatus.RUNNING - _background_tasks[task_id].started_at = datetime.now() - result_holder = _background_tasks[task_id] + result.status = SubagentStatus.RUNNING + result.started_at = datetime.now() try: # Submit execution directly to the persistent isolated loop so the # background path does not create a temporary loop via execute(). execution_future = _submit_to_isolated_loop_in_context( parent_context, - lambda: self._aexecute(task, result_holder), + lambda: self._aexecute(task, result), ) try: # Wait for execution with timeout @@ -1234,26 +1244,24 @@ class SubagentExecutor: except FuturesTimeoutError: logger.error(f"[trace={self.trace_id}] Subagent {self.config.name} execution timed out after {self.config.timeout_seconds}s") # Signal cooperative cancellation and cancel the future - result_holder.cancel_event.set() - result_holder.try_set_terminal( + result.cancel_event.set() + result.try_set_terminal( SubagentStatus.TIMED_OUT, error=f"Execution timed out after {self.config.timeout_seconds} seconds", ) execution_future.cancel() except Exception as e: logger.exception(f"[trace={self.trace_id}] Subagent {self.config.name} async execution failed") - with _background_tasks_lock: - task_result = _background_tasks[task_id] - task_result.try_set_terminal(SubagentStatus.FAILED, error=str(e)) + result.try_set_terminal(SubagentStatus.FAILED, error=str(e)) _scheduler_pool.submit(run_task) - return task_id + return execution_id MAX_CONCURRENT_SUBAGENTS = 3 -def request_cancel_background_task(task_id: str) -> None: +def request_cancel_background_task(execution_id: str) -> None: """Signal a running background task to stop. Sets the cancel_event on the task, which is checked cooperatively @@ -1262,26 +1270,26 @@ def request_cancel_background_task(task_id: str) -> None: — to stop at the next iteration boundary. Args: - task_id: The task ID to cancel. + execution_id: The execution ID returned by execute_async. """ with _background_tasks_lock: - result = _background_tasks.get(task_id) + result = _background_tasks.get(execution_id) if result is not None: result.cancel_event.set() - logger.info("Requested cancellation for background task %s", task_id) + logger.info("Requested cancellation for background execution %s", execution_id) -def get_background_task_result(task_id: str) -> SubagentResult | None: +def get_background_task_result(execution_id: str) -> SubagentResult | None: """Get the result of a background task. Args: - task_id: The task ID returned by execute_async. + execution_id: The execution ID returned by execute_async. Returns: SubagentResult if found, None otherwise. """ with _background_tasks_lock: - return _background_tasks.get(task_id) + return _background_tasks.get(execution_id) def list_background_tasks() -> list[SubagentResult]: @@ -1294,7 +1302,7 @@ def list_background_tasks() -> list[SubagentResult]: return list(_background_tasks.values()) -def cleanup_background_task(task_id: str) -> None: +def cleanup_background_task(execution_id: str) -> None: """Remove a completed task from background tasks. Should be called by task_tool after it finishes polling and returns the result. @@ -1304,23 +1312,23 @@ def cleanup_background_task(task_id: str) -> None: to avoid race conditions with the background executor still updating the task entry. Args: - task_id: The task ID to remove. + execution_id: The execution ID to remove. """ with _background_tasks_lock: - result = _background_tasks.get(task_id) + result = _background_tasks.get(execution_id) if result is None: # Nothing to clean up; may have been removed already. - logger.debug("Requested cleanup for unknown background task %s", task_id) + logger.debug("Requested cleanup for unknown background execution %s", execution_id) return # Only clean up tasks that are in a terminal state to avoid races with # the background executor still updating the task entry. if result.status.is_terminal or result.completed_at is not None: - del _background_tasks[task_id] - logger.debug("Cleaned up background task: %s", task_id) + del _background_tasks[execution_id] + logger.debug("Cleaned up background execution: %s", execution_id) else: logger.debug( - "Skipping cleanup for non-terminal background task %s (status=%s)", - task_id, + "Skipping cleanup for non-terminal background execution %s (status=%s)", + execution_id, result.status.value if hasattr(result.status, "value") else result.status, ) diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index b7a9cd2c9..a7136d2d5 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -40,38 +40,16 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -# Cache subagent token usage by tool_call_id so TokenUsageMiddleware can -# write it back to the triggering AIMessage's usage_metadata. -_subagent_usage_cache: dict[str, dict[str, int]] = {} - - -def _token_usage_cache_enabled(app_config: "AppConfig | None") -> bool: - if app_config is None: - try: - app_config = get_app_config() - except FileNotFoundError: - return False - return bool(getattr(getattr(app_config, "token_usage", None), "enabled", False)) - - -def _cache_subagent_usage(tool_call_id: str, usage: dict | None, *, enabled: bool = True) -> None: - if enabled and usage: - _subagent_usage_cache[tool_call_id] = usage - - -def pop_cached_subagent_usage(tool_call_id: str) -> dict | None: - return _subagent_usage_cache.pop(tool_call_id, None) - def _is_subagent_terminal(result: Any) -> bool: """Return whether a background subagent result is safe to clean up.""" return result.status in {SubagentStatus.COMPLETED, SubagentStatus.FAILED, SubagentStatus.CANCELLED, SubagentStatus.TIMED_OUT} or getattr(result, "completed_at", None) is not None -async def _await_subagent_terminal(task_id: str, max_polls: int) -> Any | None: +async def _await_subagent_terminal(execution_id: str, max_polls: int) -> Any | None: """Poll until the background subagent reaches a terminal status or we run out of polls.""" for _ in range(max_polls): - result = get_background_task_result(task_id) + result = get_background_task_result(execution_id) if result is None: return None if _is_subagent_terminal(result): @@ -80,36 +58,36 @@ async def _await_subagent_terminal(task_id: str, max_polls: int) -> Any | None: return None -async def _deferred_cleanup_subagent_task(task_id: str, trace_id: str, max_polls: int) -> None: +async def _deferred_cleanup_subagent_task(execution_id: str, trace_id: str, max_polls: int) -> None: """Keep polling a cancelled subagent until it can be safely removed.""" cleanup_poll_count = 0 while True: - result = get_background_task_result(task_id) + result = get_background_task_result(execution_id) if result is None: return if _is_subagent_terminal(result): - cleanup_background_task(task_id) + cleanup_background_task(execution_id) return if cleanup_poll_count >= max_polls: - logger.warning(f"[trace={trace_id}] Deferred cleanup for task {task_id} timed out after {cleanup_poll_count} polls") + logger.warning(f"[trace={trace_id}] Deferred cleanup for execution {execution_id} timed out after {cleanup_poll_count} polls") return await asyncio.sleep(5) cleanup_poll_count += 1 -def _log_cleanup_failure(cleanup_task: asyncio.Task[None], *, trace_id: str, task_id: str) -> None: +def _log_cleanup_failure(cleanup_task: asyncio.Task[None], *, trace_id: str, execution_id: str) -> None: if cleanup_task.cancelled(): return exc = cleanup_task.exception() if exc is not None: - logger.error(f"[trace={trace_id}] Deferred cleanup failed for task {task_id}: {exc}") + logger.error(f"[trace={trace_id}] Deferred cleanup failed for execution {execution_id}: {exc}") -def _schedule_deferred_subagent_cleanup(task_id: str, trace_id: str, max_polls: int) -> None: - logger.debug(f"[trace={trace_id}] Scheduling deferred cleanup for cancelled task {task_id}") - cleanup_task = asyncio.create_task(_deferred_cleanup_subagent_task(task_id, trace_id, max_polls)) - cleanup_task.add_done_callback(lambda task: _log_cleanup_failure(task, trace_id=trace_id, task_id=task_id)) +def _schedule_deferred_subagent_cleanup(execution_id: str, trace_id: str, max_polls: int) -> None: + logger.debug(f"[trace={trace_id}] Scheduling deferred cleanup for cancelled execution {execution_id}") + cleanup_task = asyncio.create_task(_deferred_cleanup_subagent_task(execution_id, trace_id, max_polls)) + cleanup_task.add_done_callback(lambda task: _log_cleanup_failure(task, trace_id=trace_id, execution_id=execution_id)) def _find_usage_recorder(runtime: Any) -> Any | None: @@ -284,7 +262,6 @@ async def task_tool( subagent_type: The type of subagent to use. ALWAYS PROVIDE THIS PARAMETER THIRD. """ runtime_app_config = _get_runtime_app_config(runtime) - cache_token_usage = _token_usage_cache_enabled(runtime_app_config) available_subagent_names = get_available_subagent_names(app_config=runtime_app_config) if runtime_app_config is not None else get_available_subagent_names() # Get subagent configuration @@ -424,9 +401,9 @@ async def task_tool( executor_kwargs["extensions"] = run_extensions executor = SubagentExecutor(**executor_kwargs) - # Start background execution (always async to prevent blocking) - # Use tool_call_id as task_id for better traceability - task_id = executor.execute_async(prompt, task_id=tool_call_id) + # Keep the provider tool-call ID for stream/message correlation, but use a + # server-generated execution ID for process-wide background task control. + execution_id = executor.execute_async(prompt, task_id=tool_call_id) # Poll for task completion in backend (removes need for LLM to poll) poll_count = 0 @@ -435,14 +412,14 @@ async def task_tool( # Polling timeout: execution timeout + 60s buffer, checked every 5s max_poll_count = (config.timeout_seconds + 60) // 5 - logger.info(f"[trace={trace_id}] Started background task {task_id} (subagent={subagent_type}, timeout={config.timeout_seconds}s, polling_limit={max_poll_count} polls)") + logger.info(f"[trace={trace_id}] Started background task {tool_call_id} (execution_id={execution_id}, subagent={subagent_type}, timeout={config.timeout_seconds}s, polling_limit={max_poll_count} polls)") writer = get_stream_writer() # Send Task Started message' await aemit_custom_event( { "type": "task_started", - "task_id": task_id, + "task_id": tool_call_id, "description": description, "model_name": effective_model, }, @@ -451,16 +428,16 @@ async def task_tool( try: while True: - result = get_background_task_result(task_id) + result = get_background_task_result(execution_id) if result is None: - logger.error(f"[trace={trace_id}] Task {task_id} not found in background tasks") + logger.error(f"[trace={trace_id}] Task {tool_call_id} execution {execution_id} not found in background tasks") await aemit_custom_event( - {"type": "task_failed", "task_id": task_id, "error": "Task disappeared from background tasks"}, + {"type": "task_failed", "task_id": tool_call_id, "error": "Task disappeared from background tasks"}, writer=writer, ) - cleanup_background_task(task_id) - error = f"Task {task_id} disappeared from background tasks" + cleanup_background_task(execution_id) + error = f"Task {tool_call_id} disappeared from background tasks" return _task_result_command( tool_call_id=tool_call_id, status="failed", @@ -469,7 +446,7 @@ async def task_tool( # Log status changes for debugging if result.status != last_status: - logger.info(f"[trace={trace_id}] Task {task_id} status: {result.status.value}") + logger.info(f"[trace={trace_id}] Task {tool_call_id} execution {execution_id} status: {result.status.value}") last_status = result.status # The collector publishes cumulative records. Reuse one snapshot for @@ -487,7 +464,7 @@ async def task_tool( await aemit_custom_event( { "type": "task_running", - "task_id": task_id, + "task_id": tool_call_id, "message": message, "message_index": i + 1, # 1-based index for display "total_messages": current_message_count, @@ -496,25 +473,24 @@ async def task_tool( }, writer=writer, ) - logger.info(f"[trace={trace_id}] Task {task_id} sent message #{i + 1}/{current_message_count}") + logger.info(f"[trace={trace_id}] Task {tool_call_id} sent message #{i + 1}/{current_message_count}") last_message_count = current_message_count # Check if task completed, failed, or timed out if result.status == SubagentStatus.COMPLETED: - _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) await aemit_custom_event( { "type": "task_completed", - "task_id": task_id, + "task_id": tool_call_id, "result": result.result, "usage": usage, "model_name": effective_model, }, writer=writer, ) - logger.info(f"[trace={trace_id}] Task {task_id} completed after {poll_count} polls") - cleanup_background_task(task_id) + logger.info(f"[trace={trace_id}] Task {tool_call_id} completed after {poll_count} polls") + cleanup_background_task(execution_id) # stop_reason carries a guardrail cap (token_capped / turn_capped) # when the run was ended early but still produced a final answer # — the work survives on result_brief like a clean success. @@ -527,20 +503,19 @@ async def task_tool( usage=usage, ) elif result.status == SubagentStatus.FAILED: - _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) await aemit_custom_event( { "type": "task_failed", - "task_id": task_id, + "task_id": tool_call_id, "error": result.error, "usage": usage, "model_name": effective_model, }, writer=writer, ) - logger.error(f"[trace={trace_id}] Task {task_id} failed: {result.error}") - cleanup_background_task(task_id) + logger.error(f"[trace={trace_id}] Task {tool_call_id} failed: {result.error}") + cleanup_background_task(execution_id) # A turn-capped run with no usable output surfaces as failed + # stop_reason=turn_capped; the cap note lets the lead tell "out # of budget" from "broken subagent". @@ -553,20 +528,19 @@ async def task_tool( usage=usage, ) elif result.status == SubagentStatus.CANCELLED: - _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) await aemit_custom_event( { "type": "task_cancelled", - "task_id": task_id, + "task_id": tool_call_id, "error": result.error, "usage": usage, "model_name": effective_model, }, writer=writer, ) - logger.info(f"[trace={trace_id}] Task {task_id} cancelled: {result.error}") - cleanup_background_task(task_id) + logger.info(f"[trace={trace_id}] Task {tool_call_id} cancelled: {result.error}") + cleanup_background_task(execution_id) return _task_result_command( tool_call_id=tool_call_id, status="cancelled", @@ -575,20 +549,19 @@ async def task_tool( usage=usage, ) elif result.status == SubagentStatus.TIMED_OUT: - _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) _report_subagent_usage(runtime, result) await aemit_custom_event( { "type": "task_timed_out", - "task_id": task_id, + "task_id": tool_call_id, "error": result.error, "usage": usage, "model_name": effective_model, }, writer=writer, ) - logger.warning(f"[trace={trace_id}] Task {task_id} timed out: {result.error}") - cleanup_background_task(task_id) + logger.warning(f"[trace={trace_id}] Task {tool_call_id} timed out: {result.error}") + cleanup_background_task(execution_id) return _task_result_command( tool_call_id=tool_call_id, status="timed_out", @@ -606,14 +579,13 @@ async def task_tool( # This catches edge cases where the background task gets stuck if poll_count > max_poll_count: timeout_minutes = config.timeout_seconds // 60 - logger.error(f"[trace={trace_id}] Task {task_id} polling timed out after {poll_count} polls (should have been caught by thread pool timeout)") + logger.error(f"[trace={trace_id}] Task {tool_call_id} polling timed out after {poll_count} polls (should have been caught by thread pool timeout)") _report_subagent_usage(runtime, result) usage = _summarize_usage(getattr(result, "token_usage_records", None)) - _cache_subagent_usage(tool_call_id, usage, enabled=cache_token_usage) await aemit_custom_event( { "type": "task_timed_out", - "task_id": task_id, + "task_id": tool_call_id, "usage": usage, "model_name": effective_model, }, @@ -622,8 +594,8 @@ async def task_tool( # The task may still be running in the background. Signal cooperative # cancellation and schedule deferred cleanup to remove the entry from # _background_tasks once the background thread reaches a terminal state. - request_cancel_background_task(task_id) - _schedule_deferred_subagent_cleanup(task_id, trace_id, max_poll_count) + request_cancel_background_task(execution_id) + _schedule_deferred_subagent_cleanup(execution_id, trace_id, max_poll_count) message = f"Task polling timed out after {timeout_minutes} minutes. This may indicate the background task is stuck. Status: {result.status.value}" return _task_result_command( tool_call_id=tool_call_id, @@ -634,27 +606,23 @@ async def task_tool( ) except asyncio.CancelledError: # Signal the background subagent thread to stop cooperatively. - request_cancel_background_task(task_id) + request_cancel_background_task(execution_id) # Wait (shielded) for the subagent to reach a terminal state so the # final token usage snapshot is reported to the parent RunJournal # before the parent worker persists get_completion_data(). terminal_result = None try: - terminal_result = await asyncio.shield(_await_subagent_terminal(task_id, max_poll_count)) + terminal_result = await asyncio.shield(_await_subagent_terminal(execution_id, max_poll_count)) except asyncio.CancelledError: pass # Report whatever the subagent collected (even if we timed out). - final_result = terminal_result or get_background_task_result(task_id) + final_result = terminal_result or get_background_task_result(execution_id) if final_result is not None: _report_subagent_usage(runtime, final_result) if final_result is not None and _is_subagent_terminal(final_result): - cleanup_background_task(task_id) + cleanup_background_task(execution_id) else: - _schedule_deferred_subagent_cleanup(task_id, trace_id, max_poll_count) - _subagent_usage_cache.pop(tool_call_id, None) - raise - except Exception: - _subagent_usage_cache.pop(tool_call_id, None) + _schedule_deferred_subagent_cleanup(execution_id, trace_id, max_poll_count) raise diff --git a/backend/tests/test_custom_events.py b/backend/tests/test_custom_events.py index 84c977e8b..358407140 100644 --- a/backend/tests/test_custom_events.py +++ b/backend/tests/test_custom_events.py @@ -149,7 +149,6 @@ async def test_real_task_tool_events_reach_astream_events(monkeypatch): monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _name: config) monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _task_id: completed) monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda _task_id: None) - monkeypatch.setattr(task_tool_module, "_token_usage_cache_enabled", lambda _config: False) monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **_kwargs: []) agent = create_agent( diff --git a/backend/tests/test_extension_task_store_runtime.py b/backend/tests/test_extension_task_store_runtime.py index 568aadd1f..a022b912b 100644 --- a/backend/tests/test_extension_task_store_runtime.py +++ b/backend/tests/test_extension_task_store_runtime.py @@ -161,14 +161,19 @@ def _subagent_env(): sys.modules["deerflow.skills.storage"] = storage_module from deerflow.subagents.config import SubagentConfig - from deerflow.subagents.executor import SubagentExecutor + from deerflow.subagents.executor import SubagentExecutor, SubagentResult, SubagentStatus executor_module = sys.modules["deerflow.subagents.executor"] executor_module.get_app_config = lambda: SimpleNamespace( tool_search=SimpleNamespace(enabled=False), authorization=SimpleNamespace(enabled=False), ) - yield SimpleNamespace(SubagentConfig=SubagentConfig, SubagentExecutor=SubagentExecutor) + yield SimpleNamespace( + SubagentConfig=SubagentConfig, + SubagentExecutor=SubagentExecutor, + SubagentResult=SubagentResult, + SubagentStatus=SubagentStatus, + ) finally: for name, original in original_modules.items(): if original is None: @@ -197,7 +202,7 @@ class _CapturingSubagent: yield {"messages": [AIMessage(content="done")]} -async def _run_subagent(monkeypatch, env, *, seen: dict): +async def _run_subagent(monkeypatch, env, *, seen: dict, result_holder=None): async def _initial_state(self, task): return ({}, [], None) @@ -205,7 +210,7 @@ async def _run_subagent(monkeypatch, env, *, seen: dict): monkeypatch.setattr(env.SubagentExecutor, "_create_agent", lambda self, tools, **kwargs: _CapturingSubagent(seen)) config = env.SubagentConfig(name="researcher", description="d", system_prompt="p", tools=[]) executor = env.SubagentExecutor(config=config, tools=[], thread_id="thread-1", run_id=None) - return await executor._aexecute("do the thing") + return await executor._aexecute("do the thing", result_holder=result_holder) @pytest.mark.anyio @@ -305,6 +310,33 @@ async def test_subagent_middleware_without_parent_run_receives_its_task_store(mo assert store.scope_id == result.task_id +@pytest.mark.anyio +async def test_async_subagent_task_store_preserves_external_correlation_scope(monkeypatch, _isolated_extensions, _subagent_env): + registry = ExtensionRegistry() + with registry.attributed_to("demo:install"): + registry.middlewares(_MiddlewareContributor()) + set_loaded_extensions(registry.build()) + seen: dict = {} + result_holder = _subagent_env.SubagentResult( + task_id="server-execution-id", + external_task_id="provider-tool-call-id", + trace_id="trace-1", + status=_subagent_env.SubagentStatus.RUNNING, + ) + + result = await _run_subagent( + monkeypatch, + _subagent_env, + seen=seen, + result_holder=result_holder, + ) + + store = (seen.get("context") or {}).get(EXTENSION_TASK_STORE_KEY) + assert isinstance(store, ExtensionData) + assert result.task_id == "server-execution-id" + assert store.scope_id == "provider-tool-call-id" + + @pytest.mark.anyio async def test_subagent_builder_receives_the_same_extension_snapshot_as_the_store(monkeypatch, _isolated_extensions, _subagent_env): registry = ExtensionRegistry() diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index 5881d3dad..153a9017c 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -2344,6 +2344,81 @@ class TestCooperativeCancellation: assert result.result == "done: Task" assert result.error is None + def test_execute_async_isolates_duplicate_external_task_ids(self, executor_module, classes, base_config): + """Concurrent runs must not share registry entries when provider IDs collide.""" + import concurrent.futures + + SubagentExecutor = classes["SubagentExecutor"] + SubagentStatus = classes["SubagentStatus"] + scheduled: list = [] + + def capture_submission(fn, *args, **kwargs): + scheduled.append(lambda: fn(*args, **kwargs)) + return concurrent.futures.Future() + + def run_coroutine(_context, coroutine_factory): + future = concurrent.futures.Future() + try: + future.set_result(asyncio.run(coroutine_factory())) + except Exception as exc: + future.set_exception(exc) + return future + + async def complete_a(_task, result_holder=None): + result_holder.status = SubagentStatus.COMPLETED + result_holder.result = "done-a" + result_holder.completed_at = datetime.now() + return result_holder + + async def complete_b(_task, result_holder=None): + result_holder.status = SubagentStatus.COMPLETED + result_holder.result = "done-b" + result_holder.completed_at = datetime.now() + return result_holder + + executor_a = SubagentExecutor( + config=base_config, + tools=[], + thread_id="thread-a", + trace_id="trace-a", + ) + executor_b = SubagentExecutor( + config=base_config, + tools=[], + thread_id="thread-b", + trace_id="trace-b", + ) + + with ( + patch.object(executor_module._scheduler_pool, "submit", side_effect=capture_submission), + patch.object(executor_module, "_submit_to_isolated_loop_in_context", side_effect=run_coroutine), + patch.object(executor_a, "_aexecute", side_effect=complete_a), + patch.object(executor_b, "_aexecute", side_effect=complete_b), + ): + execution_a = executor_a.execute_async("Task A", task_id="same-provider-tool-call-id") + execution_b = executor_b.execute_async("Task B", task_id="same-provider-tool-call-id") + + assert execution_a != execution_b + assert executor_module._background_tasks[execution_a].trace_id == "trace-a" + assert executor_module._background_tasks[execution_b].trace_id == "trace-b" + + for run_task in scheduled: + run_task() + + assert executor_module._background_tasks[execution_a].result == "done-a" + assert executor_module._background_tasks[execution_b].result == "done-b" + + result_a = executor_module._background_tasks[execution_a] + result_b = executor_module._background_tasks[execution_b] + executor_module.request_cancel_background_task(execution_a) + assert result_a.cancel_event.is_set() + assert not result_b.cancel_event.is_set() + + executor_module.cleanup_background_task(execution_a) + assert execution_a not in executor_module._background_tasks + assert executor_module._background_tasks[execution_b] is result_b + executor_module.cleanup_background_task(execution_b) + def test_execute_async_propagates_user_context_to_isolated_loop(self, executor_module, classes, base_config): """Regression: background subagent execution must keep request user context.""" import concurrent.futures diff --git a/backend/tests/test_task_tool_core_logic.py b/backend/tests/test_task_tool_core_logic.py index a6535c6ae..a220d1c11 100644 --- a/backend/tests/test_task_tool_core_logic.py +++ b/backend/tests/test_task_tool_core_logic.py @@ -568,6 +568,8 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): events = [] dispatched_events = [] captured = {} + polled_execution_ids = [] + cleaned_execution_ids = [] get_available_tools = MagicMock(return_value=["tool-a", "tool-b"]) async def fake_emit_custom_event(payload, *, writer): @@ -581,7 +583,7 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): def execute_async(self, prompt, task_id=None): captured["prompt"] = prompt captured["task_id"] = task_id - return task_id or "generated-task-id" + return "execution-456" # Simulate two polling rounds: first running (with one message), then completed. responses = iter( @@ -599,7 +601,12 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): monkeypatch.setattr(task_tool_module, "SubagentExecutor", DummyExecutor) monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) - monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: next(responses)) + def get_result(execution_id): + polled_execution_ids.append(execution_id) + return next(responses) + + monkeypatch.setattr(task_tool_module, "get_background_task_result", get_result) + monkeypatch.setattr(task_tool_module, "cleanup_background_task", cleaned_execution_ids.append) monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: events.append) monkeypatch.setattr(task_tool_module, "aemit_custom_event", fake_emit_custom_event) monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) @@ -631,6 +638,9 @@ def test_task_tool_emits_running_and_completed_events(monkeypatch): event_types = [e["type"] for e in events] assert event_types == ["task_started", "task_running", "task_running", "task_completed"] assert dispatched_events == events + assert polled_execution_ids == ["execution-456", "execution-456"] + assert cleaned_execution_ids == ["execution-456"] + assert {event["task_id"] for event in events} == {"tc-123"} assert events[0]["model_name"] == "ark-model" assert events[-1]["result"] == "all done" @@ -1491,7 +1501,7 @@ def test_cancellation_wait_uses_subagent_polling_budget(monkeypatch): def test_cancellation_calls_request_cancel(monkeypatch): - """Verify CancelledError path calls request_cancel_background_task(task_id).""" + """Verify CancelledError path cancels the server-generated execution ID.""" config = _make_subagent_config() events = [] cancel_requests = [] @@ -1503,7 +1513,7 @@ def test_cancellation_calls_request_cancel(monkeypatch): monkeypatch.setattr( task_tool_module, "SubagentExecutor", - type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: task_id}), + type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: "execution-cancel"}), ) monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _: config) @@ -1535,7 +1545,7 @@ def test_cancellation_calls_request_cancel(monkeypatch): tool_call_id="tc-cancel-request", ) - assert cancel_requests == ["tc-cancel-request"] + assert cancel_requests == ["execution-cancel"] def test_task_tool_returns_cancelled_message(monkeypatch): @@ -1843,76 +1853,3 @@ def test_terminal_event_usage_none_when_no_records(monkeypatch): completed = [e for e in events if e["type"] == "task_completed"] assert len(completed) == 1 assert completed[0]["usage"] is None - - -def test_subagent_usage_cache_is_skipped_when_config_file_is_missing(monkeypatch): - monkeypatch.setattr( - task_tool_module, - "get_app_config", - MagicMock(side_effect=FileNotFoundError("missing config")), - ) - - assert task_tool_module._token_usage_cache_enabled(None) is False - - -def test_subagent_usage_cache_is_skipped_when_token_usage_is_disabled(monkeypatch): - config = _make_subagent_config() - app_config = SimpleNamespace(token_usage=SimpleNamespace(enabled=False)) - runtime = _make_runtime(app_config=app_config) - records = [{"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}] - result = _make_result(FakeSubagentStatus.COMPLETED, result="done", token_usage_records=records) - - task_tool_module._subagent_usage_cache.clear() - monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) - monkeypatch.setattr(task_tool_module, "get_available_subagent_names", lambda *, app_config: ["general-purpose"]) - monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _, *, app_config: config) - monkeypatch.setattr( - task_tool_module, - "SubagentExecutor", - type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: task_id}), - ) - monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _: result) - monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _: None) - monkeypatch.setattr(task_tool_module, "_report_subagent_usage", lambda *_: None) - monkeypatch.setattr(task_tool_module, "cleanup_background_task", lambda _: None) - monkeypatch.setattr("deerflow.tools.get_available_tools", MagicMock(return_value=[])) - - _run_task_tool( - runtime=runtime, - description="test", - prompt="do work", - subagent_type="general-purpose", - tool_call_id="tc-disabled-cache", - ) - - assert task_tool_module.pop_cached_subagent_usage("tc-disabled-cache") is None - - -def test_subagent_usage_cache_is_cleared_when_polling_raises(monkeypatch): - config = _make_subagent_config() - app_config = SimpleNamespace(token_usage=SimpleNamespace(enabled=True)) - runtime = _make_runtime(app_config=app_config) - - task_tool_module._subagent_usage_cache["tc-error"] = {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2} - monkeypatch.setattr(task_tool_module, "SubagentStatus", FakeSubagentStatus) - monkeypatch.setattr(task_tool_module, "get_available_subagent_names", lambda *, app_config: ["general-purpose"]) - monkeypatch.setattr(task_tool_module, "get_subagent_config", lambda _, *, app_config: config) - monkeypatch.setattr( - task_tool_module, - "SubagentExecutor", - type("DummyExecutor", (), {"__init__": lambda self, **kwargs: None, "execute_async": lambda self, prompt, task_id=None: task_id}), - ) - monkeypatch.setattr(task_tool_module, "get_background_task_result", MagicMock(side_effect=RuntimeError("poll failed"))) - monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _: None) - monkeypatch.setattr("deerflow.tools.get_available_tools", MagicMock(return_value=[])) - - with pytest.raises(RuntimeError, match="poll failed"): - _run_task_tool( - runtime=runtime, - description="test", - prompt="do work", - subagent_type="general-purpose", - tool_call_id="tc-error", - ) - - assert task_tool_module.pop_cached_subagent_usage("tc-error") is None diff --git a/backend/tests/test_token_usage_middleware.py b/backend/tests/test_token_usage_middleware.py index fcfbe27e6..545c57c5f 100644 --- a/backend/tests/test_token_usage_middleware.py +++ b/backend/tests/test_token_usage_middleware.py @@ -1,16 +1,17 @@ """Tests for TokenUsageMiddleware attribution annotations.""" -import importlib import logging from unittest.mock import MagicMock from langchain_core.messages import AIMessage, ToolMessage +from langgraph.graph.message import add_messages from deerflow.agents.middlewares.token_usage_middleware import ( TOKEN_USAGE_ATTRIBUTION_KEY, TokenUsageMiddleware, _build_todo_actions, ) +from deerflow.subagents.status_contract import SUBAGENT_TOKEN_USAGE_KEY def _make_runtime(): @@ -235,7 +236,7 @@ class TestTokenUsageMiddleware: } ] - def test_merges_subagent_usage_by_message_position_when_ai_message_ids_are_missing(self, monkeypatch): + def test_merges_subagent_usage_by_message_position_when_ai_message_ids_are_missing(self): middleware = TokenUsageMiddleware() first_dispatch = AIMessage( content="", @@ -252,21 +253,18 @@ class TestTokenUsageMiddleware: first_dispatch, ToolMessage(content="first", tool_call_id="task:first"), second_dispatch, - ToolMessage(content="second-a", tool_call_id="task:second-a"), - ToolMessage(content="second-b", tool_call_id="task:second-b"), + ToolMessage( + content="second-a", + tool_call_id="task:second-a", + additional_kwargs={SUBAGENT_TOKEN_USAGE_KEY: {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}}, + ), + ToolMessage( + content="second-b", + tool_call_id="task:second-b", + additional_kwargs={SUBAGENT_TOKEN_USAGE_KEY: {"input_tokens": 20, "output_tokens": 7, "total_tokens": 27}}, + ), AIMessage(content="done"), ] - cached_usage = { - "task:second-a": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, - "task:second-b": {"input_tokens": 20, "output_tokens": 7, "total_tokens": 27}, - } - - task_tool_module = importlib.import_module("deerflow.tools.builtins.task_tool") - monkeypatch.setattr( - task_tool_module, - "pop_cached_subagent_usage", - lambda tool_call_id: cached_usage.pop(tool_call_id, None), - ) result = middleware.after_model({"messages": messages}, _make_runtime()) @@ -281,6 +279,78 @@ class TestTokenUsageMiddleware: "total_tokens": 42, } + def test_reused_tool_call_id_keeps_usage_scoped_to_each_run_history(self): + middleware = TokenUsageMiddleware() + tool_call_id = "reused-provider-tool-call-id" + + def apply_usage(usage): + dispatch = AIMessage( + content="", + tool_calls=[{"id": tool_call_id, "name": "task", "args": {}}], + ) + messages = [ + dispatch, + ToolMessage( + content="task result", + tool_call_id=tool_call_id, + additional_kwargs={SUBAGENT_TOKEN_USAGE_KEY: usage}, + ), + AIMessage(content="done"), + ] + + result = middleware.after_model({"messages": messages}, _make_runtime()) + + assert result is not None + return next(message.usage_metadata for message in result["messages"] if getattr(message, "usage_metadata", None)) + + assert apply_usage({"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}) == { + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + } + assert apply_usage({"input_tokens": 90, "output_tokens": 8, "total_tokens": 98}) == { + "input_tokens": 90, + "output_tokens": 8, + "total_tokens": 98, + } + + def test_subagent_usage_attribution_is_idempotent_when_state_is_reprocessed(self): + middleware = TokenUsageMiddleware() + dispatch = AIMessage( + id="dispatch-message", + content="", + tool_calls=[{"id": "task:replayed", "name": "task", "args": {}}], + ) + tool_result = ToolMessage( + id="tool-message", + content="task result", + tool_call_id="task:replayed", + additional_kwargs={ + SUBAGENT_TOKEN_USAGE_KEY: { + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + } + }, + ) + final = AIMessage(id="final-message", content="done") + messages = [dispatch, tool_result, final] + + first_update = middleware.after_model({"messages": messages}, _make_runtime()) + + assert first_update is not None + checkpoint_messages = add_messages(messages, first_update["messages"]) + updated_dispatch = next(message for message in checkpoint_messages if message.id == dispatch.id) + assert updated_dispatch.usage_metadata == { + "input_tokens": 10, + "output_tokens": 2, + "total_tokens": 12, + } + + second_update = middleware.after_model({"messages": checkpoint_messages}, _make_runtime()) + + assert second_update is None + class TestBuildTodoActions: def test_duplicate_content_emits_todo_remove(self):