Reset contextvar tokens in resolved_dependencies

The resolved_dependencies function was setting contextvars without
capturing reset tokens, which left stale references after the context
manager exited. This could cause issues in reentrant or sequential
calls, like "stack is closed" errors when the AsyncExitStack was
already closed but still referenced.

Now we capture tokens from all contextvar.set() calls and reset them
in finally blocks to restore prior state. This follows the same
pattern suggested in the fastmcp PR review:
jlowin/fastmcp#2318

Added tests that verify contextvars are properly isolated and cleaned
up between task executions, including a test for reentrant calls to
resolved_dependencies that would previously fail with stale context.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: chrisguidry

src/docket/dependencies.py

@@ -697,38 +697,50 @@ def __init__(self, parameter: str, error: Exception) -> None:
 async def resolved_dependencies(
     worker: "Worker", execution: Execution
 ) -> AsyncGenerator[dict[str, Any], None]:
-    # Set context variables once at the beginning
-    Dependency.docket.set(worker.docket)
-    Dependency.worker.set(worker)
-    Dependency.execution.set(execution)
-
-    _Depends.cache.set({})
-
-    async with AsyncExitStack() as stack:
-        _Depends.stack.set(stack)
-
-        arguments: dict[str, Any] = {}
-
-        parameters = get_dependency_parameters(execution.function)
-        for parameter, dependency in parameters.items():
-            kwargs = execution.kwargs
-            if parameter in kwargs:
-                arguments[parameter] = kwargs[parameter]
-                continue
-
-            # Special case for TaskArguments, they are "magical" and infer the parameter
-            # they refer to from the parameter name (unless otherwise specified).  At
-            # the top-level task function call, it doesn't make sense to specify one
-            # _without_ a parameter name, so we'll call that a failed dependency.
-            if isinstance(dependency, _TaskArgument) and not dependency.parameter:
-                arguments[parameter] = FailedDependency(
-                    parameter, ValueError("No parameter name specified")
-                )
-                continue
-
+    # Capture tokens for all contextvar sets to ensure proper cleanup
+    docket_token = Dependency.docket.set(worker.docket)
+    worker_token = Dependency.worker.set(worker)
+    execution_token = Dependency.execution.set(execution)
+    cache_token = _Depends.cache.set({})
+
+    try:
+        async with AsyncExitStack() as stack:
+            stack_token = _Depends.stack.set(stack)
             try:
-                arguments[parameter] = await stack.enter_async_context(dependency)
-            except Exception as error:
-                arguments[parameter] = FailedDependency(parameter, error)
-
-        yield arguments
+                arguments: dict[str, Any] = {}
+
+                parameters = get_dependency_parameters(execution.function)
+                for parameter, dependency in parameters.items():
+                    kwargs = execution.kwargs
+                    if parameter in kwargs:
+                        arguments[parameter] = kwargs[parameter]
+                        continue
+
+                    # Special case for TaskArguments, they are "magical" and infer the parameter
+                    # they refer to from the parameter name (unless otherwise specified).  At
+                    # the top-level task function call, it doesn't make sense to specify one
+                    # _without_ a parameter name, so we'll call that a failed dependency.
+                    if (
+                        isinstance(dependency, _TaskArgument)
+                        and not dependency.parameter
+                    ):
+                        arguments[parameter] = FailedDependency(
+                            parameter, ValueError("No parameter name specified")
+                        )
+                        continue
+
+                    try:
+                        arguments[parameter] = await stack.enter_async_context(
+                            dependency
+                        )
+                    except Exception as error:
+                        arguments[parameter] = FailedDependency(parameter, error)
+
+                yield arguments
+            finally:
+                _Depends.stack.reset(stack_token)
+    finally:
+        _Depends.cache.reset(cache_token)
+        Dependency.execution.reset(execution_token)
+        Dependency.worker.reset(worker_token)
+        Dependency.docket.reset(docket_token)

tests/test_dependencies.py

@@ -5,7 +5,16 @@
 import pytest
 
 from docket import CurrentDocket, CurrentWorker, Docket, Worker
-from docket.dependencies import Depends, ExponentialRetry, Retry, TaskArgument
+from docket.dependencies import (
+    Depends,
+    Dependency,
+    ExponentialRetry,
+    Retry,
+    TaskArgument,
+    _Depends,  # type: ignore[attr-defined]
+    resolved_dependencies,
+)
+from docket.execution import Execution
 
 
 async def test_dependencies_may_be_duplicated(docket: Docket, worker: Worker):
@@ -449,3 +458,216 @@ async def dependent_task(result: int = Depends(sync_adder)):
     await worker.run_until_finished()
 
     assert called
+
+
+async def test_contextvar_isolation_between_tasks(docket: Docket, worker: Worker):
+    """Contextvars should be isolated between sequential task executions"""
+    executions_seen: list[tuple[str, Execution]] = []
+
+    async def first_task(a: str):
+        # Capture the execution context during first task
+        execution = Dependency.execution.get()
+        executions_seen.append(("first", execution))
+        assert a == "first"
+
+    async def second_task(b: str):
+        # Capture the execution context during second task
+        execution = Dependency.execution.get()
+        executions_seen.append(("second", execution))
+        assert b == "second"
+
+        # The execution should be different from the first task
+        first_execution = executions_seen[0][1]
+        assert execution is not first_execution
+        assert execution.kwargs["b"] == "second"
+        assert first_execution.kwargs["a"] == "first"
+
+    await docket.add(first_task)(a="first")
+    await docket.add(second_task)(b="second")
+    await worker.run_until_finished()
+
+    assert len(executions_seen) == 2
+    assert executions_seen[0][0] == "first"
+    assert executions_seen[1][0] == "second"
+
+
+async def test_contextvar_cleanup_after_task(docket: Docket, worker: Worker):
+    """Contextvars should be reset after task execution completes"""
+    captured_stack = None
+    captured_cache = None
+
+    async def capture_task():
+        nonlocal captured_stack, captured_cache
+        # Capture references during task execution
+        captured_stack = _Depends.stack.get()
+        captured_cache = _Depends.cache.get()
+
+    await docket.add(capture_task)()
+    await worker.run_until_finished()
+
+    # After the task completes, the contextvars should be reset
+    # Attempting to get them should raise LookupError
+    with pytest.raises(LookupError):
+        _Depends.stack.get()
+
+    with pytest.raises(LookupError):
+        _Depends.cache.get()
+
+    with pytest.raises(LookupError):
+        Dependency.execution.get()
+
+    with pytest.raises(LookupError):
+        Dependency.worker.get()
+
+    with pytest.raises(LookupError):
+        Dependency.docket.get()
+
+
+async def test_dependency_cache_isolated_between_tasks(docket: Docket, worker: Worker):
+    """Dependency cache should be fresh for each task, not reused"""
+    call_counts = {"task1": 0, "task2": 0}
+
+    def dependency_for_task1() -> str:
+        call_counts["task1"] += 1
+        return f"task1-call-{call_counts['task1']}"
+
+    def dependency_for_task2() -> str:
+        call_counts["task2"] += 1
+        return f"task2-call-{call_counts['task2']}"
+
+    async def first_task(val: str = Depends(dependency_for_task1)):
+        assert val == "task1-call-1"
+
+    async def second_task(val: str = Depends(dependency_for_task2)):
+        assert val == "task2-call-1"
+
+    # Run tasks sequentially
+    await docket.add(first_task)()
+    await worker.run_until_finished()
+
+    await docket.add(second_task)()
+    await worker.run_until_finished()
+
+    # Each dependency should have been called once (no cache leakage between tasks)
+    assert call_counts["task1"] == 1
+    assert call_counts["task2"] == 1
+
+
+async def test_async_exit_stack_cleanup(docket: Docket, worker: Worker):
+    """AsyncExitStack should be properly cleaned up after task execution"""
+    from contextlib import asynccontextmanager
+
+    cleanup_called: list[str] = []
+
+    @asynccontextmanager
+    async def tracked_resource():
+        try:
+            yield "resource"
+        finally:
+            cleanup_called.append("cleaned")
+
+    async def task_with_context(res: str = Depends(tracked_resource)):
+        assert res == "resource"
+        assert len(cleanup_called) == 0  # Not cleaned up yet
+
+    await docket.add(task_with_context)()
+    await worker.run_until_finished()
+
+    # After task completes, cleanup should have been called
+    assert cleanup_called == ["cleaned"]
+
+
+async def test_contextvar_reset_on_reentrant_call(docket: Docket, worker: Worker):
+    """Contextvars should be properly reset on reentrant calls to resolved_dependencies"""
+
+    # Create two mock executions
+    async def task1():
+        pass
+
+    async def task2():
+        pass
+
+    execution1 = Execution(
+        key="task1-key",
+        function=task1,
+        args=(),
+        kwargs={},
+        attempt=1,
+        when=datetime.now(timezone.utc),
+    )
+
+    execution2 = Execution(
+        key="task2-key",
+        function=task2,
+        args=(),
+        kwargs={},
+        attempt=1,
+        when=datetime.now(timezone.utc),
+    )
+
+    # Capture contextvars from first call
+    captured_exec1 = None
+    captured_stack1 = None
+
+    async with resolved_dependencies(worker, execution1):
+        captured_exec1 = Dependency.execution.get()
+        captured_stack1 = _Depends.stack.get()
+        assert captured_exec1 is execution1
+
+    # After exiting, contextvars should be reset (raise LookupError)
+    try:
+        current_exec = Dependency.execution.get()
+        # If we get here without LookupError, check if it's stale
+        assert current_exec is not execution1, (
+            "Contextvar still points to old execution!"
+        )
+    except LookupError:
+        # Expected - contextvar was properly reset
+        pass
+
+    # Now make a second call - should not see values from first call
+    async with resolved_dependencies(worker, execution2):
+        captured_exec2 = Dependency.execution.get()
+        captured_stack2 = _Depends.stack.get()
+        assert captured_exec2 is execution2
+        assert captured_exec2 is not captured_exec1
+        # Stacks should be different objects
+        assert captured_stack2 is not captured_stack1
+
+
+async def test_contextvar_not_leaked_to_caller(docket: Docket):
+    """Verify contextvars don't leak outside resolved_dependencies context"""
+    # Before calling resolved_dependencies, contextvars should not be set
+    with pytest.raises(LookupError):
+        Dependency.execution.get()
+
+    async def dummy_task():
+        pass
+
+    execution = Execution(
+        key="test-key",
+        function=dummy_task,
+        args=(),
+        kwargs={},
+        attempt=1,
+        when=datetime.now(timezone.utc),
+    )
+
+    from docket.worker import Worker
+
+    async with Docket("test-contextvar-leak", url="memory://leak-test") as test_docket:
+        async with Worker(test_docket) as test_worker:
+            # Use resolved_dependencies
+            async with resolved_dependencies(test_worker, execution):
+                # Inside context, we should be able to get values
+                assert Dependency.execution.get() is execution
+
+            # After exiting context, contextvars should be cleaned up
+            with pytest.raises(LookupError):
+                Dependency.execution.get()
+
+            with pytest.raises(LookupError):
+                _Depends.stack.get()
+
+            with pytest.raises(LookupError):
+                _Depends.cache.get()