diff --git a/backend/app/gateway/routers/runs.py b/backend/app/gateway/routers/runs.py index 846913b2f..c16aa66a4 100644 --- a/backend/app/gateway/routers/runs.py +++ b/backend/app/gateway/routers/runs.py @@ -16,7 +16,7 @@ from app.gateway.authz import require_permission from app.gateway.deps import get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge from app.gateway.pagination import trim_run_message_page from app.gateway.run_models import RunCreateRequest -from app.gateway.services import build_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion +from app.gateway.services import abuild_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion from deerflow.runtime import serialize_channel_values_for_api from deerflow.utils.thread_id import resolve_thread_id @@ -76,7 +76,7 @@ async def stateless_wait(body: RunCreateRequest, request: Request) -> dict: if completed: try: - accessor, config = build_checkpoint_state_accessor( + accessor, config = await abuild_checkpoint_state_accessor( request, thread_id=thread_id, assistant_id=body.assistant_id, diff --git a/backend/app/gateway/routers/thread_runs.py b/backend/app/gateway/routers/thread_runs.py index 399b9e6fb..551c7a81f 100644 --- a/backend/app/gateway/routers/thread_runs.py +++ b/backend/app/gateway/routers/thread_runs.py @@ -42,7 +42,7 @@ from app.gateway.deps import get_current_user, get_feedback_repo, get_run_event_ from app.gateway.internal_auth import get_trusted_internal_owner_user_id from app.gateway.pagination import trim_run_message_page from app.gateway.run_models import RunCreateRequest -from app.gateway.services import build_checkpoint_state_accessor, build_thread_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion +from app.gateway.services import abuild_checkpoint_state_accessor, build_thread_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion from app.gateway.utils import sanitize_log_param from deerflow.agents.middlewares.dynamic_context_middleware import strip_injected_user_message_id_suffix from deerflow.authz.sandbox_authz import safe_app_config_async @@ -1038,7 +1038,7 @@ async def wait_run( # thread head may be a later run, so do not claim it as this run's result. if completed and not reused: try: - accessor, config = build_checkpoint_state_accessor( + accessor, config = await abuild_checkpoint_state_accessor( request, thread_id=thread_id, assistant_id=body.assistant_id, diff --git a/backend/app/gateway/routers/threads.py b/backend/app/gateway/routers/threads.py index 6441a3739..28b75ff50 100644 --- a/backend/app/gateway/routers/threads.py +++ b/backend/app/gateway/routers/threads.py @@ -35,7 +35,7 @@ from app.gateway.checkpoint_lineage import ( from app.gateway.deps import get_checkpointer, get_run_event_store, get_run_manager from app.gateway.internal_auth import get_trusted_internal_owner_user_id from app.gateway.services import ( - build_checkpoint_state_accessor, + abuild_checkpoint_state_accessor, build_checkpoint_state_mutation_accessor, build_thread_checkpoint_state_accessor, build_thread_checkpoint_state_mutation_accessor, @@ -926,7 +926,7 @@ async def _branch_thread_with_reservation( source_metadata = source_record.get("metadata") or {} if source_metadata.get(_SIDECAR_METADATA_KEY) is True: raise HTTPException(status_code=409, detail="Branching is only available in the main conversation.") - source_accessor, source_config = build_checkpoint_state_accessor( + source_accessor, source_config = await abuild_checkpoint_state_accessor( request, thread_id=thread_id, assistant_id=source_record.get("assistant_id"), @@ -1222,7 +1222,7 @@ async def get_thread(thread_id: ThreadId, request: Request) -> ThreadResponse: checkpointer = get_checkpointer(request) record: dict | None = await thread_store.get(thread_id) try: - accessor, config = build_checkpoint_state_accessor( + accessor, config = await abuild_checkpoint_state_accessor( request, thread_id=thread_id, assistant_id=record.get("assistant_id") if record is not None else None, diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 508de8469..e9fa986ec 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -11,6 +11,7 @@ import asyncio import json import logging import re +import threading from collections.abc import AsyncIterator, Mapping from contextlib import asynccontextmanager from types import SimpleNamespace @@ -69,6 +70,7 @@ from deerflow.runtime.checkpoint_state import graph_state_schema from deerflow.runtime.events.message_identity import MESSAGE_SEQ_KEY from deerflow.runtime.goal import goal_thread_lock from deerflow.runtime.journal import build_checkpoint_history_seed_events +from deerflow.runtime.keyed_lock import KeyedLockTable from deerflow.runtime.runs.naming import resolve_root_run_name from deerflow.runtime.secret_context import ( LegacyRunMetadataSecretError, @@ -80,6 +82,7 @@ from deerflow.runtime.user_context import reset_current_user, set_current_user from deerflow.sandbox.lease import SANDBOX_SERVER_OWNED_CONTEXT_KEYS from deerflow.subagents.status_contract import SUBAGENT_ACCEPTANCE_VERDICT_KEY, SUBAGENT_RECEIPT_VERDICT_KEY, SUBAGENT_TOOL_RECEIPTS_KEY from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_context, ensure_trace_id +from deerflow.utils.assembly_io import run_assembly from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY from deerflow.utils.thread_id import validate_thread_id @@ -906,6 +909,8 @@ def build_checkpoint_state_mutation_accessor( # a restart. _STATE_ACCESSOR_GRAPH_CACHE_MAX = 64 _state_accessor_graph_cache: dict[tuple[str | None, str, int | None], tuple[Any, Any, Any]] = {} +_state_accessor_graph_cache_lock = threading.Lock() +_state_accessor_graph_build_locks = KeyedLockTable[tuple[str | None, str, int | None]]() def _accessor_graph_cache_max(app_config: Any) -> int: @@ -916,25 +921,52 @@ def _accessor_graph_cache_max(app_config: Any) -> int: ) -def _state_accessor_graph(agent_factory: Any, assistant_id: str | None, mode: str, snapshot_frequency: int | None, config: dict[str, Any]) -> Any: - app_config = (config.get("context") or {}).get("app_config") - key = (assistant_id, mode, snapshot_frequency) - cached = _state_accessor_graph_cache.get(key) - if cached is not None and cached[0] is agent_factory and cached[1] is app_config: - return cached[2] - if len(_state_accessor_graph_cache) >= _accessor_graph_cache_max(app_config): - _state_accessor_graph_cache.clear() +def _cached_state_accessor_graph(key: tuple[str | None, str, int | None], agent_factory: Any, app_config: Any) -> Any | None: + with _state_accessor_graph_cache_lock: + cached = _state_accessor_graph_cache.get(key) + if cached is not None and cached[0] is agent_factory and cached[1] is app_config: + return cached[2] + return None + + +def _cache_state_accessor_graph(key: tuple[str | None, str, int | None], agent_factory: Any, app_config: Any, graph: Any) -> None: + with _state_accessor_graph_cache_lock: + if len(_state_accessor_graph_cache) >= _accessor_graph_cache_max(app_config): + _state_accessor_graph_cache.clear() + _state_accessor_graph_cache[key] = (agent_factory, app_config, graph) + + +def _build_state_accessor_graph(agent_factory: Any, config: dict[str, Any]) -> Any: agent_result = agent_factory(config=config) try: from deerflow.agents.lead_agent.agent import unwrap_agent_graph - graph = unwrap_agent_graph(agent_result) + return unwrap_agent_graph(agent_result) except Exception: # A custom factory must keep working even if importing the lead # assembly type fails. - graph = agent_result - _state_accessor_graph_cache[key] = (agent_factory, app_config, graph) - return graph + return agent_result + + +def _state_accessor_graph(agent_factory: Any, assistant_id: str | None, mode: str, snapshot_frequency: int | None, config: dict[str, Any]) -> Any: + app_config = (config.get("context") or {}).get("app_config") + key = (assistant_id, mode, snapshot_frequency) + cached = _cached_state_accessor_graph(key, agent_factory, app_config) + if cached is not None: + return cached + + # Construction runs on assembly-pool threads, so same-key cold misses are + # serialized with a thread lock. The re-check under the lock makes + # overlapping first readers run the factory exactly once; a waiter whose + # factory or app-config identity changed while it waited still rebuilds, + # preserving identity-based cache invalidation. + with _state_accessor_graph_build_locks.hold(key): + cached = _cached_state_accessor_graph(key, agent_factory, app_config) + if cached is not None: + return cached + graph = _build_state_accessor_graph(agent_factory, config) + _cache_state_accessor_graph(key, agent_factory, app_config, graph) + return graph class _RawCheckpointSnapshot: @@ -1059,6 +1091,34 @@ def build_checkpoint_state_accessor( return accessor, config +async def abuild_checkpoint_state_accessor( + request: Request, + *, + thread_id: str, + assistant_id: str | None = None, + checkpoint_id: str | None = None, +) -> tuple[CheckpointStateAccessor, dict[str, Any]]: + """Async variant of :func:`build_checkpoint_state_accessor`. + + Identical accessor construction, but the agent-factory assembly — which + re-enters ``get_available_tools()`` and may block on MCP cache + initialization — runs off-loop on the dedicated assembly pool so the + Gateway event loop keeps making progress (issue #5172). Repeat calls hit + ``_state_accessor_graph_cache`` and only pay the thread hop; overlapping + cold readers with the same cache key are serialized per key so the + factory runs exactly once, and a reader whose factory or app-config + identity changed while it waited rebuilds instead of reusing the + winner's graph. + """ + return await run_assembly( + build_checkpoint_state_accessor, + request, + thread_id=thread_id, + assistant_id=assistant_id, + checkpoint_id=checkpoint_id, + ) + + async def resolve_thread_assistant_id( request: Request, thread_id: str, @@ -1098,7 +1158,7 @@ async def build_thread_checkpoint_state_accessor( ``AgentMiddleware.state_schema`` from the response. """ assistant_id = await resolve_thread_assistant_id(request, thread_id, fail_closed=fail_closed) - return build_checkpoint_state_accessor( + return await abuild_checkpoint_state_accessor( request, thread_id=thread_id, assistant_id=assistant_id, @@ -1229,7 +1289,7 @@ async def ensure_checkpoint_history_seeded( if await get_checkpointer(request).aget_tuple(checkpoint_config) is None: return - accessor, config = build_checkpoint_state_accessor( + accessor, config = await abuild_checkpoint_state_accessor( request, thread_id=thread_id, assistant_id=assistant_id, diff --git a/backend/packages/harness/deerflow/runtime/keyed_lock.py b/backend/packages/harness/deerflow/runtime/keyed_lock.py index c6a3d56dc..9b300c88e 100644 --- a/backend/packages/harness/deerflow/runtime/keyed_lock.py +++ b/backend/packages/harness/deerflow/runtime/keyed_lock.py @@ -1,12 +1,17 @@ -"""Async per-key serialization with waiter-aware entry reclamation.""" +"""Per-key serialization with waiter-aware entry reclamation. + +:func:`AsyncKeyedLockTable` serves asyncio callers; +:class:`KeyedLockTable` is the thread-side counterpart for blocking +critical sections running on worker threads. +""" from __future__ import annotations import asyncio import threading import weakref -from collections.abc import AsyncIterator, Hashable -from contextlib import asynccontextmanager +from collections.abc import AsyncIterator, Hashable, Iterator +from contextlib import asynccontextmanager, contextmanager from dataclasses import dataclass @@ -85,3 +90,58 @@ class AsyncKeyedLockTable[KeyT: Hashable]: entries.pop(key) if not entries and self._entries_by_loop.get(loop) is entries: self._entries_by_loop.pop(loop, None) + + +@dataclass(slots=True) +class _ThreadEntry: + lock: threading.Lock + participants: int = 0 # current holder plus queued waiters + + +class KeyedLockTable[KeyT: Hashable]: + """Serialize same-key blocking work across worker threads. + + Thread-side counterpart of :class:`AsyncKeyedLockTable`: the guard + protects only the registry, and the blocking critical section never + holds it. Participants are counted before acquiring the lock, so an + entry stays discoverable until its final holder or waiter leaves and a + new caller cannot create a second lock that bypasses an already queued + waiter. Entries whose last participant leaves are reclaimed so idle + keys do not accumulate. + """ + + def __init__(self) -> None: + self._guard = threading.Lock() + self._entries: dict[KeyT, _ThreadEntry] = {} + + @contextmanager + def hold(self, key: KeyT) -> Iterator[None]: + """Hold the lock for ``key``, blocking the calling thread.""" + entry = self._checkout(key) + acquired = False + try: + entry.lock.acquire() + acquired = True + yield + finally: + if acquired: + entry.lock.release() + self._checkin(key, entry) + + def _checkout(self, key: KeyT) -> _ThreadEntry: + with self._guard: + entry = self._entries.get(key) + if entry is None: + entry = _ThreadEntry(lock=threading.Lock()) + self._entries[key] = entry + entry.participants += 1 + return entry + + def _checkin(self, key: KeyT, entry: _ThreadEntry) -> None: + with self._guard: + entry.participants -= 1 + if entry.participants != 0 or entry.lock.locked(): + return + if self._entries.get(key) is not entry: + return + self._entries.pop(key) diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index f4ff71f93..b6037d95e 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -81,6 +81,7 @@ from deerflow.runtime.user_context import get_current_user, get_effective_user_i from deerflow.sandbox.lease import SANDBOX_SERVER_OWNED_CONTEXT_KEYS from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, ensure_trace_id from deerflow.tracing import inject_langfuse_metadata +from deerflow.utils.assembly_io import run_assembly from deerflow.utils.messages import message_to_text from deerflow.workspace_changes import capture_workspace_snapshot, get_changed_output_paths, record_workspace_changes from deerflow.workspace_changes.types import WorkspaceSnapshot @@ -1087,7 +1088,11 @@ async def run_agent( from deerflow.extensions import bind_agent_build_extensions with bind_agent_build_extensions(extensions): - agent = _agent_graph(agent_factory(**agent_factory_kwargs)) + # Assemble off-loop: agent construction re-enters + # get_available_tools(), which may block on MCP cache + # initialization — it must not stall the calling event loop + # (issue #5172). + agent = _agent_graph(await run_assembly(agent_factory, **agent_factory_kwargs)) accessor = CheckpointStateAccessor.bind( agent, diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index 2abcc2b47..d7438e63b 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -21,6 +21,7 @@ from deerflow.subagents.executor import ( get_background_task_result, request_cancel_background_task, ) +from deerflow.utils.assembly_io import run_assembly logger = logging.getLogger(__name__) @@ -195,13 +196,33 @@ class SubagentBatchService: spec.get("parent_model"), app_config=app_config, ) - tools = get_available_tools( + # Assemble off-loop: tool assembly may block on MCP cache + # initialization, which must not stall the calling event loop (issue #5172). + tools = await run_assembly( + get_available_tools, groups=spec.get("tool_groups"), model_name=effective_model, subagent_enabled=False, include_upload_tool=False, app_config=app_config, ) + # Revalidate durable state before launching: cancel_batch may have + # terminalized this item (or its lease may have been lost) while + # assembly blocked in the worker thread — the poll loop's checks + # only start after execute_async(), so launching without this + # check would run work the user already cancelled. + lease = await self._repository.renew_item_lease( + item_id, + lease_owner=self._lease_owner, + lease_seconds=self._config.lease_seconds, + now=datetime.now(UTC), + ) + if not lease["valid"]: + logger.info( + "Durable batch item %s cancelled or lease lost during tool assembly; skipping launch", + item_id, + ) + return executor = SubagentExecutor( config=config, tools=tools, diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index ca0f5e27d..261c13646 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -43,6 +43,7 @@ from deerflow.subagents.status_contract import ( ) from deerflow.tools.types import Runtime from deerflow.trace_context import DEERFLOW_TRACE_METADATA_KEY, resolve_trace_id +from deerflow.utils.assembly_io import run_assembly from deerflow.utils.custom_events import aemit_custom_event if TYPE_CHECKING: @@ -881,7 +882,9 @@ async def task_tool( } if resolved_app_config is not None: available_tools_kwargs["app_config"] = resolved_app_config - tools = get_available_tools(**available_tools_kwargs) + # Assemble off-loop: tool assembly may block on MCP cache initialization, + # which must not stall the calling event loop (issue #5172). + tools = await run_assembly(get_available_tools, **available_tools_kwargs) # Create executor executor_kwargs = { diff --git a/backend/packages/harness/deerflow/tools/tools.py b/backend/packages/harness/deerflow/tools/tools.py index b59861932..6b885bd41 100644 --- a/backend/packages/harness/deerflow/tools/tools.py +++ b/backend/packages/harness/deerflow/tools/tools.py @@ -1,4 +1,5 @@ import logging +import threading from langchain.tools import BaseTool @@ -49,10 +50,22 @@ def _is_host_bash_tool(tool: object) -> bool: return False +_sync_invocable_tool_lock = threading.Lock() + + def _ensure_sync_invocable_tool(tool: BaseTool) -> BaseTool: - """Attach a sync wrapper to async-only tools used by sync agent callers.""" - if getattr(tool, "func", None) is None and getattr(tool, "coroutine", None) is not None: - tool.func = make_sync_tool_wrapper(tool.coroutine, tool.name) + """Attach a sync wrapper to async-only tools used by sync agent callers. + + The wrapped objects are process-wide singletons (BUILTIN_TOOLS / + SUBAGENT_TOOLS / MCP cache entries) and tool assembly may now run on + worker threads concurrently; double-checked locking makes the in-place + ``tool.func`` wrap explicitly single-shot instead of incidental. + """ + if getattr(tool, "func", None) is not None or getattr(tool, "coroutine", None) is None: + return tool + with _sync_invocable_tool_lock: + if getattr(tool, "func", None) is None: + tool.func = make_sync_tool_wrapper(tool.coroutine, tool.name) return tool diff --git a/backend/packages/harness/deerflow/utils/AGENTS.md b/backend/packages/harness/deerflow/utils/AGENTS.md new file mode 100644 index 000000000..9747410ce --- /dev/null +++ b/backend/packages/harness/deerflow/utils/AGENTS.md @@ -0,0 +1,3 @@ +### Agent / Tool Assembly Off-Load + +Tool and agent assembly re-enters `get_available_tools()` and may block on MCP discovery, so the four async assembly entry points — `run_agent`'s `agent_factory` call, `task_tool`, durable batch `_execute_item`, and `abuild_checkpoint_state_accessor` — dispatch through `deerflow.utils.assembly_io.run_assembly`, a dedicated ContextVar-preserving bounded executor (`DEER_FLOW_ASSEMBLY_WORKERS`, default 8) rather than the loop's default executor; a hung MCP server therefore parks an assembly worker instead of queueing unrelated default-executor work, and the pool logs a warning when pending assemblies exceed the worker count. `tests/blocking_io/test_tool_assembly_offloop.py` pins all four offloads plus the ContextVar propagation. diff --git a/backend/packages/harness/deerflow/utils/assembly_io.py b/backend/packages/harness/deerflow/utils/assembly_io.py new file mode 100644 index 000000000..a654a9a8f --- /dev/null +++ b/backend/packages/harness/deerflow/utils/assembly_io.py @@ -0,0 +1,122 @@ +"""Dedicated async offload helper for agent/tool assembly work.""" + +from __future__ import annotations + +import asyncio +import atexit +import contextvars +import logging +import os +import threading +import time +from collections.abc import Callable +from concurrent.futures import Future, ThreadPoolExecutor + +logger = logging.getLogger(__name__) + + +def _default_assembly_workers() -> int: + raw = os.getenv("DEER_FLOW_ASSEMBLY_WORKERS") + if raw: + try: + workers = int(raw) + if workers > 0: + return workers + except ValueError: + pass + logger.warning("Invalid DEER_FLOW_ASSEMBLY_WORKERS value; using default assembly worker count") + return 8 + + +_ASSEMBLY_WORKERS = _default_assembly_workers() +_ASSEMBLY_EXECUTOR = ThreadPoolExecutor(max_workers=_ASSEMBLY_WORKERS, thread_name_prefix="assembly") + +# Pending (submitted, unfinished) assembly count. Increments on the event +# loop before dispatch and decrements from the dispatched work item's +# `finally` on a pool thread; guarded for multi-loop test environments. +_pending_assemblies = 0 +_pending_lock = threading.Lock() +_last_starvation_log = 0.0 +_STARVATION_LOG_INTERVAL_SECONDS = 30.0 + + +def _shutdown_assembly_executor() -> None: + _ASSEMBLY_EXECUTOR.shutdown(wait=False, cancel_futures=True) + + +atexit.register(_shutdown_assembly_executor) + + +async def run_assembly[**P, T](func: Callable[P, T], /, *args: P.args, **kwargs: P.kwargs) -> T: + """Run blocking agent/tool assembly on the dedicated assembly pool. + + Tool and agent assembly re-enters ``get_available_tools()``, which may + block for the full MCP discovery duration (a slow or hung stdio server + parks a worker until the MCP timeout). Dispatching those hops onto the + loop's **default** executor lets a few parked assemblies queue every other + ``asyncio.to_thread`` / ``run_in_executor(None, ...)`` caller behind them + and reintroduce a loop-wide stall through a different door; this pool + keeps the capacity explicit and bounds the blast radius, mirroring + ``utils/file_io.py`` and ``tools/sync.py``. + + ``asyncio.to_thread`` copies ``ContextVar`` values automatically; raw + ``loop.run_in_executor`` does not. Copy the current context explicitly so + agent-assembly helpers such as ``bind_agent_build_extensions`` keep + working inside the worker thread. + """ + global _pending_assemblies, _last_starvation_log + + with _pending_lock: + _pending_assemblies += 1 + pending = _pending_assemblies + # Saturation is invisible otherwise: workers parked on a hung MCP + # server leave later assemblies queued indefinitely while the loop + # stays healthy. Warn at most once per interval while starved. + warn = pending > _ASSEMBLY_WORKERS and (time.monotonic() - _last_starvation_log) > _STARVATION_LOG_INTERVAL_SECONDS + if warn: + _last_starvation_log = time.monotonic() + if warn: + logger.warning( + "Assembly pool saturated: %d pending assemblies on %d workers; agent assembly is starved (likely a hung MCP server)", + pending, + _ASSEMBLY_WORKERS, + ) + + loop = asyncio.get_running_loop() + ctx = contextvars.copy_context() + + def _work() -> T: + # The decrement must ride the dispatched work item, not the asyncio + # future: if the submitting loop is closed while the worker is still + # running, the future never resolves and an asyncio-future callback + # would never fire, ratcheting the count up permanently and eventually + # firing the starvation warning with no starvation behind it. + global _pending_assemblies + try: + return ctx.run(func, *args, **kwargs) + finally: + with _pending_lock: + _pending_assemblies -= 1 + + try: + executor_future = _ASSEMBLY_EXECUTOR.submit(_work) + except Exception: + # Nothing was dispatched, so the work item's finally will never run; + # release the slot claimed above or the count wedges. + with _pending_lock: + _pending_assemblies -= 1 + raise + + def _release_if_cancelled(future: Future[T]) -> None: + # Exactly-once cleanup for jobs cancelled while still queued: + # ``Future.cancel`` only succeeds before the executor starts the + # item, so ``cancelled()`` is true precisely when ``_work()`` never + # ran and its finally-block decrement never will. Runs in whichever + # thread cancels or completes the job — loop-independent. + global _pending_assemblies + if future.cancelled(): + with _pending_lock: + _pending_assemblies -= 1 + + executor_future.add_done_callback(_release_if_cancelled) + return await asyncio.wrap_future(executor_future, loop=loop) diff --git a/backend/tests/blocking_io/test_tool_assembly_offloop.py b/backend/tests/blocking_io/test_tool_assembly_offloop.py new file mode 100644 index 000000000..c683211a4 --- /dev/null +++ b/backend/tests/blocking_io/test_tool_assembly_offloop.py @@ -0,0 +1,351 @@ +"""Regression: tool assembly runs off the event loop (issue #5172). + +``get_available_tools()`` may block on MCP cache initialization while it runs +on async agent-assembly paths. The offload dispatches the (unchanged, +synchronous) assembly to the dedicated assembly pool (``asyncio.to_thread``' +s default-executor alternative) at the async entry points: ``task_tool``, +``SubagentBatchService._execute_item``, the Gateway run worker's agent +construction (``run_agent`` -> ``agent_factory`` -> lead-agent assembly), and +the checkpoint state-accessor build (``abuild_checkpoint_state_accessor`` -> +``build_thread_checkpoint_state_accessor``). + +Under the strict Blockbuster context (this directory's conftest), any +blocking IO reached from ``deerflow.*`` while on the event loop raises +``BlockingError``. ``get_available_tools`` is injected here as a **blocking +probe** (real file IO): what must be pinned is that the assembly call never +executes on the event loop, not that today's assembly happens to be cheap — +a slow or hung stdio MCP server turns the same call into a full-loop stall. +If an entry point is flattened back to a plain call, the main test fails; +the meta-check below proves the probe has teeth by calling it directly on +the loop. +""" + +from __future__ import annotations + +import importlib +import json +import threading +from enum import Enum +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from langchain_core.messages import ToolMessage + +from deerflow.config.extensions_config import ExtensionsConfig +from deerflow.extensions import get_agent_build_extensions +from deerflow.runtime.events.store.memory import MemoryRunEventStore +from deerflow.runtime.runs.manager import RunManager +from deerflow.runtime.runs.worker import RunContext, run_agent +from deerflow.subagents.config import SubagentConfig + +# importlib.import_module binds the real module: the package attribute +# ``deerflow.tools.builtins.task_tool`` is shadowed by the StructuredTool. +task_tool_module = importlib.import_module("deerflow.tools.builtins.task_tool") +batch_service_module = importlib.import_module("deerflow.subagents.batch_service") +# Imported at module scope: the first import of app.gateway.services pulls in +# fastapi/pydantic, whose one-time metadata reads must not run inside a gated +# test item. +gateway_services = importlib.import_module("app.gateway.services") + +pytestmark = pytest.mark.asyncio + + +class _FakeSubagentStatus(Enum): + COMPLETED = "completed" + FAILED = "failed" + RUNNING = "running" + + @property + def is_terminal(self) -> bool: + return self is not _FakeSubagentStatus.RUNNING + + +def _blocking_probe_tools(probe_file: Path, observed_threads: list | None = None): + """A ``get_available_tools`` replacement performing real blocking file IO.""" + + def get_tools(**_kwargs): + # Real filesystem IO: trips the strict gate when it runs on the loop. + body = probe_file.read_text(encoding="utf-8") + if observed_threads is not None: + observed_threads.append(threading.current_thread()) + return [body] + + return get_tools + + +def _completed_result() -> SimpleNamespace: + return SimpleNamespace( + status=_FakeSubagentStatus.COMPLETED, + ai_messages=[], + result="done", + error=None, + stop_reason=None, + token_usage_records=[], + usage_reported=False, + tool_receipts=None, + bash_executions=None, + ) + + +class _DummyExecutor: + def __init__(self, **_kwargs): + pass + + def execute_async(self, _prompt, task_id=None): + return task_id or "generated-task-id" + + +async def test_task_tool_assembles_off_loop(monkeypatch, tmp_path): + """task_tool dispatches get_available_tools to a worker thread.""" + (tmp_path / "probe.txt").write_text("probe body", encoding="utf-8") + observed_threads: list = [] + monkeypatch.setattr( + "deerflow.tools.get_available_tools", + _blocking_probe_tools(tmp_path / "probe.txt", observed_threads), + ) + monkeypatch.setattr(task_tool_module, "SubagentStatus", _FakeSubagentStatus) + monkeypatch.setattr(task_tool_module, "SubagentExecutor", _DummyExecutor) + monkeypatch.setattr( + task_tool_module, + "get_subagent_config", + lambda _name: SubagentConfig( + name="general-purpose", + description="General helper", + system_prompt="Base system prompt", + max_turns=50, + timeout_seconds=10, + ), + ) + monkeypatch.setattr(task_tool_module, "get_available_subagent_names", lambda **_kwargs: ["general-purpose"]) + monkeypatch.setattr(task_tool_module, "get_background_task_result", lambda _task_id: _completed_result()) + monkeypatch.setattr(task_tool_module, "get_stream_writer", lambda: lambda _event: None) + + async def _no_sleep(_: float) -> None: + return None + + monkeypatch.setattr(task_tool_module.asyncio, "sleep", _no_sleep) + + workspace = tmp_path / "user-data" / "workspace" + workspace.mkdir(parents=True, exist_ok=True) + runtime = SimpleNamespace( + state={ + "sandbox": {"sandbox_id": "local"}, + "thread_data": { + "workspace_path": str(workspace), + "uploads_path": str(tmp_path / "user-data" / "uploads"), + "outputs_path": str(tmp_path / "user-data" / "outputs"), + }, + }, + context={"thread_id": "thread-1"}, + config={"metadata": {"model_name": "ark-model", "trace_id": "trace-1"}}, + ) + + tool = task_tool_module.task_tool + invoke = getattr(tool, "coroutine", None) or getattr(tool, "func", None) + assert invoke is not None + command = await invoke( + runtime=runtime, + description="test", + prompt="p", + subagent_type="general-purpose", + tool_call_id="tc-offloop", + ) + + messages = command.update["messages"] + assert len(messages) == 1 + assert isinstance(messages[0], ToolMessage) + assert observed_threads, "tool assembly must be invoked" + assert all(thread is not threading.main_thread() for thread in observed_threads) + + +async def test_batch_item_assembles_off_loop(monkeypatch, tmp_path): + """SubagentBatchService._execute_item dispatches assembly to a worker thread.""" + (tmp_path / "probe.txt").write_text("probe body", encoding="utf-8") + observed_threads: list = [] + monkeypatch.setattr( + "deerflow.tools.get_available_tools", + _blocking_probe_tools(tmp_path / "probe.txt", observed_threads), + ) + monkeypatch.setattr(batch_service_module, "SubagentStatus", _FakeSubagentStatus) + monkeypatch.setattr(batch_service_module, "SubagentExecutor", _DummyExecutor) + monkeypatch.setattr( + batch_service_module, + "get_background_task_result", + lambda _execution_id: _completed_result(), + ) + monkeypatch.setattr( + batch_service_module, + "request_cancel_background_task", + lambda _execution_id: None, + ) + monkeypatch.setattr( + batch_service_module, + "resolve_subagent_model_name", + lambda *_args, **_kwargs: "test-model", + ) + + service = batch_service_module.SubagentBatchService( + repository=SimpleNamespace( + mark_item_running=None, + renew_item_lease=None, + finalize_item=None, + ), + config=SimpleNamespace( + lease_seconds=10.0, + poll_interval_seconds=1.0, + max_result_chars=1000, + result_preview_max_chars=200, + ), + runtime_config=SimpleNamespace(), + app_config=SimpleNamespace(), + execution_capacity=None, + ) + + finalize_calls: list[dict] = [] + + async def _finalize_item(item_id, **kwargs): + finalize_calls.append({"item_id": item_id, **kwargs}) + + async def _renew_item_lease(item_id, **_kwargs): + return {"valid": True, "cancel_requested": False} + + service._repository = SimpleNamespace(finalize_item=_finalize_item, renew_item_lease=_renew_item_lease) + + item = { + "id": "item-1", + "item_key": "key-1", + "prompt": "do the thing", + "batch": { + "id": "batch-1", + "thread_id": "thread-1", + "user_id": "user-1", + "run_id": None, + "execution_spec": { + "subagent_config": { + "name": "general-purpose", + "description": "General helper", + "system_prompt": "Base system prompt", + "model": "test-model", + "max_turns": 5, + "timeout_seconds": 10, + }, + }, + }, + } + + await service._execute_item(item) + + assert len(finalize_calls) == 1 + assert finalize_calls[0]["item_id"] == "item-1" + assert finalize_calls[0]["succeeded"] is True + assert observed_threads, "tool assembly must be invoked" + assert all(thread is not threading.main_thread() for thread in observed_threads) + + +async def test_run_agent_assembles_off_loop(monkeypatch, tmp_path): + """run_agent dispatches agent_factory (lead-agent assembly) to a worker thread.""" + cfg = tmp_path / "extensions_config.json" + cfg.write_text(json.dumps({"mcpServers": {}, "skills": {}}), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + observed_threads: list = [] + # Sentinel bound via ctx.extensions: pins that run_assembly() preserves + # ContextVars, so bind_agent_build_extensions reaches the factory. Dropping + # the ctx.run in run_assembly makes the factory observe the startup + # fallback instead, with no error — this is the regression nothing else + # in the suite catches. + sentinel_extensions = SimpleNamespace( + id="sentinel-extensions", + needs_task_store=False, + has_task_lifecycle=False, + ) + observed_extensions: list = [] + + class _DummyStreamAgent: + async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + yield {"messages": []} + + def _factory(*, config): + observed_threads.append(threading.current_thread()) + observed_extensions.append(get_agent_build_extensions()) + # Real production blocking read (executed inside a deerflow.* frame): + # trips the strict gate when the factory runs on the loop. + ExtensionsConfig.from_file() + return _DummyStreamAgent() + + run_manager = RunManager() + record = await run_manager.create("thread-1") + + await run_agent( + SimpleNamespace(publish=AsyncMock(), publish_end=AsyncMock(), cleanup=AsyncMock()), + run_manager, + record, + ctx=RunContext(checkpointer=None, event_store=MemoryRunEventStore(), extensions=sentinel_extensions), + agent_factory=_factory, + graph_input={}, + config={}, + ) + + assert observed_threads, "agent assembly must be invoked" + assert all(thread is not threading.main_thread() for thread in observed_threads) + assert observed_extensions == [sentinel_extensions], "the factory must observe the run-bound extension snapshot, not the startup fallback" + + +async def test_state_accessor_build_assembles_off_loop(monkeypatch, tmp_path): + """abuild_checkpoint_state_accessor dispatches assembly to the assembly pool.""" + cfg = tmp_path / "extensions_config.json" + cfg.write_text(json.dumps({"mcpServers": {}, "skills": {}}), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + observed_threads: list = [] + + ctx = SimpleNamespace( + checkpointer=None, + store=None, + checkpoint_channel_mode="full", + checkpoint_snapshot_frequency=None, + app_config=None, + ) + monkeypatch.setattr(gateway_services, "get_run_context", lambda _request: ctx) + + async def _no_assistant(_request, _thread_id, **_kwargs): + return None + + monkeypatch.setattr(gateway_services, "resolve_thread_assistant_id", _no_assistant) + + def _resolve_factory(_assistant_id): + # A fresh factory per resolution: the accessor graph cache validates + # the factory identity, so this always misses and always reaches the + # probe regardless of what earlier tests left cached. + def _factory(*, config): + observed_threads.append(threading.current_thread()) + # Real production blocking read (executed inside a deerflow.* frame): + # trips the strict gate when the factory runs on the loop. + ExtensionsConfig.from_file() + return SimpleNamespace() + + return _factory + + monkeypatch.setattr(gateway_services, "resolve_agent_factory", _resolve_factory) + + await gateway_services.build_thread_checkpoint_state_accessor(SimpleNamespace(), thread_id="thread-1") + + assert observed_threads, "agent assembly must be invoked" + assert all(thread is not threading.main_thread() for thread in observed_threads) + + +async def test_extensions_config_read_trips_the_gate(monkeypatch, tmp_path): + """Meta-check: reading the extensions config from ``deerflow.*`` code on + the event loop must raise BlockingError — the exact syscall class issue + #5172 is about — so the anchors above cannot go vacuously green. (The + probe's own ``read_text`` trips through the same gate, proven here with + the production reader instead of a test-file stack, which the + ``scanned_modules`` filter would ignore.)""" + from blockbuster import BlockingError + + cfg = tmp_path / "extensions_config.json" + cfg.write_text(json.dumps({"mcpServers": {}, "skills": {}}), encoding="utf-8") + monkeypatch.setenv("DEER_FLOW_EXTENSIONS_CONFIG_PATH", str(cfg)) + + with pytest.raises(BlockingError): + ExtensionsConfig.from_file() diff --git a/backend/tests/test_agent_guidance_check.py b/backend/tests/test_agent_guidance_check.py index deeaf9d4f..b87fa3d10 100644 --- a/backend/tests/test_agent_guidance_check.py +++ b/backend/tests/test_agent_guidance_check.py @@ -30,6 +30,7 @@ EXPECTED_GUIDANCE_PATHS = { "backend/packages/harness/deerflow/tools/AGENTS.md", "backend/packages/harness/deerflow/tracing/AGENTS.md", "backend/packages/harness/deerflow/tui/AGENTS.md", + "backend/packages/harness/deerflow/utils/AGENTS.md", "frontend/src/AGENTS.md", "scripts/AGENTS.md", } diff --git a/backend/tests/test_assembly_io.py b/backend/tests/test_assembly_io.py new file mode 100644 index 000000000..9ce489233 --- /dev/null +++ b/backend/tests/test_assembly_io.py @@ -0,0 +1,114 @@ +"""Unit tests for the assembly pool's pending counter. + +The starvation warning in :func:`deerflow.utils.assembly_io.run_assembly` +fires once the pending (submitted, unfinished) count exceeds the worker +count. Nothing else in the suite reads ``_pending_assemblies``, so a drift +in the decrement would silently ratchet the count up and eventually fire +the warning with no starvation behind it — pin the three behaviors here. +""" + +from __future__ import annotations + +import asyncio +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest + +import deerflow.utils.assembly_io as assembly_io + + +def test_pending_count_returns_to_zero_after_healthy_call() -> None: + async def main() -> str: + return await assembly_io.run_assembly(lambda: "ok") + + assert asyncio.run(main()) == "ok" + assert assembly_io._pending_assemblies == 0 + + +def test_abandoned_loop_does_not_wedge_the_counter() -> None: + """A submitting loop that dies while its worker is still parked must not + wedge the count: the decrement rides the dispatched work item's ``finally`` + (pool thread), not the asyncio future's done callback (submitting loop).""" + worker_started = threading.Event() + worker_release = threading.Event() + + def _parked() -> str: + worker_started.set() + worker_release.wait(timeout=10) + return "done" + + loop = asyncio.new_event_loop() + errors: list[BaseException] = [] + + def _run() -> None: + try: + loop.run_until_complete(assembly_io.run_assembly(_parked)) + except BaseException as exc: # the abandoned submission is cancelled/raises + errors.append(exc) + + thread = threading.Thread(target=_run, daemon=True) + thread.start() + assert worker_started.wait(timeout=5) + + # Abandon the submission loop: stop it while the coroutine is still + # awaiting the dispatched work, so its future can never resolve and the + # old done-callback decrement would never fire. + loop.call_soon_threadsafe(loop.stop) + worker_release.set() + thread.join(timeout=10) + assert not thread.is_alive() + loop.close() + + # Give the pool worker time to finish its finally-block decrement. + deadline = time.monotonic() + 5 + while assembly_io._pending_assemblies != 0 and time.monotonic() < deadline: + time.sleep(0.01) + assert assembly_io._pending_assemblies == 0, errors or "counter was not decremented by the dispatched work item" + + +def test_queued_cancellation_releases_the_pending_count() -> None: + """A job cancelled while still queued never runs its work item, so the + dispatched finally never fires; the cancelled-future cleanup must + release the slot exactly once instead (WillemJiang, PR #5224 review).""" + worker_started = threading.Event() + worker_release = threading.Event() + + def _parked() -> str: + worker_started.set() + worker_release.wait(timeout=10) + return "done" + + pool = ThreadPoolExecutor(max_workers=1) + original_executor = assembly_io._ASSEMBLY_EXECUTOR + assembly_io._ASSEMBLY_EXECUTOR = pool + try: + + async def main() -> None: + loop = asyncio.get_running_loop() + first = asyncio.ensure_future(assembly_io.run_assembly(_parked)) + assert await loop.run_in_executor(None, worker_started.wait, 5.0) + assert assembly_io._pending_assemblies == 1 + + # The single worker is parked in the first call, so the second + # submission is queued, not started. + second = asyncio.ensure_future(assembly_io.run_assembly(lambda: "queued")) + deadline = time.monotonic() + 5 + while assembly_io._pending_assemblies != 2 and time.monotonic() < deadline: + await asyncio.sleep(0.01) + assert assembly_io._pending_assemblies == 2 + + second.cancel() + with pytest.raises(asyncio.CancelledError): + await second + assert assembly_io._pending_assemblies == 1, "cancelled queued job must release its slot without running _work" + + worker_release.set() + assert await first == "done" + assert assembly_io._pending_assemblies == 0 + + asyncio.run(main()) + finally: + assembly_io._ASSEMBLY_EXECUTOR = original_executor + pool.shutdown(wait=True) diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index f017504ce..95af63958 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -990,6 +990,67 @@ def test_state_accessor_graph_cache_honors_configured_cap(): gateway_services._state_accessor_graph_cache.clear() +def test_state_accessor_graph_serializes_same_key_cold_construction(): + """Overlapping first reads with the same factory object and app-config + identity must run the factory exactly once (per-key construction + serialization, PR #5224 review), while a changed factory identity still + rebuilds instead of reusing the stored graph.""" + import threading + import time + from typing import Any + + from app.gateway import services as gateway_services + + builds = [] + first_inside = threading.Event() + release_first = threading.Event() + + def slow_factory(*, config): + graph = object() + builds.append(graph) + first_inside.set() + release_first.wait(timeout=10) + return graph + + gateway_services._state_accessor_graph_cache.clear() + results: list[Any] = [] + errors: list[BaseException] = [] + + def reader() -> None: + try: + results.append(gateway_services._state_accessor_graph(slow_factory, None, "full", None, {})) + except BaseException as exc: # pragma: no cover - surfaced below + errors.append(exc) + + try: + first = threading.Thread(target=reader) + second = threading.Thread(target=reader) + first.start() + assert first_inside.wait(timeout=5) + second.start() + # The second reader blocks on the per-key lock while the first is + # still inside the factory: no duplicate construction. + deadline = time.monotonic() + 5 + while second.is_alive() and time.monotonic() < deadline: + time.sleep(0.01) + assert len(builds) == 1, errors + assert second.is_alive(), "second cold reader must wait for the in-flight construction" + + release_first.set() + first.join(timeout=10) + second.join(timeout=10) + assert not (first.is_alive() or second.is_alive()) + assert len(builds) == 1 + assert len(results) == 2 and results[0] is results[1] + + # A different factory object is an identity change: rebuild, not reuse. + other = gateway_services._state_accessor_graph(lambda *, config: object(), None, "full", None, {}) + assert other is not results[0] + assert len(builds) == 1 + finally: + gateway_services._state_accessor_graph_cache.clear() + + def test_build_run_config_configurable_custom_agent_dual_writes_agent_name(): """Regression for issue #3549: even when the caller uses the legacy ``configurable`` path, ``agent_name`` must also land in diff --git a/backend/tests/test_runtime_keyed_lock.py b/backend/tests/test_runtime_keyed_lock.py index 34d9e1855..361bf2f4c 100644 --- a/backend/tests/test_runtime_keyed_lock.py +++ b/backend/tests/test_runtime_keyed_lock.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import gc import threading +import time import weakref from concurrent.futures import ThreadPoolExecutor from uuid import uuid4 @@ -198,3 +199,104 @@ def test_same_key_is_independent_across_event_loops() -> None: futures = [executor.submit(run_loop) for _ in range(2)] for future in futures: future.result(timeout=3) + + +def test_keyed_lock_table_late_arrival_cannot_bypass_queued_waiter() -> None: + """Thread-side counterpart of the async table's late-arrival test. + + Overlapping ``hold()`` calls on one key serialize the bodies, and a + queued waiter keeps the entry alive: an arrival that shows up while + waiters are still queued must join the live entry instead of creating + a second lock and entering concurrently (which is what an early + reclamation would allow). + """ + from deerflow.runtime.keyed_lock import KeyedLockTable + + table = KeyedLockTable[str]() + release_first = threading.Event() + release_second = threading.Event() + release_third = threading.Event() + first_entered = threading.Event() + second_started = threading.Event() + second_entered = threading.Event() + third_started = threading.Event() + third_entered = threading.Event() + fourth_started = threading.Event() + fourth_entered = threading.Event() + guard = threading.Lock() + active = 0 + max_active = 0 + entered: list[str] = [] + + def participant(name: str, started: threading.Event | None, entered_event: threading.Event, release: threading.Event | None) -> None: + nonlocal active, max_active + if started is not None: + started.set() + with table.hold("key"): + with guard: + active += 1 + max_active = max(max_active, active) + entered.append(name) + entered_event.set() + try: + if release is not None: + release.wait(timeout=10) + finally: + with guard: + active -= 1 + + first = threading.Thread(target=participant, args=("first", None, first_entered, release_first)) + first.start() + assert first_entered.wait(timeout=5) + + second = threading.Thread(target=participant, args=("second", second_started, second_entered, release_second)) + second.start() + assert second_started.wait(timeout=5) + assert not second_entered.is_set(), "second must queue while the first still holds the entry" + + third = threading.Thread(target=participant, args=("third", third_started, third_entered, release_third)) + third.start() + assert third_started.wait(timeout=5) + time.sleep(0.05) + assert not third_entered.is_set(), "third must queue on the live entry, not enter" + + # Hand the entry off: first leaves, exactly one queued waiter gets in + # and parks inside its body on its release event. + release_first.set() + first.join(timeout=10) + deadline = time.monotonic() + 5 + while not (second_entered.is_set() or third_entered.is_set()) and time.monotonic() < deadline: + time.sleep(0.01) + assert second_entered.is_set() != third_entered.is_set(), "exactly one waiter holds the entry" + + # A late arrival while a waiter is still queued must join the live + # entry (and therefore stay out until that waiter leaves), never enter + # through a freshly created second lock. + fourth = threading.Thread(target=participant, args=("fourth", fourth_started, fourth_entered, None)) + fourth.start() + assert fourth_started.wait(timeout=5) + time.sleep(0.05) + assert not fourth_entered.is_set(), "late arrival bypassed a queued waiter" + + release_second.set() + release_third.set() + for thread in (second, third, fourth): + thread.join(timeout=10) + assert not thread.is_alive() + + assert max_active == 1 + assert len(entered) == 4 and set(entered) == {"first", "second", "third", "fourth"} + assert table._entries == {}, "the last check-in must pop the entry" + + +def test_keyed_lock_table_many_unique_keys_are_reclaimed() -> None: + from deerflow.runtime.keyed_lock import KeyedLockTable + + table = KeyedLockTable[int]() + + for key in range(1000): + with table.hold(key): + assert len(table._entries) == 1 + assert key not in table._entries + + assert table._entries == {} diff --git a/backend/tests/test_subagent_batch_service.py b/backend/tests/test_subagent_batch_service.py index 6a1aab722..2ac76bfd2 100644 --- a/backend/tests/test_subagent_batch_service.py +++ b/backend/tests/test_subagent_batch_service.py @@ -101,6 +101,9 @@ async def test_execute_item_marks_real_running_then_persists_terminal_result(mon result.result = "done" return True + async def renew_item_lease(self, *_args, **_kwargs): + return {"valid": True, "cancel_requested": False} + async def finalize_item(self, *_args, **kwargs): self.finalized = kwargs return True @@ -176,7 +179,10 @@ async def test_execute_item_polls_completion_without_waiting_for_lease_renewal(m raise AssertionError("a task that completes between polls need not expose running") async def renew_item_lease(self, *_args, **_kwargs): - raise AssertionError("short completion must not wait for lease renewal") + # Exactly one pre-launch revalidation is expected; the poll loop + # must not renew for a task that completes between polls. + self.renews = getattr(self, "renews", 0) + 1 + return {"valid": True, "cancel_requested": False} async def finalize_item(self, *_args, **kwargs): self.finalized = kwargs @@ -220,6 +226,7 @@ async def test_execute_item_polls_completion_without_waiting_for_lease_renewal(m assert repository.finalized is not None assert repository.finalized["result"] == "fast result" + assert repository.renews == 1 @pytest.mark.asyncio @@ -254,6 +261,9 @@ async def test_executor_admission_failure_requeues_instead_of_finalizing(monkeyp } ] + async def renew_item_lease(self, *_args, **_kwargs): + return {"valid": True, "cancel_requested": False} + async def requeue_item_after_admission_failure(self, item_id, **kwargs): self.requeued = (item_id, kwargs) return True @@ -290,3 +300,115 @@ async def test_executor_admission_failure_requeues_instead_of_finalizing(monkeyp assert repository.requeued is not None assert repository.requeued[0] == "item-1" assert repository.finalized is False + + +@pytest.mark.asyncio +async def test_cancel_during_tool_assembly_skips_launch(monkeypatch, tmp_path) -> None: + """A batch cancelled while tool assembly is blocked must not launch. + + Regression (review of 249dba82): ``_execute_item()`` called + ``executor.execute_async()`` unconditionally after assembly, so + ``cancel_batch()`` landing while assembly was blocked in the worker thread + terminalized the durable item but could not stop the not-yet-started + execution, and the orphaned launch still invoked the model. + """ + import threading + from datetime import UTC, datetime + + from deerflow.config.database_config import DatabaseConfig + from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config + from deerflow.persistence.subagent_batches import SubagentBatchRepository + + await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) + try: + repository = SubagentBatchRepository(get_session_factory()) + await repository.create_batch( + batch_id="batch-1", + user_id="user-1", + thread_id="thread-1", + run_id="run-1", + tool_call_id="call-1", + submission_key="run-1:call-1", + title="Cancelled during assembly", + subagent_type="general-purpose", + items=[{"key": "record-1", "prompt": "Process record 1"}], + max_live_items=2, + max_running_items=1, + max_attempts=2, + execution_spec={ + "subagent_config": { + "name": "general-purpose", + "description": "test", + "system_prompt": "sys", + }, + }, + ) + service = SubagentBatchService( + repository=repository, + config=SubagentBatchesConfig(lease_seconds=60), + runtime_config=SubagentRuntimeConfig(max_running=3), + app_config=SimpleNamespace(), + ) + claimed = await repository.claim_items( + now=datetime.now(UTC), + lease_owner=service._lease_owner, + lease_seconds=60, + limit=10, + ) + assert len(claimed) == 1 + item = claimed[0] + assert item["batch"]["id"] == "batch-1" + + assembly_started = threading.Event() + assembly_release = threading.Event() + + def _blocking_assembly(**_kwargs): + # Real blocking wait in the assembly-pool worker: parks assembly + # until the test has cancelled the batch. + assembly_started.set() + assembly_release.wait(timeout=15) + return [] + + monkeypatch.setattr("deerflow.tools.get_available_tools", _blocking_assembly) + monkeypatch.setattr(service_module, "SubagentStatus", FakeStatus) + monkeypatch.setattr(service_module, "resolve_subagent_model_name", lambda *_a, **_k: "test-model") + monkeypatch.setattr(service_module, "request_cancel_background_task", lambda _execution_id: None) + + launched: list[str] = [] + + class _Executor: + def __init__(self, **_kwargs) -> None: + pass + + def execute_async(self, _prompt, task_id=None): + launched.append(task_id or "generated") + return task_id or "generated" + + monkeypatch.setattr(service_module, "SubagentExecutor", _Executor) + monkeypatch.setattr( + service_module, + "get_background_task_result", + lambda _execution_id: SimpleNamespace( + status=FakeStatus.COMPLETED, + result="done", + error=None, + stop_reason=None, + token_usage_records=[], + ), + ) + + item_task = asyncio.create_task(service._execute_item(item)) + assert await asyncio.to_thread(assembly_started.wait, 15) + + cancelled = await service.cancel_batch(batch_id="batch-1", user_id="user-1") + assert cancelled is not None + + assembly_release.set() + await asyncio.wait_for(item_task, timeout=10) + + assert launched == [], "cancelled work must not launch after assembly" + batch = await repository.get_batch("batch-1", user_id="user-1") + assert batch is not None + assert batch["counts"]["cancelled"] == 1 + finally: + await close_engine() diff --git a/backend/tests/test_thread_regenerate_prepare.py b/backend/tests/test_thread_regenerate_prepare.py index 59eba4ff9..1f8a324d8 100644 --- a/backend/tests/test_thread_regenerate_prepare.py +++ b/backend/tests/test_thread_regenerate_prepare.py @@ -145,6 +145,7 @@ class FakeAccessor: @pytest.fixture(autouse=True) def _patch_checkpoint_accessor(monkeypatch): + from app.gateway import services from app.gateway.routers import thread_runs def build_accessor(request, *, thread_id, assistant_id=None, checkpoint_id=None): @@ -156,7 +157,7 @@ def _patch_checkpoint_accessor(monkeypatch): async def build_thread_accessor(request, *, thread_id, checkpoint_id=None): return build_accessor(request, thread_id=thread_id, checkpoint_id=checkpoint_id) - monkeypatch.setattr(thread_runs, "build_checkpoint_state_accessor", build_accessor) + monkeypatch.setattr(services, "build_checkpoint_state_accessor", build_accessor) monkeypatch.setattr(thread_runs, "build_thread_checkpoint_state_accessor", build_thread_accessor) @@ -195,6 +196,7 @@ def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"): def test_run_wait_readers_return_materialized_final_values() -> None: + from app.gateway import services from app.gateway.routers import runs, thread_runs snapshot = SimpleNamespace( @@ -237,7 +239,7 @@ def test_run_wait_readers_return_materialized_final_values() -> None: patch.object(thread_runs, "get_run_manager", return_value=object()), patch.object(thread_runs, "start_run", AsyncMock(return_value=record)), patch.object( - thread_runs, + services, "build_checkpoint_state_accessor", create=True, return_value=(accessor, snapshot.config), @@ -246,7 +248,7 @@ def test_run_wait_readers_return_materialized_final_values() -> None: patch.object(runs, "get_run_manager", return_value=object()), patch.object(runs, "start_run", AsyncMock(return_value=record)), patch.object( - runs, + services, "build_checkpoint_state_accessor", create=True, return_value=(accessor, snapshot.config), @@ -263,6 +265,7 @@ def test_run_wait_readers_return_materialized_final_values() -> None: def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None: + from app.gateway import services from app.gateway.routers import runs, thread_runs snapshot = SimpleNamespace( @@ -291,7 +294,7 @@ def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None: patch.object(thread_runs, "get_run_manager", return_value=object()), patch.object(thread_runs, "start_run", AsyncMock(return_value=record)), patch.object( - thread_runs, + services, "build_checkpoint_state_accessor", return_value=(accessor, snapshot.config), ), @@ -299,7 +302,7 @@ def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None: patch.object(runs, "get_run_manager", return_value=object()), patch.object(runs, "start_run", AsyncMock(return_value=record)), patch.object( - runs, + services, "build_checkpoint_state_accessor", return_value=(accessor, snapshot.config), ), @@ -317,6 +320,7 @@ def test_run_wait_readers_preserve_terminal_error_without_checkpoint() -> None: @pytest.mark.parametrize("route_name", ["thread", "stateless"]) def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(route_name: str) -> None: + from app.gateway import services from app.gateway.routers import runs, thread_runs record = SimpleNamespace( @@ -336,7 +340,7 @@ def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(ro patch.object(thread_runs, "get_run_manager", return_value=object()), patch.object(thread_runs, "start_run", AsyncMock(return_value=record)), patch.object( - thread_runs, + services, "build_checkpoint_state_accessor", side_effect=RuntimeError("graph construction failed"), ), @@ -348,7 +352,7 @@ def test_run_wait_readers_preserve_terminal_error_when_accessor_builder_fails(ro patch.object(runs, "get_run_manager", return_value=object()), patch.object(runs, "start_run", AsyncMock(return_value=record)), patch.object( - runs, + services, "build_checkpoint_state_accessor", side_effect=RuntimeError("graph construction failed"), ), diff --git a/backend/tests/test_thread_run_idempotency.py b/backend/tests/test_thread_run_idempotency.py index cd4538100..8367756f3 100644 --- a/backend/tests/test_thread_run_idempotency.py +++ b/backend/tests/test_thread_run_idempotency.py @@ -12,6 +12,7 @@ from _router_auth_helpers import call_unwrapped, make_authed_test_app from fastapi import HTTPException from fastapi.testclient import TestClient +from app.gateway import services from app.gateway.auth.models import User from app.gateway.routers import thread_runs from app.gateway.run_models import RunCreateRequest @@ -215,7 +216,7 @@ def test_wait_reused_store_only_run_does_not_return_stale_checkpoint(monkeypatch monkeypatch.setattr(thread_runs, "start_run", fake_start_run) monkeypatch.setattr( - thread_runs, + services, "build_checkpoint_state_accessor", lambda *args, **kwargs: (SimpleNamespace(aget=fake_aget), {}), ) @@ -262,7 +263,7 @@ def test_wait_reused_completed_run_does_not_return_later_checkpoint(monkeypatch) monkeypatch.setattr(thread_runs, "start_run", fake_start_run) monkeypatch.setattr( - thread_runs, + services, "build_checkpoint_state_accessor", lambda *args, **kwargs: (SimpleNamespace(aget=fake_aget), {}), ) @@ -320,7 +321,7 @@ async def test_wait_original_request_keeps_checkpoint_when_retry_overlaps(): patch.object(thread_runs, "get_stream_bridge", return_value=bridge), patch.object(thread_runs, "get_run_manager", return_value=MagicMock()), patch.object( - thread_runs, + services, "build_checkpoint_state_accessor", lambda *args, **kwargs: (SimpleNamespace(aget=fake_aget), {}), ), @@ -386,7 +387,7 @@ async def test_wait_peer_refreshes_status_after_owner_completes(): patch.object(thread_runs, "get_stream_bridge", return_value=bridge), patch.object(thread_runs, "get_run_manager", return_value=peer), patch.object( - thread_runs, + services, "build_checkpoint_state_accessor", side_effect=AssertionError("reused wait must not read latest checkpoint"), ), @@ -769,7 +770,7 @@ async def test_wait_retry_after_later_run_does_not_return_later_checkpoint(monke monkeypatch.setattr(thread_runs, "start_run", fake_start_run) monkeypatch.setattr( - thread_runs, + services, "build_checkpoint_state_accessor", lambda *args, **kwargs: (SimpleNamespace(aget=fake_aget), {}), ) diff --git a/backend/tests/test_threads_router.py b/backend/tests/test_threads_router.py index b221d1b11..48173212e 100644 --- a/backend/tests/test_threads_router.py +++ b/backend/tests/test_threads_router.py @@ -269,7 +269,10 @@ def _patch_checkpoint_state_builder(monkeypatch): async def _mutation_boundary(request, *, thread_id, as_node, checkpoint_id=None): return _mutation_builder(request, thread_id=thread_id, as_node=as_node, checkpoint_id=checkpoint_id) - monkeypatch.setattr(threads, "build_checkpoint_state_accessor", _builder) + async def _abuild(request, *, thread_id, assistant_id=None, checkpoint_id=None): + return _builder(request, thread_id=thread_id, assistant_id=assistant_id, checkpoint_id=checkpoint_id) + + monkeypatch.setattr(threads, "abuild_checkpoint_state_accessor", _abuild) monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", _mutation_builder) monkeypatch.setattr(threads, "build_thread_checkpoint_state_accessor", _read_boundary) monkeypatch.setattr(threads, "build_thread_checkpoint_state_mutation_accessor", _mutation_boundary) @@ -1171,9 +1174,8 @@ def test_latest_thread_readers_use_materialized_snapshot_values() -> None: with ( patch( - "app.gateway.routers.threads.build_checkpoint_state_accessor", - create=True, - return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}), + "app.gateway.routers.threads.abuild_checkpoint_state_accessor", + new=AsyncMock(return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})), ), patch( "app.gateway.routers.threads.build_thread_checkpoint_state_accessor", @@ -1228,8 +1230,8 @@ def test_get_thread_status_uses_raw_pending_writes_for_materialized_checkpoint() with ( patch( - "app.gateway.routers.threads.build_checkpoint_state_accessor", - return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}), + "app.gateway.routers.threads.abuild_checkpoint_state_accessor", + new=AsyncMock(return_value=(accessor, {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})), ), TestClient(app) as client, ): @@ -1272,8 +1274,8 @@ def test_get_thread_preserves_metadata_status_without_checkpoint(stored_status: with ( patch( - "app.gateway.routers.threads.build_checkpoint_state_accessor", - return_value=(accessor, snapshot.config), + "app.gateway.routers.threads.abuild_checkpoint_state_accessor", + new=AsyncMock(return_value=(accessor, snapshot.config)), ), TestClient(app) as client, ): @@ -2745,7 +2747,7 @@ def test_branch_thread_uses_materialized_history_and_overwrites_fresh_seed(monke } } - monkeypatch.setattr(threads, "build_checkpoint_state_accessor", build_accessor) + monkeypatch.setattr(threads, "abuild_checkpoint_state_accessor", AsyncMock(side_effect=build_accessor)) monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", build_mutation_accessor) with TestClient(app) as client: @@ -2864,7 +2866,7 @@ def test_branch_thread_preserves_unlinked_legacy_histories( } } - monkeypatch.setattr(threads, "build_checkpoint_state_accessor", build_accessor) + monkeypatch.setattr(threads, "abuild_checkpoint_state_accessor", AsyncMock(side_effect=build_accessor)) monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", build_mutation_accessor) with TestClient(app) as client: @@ -2993,7 +2995,7 @@ def test_branch_thread_real_mutation_graph_finishes_without_scheduling(monkeypat } real_mutation_builder = gateway_services.build_checkpoint_state_mutation_accessor - monkeypatch.setattr(threads, "build_checkpoint_state_accessor", source_builder) + monkeypatch.setattr(threads, "abuild_checkpoint_state_accessor", AsyncMock(side_effect=source_builder)) monkeypatch.setattr( threads, "build_checkpoint_state_mutation_accessor", @@ -3066,7 +3068,7 @@ def _wire_extension_agent(monkeypatch, app, checkpointer, mode): ctx = SimpleNamespace(checkpointer=checkpointer, store=None, checkpoint_channel_mode=mode, app_config=None) monkeypatch.setattr(gateway_services, "get_run_context", lambda _request: ctx) monkeypatch.setattr(gateway_services, "resolve_agent_factory", selective_factory) - monkeypatch.setattr(threads, "build_checkpoint_state_accessor", gateway_services.build_checkpoint_state_accessor) + monkeypatch.setattr(threads, "abuild_checkpoint_state_accessor", gateway_services.abuild_checkpoint_state_accessor) monkeypatch.setattr(threads, "build_checkpoint_state_mutation_accessor", gateway_services.build_checkpoint_state_mutation_accessor) monkeypatch.setattr(threads, "build_thread_checkpoint_state_accessor", gateway_services.build_thread_checkpoint_state_accessor) monkeypatch.setattr(threads, "build_thread_checkpoint_state_mutation_accessor", gateway_services.build_thread_checkpoint_state_mutation_accessor)