Commit graph

7 commits

Author SHA1 Message Date
Chris Guidry
d41bcb2c9e
Unify background task context forwarding, fix concurrent dependency bugs (#3710)
* Unify background task context forwarding and fix concurrent dependency 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) <noreply@anthropic.com>

* Address code review feedback

- _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) <noreply@anthropic.com>

* Use pop-on-access transfer buffer instead of bounded LRU cache for snapshots

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) <noreply@anthropic.com>

* Drop in-memory transfer buffer, use sync Redis for all backends

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) <noreply@anthropic.com>

* Move snapshot operations to TaskContextSnapshot methods

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) <noreply@anthropic.com>

* Trim implementation details from register_with_docket docstrings

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Clarify docket lookup comment in submit_to_docket

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

* Restore docket/worker ContextVar bridge in Context.__aenter__

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) <noreply@anthropic.com>

* Key snapshot cache by task_id to prevent cross-task context leakage

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) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-03 10:48:27 -04:00
Guillaume FORTAINE
40d3190317
fix: propagate origin_request_id to background task workers (#3175)
* Fix background Context request correlation

* Make OptionalCurrentContext type-safe

Refactor OptionalCurrentContext to wrap CurrentContext instead of overriding __aenter__ with a wider return type. Adds a background-task origin_request_id round-trip test and applies ruff formatting.
2026-02-22 11:16:30 -05:00
Guillaume FORTAINE
263e0bf6e0
fix: snapshot access token for background tasks (#3095) (#3138)
Co-authored-by: cristiangreco94 <cristiangreco94@users.noreply.github.com>
2026-02-11 10:20:14 -05:00
Chris Guidry
361eb08f42 Relay task elicitation through standard MCP protocol
When a background task calls ctx.elicit(), the notification subscriber now
detects the input_required notification and sends a standard elicitation/create
request to the client via session.elicit(). The client's elicitation_handler
fires, and the relay pushes the response to Redis for the blocked worker.

This means clients can respond to background task elicitation using the same
elicitation_handler they'd use for any other elicitation — no need to interact
with Redis or call handle_task_input() directly.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-10 14:58:19 -05:00
Guillaume FORTAINE
3d665d48c0 fix: stabilize task notification integration tests
🤖 Generated with Codex
2026-02-08 15:20:22 +01:00
Guillaume FORTAINE
1f84e5055e test: rewrite as integration tests with zero mocks
Replace 1300+ lines of mock-heavy unit tests with 391 lines of integration
tests using real Client(mcp) connections and memory:// Docket backend.

- test_context_background_task.py: 17 tests covering report_progress delta
  tracking, elicitation flow, edge cases, and fail-fast on push failure
- test_notifications.py: 2 E2E tests for notification queue lifecycle
2026-02-08 01:59:07 +01:00
Guillaume FORTAINE
08974e50d9
feat(context): Add background task support for Context (SEP-1686) (#2905) 2026-02-02 19:28:42 -05:00