mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 07:09:11 +02:00
Skip expired snapshot tokens; bound task wait polls by deadline
A queued task can outlive its submitter's token expiry: install the snapshot token only if still valid, matching the SDK bearer check, so a delayed task never runs under credentials a live request would reject. ToolTask.wait now bounds each tasks/get by the remaining deadline so a stalled poll cannot block past the caller's timeout.
This commit is contained in:
parent
1c7ade215b
commit
110943fc61
3 changed files with 49 additions and 2 deletions
|
|
@ -378,7 +378,18 @@ class ToolTask:
|
|||
deadline = loop.time() + timeout
|
||||
backoff = MIN_POLL_INTERVAL
|
||||
while True:
|
||||
current = await self.status()
|
||||
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"
|
||||
)
|
||||
# Bound the request itself by the remaining deadline: a stalled
|
||||
# `tasks/get` must not block past the caller's timeout waiting for
|
||||
# the session-wide default before the deadline is next checked.
|
||||
current = await _send_get(
|
||||
self._session, self.task_id, read_timeout_seconds=remaining
|
||||
)
|
||||
if state is not None:
|
||||
if current.status == state:
|
||||
return current
|
||||
|
|
|
|||
|
|
@ -420,13 +420,20 @@ def _apply_snapshot_to_context(snapshot: TaskContextSnapshot) -> None:
|
|||
to the tool the same way the snapshot var already does.
|
||||
"""
|
||||
if snapshot.access_token_json is not None:
|
||||
import time
|
||||
|
||||
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))
|
||||
# A task may sit queued past its submitter's token expiry. Install it
|
||||
# only if still valid — mirroring the SDK's bearer check — so a delayed
|
||||
# task never runs under credentials a live request would reject (401).
|
||||
# An expired token leaves the worker unauthenticated, the honest state.
|
||||
if token.expires_at is None or token.expires_at >= int(time.time()):
|
||||
auth_context_var.set(AuthenticatedUser(token))
|
||||
|
||||
if snapshot.http_headers:
|
||||
from starlette.requests import Request
|
||||
|
|
|
|||
|
|
@ -118,6 +118,35 @@ def test_apply_snapshot_restores_auth_and_headers_in_clean_context():
|
|||
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.
|
||||
|
||||
A task may sit queued past its submitter's token expiry. A live request with
|
||||
an expired bearer token is rejected (401), so restoring one as authenticated
|
||||
would let a delayed task run under credentials that should now be treated as
|
||||
unauthenticated. The headers still restore — only the auth token is dropped.
|
||||
"""
|
||||
expired = AccessToken(
|
||||
token="jwt-expired",
|
||||
client_id="remote-client",
|
||||
scopes=["read"],
|
||||
expires_at=1, # 1970 — long past
|
||||
)
|
||||
snapshot = TaskContextSnapshot(
|
||||
access_token_json=expired.model_dump_json(),
|
||||
http_headers={"x-trace-id": "abc123"},
|
||||
)
|
||||
|
||||
def run_in_clean_worker_context() -> None:
|
||||
assert get_access_token() is None
|
||||
_apply_snapshot_to_context(snapshot)
|
||||
assert get_access_token() is None
|
||||
# Non-auth context still restores independently of the token.
|
||||
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