fix: snapshot access token for background tasks (#3095)#3138
fix: snapshot access token for background tasks (#3095)#3138jlowin merged 1 commit intoPrefectHQ:mainfrom
Conversation
marvin-context-protocol
Bot
added
bug
Test Failure AnalysisSummary: The static analysis workflow failed due to code formatting issues detected by Root Cause: Two files have trailing blank lines that don't match the project's formatting standards:
Suggested Solution: Run
Note: There are also Detailed AnalysisFrom the CI logs: The specific changes needed: diff --git a/src/fastmcp/server/dependencies.py b/src/fastmcp/server/dependencies.py
index fabed02..a1c2084 100644
--- a/src/fastmcp/server/dependencies.py
+++ b/src/fastmcp/server/dependencies.py
@@ -1058,7 +1058,6 @@ def CurrentHeaders() -> dict[str, str]:
return cast(dict[str, str], _CurrentHeaders())
-
# --- Progress dependency ---
diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py
index 61473e1..400e4c3 100644
--- a/tests/server/tasks/test_context_background_task.py
+++ b/tests/server/tasks/test_context_background_task.py
@@ -452,4 +452,3 @@ class TestLifespanContextInBackgroundTasks:
ctx = Context(mcp, task_id="test-task")
assert ctx.request_context is None
assert ctx.lifespan_context == {}
-Related Files
|
chrisguidry
left a comment
chrisguidry
left a comment
There was a problem hiding this comment.
Looks great! Thank you both!
WalkthroughAdds Redis-backed snapshot and restoration of access tokens for background Docket workers: submit_to_docket now saves the current access token (if any) to Redis alongside task metadata; get_access_token and related context managers (_CurrentContext, _CurrentAccessToken) gain logic to restore that snapshot into a ContextVar when running in background/task contexts and to clean it up on exit. Also, lifespan_context now reads fastmcp._lifespan_result when request_context is None to surface lifespan data in background tasks. Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
No actionable comments were generated in the recent review. 🎉 Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e200d89ff5
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| if access_token is None: | ||
| task_token = _task_access_token.get() |
There was a problem hiding this comment.
get_access_token() now falls back to _task_access_token, but this ContextVar is only populated when _restore_task_access_token() runs via _CurrentContext or _CurrentAccessToken. A background task that calls get_access_token() directly and does not inject ctx: Context/CurrentAccessToken() never triggers restoration, so this branch still returns None even though submit_to_docket() persisted a token for the task.
Useful? React with 👍 / 👎.
| self._access_token_cv_token = await _restore_task_access_token( | ||
| task_info.session_id, task_info.task_id | ||
| ) | ||
| token = get_access_token() |
There was a problem hiding this comment.
_CurrentAccessToken.__aenter__ stores self._access_token_cv_token before confirming the restored token is usable; if the restored token is expired, get_access_token() returns None and __aenter__ raises. Because failed __aenter__ calls are not paired with __aexit__, this stale token remains on the reused dependency instance and a later invocation can call _task_access_token.reset() with a token from another context, raising ValueError during cleanup.
Useful? React with 👍 / 👎.
gfortaine
force-pushed
the
fix/access-token-background-tasks
branch
from
e200d89 to
8ee0b7a
Compare
Test Failure AnalysisSummary: Static analysis failed because detected a formatting issue - there's a trailing blank line at the end of that needs to be removed. Root Cause: The file has an extra blank line after line 442 (after the last assertion in the test). Ruff's formatter automatically removed it during the CI run, causing the static_analysis job to fail because files were modified by the hook. Suggested Solution: Remove the trailing blank line at the end of . You can fix this by:
Or simply remove the blank line manually at the end of the file. Detailed AnalysisFrom the workflow logs: The diff shows: diff --git a/tests/server/tasks/test_context_background_task.py b/tests/server/tasks/test_context_background_task.py
index 53b94aa..c7eb9e9 100644
--- a/tests/server/tasks/test_context_background_task.py
+++ b/tests/server/tasks/test_context_background_task.py
@@ -440,4 +440,3 @@ class TestLifespanContextInBackgroundTasks:
ctx = Context(mcp, task_id="test-task")
assert ctx.request_context is None
assert ctx.lifespan_context == {}
-The Related Files
|
|
@gfortaine I can't push to your branch but if you run |
Snapshot the current access token at task submission time and restore it in Docket workers via ContextVar. This makes get_access_token() return the caller's token inside background tasks. Changes: - handlers.py: store token in Redis at submit_to_docket() time - dependencies.py: add _restore_task_access_token() to read from Redis, add fallback in get_access_token() for _task_access_token ContextVar, wire restoration into _CurrentContext and _CurrentAccessToken - context.py: fall back to server lifespan_result when request_context is unavailable in background tasks - tests: integration tests using Client(mcp) + memory:// Docket Review feedback addressed (chrisguidry): - De-mocked test suite: replaced MagicMock/patch tests with real Client(mcp) + fakeredis integration tests - No try/finally in tests: ContextVar cleanup via cv_token pattern CodeRabbit feedback addressed: - Replaced silent except-pass with logger.warning(exc_info=True) - Dropped DocketDependency.docket internal API fallback - Hoisted datetime import to module level - Typed _access_token_cv_token as Token[AccessToken | None] | None Co-authored-by: cristiangreco94 <cristiangreco94@users.noreply.github.com>
gfortaine
force-pushed
the
fix/access-token-background-tasks
branch
from
8ee0b7a to
9c61a75
Compare
|
Thank you! |
3 participants
Summary
Makes
get_access_token()return the caller's token inside Docket background tasks. The token is snapshotted to Redis at submission time and restored into a ContextVar when the worker picks up the task.Also fixes
lifespan_contextreturning{}in background tasks by falling back to the server's lifespan result whenrequest_contextis unavailable.Closes #3095. Supersedes #3121.
Changes
handlers.py— snapshot token at submit timeget_access_token()insubmit_to_docket()and store the result in Redis alongside other task metadatadependencies.py— restore token in worker_task_access_tokenContextVar for the restored token_restore_task_access_token()reads from Redis and sets the ContextVarget_access_token()falls back to_task_access_tokenwhen HTTP request and SDK context var are both unavailable_CurrentContext.__aenter__()restores the token when entering background task context_CurrentAccessToken.__aenter__()restores the token as fallback when_CurrentContexthasn't run (noctx: Contextin signature)Nonecontext.py— lifespan fallbacklifespan_contextfalls back toserver._lifespan_resultwhenrequest_contextisNoneTests — integration, zero mocks
test_token_round_trips_through_background_task: E2E withClient(mcp)+memory://Docket — token injected via SDK auth context, verified inside workertest_no_token_when_unauthenticated: no auth →Nonein workertest_expired_token_returns_none,test_valid_token_with_future_expiry,test_token_without_expiry_always_valid: expiration edge casestest_lifespan_context_falls_back_to_server_result,test_lifespan_context_returns_empty_dict_when_no_lifespan: lifespan fallbackReview feedback addressed
@chrisguidry's comments on #3121:
MagicMock/patchtests with realClient(mcp)+ fakeredis integration tests (same pattern as feat: distributed notification queue + BLPOP elicitation for background tasks #2906 andtest_task_elicitation_relay.py)try/finallyblocks in tests — ContextVar cleanup viacv_tokenpatterntest_lifespan_context_still_uses_request_context_when_availabletest that relied onpatch.objectCodeRabbit feedback:
except Exception: passwithlogger.warning(..., exc_info=True)(fixes S110/BLE001)DocketDependency.docket.get()fallback — uses only_current_docketContextVar (no unstable internal API)from datetime import datetime, timezoneto module-level imports_access_token_cv_tokenasToken[AccessToken | None] | Noneinstead ofAnyCredit
Based on the initial work by @cristiangreco94 in #3121. The core design (snapshot at submit → restore in worker → ContextVar fallback) is his — this PR rebases onto current
main, resolves conflicts from #2906/#3136, rewrites tests as integration, and addresses review feedback.