From 110943fc61033d8d3e2e18b427cf82e9f277874d Mon Sep 17 00:00:00 2001
From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com>
Date: Thu, 23 Jul 2026 08:12:42 -0400
Subject: [PATCH] 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.
---
 fastmcp_tasks/fastmcp_tasks/client.py       | 13 ++++++++-
 fastmcp_tasks/fastmcp_tasks/context.py      |  9 ++++++-
 tests/tasks/server/test_snapshot_restore.py | 29 +++++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/fastmcp_tasks/fastmcp_tasks/client.py b/fastmcp_tasks/fastmcp_tasks/client.py
index 749a9abd2..34012b3e7 100644
--- a/fastmcp_tasks/fastmcp_tasks/client.py
+++ b/fastmcp_tasks/fastmcp_tasks/client.py
@@ -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
diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py
index d59a85c5c..f48f974da 100644
--- a/fastmcp_tasks/fastmcp_tasks/context.py
+++ b/fastmcp_tasks/fastmcp_tasks/context.py
@@ -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
diff --git a/tests/tasks/server/test_snapshot_restore.py b/tests/tasks/server/test_snapshot_restore.py
index 9af2df0cd..6fe5014fc 100644
--- a/tests/tasks/server/test_snapshot_restore.py
+++ b/tests/tasks/server/test_snapshot_restore.py
@@ -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."""