fix(mcp): scope sessions and task access by thread incarnation (#5556)

* fix(mcp): scope sessions and task access by thread incarnation

* fix(mcp): preserve thread incarnation in delegated subagents

* fix(mcp): preserve incarnation in durable batches

* fix(mcp): bind standalone graph lifecycle context

* fix(studio): preserve implicit thread creation metadata

---------

Co-authored-by: CorgiBoyG <CorgiBoyG@users.noreply.github.com>
Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
RongJie G 2026-09-23 08:33:20 +08:00 committed by GitHub
parent d3f44fd159
commit e01314442c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
48 changed files with 2158 additions and 190 deletions

View File

@ -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 The documented command uses LangGraph's file-based custom-app loader, which is
also covered directly by DeerFlow's regression tests. 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 For workflows that invoke `backend/langgraph.json` through LangGraph Studio or
a direct LangGraph Server, DeerFlow consumes the authenticated identity a direct LangGraph Server, DeerFlow consumes the authenticated identity
published by that runtime and uses it for custom-agent configuration/SOUL, user 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. 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 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. 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 `<server_name>_` 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. MCP tool names are prefixed with `<server_name>_` 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. 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.

View File

@ -1,5 +1,7 @@
### Gateway API (`app/gateway/`) ### Gateway API (`app/gateway/`)
Studio retains sanitized creation metadata.
Capability Center's `business` adapter validates only the bundled provider's Capability Center's `business` adapter validates only the bundled provider's
credential fields and creates a normal MCP connection. The MCP API accepts the credential fields and creates a normal MCP connection. The MCP API accepts the
exact isolated interpreter/module/provider launcher generated by exact isolated interpreter/module/provider launcher generated by

View File

@ -15,13 +15,21 @@ Two layers:
""" """
import secrets import secrets
from contextvars import ContextVar
from uuid import uuid4
from langgraph_sdk import Auth from langgraph_sdk import Auth
from starlette.exceptions import HTTPException
from app.gateway.auth.errors import TokenError from app.gateway.auth.errors import TokenError
from app.gateway.auth.jwt import decode_token from app.gateway.auth.jwt import decode_token
from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID, is_auth_disabled from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID, is_auth_disabled
from app.gateway.deps import get_local_provider 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() auth = Auth()
@ -32,6 +40,160 @@ _STUDIO_USER_TYPE = getattr(Auth.types, "StudioUser", None)
# Methods that require CSRF validation (state-changing per RFC 7231). # Methods that require CSRF validation (state-changing per RFC 7231).
_CSRF_METHODS = frozenset({"POST", "PUT", "DELETE", "PATCH"}) _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: 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 # boundary. The standalone pre-runtime persistence repair also scrubs this
# marker from legacy active rows and their version history before normal # marker from legacy active rows and their version history before normal
# version selection becomes available. # version selection becomes available.
metadata = value.setdefault("metadata", {}) metadata = _metadata(value)
metadata["user_id"] = ctx.user.identity 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"}: if ctx.resource == "assistants" and ctx.action in {"create", "update"}:
metadata["created_by"] = "user" metadata["created_by"] = "user"

View File

@ -7,7 +7,8 @@ from typing import Any
from fastapi import APIRouter, HTTPException, Query, Request from fastapi import APIRouter, HTTPException, Query, Request
from app.gateway.authz import require_permission 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 from deerflow.utils.thread_id import ThreadId
router = APIRouter(prefix="/api/threads/{thread_id}/mcp-tasks", tags=["mcp-tasks"]) 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 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("") @router.get("")
@require_permission("threads", "read", owner_check=True) @require_permission("threads", "read", owner_check=True)
async def list_mcp_tasks( async def list_mcp_tasks(
@ -73,9 +95,15 @@ async def list_mcp_tasks(
repository = get_mcp_task_repo(request) repository = get_mcp_task_repo(request)
service = get_mcp_task_service(request) service = get_mcp_task_service(request)
user_id = await _current_user_id(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( records = await repository.list_by_thread(
thread_id, thread_id,
user_id=user_id, user_id=user_id,
thread_incarnation=thread_incarnation,
limit=limit, limit=limit,
) )
threshold = service.tracking_degraded_after_errors threshold = service.tracking_degraded_after_errors
@ -92,8 +120,18 @@ async def get_mcp_task(
repository = get_mcp_task_repo(request) repository = get_mcp_task_repo(request)
service = get_mcp_task_service(request) service = get_mcp_task_service(request)
user_id = await _current_user_id(request) user_id = await _current_user_id(request)
record = await repository.get(task_id, user_id=user_id) thread_incarnation = await _current_thread_incarnation(
if record is None or record["thread_id"] != thread_id: 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") raise HTTPException(status_code=404, detail="MCP task not found")
return _detail( return _detail(
record, record,
@ -110,6 +148,11 @@ async def cancel_mcp_task(
) -> dict[str, Any]: ) -> dict[str, Any]:
service = get_mcp_task_service(request) service = get_mcp_task_service(request)
user_id = await _current_user_id(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): if not getattr(request.app.state, "mcp_tasks_available", False):
# The service exists whenever SQL persistence is configured, but the # The service exists whenever SQL persistence is configured, but the
# background loop that owns the remote cancel call only runs when # background loop that owns the remote cancel call only runs when
@ -120,6 +163,7 @@ async def cancel_mcp_task(
task_id=task_id, task_id=task_id,
thread_id=thread_id, thread_id=thread_id,
user_id=user_id, user_id=user_id,
thread_incarnation=thread_incarnation,
) )
if record is None: if record is None:
raise HTTPException(status_code=404, detail="MCP task not found") raise HTTPException(status_code=404, detail="MCP task not found")

View File

@ -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.app_config import get_app_config
from deerflow.config.database_config import resolve_checkpoint_graph_cache_max 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.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.projects.context import PROJECT_CONTEXT_MESSAGE_MARKER, resolve_project_context
from deerflow.runtime import ( from deerflow.runtime import (
END_SENTINEL, END_SENTINEL,
@ -206,7 +210,7 @@ async def _ensure_thread_metadata(
*, *,
owner_user_id: str | None, owner_user_id: str | None,
require_existing_thread: bool = False, require_existing_thread: bool = False,
) -> None: ) -> dict[str, Any]:
"""Ensure an admitted run's thread exists without delaying task attachment.""" """Ensure an admitted run's thread exists without delaying task attachment."""
thread_store = run_ctx.thread_store thread_store = run_ctx.thread_store
existing = await thread_store.get(record.thread_id) 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. # /threads/{id}/move — so the key must not persist either.
if key not in (DEERFLOW_TRACE_METADATA_KEY, THREAD_PROJECT_METADATA_KEY) if key not in (DEERFLOW_TRACE_METADATA_KEY, THREAD_PROJECT_METADATA_KEY)
} }
await thread_store.create( existing = await thread_store.create(
record.thread_id, record.thread_id,
assistant_id=record.assistant_id, assistant_id=record.assistant_id,
metadata=metadata, metadata=metadata,
) )
return return existing
async def _terminal_record_stream_missing(bridge: StreamBridge, record: RunRecord) -> bool: 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", "__run_tool_progress_recorder",
"langgraph_auth_user", "langgraph_auth_user",
"langgraph_auth_user_id", "langgraph_auth_user_id",
THREAD_INCARNATION_METADATA_GUARD_KEY,
# Server-owned pinned project snapshot (spec §7.1): resolved once # Server-owned pinned project snapshot (spec §7.1): resolved once
# at admission from threads_meta; a client-supplied value must # at admission from threads_meta; a client-supplied value must
# never survive in either run-config section. # never survive in either run-config section.
@ -1916,6 +1921,7 @@ async def start_run(
abort_task = asyncio.create_task(record.abort_event.wait()) abort_task = asyncio.create_task(record.abort_event.wait())
metadata_failure_logged = False metadata_failure_logged = False
metadata_failure: Exception | None = None metadata_failure: Exception | None = None
metadata_record: dict[str, Any] | None = None
try: try:
done, _ = await asyncio.wait( done, _ = await asyncio.wait(
(metadata_task, abort_task), (metadata_task, abort_task),
@ -1924,7 +1930,7 @@ async def start_run(
) )
if metadata_task in done: if metadata_task in done:
try: try:
metadata_task.result() metadata_record = metadata_task.result()
except asyncio.CancelledError: except asyncio.CancelledError:
pass pass
except Exception as exc: except Exception as exc:
@ -1946,8 +1952,20 @@ async def start_run(
metadata_failure = TimeoutError("Timed out verifying existing thread metadata") metadata_failure = TimeoutError("Timed out verifying existing thread metadata")
finally: finally:
if metadata_task.done(): if metadata_task.done():
if not metadata_failure_logged: if metadata_record is None and not metadata_failure_logged:
_log_thread_metadata_task_result(metadata_task, thread_id=thread_id) 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: else:
metadata_task.cancel() metadata_task.cancel()
metadata_task.add_done_callback( metadata_task.add_done_callback(
@ -1968,6 +1986,28 @@ async def start_run(
# or strict verification failure: # or strict verification failure:
# its startup barrier is the single path that turns pending # its startup barrier is the single path that turns pending
# cancellation into no-agent-construction plus publish_end. # 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( await run_agent(
bridge, bridge,
run_mgr, run_mgr,
@ -1981,6 +2021,7 @@ async def start_run(
interrupt_before=body.interrupt_before, interrupt_before=body.interrupt_before,
interrupt_after=body.interrupt_after, interrupt_after=body.interrupt_after,
knowledge_scope=admitted_knowledge_scope, knowledge_scope=admitted_knowledge_scope,
**incarnation_kwargs,
) )
try: try:

View File

@ -310,6 +310,7 @@ class McpTaskService:
thread_id=request.thread_id, thread_id=request.thread_id,
server_name=request.server_name, server_name=request.server_name,
remote_task_id=submission.remote_task_id, remote_task_id=submission.remote_task_id,
thread_incarnation=request.thread_incarnation,
driver_data=driver_data, driver_data=driver_data,
) )
try: try:
@ -321,6 +322,7 @@ class McpTaskService:
task_id=local_task_id, task_id=local_task_id,
user_id=request.user_id, user_id=request.user_id,
thread_id=request.thread_id, thread_id=request.thread_id,
expected_thread_incarnation=request.thread_incarnation,
run_id=request.run_id, run_id=request.run_id,
tool_call_id=request.tool_call_id, tool_call_id=request.tool_call_id,
server_name=request.server_name, server_name=request.server_name,
@ -692,12 +694,14 @@ class McpTaskService:
*, *,
thread_id: str, thread_id: str,
user_id: str, user_id: str,
thread_incarnation: str | None,
limit: int = 50, limit: int = 50,
active_only: bool = False, active_only: bool = False,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
return await self._repository.list_by_thread( return await self._repository.list_by_thread(
thread_id, thread_id,
user_id=user_id, user_id=user_id,
thread_incarnation=thread_incarnation,
limit=limit, limit=limit,
active_only=active_only, active_only=active_only,
) )
@ -708,12 +712,14 @@ class McpTaskService:
task_id: str, task_id: str,
thread_id: str, thread_id: str,
user_id: str, user_id: str,
thread_incarnation: str | None,
now: datetime | None = None, now: datetime | None = None,
) -> dict[str, Any] | None: ) -> dict[str, Any] | None:
return await self._repository.request_cancel( return await self._repository.request_cancel(
task_id, task_id,
user_id=user_id, user_id=user_id,
thread_id=thread_id, thread_id=thread_id,
thread_incarnation=thread_incarnation,
requested_at=now or datetime.now(UTC), requested_at=now or datetime.now(UTC),
) )
@ -722,9 +728,15 @@ class McpTaskService:
*, *,
thread_id: str, thread_id: str,
user_id: str, user_id: str,
thread_incarnation: str | None,
task: str | None = None, task: str | None = None,
) -> dict[str, Any]: ) -> 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: if task:
normalized = task.casefold().strip() normalized = task.casefold().strip()
matches = [item for item in active if item["id"] == task or str(item.get("task_name") or "").casefold() == normalized] 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"], task_id=matches[0]["id"],
thread_id=thread_id, thread_id=thread_id,
user_id=user_id, user_id=user_id,
thread_incarnation=thread_incarnation,
) )
if result is None: if result is None:
raise LookupError("The selected background task no longer exists") raise LookupError("The selected background task no longer exists")

View File

@ -1117,6 +1117,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
middleware=normalize_middleware_state_schemas(middlewares, mode), middleware=normalize_middleware_state_schemas(middlewares, mode),
system_prompt=system_prompt, system_prompt=system_prompt,
state_schema=get_thread_state_schema(mode), state_schema=get_thread_state_schema(mode),
context_schema=dict,
) )
return _complete_assembly( return _complete_assembly(
config=config, config=config,
@ -1249,6 +1250,7 @@ def _assemble_lead_agent(config: RunnableConfig, *, app_config: AppConfig) -> Le
middleware=normalize_middleware_state_schemas(middlewares, mode), middleware=normalize_middleware_state_schemas(middlewares, mode),
system_prompt=system_prompt, system_prompt=system_prompt,
state_schema=get_thread_state_schema(mode), state_schema=get_thread_state_schema(mode),
context_schema=dict,
) )
return _complete_assembly( return _complete_assembly(
config=config, config=config,

View File

@ -52,6 +52,7 @@ from deerflow.config.extensions_config import (
) )
from deerflow.config.paths import get_paths from deerflow.config.paths import get_paths
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig 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.models import create_chat_model
from deerflow.runtime import CheckpointStateAccessor from deerflow.runtime import CheckpointStateAccessor
from deerflow.runtime.checkpoint_mode import ( 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._available_skills = set(available_skills) if available_skills is not None else None
self._middlewares = list(middlewares) if middlewares else [] self._middlewares = list(middlewares) if middlewares else []
self._environment = environment self._environment = environment
self._thread_incarnations: dict[str, str] = {}
# Lazy agent — created on first call, recreated when config changes. # Lazy agent — created on first call, recreated when config changes.
self._agent = None self._agent = None
@ -948,7 +950,15 @@ class DeerFlowClient:
config["callbacks"] = [*existing_callbacks, *tracing_callbacks] config["callbacks"] = [*existing_callbacks, *tracing_callbacks]
run_id = str(uuid.uuid4()) 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: for key in _EMBEDDED_AUTHORIZATION_CONTEXT_KEYS:
if key in kwargs: if key in kwargs:
context[key] = kwargs[key] context[key] = kwargs[key]

View File

@ -3,7 +3,7 @@
- Uses `langchain-mcp-adapters` `MultiServerMCPClient` for multi-server management - 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. - **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. - **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.<server>.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.<server>.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. - **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()` - **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()`
- **Loop-isolated stdio sessions**: Both the live registry and in-flight creations - **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 `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 references `tool_search` when a hinted MCP tool is currently deferred; do not
add a parallel routing middleware for PR1-style preference hints. 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. - **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. - **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. - **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: - **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:

View File

@ -17,15 +17,11 @@ from deerflow.mcp.headers import apply_header_overrides
from deerflow.mcp.interceptors import build_mcp_tool_interceptors from deerflow.mcp.interceptors import build_mcp_tool_interceptors
from deerflow.mcp.oauth import OAuthTokenManager, build_oauth_tool_interceptor 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.session_pool import MCPSessionPool, call_pooled_session_tool, get_session_pool
from deerflow.mcp_scope import mcp_session_scope_key
logger = logging.getLogger(__name__) 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( def _prepare_stdio_connection(
connection: dict[str, Any], connection: dict[str, Any],
*, *,
@ -92,6 +88,7 @@ class McpTaskToolCaller:
arguments: dict[str, Any], arguments: dict[str, Any],
user_id: str, user_id: str,
thread_id: str, thread_id: str,
thread_incarnation: str | None = None,
request_scoped_headers: bool = False, request_scoped_headers: bool = False,
) -> Any: ) -> Any:
"""Call a raw MCP tool. """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") 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) connection = build_server_params(server_name, server_config)
transport = connection.get("transport", "stdio") 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": if transport == "stdio":
connection = await asyncio.to_thread( connection = await asyncio.to_thread(

View File

@ -93,6 +93,7 @@ class TaskReference:
server_name: str server_name: str
remote_task_id: str remote_task_id: str
driver_data: dict[str, Any] = field(default_factory=dict) driver_data: dict[str, Any] = field(default_factory=dict)
thread_incarnation: str | None = None
@classmethod @classmethod
def from_record(cls, record: dict[str, Any]) -> TaskReference: def from_record(cls, record: dict[str, Any]) -> TaskReference:
@ -102,6 +103,9 @@ class TaskReference:
thread_id=record["thread_id"], thread_id=record["thread_id"],
server_name=record["server_name"], server_name=record["server_name"],
remote_task_id=record["remote_task_id"], 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 {}), driver_data=dict(record.get("driver_data") or {}),
) )
@ -119,6 +123,7 @@ class TaskSubmitRequest:
arguments: dict[str, Any] arguments: dict[str, Any]
driver_data: dict[str, Any] = field(default_factory=dict) driver_data: dict[str, Any] = field(default_factory=dict)
local_task_id: str | None = None local_task_id: str | None = None
thread_incarnation: str | None = None
def __post_init__(self) -> None: def __post_init__(self) -> None:
_validate_storage_text( _validate_storage_text(

View File

@ -32,6 +32,7 @@ class McpTaskToolCaller(Protocol):
arguments: dict[str, Any], arguments: dict[str, Any],
user_id: str, user_id: str,
thread_id: str, thread_id: str,
thread_incarnation: str | None = None,
request_scoped_headers: bool = False, request_scoped_headers: bool = False,
) -> Any: ... ) -> Any: ...
@ -159,6 +160,7 @@ class OrdinaryMcpTaskDriver:
arguments=request.arguments, arguments=request.arguments,
user_id=request.user_id, user_id=request.user_id,
thread_id=request.thread_id, thread_id=request.thread_id,
thread_incarnation=request.thread_incarnation,
# Submit alone is awaited inside the Agent run, so it is the one # Submit alone is awaited inside the Agent run, so it is the one
# durable-task call that can carry the run's request-scoped # durable-task call that can carry the run's request-scoped
# credentials; status and cancel run after that run ended. # credentials; status and cancel run after that run ended.
@ -184,6 +186,7 @@ class OrdinaryMcpTaskDriver:
arguments={"task_id": task.remote_task_id}, arguments={"task_id": task.remote_task_id},
user_id=task.user_id, user_id=task.user_id,
thread_id=task.thread_id, thread_id=task.thread_id,
thread_incarnation=task.thread_incarnation,
) )
payload = _parse( payload = _parse(
_StatusPayload, _StatusPayload,
@ -202,6 +205,7 @@ class OrdinaryMcpTaskDriver:
arguments={"task_id": task.remote_task_id}, arguments={"task_id": task.remote_task_id},
user_id=task.user_id, user_id=task.user_id,
thread_id=task.thread_id, thread_id=task.thread_id,
thread_incarnation=task.thread_incarnation,
) )
payload = _parse( payload = _parse(
_CancelPayload, _CancelPayload,

View File

@ -26,6 +26,7 @@ class McpTaskSubmitter(Protocol):
*, *,
thread_id: str, thread_id: str,
user_id: str, user_id: str,
thread_incarnation: str | None,
limit: int = 50, limit: int = 50,
active_only: bool = False, active_only: bool = False,
) -> list[dict[str, Any]]: ... ) -> list[dict[str, Any]]: ...
@ -35,6 +36,7 @@ class McpTaskSubmitter(Protocol):
*, *,
thread_id: str, thread_id: str,
user_id: str, user_id: str,
thread_incarnation: str | None,
task: str | None = None, task: str | None = None,
) -> dict[str, Any]: ... ) -> dict[str, Any]: ...

View File

@ -29,6 +29,7 @@ from deerflow.mcp.tasks.runtime import (
get_mcp_task_submitter, get_mcp_task_submitter,
validate_mcp_task_config_snapshot, 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.reflection import resolve_variable
from deerflow.runtime.user_context import resolve_runtime_user_id from deerflow.runtime.user_context import resolve_runtime_user_id
from deerflow.tools.mcp_metadata import tag_mcp_routing, tag_mcp_tool 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. """Wrap an MCP tool so it reuses a persistent session from the pool.
Replaces the per-call session creation with pool-managed sessions scoped 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 (e.g. Playwright) keep their state across tool calls within the same thread
while staying isolated per user. 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 # 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 # 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. # 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) session_connection = dict(connection)
# cwd/temp pinning and the workspace snapshot only matter for stdio # cwd/temp pinning and the workspace snapshot only matter for stdio
# servers, which run as local subprocesses writing to a real filesystem. # 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() submitter = get_mcp_task_submitter()
thread_id = _extract_thread_id(runtime) thread_id = _extract_thread_id(runtime)
user_id = resolve_runtime_user_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 {} context = runtime.context if runtime is not None and runtime.context else {}
run_id = context.get("run_id") run_id = context.get("run_id")
tool_call_id = getattr(runtime, "tool_call_id", None) if runtime is not None else None 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( request=TaskSubmitRequest(
user_id=user_id, user_id=user_id,
thread_id=thread_id, thread_id=thread_id,
thread_incarnation=thread_incarnation,
run_id=str(run_id) if run_id is not None else None, 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, tool_call_id=str(tool_call_id) if tool_call_id is not None else None,
server_name=server_name, server_name=server_name,

View File

@ -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

View File

@ -1,4 +1,13 @@
from deerflow.persistence.mcp_tasks.model import McpTaskRow 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",
]

View File

@ -6,7 +6,7 @@ import uuid
from datetime import UTC, datetime, timedelta from datetime import UTC, datetime, timedelta
from typing import Any 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.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker 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.""" """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: def _is_remote_task_unique_conflict(exc: IntegrityError) -> bool:
original = exc.orig original = exc.orig
diagnostic = getattr(original, "diag", None) 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 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: class McpTaskRepository:
"""Durable source of truth for long-running MCP task lifecycle state.""" """Durable source of truth for long-running MCP task lifecycle state."""
@ -103,9 +165,11 @@ class McpTaskRepository:
self._sf = session_factory self._sf = session_factory
@staticmethod @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 = 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: for key in _TIMESTAMP_FIELDS:
if data.get(key) is not None: if data.get(key) is not None:
data[key] = coerce_iso(data[key]) data[key] = coerce_iso(data[key])
@ -117,6 +181,7 @@ class McpTaskRepository:
task_id: str, task_id: str,
user_id: str, user_id: str,
thread_id: str, thread_id: str,
expected_thread_incarnation: str | None,
run_id: str | None, run_id: str | None,
tool_call_id: str | None, tool_call_id: str | None,
server_name: str, server_name: str,
@ -160,19 +225,14 @@ class McpTaskRepository:
) )
_record_event_if_changed(row, tracking_degraded=False, now=now) _record_event_if_changed(row, tracking_degraded=False, now=now)
async with self._sf() as session: async with self._sf() as session:
matching_thread = select(ThreadMetaRow.incarnation).where( if not await _lock_current_thread_incarnation(
ThreadMetaRow.thread_id == thread_id, session,
or_(ThreadMetaRow.user_id == user_id, ThreadMetaRow.user_id.is_(None)), user_id=user_id,
) thread_id=thread_id,
if session.get_bind().dialect.name == "sqlite": thread_incarnation=expected_thread_incarnation,
# Keep lookup and write in one SQLite statement. A preliminary ):
# read would leave a delete/recreate window before the INSERT. raise McpTaskThreadMismatchError("MCP task submission crossed a thread lifecycle boundary")
row.thread_incarnation = matching_thread.scalar_subquery() row.thread_incarnation = expected_thread_incarnation
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()
session.add(row) session.add(row)
try: try:
await session.commit() await session.commit()
@ -184,24 +244,47 @@ class McpTaskRepository:
await session.refresh(row) await session.refresh(row)
return self._row_to_dict(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: async with self._sf() as session:
row = await session.get(McpTaskRow, task_id) stmt = select(McpTaskRow).where(
if row is None or row.user_id != user_id: McpTaskRow.id == task_id,
return None McpTaskRow.user_id == user_id,
return self._row_to_dict(row) 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( async def list_by_thread(
self, self,
thread_id: str, thread_id: str,
*, *,
user_id: str, user_id: str,
thread_incarnation: str | None,
limit: int = 50, limit: int = 50,
active_only: bool = False, active_only: bool = False,
) -> list[dict[str, Any]]: ) -> list[dict[str, Any]]:
stmt = select(McpTaskRow).where( stmt = select(McpTaskRow).where(
McpTaskRow.thread_id == thread_id, McpTaskRow.thread_id == thread_id,
McpTaskRow.user_id == user_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: if active_only:
stmt = stmt.where(McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES)) stmt = stmt.where(McpTaskRow.status.in_(_POLLABLE_STATUS_VALUES))
@ -245,7 +328,7 @@ class McpTaskRepository:
row.poll_attempt_count += 1 row.poll_attempt_count += 1
row.updated_at = now row.updated_at = now
await session.commit() 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( async def apply_snapshot(
self, self,
@ -387,16 +470,31 @@ class McpTaskRepository:
*, *,
user_id: str, user_id: str,
thread_id: str, thread_id: str,
thread_incarnation: str | None,
requested_at: datetime, requested_at: datetime,
) -> dict[str, Any] | None: ) -> dict[str, Any] | None:
"""Persist a user-scoped cancellation request without exposing the remote id.""" """Persist a user-scoped cancellation request without exposing the remote id."""
async with self._sf() as session: 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 = ( stmt = (
select(McpTaskRow) select(McpTaskRow)
.where( .where(
McpTaskRow.id == task_id, McpTaskRow.id == task_id,
McpTaskRow.user_id == user_id, McpTaskRow.user_id == user_id,
McpTaskRow.thread_id == thread_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() .with_for_update()
) )
@ -447,7 +545,7 @@ class McpTaskRepository:
row.cancel_attempt_count = int(row.cancel_attempt_count or 0) + 1 row.cancel_attempt_count = int(row.cancel_attempt_count or 0) + 1
row.updated_at = now row.updated_at = now
await session.commit() 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( async def apply_cancel_snapshot(
self, self,

View File

@ -42,6 +42,10 @@ from deerflow.config.app_config import AppConfig
from deerflow.config.database_config import CheckpointChannelMode from deerflow.config.database_config import CheckpointChannelMode
from deerflow.constants import CONVERSATION_READER_CONTEXT_KEY, TOOL_RESULTS_DIRNAME from deerflow.constants import CONVERSATION_READER_CONTEXT_KEY, TOOL_RESULTS_DIRNAME
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope 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 ( from deerflow.runtime.checkpoint_mode import (
aensure_checkpoint_mode_compatible, aensure_checkpoint_mode_compatible,
inject_checkpoint_mode, inject_checkpoint_mode,
@ -99,6 +103,7 @@ from .naming import resolve_root_run_name
from .schemas import RunStatus from .schemas import RunStatus
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
_THREAD_INCARNATION_UNSET = object()
_checkpoint_locks = AsyncKeyedLockTable[str]() _checkpoint_locks = AsyncKeyedLockTable[str]()
@ -514,6 +519,8 @@ _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: Final[frozenset[str]] = (
CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY, CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY,
DEERFLOW_TRACE_METADATA_KEY, DEERFLOW_TRACE_METADATA_KEY,
CONVERSATION_READER_CONTEXT_KEY, CONVERSATION_READER_CONTEXT_KEY,
THREAD_INCARNATION_CONTEXT_KEY,
THREAD_INCARNATION_METADATA_GUARD_KEY,
"is_subagent", "is_subagent",
"agent_id", "agent_id",
"__run_loop_detection_recorder", "__run_loop_detection_recorder",
@ -538,6 +545,8 @@ def _build_runtime_context(
task_store: Any | None = None, task_store: Any | None = None,
extensions: Any | None = None, extensions: Any | None = None,
conversation_reader: Any | None = None, conversation_reader: Any | None = None,
*,
thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET,
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Build the dict that becomes ``ToolRuntime.context`` for the run. """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. ``langgraph.pregel.main`` where ``parent_runtime.merge(...)`` is invoked.
""" """
runtime_ctx: dict[str, Any] = {"thread_id": thread_id, "run_id": run_id} 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): if isinstance(caller_context, dict):
for key, value in caller_context.items(): for key, value in caller_context.items():
if key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS: if key in _SERVER_OWNED_RUNTIME_CONTEXT_KEYS:
@ -806,6 +817,7 @@ async def run_agent(
agent_factory: Any, agent_factory: Any,
graph_input: dict, graph_input: dict,
config: dict, config: dict,
thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET,
stream_modes: list[str] | None = None, stream_modes: list[str] | None = None,
stream_subgraphs: bool = False, stream_subgraphs: bool = False,
interrupt_before: list[str] | Literal["*"] | None = None, interrupt_before: list[str] | Literal["*"] | None = None,
@ -946,9 +958,13 @@ async def run_agent(
# cancellation cannot strand a pending RunRecord or stream subscriber. # cancellation cannot strand a pending RunRecord or stream subscriber.
if ctx.mcp_task_repo is not None and record.user_id is not None: if ctx.mcp_task_repo is not None and record.user_id is not None:
try: 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( task_rows = await ctx.mcp_task_repo.list_by_thread(
thread_id, thread_id,
user_id=record.user_id, user_id=record.user_id,
thread_incarnation=thread_incarnation,
limit=20, limit=20,
) )
graph_input = { graph_input = {
@ -1074,6 +1090,7 @@ async def run_agent(
task_store, task_store,
extensions, extensions,
ctx.conversation_reader, ctx.conversation_reader,
thread_incarnation=thread_incarnation,
) )
# Bind every checkpoint produced by this run to the effective agent # Bind every checkpoint produced by this run to the effective agent
# identity that produced its state. Manual compaction uses only this # identity that produced its state. Manual compaction uses only this

View File

@ -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. **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`. **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 **Durable batch acceptance**: `batch_task` normalizes optional per-item criteria

View File

@ -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_batches_config import SubagentBatchesConfig
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
from deerflow.extensions import LoadedExtensions, get_loaded_extensions 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_acceptance import check_batch_acceptance
from deerflow.subagents.batch_runtime import BatchSubmitRequest from deerflow.subagents.batch_runtime import BatchSubmitRequest
from deerflow.subagents.capacity import SubagentExecutionCapacity from deerflow.subagents.capacity import SubagentExecutionCapacity
@ -233,6 +234,9 @@ class SubagentBatchService:
item_id, item_id,
) )
return return
executor_kwargs = {}
if THREAD_INCARNATION_CONTEXT_KEY in spec:
executor_kwargs[THREAD_INCARNATION_CONTEXT_KEY] = spec[THREAD_INCARNATION_CONTEXT_KEY]
executor = SubagentExecutor( executor = SubagentExecutor(
config=config, config=config,
tools=tools, tools=tools,
@ -251,6 +255,7 @@ class SubagentBatchService:
execution_capacity=self._execution_capacity, execution_capacity=self._execution_capacity,
extensions=self._extensions, extensions=self._extensions,
acceptance_criteria=item.get("acceptance_criteria"), 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']}" 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) execution_id = executor.execute_async(prompt, task_id=item_id)

View File

@ -37,6 +37,7 @@ from deerflow.authz.principal import normalize_authz_attributes
from deerflow.config import get_app_config from deerflow.config import get_app_config
from deerflow.config.app_config import AppConfig from deerflow.config.app_config import AppConfig
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope 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.models import create_chat_model
from deerflow.runtime.runs.stream_cleanup import close_agent_stream from deerflow.runtime.runs.stream_cleanup import close_agent_stream
from deerflow.runtime.user_context import DEFAULT_USER_ID from deerflow.runtime.user_context import DEFAULT_USER_ID
@ -775,6 +776,9 @@ def _filter_tools(
return filtered return filtered
_THREAD_INCARNATION_UNSET = object()
class SubagentExecutor: class SubagentExecutor:
"""Executor for running subagents.""" """Executor for running subagents."""
@ -806,6 +810,7 @@ class SubagentExecutor:
tool_promotion_recorder: Any | None = None, tool_promotion_recorder: Any | None = None,
tool_progress_recorder: Any | None = None, tool_progress_recorder: Any | None = None,
context_snapshot: ParentContextSnapshot | None = None, context_snapshot: ParentContextSnapshot | None = None,
thread_incarnation: str | None | object = _THREAD_INCARNATION_UNSET,
): ):
"""Initialize the executor. """Initialize the executor.
@ -822,6 +827,8 @@ class SubagentExecutor:
run. Seeded into the child graph state so ``list_uploaded_files`` run. Seeded into the child graph state so ``list_uploaded_files``
can exclude them from historical-upload results. can exclude them from historical-upload results.
thread_id: Thread ID for sandbox operations. 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. trace_id: Trace ID from parent for distributed tracing.
user_id: User ID captured from the parent tool's runtime context. user_id: User ID captured from the parent tool's runtime context.
When None, the tracing layer falls back to DEFAULT_USER_ID. 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.uploaded_files = deepcopy(uploaded_files) if uploaded_files is not None else None
self.context_snapshot = context_snapshot self.context_snapshot = context_snapshot
self.thread_id = thread_id self.thread_id = thread_id
self.thread_incarnation = thread_incarnation
# Generate trace_id if not provided (for top-level calls) # Generate trace_id if not provided (for top-level calls)
self.trace_id = trace_id or str(uuid.uuid4())[:8] self.trace_id = trace_id or str(uuid.uuid4())[:8]
self.user_id = user_id self.user_id = user_id
@ -1548,6 +1556,8 @@ class SubagentExecutor:
context: dict[str, Any] = {} context: dict[str, Any] = {}
if self.thread_id: if self.thread_id:
context["thread_id"] = 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: if self.app_config is not None:
context["app_config"] = self.app_config context["app_config"] = self.app_config
# Propagate guardrail attribution so delegated tool calls are # Propagate guardrail attribution so delegated tool calls are

View File

@ -8,6 +8,7 @@ from langchain.tools import tool
from deerflow.agents.middlewares.input_sanitization_middleware import neutralize_untrusted_tags from deerflow.agents.middlewares.input_sanitization_middleware import neutralize_untrusted_tags
from deerflow.mcp.tasks.runtime import get_mcp_task_submitter 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.builtins.list_uploaded_files_tool import _resolve_thread_id, _resolve_user_id
from deerflow.tools.types import Runtime from deerflow.tools.types import Runtime
@ -36,6 +37,7 @@ async def _list_background_tasks_impl(
records = await get_mcp_task_submitter().list_tasks( records = await get_mcp_task_submitter().list_tasks(
thread_id=thread_id, thread_id=thread_id,
user_id=_resolve_user_id(runtime), user_id=_resolve_user_id(runtime),
thread_incarnation=runtime_thread_incarnation(runtime),
limit=max(1, min(limit, 50)), limit=max(1, min(limit, 50)),
active_only=active_only, active_only=active_only,
) )
@ -72,6 +74,7 @@ async def cancel_background_task(
record = await get_mcp_task_submitter().cancel_matching_task( record = await get_mcp_task_submitter().cancel_matching_task(
thread_id=thread_id, thread_id=thread_id,
user_id=_resolve_user_id(runtime), user_id=_resolve_user_id(runtime),
thread_incarnation=runtime_thread_incarnation(runtime),
task=task, task=task,
) )
except (LookupError, ValueError) as exc: except (LookupError, ValueError) as exc:

View File

@ -15,6 +15,11 @@ from pydantic import BaseModel, Field
from deerflow.authz.principal import normalize_authz_attributes from deerflow.authz.principal import normalize_authz_attributes
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope 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.runtime.user_context import resolve_runtime_user_id
from deerflow.subagents.batch_runtime import ( from deerflow.subagents.batch_runtime import (
BatchItemInput, BatchItemInput,
@ -179,6 +184,8 @@ async def batch_task(
return _result(tool_call_id, content="Batch item keys must be unique.", error=True) 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 {} 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 {} metadata = runtime.config.get("metadata", {}) if runtime is not None else {}
app_config = _batch_app_config(runtime) app_config = _batch_app_config(runtime)
allowed_subagents = metadata.get("allowed_subagents") allowed_subagents = metadata.get("allowed_subagents")
@ -214,6 +221,8 @@ async def batch_task(
"is_internal": context.get("is_internal") is True, "is_internal": context.get("is_internal") is True,
"authz_attributes": normalize_authz_attributes(context.get("authz_attributes")), "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: if KNOWLEDGE_SCOPE_RUNTIME_KEY in context:
execution_spec["knowledge_scope"] = execution_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY]) execution_spec["knowledge_scope"] = execution_scope(context[KNOWLEDGE_SCOPE_RUNTIME_KEY])
try: try:

View File

@ -23,6 +23,11 @@ from deerflow.community.ragflow.sources import cited_source_artifact
from deerflow.config import get_app_config from deerflow.config import get_app_config
from deerflow.extensions import resolve_run_extensions from deerflow.extensions import resolve_run_extensions
from deerflow.knowledge_scope import KNOWLEDGE_SCOPE_RUNTIME_KEY, execution_scope 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.runtime.user_context import resolve_runtime_user_id
from deerflow.sandbox.security import LOCAL_BASH_SUBAGENT_DISABLED_MESSAGE, is_host_bash_allowed 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 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). # tool call delegated to a subagent (user_role=None).
parent_context = runtime.context if runtime is not None else None parent_context = runtime.context if runtime is not None else None
parent_context = parent_context if isinstance(parent_context, dict) else {} 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") user_role = parent_context.get("user_role")
oauth_provider = parent_context.get("oauth_provider") oauth_provider = parent_context.get("oauth_provider")
oauth_id = parent_context.get("oauth_id") oauth_id = parent_context.get("oauth_id")
@ -942,6 +949,10 @@ async def task_tool(
# system-channel authority over framework instructions. # system-channel authority over framework instructions.
"acceptance_criteria": acceptance_criteria, "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: if context_snapshot is not None:
executor_kwargs["context_snapshot"] = context_snapshot executor_kwargs["context_snapshot"] = context_snapshot
middleware_recorder = None middleware_recorder = None

View File

@ -239,6 +239,7 @@ class TestLeadAgentAssembly:
assembly = assemble_lead_agent({"configurable": {"thread_id": "t-1"}}) assembly = assemble_lead_agent({"configurable": {"thread_id": "t-1"}})
assert isinstance(assembly, LeadAgentAssembly) assert isinstance(assembly, LeadAgentAssembly)
assert assembly.graph is not None assert assembly.graph is not None
assert assembly.graph["context_schema"] is dict
assert assembly.descriptor.effective_model assert assembly.descriptor.effective_model
assert assembly.descriptor.fingerprint assert assembly.descriptor.fingerprint
@ -653,6 +654,7 @@ class TestCustomAgentModelSettingsReachTheDescriptor:
TestLeadAgentAssembly._isolate_from_the_ambient_config(monkeypatch) TestLeadAgentAssembly._isolate_from_the_ambient_config(monkeypatch)
with bind_agent_build_extensions(TestLeadAgentAssembly._extensions_with_an_agent_assembly_observer()): with bind_agent_build_extensions(TestLeadAgentAssembly._extensions_with_an_agent_assembly_observer()):
assembly = assemble_lead_agent({"configurable": {"thread_id": "t-bootstrap", "is_bootstrap": True}}) 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 assert "temperature" not in assembly.descriptor.model_parameters

View File

@ -18,7 +18,11 @@ def _clear_submitter():
def _runtime(): def _runtime():
return SimpleNamespace( 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={}, state={},
config={}, config={},
) )
@ -53,6 +57,7 @@ async def test_list_background_tasks_returns_only_safe_local_fields():
manager.list_tasks.assert_awaited_once_with( manager.list_tasks.assert_awaited_once_with(
thread_id="thread-1", thread_id="thread-1",
user_id="user-1", user_id="user-1",
thread_incarnation="incarnation-1",
limit=20, limit=20,
active_only=False, 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( manager.cancel_matching_task.assert_awaited_once_with(
thread_id="thread-1", thread_id="thread-1",
user_id="user-1", user_id="user-1",
thread_incarnation="incarnation-1",
task="report", 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()

View File

@ -6,26 +6,34 @@ import pytest
from langchain_core.messages import ToolMessage from langchain_core.messages import ToolMessage
from langgraph.types import Command 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.subagents.config import SubagentConfig
from deerflow.tools.builtins.batch_task_tool import BatchTaskItem from deerflow.tools.builtins.batch_task_tool import BatchTaskItem
tool_module = importlib.import_module("deerflow.tools.builtins.batch_task_tool") 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( return SimpleNamespace(
state={}, state={},
context={ context=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"],
},
},
config={ config={
"metadata": { "metadata": {
"model_name": "model-a", "model_name": "model-a",
@ -45,7 +53,19 @@ def _message(command: Command) -> ToolMessage:
@pytest.mark.asyncio @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 = AsyncMock()
submitter.submit.return_value = { submitter.submit.return_value = {
"id": "subagent-batch-1", "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( command = await tool_module.batch_task.coroutine(
runtime=_runtime(), runtime=_runtime(thread_incarnation),
title="Process records", title="Process records",
items=[ items=[
BatchTaskItem(key="record-1", prompt="Process one"), BatchTaskItem(key="record-1", prompt="Process one"),
@ -93,10 +113,35 @@ async def test_batch_task_is_explicit_idempotent_submission(monkeypatch) -> None
"mode": "selected", "mode": "selected",
"dataset_ids": ["dataset-1"], "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 message.additional_kwargs["subagent_batch_id"] == "subagent-batch-1"
assert "running independently" in message.content 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 @pytest.mark.asyncio
async def test_batch_task_rejects_duplicate_item_keys_without_submitting(monkeypatch) -> None: async def test_batch_task_rejects_duplicate_item_keys_without_submitting(monkeypatch) -> None:
submitter = AsyncMock() submitter = AsyncMock()

View File

@ -1909,6 +1909,7 @@ async def test_pending_cancel_bypasses_thread_metadata_and_logs_failure(_stub_ap
await asyncio.sleep(0) await asyncio.sleep(0)
assert "thread metadata store failed after cancellation" in caplog.text assert "thread metadata store failed after cancellation" in caplog.text
assert "MCP access will fail closed" not in caplog.text
@pytest.mark.asyncio @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 record.status == RunStatus.running
assert (await run_manager.get(record.run_id)).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 "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(): 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_loop_detection_recorder": "forged",
"__run_tool_promotion_recorder": "forged", "__run_tool_promotion_recorder": "forged",
"__run_tool_progress_recorder": "forged", "__run_tool_progress_recorder": "forged",
"__deerflow_thread_incarnation_metadata_guard": True,
} }
config = build_run_config( config = build_run_config(
"thread-1", "thread-1",

View File

@ -6,9 +6,10 @@ Validates that the LangGraph auth layer enforces the same rules as Gateway:
import asyncio import asyncio
import os import os
import sys
from datetime import timedelta from datetime import timedelta
from pathlib import Path from pathlib import Path
from types import SimpleNamespace from types import ModuleType, SimpleNamespace
from unittest.mock import AsyncMock, patch from unittest.mock import AsyncMock, patch
from uuid import uuid4 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 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.config import AuthConfig, set_auth_config
from app.gateway.auth.jwt import create_access_token, decode_token from app.gateway.auth.jwt import create_access_token, decode_token
from app.gateway.auth.models import User from app.gateway.auth.models import User
from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID from app.gateway.auth_disabled import AUTH_DISABLED_USER_ID
from app.gateway.langgraph_auth import add_owner_filter, authenticate 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 ─────────────────────────────────────────────────────────────── # ── Helpers ───────────────────────────────────────────────────────────────
@ -235,6 +241,267 @@ def test_filter_with_empty_metadata():
assert result == {"user_id": "user-z"} 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"]) @pytest.mark.parametrize("action", ["read", "search"])
def test_studio_user_assistant_discovery_includes_system_and_studio_owned_assistants(action): def test_studio_user_assistant_discovery_includes_system_and_studio_owned_assistants(action):
value = {} value = {}

View File

@ -10,6 +10,7 @@ import subprocess
import sys import sys
import time import time
from collections.abc import Iterator from collections.abc import Iterator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager from contextlib import contextmanager
from pathlib import Path from pathlib import Path
from uuid import uuid4 from uuid import uuid4
@ -17,15 +18,85 @@ from uuid import uuid4
import httpx import httpx
import pytest import pytest
from deerflow.mcp_scope import mcp_session_scope_key
BACKEND_DIR = Path(__file__).resolve().parents[1] BACKEND_DIR = Path(__file__).resolve().parents[1]
_GRAPH_SOURCE = """ _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) from deerflow.mcp import tools as mcp_tools
builder.add_node("noop", lambda state: {})
builder.add_edge(START, "noop")
builder.add_edge("noop", END) 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() graph = builder.compile()
""".lstrip() """.lstrip()
@ -60,6 +131,48 @@ def _free_port() -> int:
return int(sock.getsockname()[1]) 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 @contextmanager
def _running_studio_server( def _running_studio_server(
runtime_dir: Path, runtime_dir: Path,
@ -121,7 +234,7 @@ def _running_studio_server(
) )
base_url = f"http://127.0.0.1:{port}" base_url = f"http://127.0.0.1:{port}"
deadline = time.monotonic() + 45 deadline = time.monotonic() + 90
last_error: Exception | None = None last_error: Exception | None = None
while time.monotonic() < deadline and process.poll() is None: while time.monotonic() < deadline and process.poll() is None:
try: try:
@ -137,7 +250,11 @@ def _running_studio_server(
time.sleep(0.1) time.sleep(0.1)
else: else:
process.terminate() 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')}") pytest.fail(f"LangGraph dev server failed to start ({last_error!r}).\n{log_path.read_text(encoding='utf-8')}")
client = httpx.Client( client = httpx.Client(
@ -224,6 +341,217 @@ def test_studio_can_get_and_search_registered_system_assistant(
assert response.json()["metadata"]["created_by"] == "system" 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( def test_studio_update_cannot_forge_system_provenance(
studio_client: httpx.Client, studio_client: httpx.Client,
): ):

View File

@ -13,8 +13,9 @@ from deerflow.mcp.session_pool import MCPSessionPool
@pytest.mark.asyncio @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"]) @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 = tmp_path / "mcp server"
server_dir.mkdir() server_dir.mkdir()
(server_dir / "marker.txt").write_text("configured-directory", encoding="utf-8") (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")) monkeypatch.setattr(mcp_tools, "get_paths", lambda: Paths(tmp_path / "deerflow"))
pool = MCPSessionPool() pool = MCPSessionPool()
monkeypatch.setattr(mcp_tools, "get_session_pool", lambda: pool) 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: try:
tools = await mcp_tools.get_mcp_tools() tools = await mcp_tools.get_mcp_tools()
@ -66,12 +67,13 @@ mcp.run(transport="stdio")
@pytest.mark.asyncio @pytest.mark.asyncio
@pytest.mark.parametrize("thread_incarnation", [None, "incarnation-1"], ids=["legacy", "versioned"])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"cwd_config", "cwd_config",
[{}, {"cwd": None}, {"cwd": ""}, {"cwd": "$TEST_UNSET_MCP_CWD"}], [{}, {"cwd": None}, {"cwd": ""}, {"cwd": "$TEST_UNSET_MCP_CWD"}],
ids=["omitted", "null", "empty", "unset-env"], 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 = tmp_path / "gateway"
launch_dir.mkdir() launch_dir.mkdir()
monkeypatch.chdir(launch_dir) monkeypatch.chdir(launch_dir)
@ -102,7 +104,7 @@ mcp.run(transport="stdio")
monkeypatch.setattr(mcp_tools, "get_paths", lambda: paths) monkeypatch.setattr(mcp_tools, "get_paths", lambda: paths)
pool = MCPSessionPool() pool = MCPSessionPool()
monkeypatch.setattr(mcp_tools, "get_session_pool", lambda: pool) 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: try:
tools = await mcp_tools.get_mcp_tools() tools = await mcp_tools.get_mcp_tools()

View File

@ -7,6 +7,7 @@ import stat
import sys import sys
import threading import threading
import weakref import weakref
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch from unittest.mock import AsyncMock, MagicMock, patch
import anyio import anyio
@ -15,6 +16,11 @@ from mcp.shared.exceptions import McpError
from mcp.types import CONNECTION_CLOSED, CallToolResult, ErrorData, TextContent 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.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) @pytest.fixture(autouse=True)
@ -24,6 +30,75 @@ def _reset_pool():
reset_session_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 # MCPSessionPool unit tests
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@ -468,7 +543,7 @@ mcp.run(transport="stdio")
"args": ["-c", server, str(marker)], "args": ["-c", server, str(marker)],
} }
runtime = MagicMock() runtime = MagicMock()
runtime.context = {"thread_id": "thread", "user_id": "user"} runtime.context = {"thread_id": "thread", "user_id": "user", "thread_incarnation": "incarnation-1"}
runtime.config = {} runtime.config = {}
with patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)): 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) await wrapped.coroutine(runtime=runtime)
assert exc_info.value.error.code == CONNECTION_CLOSED 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) 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)), patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)),
pytest.raises(type(transport_error)), 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) 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)), patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)),
pytest.raises(McpError, match="Connection closed"), 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) 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)), patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)),
pytest.raises(McpError, match=str(error)), 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() 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)), patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)),
pytest.raises(anyio.ClosedResourceError) as exc_info, 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 assert exc_info.value is error
pool.close_session_if_current.assert_awaited_once_with("srv", "test-user-autouse:default", session) 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)), patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)),
pytest.raises(ToolException, match="invalid input"), 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() 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)), patch("deerflow.mcp.tools.get_paths", return_value=Paths(tmp_path)),
pytest.raises(RuntimeError, match="interceptor failed"), 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() session.call_tool.assert_not_awaited()
pool.close_session_if_current.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", "srv",
{"transport": "stdio", "command": "x", "args": []}, {"transport": "stdio", "command": "x", "args": []},
) )
first_call = asyncio.create_task(wrapped.coroutine(value=1)) runtime = _legacy_tool_runtime()
late_call = asyncio.create_task(wrapped.coroutine(value=2)) 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) await asyncio.wait_for(both_started.wait(), timeout=1)
first_failure.set() first_failure.set()
with pytest.raises(anyio.ClosedResourceError): with pytest.raises(anyio.ClosedResourceError):
await first_call await first_call
await wrapped.coroutine(value=3) await wrapped.coroutine(runtime=runtime, value=3)
late_failure.set() late_failure.set()
with pytest.raises(anyio.ClosedResourceError): with pytest.raises(anyio.ClosedResourceError):
await late_call 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. # Simulate a tool call with a runtime context containing thread_id.
mock_runtime = MagicMock() mock_runtime = MagicMock()
mock_runtime.context = {"thread_id": "thread-42"} mock_runtime.context = {"thread_id": "thread-42", "thread_incarnation": "incarnation-1"}
mock_runtime.config = {} mock_runtime.config = {}
await wrapped.coroutine(runtime=mock_runtime, url="https://example.com") 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) paths = Paths(tmp_path)
connection = {"transport": "stdio", "command": "pw", "args": [], "env": {"KEEP": "1"}} connection = {"transport": "stdio", "command": "pw", "args": [], "env": {"KEEP": "1"}}
mock_runtime = MagicMock() 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 = {} mock_runtime.config = {}
with ( with (
@ -868,7 +949,7 @@ async def test_session_pool_tool_does_not_override_explicit_tmpdir(tmp_path):
paths = Paths(tmp_path) paths = Paths(tmp_path)
connection = {"transport": "stdio", "command": "pw", "args": [], "env": {"TMPDIR": "/operator/tmp"}} connection = {"transport": "stdio", "command": "pw", "args": [], "env": {"TMPDIR": "/operator/tmp"}}
mock_runtime = MagicMock() 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 = {} mock_runtime.config = {}
with ( with (
@ -915,7 +996,7 @@ async def test_session_pool_tool_does_not_override_explicit_cwd(tmp_path):
paths = Paths(tmp_path) paths = Paths(tmp_path)
connection = {"transport": "stdio", "command": "pw", "args": [], "cwd": operator_cwd} connection = {"transport": "stdio", "command": "pw", "args": [], "cwd": operator_cwd}
mock_runtime = MagicMock() 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 = {} mock_runtime.config = {}
with ( with (
@ -962,7 +1043,7 @@ async def test_session_pool_tool_skips_fs_work_for_non_stdio_transport(tmp_path)
paths = Paths(tmp_path) paths = Paths(tmp_path)
connection = {"transport": "sse", "url": "http://localhost:9000/sse", "env": {"KEEP": "1"}} connection = {"transport": "sse", "url": "http://localhost:9000/sse", "env": {"KEEP": "1"}}
mock_runtime = MagicMock() 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 = {} mock_runtime.config = {}
with ( with (
@ -1015,7 +1096,7 @@ async def test_session_pool_tool_skips_after_walk_when_no_text_content(tmp_path)
paths = Paths(tmp_path) paths = Paths(tmp_path)
connection = {"transport": "stdio", "command": "pw", "args": []} connection = {"transport": "stdio", "command": "pw", "args": []}
mock_runtime = MagicMock() 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 = {} mock_runtime.config = {}
with ( with (
@ -1061,7 +1142,7 @@ async def test_session_pool_tool_runs_after_walk_when_text_content_present(tmp_p
paths = Paths(tmp_path) paths = Paths(tmp_path)
connection = {"transport": "stdio", "command": "pw", "args": []} connection = {"transport": "stdio", "command": "pw", "args": []}
mock_runtime = MagicMock() 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 = {} mock_runtime.config = {}
with ( with (
@ -1113,7 +1194,7 @@ async def test_session_pool_tool_forwards_interceptor_headers():
{"transport": "stdio", "command": "x", "args": []}, {"transport": "stdio", "command": "x", "args": []},
tool_interceptors=[header_interceptor], 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"}}) 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": []}, {"transport": "stdio", "command": "x", "args": []},
tool_interceptors=[secret_header_interceptor], 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( mock_session.call_tool.assert_awaited_once_with(
"act", "act",
@ -1208,7 +1289,7 @@ async def test_session_pool_tool_no_headers_omits_meta():
{"transport": "stdio", "command": "x", "args": []}, {"transport": "stdio", "command": "x", "args": []},
tool_interceptors=[passthrough_interceptor], 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}) 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": []}, {"transport": "stdio", "command": "x", "args": []},
tool_interceptors=[invalid_header_interceptor], 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}) mock_session.call_tool.assert_awaited_once_with("act", {"x": 1})
assert "unsupported type" in caplog.text 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": []}) wrapped = _make_session_pool_tool(original_tool, "server", {"transport": "stdio", "command": "x", "args": []})
mock_runtime = MagicMock() mock_runtime = MagicMock()
mock_runtime.context = {} mock_runtime.context = {"thread_incarnation": "incarnation-1"}
mock_runtime.config = {"configurable": {"thread_id": "from-config"}} mock_runtime.config = {"configurable": {"thread_id": "from-config"}}
await wrapped.coroutine(runtime=mock_runtime, x=1) await wrapped.coroutine(runtime=mock_runtime, x=1)
# Verify the session was created with the correct scope key. # Verify the session was created with the canonical versioned JSON scope;
# The scope key is "{user_id}:{thread_id}"; the autouse fixture sets # the autouse fixture sets the effective user to "test-user-autouse".
# the effective user to "test-user-autouse".
pool = get_session_pool() 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 @pytest.mark.asyncio

View File

@ -289,7 +289,12 @@ def test_func_patched_mcp_tool_keeps_toolnode_runtime_injection(tmp_path):
graph.compile().ainvoke( graph.compile().ainvoke(
{"messages": [ai]}, {"messages": [ai]},
config={"configurable": {"thread_id": "T"}}, 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",
},
) )
) )

View File

@ -2,7 +2,14 @@ from datetime import timedelta
import pytest 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(): 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(): def test_driver_registry_rejects_duplicate_names():
registry = McpTaskDriverRegistry() registry = McpTaskDriverRegistry()
driver = object() driver = object()

View File

@ -28,6 +28,7 @@ def _request() -> TaskSubmitRequest:
return TaskSubmitRequest( return TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -47,6 +48,7 @@ def _reference() -> TaskReference:
local_task_id="local-1", local_task_id="local-1",
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
server_name="reports", server_name="reports",
remote_task_id="remote-1", remote_task_id="remote-1",
driver_data={ driver_data={
@ -74,6 +76,7 @@ async def test_submit_uses_structured_content_and_keeps_remote_id_out_of_driver_
"arguments": {"topic": "MCP"}, "arguments": {"topic": "MCP"},
"user_id": "user-1", "user_id": "user-1",
"thread_id": "thread-1", "thread_id": "thread-1",
"thread_incarnation": "incarnation-1",
# Submit is the one durable-task call awaited inside the Agent run, # Submit is the one durable-task call awaited inside the Agent run,
# so it is the only one that may carry request-scoped credentials. # so it is the only one that may carry request-scoped credentials.
"request_scoped_headers": True, "request_scoped_headers": True,

View File

@ -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.engine import close_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.mcp_tasks import McpTaskRepository from deerflow.persistence.mcp_tasks import McpTaskRepository
from deerflow.persistence.thread_meta.model import ThreadMetaRow
@pytest_asyncio.fixture(autouse=True) @pytest_asyncio.fixture(autouse=True)
@ -67,6 +68,7 @@ def _request(remote_id: str) -> TaskSubmitRequest:
return TaskSubmitRequest( return TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", 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 @pytest.mark.asyncio
async def test_submit_poll_restart_recovery_complete_and_fail(tmp_path) -> None: 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))) await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
session_factory = get_session_factory() session_factory = get_session_factory()
assert session_factory is not None assert session_factory is not None
repo = McpTaskRepository(session_factory) repo = McpTaskRepository(session_factory)
await _create_thread(repo)
fake_server = FakeMcpServer() fake_server = FakeMcpServer()
submitted_at = datetime.now(UTC) submitted_at = datetime.now(UTC)
fake_server.status_results.extend( 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) restarted_process = _service(repo, fake_server)
await restarted_process.run_once(now=datetime.now(UTC) + timedelta(seconds=2)) 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 is not None
assert completed["status"] == "completed" assert completed["status"] == "completed"
assert completed["result"] == {"report": "ready"} 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)) 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 is not None
assert failed["status"] == "failed" assert failed["status"] == "failed"
assert failed["error"] == "report generation 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() session_factory = get_session_factory()
assert session_factory is not None assert session_factory is not None
repo = McpTaskRepository(session_factory) repo = McpTaskRepository(session_factory)
await _create_thread(repo)
fake_server = FakeMcpServer() fake_server = FakeMcpServer()
service = _service(repo, fake_server) service = _service(repo, fake_server)
submitted_at = datetime.now(UTC) 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)) 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 is not None
assert retrying["status"] == "submitted" assert retrying["status"] == "submitted"
assert retrying["consecutive_poll_error_count"] == 1 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)) 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 is not None
assert failed["status"] == "failed" assert failed["status"] == "failed"
assert failed["error"] == "report generation failed" assert failed["error"] == "report generation failed"

View File

@ -12,7 +12,7 @@ from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from deerflow.config.database_config import DatabaseConfig 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.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.mcp_tasks.model import McpTaskRow
from deerflow.persistence.thread_meta import ThreadMetaRepository from deerflow.persistence.thread_meta import ThreadMetaRepository
from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.persistence.thread_meta.model import ThreadMetaRow
@ -54,11 +54,17 @@ async def postgres_repositories():
await close_engine() 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( await repo.create(
task_id=task_id, task_id=task_id,
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
expected_thread_incarnation=thread_incarnation,
run_id=None, run_id=None,
tool_call_id=None, tool_call_id=None,
server_name="reports", 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 # PostgreSQL grants the mutation's earlier queued row-lock request
# before this later FOR SHARE request. The task therefore observes the # before this later FOR SHARE request. The task therefore observes the
# committed delete/owner change rather than the pre-mutation row. # 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) await asyncio.sleep(0.1)
assert not create_task.done() assert not create_task.done()
await blocker.commit() await blocker.commit()
await asyncio.wait_for(mutation_task, timeout=5) 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: async with session_factory() as session:
task = await session.get(McpTaskRow, f"task-{mutation}") task = await session.get(McpTaskRow, f"task-{mutation}")
assert task is not None assert task is None
assert task.thread_incarnation is None
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_postgres_task_create_uses_share_lock(postgres_repositories) -> None: async def test_postgres_task_create_uses_share_lock(postgres_repositories) -> None:
thread_repo, task_repo, _session_factory = postgres_repositories 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() engine = get_engine()
assert engine is not None assert engine is not None
statements: list[str] = [] 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) event.listen(engine.sync_engine, "before_cursor_execute", capture_statement)
try: try:
await _create_task(task_repo, "task-share-lock") await _create_task(
task_repo,
"task-share-lock",
thread_incarnation=created["incarnation"],
)
finally: finally:
event.remove(engine.sync_engine, "before_cursor_execute", capture_statement) 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() await legacy_writer.commit()
event.listen(engine.sync_engine, "before_cursor_execute", observe_owner_update) 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 owner_update = None
try: try:
await asyncio.wait_for(task_commit_entered.wait(), timeout=5) await asyncio.wait_for(task_commit_entered.wait(), timeout=5)

View File

@ -11,7 +11,11 @@ from sqlalchemy.ext.asyncio import AsyncSession
from deerflow.config.database_config import DatabaseConfig 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.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.mcp_tasks.model import McpTaskRow
from deerflow.persistence.thread_meta.model import ThreadMetaRow from deerflow.persistence.thread_meta.model import ThreadMetaRow
@ -36,11 +40,26 @@ async def _create_working_task(
now: datetime, now: datetime,
user_id: str = "user-1", user_id: str = "user-1",
remote_task_id: str | None = None, remote_task_id: str | None = None,
thread_incarnation: str | None = None,
) -> dict: ) -> 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( return await repo.create(
task_id=task_id, task_id=task_id,
user_id=user_id, user_id=user_id,
thread_id="thread-1", thread_id="thread-1",
expected_thread_incarnation=thread_incarnation,
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", 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) await _create_working_task(repo, task_id=task_id, now=now)
claim = repo.claim_due_tasks claim = repo.claim_due_tasks
if operation == "apply_cancel_snapshot": 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 claim = repo.claim_cancel_requests
first = await claim(now=now, lease_owner="worker-1", lease_seconds=60, limit=1) 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"]} 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) 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 len(second) == 1
assert second[0]["lease_token"] != first[0]["lease_token"] 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() resume.set()
applied = await asyncio.wait_for(pending, timeout=5) applied = await asyncio.wait_for(pending, timeout=5)
# Check the entire row, including scheduling, errors, results and event # Check the entire row, including scheduling, errors, results and event
# versions, not just the new lease: stale work must have no side effects. # 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 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) second = await asyncio.wait_for(repo.claim_notification_work(now=now + timedelta(seconds=61), **claim_kwargs), timeout=5)
assert len(second) == 1 assert len(second) == 1
assert second[0]["notification_lease_token"] != first[0]["notification_lease_token"] 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() resume.set()
applied = await asyncio.wait_for(pending, timeout=5) 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 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.asyncio
@pytest.mark.parametrize( @pytest.mark.parametrize(
("thread_owner", "expected_incarnation"), "thread_owner",
[ [
("user-1", "matching-owner"), "user-1",
(None, "shared-thread"), None,
("user-2", None),
], ],
) )
async def test_create_atomically_copies_accessible_thread_incarnation( async def test_create_atomically_copies_accessible_thread_incarnation(
tmp_path, tmp_path,
thread_owner, thread_owner,
expected_incarnation,
): ):
repo = await _make_repo(tmp_path) repo = await _make_repo(tmp_path)
now = datetime.now(UTC) now = datetime.now(UTC)
incarnation = expected_incarnation or "different-owner" incarnation = "matching-incarnation"
async with repo._sf() as session: async with repo._sf() as session:
session.add( session.add(
ThreadMetaRow( ThreadMetaRow(
@ -269,25 +288,49 @@ async def test_create_atomically_copies_accessible_thread_incarnation(
) )
await session.commit() 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 assert "thread_incarnation" not in task
async with repo._sf() as session: async with repo._sf() as session:
row = await session.get(McpTaskRow, "new-writer") row = await session.get(McpTaskRow, "new-writer")
assert row is not None assert row is not None
assert row.thread_incarnation == expected_incarnation assert row.thread_incarnation == incarnation
@pytest.mark.asyncio @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) 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: async with repo._sf() as session:
row = await session.get(McpTaskRow, "missing-thread") row = await session.get(McpTaskRow, "missing-thread")
assert row is not None assert row is None
assert row.thread_incarnation is None
@pytest.mark.asyncio @pytest.mark.asyncio
@ -310,9 +353,9 @@ async def test_create_observes_delete_and_recreate_at_insert_boundary(tmp_path):
engine = get_engine() engine = get_engine()
assert engine is not None assert engine is not None
replaced = False replaced = False
insert_statement = None lock_statement = None
def replace_thread_before_task_insert( def replace_thread_before_scope_lock(
_conn, _conn,
_cursor, _cursor,
statement, statement,
@ -320,11 +363,11 @@ async def test_create_observes_delete_and_recreate_at_insert_boundary(tmp_path):
_context, _context,
_executemany, _executemany,
): ):
nonlocal insert_statement, replaced nonlocal lock_statement, replaced
if replaced or not statement.lstrip().upper().startswith("INSERT INTO MCP_TASKS"): if replaced or "UPDATE THREADS_META" not in statement.upper():
return return
replaced = True replaced = True
insert_statement = statement lock_statement = statement
with contextlib.closing(sqlite3.connect(tmp_path / "deerflow.db")) as connection: with contextlib.closing(sqlite3.connect(tmp_path / "deerflow.db")) as connection:
with connection: with connection:
connection.execute("DELETE FROM threads_meta WHERE thread_id = ?", ("thread-1",)) 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: 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: 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 replaced is True
assert insert_statement is not None assert lock_statement is not None
normalized_insert = " ".join(insert_statement.upper().split()) normalized_lock = " ".join(lock_statement.upper().split())
assert "SELECT THREADS_META.INCARNATION" in normalized_insert assert "UPDATE THREADS_META SET INCARNATION = INCARNATION" in normalized_lock
assert "INSERT INTO MCP_TASKS" in normalized_insert assert "INCARNATION IS ?" in normalized_lock
async with repo._sf() as session: async with repo._sf() as session:
row = await session.get(McpTaskRow, "racing-task") row = await session.get(McpTaskRow, "racing-task")
assert row is not None assert row is None
assert row.thread_incarnation == "replacement-incarnation"
@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 @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", 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( other_user = await _create_working_task(
repo, repo,
task_id="task-remote-3", 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 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 is not None
assert stored["status"] == "completed" assert stored["status"] == "completed"
assert stored["result"] == {"report": "ready"} 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 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 is not None
assert stored["status"] == "working" assert stored["status"] == "working"
assert stored["result"] is None 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 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 is not None
assert stored["input_required"] == {"prompt": "Approve deployment?"} assert stored["input_required"] == {"prompt": "Approve deployment?"}
assert stored["notification_status"] == "pending" assert stored["notification_status"] == "pending"
@ -596,7 +811,7 @@ async def test_release_claim_retries_transient_poll_failure(tmp_path):
) )
assert released is True 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 is not None
assert stored["status"] == "working" assert stored["status"] == "working"
assert stored["last_poll_error"] == "temporary network failure" 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 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 is not None
assert stored["lease_owner"] == "worker-1" assert stored["lease_owner"] == "worker-1"
assert stored["lease_token"] == new_token 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 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 is not None
assert stored["lease_owner"] == "worker-1" assert stored["lease_owner"] == "worker-1"
assert stored["lease_token"] == new_token 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 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 is not None
assert stored["lease_owner"] == "worker-1" assert stored["lease_owner"] == "worker-1"
assert stored["lease_token"] == new_token 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 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 is not None
assert stored["notification_lease_owner"] == "notifier" assert stored["notification_lease_owner"] == "notifier"
assert stored["notification_lease_token"] == new_notify_token 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, next_poll_at=retry_at,
error="temporary network failure", 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 assert before is not None
reclaimed = await repo.claim_due_tasks(now=retry_at, lease_owner="worker-2", lease_seconds=60, limit=10) 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 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 is not None
assert stored["next_poll_at"] == before["next_poll_at"] assert stored["next_poll_at"] == before["next_poll_at"]
assert stored["last_poll_error"] == before["last_poll_error"] 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 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 is not None
assert stored["lease_owner"] == "worker-current" assert stored["lease_owner"] == "worker-current"
assert stored["lease_expires_at"] is not None 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), next_poll_at=now - timedelta(seconds=1),
error="temporary network failure", 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 is not None
assert stored["consecutive_poll_error_count"] == expected_errors 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 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 is not None
assert stored["consecutive_poll_error_count"] == 0 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, next_poll_at=None,
polled_at=now, 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 is not None
assert changed["event_version"] == 2 assert changed["event_version"] == 2
assert changed["dispatch_version"] == 1 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", error="Agent run failed",
now=now, 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 is not None
assert failed["notification_status"] == "retry" assert failed["notification_status"] == "retry"
assert failed["dispatch_attempt"] == 1 assert failed["dispatch_attempt"] == 1
@ -1083,7 +1298,7 @@ async def test_unexpected_notification_failure_releases_lease_without_changing_p
error="run store unavailable", 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 is not None
assert stored["notification_status"] == "claimed" assert stored["notification_status"] == "claimed"
assert stored["notification_lease_owner"] is None 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, 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 is not None
assert stored["notification_status"] == "pending" assert stored["notification_status"] == "pending"
assert stored["notification_attempt_count"] == 1 assert stored["notification_attempt_count"] == 1
@ -1185,7 +1400,7 @@ async def test_permanent_notification_failure_is_not_reclaimed(tmp_path):
now=now, 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 is not None
assert stored["notification_status"] == "dead_letter" assert stored["notification_status"] == "dead_letter"
assert stored["notification_attempt_count"] == 1 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, 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 is not None
assert stored["notification_status"] == "dead_letter" assert stored["notification_status"] == "dead_letter"
assert stored["notification_run_id"] is None 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, 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 is not None
assert stored["notification_status"] == "pending" assert stored["notification_status"] == "pending"
assert stored["notification_attempt_count"] == 0 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", "task-cancel",
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
requested_at=now, requested_at=now,
) )
assert requested is not None 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", "task-cancel",
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
requested_at=now + timedelta(seconds=1), requested_at=now + timedelta(seconds=1),
) )
assert repeated is not None 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, input_required=None,
completed_at=now, 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 is not None
assert stored["status"] == "cancelled" assert stored["status"] == "cancelled"
assert stored["notification_status"] == "pending" 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 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 is not None
assert stored["lease_owner"] == "worker-same" assert stored["lease_owner"] == "worker-same"
assert stored["lease_token"] == reclaimed[0]["lease_token"] 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", "task-late-cancel",
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
requested_at=now, 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 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 is not None
assert stored["lease_owner"] == "worker-same" assert stored["lease_owner"] == "worker-same"
assert stored["lease_token"] == reclaimed[0]["lease_token"] 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 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 is not None
assert stored["notification_lease_owner"] == "notifier-same" assert stored["notification_lease_owner"] == "notifier-same"
assert stored["notification_lease_token"] == reclaimed[0]["notification_lease_token"] 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 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 is not None
assert stored["lease_owner"] == "worker-same" assert stored["lease_owner"] == "worker-same"
assert stored["lease_token"] == reclaimed[0]["lease_token"] assert stored["lease_token"] == reclaimed[0]["lease_token"]

View File

@ -19,7 +19,10 @@ from deerflow.mcp.tasks import (
TaskSubmitRequest, TaskSubmitRequest,
) )
from deerflow.mcp.tasks.ordinary import McpTaskProtocolError 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.manager import ConflictError
from deerflow.runtime.runs.schemas import RunStatus from deerflow.runtime.runs.schemas import RunStatus
@ -85,6 +88,12 @@ class DuplicateCreateRepository(FakeRepository):
raise DuplicateMcpRemoteTaskError("already tracked") raise DuplicateMcpRemoteTaskError("already tracked")
class DriftedThreadCreateRepository(FakeRepository):
async def create(self, **kwargs):
self.created.append(kwargs)
raise McpTaskThreadMismatchError("thread incarnation changed")
class FakeDriver: class FakeDriver:
def __init__( def __init__(
self, self,
@ -202,6 +211,7 @@ async def test_submit_persists_remote_handle_before_returning():
request = TaskSubmitRequest( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -214,6 +224,7 @@ async def test_submit_persists_remote_handle_before_returning():
assert created["remote_task_id"] == "remote-1" assert created["remote_task_id"] == "remote-1"
persisted = repo.created[0] persisted = repo.created[0]
assert persisted["expected_thread_incarnation"] == "incarnation-1"
assert persisted["next_poll_at"] == now + timedelta(seconds=9) assert persisted["next_poll_at"] == now + timedelta(seconds=9)
assert persisted["driver_data"] == {"submit_tool": "submit", "status_tool": "status"} assert persisted["driver_data"] == {"submit_tool": "submit", "status_tool": "status"}
assert driver.submit_calls[0].local_task_id == created["id"] 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( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", 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 @pytest.mark.asyncio
async def test_submit_cancellation_during_persistence_cancels_remote_task(): async def test_submit_cancellation_during_persistence_cancels_remote_task():
repo = BlockingCreateRepository() repo = BlockingCreateRepository()
@ -286,6 +339,7 @@ async def test_submit_cancellation_during_persistence_cancels_remote_task():
request = TaskSubmitRequest( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -329,6 +383,7 @@ async def test_submit_repeated_cancellation_does_not_interrupt_compensation():
request = TaskSubmitRequest( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -376,6 +431,7 @@ async def test_submit_stops_waiting_for_hung_compensation_without_cancelling_it(
request = TaskSubmitRequest( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -426,6 +482,7 @@ async def test_submit_cancellation_preserves_cancelled_error_when_compensation_f
request = TaskSubmitRequest( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -472,6 +529,7 @@ async def test_submit_cancels_remote_task_when_its_id_exceeds_storage_limit():
request=TaskSubmitRequest( request=TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -512,6 +570,7 @@ async def test_duplicate_remote_handle_is_rejected_without_cancelling_existing_t
request=TaskSubmitRequest( request=TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-2", thread_id="thread-2",
thread_incarnation="incarnation-2",
run_id="run-2", run_id="run-2",
tool_call_id="call-2", tool_call_id="call-2",
server_name="reports", server_name="reports",
@ -548,6 +607,7 @@ async def test_cancel_task_persists_request_without_calling_remote():
task_id="task-1", task_id="task-1",
thread_id="thread-1", thread_id="thread-1",
user_id="user-1", user_id="user-1",
thread_incarnation="incarnation-1",
) )
assert result == record assert result == record
@ -961,6 +1021,7 @@ async def test_submit_preserves_persistence_error_when_compensation_cancel_fails
request = TaskSubmitRequest( request = TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",
@ -1166,6 +1227,7 @@ async def test_persisted_snapshot_errors_are_bounded_on_submit_and_poll():
request=TaskSubmitRequest( request=TaskSubmitRequest(
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation="incarnation-1",
run_id="run-1", run_id="run-1",
tool_call_id="call-1", tool_call_id="call-1",
server_name="reports", server_name="reports",

View File

@ -15,7 +15,8 @@ from mcp.types import CONNECTION_CLOSED, ErrorData
from deerflow.config.extensions_config import ExtensionsConfig from deerflow.config.extensions_config import ExtensionsConfig
from deerflow.config.paths import Paths from deerflow.config.paths import Paths
from deerflow.mcp.session_pool import MCPSessionPool 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: def _config() -> ExtensionsConfig:
@ -78,8 +79,21 @@ async def _assert_configured_timeout(awaitable: Coroutine[Any, Any, Any]) -> Non
await task await task
def test_task_session_scope_includes_user_and_thread() -> None: def test_task_session_scope_includes_user_thread_and_incarnation() -> None:
assert mcp_task_session_scope_key(user_id="user-1", thread_id="thread-1") == "user-1:thread-1" 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 @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"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
assert actual is result 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"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
pool.close_session_if_current.assert_awaited_once_with( 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"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
pool.close_session_if_current.assert_not_awaited() 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"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
session.call_tool.assert_not_awaited() session.call_tool.assert_not_awaited()
@ -277,6 +295,7 @@ mcp.run(transport="stdio")
arguments={}, arguments={},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
server_config.tool_call_timeout = 0.05 server_config.tool_call_timeout = 0.05
@ -287,6 +306,7 @@ mcp.run(transport="stdio")
arguments={"task_id": submitted.structuredContent["task_id"]}, arguments={"task_id": submitted.structuredContent["task_id"]},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
assert exc_info.value.error.code == 408 assert exc_info.value.error.code == 408
@ -298,6 +318,7 @@ mcp.run(transport="stdio")
arguments={"task_id": submitted.structuredContent["task_id"]}, arguments={"task_id": submitted.structuredContent["task_id"]},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
finally: finally:
await pool.close_all() await pool.close_all()
@ -337,6 +358,7 @@ async def test_stdio_task_session_initialization_respects_configured_timeout() -
arguments={"task_id": "remote-1"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
pool.close_session.assert_not_awaited() 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"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
assert actual is result assert actual is result
@ -420,6 +443,7 @@ async def test_remote_task_session_initialization_respects_configured_timeout(tr
arguments={"task_id": "remote-1"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-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"}, arguments={"task_id": "remote-1"},
user_id="user-1", user_id="user-1",
thread_id="thread-1", thread_id="thread-1",
thread_incarnation=None,
) )
) )

View File

@ -140,7 +140,7 @@ async def test_submit_wrapper_persists_before_returning_only_the_local_handle()
) )
submit_tool = configured[0] submit_tool = configured[0]
runtime = SimpleNamespace( 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={}, config={},
tool_call_id="call-1", tool_call_id="call-1",
) )

View File

@ -14,13 +14,16 @@ class FakeRepository:
self.list_calls = [] self.list_calls = []
self.get_calls = [] self.get_calls = []
async def list_by_thread(self, 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, limit)) self.list_calls.append((thread_id, user_id, thread_incarnation, limit))
return list(self.rows) return list(self.rows)
async def get(self, task_id, *, user_id): async def get(self, task_id, *, user_id, thread_id, thread_incarnation):
self.get_calls.append((task_id, user_id)) 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), None) 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): def _record(**overrides):
@ -56,6 +59,15 @@ def _request(repo):
state=SimpleNamespace( state=SimpleNamespace(
mcp_task_repo=repo, mcp_task_repo=repo,
mcp_task_service=SimpleNamespace(tracking_degraded_after_errors=3), 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 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 @pytest.mark.asyncio
async def test_list_returns_only_safe_current_user_thread_fields(monkeypatch) -> None: async def test_list_returns_only_safe_current_user_thread_fields(monkeypatch) -> None:
repo = FakeRepository([_record()]) repo = FakeRepository([_record()])
@ -78,7 +177,7 @@ async def test_list_returns_only_safe_current_user_thread_fields(monkeypatch) ->
limit=25, limit=25,
) )
assert repo.list_calls == [("thread-1", "user-1", 25)] assert repo.list_calls == [("thread-1", "user-1", "incarnation-1", 25)]
assert response == [ assert response == [
{ {
"task_id": "mcp-task-1", "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", task_id="mcp-task-1",
thread_id="thread-1", thread_id="thread-1",
user_id="user-1", user_id="user-1",
thread_incarnation="incarnation-1",
) )
assert response["status"] == "working" assert response["status"] == "working"
assert response["cancel_requested"] is True assert response["cancel_requested"] is True

View File

@ -184,8 +184,8 @@ async def test_run_agent_cleans_up_when_mcp_task_projection_is_cancelled():
projection_started = asyncio.Event() projection_started = asyncio.Event()
class BlockingTaskRepository: class BlockingTaskRepository:
async def list_by_thread(self, thread_id, *, user_id, limit): async def list_by_thread(self, thread_id, *, user_id, thread_incarnation, limit):
del thread_id, user_id, limit del thread_id, user_id, thread_incarnation, limit
projection_started.set() projection_started.set()
await asyncio.Event().wait() 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, agent_factory=agent_factory,
graph_input={}, graph_input={},
config={}, config={},
thread_incarnation=None,
) )
) )
await asyncio.wait_for(projection_started.wait(), timeout=1) 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" 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(): def test_build_runtime_context_ignores_caller_pre_existing_message_ids():
caller_context = {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}} caller_context = {CURRENT_RUN_PRE_EXISTING_MESSAGE_IDS_KEY: {"spoofed"}}

View File

@ -4,9 +4,12 @@ import pytest
import pytest_asyncio import pytest_asyncio
from deerflow.config.database_config import DatabaseConfig 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.engine import close_engine, get_session_factory, init_engine_from_config
from deerflow.persistence.subagent_batches import SubagentBatchRepository from deerflow.persistence.subagent_batches import SubagentBatchRepository
_MISSING = object()
@pytest_asyncio.fixture(autouse=True) @pytest_asyncio.fixture(autouse=True)
async def _close_engine() -> None: 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 == [] 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 @pytest.mark.asyncio
async def test_expired_lease_is_recovered_with_stable_item_identity(tmp_path) -> None: async def test_expired_lease_is_recovered_with_stable_item_identity(tmp_path) -> None:
repo = await _repo(tmp_path) repo = await _repo(tmp_path)

View File

@ -7,11 +7,14 @@ import pytest
from deerflow.config.subagent_batches_config import SubagentBatchesConfig from deerflow.config.subagent_batches_config import SubagentBatchesConfig
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig 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 import batch_service as service_module
from deerflow.subagents.batch_runtime import BatchSubmitRequest from deerflow.subagents.batch_runtime import BatchSubmitRequest
from deerflow.subagents.batch_service import SubagentBatchService from deerflow.subagents.batch_service import SubagentBatchService
from deerflow.subagents.capacity import SubagentExecutionCapacity from deerflow.subagents.capacity import SubagentExecutionCapacity
_MISSING = object()
class FakeStatus(Enum): class FakeStatus(Enum):
PENDING = "pending" PENDING = "pending"
@ -116,7 +119,19 @@ async def test_submit_defaults_only_the_limits_the_caller_omitted(overrides: dic
@pytest.mark.asyncio @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( result = SimpleNamespace(
status=FakeStatus.RUNNING, status=FakeStatus.RUNNING,
result=None, result=None,
@ -125,6 +140,10 @@ async def test_execute_item_marks_real_running_then_persists_terminal_result(mon
token_usage_records=None, 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: class Repository:
def __init__(self) -> None: def __init__(self) -> None:
self.marked_running = False 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", "thread_id": "thread-1",
"user_id": "user-1", "user_id": "user-1",
"run_id": "run-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", "mode": "selected",
"dataset_ids": ["dataset-1"], "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 @pytest.mark.asyncio

View File

@ -4177,6 +4177,71 @@ class _FakeStreamAgent:
yield # pragma: no cover - make this an async generator 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: class TestSubagentCheckpointLineage:
"""Keep delegated graphs on the parent run's checkpoint lineage.""" """Keep delegated graphs on the parent run's checkpoint lineage."""

View File

@ -3265,6 +3265,34 @@ def _capture_executor_call(monkeypatch, **call_kwargs):
return captured["executor_kwargs"], captured["prompt"] 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): def test_task_tool_forwards_acceptance_criteria_to_executor(monkeypatch):
"""RFC #4651 PR3: criteria travel via the executor constructor; the """RFC #4651 PR3: criteria travel via the executor constructor; the
executor appends them to the subagent's task HumanMessage as untrusted executor appends them to the subagent's task HumanMessage as untrusted