diff --git a/README.md b/README.md index 63942d95c..7139b1243 100644 --- a/README.md +++ b/README.md @@ -547,6 +547,11 @@ warning if the persisted-store contract no longer matches its expectations. The documented command uses LangGraph's file-based custom-app loader, which is also covered directly by DeerFlow's regression tests. +Standalone runs using `if_not_exists="create"` retain config and run metadata +on the newly created thread, including searchable tags; run metadata takes +precedence for duplicate keys. Thread ownership and MCP incarnation remain +server-owned, and later runs do not replace the thread's creation metadata. + For workflows that invoke `backend/langgraph.json` through LangGraph Studio or a direct LangGraph Server, DeerFlow consumes the authenticated identity published by that runtime and uses it for custom-agent configuration/SOUL, user @@ -612,6 +617,7 @@ already received by the browser, without an additional secret-redaction layer. DeerFlow supports configurable MCP servers and skills to extend its capabilities. For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`). For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`; durable background-task calls honor the same setting for HTTP/SSE servers as well. +Ordinary `task` subagents retain the parent run's captured thread incarnation for MCP calls, including legacy threads, so delegation preserves the same lifecycle scope. MCP tool names are prefixed with `_` by default to prevent collisions across servers. If a server already namespaces its own tools, set `tool_name_prefix: false` on that server in `extensions_config.json` to keep the original names. Disable the prefix only when the resulting names remain unique across all enabled servers. Signed-in users' notification toggle, default model, conversation mode, and reasoning effort are saved to their account and restored on other browsers or after clearing browser storage. Browser notification permission still needs to be granted on each device. Changes retry after network failures; unsent changes survive a reload in the same tab. Concurrent edits to different fields are preserved; for the same field, the last server write wins. Existing unscoped browser preferences are not uploaded automatically because they have no account owner; reselect those settings once after upgrading. Static demos and auth-disabled development keep browser-local settings. Thread-specific model overrides and other display preferences remain local. diff --git a/backend/app/gateway/AGENTS.md b/backend/app/gateway/AGENTS.md index 341857948..da62ee006 100644 --- a/backend/app/gateway/AGENTS.md +++ b/backend/app/gateway/AGENTS.md @@ -1,5 +1,7 @@ ### Gateway API (`app/gateway/`) +Studio retains sanitized creation metadata. + Capability Center's `business` adapter validates only the bundled provider's credential fields and creates a normal MCP connection. The MCP API accepts the exact isolated interpreter/module/provider launcher generated by diff --git a/backend/app/gateway/langgraph_auth.py b/backend/app/gateway/langgraph_auth.py index 8143d5e3f..d93acd13d 100644 --- a/backend/app/gateway/langgraph_auth.py +++ b/backend/app/gateway/langgraph_auth.py @@ -15,13 +15,21 @@ Two layers: """ import secrets +from contextvars import ContextVar +from uuid import uuid4 from langgraph_sdk import Auth +from starlette.exceptions import HTTPException from app.gateway.auth.errors import TokenError from app.gateway.auth.jwt import decode_token from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID, is_auth_disabled from app.gateway.deps import get_local_provider +from deerflow.mcp_scope import ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, + is_valid_thread_incarnation, +) auth = Auth() @@ -32,6 +40,160 @@ _STUDIO_USER_TYPE = getattr(Auth.types, "StudioUser", None) # Methods that require CSRF validation (state-changing per RFC 7231). _CSRF_METHODS = frozenset({"POST", "PUT", "DELETE", "PATCH"}) +_allow_thread_incarnation_write: ContextVar[bool] = ContextVar( + "deerflow_allow_standalone_thread_incarnation_write", + default=False, +) +_MISSING = object() + + +def _metadata(value: dict) -> dict: + metadata = value.get("metadata") + if not isinstance(metadata, dict): + metadata = {} + value["metadata"] = metadata + return metadata + + +def _scrub_run_incarnation(value: dict) -> None: + lifecycle_keys = ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, + ) + metadata = _metadata(value) + for key in lifecycle_keys: + metadata.pop(key, None) + kwargs = value.get("kwargs") + if not isinstance(kwargs, dict): + kwargs = {} + value["kwargs"] = kwargs + context = kwargs.get("context") + if not isinstance(context, dict): + context = {} + kwargs["context"] = context + runtime_keys = (*lifecycle_keys, "user_id", "thread_id", "run_id") + for key in runtime_keys: + context.pop(key, None) + config = kwargs.get("config") + if not isinstance(config, dict): + config = {} + kwargs["config"] = config + for key in ("context", "metadata", "configurable"): + section = config.get(key) + if isinstance(section, dict): + for protected_key in runtime_keys: + section.pop(protected_key, None) + + +async def _read_standalone_thread(thread_id, ctx) -> dict | None: + from langgraph_runtime.database import connect + from langgraph_runtime.ops import Threads + + try: + async with connect() as conn: + rows = await Threads.get(conn, thread_id, ctx=ctx) + return await anext(rows, None) + except HTTPException as exc: + if exc.status_code == 404: + return None + raise + + +async def _ensure_standalone_thread_incarnation( + thread_id, + ctx, + *, + create_if_missing: bool, + creation_metadata: dict | None = None, +) -> str | None | object: + thread = await _read_standalone_thread(thread_id, ctx) + if thread is None: + if not create_if_missing: + return _MISSING + incarnation = uuid4().hex + else: + metadata = thread.get("metadata") + stored = metadata.get(THREAD_INCARNATION_CONTEXT_KEY, _MISSING) if isinstance(metadata, dict) else _MISSING + if stored is _MISSING: + # Match the rollout contract used by the Gateway and embedded + # runtime: a persisted pre-incarnation thread remains in the + # explicit legacy generation. A read/patch backfill cannot fence + # deletion plus same-ID recreation through LangGraph's public API. + return None + if not is_valid_thread_incarnation(stored): + raise RuntimeError("Standalone LangGraph thread has an invalid incarnation") + return stored + + from langgraph_runtime.database import connect + from langgraph_runtime.ops import Threads + + token = _allow_thread_incarnation_write.set(True) + try: + async with connect() as conn: + rows = await Threads.put( + conn, + thread_id, + metadata={**(creation_metadata or {}), THREAD_INCARNATION_CONTEXT_KEY: incarnation}, + if_exists="do_nothing", + ctx=ctx, + ) + if await anext(rows, None) is None: + raise RuntimeError("Standalone LangGraph thread incarnation was not persisted") + finally: + _allow_thread_incarnation_write.reset(token) + + persisted = await _read_standalone_thread(thread_id, ctx) + if persisted is None: + raise RuntimeError("Standalone LangGraph thread incarnation could not be verified") + persisted_metadata = persisted.get("metadata") + value = persisted_metadata.get(THREAD_INCARNATION_CONTEXT_KEY, _MISSING) if isinstance(persisted_metadata, dict) else _MISSING + if value is _MISSING: + # A mixed-version peer may have won the do-nothing create with a + # legacy thread. Keep that persisted lifecycle on the legacy scope. + return None + if not is_valid_thread_incarnation(value): + raise RuntimeError("Standalone LangGraph thread incarnation could not be verified") + return value + + +async def _bind_standalone_run_incarnation(ctx, value: dict) -> None: + _scrub_run_incarnation(value) + thread_id = value.get("thread_id") + if thread_id is None: + incarnation: str | None = None + else: + # Pre-creation bypasses LangGraph's implicit-create metadata merge. + # Preserve its precedence using only the already-sanitized metadata; + # the helper adds the server incarnation last and never updates an + # existing thread (including a concurrent creation winner). + config_metadata = value["kwargs"]["config"].get("metadata") + creation_metadata = { + **(config_metadata if isinstance(config_metadata, dict) else {}), + **value["metadata"], + } + incarnation = await _ensure_standalone_thread_incarnation( + thread_id, + ctx, + create_if_missing=value.get("if_not_exists") == "create", + creation_metadata=creation_metadata, + ) + if incarnation is _MISSING: + return + assert is_valid_thread_incarnation(incarnation) + + context = value["kwargs"]["context"] + context["user_id"] = ctx.user.identity + if thread_id is not None: + context["thread_id"] = str(thread_id) + else: + context.pop("thread_id", None) + run_id = value.get("run_id") + if run_id is not None: + context["run_id"] = str(run_id) + else: + context.pop("run_id", None) + context[THREAD_INCARNATION_CONTEXT_KEY] = incarnation + context[THREAD_INCARNATION_METADATA_GUARD_KEY] = True def _check_csrf(request) -> None: @@ -131,8 +293,18 @@ async def add_owner_filter(ctx: Auth.types.AuthContext, value: dict): # boundary. The standalone pre-runtime persistence repair also scrubs this # marker from legacy active rows and their version history before normal # version selection becomes available. - metadata = value.setdefault("metadata", {}) + metadata = _metadata(value) metadata["user_id"] = ctx.user.identity + if ctx.resource == "threads" and ctx.action == "create_run": + await _bind_standalone_run_incarnation(ctx, value) + elif ctx.resource == "threads": + if ctx.action == "create": + if not _allow_thread_incarnation_write.get(): + metadata[THREAD_INCARNATION_CONTEXT_KEY] = uuid4().hex + elif not is_valid_thread_incarnation(metadata.get(THREAD_INCARNATION_CONTEXT_KEY, _MISSING)): + raise RuntimeError("Standalone LangGraph internal thread create has an invalid incarnation") + elif ctx.action == "update" and not _allow_thread_incarnation_write.get(): + metadata.pop(THREAD_INCARNATION_CONTEXT_KEY, None) if ctx.resource == "assistants" and ctx.action in {"create", "update"}: metadata["created_by"] = "user" diff --git a/backend/app/gateway/routers/mcp_tasks.py b/backend/app/gateway/routers/mcp_tasks.py index 3d0e5a2eb..841d7e717 100644 --- a/backend/app/gateway/routers/mcp_tasks.py +++ b/backend/app/gateway/routers/mcp_tasks.py @@ -7,7 +7,8 @@ from typing import Any from fastapi import APIRouter, HTTPException, Query, Request from app.gateway.authz import require_permission -from app.gateway.deps import get_current_user, get_mcp_task_repo, get_mcp_task_service +from app.gateway.deps import get_current_user, get_mcp_task_repo, get_mcp_task_service, get_thread_store +from deerflow.mcp_scope import is_valid_thread_incarnation from deerflow.utils.thread_id import ThreadId router = APIRouter(prefix="/api/threads/{thread_id}/mcp-tasks", tags=["mcp-tasks"]) @@ -63,6 +64,27 @@ async def _current_user_id(request: Request) -> str: return user_id +async def _current_thread_incarnation( + request: Request, + *, + thread_id: str, + user_id: str, +) -> str | None: + """Capture the server-owned thread generation for the repository CAS.""" + thread_store = get_thread_store(request) + record = await thread_store.get(thread_id, user_id=user_id) + if record is None: + unscoped = await thread_store.get(thread_id, user_id=None) + if unscoped is not None and unscoped.get("user_id") is None: + record = unscoped + if record is None: + raise HTTPException(status_code=404, detail="Thread not found") + incarnation = record.get("incarnation") + if not is_valid_thread_incarnation(incarnation): + raise HTTPException(status_code=404, detail="Thread not found") + return incarnation + + @router.get("") @require_permission("threads", "read", owner_check=True) async def list_mcp_tasks( @@ -73,9 +95,15 @@ async def list_mcp_tasks( repository = get_mcp_task_repo(request) service = get_mcp_task_service(request) user_id = await _current_user_id(request) + thread_incarnation = await _current_thread_incarnation( + request, + thread_id=thread_id, + user_id=user_id, + ) records = await repository.list_by_thread( thread_id, user_id=user_id, + thread_incarnation=thread_incarnation, limit=limit, ) threshold = service.tracking_degraded_after_errors @@ -92,8 +120,18 @@ async def get_mcp_task( repository = get_mcp_task_repo(request) service = get_mcp_task_service(request) user_id = await _current_user_id(request) - record = await repository.get(task_id, user_id=user_id) - if record is None or record["thread_id"] != thread_id: + thread_incarnation = await _current_thread_incarnation( + request, + thread_id=thread_id, + user_id=user_id, + ) + record = await repository.get( + task_id, + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ) + if record is None: raise HTTPException(status_code=404, detail="MCP task not found") return _detail( record, @@ -110,6 +148,11 @@ async def cancel_mcp_task( ) -> dict[str, Any]: service = get_mcp_task_service(request) user_id = await _current_user_id(request) + thread_incarnation = await _current_thread_incarnation( + request, + thread_id=thread_id, + user_id=user_id, + ) if not getattr(request.app.state, "mcp_tasks_available", False): # The service exists whenever SQL persistence is configured, but the # background loop that owns the remote cancel call only runs when @@ -120,6 +163,7 @@ async def cancel_mcp_task( task_id=task_id, thread_id=thread_id, user_id=user_id, + thread_incarnation=thread_incarnation, ) if record is None: raise HTTPException(status_code=404, detail="MCP task not found") diff --git a/backend/app/gateway/services.py b/backend/app/gateway/services.py index 3db08f8b6..38d759127 100644 --- a/backend/app/gateway/services.py +++ b/backend/app/gateway/services.py @@ -49,6 +49,10 @@ from deerflow.config.agents_config import load_agent_config from deerflow.config.app_config import get_app_config from deerflow.config.database_config import resolve_checkpoint_graph_cache_max from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_KEY, KNOWLEDGE_SCOPE_RUNTIME_KEY +from deerflow.mcp_scope import ( + THREAD_INCARNATION_METADATA_GUARD_KEY, + is_valid_thread_incarnation, +) from deerflow.projects.context import PROJECT_CONTEXT_MESSAGE_MARKER, resolve_project_context from deerflow.runtime import ( END_SENTINEL, @@ -206,7 +210,7 @@ async def _ensure_thread_metadata( *, owner_user_id: str | None, require_existing_thread: bool = False, -) -> None: +) -> dict[str, Any]: """Ensure an admitted run's thread exists without delaying task attachment.""" thread_store = run_ctx.thread_store existing = await thread_store.get(record.thread_id) @@ -231,12 +235,12 @@ async def _ensure_thread_metadata( # /threads/{id}/move — so the key must not persist either. if key not in (DEERFLOW_TRACE_METADATA_KEY, THREAD_PROJECT_METADATA_KEY) } - await thread_store.create( + existing = await thread_store.create( record.thread_id, assistant_id=record.assistant_id, metadata=metadata, ) - return + return existing async def _terminal_record_stream_missing(bridge: StreamBridge, record: RunRecord) -> bool: @@ -594,6 +598,7 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: frozenset[str] = ( "__run_tool_progress_recorder", "langgraph_auth_user", "langgraph_auth_user_id", + THREAD_INCARNATION_METADATA_GUARD_KEY, # Server-owned pinned project snapshot (spec §7.1): resolved once # at admission from threads_meta; a client-supplied value must # never survive in either run-config section. @@ -1916,6 +1921,7 @@ async def start_run( abort_task = asyncio.create_task(record.abort_event.wait()) metadata_failure_logged = False metadata_failure: Exception | None = None + metadata_record: dict[str, Any] | None = None try: done, _ = await asyncio.wait( (metadata_task, abort_task), @@ -1924,7 +1930,7 @@ async def start_run( ) if metadata_task in done: try: - metadata_task.result() + metadata_record = metadata_task.result() except asyncio.CancelledError: pass except Exception as exc: @@ -1946,8 +1952,20 @@ async def start_run( metadata_failure = TimeoutError("Timed out verifying existing thread metadata") finally: if metadata_task.done(): - if not metadata_failure_logged: - _log_thread_metadata_task_result(metadata_task, thread_id=thread_id) + if metadata_record is None and not metadata_failure_logged: + try: + metadata_record = metadata_task.result() + except asyncio.CancelledError: + pass + except Exception as exc: + metadata_failure_logged = True + metadata_failure = exc + logger.warning( + "Failed to ensure thread_meta for %s%s", + sanitize_log_param(thread_id), + "" if require_existing_thread else " (non-fatal)", + exc_info=True, + ) else: metadata_task.cancel() metadata_task.add_done_callback( @@ -1968,6 +1986,28 @@ async def start_run( # or strict verification failure: # its startup barrier is the single path that turns pending # cancellation into no-agent-construction plus publish_end. + incarnation_kwargs: dict[str, str | None] = {} + if metadata_record is None: + if not record.abort_event.is_set(): + logger.warning( + "Thread metadata for %s is unavailable; MCP access will fail closed", + sanitize_log_param(thread_id), + ) + else: + if "incarnation" not in metadata_record: + logger.warning( + "Thread metadata for %s has no incarnation; MCP access will fail closed", + sanitize_log_param(thread_id), + ) + else: + incarnation = metadata_record["incarnation"] + if is_valid_thread_incarnation(incarnation): + incarnation_kwargs["thread_incarnation"] = incarnation + else: + logger.warning( + "Thread metadata for %s has an invalid incarnation; MCP access will fail closed", + sanitize_log_param(thread_id), + ) await run_agent( bridge, run_mgr, @@ -1981,6 +2021,7 @@ async def start_run( interrupt_before=body.interrupt_before, interrupt_after=body.interrupt_after, knowledge_scope=admitted_knowledge_scope, + **incarnation_kwargs, ) try: diff --git a/backend/app/mcp_tasks/service.py b/backend/app/mcp_tasks/service.py index 16a414950..24a7fbd9a 100644 --- a/backend/app/mcp_tasks/service.py +++ b/backend/app/mcp_tasks/service.py @@ -310,6 +310,7 @@ class McpTaskService: thread_id=request.thread_id, server_name=request.server_name, remote_task_id=submission.remote_task_id, + thread_incarnation=request.thread_incarnation, driver_data=driver_data, ) try: @@ -321,6 +322,7 @@ class McpTaskService: task_id=local_task_id, user_id=request.user_id, thread_id=request.thread_id, + expected_thread_incarnation=request.thread_incarnation, run_id=request.run_id, tool_call_id=request.tool_call_id, server_name=request.server_name, @@ -692,12 +694,14 @@ class McpTaskService: *, thread_id: str, user_id: str, + thread_incarnation: str | None, limit: int = 50, active_only: bool = False, ) -> list[dict[str, Any]]: return await self._repository.list_by_thread( thread_id, user_id=user_id, + thread_incarnation=thread_incarnation, limit=limit, active_only=active_only, ) @@ -708,12 +712,14 @@ class McpTaskService: task_id: str, thread_id: str, user_id: str, + thread_incarnation: str | None, now: datetime | None = None, ) -> dict[str, Any] | None: return await self._repository.request_cancel( task_id, user_id=user_id, thread_id=thread_id, + thread_incarnation=thread_incarnation, requested_at=now or datetime.now(UTC), ) @@ -722,9 +728,15 @@ class McpTaskService: *, thread_id: str, user_id: str, + thread_incarnation: str | None, task: str | None = None, ) -> dict[str, Any]: - active = await self.list_tasks(thread_id=thread_id, user_id=user_id, active_only=True) + active = await self.list_tasks( + thread_id=thread_id, + user_id=user_id, + thread_incarnation=thread_incarnation, + active_only=True, + ) if task: normalized = task.casefold().strip() matches = [item for item in active if item["id"] == task or str(item.get("task_name") or "").casefold() == normalized] @@ -739,6 +751,7 @@ class McpTaskService: task_id=matches[0]["id"], thread_id=thread_id, user_id=user_id, + thread_incarnation=thread_incarnation, ) if result is None: raise LookupError("The selected background task no longer exists") diff --git a/backend/packages/harness/deerflow/agents/lead_agent/agent.py b/backend/packages/harness/deerflow/agents/lead_agent/agent.py index 2ca6a983d..e8a215866 100644 --- a/backend/packages/harness/deerflow/agents/lead_agent/agent.py +++ b/backend/packages/harness/deerflow/agents/lead_agent/agent.py @@ -1117,6 +1117,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le middleware=normalize_middleware_state_schemas(middlewares, mode), system_prompt=system_prompt, state_schema=get_thread_state_schema(mode), + context_schema=dict, ) return _complete_assembly( config=config, @@ -1249,6 +1250,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le middleware=normalize_middleware_state_schemas(middlewares, mode), system_prompt=system_prompt, state_schema=get_thread_state_schema(mode), + context_schema=dict, ) return _complete_assembly( config=config, diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index f4af2c220..3776707da 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -52,6 +52,7 @@ from deerflow.config.extensions_config import ( ) from deerflow.config.paths import get_paths from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig +from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY from deerflow.models import create_chat_model from deerflow.runtime import CheckpointStateAccessor from deerflow.runtime.checkpoint_mode import ( @@ -240,6 +241,7 @@ class DeerFlowClient: self._available_skills = set(available_skills) if available_skills is not None else None self._middlewares = list(middlewares) if middlewares else [] self._environment = environment + self._thread_incarnations: dict[str, str] = {} # Lazy agent — created on first call, recreated when config changes. self._agent = None @@ -948,7 +950,15 @@ class DeerFlowClient: config["callbacks"] = [*existing_callbacks, *tracing_callbacks] run_id = str(uuid.uuid4()) - context: dict[str, Any] = {"thread_id": thread_id, "run_id": run_id} + thread_incarnations = getattr(self, "_thread_incarnations", None) + if thread_incarnations is None: + thread_incarnations = self._thread_incarnations = {} + thread_incarnation = thread_incarnations.setdefault(thread_id, uuid.uuid4().hex) + context: dict[str, Any] = { + "thread_id": thread_id, + "run_id": run_id, + THREAD_INCARNATION_CONTEXT_KEY: thread_incarnation, + } for key in _EMBEDDED_AUTHORIZATION_CONTEXT_KEYS: if key in kwargs: context[key] = kwargs[key] diff --git a/backend/packages/harness/deerflow/mcp/AGENTS.md b/backend/packages/harness/deerflow/mcp/AGENTS.md index b0d652b7f..030d5fd2d 100644 --- a/backend/packages/harness/deerflow/mcp/AGENTS.md +++ b/backend/packages/harness/deerflow/mcp/AGENTS.md @@ -3,7 +3,7 @@ - Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management - **Long-running task foundation**: `mcp/tasks/` defines the protocol-neutral `McpTaskDriver` contract and normalized `TaskSnapshot` states (`submitted`, `working`, `input_required`, `completed`, `failed`, `cancelled`). A driver-supplied `poll_after_seconds` must be a finite positive number, validated at the `TaskSnapshot` boundary so every driver is held to the same invariant rather than each one guarding the consumer that turns the interval into a `timedelta`. `persistence/mcp_tasks/` owns the durable remote-handle mapping, poll schedule, notification state, lease owner, and separate consecutive poll/delivery error counters; `app/mcp_tasks/McpTaskService` performs status, cancellation, and notification work outside the Agent/LLM loop. Notification retries keep their idempotency attempt separate from the delivery-failure count, use capped exponential backoff, and stop after five failures; strict existing-thread admission dead-letters a deleted/mismatched target immediately. A status result is applied only when the worker still owns an unexpired lease, so a stale result cannot be written after expiry even before another worker reclaims the row. Every claim carries a fresh per-claim token, and every poll/cancel/notification release or snapshot apply is fenced by the owner plus that token, so a release that completes after the same service reclaims the row cannot clear the newer lease. Poll, cancel, and notification claims also have separate phase-level single-flight owners that protect only an ambiguous claim outcome: after an uncancelled timeout or caller cancellation, later scans do not start another claim while the outcome is unknown. Once the claim resolves, the phase owner is released immediately, and any returned rows are released through bounded, service-owned background work; per-claim token fencing prevents a late release from mutating a newer claim generation, so a stuck release does not stall the whole phase. While the claim outcome remains unresolved, later scans skip the phase and emit a warning rather than overlapping an ambiguous database claim. Once the claim resolves, the handoff may continue releasing returned rows in the background without blocking later scans. Lease expiry remains the cross-process crash fallback rather than the normal recovery for a live service's late claim. Poll timestamps and retry schedules are based on the remote call's completion time rather than the scan start. If submission succeeds but persistence fails or the caller is cancelled while persistence is in flight, the service best-effort cancels the remote task and preserves the original error or cancellation if that compensation also fails. The exact `uq_mcp_tasks_user_server_remote` conflict is different: an existing durable row already owns the remote handle, so the conflict surfaces without cancelling that tracked task. Unexpected per-task poll failures are isolated from sibling claims and remain recoverable through lease expiry; Gateway shutdown cancels the poller so a hung external status call cannot block process exit. Routine cancellation releases clear only the lease/token and preserve any pre-existing `last_cancel_error` or `notification_error`; cancellation itself is not persisted as a task failure diagnostic. Cancelling an in-flight poll releases only the owner- and per-claim-token-fenced lease and preserves its preclaim schedule and poll-failure state; real poll failures retain exponential backoff and tracking-degradation behavior. `input_required` and terminal states stop polling and become `notification_status=pending` for later Agent/UI delivery. Durable recovery requires a SQL database backend (`sqlite` or `postgres`); the in-memory backend leaves the repository/service unavailable. The runtime is startup-configured by `mcp_tasks` and disabled by default until a concrete driver is registered; this foundation does not alter ordinary MCP tool behavior on its own. - **Runtime availability boundary**: the installed process-local submitter is the source of truth for durable task-management tool exposure. `mcp_tasks` is startup-only; changing it on disk does not alter the live toolset until the Gateway restarts. -- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers..task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same `(server_name, user_id:thread_id)` stdio session scope; HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task. +- **Long-running ordinary task driver**: `extensions_config.json -> mcpServers..task_toolsets` binds exact raw submit/status/cancel names; one raw tool may occupy only one role across that server's groups. `mcp/tools.py` hides status/cancel and replaces submit with a wrapper that returns only the local task ID after persistence. `ordinary.py` reads only MCP `structuredContent`, maps remote `running` to `working`, and treats `error_code=task_not_found` or malformed structured output as permanent failure. A status call with `isError=true` is a retryable call failure: the first text content block is retained as a bounded diagnostic, while a permanent remote-task outcome must arrive in a normal result with structured `status=failed`. `task_tool_caller.py` restores the same canonical `(server_name, versioned user/thread/incarnation scope)` stdio session as ordinary calls; non-NULL incarnations use an unambiguous versioned encoding, while an explicitly captured legacy NULL retains the pre-activation `user_id:thread_id` scope. A runtime with a missing/invalid incarnation key fails closed. HTTP/SSE calls remain ephemeral, apply `session_init_timeout` to initialization and `tool_call_timeout` to task calls, and support server-level OAuth refresh outside an Agent run. `McpTaskService` exponentially backs off transient status/cancel errors without a maximum attempt count, derives API `tracking_degraded` from the consecutive-error threshold, keeps `input_required` on a slower poll, and caps finite positive remote poll hints at 24 hours. Task-enabled server runtime/binding configuration and `mcpInterceptors` are frozen to the Gateway startup snapshot; hot drift fails clearly before tool discovery can diverge from background calls, while presentation-only fields and non-task servers remain reloadable. Configured task toolsets fail startup when the runtime is disabled or persistence is memory. Users still cannot submit an answer back to an `input_required` remote task. - **Durable task payload bounds**: persisted task errors are capped at 4,000 characters. `input_required` and `result_artifact` must each serialize as valid JSON within 64 KiB; an invalid or oversized payload becomes a permanent protocol failure rather than being truncated and changing its semantics. Remote task IDs/task names are limited to 255 characters and task-enabled server names to 128, matching the SQL schema; an oversized submitted remote ID is rejected only after the Service has the handle so compensation cancellation still runs. Oversized results retain the existing bounded preview/truncation/artifact behavior. - **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()` - **Loop-isolated stdio sessions**: Both the live registry and in-flight creations @@ -36,9 +36,9 @@ `tools/builtins/tool_search.py::get_mcp_routing_hints_prompt_section`, which references `tool_search` when a hinted MCP tool is currently deferred; do not add a parallel routing middleware for PR1-style preference hints. -- **Stdio file outputs**: Persistent stdio sessions are scoped by `user_id:thread_id`. For stdio transports only, DeerFlow pins the subprocess default `cwd` to the thread workspace and `TMPDIR`/`TMP`/`TEMP` to `workspace/.mcp/tmp/`, unless the operator explicitly configured `cwd` or temp env values. `.mcp` is a DeerFlow-owned internal namespace: its temporary/debug files remain addressable when returned by a tool but are excluded from run workspace-change summaries — by directory name at any depth, consistent with the other reserved names in `EXCLUDED_DIR_NAMES` (`.git`, `node_modules`, …) and robust if a server ever creates a relative `.mcp` from a different cwd. Both launch paths pin it at the workspace root today. SSE/HTTP transports skip this filesystem prep entirely. +- **Stdio file outputs**: Persistent stdio sessions use the canonical versioned user/thread/incarnation scope (legacy NULL incarnations use the exact `user_id:thread_id` key). For stdio transports only, DeerFlow pins the subprocess default `cwd` to the thread workspace and `TMPDIR`/`TMP`/`TEMP` to `workspace/.mcp/tmp/`, unless the operator explicitly configured `cwd` or temp env values. `.mcp` is a DeerFlow-owned internal namespace: its temporary/debug files remain addressable when returned by a tool but are excluded from run workspace-change summaries — by directory name at any depth, consistent with the other reserved names in `EXCLUDED_DIR_NAMES` (`.git`, `node_modules`, …) and robust if a server ever creates a relative `.mcp` from a different cwd. Both launch paths pin it at the workspace root today. SSE/HTTP transports skip this filesystem prep entirely. - **Configured stdio working directory**: A non-empty `McpServerConfig.cwd` is forwarded by `build_server_params` for discovery, pooled Agent calls, and durable task calls. Omit the connection key when the setting is `None` or empty (including unresolved environment references), so the runtime's thread-workspace default still applies; HTTP/SSE connections must not receive it. `tests/test_mcp_cwd.py` exercises the JSON-config-to-subprocess path with relative entrypoints, relative tool inputs, and discovery/pool defaults for omitted, null, empty, and unresolved values. -- **Stdio disconnect recovery**: Ordinary Agent tool calls and durable task submit/status/cancel calls that receive the MCP SDK's explicit `Connection closed` error or an AnyIO closed-stream error evict only that `(server_name, user_id:thread_id)` session when the registered entry is still the same `ClientSession` that failed. A late error from an old concurrent call cannot evict its replacement or a new in-flight creation. The failing call still surfaces its original error and is never replayed automatically; a later retry creates a fresh subprocess/session. Protocol timeouts, normal `isError=true` tool results, and interceptor failures do not evict a healthy stateful session. +- **Stdio disconnect recovery**: Ordinary Agent tool calls and durable task submit/status/cancel calls that receive the MCP SDK's explicit `Connection closed` error or an AnyIO closed-stream error evict only that `(server_name, incarnation scope)` session when the registered entry is still the same `ClientSession` that failed. A late error from an old concurrent call cannot evict its replacement or a new in-flight creation. The failing call still surfaces its original error and is never replayed automatically; a later retry creates a fresh subprocess/session. Protocol timeouts, normal `isError=true` tool results, and interceptor failures do not evict a healthy stateful session. - **Stdio path translation**: MCP-returned local file references are not copied. If a `ResourceLink` or conservative free-text path resolves to an existing file inside the thread's mounted user-data tree, it is translated deterministically to `/mnt/user-data/...`; paths outside that tree remain unchanged. - **Runtime updates**: Gateway API saves to extensions_config.json; the Gateway-embedded runtime detects changes via the resolved-path + content-signature check above, so multi-worker / stale-mtime deployments still pick up an added/removed MCP server without a restart (`PUT /api/mcp/config` keeps whole-payload validation, while `PATCH /api/mcp/config` changes only one server's `enabled` field, normalizes the same `type`/MCP-spec `transport` alias as the runtime config model, and validates the target only when enabling it; either endpoint's reset clears the cache only in its own worker). MCP, skill, and embedded-client writers hold the process-local `extensions_config_write_lock` plus the sidecar advisory `extensions_config_file_lock` for the complete read-modify-write/reload cycle, then share `atomic_write_extensions_config()`, which writes and fsyncs a same-directory temporary file before `os.replace()` and preserves an existing file's mode and symlink target; failed serialization or replacement leaves the prior config intact and cleans up the temporary file. - **Stdio launch policy at the HTTP boundary** (`routers/mcp.py::_validate_mcp_update_request`, shared by `PUT` and the enable branch of `PATCH`): a config file may express anything, but the API is untrusted input, so an API-registered stdio server must (a) name a bare executable from the allowlist — `_DEFAULT_MCP_STDIO_COMMAND_ALLOWLIST` = `{npx, uvx}`, extended by `DEER_FLOW_MCP_STDIO_COMMAND_ALLOWLIST`, with path separators, whitespace, and shell metacharacters rejected in `command`; (b) carry no `args` flag in `_ARBITRARY_EXEC_ARGS`; and (c) set no `env` name in `_CODE_INJECTING_ENV_VARS`. Checks (b) and (c) exist because the command check alone names a binary without constraining what that binary runs. The `env` denylist applies to **every** allowlisted command, and both denylists match `--flag=value` as well as `--flag value`. The `args` denylist's **scope depends on the command**, because where a launcher stops parsing its own flags is what decides whether a token is an exec flag at all: diff --git a/backend/packages/harness/deerflow/mcp/task_tool_caller.py b/backend/packages/harness/deerflow/mcp/task_tool_caller.py index fc4420742..8ff2abddc 100644 --- a/backend/packages/harness/deerflow/mcp/task_tool_caller.py +++ b/backend/packages/harness/deerflow/mcp/task_tool_caller.py @@ -17,15 +17,11 @@ from deerflow.mcp.headers import apply_header_overrides from deerflow.mcp.interceptors import build_mcp_tool_interceptors from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool +from deerflow.mcp_scope import mcp_session_scope_key logger = logging.getLogger(__name__) -def mcp_task_session_scope_key(*, user_id: str, thread_id: str) -> str: - """Keep background calls in the same per-user/per-thread session scope.""" - return f"{user_id}:{thread_id}" - - def _prepare_stdio_connection( connection: dict[str, Any], *, @@ -92,6 +88,7 @@ class McpTaskToolCaller: arguments: dict[str, Any], user_id: str, thread_id: str, + thread_incarnation: str | None = None, request_scoped_headers: bool = False, ) -> Any: """Call a raw MCP tool. @@ -107,7 +104,11 @@ class McpTaskToolCaller: raise LookupError(f"MCP task server {server_name!r} is missing or disabled in the startup configuration") connection = build_server_params(server_name, server_config) transport = connection.get("transport", "stdio") - scope_key = mcp_task_session_scope_key(user_id=user_id, thread_id=thread_id) + scope_key = mcp_session_scope_key( + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ) if transport == "stdio": connection = await asyncio.to_thread( diff --git a/backend/packages/harness/deerflow/mcp/tasks/models.py b/backend/packages/harness/deerflow/mcp/tasks/models.py index bbced4d6f..f663ce483 100644 --- a/backend/packages/harness/deerflow/mcp/tasks/models.py +++ b/backend/packages/harness/deerflow/mcp/tasks/models.py @@ -93,6 +93,7 @@ class TaskReference: server_name: str remote_task_id: str driver_data: dict[str, Any] = field(default_factory=dict) + thread_incarnation: str | None = None @classmethod def from_record(cls, record: dict[str, Any]) -> TaskReference: @@ -102,6 +103,9 @@ class TaskReference: thread_id=record["thread_id"], server_name=record["server_name"], remote_task_id=record["remote_task_id"], + # Mixed-version/custom repositories may still emit the legacy + # shape; preserve its explicit NULL session scope during rollout. + thread_incarnation=record.get("_thread_incarnation"), driver_data=dict(record.get("driver_data") or {}), ) @@ -119,6 +123,7 @@ class TaskSubmitRequest: arguments: dict[str, Any] driver_data: dict[str, Any] = field(default_factory=dict) local_task_id: str | None = None + thread_incarnation: str | None = None def __post_init__(self) -> None: _validate_storage_text( diff --git a/backend/packages/harness/deerflow/mcp/tasks/ordinary.py b/backend/packages/harness/deerflow/mcp/tasks/ordinary.py index 6faea18a8..c80e172f7 100644 --- a/backend/packages/harness/deerflow/mcp/tasks/ordinary.py +++ b/backend/packages/harness/deerflow/mcp/tasks/ordinary.py @@ -32,6 +32,7 @@ class McpTaskToolCaller(Protocol): arguments: dict[str, Any], user_id: str, thread_id: str, + thread_incarnation: str | None = None, request_scoped_headers: bool = False, ) -> Any: ... @@ -159,6 +160,7 @@ class OrdinaryMcpTaskDriver: arguments=request.arguments, user_id=request.user_id, thread_id=request.thread_id, + thread_incarnation=request.thread_incarnation, # Submit alone is awaited inside the Agent run, so it is the one # durable-task call that can carry the run's request-scoped # credentials; status and cancel run after that run ended. @@ -184,6 +186,7 @@ class OrdinaryMcpTaskDriver: arguments={"task_id": task.remote_task_id}, user_id=task.user_id, thread_id=task.thread_id, + thread_incarnation=task.thread_incarnation, ) payload = _parse( _StatusPayload, @@ -202,6 +205,7 @@ class OrdinaryMcpTaskDriver: arguments={"task_id": task.remote_task_id}, user_id=task.user_id, thread_id=task.thread_id, + thread_incarnation=task.thread_incarnation, ) payload = _parse( _CancelPayload, diff --git a/backend/packages/harness/deerflow/mcp/tasks/runtime.py b/backend/packages/harness/deerflow/mcp/tasks/runtime.py index c60fd15a9..d118fe70d 100644 --- a/backend/packages/harness/deerflow/mcp/tasks/runtime.py +++ b/backend/packages/harness/deerflow/mcp/tasks/runtime.py @@ -26,6 +26,7 @@ class McpTaskSubmitter(Protocol): *, thread_id: str, user_id: str, + thread_incarnation: str | None, limit: int = 50, active_only: bool = False, ) -> list[dict[str, Any]]: ... @@ -35,6 +36,7 @@ class McpTaskSubmitter(Protocol): *, thread_id: str, user_id: str, + thread_incarnation: str | None, task: str | None = None, ) -> dict[str, Any]: ... diff --git a/backend/packages/harness/deerflow/mcp/tools.py b/backend/packages/harness/deerflow/mcp/tools.py index 101a5d2a3..78a4d02f8 100644 --- a/backend/packages/harness/deerflow/mcp/tools.py +++ b/backend/packages/harness/deerflow/mcp/tools.py @@ -29,6 +29,7 @@ from deerflow.mcp.tasks.runtime import ( get_mcp_task_submitter, validate_mcp_task_config_snapshot, ) +from deerflow.mcp_scope import mcp_session_scope_key, runtime_thread_incarnation from deerflow.reflection import resolve_variable from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool @@ -488,7 +489,7 @@ def _make_session_pool_tool( """Wrap an MCP tool so it reuses a persistent session from the pool. Replaces the per-call session creation with pool-managed sessions scoped - by ``(server_name, user_id:thread_id)``. This ensures stateful MCP servers + by ``(server_name, user/thread/incarnation)``. This ensures stateful MCP servers (e.g. Playwright) keep their state across tool calls within the same thread while staying isolated per user. @@ -513,7 +514,12 @@ def _make_session_pool_tool( # Scope the pooled session by user *and* thread. Filesystem isolation is # per-(user_id, thread_id), so a thread_id alone could otherwise let two # users with a colliding thread_id share one stateful MCP session. - scope_key = f"{user_id}:{thread_id}" + thread_incarnation = runtime_thread_incarnation(runtime) + scope_key = mcp_session_scope_key( + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ) session_connection = dict(connection) # cwd/temp pinning and the workspace snapshot only matter for stdio # servers, which run as local subprocesses writing to a real filesystem. @@ -676,6 +682,7 @@ def _make_background_submit_tool( submitter = get_mcp_task_submitter() thread_id = _extract_thread_id(runtime) user_id = resolve_runtime_user_id(runtime) + thread_incarnation = runtime_thread_incarnation(runtime) context = runtime.context if runtime is not None and runtime.context else {} run_id = context.get("run_id") tool_call_id = getattr(runtime, "tool_call_id", None) if runtime is not None else None @@ -684,6 +691,7 @@ def _make_background_submit_tool( request=TaskSubmitRequest( user_id=user_id, thread_id=thread_id, + thread_incarnation=thread_incarnation, run_id=str(run_id) if run_id is not None else None, tool_call_id=str(tool_call_id) if tool_call_id is not None else None, server_name=server_name, diff --git a/backend/packages/harness/deerflow/mcp_scope.py b/backend/packages/harness/deerflow/mcp_scope.py new file mode 100644 index 000000000..460696eba --- /dev/null +++ b/backend/packages/harness/deerflow/mcp_scope.py @@ -0,0 +1,64 @@ +"""Canonical MCP session scope construction.""" + +from __future__ import annotations + +import json +from typing import Any, TypeGuard + +THREAD_INCARNATION_CONTEXT_KEY = "thread_incarnation" +THREAD_INCARNATION_METADATA_GUARD_KEY = "__deerflow_thread_incarnation_metadata_guard" +_MISSING = object() + + +def is_valid_thread_incarnation(value: object) -> TypeGuard[str | None]: + """Return whether *value* is a supported legacy or versioned incarnation.""" + return value is None or (isinstance(value, str) and bool(value)) + + +def mcp_session_scope_key( + *, + user_id: str, + thread_id: str, + thread_incarnation: str | None = None, +) -> str: + """Return the canonical user/thread/incarnation session key. + + A legacy NULL incarnation retains the pre-activation scope so rolling + upgrades do not split an existing legacy session. + """ + if not is_valid_thread_incarnation(thread_incarnation): + raise RuntimeError("MCP session scope requires a non-empty thread incarnation") + scope = f"{user_id}:{thread_id}" + if thread_incarnation is None: + return scope + # A JSON tuple is an unambiguous, versioned encoding even when an opaque + # user/thread id contains the delimiter used by the legacy scope. + return "v2:" + json.dumps( + [user_id, thread_id, thread_incarnation], + ensure_ascii=True, + separators=(",", ":"), + ) + + +def runtime_thread_incarnation(runtime: Any | None) -> str | None: + """Read the server-owned incarnation from ``ToolRuntime.context``.""" + if runtime is None: + # Direct tool invocation has no Agent thread lifecycle and retains the + # legacy scope. An actual runtime with a missing key remains invalid. + return None + context = getattr(runtime, "context", None) + if not isinstance(context, dict) or THREAD_INCARNATION_CONTEXT_KEY not in context: + raise RuntimeError("MCP tool execution requires a server-owned thread incarnation") + value = context[THREAD_INCARNATION_CONTEXT_KEY] + if not is_valid_thread_incarnation(value): + raise RuntimeError("MCP tool execution received an invalid thread incarnation") + if context.get(THREAD_INCARNATION_METADATA_GUARD_KEY) is True: + config = getattr(runtime, "config", None) + metadata = config.get("metadata") if isinstance(config, dict) else None + persisted = metadata.get(THREAD_INCARNATION_CONTEXT_KEY, _MISSING) if isinstance(metadata, dict) else _MISSING + if persisted is _MISSING: + if value is not None: + raise RuntimeError("MCP tool execution received a stale thread incarnation") + elif not is_valid_thread_incarnation(persisted) or persisted != value: + raise RuntimeError("MCP tool execution received a stale thread incarnation") + return value diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/__init__.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/__init__.py index 5de52b3fc..12d53bd1a 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/__init__.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/__init__.py @@ -1,4 +1,13 @@ from deerflow.persistence.mcp_tasks.model import McpTaskRow -from deerflow.persistence.mcp_tasks.sql import DuplicateMcpRemoteTaskError, McpTaskRepository +from deerflow.persistence.mcp_tasks.sql import ( + DuplicateMcpRemoteTaskError, + McpTaskRepository, + McpTaskThreadMismatchError, +) -__all__ = ["DuplicateMcpRemoteTaskError", "McpTaskRepository", "McpTaskRow"] +__all__ = [ + "DuplicateMcpRemoteTaskError", + "McpTaskRepository", + "McpTaskRow", + "McpTaskThreadMismatchError", +] diff --git a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py index 604fc3178..5145a833c 100644 --- a/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py +++ b/backend/packages/harness/deerflow/persistence/mcp_tasks/sql.py @@ -6,7 +6,7 @@ import uuid from datetime import UTC, datetime, timedelta from typing import Any -from sqlalchemy import case, or_, select, update +from sqlalchemy import case, or_, select, text, update from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker @@ -87,6 +87,10 @@ class DuplicateMcpRemoteTaskError(RuntimeError): """The current user already tracks this server's remote task handle.""" +class McpTaskThreadMismatchError(RuntimeError): + """The submitting run no longer owns the current thread incarnation.""" + + def _is_remote_task_unique_conflict(exc: IntegrityError) -> bool: original = exc.orig diagnostic = getattr(original, "diag", None) @@ -96,6 +100,64 @@ def _is_remote_task_unique_conflict(exc: IntegrityError) -> bool: return "uq_mcp_tasks_user_server_remote" in message or "mcp_tasks.user_id, mcp_tasks.server_name, mcp_tasks.remote_task_id" in message +def _matches_current_thread_incarnation( + *, + user_id: str, + thread_id: str, + thread_incarnation: str | None, +): + """Atomically match task, current thread, and caller-captured incarnation.""" + return ( + select(ThreadMetaRow.thread_id) + .where( + McpTaskRow.thread_id == thread_id, + McpTaskRow.thread_incarnation.is_not_distinct_from(thread_incarnation), + ThreadMetaRow.thread_id == thread_id, + ThreadMetaRow.incarnation.is_not_distinct_from(thread_incarnation), + or_(ThreadMetaRow.user_id == user_id, ThreadMetaRow.user_id.is_(None)), + ) + .exists() + ) + + +async def _lock_current_thread_incarnation( + session: AsyncSession, + *, + user_id: str, + thread_id: str, + thread_incarnation: str | None, +) -> bool: + """Lock the expected thread incarnation until the transaction ends.""" + conditions = ( + ThreadMetaRow.thread_id == thread_id, + ThreadMetaRow.incarnation.is_not_distinct_from(thread_incarnation), + or_(ThreadMetaRow.user_id == user_id, ThreadMetaRow.user_id.is_(None)), + ) + if session.get_bind().dialect.name == "sqlite": + # SQLite has no row-level SELECT lock. A no-op UPDATE acquires its + # database writer lock before we inspect or mutate an MCP task row. + # Raw SQL avoids firing ThreadMetaRow.updated_at's ORM onupdate hook. + result = await session.execute( + text( + """ + UPDATE threads_meta + SET incarnation = incarnation + WHERE thread_id = :thread_id + AND incarnation IS :thread_incarnation + AND (user_id = :user_id OR user_id IS NULL) + """ + ), + { + "thread_id": thread_id, + "thread_incarnation": thread_incarnation, + "user_id": user_id, + }, + ) + return result.rowcount == 1 + result = await session.execute(select(ThreadMetaRow.thread_id).where(*conditions).with_for_update(read=True)) + return result.one_or_none() is not None + + class McpTaskRepository: """Durable source of truth for long-running MCP task lifecycle state.""" @@ -103,9 +165,11 @@ class McpTaskRepository: self._sf = session_factory @staticmethod - def _row_to_dict(row: McpTaskRow) -> dict[str, Any]: + def _row_to_dict(row: McpTaskRow, *, include_internal: bool = False) -> dict[str, Any]: data = row.to_dict() - data.pop("thread_incarnation", None) + thread_incarnation = data.pop("thread_incarnation", None) + if include_internal: + data["_thread_incarnation"] = thread_incarnation for key in _TIMESTAMP_FIELDS: if data.get(key) is not None: data[key] = coerce_iso(data[key]) @@ -117,6 +181,7 @@ class McpTaskRepository: task_id: str, user_id: str, thread_id: str, + expected_thread_incarnation: str | None, run_id: str | None, tool_call_id: str | None, server_name: str, @@ -160,19 +225,14 @@ class McpTaskRepository: ) _record_event_if_changed(row, tracking_degraded=False, now=now) async with self._sf() as session: - matching_thread = select(ThreadMetaRow.incarnation).where( - ThreadMetaRow.thread_id == thread_id, - or_(ThreadMetaRow.user_id == user_id, ThreadMetaRow.user_id.is_(None)), - ) - if session.get_bind().dialect.name == "sqlite": - # Keep lookup and write in one SQLite statement. A preliminary - # read would leave a delete/recreate window before the INSERT. - row.thread_incarnation = matching_thread.scalar_subquery() - else: - # FOR SHARE also conflicts with the FOR NO KEY UPDATE lock taken - # by an older writer's plain owner UPDATE. KEY SHARE would not, - # leaving a mixed-version ownership race before this INSERT. - row.thread_incarnation = (await session.execute(matching_thread.with_for_update(read=True))).scalar_one_or_none() + if not await _lock_current_thread_incarnation( + session, + user_id=user_id, + thread_id=thread_id, + thread_incarnation=expected_thread_incarnation, + ): + raise McpTaskThreadMismatchError("MCP task submission crossed a thread lifecycle boundary") + row.thread_incarnation = expected_thread_incarnation session.add(row) try: await session.commit() @@ -184,24 +244,47 @@ class McpTaskRepository: await session.refresh(row) return self._row_to_dict(row) - async def get(self, task_id: str, *, user_id: str) -> dict[str, Any] | None: + async def get( + self, + task_id: str, + *, + user_id: str, + thread_id: str, + thread_incarnation: str | None, + ) -> dict[str, Any] | None: async with self._sf() as session: - row = await session.get(McpTaskRow, task_id) - if row is None or row.user_id != user_id: - return None - return self._row_to_dict(row) + stmt = select(McpTaskRow).where( + McpTaskRow.id == task_id, + McpTaskRow.user_id == user_id, + McpTaskRow.thread_id == thread_id, + McpTaskRow.thread_incarnation.is_not_distinct_from(thread_incarnation), + _matches_current_thread_incarnation( + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ), + ) + row = (await session.execute(stmt)).scalar_one_or_none() + return self._row_to_dict(row) if row is not None else None async def list_by_thread( self, thread_id: str, *, user_id: str, + thread_incarnation: str | None, limit: int = 50, active_only: bool = False, ) -> list[dict[str, Any]]: stmt = select(McpTaskRow).where( McpTaskRow.thread_id == thread_id, McpTaskRow.user_id == user_id, + McpTaskRow.thread_incarnation.is_not_distinct_from(thread_incarnation), + _matches_current_thread_incarnation( + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ), ) if active_only: stmt = stmt.where(McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES)) @@ -245,7 +328,7 @@ class McpTaskRepository: row.poll_attempt_count += 1 row.updated_at = now await session.commit() - return [self._row_to_dict(row) for row in rows] + return [self._row_to_dict(row, include_internal=True) for row in rows] async def apply_snapshot( self, @@ -387,16 +470,31 @@ class McpTaskRepository: *, user_id: str, thread_id: str, + thread_incarnation: str | None, requested_at: datetime, ) -> dict[str, Any] | None: """Persist a user-scoped cancellation request without exposing the remote id.""" async with self._sf() as session: + if not await _lock_current_thread_incarnation( + session, + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ): + await session.rollback() + return None stmt = ( select(McpTaskRow) .where( McpTaskRow.id == task_id, McpTaskRow.user_id == user_id, McpTaskRow.thread_id == thread_id, + McpTaskRow.thread_incarnation.is_not_distinct_from(thread_incarnation), + _matches_current_thread_incarnation( + user_id=user_id, + thread_id=thread_id, + thread_incarnation=thread_incarnation, + ), ) .with_for_update() ) @@ -447,7 +545,7 @@ class McpTaskRepository: row.cancel_attempt_count = int(row.cancel_attempt_count or 0) + 1 row.updated_at = now await session.commit() - return [self._row_to_dict(row) for row in rows] + return [self._row_to_dict(row, include_internal=True) for row in rows] async def apply_cancel_snapshot( self, diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 5288f2dc6..14d4c9d02 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -42,6 +42,10 @@ from deerflow.config.app_config import AppConfig from deerflow.config.database_config import CheckpointChannelMode from deerflow.constants import CONVERSATION_READER_CONTEXT_KEY, TOOL_RESULTS_DIRNAME from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope +from deerflow.mcp_scope import ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, +) from deerflow.runtime.checkpoint_mode import ( aensure_checkpoint_mode_compatible, inject_checkpoint_mode, @@ -99,6 +103,7 @@ from .naming import resolve_root_run_name from .schemas import RunStatus logger = logging.getLogger(__name__) +_THREAD_INCARNATION_UNSET = object() _checkpoint_locks = AsyncKeyedLockTable[str]() @@ -514,6 +519,8 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = ( CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY, DEERFLOW_TRACE_METADATA_KEY, CONVERSATION_READER_CONTEXT_KEY, + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, "is_subagent", "agent_id", "__run_loop_detection_recorder", @@ -538,6 +545,8 @@ def _build_runtime_context( task_store: Any | None = None, extensions: Any | None = None, conversation_reader: Any | None = None, + *, + thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET, ) -> dict[str, Any]: """Build the dict that becomes ``ToolRuntime.context`` for the run. @@ -552,6 +561,8 @@ def _build_runtime_context( ``langgraph.pregel.main`` where ``parent_runtime.merge(...)`` is invoked. """ runtime_ctx: dict[str, Any] = {"thread_id": thread_id, "run_id": run_id} + if thread_incarnation is not _THREAD_INCARNATION_UNSET: + runtime_ctx[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation if isinstance(caller_context, dict): for key, value in caller_context.items(): if key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: @@ -806,6 +817,7 @@ async def run_agent( agent_factory: Any, graph_input: dict, config: dict, + thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET, stream_modes: list[str] | None = None, stream_subgraphs: bool = False, interrupt_before: list[str] | Literal["*"] | None = None, @@ -946,9 +958,13 @@ async def run_agent( # cancellation cannot strand a pending RunRecord or stream subscriber. if ctx.mcp_task_repo is not None and record.user_id is not None: try: + if thread_incarnation is _THREAD_INCARNATION_UNSET: + raise RuntimeError("MCP task projection requires a server-owned thread incarnation") + assert thread_incarnation is None or isinstance(thread_incarnation, str) task_rows = await ctx.mcp_task_repo.list_by_thread( thread_id, user_id=record.user_id, + thread_incarnation=thread_incarnation, limit=20, ) graph_input = { @@ -1074,6 +1090,7 @@ async def run_agent( task_store, extensions, ctx.conversation_reader, + thread_incarnation=thread_incarnation, ) # Bind every checkpoint produced by this run to the effective agent # identity that produced its state. Manual compaction uses only this diff --git a/backend/packages/harness/deerflow/subagents/AGENTS.md b/backend/packages/harness/deerflow/subagents/AGENTS.md index c8be8a966..ba23aecac 100644 --- a/backend/packages/harness/deerflow/subagents/AGENTS.md +++ b/backend/packages/harness/deerflow/subagents/AGENTS.md @@ -7,6 +7,8 @@ FIFOs remain rejected. **Context**: Capture after validation, before setup. Keep genuine replies, even hidden clarifications; exclude framework state and unpaired calls. Mark unserializable media as omitted. +`task` preserves the host incarnation, including null/missing/invalid values; never recapture it. + **Direct runtime shutdown**: `SubagentRuntime.stop()` holds its lifecycle lock until the owned service stop task terminates, then propagates the first caller cancellation with any service failure/cancellation as its cause. The drain is intentionally unbounded: repository awaits and child cleanup must terminate; a timeout must not detach still-owned work. Keep terminal-outcome and repeated-cancellation coverage in `tests/test_subagent_runtime.py`. **Durable batch acceptance**: `batch_task` normalizes optional per-item criteria diff --git a/backend/packages/harness/deerflow/subagents/batch_service.py b/backend/packages/harness/deerflow/subagents/batch_service.py index 6cd387ce9..e46c9b6c6 100644 --- a/backend/packages/harness/deerflow/subagents/batch_service.py +++ b/backend/packages/harness/deerflow/subagents/batch_service.py @@ -11,6 +11,7 @@ from deerflow.config.app_config import AppConfig, get_app_config from deerflow.config.subagent_batches_config import SubagentBatchesConfig from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig from deerflow.extensions import LoadedExtensions, get_loaded_extensions +from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY from deerflow.subagents.batch_acceptance import check_batch_acceptance from deerflow.subagents.batch_runtime import BatchSubmitRequest from deerflow.subagents.capacity import SubagentExecutionCapacity @@ -233,6 +234,9 @@ class SubagentBatchService: item_id, ) return + executor_kwargs = {} + if THREAD_INCARNATION_CONTEXT_KEY in spec: + executor_kwargs[THREAD_INCARNATION_CONTEXT_KEY] = spec[THREAD_INCARNATION_CONTEXT_KEY] executor = SubagentExecutor( config=config, tools=tools, @@ -251,6 +255,7 @@ class SubagentBatchService: execution_capacity=self._execution_capacity, extensions=self._extensions, acceptance_criteria=item.get("acceptance_criteria"), + **executor_kwargs, ) prompt = f"Durable batch item key: {item['item_key']}\nThis item may be retried after a worker crash. Keep side effects idempotent and use the item key as the idempotency identity.\n\n{item['prompt']}" execution_id = executor.execute_async(prompt, task_id=item_id) diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index af3005f41..4d7651004 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -37,6 +37,7 @@ from deerflow.authz.principal import normalize_authz_attributes from deerflow.config import get_app_config from deerflow.config.app_config import AppConfig from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope +from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY from deerflow.models import create_chat_model from deerflow.runtime.runs.stream_cleanup import close_agent_stream from deerflow.runtime.user_context import DEFAULT_USER_ID @@ -775,6 +776,9 @@ def _filter_tools( return filtered +_THREAD_INCARNATION_UNSET = object() + + class SubagentExecutor: """Executor for running subagents.""" @@ -806,6 +810,7 @@ class SubagentExecutor: tool_promotion_recorder: Any | None = None, tool_progress_recorder: Any | None = None, context_snapshot: ParentContextSnapshot | None = None, + thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET, ): """Initialize the executor. @@ -822,6 +827,8 @@ class SubagentExecutor: run. Seeded into the child graph state so ``list_uploaded_files`` can exclude them from historical-upload results. thread_id: Thread ID for sandbox operations. + thread_incarnation: Server-captured parent lifecycle. Explicit None + preserves legacy scope; omission stays absent so MCP fails closed. trace_id: Trace ID from parent for distributed tracing. user_id: User ID captured from the parent tool's runtime context. When None, the tracing layer falls back to DEFAULT_USER_ID. @@ -879,6 +886,7 @@ class SubagentExecutor: self.uploaded_files = deepcopy(uploaded_files) if uploaded_files is not None else None self.context_snapshot = context_snapshot self.thread_id = thread_id + self.thread_incarnation = thread_incarnation # Generate trace_id if not provided (for top-level calls) self.trace_id = trace_id or str(uuid.uuid4())[:8] self.user_id = user_id @@ -1548,6 +1556,8 @@ class SubagentExecutor: context: dict[str, Any] = {} if self.thread_id: context["thread_id"] = self.thread_id + if self.thread_incarnation is not _THREAD_INCARNATION_UNSET: + context[THREAD_INCARNATION_CONTEXT_KEY] = self.thread_incarnation if self.app_config is not None: context["app_config"] = self.app_config # Propagate guardrail attribution so delegated tool calls are diff --git a/backend/packages/harness/deerflow/tools/builtins/background_tasks_tool.py b/backend/packages/harness/deerflow/tools/builtins/background_tasks_tool.py index b019e09e7..3e5c6e0e7 100644 --- a/backend/packages/harness/deerflow/tools/builtins/background_tasks_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/background_tasks_tool.py @@ -8,6 +8,7 @@ from langchain.tools import tool from deerflow.agents.middlewares.input_sanitization_middleware import neutralize_untrusted_tags from deerflow.mcp.tasks.runtime import get_mcp_task_submitter +from deerflow.mcp_scope import runtime_thread_incarnation from deerflow.tools.builtins.list_uploaded_files_tool import _resolve_thread_id, _resolve_user_id from deerflow.tools.types import Runtime @@ -36,6 +37,7 @@ async def _list_background_tasks_impl( records = await get_mcp_task_submitter().list_tasks( thread_id=thread_id, user_id=_resolve_user_id(runtime), + thread_incarnation=runtime_thread_incarnation(runtime), limit=max(1, min(limit, 50)), active_only=active_only, ) @@ -72,6 +74,7 @@ async def cancel_background_task( record = await get_mcp_task_submitter().cancel_matching_task( thread_id=thread_id, user_id=_resolve_user_id(runtime), + thread_incarnation=runtime_thread_incarnation(runtime), task=task, ) except (LookupError, ValueError) as exc: diff --git a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py index eb8bbc884..1cbadd32f 100644 --- a/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/batch_task_tool.py @@ -15,6 +15,11 @@ from pydantic import BaseModel, Field from deerflow.authz.principal import normalize_authz_attributes from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope +from deerflow.mcp_scope import ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, + runtime_thread_incarnation, +) from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.subagents.batch_runtime import ( BatchItemInput, @@ -179,6 +184,8 @@ async def batch_task( return _result(tool_call_id, content="Batch item keys must be unique.", error=True) context = runtime.context if runtime is not None and isinstance(runtime.context, dict) else {} + if context.get(THREAD_INCARNATION_METADATA_GUARD_KEY) is True: + runtime_thread_incarnation(runtime) metadata = runtime.config.get("metadata", {}) if runtime is not None else {} app_config = _batch_app_config(runtime) allowed_subagents = metadata.get("allowed_subagents") @@ -214,6 +221,8 @@ async def batch_task( "is_internal": context.get("is_internal") is True, "authz_attributes": normalize_authz_attributes(context.get("authz_attributes")), } + if THREAD_INCARNATION_CONTEXT_KEY in context: + execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = context[THREAD_INCARNATION_CONTEXT_KEY] if KNOWLEDGE_SCOPE_RUNTIME_KEY in context: execution_spec["knowledge_scope"] = execution_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY]) try: diff --git a/backend/packages/harness/deerflow/tools/builtins/task_tool.py b/backend/packages/harness/deerflow/tools/builtins/task_tool.py index b8f3d2b65..0ff4f8ae5 100644 --- a/backend/packages/harness/deerflow/tools/builtins/task_tool.py +++ b/backend/packages/harness/deerflow/tools/builtins/task_tool.py @@ -23,6 +23,11 @@ from deerflow.community.ragflow.sources import cited_source_artifact from deerflow.config import get_app_config from deerflow.extensions import resolve_run_extensions from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope +from deerflow.mcp_scope import ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, + runtime_thread_incarnation, +) from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.sandbox.security import LOCAL_BASH_SUBAGENT_DISABLED_MESSAGE, is_host_bash_allowed from deerflow.subagents import SubagentExecutor, get_available_subagent_names, get_subagent_config @@ -850,6 +855,8 @@ async def task_tool( # tool call delegated to a subagent (user_role=None). parent_context = runtime.context if runtime is not None else None parent_context = parent_context if isinstance(parent_context, dict) else {} + if parent_context.get(THREAD_INCARNATION_METADATA_GUARD_KEY) is True: + runtime_thread_incarnation(runtime) user_role = parent_context.get("user_role") oauth_provider = parent_context.get("oauth_provider") oauth_id = parent_context.get("oauth_id") @@ -942,6 +949,10 @@ async def task_tool( # system-channel authority over framework instructions. "acceptance_criteria": acceptance_criteria, } + # Carry the host-captured lifecycle, including legacy None, without + # inventing a legacy scope for missing context or re-reading thread state. + if THREAD_INCARNATION_CONTEXT_KEY in parent_context: + executor_kwargs["thread_incarnation"] = parent_context[THREAD_INCARNATION_CONTEXT_KEY] if context_snapshot is not None: executor_kwargs["context_snapshot"] = context_snapshot middleware_recorder = None diff --git a/backend/tests/test_agent_assembly_descriptor.py b/backend/tests/test_agent_assembly_descriptor.py index 0c350d466..22597eca4 100644 --- a/backend/tests/test_agent_assembly_descriptor.py +++ b/backend/tests/test_agent_assembly_descriptor.py @@ -239,6 +239,7 @@ class TestLeadAgentAssembly: assembly = assemble_lead_agent({"configurable": {"thread_id": "t-1"}}) assert isinstance(assembly, LeadAgentAssembly) assert assembly.graph is not None + assert assembly.graph["context_schema"] is dict assert assembly.descriptor.effective_model assert assembly.descriptor.fingerprint @@ -653,6 +654,7 @@ class TestCustomAgentModelSettingsReachTheDescriptor: TestLeadAgentAssembly._isolate_from_the_ambient_config(monkeypatch) with bind_agent_build_extensions(TestLeadAgentAssembly._extensions_with_an_agent_assembly_observer()): assembly = assemble_lead_agent({"configurable": {"thread_id": "t-bootstrap", "is_bootstrap": True}}) + assert assembly.graph["context_schema"] is dict assert "temperature" not in assembly.descriptor.model_parameters diff --git a/backend/tests/test_background_tasks_tool.py b/backend/tests/test_background_tasks_tool.py index 0db19d2ef..6d5f00953 100644 --- a/backend/tests/test_background_tasks_tool.py +++ b/backend/tests/test_background_tasks_tool.py @@ -18,7 +18,11 @@ def _clear_submitter(): def _runtime(): return SimpleNamespace( - context={"thread_id": "thread-1", "user_id": "user-1"}, + context={ + "thread_id": "thread-1", + "user_id": "user-1", + "thread_incarnation": "incarnation-1", + }, state={}, config={}, ) @@ -53,6 +57,7 @@ async def test_list_background_tasks_returns_only_safe_local_fields(): manager.list_tasks.assert_awaited_once_with( thread_id="thread-1", user_id="user-1", + thread_incarnation="incarnation-1", limit=20, active_only=False, ) @@ -83,5 +88,19 @@ async def test_cancel_background_task_uses_current_user_and_thread(): manager.cancel_matching_task.assert_awaited_once_with( thread_id="thread-1", user_id="user-1", + thread_incarnation="incarnation-1", task="report", ) + + +@pytest.mark.asyncio +async def test_background_task_tools_fail_closed_without_server_owned_incarnation(): + manager = SimpleNamespace(list_tasks=AsyncMock()) + set_mcp_task_submitter(manager) + runtime = _runtime() + runtime.context.pop("thread_incarnation") + + with pytest.raises(RuntimeError, match="server-owned thread incarnation"): + await _list_background_tasks_impl(runtime) + + manager.list_tasks.assert_not_awaited() diff --git a/backend/tests/test_batch_task_tool.py b/backend/tests/test_batch_task_tool.py index 51f2abea1..7ecc90aeb 100644 --- a/backend/tests/test_batch_task_tool.py +++ b/backend/tests/test_batch_task_tool.py @@ -6,26 +6,34 @@ import pytest from langchain_core.messages import ToolMessage from langgraph.types import Command +from deerflow.mcp_scope import ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, +) from deerflow.subagents.config import SubagentConfig from deerflow.tools.builtins.batch_task_tool import BatchTaskItem tool_module = importlib.import_module("deerflow.tools.builtins.batch_task_tool") +_MISSING = object() -def _runtime(): +def _runtime(thread_incarnation=_MISSING): + context = { + "thread_id": "thread-1", + "run_id": "run-1", + "user_id": "user-1", + "user_role": "member", + "__knowledge_scope_execution": { + "version": 1, + "mode": "selected", + "dataset_ids": ["dataset-1"], + }, + } + if thread_incarnation is not _MISSING: + context[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation return SimpleNamespace( state={}, - context={ - "thread_id": "thread-1", - "run_id": "run-1", - "user_id": "user-1", - "user_role": "member", - "__knowledge_scope_execution": { - "version": 1, - "mode": "selected", - "dataset_ids": ["dataset-1"], - }, - }, + context=context, config={ "metadata": { "model_name": "model-a", @@ -45,7 +53,19 @@ def _message(command: Command) -> ToolMessage: @pytest.mark.asyncio -async def test_batch_task_is_explicit_idempotent_submission(monkeypatch) -> None: +@pytest.mark.parametrize( + ("thread_incarnation", "expected_present"), + [ + ("incarnation-1", True), + (None, True), + (_MISSING, False), + ], +) +async def test_batch_task_is_explicit_idempotent_submission( + monkeypatch, + thread_incarnation, + expected_present, +) -> None: submitter = AsyncMock() submitter.submit.return_value = { "id": "subagent-batch-1", @@ -68,7 +88,7 @@ async def test_batch_task_is_explicit_idempotent_submission(monkeypatch) -> None ) command = await tool_module.batch_task.coroutine( - runtime=_runtime(), + runtime=_runtime(thread_incarnation), title="Process records", items=[ BatchTaskItem(key="record-1", prompt="Process one"), @@ -93,10 +113,35 @@ async def test_batch_task_is_explicit_idempotent_submission(monkeypatch) -> None "mode": "selected", "dataset_ids": ["dataset-1"], } + assert (THREAD_INCARNATION_CONTEXT_KEY in request.execution_spec) is expected_present + if expected_present: + assert request.execution_spec[THREAD_INCARNATION_CONTEXT_KEY] is thread_incarnation assert message.additional_kwargs["subagent_batch_id"] == "subagent-batch-1" assert "running independently" in message.content +@pytest.mark.asyncio +async def test_batch_task_rejects_stale_standalone_thread_incarnation( + monkeypatch, +) -> None: + submitter = AsyncMock() + monkeypatch.setattr(tool_module, "get_subagent_batch_submitter", lambda: submitter) + runtime = _runtime("incarnation-1") + runtime.context[THREAD_INCARNATION_METADATA_GUARD_KEY] = True + runtime.config["metadata"][THREAD_INCARNATION_CONTEXT_KEY] = "incarnation-2" + + with pytest.raises(RuntimeError, match="stale thread incarnation"): + await tool_module.batch_task.coroutine( + runtime=runtime, + title="Stale lifecycle", + items=[BatchTaskItem(key="record-1", prompt="Process one")], + subagent_type="general-purpose", + tool_call_id="call-1", + ) + + submitter.submit.assert_not_awaited() + + @pytest.mark.asyncio async def test_batch_task_rejects_duplicate_item_keys_without_submitting(monkeypatch) -> None: submitter = AsyncMock() diff --git a/backend/tests/test_gateway_services.py b/backend/tests/test_gateway_services.py index b4a5c1f88..c19e2a8e5 100644 --- a/backend/tests/test_gateway_services.py +++ b/backend/tests/test_gateway_services.py @@ -1909,6 +1909,7 @@ async def test_pending_cancel_bypasses_thread_metadata_and_logs_failure(_stub_ap await asyncio.sleep(0) assert "thread metadata store failed after cancellation" in caplog.text + assert "MCP access will fail closed" not in caplog.text @pytest.mark.asyncio @@ -1960,6 +1961,7 @@ async def test_thread_metadata_timeout_logs_and_run_still_starts(_stub_app_confi assert record.status == RunStatus.running assert (await run_manager.get(record.run_id)).status == RunStatus.running assert "Timed out ensuring thread_meta for thread-timeout-meta" in caplog.text + assert "Thread metadata for thread-timeout-meta is unavailable; MCP access will fail closed" in caplog.text def test_context_merges_into_configurable(): @@ -3999,6 +4001,7 @@ def test_strip_internal_context_keys_scrubs_audit_attribution_and_recorders(): "__run_loop_detection_recorder": "forged", "__run_tool_promotion_recorder": "forged", "__run_tool_progress_recorder": "forged", + "__deerflow_thread_incarnation_metadata_guard": True, } config = build_run_config( "thread-1", diff --git a/backend/tests/test_langgraph_auth.py b/backend/tests/test_langgraph_auth.py index 9a0e6af5b..804eebc7d 100644 --- a/backend/tests/test_langgraph_auth.py +++ b/backend/tests/test_langgraph_auth.py @@ -6,9 +6,10 @@ Validates that the LangGraph auth layer enforces the same rules as Gateway: import asyncio import os +import sys from datetime import timedelta from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from unittest.mock import AsyncMock, patch from uuid import uuid4 @@ -18,11 +19,16 @@ os.environ.setdefault("AUTH_JWT_SECRET", "test-secret-key-for-langgraph-auth-tes from langgraph_sdk import Auth +from app.gateway import langgraph_auth as auth_module from app.gateway.auth.config import AuthConfig, set_auth_config from app.gateway.auth.jwt import create_access_token, decode_token from app.gateway.auth.models import User from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID from app.gateway.langgraph_auth import add_owner_filter, authenticate +from deerflow.mcp_scope import ( + THREAD_INCARNATION_CONTEXT_KEY, + THREAD_INCARNATION_METADATA_GUARD_KEY, +) # ── Helpers ─────────────────────────────────────────────────────────────── @@ -235,6 +241,267 @@ def test_filter_with_empty_metadata(): assert result == {"user_id": "user-z"} +def test_thread_create_overwrites_client_incarnation(): + value = {"metadata": {THREAD_INCARNATION_CONTEXT_KEY: "attacker"}} + + asyncio.run(add_owner_filter(_make_ctx("user-a"), value)) + + incarnation = value["metadata"][THREAD_INCARNATION_CONTEXT_KEY] + assert incarnation != "attacker" + assert isinstance(incarnation, str) and incarnation + + +@pytest.mark.parametrize("value", ["attacker", "", None, False]) +def test_thread_update_cannot_change_incarnation(value): + request = {"metadata": {THREAD_INCARNATION_CONTEXT_KEY: value}} + + asyncio.run( + add_owner_filter( + _make_ctx("user-a", action="update"), + request, + ) + ) + + assert THREAD_INCARNATION_CONTEXT_KEY not in request["metadata"] + + +def test_run_admission_uses_persisted_incarnation_and_scrubs_client_values(): + thread_id = uuid4() + run_id = uuid4() + value = { + "thread_id": thread_id, + "run_id": run_id, + "metadata": { + THREAD_INCARNATION_CONTEXT_KEY: "metadata-attacker", + THREAD_INCARNATION_METADATA_GUARD_KEY: True, + }, + "kwargs": { + "context": { + THREAD_INCARNATION_CONTEXT_KEY: "context-attacker", + THREAD_INCARNATION_METADATA_GUARD_KEY: False, + "user_id": "attacker", + "thread_id": "attacker", + "run_id": "attacker", + }, + "config": { + "context": { + THREAD_INCARNATION_CONTEXT_KEY: "config-context-attacker", + THREAD_INCARNATION_METADATA_GUARD_KEY: False, + }, + "metadata": { + THREAD_INCARNATION_CONTEXT_KEY: "config-metadata-attacker", + THREAD_INCARNATION_METADATA_GUARD_KEY: False, + }, + "configurable": { + THREAD_INCARNATION_CONTEXT_KEY: "configurable-attacker", + THREAD_INCARNATION_METADATA_GUARD_KEY: False, + }, + }, + }, + "if_not_exists": "reject", + } + + with patch.object( + auth_module, + "_read_standalone_thread", + AsyncMock( + return_value={ + "metadata": {THREAD_INCARNATION_CONTEXT_KEY: "server-incarnation"}, + } + ), + ): + asyncio.run( + add_owner_filter( + _make_ctx("user-a", action="create_run"), + value, + ) + ) + + assert value["kwargs"]["context"][THREAD_INCARNATION_CONTEXT_KEY] == "server-incarnation" + assert value["kwargs"]["context"][THREAD_INCARNATION_METADATA_GUARD_KEY] is True + assert value["kwargs"]["context"]["user_id"] == "user-a" + assert value["kwargs"]["context"]["thread_id"] == str(thread_id) + assert value["kwargs"]["context"]["run_id"] == str(run_id) + assert THREAD_INCARNATION_CONTEXT_KEY not in value["metadata"] + assert THREAD_INCARNATION_METADATA_GUARD_KEY not in value["metadata"] + assert THREAD_INCARNATION_CONTEXT_KEY not in value["kwargs"]["config"]["context"] + assert THREAD_INCARNATION_METADATA_GUARD_KEY not in value["kwargs"]["config"]["context"] + assert "user_id" not in value["kwargs"]["config"]["context"] + assert "thread_id" not in value["kwargs"]["config"]["context"] + assert "run_id" not in value["kwargs"]["config"]["context"] + assert THREAD_INCARNATION_CONTEXT_KEY not in value["kwargs"]["config"]["metadata"] + assert THREAD_INCARNATION_METADATA_GUARD_KEY not in value["kwargs"]["config"]["metadata"] + assert THREAD_INCARNATION_CONTEXT_KEY not in value["kwargs"]["config"]["configurable"] + assert THREAD_INCARNATION_METADATA_GUARD_KEY not in value["kwargs"]["config"]["configurable"] + + +def test_existing_legacy_thread_uses_explicit_none_without_backfill(): + with patch.object( + auth_module, + "_read_standalone_thread", + AsyncMock(return_value={"metadata": {"legacy": True}}), + ): + incarnation = asyncio.run( + auth_module._ensure_standalone_thread_incarnation( + uuid4(), + _make_ctx("user-a", action="create_run"), + create_if_missing=False, + ) + ) + + assert incarnation is None + + +def test_implicit_create_accepts_legacy_thread_created_by_mixed_version_peer(): + class _Connection: + async def __aenter__(self): + return object() + + async def __aexit__(self, *_args): + return None + + async def _put_rows(): + yield {"metadata": {"legacy": True}} + + put = AsyncMock(return_value=_put_rows()) + runtime_package = ModuleType("langgraph_runtime") + runtime_package.__path__ = [] + database_module = ModuleType("langgraph_runtime.database") + database_module.connect = _Connection + ops_module = ModuleType("langgraph_runtime.ops") + ops_module.Threads = SimpleNamespace(put=put) + + with ( + patch.object( + auth_module, + "_read_standalone_thread", + AsyncMock(side_effect=[None, {"metadata": {"legacy": True}}]), + ), + patch.dict( + sys.modules, + { + "langgraph_runtime": runtime_package, + "langgraph_runtime.database": database_module, + "langgraph_runtime.ops": ops_module, + }, + ), + ): + incarnation = asyncio.run( + auth_module._ensure_standalone_thread_incarnation( + uuid4(), + _make_ctx("user-a", action="create_run"), + create_if_missing=True, + ) + ) + + assert incarnation is None + put.assert_awaited_once() + + +@pytest.mark.parametrize("persisted_incarnation", [None, "versioned-incarnation"]) +def test_run_admission_preserves_persisted_incarnation_value(persisted_incarnation): + value = { + "thread_id": uuid4(), + "metadata": {}, + "kwargs": {}, + "if_not_exists": "reject", + } + + with patch.object( + auth_module, + "_ensure_standalone_thread_incarnation", + AsyncMock(return_value=persisted_incarnation), + ): + asyncio.run( + add_owner_filter( + _make_ctx("user-a", action="create_run"), + value, + ) + ) + + assert THREAD_INCARNATION_CONTEXT_KEY not in value["metadata"] + assert value["kwargs"]["context"][THREAD_INCARNATION_CONTEXT_KEY] is persisted_incarnation + assert value["kwargs"]["context"][THREAD_INCARNATION_METADATA_GUARD_KEY] is True + + +def test_run_admission_does_not_invent_incarnation_for_missing_rejected_thread(): + value = { + "thread_id": uuid4(), + "metadata": {THREAD_INCARNATION_CONTEXT_KEY: "attacker"}, + "kwargs": { + "context": {THREAD_INCARNATION_CONTEXT_KEY: "attacker"}, + }, + "if_not_exists": "reject", + } + + with patch.object( + auth_module, + "_ensure_standalone_thread_incarnation", + AsyncMock(return_value=auth_module._MISSING), + ): + asyncio.run( + add_owner_filter( + _make_ctx("user-a", action="create_run"), + value, + ) + ) + + assert THREAD_INCARNATION_CONTEXT_KEY not in value["metadata"] + assert THREAD_INCARNATION_CONTEXT_KEY not in value["kwargs"]["context"] + assert THREAD_INCARNATION_METADATA_GUARD_KEY not in value["kwargs"]["context"] + + +def test_run_admission_rejects_invalid_persisted_incarnation(): + value = { + "thread_id": uuid4(), + "metadata": {}, + "kwargs": {}, + "if_not_exists": "reject", + } + + with ( + patch.object( + auth_module, + "_read_standalone_thread", + AsyncMock( + return_value={ + "metadata": {THREAD_INCARNATION_CONTEXT_KEY: ""}, + } + ), + ), + pytest.raises( + RuntimeError, + match="invalid incarnation", + ), + ): + asyncio.run( + add_owner_filter( + _make_ctx("user-a", action="create_run"), + value, + ) + ) + + +def test_temporary_run_gets_explicit_legacy_incarnation(): + value = { + "thread_id": None, + "metadata": {}, + "kwargs": { + "context": {THREAD_INCARNATION_CONTEXT_KEY: "attacker"}, + }, + "if_not_exists": "reject", + } + + asyncio.run( + add_owner_filter( + _make_ctx("user-a", action="create_run"), + value, + ) + ) + + assert value["kwargs"]["context"][THREAD_INCARNATION_CONTEXT_KEY] is None + + @pytest.mark.parametrize("action", ["read", "search"]) def test_studio_user_assistant_discovery_includes_system_and_studio_owned_assistants(action): value = {} diff --git a/backend/tests/test_langgraph_studio_routes.py b/backend/tests/test_langgraph_studio_routes.py index 403b99747..a3c4aba46 100644 --- a/backend/tests/test_langgraph_studio_routes.py +++ b/backend/tests/test_langgraph_studio_routes.py @@ -10,6 +10,7 @@ import subprocess import sys import time from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager from pathlib import Path from uuid import uuid4 @@ -17,15 +18,85 @@ from uuid import uuid4 import httpx import pytest +from deerflow.mcp_scope import mcp_session_scope_key + BACKEND_DIR = Path(__file__).resolve().parents[1] _GRAPH_SOURCE = """ -from langgraph.graph import END, START, StateGraph +from langchain_core.messages import AIMessage +from langchain_core.tools import StructuredTool +from langgraph.prebuilt import ToolNode +from langgraph.graph import END, START, MessagesState, StateGraph +from langgraph.runtime import Runtime +from mcp.types import CallToolResult, TextContent -builder = StateGraph(dict) -builder.add_node("noop", lambda state: {}) -builder.add_edge(START, "noop") -builder.add_edge("noop", END) +from deerflow.mcp import tools as mcp_tools + + +class _FakePool: + async def get_session(self, _server_name, scope_key, _connection): + return object() + + +_POOL = _FakePool() +mcp_tools.get_session_pool = lambda: _POOL + + +async def _unused_probe(): + raise AssertionError("the pooled wrapper must replace this implementation") + + +async def _call_remote( + _session, + _pool, + *, + scope_key, + **_kwargs, +): + return CallToolResult( + content=[TextContent(type="text", text=scope_key)], + ) + + +mcp_tools.call_pooled_session_tool = _call_remote +pooled_probe = mcp_tools._make_session_pool_tool( + StructuredTool( + name="probe", + description="Return the standalone run's pooled MCP session scope.", + args_schema={"type": "object", "properties": {}}, + coroutine=_unused_probe, + ), + "test-server", + {"transport": "streamable_http", "url": "http://unused.invalid/mcp"}, +) + +builder = StateGraph(MessagesState, context_schema=dict) + + +def request_probe(_state, runtime: Runtime): + assert runtime.context["preserved_context_probe"] == "kept" + return { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "name": "probe", + "args": {}, + "id": "probe-call", + "type": "tool_call", + } + ], + ) + ] + } + + +builder.add_node("request_probe", request_probe) +builder.add_node("tools", ToolNode([pooled_probe])) +builder.add_edge(START, "request_probe") +builder.add_edge("request_probe", "tools") +builder.add_edge("tools", END) graph = builder.compile() """.lstrip() @@ -60,6 +131,48 @@ def _free_port() -> int: return int(sock.getsockname()[1]) +def _run_scope( + client: httpx.Client, + thread_id: str, + *, + context: dict | None = None, + if_not_exists: str | None = None, +) -> str: + payload = { + "assistant_id": "test_graph", + "input": {"messages": []}, + "context": {"preserved_context_probe": "kept", **(context or {})}, + } + if if_not_exists is not None: + payload["if_not_exists"] = if_not_exists + response = client.post( + f"/threads/{thread_id}/runs/wait", + json=payload, + ) + assert response.status_code == 200, response.text + messages = response.json()["messages"] + assert messages[-1]["type"] == "tool" + assert messages[-1]["status"] == "success" + return messages[-1]["content"][0]["text"] + + +def _run_stateless_scope(client: httpx.Client, run_id: str) -> str: + response = client.post( + "/runs/wait", + json={ + "assistant_id": "test_graph", + "run_id": run_id, + "input": {"messages": []}, + "context": {"preserved_context_probe": "kept"}, + }, + ) + assert response.status_code == 200, response.text + messages = response.json()["messages"] + assert messages[-1]["type"] == "tool" + assert messages[-1]["status"] == "success" + return messages[-1]["content"][0]["text"] + + @contextmanager def _running_studio_server( runtime_dir: Path, @@ -121,7 +234,7 @@ def _running_studio_server( ) base_url = f"http://127.0.0.1:{port}" - deadline = time.monotonic() + 45 + deadline = time.monotonic() + 90 last_error: Exception | None = None while time.monotonic() < deadline and process.poll() is None: try: @@ -137,7 +250,11 @@ def _running_studio_server( time.sleep(0.1) else: process.terminate() - process.wait(timeout=10) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) pytest.fail(f"LangGraph dev server failed to start ({last_error!r}).\n{log_path.read_text(encoding='utf-8')}") client = httpx.Client( @@ -224,6 +341,217 @@ def test_studio_can_get_and_search_registered_system_assistant( assert response.json()["metadata"]["created_by"] == "system" +def test_studio_registered_graph_supplies_server_owned_mcp_incarnation( + studio_client: httpx.Client, +): + thread_id = str(uuid4()) + response = studio_client.post( + "/threads", + json={ + "thread_id": thread_id, + "metadata": {"thread_incarnation": "attacker"}, + }, + ) + assert response.status_code == 200, response.text + created = response.json() + incarnation = created["metadata"]["thread_incarnation"] + assert incarnation != "attacker" + expected_scope = mcp_session_scope_key( + user_id="langgraph-studio-user", + thread_id=thread_id, + thread_incarnation=incarnation, + ) + assert ( + _run_scope( + studio_client, + thread_id, + context={ + "thread_incarnation": "attacker", + "__deerflow_thread_incarnation_metadata_guard": False, + "user_id": "attacker", + "thread_id": "attacker", + "run_id": "attacker", + }, + ) + == expected_scope + ) + assert _run_scope(studio_client, thread_id) == expected_scope + + response = studio_client.patch( + f"/threads/{thread_id}", + json={"metadata": {"thread_incarnation": "attacker"}}, + ) + assert response.status_code == 200, response.text + assert response.json()["metadata"]["thread_incarnation"] == incarnation + assert _run_scope(studio_client, thread_id) == expected_scope + + response = studio_client.delete(f"/threads/{thread_id}") + assert response.status_code == 204, response.text + response = studio_client.post( + "/threads", + json={ + "thread_id": thread_id, + "metadata": {"thread_incarnation": "attacker"}, + }, + ) + assert response.status_code == 200, response.text + replacement_incarnation = response.json()["metadata"]["thread_incarnation"] + assert replacement_incarnation not in {"attacker", incarnation} + assert _run_scope(studio_client, thread_id) == mcp_session_scope_key( + user_id="langgraph-studio-user", + thread_id=thread_id, + thread_incarnation=replacement_incarnation, + ) + + +def test_studio_implicit_thread_creation_persists_mcp_incarnation( + studio_client: httpx.Client, +): + thread_id = str(uuid4()) + with ThreadPoolExecutor(max_workers=6) as executor: + scopes = list( + executor.map( + lambda _index: _run_scope( + studio_client, + thread_id, + context={"thread_incarnation": "attacker"}, + if_not_exists="create", + ), + range(6), + ) + ) + assert len(set(scopes)) == 1 + scope = scopes[0] + + response = studio_client.get(f"/threads/{thread_id}") + assert response.status_code == 200, response.text + incarnation = response.json()["metadata"]["thread_incarnation"] + assert incarnation != "attacker" + assert scope == mcp_session_scope_key( + user_id="langgraph-studio-user", + thread_id=thread_id, + thread_incarnation=incarnation, + ) + assert _run_scope(studio_client, thread_id) == scope + + +def test_studio_implicit_thread_creation_preserves_searchable_metadata( + studio_client: httpx.Client, +): + thread_id = str(uuid4()) + project_tag = f"project-{uuid4()}" + payload = { + "assistant_id": "test_graph", + "input": {"messages": []}, + "if_not_exists": "create", + "context": {"preserved_context_probe": "kept"}, + "metadata": { + "title": "Run title", + "project_tag": project_tag, + "user_id": "attacker", + "thread_incarnation": "attacker", + "__deerflow_thread_incarnation_metadata_guard": False, + }, + "config": { + "metadata": { + "title": "Config title", + "config_tag": "retained", + "user_id": "config-attacker", + "thread_incarnation": "config-attacker", + "__deerflow_thread_incarnation_metadata_guard": False, + }, + }, + } + response = studio_client.post(f"/threads/{thread_id}/runs/wait", json=payload) + assert response.status_code == 200, response.text + tool_message = response.json()["messages"][-1] + assert tool_message["status"] == "success" + + response = studio_client.get(f"/threads/{thread_id}") + assert response.status_code == 200, response.text + metadata = response.json()["metadata"] + assert metadata["title"] == "Run title" + assert metadata["project_tag"] == project_tag + assert metadata["config_tag"] == "retained" + assert metadata["user_id"] == "langgraph-studio-user" + assert metadata["thread_incarnation"] not in {"attacker", "config-attacker"} + assert "__deerflow_thread_incarnation_metadata_guard" not in metadata + assert tool_message["content"][0]["text"] == mcp_session_scope_key( + user_id="langgraph-studio-user", + thread_id=thread_id, + thread_incarnation=metadata["thread_incarnation"], + ) + + response = studio_client.post("/threads/search", json={"metadata": {"project_tag": project_tag}}) + assert response.status_code == 200, response.text + assert [thread["thread_id"] for thread in response.json()] == [thread_id] + + # A later run must not overwrite the existing thread's creation metadata. + payload["metadata"]["title"] = "Later run title" + payload["metadata"]["project_tag"] = "later-project" + response = studio_client.post(f"/threads/{thread_id}/runs/wait", json=payload) + assert response.status_code == 200, response.text + response = studio_client.get(f"/threads/{thread_id}") + assert response.status_code == 200, response.text + assert response.json()["metadata"] == metadata + + +def test_studio_stateless_runs_get_distinct_mcp_scopes( + studio_client: httpx.Client, +): + first_run_id = str(uuid4()) + second_run_id = str(uuid4()) + + first_scope = _run_stateless_scope(studio_client, first_run_id) + second_scope = _run_stateless_scope(studio_client, second_run_id) + + assert first_scope != second_scope + assert mcp_session_scope_key( + user_id="langgraph-studio-user", + thread_id="default", + thread_incarnation=None, + ) not in {first_scope, second_scope} + + +def test_studio_legacy_thread_uses_explicit_legacy_scope_without_backfill( + tmp_path: Path, +): + thread_id = str(uuid4()) + with _running_studio_server( + tmp_path, + auth_source=_LEGACY_AUTH_SHIM, + ) as legacy_client: + response = legacy_client.post( + "/threads", + json={"thread_id": thread_id, "metadata": {"legacy": True}}, + ) + assert response.status_code == 200, response.text + assert "thread_incarnation" not in response.json()["metadata"] + + with _running_studio_server( + tmp_path, + auth_source=_CURRENT_AUTH_SHIM, + ) as current_client: + with ThreadPoolExecutor(max_workers=2) as executor: + scopes = list( + executor.map( + lambda _index: _run_scope(current_client, thread_id), + range(2), + ) + ) + assert len(set(scopes)) == 1 + scope = scopes[0] + response = current_client.get(f"/threads/{thread_id}") + assert response.status_code == 200, response.text + assert "thread_incarnation" not in response.json()["metadata"] + assert scope == mcp_session_scope_key( + user_id="langgraph-studio-user", + thread_id=thread_id, + thread_incarnation=None, + ) + assert _run_scope(current_client, thread_id) == scope + + def test_studio_update_cannot_forge_system_provenance( studio_client: httpx.Client, ): diff --git a/backend/tests/test_mcp_cwd.py b/backend/tests/test_mcp_cwd.py index 927339795..19831330d 100644 --- a/backend/tests/test_mcp_cwd.py +++ b/backend/tests/test_mcp_cwd.py @@ -13,8 +13,9 @@ from deerflow.mcp.session_pool import MCPSessionPool @pytest.mark.asyncio +@pytest.mark.parametrize("thread_incarnation", [None, "incarnation-1"], ids=["legacy", "versioned"]) @pytest.mark.parametrize("relative_script", [True, False], ids=["relative-entrypoint", "relative-tool-input"]) -async def test_stdio_cwd_from_config_reaches_discovery_and_tool_calls(tmp_path, monkeypatch, relative_script): +async def test_stdio_cwd_from_config_reaches_discovery_and_tool_calls(tmp_path, monkeypatch, relative_script, thread_incarnation): server_dir = tmp_path / "mcp server" server_dir.mkdir() (server_dir / "marker.txt").write_text("configured-directory", encoding="utf-8") @@ -54,7 +55,7 @@ mcp.run(transport="stdio") monkeypatch.setattr(mcp_tools, "get_paths", lambda: Paths(tmp_path / "deerflow")) pool = MCPSessionPool() monkeypatch.setattr(mcp_tools, "get_session_pool", lambda: pool) - runtime = SimpleNamespace(context={"thread_id": "thread", "user_id": "user"}, config={}) + runtime = SimpleNamespace(context={"thread_id": "thread", "user_id": "user", "thread_incarnation": thread_incarnation}, config={}) try: tools = await mcp_tools.get_mcp_tools() @@ -66,12 +67,13 @@ mcp.run(transport="stdio") @pytest.mark.asyncio +@pytest.mark.parametrize("thread_incarnation", [None, "incarnation-1"], ids=["legacy", "versioned"]) @pytest.mark.parametrize( "cwd_config", [{}, {"cwd": None}, {"cwd": ""}, {"cwd": "$TEST_UNSET_MCP_CWD"}], ids=["omitted", "null", "empty", "unset-env"], ) -async def test_empty_stdio_cwd_preserves_default_working_directories(tmp_path, monkeypatch, cwd_config): +async def test_empty_stdio_cwd_preserves_default_working_directories(tmp_path, monkeypatch, cwd_config, thread_incarnation): launch_dir = tmp_path / "gateway" launch_dir.mkdir() monkeypatch.chdir(launch_dir) @@ -102,7 +104,7 @@ mcp.run(transport="stdio") monkeypatch.setattr(mcp_tools, "get_paths", lambda: paths) pool = MCPSessionPool() monkeypatch.setattr(mcp_tools, "get_session_pool", lambda: pool) - runtime = SimpleNamespace(context={"thread_id": "thread", "user_id": "user"}, config={}) + runtime = SimpleNamespace(context={"thread_id": "thread", "user_id": "user", "thread_incarnation": thread_incarnation}, config={}) try: tools = await mcp_tools.get_mcp_tools() diff --git a/backend/tests/test_mcp_session_pool.py b/backend/tests/test_mcp_session_pool.py index 44d743dfd..3d8747ec2 100644 --- a/backend/tests/test_mcp_session_pool.py +++ b/backend/tests/test_mcp_session_pool.py @@ -7,6 +7,7 @@ import stat import sys import threading import weakref +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import anyio @@ -15,6 +16,11 @@ from mcp.shared.exceptions import McpError from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent from deerflow.mcp.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool, reset_session_pool +from deerflow.mcp_scope import ( + THREAD_INCARNATION_METADATA_GUARD_KEY, + mcp_session_scope_key, + runtime_thread_incarnation, +) @pytest.fixture(autouse=True) @@ -24,6 +30,75 @@ def _reset_pool(): reset_session_pool() +def _legacy_tool_runtime(*, thread_id: str = "default"): + return SimpleNamespace( + context={"thread_id": thread_id, "thread_incarnation": None}, + config={}, + ) + + +def test_runtime_incarnation_matches_server_thread_metadata(): + runtime = SimpleNamespace( + context={ + "thread_incarnation": "incarnation-1", + THREAD_INCARNATION_METADATA_GUARD_KEY: True, + }, + config={"metadata": {"thread_incarnation": "incarnation-1"}}, + ) + + assert runtime_thread_incarnation(runtime) == "incarnation-1" + + +@pytest.mark.parametrize("metadata_value", ["incarnation-2", "", False, None]) +def test_runtime_incarnation_rejects_stale_or_invalid_server_thread_metadata( + metadata_value, +): + runtime = SimpleNamespace( + context={ + "thread_incarnation": "incarnation-1", + THREAD_INCARNATION_METADATA_GUARD_KEY: True, + }, + config={"metadata": {"thread_incarnation": metadata_value}}, + ) + + with pytest.raises(RuntimeError, match="stale thread incarnation"): + runtime_thread_incarnation(runtime) + + +def test_runtime_incarnation_ignores_untrusted_metadata_without_server_guard(): + runtime = SimpleNamespace( + context={"thread_incarnation": "incarnation-1"}, + config={"metadata": {"thread_incarnation": "attacker"}}, + ) + + assert runtime_thread_incarnation(runtime) == "incarnation-1" + + +def test_guarded_versioned_incarnation_requires_persisted_metadata(): + runtime = SimpleNamespace( + context={ + "thread_incarnation": "incarnation-1", + THREAD_INCARNATION_METADATA_GUARD_KEY: True, + }, + config={"metadata": {}}, + ) + + with pytest.raises(RuntimeError, match="stale thread incarnation"): + runtime_thread_incarnation(runtime) + + +def test_guarded_legacy_incarnation_allows_missing_persisted_metadata(): + runtime = SimpleNamespace( + context={ + "thread_incarnation": None, + THREAD_INCARNATION_METADATA_GUARD_KEY: True, + }, + config={"metadata": {}}, + ) + + assert runtime_thread_incarnation(runtime) is None + + # --------------------------------------------------------------------------- # MCPSessionPool unit tests # --------------------------------------------------------------------------- @@ -468,7 +543,7 @@ mcp.run(transport="stdio") "args": ["-c", server, str(marker)], } runtime = MagicMock() - runtime.context = {"thread_id": "thread", "user_id": "user"} + runtime.context = {"thread_id": "thread", "user_id": "user", "thread_incarnation": "incarnation-1"} runtime.config = {} with patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)): @@ -477,7 +552,12 @@ mcp.run(transport="stdio") await wrapped.coroutine(runtime=runtime) assert exc_info.value.error.code == CONNECTION_CLOSED - assert ("crash", "user:thread") not in {k[:2] for k in get_session_pool()._entries} + scope_key = mcp_session_scope_key( + user_id="user", + thread_id="thread", + thread_incarnation="incarnation-1", + ) + assert ("crash", scope_key) not in {k[:2] for k in get_session_pool()._entries} content, _artifact = await wrapped.coroutine(runtime=runtime) @@ -504,7 +584,7 @@ async def test_session_pool_tool_evicts_session_after_transport_disconnect(tmp_p patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(type(transport_error)), ): - await wrapped.coroutine(value=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) @@ -529,7 +609,7 @@ async def test_session_pool_tool_evicts_connection_closed_through_interceptor(tm patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(McpError, match="Connection closed"), ): - await wrapped.coroutine(value=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) @@ -553,7 +633,7 @@ async def test_session_pool_tool_keeps_session_after_nonfatal_mcp_error(tmp_path patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(McpError, match=str(error)), ): - await wrapped.coroutine(value=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_not_awaited() @@ -572,7 +652,7 @@ async def test_session_pool_tool_preserves_disconnect_error_when_eviction_fails( patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(anyio.ClosedResourceError) as exc_info, ): - await wrapped.coroutine(value=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) assert exc_info.value is error pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) @@ -636,7 +716,7 @@ async def test_session_pool_tool_keeps_session_after_tool_error_result(tmp_path) patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(ToolException, match="invalid input"), ): - await wrapped.coroutine(value=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) pool.close_session_if_current.assert_not_awaited() @@ -660,7 +740,7 @@ async def test_session_pool_tool_keeps_session_after_interceptor_error(tmp_path) patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)), pytest.raises(RuntimeError, match="interceptor failed"), ): - await wrapped.coroutine(value=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), value=1) session.call_tool.assert_not_awaited() pool.close_session_if_current.assert_not_awaited() @@ -722,15 +802,16 @@ async def test_late_disconnect_from_old_session_does_not_evict_replacement(tmp_p "srv", {"transport": "stdio", "command": "x", "args": []}, ) - first_call = asyncio.create_task(wrapped.coroutine(value=1)) - late_call = asyncio.create_task(wrapped.coroutine(value=2)) + runtime = _legacy_tool_runtime() + first_call = asyncio.create_task(wrapped.coroutine(runtime=runtime, value=1)) + late_call = asyncio.create_task(wrapped.coroutine(runtime=runtime, value=2)) await asyncio.wait_for(both_started.wait(), timeout=1) first_failure.set() with pytest.raises(anyio.ClosedResourceError): await first_call - await wrapped.coroutine(value=3) + await wrapped.coroutine(runtime=runtime, value=3) late_failure.set() with pytest.raises(anyio.ClosedResourceError): await late_call @@ -772,7 +853,7 @@ async def test_session_pool_tool_wrapping(): # Simulate a tool call with a runtime context containing thread_id. mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42"} + mock_runtime.context = {"thread_id": "thread-42", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") @@ -812,7 +893,7 @@ async def test_session_pool_tool_pins_cwd_and_temp_env(tmp_path): paths = Paths(tmp_path) connection = {"transport": "stdio", "command": "pw", "args": [], "env": {"KEEP": "1"}} mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7"} + mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} with ( @@ -868,7 +949,7 @@ async def test_session_pool_tool_does_not_override_explicit_tmpdir(tmp_path): paths = Paths(tmp_path) connection = {"transport": "stdio", "command": "pw", "args": [], "env": {"TMPDIR": "/operator/tmp"}} mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7"} + mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} with ( @@ -915,7 +996,7 @@ async def test_session_pool_tool_does_not_override_explicit_cwd(tmp_path): paths = Paths(tmp_path) connection = {"transport": "stdio", "command": "pw", "args": [], "cwd": operator_cwd} mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7"} + mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} with ( @@ -962,7 +1043,7 @@ async def test_session_pool_tool_skips_fs_work_for_non_stdio_transport(tmp_path) paths = Paths(tmp_path) connection = {"transport": "sse", "url": "http://localhost:9000/sse", "env": {"KEEP": "1"}} mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7"} + mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} with ( @@ -1015,7 +1096,7 @@ async def test_session_pool_tool_skips_after_walk_when_no_text_content(tmp_path) paths = Paths(tmp_path) connection = {"transport": "stdio", "command": "pw", "args": []} mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7"} + mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} with ( @@ -1061,7 +1142,7 @@ async def test_session_pool_tool_runs_after_walk_when_text_content_present(tmp_p paths = Paths(tmp_path) connection = {"transport": "stdio", "command": "pw", "args": []} mock_runtime = MagicMock() - mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7"} + mock_runtime.context = {"thread_id": "thread-42", "user_id": "user-7", "thread_incarnation": "incarnation-1"} mock_runtime.config = {} with ( @@ -1113,7 +1194,7 @@ async def test_session_pool_tool_forwards_interceptor_headers(): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[header_interceptor], ) - await wrapped.coroutine(runtime=None, x=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1}, meta={"headers": {"X-User-Id": "u-42"}}) @@ -1162,7 +1243,7 @@ async def test_session_pool_interceptor_reads_request_scoped_secret(): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[secret_header_interceptor], ) - await wrapped.coroutine(runtime=None, x=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with( "act", @@ -1208,7 +1289,7 @@ async def test_session_pool_tool_no_headers_omits_meta(): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[passthrough_interceptor], ) - await wrapped.coroutine(runtime=None, x=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1}) @@ -1252,7 +1333,7 @@ async def test_session_pool_tool_ignores_unsupported_header_type(caplog): {"transport": "stdio", "command": "x", "args": []}, tool_interceptors=[invalid_header_interceptor], ) - await wrapped.coroutine(runtime=None, x=1) + await wrapped.coroutine(runtime=_legacy_tool_runtime(), x=1) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1}) assert "unsupported type" in caplog.text @@ -1287,16 +1368,20 @@ async def test_session_pool_tool_extracts_thread_id(): wrapped = _make_session_pool_tool(original_tool, "server", {"transport": "stdio", "command": "x", "args": []}) mock_runtime = MagicMock() - mock_runtime.context = {} + mock_runtime.context = {"thread_incarnation": "incarnation-1"} mock_runtime.config = {"configurable": {"thread_id": "from-config"}} await wrapped.coroutine(runtime=mock_runtime, x=1) - # Verify the session was created with the correct scope key. - # The scope key is "{user_id}:{thread_id}"; the autouse fixture sets - # the effective user to "test-user-autouse". + # Verify the session was created with the canonical versioned JSON scope; + # the autouse fixture sets the effective user to "test-user-autouse". pool = get_session_pool() - assert ("server", "test-user-autouse:from-config") in {k[:2] for k in pool._entries} + expected_scope = mcp_session_scope_key( + user_id="test-user-autouse", + thread_id="from-config", + thread_incarnation="incarnation-1", + ) + assert ("server", expected_scope) in {k[:2] for k in pool._entries} @pytest.mark.asyncio diff --git a/backend/tests/test_mcp_sync_wrapper.py b/backend/tests/test_mcp_sync_wrapper.py index 89550a56d..929461e78 100644 --- a/backend/tests/test_mcp_sync_wrapper.py +++ b/backend/tests/test_mcp_sync_wrapper.py @@ -289,7 +289,12 @@ def test_func_patched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path): graph.compile().ainvoke( {"messages": [ai]}, config={"configurable": {"thread_id": "T"}}, - context={"thread_id": "T", "run_id": "run-1", "user_id": "alice"}, + context={ + "thread_id": "T", + "thread_incarnation": "incarnation-1", + "run_id": "run-1", + "user_id": "alice", + }, ) ) diff --git a/backend/tests/test_mcp_task_models.py b/backend/tests/test_mcp_task_models.py index 6486de0cb..ae822a064 100644 --- a/backend/tests/test_mcp_task_models.py +++ b/backend/tests/test_mcp_task_models.py @@ -2,7 +2,14 @@ from datetime import timedelta import pytest -from deerflow.mcp.tasks import McpTaskDriverRegistry, TaskSnapshot, TaskStatus, TaskSubmission, TaskSubmitRequest +from deerflow.mcp.tasks import ( + McpTaskDriverRegistry, + TaskReference, + TaskSnapshot, + TaskStatus, + TaskSubmission, + TaskSubmitRequest, +) def test_task_snapshot_normalizes_string_statuses(): @@ -55,6 +62,36 @@ def test_task_storage_identifiers_reject_values_longer_than_the_database_columns ) +def test_incarnation_fields_preserve_legacy_positional_construction(): + reference_driver_data = {"status_tool": "status"} + reference = TaskReference( + "task-1", + "user-1", + "thread-1", + "reports", + "remote-1", + reference_driver_data, + ) + assert reference.driver_data is reference_driver_data + assert reference.thread_incarnation is None + + request_driver_data = {"submit_tool": "submit"} + request = TaskSubmitRequest( + "user-1", + "thread-1", + "run-1", + "call-1", + "reports", + "Generate report", + {}, + request_driver_data, + "task-1", + ) + assert request.driver_data is request_driver_data + assert request.local_task_id == "task-1" + assert request.thread_incarnation is None + + def test_driver_registry_rejects_duplicate_names(): registry = McpTaskDriverRegistry() driver = object() diff --git a/backend/tests/test_mcp_task_ordinary_driver.py b/backend/tests/test_mcp_task_ordinary_driver.py index 25ec5532b..ddc3025e4 100644 --- a/backend/tests/test_mcp_task_ordinary_driver.py +++ b/backend/tests/test_mcp_task_ordinary_driver.py @@ -28,6 +28,7 @@ def _request() -> TaskSubmitRequest: return TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -47,6 +48,7 @@ def _reference() -> TaskReference: local_task_id="local-1", user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", server_name="reports", remote_task_id="remote-1", driver_data={ @@ -74,6 +76,7 @@ async def test_submit_uses_structured_content_and_keeps_remote_id_out_of_driver_ "arguments": {"topic": "MCP"}, "user_id": "user-1", "thread_id": "thread-1", + "thread_incarnation": "incarnation-1", # Submit is the one durable-task call awaited inside the Agent run, # so it is the only one that may carry request-scoped credentials. "request_scoped_headers": True, diff --git a/backend/tests/test_mcp_task_ordinary_e2e.py b/backend/tests/test_mcp_task_ordinary_e2e.py index e2ef5f375..5e84836ca 100644 --- a/backend/tests/test_mcp_task_ordinary_e2e.py +++ b/backend/tests/test_mcp_task_ordinary_e2e.py @@ -14,6 +14,7 @@ from deerflow.mcp.tasks import ( ) from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config from deerflow.persistence.mcp_tasks import McpTaskRepository +from deerflow.persistence.thread_meta.model import ThreadMetaRow @pytest_asyncio.fixture(autouse=True) @@ -67,6 +68,7 @@ def _request(remote_id: str) -> TaskSubmitRequest: return TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -80,12 +82,29 @@ def _request(remote_id: str) -> TaskSubmitRequest: ) +async def _create_thread(repo: McpTaskRepository) -> None: + now = datetime.now(UTC) + async with repo._sf() as session: + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation="incarnation-1", + user_id="user-1", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + + @pytest.mark.asyncio async def test_submit_poll_restart_recovery_complete_and_fail(tmp_path) -> None: await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path))) session_factory = get_session_factory() assert session_factory is not None repo = McpTaskRepository(session_factory) + await _create_thread(repo) fake_server = FakeMcpServer() submitted_at = datetime.now(UTC) fake_server.status_results.extend( @@ -116,7 +135,7 @@ async def test_submit_poll_restart_recovery_complete_and_fail(tmp_path) -> None: restarted_process = _service(repo, fake_server) await restarted_process.run_once(now=datetime.now(UTC) + timedelta(seconds=2)) - completed = await repo.get(created["id"], user_id="user-1") + completed = await repo.get(created["id"], user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") assert completed is not None assert completed["status"] == "completed" assert completed["result"] == {"report": "ready"} @@ -135,7 +154,7 @@ async def test_submit_poll_restart_recovery_complete_and_fail(tmp_path) -> None: ) await restarted_process.run_once(now=datetime.now(UTC)) - failed = await repo.get(failed_created["id"], user_id="user-1") + failed = await repo.get(failed_created["id"], user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") assert failed is not None assert failed["status"] == "failed" assert failed["error"] == "report generation failed" @@ -147,6 +166,7 @@ async def test_status_tool_error_retries_with_detail_before_structured_failure_t session_factory = get_session_factory() assert session_factory is not None repo = McpTaskRepository(session_factory) + await _create_thread(repo) fake_server = FakeMcpServer() service = _service(repo, fake_server) submitted_at = datetime.now(UTC) @@ -172,7 +192,7 @@ async def test_status_tool_error_retries_with_detail_before_structured_failure_t await service.run_once(now=submitted_at + timedelta(seconds=2)) - retrying = await repo.get(created["id"], user_id="user-1") + retrying = await repo.get(created["id"], user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") assert retrying is not None assert retrying["status"] == "submitted" assert retrying["consecutive_poll_error_count"] == 1 @@ -181,7 +201,7 @@ async def test_status_tool_error_retries_with_detail_before_structured_failure_t await service.run_once(now=datetime.now(UTC) + timedelta(seconds=10)) - failed = await repo.get(created["id"], user_id="user-1") + failed = await repo.get(created["id"], user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") assert failed is not None assert failed["status"] == "failed" assert failed["error"] == "report generation failed" diff --git a/backend/tests/test_mcp_task_postgres.py b/backend/tests/test_mcp_task_postgres.py index 3375f1711..d319f47e2 100644 --- a/backend/tests/test_mcp_task_postgres.py +++ b/backend/tests/test_mcp_task_postgres.py @@ -12,7 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.config.database_config import DatabaseConfig from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config -from deerflow.persistence.mcp_tasks import McpTaskRepository +from deerflow.persistence.mcp_tasks import McpTaskRepository, McpTaskThreadMismatchError from deerflow.persistence.mcp_tasks.model import McpTaskRow from deerflow.persistence.thread_meta import ThreadMetaRepository from deerflow.persistence.thread_meta.model import ThreadMetaRow @@ -54,11 +54,17 @@ async def postgres_repositories(): await close_engine() -async def _create_task(repo: McpTaskRepository, task_id: str) -> None: +async def _create_task( + repo: McpTaskRepository, + task_id: str, + *, + thread_incarnation: str, +) -> None: await repo.create( task_id=task_id, user_id="user-1", thread_id="thread-1", + expected_thread_incarnation=thread_incarnation, run_id=None, tool_call_id=None, server_name="reports", @@ -96,24 +102,30 @@ async def test_postgres_task_create_serializes_with_thread_mutation(postgres_rep # PostgreSQL grants the mutation's earlier queued row-lock request # before this later FOR SHARE request. The task therefore observes the # committed delete/owner change rather than the pre-mutation row. - create_task = asyncio.create_task(_create_task(task_repo, f"task-{mutation}")) + create_task = asyncio.create_task( + _create_task( + task_repo, + f"task-{mutation}", + thread_incarnation=created["incarnation"], + ) + ) await asyncio.sleep(0.1) assert not create_task.done() await blocker.commit() await asyncio.wait_for(mutation_task, timeout=5) - await asyncio.wait_for(create_task, timeout=5) + with pytest.raises(McpTaskThreadMismatchError): + await asyncio.wait_for(create_task, timeout=5) async with session_factory() as session: task = await session.get(McpTaskRow, f"task-{mutation}") - assert task is not None - assert task.thread_incarnation is None + assert task is None @pytest.mark.asyncio async def test_postgres_task_create_uses_share_lock(postgres_repositories) -> None: thread_repo, task_repo, _session_factory = postgres_repositories - await thread_repo.create("thread-1", user_id="user-1") + created = await thread_repo.create("thread-1", user_id="user-1") engine = get_engine() assert engine is not None statements: list[str] = [] @@ -123,7 +135,11 @@ async def test_postgres_task_create_uses_share_lock(postgres_repositories) -> No event.listen(engine.sync_engine, "before_cursor_execute", capture_statement) try: - await _create_task(task_repo, "task-share-lock") + await _create_task( + task_repo, + "task-share-lock", + thread_incarnation=created["incarnation"], + ) finally: event.remove(engine.sync_engine, "before_cursor_execute", capture_statement) @@ -173,7 +189,13 @@ async def test_postgres_repository_holds_share_lock_until_task_commit(postgres_r await legacy_writer.commit() event.listen(engine.sync_engine, "before_cursor_execute", observe_owner_update) - task_create = asyncio.create_task(_create_task(task_repo, "task-lock-lifetime")) + task_create = asyncio.create_task( + _create_task( + task_repo, + "task-lock-lifetime", + thread_incarnation=created["incarnation"], + ) + ) owner_update = None try: await asyncio.wait_for(task_commit_entered.wait(), timeout=5) diff --git a/backend/tests/test_mcp_task_repository.py b/backend/tests/test_mcp_task_repository.py index eb429f2d6..dc7f89148 100644 --- a/backend/tests/test_mcp_task_repository.py +++ b/backend/tests/test_mcp_task_repository.py @@ -11,7 +11,11 @@ from sqlalchemy.ext.asyncio import AsyncSession from deerflow.config.database_config import DatabaseConfig from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config -from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError, McpTaskRepository +from deerflow.persistence.mcp_tasks import ( + DuplicateMcpRemoteTaskError, + McpTaskRepository, + McpTaskThreadMismatchError, +) from deerflow.persistence.mcp_tasks.model import McpTaskRow from deerflow.persistence.thread_meta.model import ThreadMetaRow @@ -36,11 +40,26 @@ async def _create_working_task( now: datetime, user_id: str = "user-1", remote_task_id: str | None = None, + thread_incarnation: str | None = None, ) -> dict: + async with repo._sf() as session: + if await session.get(ThreadMetaRow, "thread-1") is None: + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation=None, + user_id=user_id, + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() return await repo.create( task_id=task_id, user_id=user_id, thread_id="thread-1", + expected_thread_incarnation=thread_incarnation, run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -108,7 +127,7 @@ async def test_interleaved_reclaim_fences_inflight_poll_and_cancel_mutations(tmp await _create_working_task(repo, task_id=task_id, now=now) claim = repo.claim_due_tasks if operation == "apply_cancel_snapshot": - await repo.request_cancel(task_id, user_id="user-1", thread_id="thread-1", requested_at=now) + await repo.request_cancel(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None, requested_at=now) claim = repo.claim_cancel_requests first = await claim(now=now, lease_owner="worker-1", lease_seconds=60, limit=1) kwargs = {"lease_owner": "worker-1", "lease_token": first[0]["lease_token"]} @@ -135,13 +154,14 @@ async def test_interleaved_reclaim_fences_inflight_poll_and_cancel_mutations(tmp second = await asyncio.wait_for(claim(now=now + timedelta(seconds=61), lease_owner="worker-1", lease_seconds=60, limit=1), timeout=5) assert len(second) == 1 assert second[0]["lease_token"] != first[0]["lease_token"] - before = await repo.get(task_id, user_id="user-1") + before = await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert before is not None resume.set() applied = await asyncio.wait_for(pending, timeout=5) # Check the entire row, including scheduling, errors, results and event # versions, not just the new lease: stale work must have no side effects. - assert await repo.get(task_id, user_id="user-1") == before + assert await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) == before assert applied is False @@ -192,11 +212,12 @@ async def test_interleaved_reclaim_fences_inflight_notification_completion(tmp_p second = await asyncio.wait_for(repo.claim_notification_work(now=now + timedelta(seconds=61), **claim_kwargs), timeout=5) assert len(second) == 1 assert second[0]["notification_lease_token"] != first[0]["notification_lease_token"] - before = await repo.get(task_id, user_id="user-1") + before = await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) + assert before is not None resume.set() applied = await asyncio.wait_for(pending, timeout=5) - assert await repo.get(task_id, user_id="user-1") == before + assert await repo.get(task_id, user_id="user-1", thread_id="thread-1", thread_incarnation=None) == before assert applied is False @@ -241,21 +262,19 @@ async def test_legacy_task_writer_leaves_thread_incarnation_null(tmp_path): @pytest.mark.asyncio @pytest.mark.parametrize( - ("thread_owner", "expected_incarnation"), + "thread_owner", [ - ("user-1", "matching-owner"), - (None, "shared-thread"), - ("user-2", None), + "user-1", + None, ], ) async def test_create_atomically_copies_accessible_thread_incarnation( tmp_path, thread_owner, - expected_incarnation, ): repo = await _make_repo(tmp_path) now = datetime.now(UTC) - incarnation = expected_incarnation or "different-owner" + incarnation = "matching-incarnation" async with repo._sf() as session: session.add( ThreadMetaRow( @@ -269,25 +288,49 @@ async def test_create_atomically_copies_accessible_thread_incarnation( ) await session.commit() - task = await _create_working_task(repo, task_id="new-writer", now=now) + task = await _create_working_task( + repo, + task_id="new-writer", + now=now, + thread_incarnation=incarnation, + ) assert "thread_incarnation" not in task async with repo._sf() as session: row = await session.get(McpTaskRow, "new-writer") assert row is not None - assert row.thread_incarnation == expected_incarnation + assert row.thread_incarnation == incarnation @pytest.mark.asyncio -async def test_create_leaves_incarnation_null_without_matching_thread(tmp_path): +async def test_create_rejects_missing_or_inaccessible_thread(tmp_path): repo = await _make_repo(tmp_path) - await _create_working_task(repo, task_id="missing-thread", now=datetime.now(UTC)) + with pytest.raises(McpTaskThreadMismatchError): + await repo.create( + task_id="missing-thread", + user_id="user-1", + thread_id="missing-thread", + expected_thread_incarnation=None, + run_id="run-1", + tool_call_id="call-1", + server_name="reports", + driver_name="fake", + remote_task_id="remote-missing", + task_name="Generate report", + status="working", + result=None, + result_preview=None, + result_truncated=False, + result_artifact=None, + error=None, + input_required=None, + next_poll_at=None, + ) async with repo._sf() as session: row = await session.get(McpTaskRow, "missing-thread") - assert row is not None - assert row.thread_incarnation is None + assert row is None @pytest.mark.asyncio @@ -310,9 +353,9 @@ async def test_create_observes_delete_and_recreate_at_insert_boundary(tmp_path): engine = get_engine() assert engine is not None replaced = False - insert_statement = None + lock_statement = None - def replace_thread_before_task_insert( + def replace_thread_before_scope_lock( _conn, _cursor, statement, @@ -320,11 +363,11 @@ async def test_create_observes_delete_and_recreate_at_insert_boundary(tmp_path): _context, _executemany, ): - nonlocal insert_statement, replaced - if replaced or not statement.lstrip().upper().startswith("INSERT INTO MCP_TASKS"): + nonlocal lock_statement, replaced + if replaced or "UPDATE THREADS_META" not in statement.upper(): return replaced = True - insert_statement = statement + lock_statement = statement with contextlib.closing(sqlite3.connect(tmp_path / "deerflow.db")) as connection: with connection: connection.execute("DELETE FROM threads_meta WHERE thread_id = ?", ("thread-1",)) @@ -346,21 +389,187 @@ async def test_create_observes_delete_and_recreate_at_insert_boundary(tmp_path): ), ) - event.listen(engine.sync_engine, "before_cursor_execute", replace_thread_before_task_insert) + event.listen(engine.sync_engine, "before_cursor_execute", replace_thread_before_scope_lock) try: - await _create_working_task(repo, task_id="racing-task", now=now) + with pytest.raises(McpTaskThreadMismatchError): + await _create_working_task( + repo, + task_id="racing-task", + now=now, + thread_incarnation="old-incarnation", + ) finally: - event.remove(engine.sync_engine, "before_cursor_execute", replace_thread_before_task_insert) + event.remove(engine.sync_engine, "before_cursor_execute", replace_thread_before_scope_lock) assert replaced is True - assert insert_statement is not None - normalized_insert = " ".join(insert_statement.upper().split()) - assert "SELECT THREADS_META.INCARNATION" in normalized_insert - assert "INSERT INTO MCP_TASKS" in normalized_insert + assert lock_statement is not None + normalized_lock = " ".join(lock_statement.upper().split()) + assert "UPDATE THREADS_META SET INCARNATION = INCARNATION" in normalized_lock + assert "INCARNATION IS ?" in normalized_lock async with repo._sf() as session: row = await session.get(McpTaskRow, "racing-task") - assert row is not None - assert row.thread_incarnation == "replacement-incarnation" + assert row is None + + +@pytest.mark.asyncio +async def test_request_cancel_rejects_delete_recreate_before_scope_lock(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + async with repo._sf() as session: + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation="old-incarnation", + user_id="user-1", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + await _create_working_task( + repo, + task_id="old-task", + now=now, + thread_incarnation="old-incarnation", + ) + + engine = get_engine() + assert engine is not None + replaced = False + + def replace_thread_before_scope_lock( + _conn, + _cursor, + statement, + _parameters, + _context, + _executemany, + ): + nonlocal replaced + if replaced or "UPDATE THREADS_META" not in statement.upper(): + return + replaced = True + with contextlib.closing(sqlite3.connect(tmp_path / "deerflow.db")) as connection: + with connection: + connection.execute("DELETE FROM threads_meta WHERE thread_id = ?", ("thread-1",)) + connection.execute( + """ + INSERT INTO threads_meta ( + thread_id, incarnation, user_id, status, metadata_json, + created_at, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + "thread-1", + "replacement-incarnation", + "user-1", + "idle", + "{}", + now.isoformat(), + now.isoformat(), + ), + ) + + event.listen(engine.sync_engine, "before_cursor_execute", replace_thread_before_scope_lock) + try: + result = await repo.request_cancel( + "old-task", + user_id="user-1", + thread_id="thread-1", + thread_incarnation="old-incarnation", + requested_at=now, + ) + finally: + event.remove(engine.sync_engine, "before_cursor_execute", replace_thread_before_scope_lock) + + assert replaced is True + assert result is None + async with repo._sf() as session: + task = await session.get(McpTaskRow, "old-task") + assert task is not None + assert task.cancel_requested_at is None + + +@pytest.mark.asyncio +async def test_user_access_is_limited_to_current_thread_incarnation(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="old-task", now=now) + + async with repo._sf() as session: + old_thread = await session.get(ThreadMetaRow, "thread-1") + assert old_thread is not None + await session.delete(old_thread) + await session.commit() + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation="replacement-incarnation", + user_id="user-1", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + + assert await repo.list_by_thread("thread-1", user_id="user-1", thread_incarnation=None) == [] + assert await repo.get("old-task", user_id="user-1", thread_id="thread-1", thread_incarnation=None) is None + assert ( + await repo.request_cancel( + "old-task", + user_id="user-1", + thread_id="thread-1", + thread_incarnation=None, + requested_at=now, + ) + is None + ) + + +@pytest.mark.asyncio +async def test_pr1_does_not_change_worker_claim_eligibility(tmp_path): + repo = await _make_repo(tmp_path) + now = datetime.now(UTC) + await _create_working_task(repo, task_id="old-task", now=now) + + async with repo._sf() as session: + old_thread = await session.get(ThreadMetaRow, "thread-1") + assert old_thread is not None + await session.delete(old_thread) + await session.commit() + session.add( + ThreadMetaRow( + thread_id="thread-1", + incarnation="replacement-incarnation", + user_id="user-1", + metadata_json={}, + created_at=now, + updated_at=now, + ) + ) + await session.commit() + + claimed = await repo.claim_due_tasks( + now=now, + lease_owner="worker-1", + lease_seconds=60, + limit=10, + ) + + assert [task["id"] for task in claimed] == ["old-task"] + assert claimed[0]["_thread_incarnation"] is None + + +@pytest.mark.asyncio +async def test_user_access_treats_legacy_null_incarnations_as_equal(tmp_path): + repo = await _make_repo(tmp_path) + task = await _create_working_task(repo, task_id="legacy-task", now=datetime.now(UTC)) + + assert "thread_incarnation" not in task + assert "_thread_incarnation" not in task + assert await repo.get("legacy-task", user_id="user-1", thread_id="thread-1", thread_incarnation=None) is not None @pytest.mark.asyncio @@ -382,6 +591,12 @@ async def test_remote_task_id_is_unique_per_user_and_server(tmp_path): remote_task_id="shared-remote-id", ) + async with repo._sf() as session: + thread = await session.get(ThreadMetaRow, "thread-1") + assert thread is not None + thread.user_id = None + await session.commit() + other_user = await _create_working_task( repo, task_id="task-remote-3", @@ -483,7 +698,7 @@ async def test_apply_snapshot_requires_current_lease_owner_and_terminalizes_task ) assert applied is True - stored = await repo.get("task-2", user_id="user-1") + stored = await repo.get("task-2", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["status"] == "completed" assert stored["result"] == {"report": "ready"} @@ -533,7 +748,7 @@ async def test_apply_snapshot_rejects_result_after_same_workers_lease_expires(tm ) assert applied is False - stored = await repo.get("task-expired", user_id="user-1") + stored = await repo.get("task-expired", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["status"] == "working" assert stored["result"] is None @@ -567,7 +782,7 @@ async def test_input_required_is_persisted_and_remains_scheduled_for_slow_pollin ) assert applied is True - stored = await repo.get("task-3", user_id="user-1") + stored = await repo.get("task-3", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["input_required"] == {"prompt": "Approve deployment?"} assert stored["notification_status"] == "pending" @@ -596,7 +811,7 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path): ) assert released is True - stored = await repo.get("task-4", user_id="user-1") + stored = await repo.get("task-4", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["status"] == "working" assert stored["last_poll_error"] == "temporary network failure" @@ -630,7 +845,7 @@ async def test_release_claim_after_same_worker_reclaim_cannot_clear_new_claim(tm ) assert released is False - stored = await repo.get("task-fence", user_id="user-1") + stored = await repo.get("task-fence", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-1" assert stored["lease_token"] == new_token @@ -667,7 +882,7 @@ async def test_apply_snapshot_after_same_worker_reclaim_cannot_clear_new_claim(t ) assert applied is False - stored = await repo.get("task-apply-fence", user_id="user-1") + stored = await repo.get("task-apply-fence", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-1" assert stored["lease_token"] == new_token @@ -703,7 +918,7 @@ async def test_apply_cancel_snapshot_after_same_worker_reclaim_cannot_clear_new_ ) assert applied is False - stored = await repo.get("task-cancel-fence", user_id="user-1") + stored = await repo.get("task-cancel-fence", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-1" assert stored["lease_token"] == new_token @@ -768,7 +983,7 @@ async def test_finish_notification_run_after_reclaim_cannot_clear_new_claim(tmp_ ) assert finished is False - stored = await repo.get("task-notify-fence", user_id="user-1") + stored = await repo.get("task-notify-fence", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_lease_owner"] == "notifier" assert stored["notification_lease_token"] == new_notify_token @@ -788,7 +1003,7 @@ async def test_release_poll_claim_after_cancellation_preserves_poll_failure_stat next_poll_at=retry_at, error="temporary network failure", ) - before = await repo.get("task-cancelled-poll", user_id="user-1") + before = await repo.get("task-cancelled-poll", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert before is not None reclaimed = await repo.claim_due_tasks(now=retry_at, lease_owner="worker-2", lease_seconds=60, limit=10) @@ -799,7 +1014,7 @@ async def test_release_poll_claim_after_cancellation_preserves_poll_failure_stat ) assert released is True - stored = await repo.get("task-cancelled-poll", user_id="user-1") + stored = await repo.get("task-cancelled-poll", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["next_poll_at"] == before["next_poll_at"] assert stored["last_poll_error"] == before["last_poll_error"] @@ -823,7 +1038,7 @@ async def test_release_poll_claim_after_cancellation_requires_current_owner(tmp_ ) assert released is False - stored = await repo.get("task-stale-cancel", user_id="user-1") + stored = await repo.get("task-stale-cancel", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-current" assert stored["lease_expires_at"] is not None @@ -844,7 +1059,7 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp next_poll_at=now - timedelta(seconds=1), error="temporary network failure", ) - stored = await repo.get("task-6", user_id="user-1") + stored = await repo.get("task-6", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["consecutive_poll_error_count"] == expected_errors @@ -865,7 +1080,7 @@ async def test_consecutive_poll_error_count_increments_and_resets_on_success(tmp ) assert applied is True - stored = await repo.get("task-6", user_id="user-1") + stored = await repo.get("task-6", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["consecutive_poll_error_count"] == 0 @@ -916,7 +1131,7 @@ async def test_notification_snapshot_is_versioned_and_not_overwritten_in_flight( next_poll_at=None, polled_at=now, ) - changed = await repo.get("task-notify", user_id="user-1") + changed = await repo.get("task-notify", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert changed is not None assert changed["event_version"] == 2 assert changed["dispatch_version"] == 1 @@ -1011,7 +1226,7 @@ async def test_notification_retry_rebuilds_a_newer_event_and_resets_its_budget(t error="Agent run failed", now=now, ) - failed = await repo.get("task-retry-latest", user_id="user-1") + failed = await repo.get("task-retry-latest", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert failed is not None assert failed["notification_status"] == "retry" assert failed["dispatch_attempt"] == 1 @@ -1083,7 +1298,7 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p error="run store unavailable", ) - stored = await repo.get("task-notify-release", user_id="user-1") + stored = await repo.get("task-notify-release", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_status"] == "claimed" assert stored["notification_lease_owner"] is None @@ -1130,7 +1345,7 @@ async def test_notification_launch_failure_counts_and_reclaims_latest_snapshot(t count_failure=True, ) - stored = await repo.get("task-launch-retry", user_id="user-1") + stored = await repo.get("task-launch-retry", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_status"] == "pending" assert stored["notification_attempt_count"] == 1 @@ -1185,7 +1400,7 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path): now=now, ) - stored = await repo.get("task-dead-letter", user_id="user-1") + stored = await repo.get("task-dead-letter", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_status"] == "dead_letter" assert stored["notification_attempt_count"] == 1 @@ -1257,7 +1472,7 @@ async def test_dispatched_notification_can_be_dead_lettered_after_retry_budget(t now=now, ) - stored = await repo.get("task-dispatched-budget", user_id="user-1") + stored = await repo.get("task-dispatched-budget", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_status"] == "dead_letter" assert stored["notification_run_id"] is None @@ -1333,7 +1548,7 @@ async def test_dead_lettering_dispatched_snapshot_preserves_newer_event(tmp_path now=now, ) - stored = await repo.get("task-dispatched-latest", user_id="user-1") + stored = await repo.get("task-dispatched-latest", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_status"] == "pending" assert stored["notification_attempt_count"] == 0 @@ -1360,6 +1575,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa "task-cancel", user_id="user-1", thread_id="thread-1", + thread_incarnation=None, requested_at=now, ) assert requested is not None @@ -1394,6 +1610,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa "task-cancel", user_id="user-1", thread_id="thread-1", + thread_incarnation=None, requested_at=now + timedelta(seconds=1), ) assert repeated is not None @@ -1413,7 +1630,7 @@ async def test_cancel_request_stops_polling_and_rejects_stale_poll_result(tmp_pa input_required=None, completed_at=now, ) - stored = await repo.get("task-cancel", user_id="user-1") + stored = await repo.get("task-cancel", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["status"] == "cancelled" assert stored["notification_status"] == "pending" @@ -1451,7 +1668,7 @@ async def test_late_poll_release_after_same_worker_reclaim_is_fenced(tmp_path): ) assert released is False - stored = await repo.get("task-late-poll", user_id="user-1") + stored = await repo.get("task-late-poll", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-same" assert stored["lease_token"] == reclaimed[0]["lease_token"] @@ -1466,6 +1683,7 @@ async def test_late_cancel_release_after_same_worker_reclaim_is_fenced(tmp_path) "task-late-cancel", user_id="user-1", thread_id="thread-1", + thread_incarnation=None, requested_at=now, ) @@ -1497,7 +1715,7 @@ async def test_late_cancel_release_after_same_worker_reclaim_is_fenced(tmp_path) ) assert released is False - stored = await repo.get("task-late-cancel", user_id="user-1") + stored = await repo.get("task-late-cancel", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-same" assert stored["lease_token"] == reclaimed[0]["lease_token"] @@ -1554,7 +1772,7 @@ async def test_late_notification_release_after_same_worker_reclaim_is_fenced(tmp ) assert released is False - stored = await repo.get("task-late-notify", user_id="user-1") + stored = await repo.get("task-late-notify", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["notification_lease_owner"] == "notifier-same" assert stored["notification_lease_token"] == reclaimed[0]["notification_lease_token"] @@ -1598,7 +1816,7 @@ async def test_late_snapshot_apply_after_same_worker_reclaim_is_fenced(tmp_path) ) assert applied is False - stored = await repo.get("task-late-apply", user_id="user-1") + stored = await repo.get("task-late-apply", user_id="user-1", thread_id="thread-1", thread_incarnation=None) assert stored is not None assert stored["lease_owner"] == "worker-same" assert stored["lease_token"] == reclaimed[0]["lease_token"] diff --git a/backend/tests/test_mcp_task_service.py b/backend/tests/test_mcp_task_service.py index 531efb149..5c0e55661 100644 --- a/backend/tests/test_mcp_task_service.py +++ b/backend/tests/test_mcp_task_service.py @@ -19,7 +19,10 @@ from deerflow.mcp.tasks import ( TaskSubmitRequest, ) from deerflow.mcp.tasks.ordinary import McpTaskProtocolError -from deerflow.persistence.mcp_tasks import DuplicateMcpRemoteTaskError +from deerflow.persistence.mcp_tasks import ( + DuplicateMcpRemoteTaskError, + McpTaskThreadMismatchError, +) from deerflow.runtime.runs.manager import ConflictError from deerflow.runtime.runs.schemas import RunStatus @@ -85,6 +88,12 @@ class DuplicateCreateRepository(FakeRepository): raise DuplicateMcpRemoteTaskError("already tracked") +class DriftedThreadCreateRepository(FakeRepository): + async def create(self, **kwargs): + self.created.append(kwargs) + raise McpTaskThreadMismatchError("thread incarnation changed") + + class FakeDriver: def __init__( self, @@ -202,6 +211,7 @@ async def test_submit_persists_remote_handle_before_returning(): request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -214,6 +224,7 @@ async def test_submit_persists_remote_handle_before_returning(): assert created["remote_task_id"] == "remote-1" persisted = repo.created[0] + assert persisted["expected_thread_incarnation"] == "incarnation-1" assert persisted["next_poll_at"] == now + timedelta(seconds=9) assert persisted["driver_data"] == {"submit_tool": "submit", "status_tool": "status"} assert driver.submit_calls[0].local_task_id == created["id"] @@ -241,6 +252,7 @@ async def test_submit_cancels_remote_task_when_persistence_fails(): request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -264,6 +276,47 @@ async def test_submit_cancels_remote_task_when_persistence_fails(): } +@pytest.mark.asyncio +async def test_submit_compensates_remote_task_when_thread_incarnation_drifted(): + repo = DriftedThreadCreateRepository() + driver = FakeDriver( + submission=TaskSubmission( + remote_task_id="remote-1", + snapshot=TaskSnapshot(status=TaskStatus.SUBMITTED), + driver_data={"cancel_tool": "cancel"}, + ) + ) + registry = McpTaskDriverRegistry() + registry.register("fake", driver) + service = McpTaskService( + repository=repo, + drivers=registry, + poll_interval_seconds=5, + lease_seconds=120, + max_concurrent_polls=3, + ) + + with pytest.raises(McpTaskThreadMismatchError, match="incarnation changed"): + await service.submit( + driver_name="fake", + request=TaskSubmitRequest( + user_id="user-1", + thread_id="thread-1", + thread_incarnation="captured-incarnation", + run_id="run-1", + tool_call_id="call-1", + server_name="reports", + task_name="Generate report", + arguments={}, + local_task_id="task-1", + ), + ) + + assert repo.created[0]["expected_thread_incarnation"] == "captured-incarnation" + assert len(driver.cancel_calls) == 1 + assert driver.cancel_calls[0].thread_incarnation == "captured-incarnation" + + @pytest.mark.asyncio async def test_submit_cancellation_during_persistence_cancels_remote_task(): repo = BlockingCreateRepository() @@ -286,6 +339,7 @@ async def test_submit_cancellation_during_persistence_cancels_remote_task(): request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -329,6 +383,7 @@ async def test_submit_repeated_cancellation_does_not_interrupt_compensation(): request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -376,6 +431,7 @@ async def test_submit_stops_waiting_for_hung_compensation_without_cancelling_it( request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -426,6 +482,7 @@ async def test_submit_cancellation_preserves_cancelled_error_when_compensation_f request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -472,6 +529,7 @@ async def test_submit_cancels_remote_task_when_its_id_exceeds_storage_limit(): request=TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -512,6 +570,7 @@ async def test_duplicate_remote_handle_is_rejected_without_cancelling_existing_t request=TaskSubmitRequest( user_id="user-1", thread_id="thread-2", + thread_incarnation="incarnation-2", run_id="run-2", tool_call_id="call-2", server_name="reports", @@ -548,6 +607,7 @@ async def test_cancel_task_persists_request_without_calling_remote(): task_id="task-1", thread_id="thread-1", user_id="user-1", + thread_incarnation="incarnation-1", ) assert result == record @@ -961,6 +1021,7 @@ async def test_submit_preserves_persistence_error_when_compensation_cancel_fails request = TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", @@ -1166,6 +1227,7 @@ async def test_persisted_snapshot_errors_are_bounded_on_submit_and_poll(): request=TaskSubmitRequest( user_id="user-1", thread_id="thread-1", + thread_incarnation="incarnation-1", run_id="run-1", tool_call_id="call-1", server_name="reports", diff --git a/backend/tests/test_mcp_task_tool_caller.py b/backend/tests/test_mcp_task_tool_caller.py index e0f8dad64..f94373d75 100644 --- a/backend/tests/test_mcp_task_tool_caller.py +++ b/backend/tests/test_mcp_task_tool_caller.py @@ -15,7 +15,8 @@ from mcp.types import CONNECTION_CLOSED, ErrorData from deerflow.config.extensions_config import ExtensionsConfig from deerflow.config.paths import Paths from deerflow.mcp.session_pool import MCPSessionPool -from deerflow.mcp.task_tool_caller import McpTaskToolCaller, mcp_task_session_scope_key +from deerflow.mcp.task_tool_caller import McpTaskToolCaller +from deerflow.mcp_scope import mcp_session_scope_key def _config() -> ExtensionsConfig: @@ -78,8 +79,21 @@ async def _assert_configured_timeout(awaitable: Coroutine[Any, Any, Any]) -> Non await task -def test_task_session_scope_includes_user_and_thread() -> None: - assert mcp_task_session_scope_key(user_id="user-1", thread_id="thread-1") == "user-1:thread-1" +def test_task_session_scope_includes_user_thread_and_incarnation() -> None: + first = mcp_session_scope_key(user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-1") + second = mcp_session_scope_key(user_id="user-1", thread_id="thread-1", thread_incarnation="incarnation-2") + + assert first == 'v2:["user-1","thread-1","incarnation-1"]' + assert second == 'v2:["user-1","thread-1","incarnation-2"]' + assert first != second + assert mcp_session_scope_key(user_id="a:b", thread_id="c", thread_incarnation="d") != mcp_session_scope_key( + user_id="a", + thread_id="b:c", + thread_incarnation="d", + ) + assert mcp_session_scope_key(user_id="user-1", thread_id="thread-1", thread_incarnation=None) == "user-1:thread-1" + with pytest.raises(RuntimeError, match="non-empty"): + mcp_session_scope_key(user_id="user-1", thread_id="thread-1", thread_incarnation="") @pytest.mark.asyncio @@ -104,6 +118,7 @@ async def test_stdio_task_call_reuses_exact_scope_and_raw_tool_name() -> None: arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) assert actual is result @@ -148,6 +163,7 @@ async def test_broken_stdio_task_session_is_evicted_for_next_poll_reconnect(disc arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) pool.close_session_if_current.assert_awaited_once_with( @@ -182,6 +198,7 @@ async def test_stdio_task_timeout_keeps_healthy_stateful_session() -> None: arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) pool.close_session_if_current.assert_not_awaited() @@ -216,6 +233,7 @@ async def test_stdio_task_interceptor_failure_keeps_healthy_session() -> None: arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) session.call_tool.assert_not_awaited() @@ -277,6 +295,7 @@ mcp.run(transport="stdio") arguments={}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) server_config.tool_call_timeout = 0.05 @@ -287,6 +306,7 @@ mcp.run(transport="stdio") arguments={"task_id": submitted.structuredContent["task_id"]}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) assert exc_info.value.error.code == 408 @@ -298,6 +318,7 @@ mcp.run(transport="stdio") arguments={"task_id": submitted.structuredContent["task_id"]}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) finally: await pool.close_all() @@ -337,6 +358,7 @@ async def test_stdio_task_session_initialization_respects_configured_timeout() - arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) pool.close_session.assert_not_awaited() @@ -368,6 +390,7 @@ async def test_http_task_call_authenticates_session_initialization() -> None: arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) assert actual is result @@ -420,6 +443,7 @@ async def test_remote_task_session_initialization_respects_configured_timeout(tr arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) ) @@ -458,6 +482,7 @@ async def test_remote_task_call_respects_configured_timeout(transport: str) -> N arguments={"task_id": "remote-1"}, user_id="user-1", thread_id="thread-1", + thread_incarnation=None, ) ) diff --git a/backend/tests/test_mcp_task_tool_wrapping.py b/backend/tests/test_mcp_task_tool_wrapping.py index 88e10659a..6cf262d6d 100644 --- a/backend/tests/test_mcp_task_tool_wrapping.py +++ b/backend/tests/test_mcp_task_tool_wrapping.py @@ -140,7 +140,7 @@ async def test_submit_wrapper_persists_before_returning_only_the_local_handle() ) submit_tool = configured[0] runtime = SimpleNamespace( - context={"thread_id": "thread-1", "run_id": "run-1"}, + context={"thread_id": "thread-1", "thread_incarnation": "incarnation-1", "run_id": "run-1"}, config={}, tool_call_id="call-1", ) diff --git a/backend/tests/test_mcp_tasks_router.py b/backend/tests/test_mcp_tasks_router.py index a96c6cdff..c2302c451 100644 --- a/backend/tests/test_mcp_tasks_router.py +++ b/backend/tests/test_mcp_tasks_router.py @@ -14,13 +14,16 @@ class FakeRepository: self.list_calls = [] self.get_calls = [] - async def list_by_thread(self, thread_id, *, user_id, limit): - self.list_calls.append((thread_id, user_id, limit)) + async def list_by_thread(self, thread_id, *, user_id, thread_incarnation, limit): + self.list_calls.append((thread_id, user_id, thread_incarnation, limit)) return list(self.rows) - async def get(self, task_id, *, user_id): - self.get_calls.append((task_id, user_id)) - return next((row for row in self.rows if row["id"] == task_id and row["user_id"] == user_id), None) + async def get(self, task_id, *, user_id, thread_id, thread_incarnation): + self.get_calls.append((task_id, user_id, thread_id, thread_incarnation)) + return next( + (row for row in self.rows if row["id"] == task_id and row["user_id"] == user_id and row["thread_id"] == thread_id), + None, + ) def _record(**overrides): @@ -56,6 +59,15 @@ def _request(repo): state=SimpleNamespace( mcp_task_repo=repo, mcp_task_service=SimpleNamespace(tracking_degraded_after_errors=3), + thread_store=SimpleNamespace( + get=AsyncMock( + return_value={ + "thread_id": "thread-1", + "user_id": "user-1", + "incarnation": "incarnation-1", + } + ) + ), ) ) ) @@ -67,6 +79,93 @@ def test_gateway_mounts_thread_scoped_mcp_task_routes() -> None: assert "/api/threads/{thread_id}/mcp-tasks/{task_id}" in paths +@pytest.mark.asyncio +async def test_current_thread_incarnation_rejects_missing_thread() -> None: + request = _request(FakeRepository([])) + request.app.state.thread_store.get = AsyncMock(side_effect=[None, None]) + + with pytest.raises(HTTPException) as exc_info: + await mcp_tasks._current_thread_incarnation( + request, + thread_id="missing-thread", + user_id="user-1", + ) + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +async def test_current_thread_incarnation_accepts_shared_fallback() -> None: + request = _request(FakeRepository([])) + request.app.state.thread_store.get = AsyncMock( + side_effect=[ + None, + { + "thread_id": "shared-thread", + "user_id": None, + "incarnation": "shared-incarnation", + }, + ] + ) + + incarnation = await mcp_tasks._current_thread_incarnation( + request, + thread_id="shared-thread", + user_id="user-1", + ) + + assert incarnation == "shared-incarnation" + assert request.app.state.thread_store.get.await_args_list[1].kwargs == {"user_id": None} + + +@pytest.mark.asyncio +async def test_current_thread_incarnation_rejects_foreign_unscoped_fallback() -> None: + request = _request(FakeRepository([])) + request.app.state.thread_store.get = AsyncMock( + side_effect=[ + None, + { + "thread_id": "foreign-thread", + "user_id": "user-2", + "incarnation": "foreign-incarnation", + }, + ] + ) + + with pytest.raises(HTTPException) as exc_info: + await mcp_tasks._current_thread_incarnation( + request, + thread_id="foreign-thread", + user_id="user-1", + ) + + assert exc_info.value.status_code == 404 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("incarnation", ["", 7, False]) +async def test_current_thread_incarnation_rejects_malformed_value( + incarnation, +) -> None: + request = _request(FakeRepository([])) + request.app.state.thread_store.get = AsyncMock( + return_value={ + "thread_id": "thread-1", + "user_id": "user-1", + "incarnation": incarnation, + } + ) + + with pytest.raises(HTTPException) as exc_info: + await mcp_tasks._current_thread_incarnation( + request, + thread_id="thread-1", + user_id="user-1", + ) + + assert exc_info.value.status_code == 404 + + @pytest.mark.asyncio async def test_list_returns_only_safe_current_user_thread_fields(monkeypatch) -> None: repo = FakeRepository([_record()]) @@ -78,7 +177,7 @@ async def test_list_returns_only_safe_current_user_thread_fields(monkeypatch) -> limit=25, ) - assert repo.list_calls == [("thread-1", "user-1", 25)] + assert repo.list_calls == [("thread-1", "user-1", "incarnation-1", 25)] assert response == [ { "task_id": "mcp-task-1", @@ -174,6 +273,7 @@ async def test_cancel_uses_service_with_exact_user_and_thread_scope(monkeypatch) task_id="mcp-task-1", thread_id="thread-1", user_id="user-1", + thread_incarnation="incarnation-1", ) assert response["status"] == "working" assert response["cancel_requested"] is True diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index d04b8bab9..64976dfcd 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -184,8 +184,8 @@ async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled(): projection_started = asyncio.Event() class BlockingTaskRepository: - async def list_by_thread(self, thread_id, *, user_id, limit): - del thread_id, user_id, limit + async def list_by_thread(self, thread_id, *, user_id, thread_incarnation, limit): + del thread_id, user_id, thread_incarnation, limit projection_started.set() await asyncio.Event().wait() @@ -210,6 +210,7 @@ async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled(): agent_factory=agent_factory, graph_input={}, config={}, + thread_incarnation=None, ) ) await asyncio.wait_for(projection_started.wait(), timeout=1) @@ -2559,6 +2560,21 @@ def test_build_runtime_context_caller_cannot_override_thread_id_or_run_id(): assert ctx["agent_name"] == "ok" +def test_build_runtime_context_uses_server_owned_thread_incarnation(): + ctx = _build_runtime_context( + "thread-1", + "run-1", + { + "thread_incarnation": "spoofed", + "__deerflow_thread_incarnation_metadata_guard": True, + }, + thread_incarnation="server-incarnation", + ) + + assert ctx["thread_incarnation"] == "server-incarnation" + assert "__deerflow_thread_incarnation_metadata_guard" not in ctx + + def test_build_runtime_context_ignores_caller_pre_existing_message_ids(): caller_context = {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}} diff --git a/backend/tests/test_subagent_batch_repository.py b/backend/tests/test_subagent_batch_repository.py index ebc25301b..32d07bc60 100644 --- a/backend/tests/test_subagent_batch_repository.py +++ b/backend/tests/test_subagent_batch_repository.py @@ -4,9 +4,12 @@ import pytest import pytest_asyncio from deerflow.config.database_config import DatabaseConfig +from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config from deerflow.persistence.subagent_batches import SubagentBatchRepository +_MISSING = object() + @pytest_asyncio.fixture(autouse=True) async def _close_engine() -> None: @@ -86,6 +89,59 @@ async def test_claim_separates_total_live_leased_and_running(tmp_path) -> None: assert while_full == [] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("thread_incarnation", "expected_present"), + [ + ("incarnation-1", True), + (None, True), + (_MISSING, False), + ], +) +async def test_claim_preserves_thread_incarnation_presence( + tmp_path, + thread_incarnation, + expected_present, +) -> None: + repo = await _repo(tmp_path) + execution_spec = { + "subagent_config": { + "name": "general-purpose", + "description": "test", + }, + } + if thread_incarnation is not _MISSING: + execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation + await repo.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="Incarnation persistence", + subagent_type="general-purpose", + items=[{"key": "item-1", "prompt": "Process item"}], + max_live_items=1, + max_running_items=1, + max_attempts=1, + execution_spec=execution_spec, + ) + + claimed = await repo.claim_items( + now=datetime.now(UTC), + lease_owner="worker-1", + lease_seconds=60, + limit=1, + ) + + assert len(claimed) == 1 + restored_spec = claimed[0]["batch"]["execution_spec"] + assert (THREAD_INCARNATION_CONTEXT_KEY in restored_spec) is expected_present + if expected_present: + assert restored_spec[THREAD_INCARNATION_CONTEXT_KEY] == thread_incarnation + + @pytest.mark.asyncio async def test_expired_lease_is_recovered_with_stable_item_identity(tmp_path) -> None: repo = await _repo(tmp_path) diff --git a/backend/tests/test_subagent_batch_service.py b/backend/tests/test_subagent_batch_service.py index ebbd12b5e..1934160fb 100644 --- a/backend/tests/test_subagent_batch_service.py +++ b/backend/tests/test_subagent_batch_service.py @@ -7,11 +7,14 @@ import pytest from deerflow.config.subagent_batches_config import SubagentBatchesConfig from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig +from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY from deerflow.subagents import batch_service as service_module from deerflow.subagents.batch_runtime import BatchSubmitRequest from deerflow.subagents.batch_service import SubagentBatchService from deerflow.subagents.capacity import SubagentExecutionCapacity +_MISSING = object() + class FakeStatus(Enum): PENDING = "pending" @@ -116,7 +119,19 @@ async def test_submit_defaults_only_the_limits_the_caller_omitted(overrides: dic @pytest.mark.asyncio -async def test_execute_item_marks_real_running_then_persists_terminal_result(monkeypatch) -> None: +@pytest.mark.parametrize( + ("thread_incarnation", "expected_present"), + [ + ("incarnation-1", True), + (None, True), + (_MISSING, False), + ], +) +async def test_execute_item_marks_real_running_then_persists_terminal_result( + monkeypatch, + thread_incarnation, + expected_present, +) -> None: result = SimpleNamespace( status=FakeStatus.RUNNING, result=None, @@ -125,6 +140,10 @@ async def test_execute_item_marks_real_running_then_persists_terminal_result(mon token_usage_records=None, ) + execution_spec = dict(_request().execution_spec) + if thread_incarnation is not _MISSING: + execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation + class Repository: def __init__(self) -> None: self.marked_running = False @@ -141,7 +160,7 @@ async def test_execute_item_marks_real_running_then_persists_terminal_result(mon "thread_id": "thread-1", "user_id": "user-1", "run_id": "run-1", - "execution_spec": _request().execution_spec, + "execution_spec": execution_spec, }, } ] @@ -198,6 +217,9 @@ async def test_execute_item_marks_real_running_then_persists_terminal_result(mon "mode": "selected", "dataset_ids": ["dataset-1"], } + assert (THREAD_INCARNATION_CONTEXT_KEY in executor_kwargs) is expected_present + if expected_present: + assert executor_kwargs[THREAD_INCARNATION_CONTEXT_KEY] is thread_incarnation @pytest.mark.asyncio diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index 3b983dfcb..1770d72bb 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -4177,6 +4177,71 @@ class _FakeStreamAgent: yield # pragma: no cover - make this an async generator +@pytest.mark.anyio +@pytest.mark.parametrize( + "scope_kwargs, expected_scope", + [ + ({"thread_incarnation": "captured-incarnation"}, 'v2:["alice","parent-thread","captured-incarnation"]'), + ({"thread_incarnation": None}, "alice:parent-thread"), + ({}, None), + ({"thread_incarnation": ""}, None), + ({"thread_incarnation": False}, None), + ({"thread_incarnation": {}}, None), + ], +) +async def test_subagent_mcp_uses_captured_thread_incarnation(classes, monkeypatch, tmp_path, scope_kwargs, expected_scope): + """A real child ToolNode must receive the parent's captured MCP scope.""" + from langchain_core.tools import StructuredTool + from langgraph.graph import END, START, MessagesState, StateGraph + from langgraph.prebuilt import ToolNode + from mcp.types import CallToolResult + + from deerflow.mcp import tools as mcp_tools + + executor_module = importlib.import_module("deerflow.subagents.executor") + monkeypatch.setattr(executor_module, "build_tracing_callbacks", lambda: []) + pool = SimpleNamespace(get_session=AsyncMock(return_value=object())) + call_remote = AsyncMock(return_value=CallToolResult(content=[], isError=False)) + monkeypatch.setattr(mcp_tools, "get_session_pool", lambda: pool) + monkeypatch.setattr(mcp_tools, "call_pooled_session_tool", call_remote) + monkeypatch.setattr(mcp_tools, "_prepare_stdio_workspace", lambda *args, **kwargs: (tmp_path, tmp_path, {})) + tool = mcp_tools._make_session_pool_tool( + StructuredTool(name="server_probe", description="Probe MCP", args_schema={"type": "object", "properties": {}}, coroutine=AsyncMock()), + "server", + {"transport": "stdio", "command": "unused"}, + ) + graph = StateGraph(MessagesState, context_schema=dict) + graph.add_node("tools", ToolNode([tool], handle_tool_errors=False)) + graph.add_edge(START, "tools") + graph.add_edge("tools", END) + child = graph.compile() + executor = classes["SubagentExecutor"]( + config=classes["SubagentConfig"](name="general-purpose", description="MCP scope test", system_prompt="Test", max_turns=5, timeout_seconds=30), + tools=[tool], + parent_model="test-model", + thread_id="parent-thread", + user_id="alice", + **scope_kwargs, + ) + + async def build_initial_state(task): + return ({"messages": [classes["AIMessage"](content="", tool_calls=[{"name": tool.name, "args": {}, "id": "probe"}])]}, [], None) + + monkeypatch.setattr(executor, "_build_initial_state", build_initial_state) + monkeypatch.setattr(executor, "_create_agent", lambda *args, **kwargs: child) + result = await executor._aexecute("probe MCP") + + if expected_scope is None: + assert result.status == classes["SubagentStatus"].FAILED, result.ai_messages + assert "thread incarnation" in result.error + pool.get_session.assert_not_awaited() + call_remote.assert_not_awaited() + else: + assert result.status == classes["SubagentStatus"].COMPLETED, result.error + call_remote.assert_awaited_once() + assert call_remote.await_args.kwargs["scope_key"] == expected_scope + + class TestSubagentCheckpointLineage: """Keep delegated graphs on the parent run's checkpoint lineage.""" diff --git a/backend/tests/test_task_tool_core_logic.py b/backend/tests/test_task_tool_core_logic.py index 681d0e494..542974783 100644 --- a/backend/tests/test_task_tool_core_logic.py +++ b/backend/tests/test_task_tool_core_logic.py @@ -3265,6 +3265,34 @@ def _capture_executor_call(monkeypatch, **call_kwargs): return captured["executor_kwargs"], captured["prompt"] +@pytest.mark.parametrize("incarnation", ["captured-incarnation", None, "", False, {}]) +def test_task_tool_forwards_captured_thread_incarnation(monkeypatch, incarnation): + runtime = _make_runtime() + runtime.context["thread_incarnation"] = incarnation + executor_kwargs, _ = _capture_executor_call(monkeypatch, runtime=runtime) + assert executor_kwargs["thread_incarnation"] is incarnation + + +def test_task_tool_does_not_invent_missing_thread_incarnation(monkeypatch): + runtime = _make_runtime() + runtime.context.pop("thread_incarnation", None) + runtime.state["thread_incarnation"] = "untrusted-state" + runtime.config.setdefault("configurable", {})["thread_incarnation"] = "untrusted-config" + executor_kwargs, _ = _capture_executor_call(monkeypatch, runtime=runtime) + assert "thread_incarnation" not in executor_kwargs + assert "thread_incarnation" not in task_tool_module.task_tool.tool_call_schema.model_fields + + +def test_task_tool_rejects_stale_standalone_thread_incarnation(monkeypatch): + runtime = _make_runtime() + runtime.context["thread_incarnation"] = "incarnation-1" + runtime.context["__deerflow_thread_incarnation_metadata_guard"] = True + runtime.config["metadata"]["thread_incarnation"] = "incarnation-2" + + with pytest.raises(RuntimeError, match="stale thread incarnation"): + _capture_executor_call(monkeypatch, runtime=runtime) + + def test_task_tool_forwards_acceptance_criteria_to_executor(monkeypatch): """RFC #4651 PR3: criteria travel via the executor constructor; the executor appends them to the subagent's task HumanMessage as untrusted diff --git a/config.example.yaml b/config.example.yaml index 2a1522ab4..b7f06ab58 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -1503,7 +1503,7 @@ uploads: # ============================================================================ # Sandbox Configuration # ============================================================================ -# Choose between local sandbox (direct execution), Docker-based AIO sandbox, +# Choose between local sandbox (direct execution), Docker-based AIO sandbox, # or BoxLite micro-VM sandbox # Option 1: Local Sandbox (Default) # Executes commands directly on the host machine