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:
Jeremiah Lowin 2026-07-23 08:12:42 -04:00
commit 110943fc61
No known key found for this signature in database
3 changed files with 49 additions and 2 deletions

View file

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

View file

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

View file

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