Skip to content

Unify background task context forwarding, fix concurrent dependency bugs#3710

Merged
jlowin merged 10 commits intomainfrom
unified-task-context
Apr 3, 2026
Merged

Unify background task context forwarding, fix concurrent dependency bugs#3710
jlowin merged 10 commits intomainfrom
unified-task-context

Conversation

@chrisguidry
Copy link
Copy Markdown
Collaborator

@chrisguidry chrisguidry commented Mar 30, 2026

We've been getting a steady trickle of edge-case reports around background tasks and contextual dependencies over the last few months (#3654, #3656, #3569). Each one pointed at a different symptom, but they all traced back to the same area: the way context is negotiated between the "frontend" server and Docket workers was grown piecemeal, with each new piece of context (access tokens, HTTP headers, origin request IDs) getting its own Redis key, its own restore function, and its own ContextVar. This made it hard to reason about what state was available where, and the shared-instance Dependency pattern made concurrent tasks stomp on each other's cleanup state.

Rather than patch each report individually, this takes a step back and reworks the whole thing as a single unified system.

Dependency subclasses (_CurrentContext, Progress, _CurrentAccessToken, etc.) are now stateless factories — __aenter__ returns a fresh per-invocation object, so concurrent tasks never share mutable state. The three individual context-snapshot Redis keys are collapsed into a single TaskContextSnapshot stored as one JSON key per task, and the three _restore_task_* functions and two ContextVars they populated are gone.

Sync functions like get_http_request() and get_access_token() now find the snapshot transparently in background tasks via a 3-tier sync fallback chain: ContextVar → in-memory dict → sync Redis GET. No function wrapping needed — _wrap_for_task_http_headers is deleted. FunctionTool registers its raw function with Docket so Docket sees and resolves all dependencies, including Docket-native ones like Retry and Timeout.

@mcp.tool(task=True)
async def my_tool(ctx: Context) -> str:
    # All of these now work identically in foreground and background,
    # regardless of which dependencies you declare
    request = get_http_request()
    token = get_access_token()
    headers = get_http_headers()
    await ctx.report_progress(50, 100, "halfway")
    return "done"

Progress is also simplified — it reads ExecutionProgress directly from Docket's current_execution ContextVar instead of creating and manually entering a DocketProgress wrapper around it.

Also fixes ProxyTool.from_mcp_tool() dropping execution.taskSupport metadata from remote tools.

Note: The Redis snapshot key format changed from three individual keys (access_token, http_headers, origin_request_id) to a single snapshot key. Tasks submitted by the old code won't have their context available to workers running the new code during a rolling deploy. This only affects the brief transition window and only loses auth/header context (not task execution itself).

Closes #3654
Closes #3656
Closes #3569

🤖 Generated with Claude Code

@marvin-context-protocol marvin-context-protocol Bot added bug Something isn't working. Reports of errors, unexpected behavior, or broken functionality. server Related to FastMCP server implementation or server-side functionality. labels Mar 30, 2026
@chrisguidry chrisguidry marked this pull request as ready for review March 30, 2026 20:17
chrisguidry and others added 8 commits March 30, 2026 16:17
…y bugs

We've been getting a steady trickle of edge-case reports around background tasks
and contextual dependencies over the last few months (#3654, #3656, #3569). Each
one pointed at a different symptom, but they all traced back to the same area:
the way context is negotiated between the "frontend" server and Docket workers
was grown piecemeal, with each new piece of context (access tokens, HTTP headers,
origin request IDs) getting its own Redis key, its own restore function, and its
own ContextVar. This made it hard to reason about what state was available where,
and the shared-instance Dependency pattern made concurrent tasks stomp on each
other's cleanup state.

This takes a step back and reworks the whole thing as a single unified system:

- Dependency subclasses (_CurrentContext, Progress, _CurrentAccessToken, etc.)
  are now stateless factories — __aenter__ returns a fresh per-invocation
  object, so concurrent tasks never share mutable state. Fixes #3654, #3656.

- The three individual context-snapshot Redis keys (access_token, http_headers,
  origin_request_id) are collapsed into a single TaskContextSnapshot stored as
  one JSON key per task. The three _restore_task_* functions and two ContextVars
  they populated are gone.

- Sync functions like get_http_request() and get_access_token() now find the
  snapshot transparently in background tasks via a 3-tier sync fallback:
  ContextVar (set by _CurrentContext for functions with deps) → in-memory dict
  (same-process workers) → sync Redis GET (out-of-process workers). No function
  wrapping needed.

- The _wrap_for_task_http_headers hack is deleted. FunctionTool registers its
  raw function with Docket so Docket sees and resolves ALL dependencies,
  including Docket-native ones like Retry and Timeout.

- ProxyTool.from_mcp_tool() now propagates execution.taskSupport metadata from
  remote tools. Fixes #3569.

- Removed redundant _current_docket/_current_worker ContextVar management from
  Context.__aenter__/__aexit__ (they're only set in the lifespan now).

Closes #3654
Closes #3656
Closes #3569

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
- _OptionalCurrentContext: guard __aexit__ against cleaning up contexts it
  didn't create (check is_background_task before delegating)
- Narrow except clauses in snapshot loading (OSError, JSONDecodeError, etc.
  instead of bare Exception)
- Fix docstrings on register_with_docket for resources/prompts/templates
- Simplify Progress: read ExecutionProgress directly from current_execution
  instead of creating and manually entering a DocketProgress wrapper

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
…apshots

The in-memory snapshot dict is a transfer mechanism, not a cache. Entries go
in at submission and come out at the worker's first access. Using pop instead
of get means the dict only holds entries during the brief submission-to-execution
window, bounded by task concurrency (~10) rather than a 10,000-entry LRU limit.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Instead of maintaining an in-memory dict to bridge the async/sync gap, use
a sync Redis client directly. For memory:// backends (fakeredis), shares the
same FakeServer instance via docket._redis.get_memory_server() so data written
by the async Docket client is visible to sync reads. For real Redis, creates a
standard sync connection. No in-process state to manage at all.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
capture(), from_json(), to_json(), save() are now classmethod/instance methods
on the dataclass instead of free functions. Deduplicates JSON parsing that was
copy-pasted between the async and sync load paths.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Servers that own the Docket (the parent) re-set _current_docket/_current_worker
from their instance attributes when entering a Context. Mounted children skip
this (their _docket is None), so they inherit the parent's value. This is needed
for ASGI deployments where ContextVars set during the lifespan don't propagate
to request handlers.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
@chrisguidry chrisguidry force-pushed the unified-task-context branch from 68e52cb to ed25d39 Compare March 30, 2026 20:19
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 68e52cb8e6

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread src/fastmcp/server/dependencies.py Outdated
Comment thread src/fastmcp/server/dependencies.py
Docket workers may reuse the same asyncio context for sequential tasks.
The ContextVar cache now stores (task_id, snapshot) tuples so stale entries
from previous tasks are automatically ignored.

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
Copy link
Copy Markdown
Member

@jlowin jlowin left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clean refactor — stateless dependency factories and the unified snapshot are the right fix for the concurrent stomping bugs. Nicely done.

@jlowin jlowin merged commit d41bcb2 into main Apr 3, 2026
9 checks passed
@jlowin jlowin deleted the unified-task-context branch April 3, 2026 14:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working. Reports of errors, unexpected behavior, or broken functionality. server Related to FastMCP server implementation or server-side functionality.

Projects

None yet

2 participants