diff --git a/fastmcp_slim/fastmcp/server/dependencies.py b/fastmcp_slim/fastmcp/server/dependencies.py index 5b34e16ef..9e2e1c268 100644 --- a/fastmcp_slim/fastmcp/server/dependencies.py +++ b/fastmcp_slim/fastmcp/server/dependencies.py @@ -207,6 +207,17 @@ def set_worker_server_resolver( _worker_server_resolver = resolver +#: Headers a background task carries from its originating request. A worker has +#: no live HTTP request — especially a Redis-backed worker in a separate process +#: — so ``get_http_request()`` correctly raises there. The tasks extension sets +#: this from the task snapshot so ``get_http_headers()`` still returns the +#: submitting request's headers without fabricating a fake ``Request`` (which +#: would make ``get_http_request()``/``CurrentRequest()`` wrongly succeed). +_background_task_headers: ContextVar[dict[str, str] | None] = ContextVar( + "fastmcp_background_task_headers", default=None +) + + # --- Docket availability check --- _DOCKET_AVAILABLE: bool | None = None @@ -494,14 +505,21 @@ def get_http_headers( headers: dict[str, str] = {} try: - request = get_http_request() - for name, value in request.headers.items(): - lower_name = name.lower() - if lower_name not in exclude_headers: - headers[lower_name] = str(value) - return headers + source: Any = get_http_request().headers.items() except RuntimeError: - return {} + # No live request: inside a background-task worker, fall back to the + # headers the task carried from its originating request (set by the + # tasks extension from the snapshot). Empty elsewhere. + task_headers = _background_task_headers.get() + if task_headers is None: + return {} + source = task_headers.items() + + for name, value in source: + lower_name = name.lower() + if lower_name not in exclude_headers: + headers[lower_name] = str(value) + return headers def get_access_token() -> AccessToken | None: diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index f48f974da..9975c088f 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/fastmcp_tasks/fastmcp_tasks/context.py @@ -144,10 +144,17 @@ class TaskContextSnapshot: http_headers: dict[str, str] | None = None origin_request_id: str | None = None session_id: str | None = None + owning_tool_name: str | None = None @classmethod - def capture(cls) -> TaskContextSnapshot: - """Capture current context for background task execution.""" + def capture(cls, owning_tool_name: str | None = None) -> TaskContextSnapshot: + """Capture current context for background task execution. + + ``owning_tool_name`` is the routable name of the tool the call targeted. + A remote worker (separate process) cannot reach the submitting process's + server map, so it re-resolves the owning (child) server from this name + against the root — see ``make_task_context``. + """ from fastmcp.server.dependencies import ( get_access_token, get_context, @@ -170,6 +177,7 @@ class TaskContextSnapshot: str(request_context.request_id) if request_context is not None else None ), session_id=session_id, + owning_tool_name=owning_tool_name, ) @classmethod @@ -186,6 +194,7 @@ class TaskContextSnapshot: http_headers=headers, origin_request_id=parsed.get("origin_request_id"), session_id=parsed.get("session_id"), + owning_tool_name=parsed.get("owning_tool_name"), ) def to_json(self) -> str: @@ -196,6 +205,7 @@ class TaskContextSnapshot: "http_headers": self.http_headers, "origin_request_id": self.origin_request_id, "session_id": self.session_id, + "owning_tool_name": self.owning_tool_name, } ) @@ -400,6 +410,8 @@ def resolve_worker_server() -> FastMCP | None: Installed as core's worker-server resolver by ``TasksExtension`` so ``get_server()``/``CurrentFastMCP()`` inside a worker resolve to the (child) server the task was submitted against, not the root that runs the worker. + The map is populated at submission (same process) and, for a remote worker, + by ``make_task_context`` re-resolving from the snapshot before the tool runs. """ task_info = get_task_context() if task_info is None: @@ -407,15 +419,43 @@ def resolve_worker_server() -> FastMCP | None: return get_task_server(task_info.task_id) +async def _resolve_owning_server( + snapshot: TaskContextSnapshot | None, +) -> FastMCP | None: + """Re-resolve a mounted task's owning child server from the root (remote worker). + + A separate worker process cannot reach the submitting process's server map, + so the owning server is recovered by looking the snapshotted tool name up on + the root: a mounted tool resolves to a ``FastMCPProviderTool`` referencing + its child server. Returns ``None`` for an unmounted tool (the root owns it) + or when the name no longer resolves, so the caller falls back to the root. + """ + if snapshot is None or snapshot.owning_tool_name is None: + return None + from fastmcp.exceptions import NotFoundError + from fastmcp.server.dependencies import get_server + from fastmcp.server.providers.fastmcp_provider import FastMCPProviderTool + + root = get_server() + try: + tool = await root.get_tool(snapshot.owning_tool_name) + except NotFoundError: + return None + if isinstance(tool, FastMCPProviderTool): + return tool._server + return None + + def _apply_snapshot_to_context(snapshot: TaskContextSnapshot) -> None: """Populate the ambient request context a worker's tool body reads. A Docket worker has no live request or SDK auth context — especially a - Redis-backed worker in a separate process. Rather than teach core's - ``get_access_token()`` / ``get_http_headers()`` about tasks, this restores - the *same* context vars a normal request would set, so those functions work - unchanged: the SDK auth context var (from the snapshotted token) and a - minimal HTTP request rebuilt from the snapshotted headers. Runs inside + Redis-backed worker in a separate process. This restores the context vars a + tool reads so ``get_access_token()`` / ``get_http_headers()`` work unchanged: + the SDK auth context var (from the snapshotted token) and core's background + task-headers var (from the snapshotted headers). It deliberately does *not* + fabricate a live ``Request``, so ``get_http_request()`` / ``CurrentRequest()`` + still raise inside a task — there is no request. Runs inside ``restore_task_snapshot`` (a Docket dependency), whose context vars propagate to the tool the same way the snapshot var already does. """ @@ -436,27 +476,9 @@ def _apply_snapshot_to_context(snapshot: TaskContextSnapshot) -> None: auth_context_var.set(AuthenticatedUser(token)) if snapshot.http_headers: - from starlette.requests import Request + from fastmcp.server.dependencies import _background_task_headers - from fastmcp.server.http import _current_http_request - - _current_http_request.set( - Request( - { - "type": "http", - "http_version": "1.1", - "method": "POST", - "scheme": "http", - "path": "/", - "raw_path": b"/", - "query_string": b"", - "headers": [ - (name.encode("latin-1"), value.encode("latin-1")) - for name, value in snapshot.http_headers.items() - ], - } - ) - ) + _background_task_headers.set(dict(snapshot.http_headers)) async def make_task_context() -> Context | None: @@ -482,8 +504,16 @@ async def make_task_context() -> Context | None: if task_info is None: return None - server = get_task_server(task_info.task_id) or get_server() snapshot = _recall_snapshot(task_info.task_id) + server = get_task_server(task_info.task_id) + if server is None: + # In-process submission map missed — this is a remote worker (separate + # process). Re-resolve the owning (child) server from the root using the + # snapshotted tool name, and register it so `CurrentFastMCP()` mid-tool + # resolves the child too. Falls back to the root when unmounted or + # unresolvable. + server = await _resolve_owning_server(snapshot) or get_server() + register_task_server(task_info.task_id, server) origin_request_id = snapshot.origin_request_id if snapshot else None ctx = Context( diff --git a/fastmcp_tasks/fastmcp_tasks/creation.py b/fastmcp_tasks/fastmcp_tasks/creation.py index 678e57881..745e7ceb7 100644 --- a/fastmcp_tasks/fastmcp_tasks/creation.py +++ b/fastmcp_tasks/fastmcp_tasks/creation.py @@ -109,7 +109,7 @@ async def create_task( created_at_key = docket.key(f"{prefix}:{task_id}:created_at") poll_interval_key = docket.key(f"{prefix}:{task_id}:poll_interval") - snapshot = TaskContextSnapshot.capture() + snapshot = TaskContextSnapshot.capture(owning_tool_name=tool.name) async with docket.redis() as redis: await redis.set(task_meta_key, task_key, ex=ttl_seconds) diff --git a/tests/tasks/server/test_snapshot_restore.py b/tests/tasks/server/test_snapshot_restore.py index 6fe5014fc..422b90114 100644 --- a/tests/tasks/server/test_snapshot_restore.py +++ b/tests/tasks/server/test_snapshot_restore.py @@ -13,6 +13,7 @@ from __future__ import annotations import contextvars from unittest.mock import patch +import pytest from fastmcp_tasks.context import ( TaskContextSnapshot, _apply_snapshot_to_context, @@ -118,6 +119,27 @@ def test_apply_snapshot_restores_auth_and_headers_in_clean_context(): contextvars.copy_context().run(run_in_clean_worker_context) +def test_apply_snapshot_headers_without_faking_a_request(): + """Snapshot headers are readable, but no live request is fabricated. + + `get_http_headers()` returns the submitting request's headers, while + `get_http_request()` still raises — there is no live request inside a + background task, and impersonating one would make `CurrentRequest()` expose + invented method/URL/client data. + """ + from fastmcp.server.dependencies import get_http_request + + snapshot = TaskContextSnapshot(http_headers={"x-trace-id": "abc123"}) + + def run_in_clean_worker_context() -> None: + _apply_snapshot_to_context(snapshot) + assert get_http_headers()["x-trace-id"] == "abc123" + with pytest.raises(RuntimeError): + get_http_request() + + contextvars.copy_context().run(run_in_clean_worker_context) + + def test_apply_snapshot_skips_expired_token(): """An expired snapshot token is not installed, so the worker is unauthenticated. diff --git a/tests/tasks/server/test_task_mount.py b/tests/tasks/server/test_task_mount.py index 15dfbee4d..c13836247 100644 --- a/tests/tasks/server/test_task_mount.py +++ b/tests/tasks/server/test_task_mount.py @@ -123,6 +123,51 @@ class TestMountedToolTasks: assert "child sync: hi" in result.content[0].text +class TestRemoteWorkerServerResolution: + """A separate worker process re-resolves the owning child from the root. + + The in-process submission map is unreachable across processes, so the worker + recovers the mounted child server from the snapshotted tool name instead of + falling back to the root (which would break child-specific state/config). + """ + + async def test_resolve_owning_server_recovers_mounted_child(self, parent_server): + import weakref + + from fastmcp_tasks.context import ( + TaskContextSnapshot, + _resolve_owning_server, + ) + + from fastmcp.server.dependencies import _current_server + + child = await parent_server.get_tool("child_multiply") + + token = _current_server.set(weakref.ref(parent_server)) + try: + snapshot = TaskContextSnapshot(owning_tool_name="child_multiply") + resolved = await _resolve_owning_server(snapshot) + assert resolved is child._server + + # A parent-owned (unmounted) tool resolves to None so the caller + # falls back to the root, and a missing name is likewise None. + assert ( + await _resolve_owning_server( + TaskContextSnapshot(owning_tool_name="parent_tool") + ) + is None + ) + assert ( + await _resolve_owning_server( + TaskContextSnapshot(owning_tool_name="does_not_exist") + ) + is None + ) + assert await _resolve_owning_server(TaskContextSnapshot()) is None + finally: + _current_server.reset(token) + + class TestMountedToolTasksNoPrefix: async def test_mounted_tool_without_prefix_works(self, child_server): parent = FastMCP("parent-no-prefix")