mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 05:54:19 +02:00
Address review feedback on SEP-2663 tasks
- Client task support is opt-in via importing fastmcp_tasks (drop the core auto-load of companion packages); a plain Client never advertises tasks. - A worker restores the submitting caller's auth token and headers from the task snapshot into the standard ambient context, so get_access_token() / get_http_headers() work in a distributed worker with no new core hooks. - worker_cli validates the loaded extension's resolved backend, not env defaults, so a constructor-configured Redis worker starts. - Thread the per-call read timeout through task polling; bound ToolTask.wait by its deadline; set_elicitation_callback rebuilds internal extensions so a later-set handler answers in-task input. - README imports TaskConfig from fastmcp.utilities.tasks. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jeremiah Lowin
parent
e5ca0269cb
commit
19c5c507cc
11 changed files with 259 additions and 58 deletions
|
|
@ -13,14 +13,17 @@ import { VersionBadge } from "/snippets/version-badge.mdx"
|
|||
Some tool calls take a while. The MCP background tasks extension lets a server run one in the background instead of holding the request open, and FastMCP's client drives the whole thing for you — most of the time you don't need to know a call was tasked at all.
|
||||
|
||||
<Note>
|
||||
**Background tasks require the modern protocol.** The tasks capability is negotiated over `2026-07-28` connections. `mode="auto"` (the client default) negotiates it automatically; `mode="legacy"` never does, so a tasked tool just runs synchronously for a legacy-pinned client. See [protocol negotiation](/clients/client#protocol-negotiation).
|
||||
**Client task support is opt-in.** Install the `fastmcp-tasks` package (`pip install "fastmcp[tasks]"`) and import it — importing `fastmcp_tasks` anywhere (which you do to use `call_tool_task`) enables task support for every `Client` in the process. Without it, a `Client` never advertises the tasks capability, so the server runs its calls synchronously and background tasks simply don't happen.
|
||||
|
||||
**Tasks also require the modern protocol.** The capability is negotiated over `2026-07-28` connections. `mode="auto"` (the client default) negotiates it automatically; `mode="legacy"` never does. See [protocol negotiation](/clients/client#protocol-negotiation).
|
||||
</Note>
|
||||
|
||||
## Transparent Calls
|
||||
|
||||
Just call the tool. If the server runs it as a background task, `call_tool` polls it to completion under the hood and returns the same result you'd get from a synchronous call — the task is invisible.
|
||||
With task support enabled, just call the tool. If the server runs it as a background task, `call_tool` polls it to completion under the hood and returns the same result you'd get from a synchronous call — the task is invisible.
|
||||
|
||||
```python
|
||||
import fastmcp_tasks # enables client task support
|
||||
from fastmcp import Client
|
||||
|
||||
async with Client(server, mode="auto") as client:
|
||||
|
|
|
|||
|
|
@ -160,10 +160,9 @@ def get_client_ip() -> str:
|
|||
```
|
||||
|
||||
<Note>
|
||||
Both raise `RuntimeError` when called outside an HTTP context (e.g., STDIO transport).
|
||||
For background tasks created from an HTTP request, FastMCP restores a minimal request
|
||||
backed by the originating request's snapshotted headers. Use HTTP Headers if you need
|
||||
graceful fallback.
|
||||
Both raise `RuntimeError` when called outside an HTTP context (e.g., STDIO transport,
|
||||
or inside a background task — there is no live request object to reconstruct there).
|
||||
Use HTTP Headers below if you need graceful fallback, including inside background tasks.
|
||||
</Note>
|
||||
|
||||
### HTTP Headers
|
||||
|
|
|
|||
|
|
@ -686,9 +686,12 @@ class Client(
|
|||
self, elicitation_callback: ElicitationHandler
|
||||
) -> None:
|
||||
"""Set the elicitation callback for the client."""
|
||||
self._session_kwargs["elicitation_callback"] = create_elicitation_callback(
|
||||
elicitation_callback
|
||||
)
|
||||
self._elicitation_callback = create_elicitation_callback(elicitation_callback)
|
||||
self._session_kwargs["elicitation_callback"] = self._elicitation_callback
|
||||
# Rebuild internal extensions (e.g. the tasks extension) so a background
|
||||
# task's in-task input is answered through the newly-set handler, not the
|
||||
# one captured when the client was constructed.
|
||||
self._session_kwargs.update(self._build_extension_kwargs())
|
||||
|
||||
def is_connected(self) -> bool:
|
||||
"""Check if the client is currently connected."""
|
||||
|
|
|
|||
|
|
@ -2,15 +2,18 @@
|
|||
|
||||
Core ships the client wiring for opt-in extensions but no extension of its own.
|
||||
A companion package (``fastmcp-tasks``) provides an extension the ``Client``
|
||||
should register *automatically* — so an ordinary ``Client(url)`` transparently
|
||||
drives a server's background tasks without the caller passing anything. The
|
||||
package cannot reach into core's ``Client`` constructor, so core exposes this
|
||||
hook instead: the package registers a factory on import, and ``Client`` folds
|
||||
the factory's extension in alongside the user's own.
|
||||
folds in automatically once the package is imported — so a caller that uses
|
||||
tasks (importing ``fastmcp_tasks`` for ``call_tool_task``, or to register the
|
||||
server extension) gets transparent client task support without passing anything
|
||||
per ``Client``. The package cannot reach into core's ``Client`` constructor, so
|
||||
core exposes this hook instead: the package registers a factory on import, and
|
||||
``Client`` folds the factory's extension in alongside the user's own.
|
||||
|
||||
This mirrors the server-side ``set_background_context_factory`` hook: core
|
||||
declares the extension point, the tasks package fills it. With no package
|
||||
imported, the registry is empty and ``Client`` behaves exactly as before.
|
||||
declares the extension point, the tasks package fills it. Task support is
|
||||
opt-in — with ``fastmcp_tasks`` unimported the registry is empty and ``Client``
|
||||
behaves exactly as core alone, so a plain ``from fastmcp import Client`` never
|
||||
advertises the tasks capability and the server never runs its calls as tasks.
|
||||
|
||||
A factory receives the client's elicitation callback (so a task resolver can
|
||||
answer in-task input prompts) and returns a ``ClientExtension`` to register, or
|
||||
|
|
@ -54,8 +57,8 @@ def build_internal_client_extensions(
|
|||
"""Build the internal extensions to fold into a ``Client`` under construction.
|
||||
|
||||
Each registered factory is invoked with the client's elicitation callback;
|
||||
factories that return ``None`` contribute nothing. Empty when no package has
|
||||
registered a factory (plain core).
|
||||
factories that return ``None`` contribute nothing. Empty when no companion
|
||||
package has registered a factory (plain core, or ``fastmcp_tasks`` unimported).
|
||||
"""
|
||||
extensions: list[ClientExtension] = []
|
||||
for factory in _internal_client_extension_factories:
|
||||
|
|
|
|||
|
|
@ -45,7 +45,7 @@ async def analyze(dataset: str) -> str:
|
|||
`task=True` is a declaration of intent — this tool *may* run as a task — while the server, per the spec, decides per call whether to actually task it. Use `TaskConfig` for finer control:
|
||||
|
||||
```python
|
||||
from fastmcp_tasks import TaskConfig
|
||||
from fastmcp.utilities.tasks import TaskConfig
|
||||
|
||||
|
||||
@mcp.tool(task=TaskConfig(mode="required"))
|
||||
|
|
|
|||
|
|
@ -69,26 +69,43 @@ _TERMINAL_STATES = frozenset({"completed", "failed", "cancelled"})
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _send_get(session: ClientSession, task_id: str) -> ClientGetTaskResult:
|
||||
async def _send_get(
|
||||
session: ClientSession,
|
||||
task_id: str,
|
||||
read_timeout_seconds: float | None = None,
|
||||
) -> ClientGetTaskResult:
|
||||
"""Send `tasks/get` and parse the detailed task response."""
|
||||
request = GetTaskRequest(params=GetTaskRequestParams(task_id=task_id))
|
||||
return await session.send_request(request, ClientGetTaskResult)
|
||||
return await session.send_request(
|
||||
request, ClientGetTaskResult, request_read_timeout_seconds=read_timeout_seconds
|
||||
)
|
||||
|
||||
|
||||
async def _send_update(
|
||||
session: ClientSession, task_id: str, input_responses: dict[str, Any]
|
||||
session: ClientSession,
|
||||
task_id: str,
|
||||
input_responses: dict[str, Any],
|
||||
read_timeout_seconds: float | None = None,
|
||||
) -> None:
|
||||
"""Send `tasks/update` delivering the caller's answers to a parked task."""
|
||||
request = UpdateTaskRequest(
|
||||
params=UpdateTaskRequestParams(task_id=task_id, input_responses=input_responses)
|
||||
)
|
||||
await session.send_request(request, mcp_types.Result)
|
||||
await session.send_request(
|
||||
request, mcp_types.Result, request_read_timeout_seconds=read_timeout_seconds
|
||||
)
|
||||
|
||||
|
||||
async def _send_cancel(session: ClientSession, task_id: str) -> None:
|
||||
async def _send_cancel(
|
||||
session: ClientSession,
|
||||
task_id: str,
|
||||
read_timeout_seconds: float | None = None,
|
||||
) -> None:
|
||||
"""Send `tasks/cancel` to cooperatively cancel a task."""
|
||||
request = CancelTaskRequest(params=CancelTaskRequestParams(task_id=task_id))
|
||||
await session.send_request(request, mcp_types.Result)
|
||||
await session.send_request(
|
||||
request, mcp_types.Result, request_read_timeout_seconds=read_timeout_seconds
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -135,6 +152,7 @@ async def _answer_input_requests(
|
|||
task_id: str,
|
||||
input_requests: dict[str, Any],
|
||||
elicitation_callback: ElicitationFnT | None,
|
||||
read_timeout_seconds: float | None = None,
|
||||
) -> None:
|
||||
"""Answer a task's outstanding input requests, then deliver via `tasks/update`.
|
||||
|
||||
|
|
@ -170,7 +188,7 @@ async def _answer_input_requests(
|
|||
by_alias=True, mode="json", exclude_none=True
|
||||
)
|
||||
|
||||
await _send_update(session, task_id, responses)
|
||||
await _send_update(session, task_id, responses, read_timeout_seconds)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -182,22 +200,29 @@ async def _drive_to_terminal(
|
|||
session: ClientSession,
|
||||
task_id: str,
|
||||
elicitation_callback: ElicitationFnT | None,
|
||||
read_timeout_seconds: float | None = None,
|
||||
) -> ClientGetTaskResult:
|
||||
"""Poll `tasks/get` until the task reaches a terminal state.
|
||||
|
||||
`working` sleeps and polls again; `input_required` answers the outstanding
|
||||
requests through the elicitation handler and re-enters; a terminal state
|
||||
(completed / failed / cancelled) is returned. Shared by the transparent
|
||||
resolver and `ToolTask.result()`.
|
||||
resolver and `ToolTask.result()`. `read_timeout_seconds`, when set, bounds
|
||||
each `tasks/get`/`tasks/update` request so a stalled poll can't outlast the
|
||||
per-call timeout the synchronous path would have honored.
|
||||
"""
|
||||
backoff = MIN_POLL_INTERVAL
|
||||
while True:
|
||||
current = await _send_get(session, task_id)
|
||||
current = await _send_get(session, task_id, read_timeout_seconds)
|
||||
if current.status in _TERMINAL_STATES:
|
||||
return current
|
||||
if current.status == "input_required":
|
||||
await _answer_input_requests(
|
||||
session, task_id, current.input_requests or {}, elicitation_callback
|
||||
session,
|
||||
task_id,
|
||||
current.input_requests or {},
|
||||
elicitation_callback,
|
||||
read_timeout_seconds,
|
||||
)
|
||||
backoff = MIN_POLL_INTERVAL
|
||||
continue
|
||||
|
|
@ -266,7 +291,10 @@ class TasksClientExtension(ClientExtension):
|
|||
schema-valid so the SDK's output-schema revalidation passes.
|
||||
"""
|
||||
final = await _drive_to_terminal(
|
||||
ctx.session, create_result.task_id, self._elicitation_callback
|
||||
ctx.session,
|
||||
create_result.task_id,
|
||||
self._elicitation_callback,
|
||||
ctx.read_timeout_seconds,
|
||||
)
|
||||
if final.status == "completed":
|
||||
return _inlined_call_tool_result(final.result)
|
||||
|
|
@ -356,13 +384,16 @@ class ToolTask:
|
|||
return current
|
||||
elif current.status in _TERMINAL_STATES:
|
||||
return current
|
||||
if loop.time() >= deadline:
|
||||
remaining = deadline - loop.time()
|
||||
if remaining <= 0:
|
||||
raise TimeoutError(
|
||||
f"Task {self.task_id} did not reach "
|
||||
f"{state or 'a terminal state'} within {timeout}s"
|
||||
)
|
||||
delay, backoff = _next_poll_delay(current.poll_interval_ms, backoff)
|
||||
await asyncio.sleep(delay)
|
||||
# Never sleep past the deadline, so `wait` returns on time rather
|
||||
# than up to one poll interval late.
|
||||
await asyncio.sleep(min(delay, remaining))
|
||||
|
||||
async def result(self) -> FastMCPCallToolResult:
|
||||
"""Drive the task to completion and return its parsed result.
|
||||
|
|
|
|||
|
|
@ -293,7 +293,12 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None:
|
|||
)
|
||||
if raw is None:
|
||||
return
|
||||
_remember_snapshot(task_id, TaskContextSnapshot.from_json(raw))
|
||||
snapshot = TaskContextSnapshot.from_json(raw)
|
||||
_remember_snapshot(task_id, snapshot)
|
||||
# Restore the ambient request context (auth token, headers) so core's
|
||||
# get_access_token()/get_http_headers() see the submitting caller inside
|
||||
# the worker, exactly as a normal request would.
|
||||
_apply_snapshot_to_context(snapshot)
|
||||
except Exception:
|
||||
_logger.warning("Failed to restore task snapshot for %s", key, exc_info=True)
|
||||
|
||||
|
|
@ -402,6 +407,51 @@ def resolve_worker_server() -> FastMCP | None:
|
|||
return get_task_server(task_info.task_id)
|
||||
|
||||
|
||||
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
|
||||
``restore_task_snapshot`` (a Docket dependency), whose context vars propagate
|
||||
to the tool the same way the snapshot var already does.
|
||||
"""
|
||||
if snapshot.access_token_json is not None:
|
||||
from mcp.server.auth.middleware.auth_context import auth_context_var
|
||||
from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
||||
|
||||
from fastmcp.server.auth import AccessToken
|
||||
|
||||
token = AccessToken.model_validate_json(snapshot.access_token_json)
|
||||
auth_context_var.set(AuthenticatedUser(token))
|
||||
|
||||
if snapshot.http_headers:
|
||||
from starlette.requests import Request
|
||||
|
||||
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()
|
||||
],
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def make_task_context() -> Context | None:
|
||||
"""Build and enter a worker ``Context`` for the current background task.
|
||||
|
||||
|
|
|
|||
|
|
@ -1,15 +1,21 @@
|
|||
"""FastMCP tasks CLI for Docket task management."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from typing import Annotated
|
||||
from typing import TYPE_CHECKING, Annotated
|
||||
|
||||
import cyclopts
|
||||
from rich.console import Console
|
||||
|
||||
from fastmcp.utilities.cli import load_and_merge_config
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
from fastmcp_tasks.settings import docket_settings
|
||||
from fastmcp.utilities.tasks import TASKS_EXTENSION_ID
|
||||
from fastmcp_tasks.settings import DocketSettings
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from fastmcp.server.server import FastMCP
|
||||
|
||||
logger = get_logger("cli.tasks")
|
||||
console = Console()
|
||||
|
|
@ -20,7 +26,32 @@ tasks_app = cyclopts.App(
|
|||
)
|
||||
|
||||
|
||||
def check_distributed_backend() -> None:
|
||||
def resolve_docket_settings(server: FastMCP) -> DocketSettings:
|
||||
"""The effective Docket settings for `server`'s registered tasks extension.
|
||||
|
||||
Reads the *registered* `TasksExtension`'s resolved settings, not the
|
||||
env-only module-level default: a server that configures
|
||||
`TasksExtension(url="redis://...")` in code has settings the environment
|
||||
alone cannot see, and checking those defaults instead would report the
|
||||
wrong backend (see #4603 review — the CLI checked before the server, and
|
||||
therefore the extension, was even loaded).
|
||||
"""
|
||||
extension = server._extensions.get(TASKS_EXTENSION_ID)
|
||||
if extension is None:
|
||||
console.print(
|
||||
f"[bold red]✗ No tasks extension registered[/bold red]\n\n"
|
||||
f"[cyan]{server.name}[/cyan] has no `TasksExtension` registered "
|
||||
"(`mcp.add_extension(TasksExtension())`), so there is nothing for "
|
||||
"this worker to serve."
|
||||
)
|
||||
sys.exit(1)
|
||||
from fastmcp_tasks.extension import TasksExtension
|
||||
|
||||
assert isinstance(extension, TasksExtension)
|
||||
return extension.docket_settings
|
||||
|
||||
|
||||
def check_distributed_backend(settings: DocketSettings) -> None:
|
||||
"""Check if Docket is configured with a distributed backend.
|
||||
|
||||
The CLI worker runs as a separate process, so it needs Redis/Valkey
|
||||
|
|
@ -29,10 +60,8 @@ def check_distributed_backend() -> None:
|
|||
Raises:
|
||||
SystemExit: If using memory:// URL
|
||||
"""
|
||||
docket_url = docket_settings.url
|
||||
|
||||
# Check for memory:// URL and provide helpful error
|
||||
if docket_url.startswith("memory://"):
|
||||
if settings.url.startswith("memory://"):
|
||||
console.print(
|
||||
"[bold red]✗ In-memory backend not supported by CLI[/bold red]\n\n"
|
||||
"Your Docket configuration uses an in-memory backend (memory://) which\n"
|
||||
|
|
@ -75,8 +104,6 @@ def worker(
|
|||
fastmcp tasks worker server.py
|
||||
fastmcp tasks worker examples/tasks/server.py
|
||||
"""
|
||||
check_distributed_backend()
|
||||
|
||||
# Load server to get task functions
|
||||
try:
|
||||
config, _resolved_spec = load_and_merge_config(server_spec)
|
||||
|
|
@ -86,15 +113,21 @@ def worker(
|
|||
# Load the server
|
||||
server = asyncio.run(config.source.load_server())
|
||||
|
||||
# Validate against the server's actual registered extension, not an
|
||||
# env-only guess — a constructor-configured Redis URL isn't visible
|
||||
# until the server (and its extension) has loaded.
|
||||
settings = resolve_docket_settings(server)
|
||||
check_distributed_backend(settings)
|
||||
|
||||
async def run_worker():
|
||||
"""Enter server lifespan and camp forever."""
|
||||
async with server._lifespan_manager():
|
||||
console.print(
|
||||
f"[bold green]✓[/bold green] Starting worker for [cyan]{server.name}[/cyan]"
|
||||
)
|
||||
console.print(f" Docket: {docket_settings.name}")
|
||||
console.print(f" Backend: {docket_settings.url}")
|
||||
console.print(f" Concurrency: {docket_settings.concurrency}")
|
||||
console.print(f" Docket: {settings.name}")
|
||||
console.print(f" Backend: {settings.url}")
|
||||
console.print(f" Concurrency: {settings.concurrency}")
|
||||
|
||||
# Server's lifespan has started its worker - just camp here forever
|
||||
while True:
|
||||
|
|
|
|||
|
|
@ -1,27 +1,48 @@
|
|||
"""Tests for the fastmcp tasks CLI."""
|
||||
|
||||
import pytest
|
||||
from fastmcp_tasks.settings import docket_settings
|
||||
from fastmcp_tasks.worker_cli import check_distributed_backend, tasks_app
|
||||
from fastmcp_tasks.settings import DocketSettings
|
||||
from fastmcp_tasks.worker_cli import (
|
||||
check_distributed_backend,
|
||||
resolve_docket_settings,
|
||||
tasks_app,
|
||||
)
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp_tasks import TasksExtension
|
||||
|
||||
|
||||
class TestResolveDocketSettings:
|
||||
"""`resolve_docket_settings` reads the server's *registered* extension."""
|
||||
|
||||
def test_reads_the_registered_extensions_settings(self):
|
||||
"""The constructor-configured URL is visible without any env var."""
|
||||
mcp = FastMCP("t")
|
||||
mcp.add_extension(TasksExtension(url="redis://example:6379/0"))
|
||||
settings = resolve_docket_settings(mcp)
|
||||
assert settings.url == "redis://example:6379/0"
|
||||
|
||||
def test_exits_when_no_tasks_extension_registered(self):
|
||||
"""A server with no TasksExtension has nothing for the CLI to serve."""
|
||||
mcp = FastMCP("t")
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
resolve_docket_settings(mcp)
|
||||
assert exc_info.value.code == 1
|
||||
|
||||
|
||||
class TestCheckDistributedBackend:
|
||||
"""Test the distributed backend checker function."""
|
||||
|
||||
def test_succeeds_with_redis_url(self, monkeypatch: pytest.MonkeyPatch):
|
||||
def test_succeeds_with_redis_url(self):
|
||||
"""Test that it succeeds with Redis URL."""
|
||||
# Docket settings moved to `fastmcp_tasks.settings.DocketSettings`
|
||||
# (env prefix `FASTMCP_DOCKET_`), so patch the settings object directly.
|
||||
monkeypatch.setattr(docket_settings, "url", "redis://localhost:6379/0")
|
||||
check_distributed_backend()
|
||||
settings = DocketSettings(url="redis://localhost:6379/0")
|
||||
check_distributed_backend(settings)
|
||||
|
||||
def test_exits_with_helpful_error_for_memory_url(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
def test_exits_with_helpful_error_for_memory_url(self):
|
||||
"""Test that it exits with helpful error for memory:// URLs."""
|
||||
monkeypatch.setattr(docket_settings, "url", "memory://test-123")
|
||||
settings = DocketSettings(url="memory://test-123")
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
check_distributed_backend()
|
||||
check_distributed_backend(settings)
|
||||
|
||||
assert isinstance(exc_info.value, SystemExit)
|
||||
assert exc_info.value.code == 1
|
||||
|
|
|
|||
|
|
@ -154,3 +154,23 @@ async def test_in_task_input_without_handler_errors(guard_server: FastMCP):
|
|||
async with Client(guard_server, mode="auto") as client:
|
||||
with pytest.raises(ToolError, match="no elicitation handler"):
|
||||
await client.call_tool("plan_dinner", {})
|
||||
|
||||
|
||||
async def test_in_task_input_answered_by_handler_set_after_construction(
|
||||
guard_server: FastMCP,
|
||||
):
|
||||
"""An elicitation handler set via set_elicitation_callback reaches in-task input.
|
||||
|
||||
The tasks client extension is built at construction; set_elicitation_callback
|
||||
must rebuild it so a later-configured handler still answers a task's input.
|
||||
"""
|
||||
|
||||
async def handle_elicitation(message, response_type, params, context):
|
||||
return DinnerPrefs(cuisine="Thai", vegetarian=True)
|
||||
|
||||
client = Client(guard_server, mode="auto")
|
||||
client.set_elicitation_callback(handle_elicitation)
|
||||
async with client:
|
||||
result = await client.call_tool("plan_dinner", {})
|
||||
|
||||
assert result.data == "Tonight: a vegetarian Thai dinner!"
|
||||
|
|
|
|||
|
|
@ -10,10 +10,12 @@ the edge cases around non-fastmcp keys and failed restores.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextvars
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastmcp_tasks.context import (
|
||||
TaskContextSnapshot,
|
||||
_apply_snapshot_to_context,
|
||||
_recall_snapshot,
|
||||
get_task_context,
|
||||
restore_task_snapshot,
|
||||
|
|
@ -23,7 +25,7 @@ from mcp.server.auth.middleware.bearer_auth import AuthenticatedUser
|
|||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.server.auth import AccessToken
|
||||
from fastmcp.server.dependencies import get_access_token
|
||||
from fastmcp.server.dependencies import get_access_token, get_http_headers
|
||||
from fastmcp_tasks import TasksExtension
|
||||
from tests.tasks.task_helpers import (
|
||||
running_task_server,
|
||||
|
|
@ -80,6 +82,42 @@ async def test_get_access_token_in_bg_task_without_context_dep():
|
|||
assert final.result["structuredContent"] == {"result": "jwt-3897"}
|
||||
|
||||
|
||||
def test_apply_snapshot_restores_auth_and_headers_in_clean_context():
|
||||
"""The cross-process path: with nothing inherited, the snapshot alone makes
|
||||
get_access_token()/get_http_headers() see the submitting caller.
|
||||
|
||||
A Redis-backed worker runs in a separate process and inherits none of the
|
||||
submitter's context vars, so contextvar inheritance (which carries the token
|
||||
on the same-process memory:// path) cannot help. Running in a fresh
|
||||
`copy_context()` with no auth/request bound simulates that worker: only
|
||||
`_apply_snapshot_to_context` populating the ambient vars makes the token and
|
||||
headers reachable.
|
||||
"""
|
||||
token = AccessToken(
|
||||
token="jwt-remote",
|
||||
client_id="remote-client",
|
||||
scopes=["read"],
|
||||
claims={"sub": "user-y"},
|
||||
)
|
||||
snapshot = TaskContextSnapshot(
|
||||
access_token_json=token.model_dump_json(),
|
||||
http_headers={"x-trace-id": "abc123"},
|
||||
)
|
||||
|
||||
def run_in_clean_worker_context() -> None:
|
||||
# Nothing bound here — no inheritance to fall back on.
|
||||
assert get_access_token() is None
|
||||
assert get_http_headers() == {}
|
||||
_apply_snapshot_to_context(snapshot)
|
||||
restored = get_access_token()
|
||||
assert restored is not None
|
||||
assert restored.token == "jwt-remote"
|
||||
assert restored.client_id == "remote-client"
|
||||
assert get_http_headers()["x-trace-id"] == "abc123"
|
||||
|
||||
contextvars.copy_context().run(run_in_clean_worker_context)
|
||||
|
||||
|
||||
async def test_restore_failure_is_nonfatal():
|
||||
"""If deserialization blows up, the task still runs to completion and
|
||||
the snapshot cache stays empty."""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue