fix: restore HTTP headers in worker execution path for background tasks (#3681)

This commit is contained in:
Jeremiah Lowin 2026-03-28 18:01:31 -04:00 committed by GitHub
commit 90bd4ffaa4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 49 additions and 3 deletions

View file

@ -916,6 +916,7 @@ class _CurrentContext(Dependency["Context"]):
_context: Context | None = None
_access_token_cv_token: Token[AccessToken | None] | None = None
_http_headers_cv_token: Token[dict[str, str] | None] | None = None
async def __aenter__(self) -> Context:
from fastmcp.server.context import Context, _current_context
@ -950,6 +951,11 @@ class _CurrentContext(Dependency["Context"]):
task_info.session_id, task_info.task_id
)
# Restore HTTP headers snapshot from Redis (#3631)
self._http_headers_cv_token = await _restore_task_http_headers(
task_info.session_id, task_info.task_id
)
return self._context
# Neither foreground nor background context available
@ -970,6 +976,10 @@ class _CurrentContext(Dependency["Context"]):
if self._access_token_cv_token is not None:
_task_access_token.reset(self._access_token_cv_token)
self._access_token_cv_token = None
# Clean up HTTP headers ContextVar
if self._http_headers_cv_token is not None:
_task_http_headers.reset(self._http_headers_cv_token)
self._http_headers_cv_token = None
# Clean up if we created a context for background task
if self._context is not None:
await self._context.__aexit__(exc_type, exc_value, traceback)

View file

@ -2,6 +2,7 @@
from __future__ import annotations
import functools
import inspect
import warnings
from collections.abc import Callable
@ -27,7 +28,12 @@ import fastmcp
from fastmcp.decorators import resolve_task_config
from fastmcp.exceptions import FastMCPDeprecationWarning
from fastmcp.server.auth.authorization import AuthCheck
from fastmcp.server.dependencies import without_injected_parameters
from fastmcp.server.dependencies import (
_restore_task_http_headers,
_task_http_headers,
get_task_context,
without_injected_parameters,
)
from fastmcp.server.tasks.config import TaskConfig
from fastmcp.tools.base import (
Tool,
@ -284,11 +290,13 @@ class FunctionTool(Tool):
"""Register this tool with docket for background execution.
FunctionTool registers the underlying function, which has the user's
Depends parameters for docket to resolve.
Depends parameters for docket to resolve. The function is wrapped to
eagerly restore HTTP headers from Redis so that get_http_request()
works even without explicit dependency injection.
"""
if not self.task_config.supports_tasks():
return
docket.register(self.fn, names=[self.key])
docket.register(_wrap_for_task_http_headers(self.fn), names=[self.key])
async def add_to_docket(
self,
@ -316,6 +324,34 @@ class FunctionTool(Tool):
return await docket.add(lookup_key, **kwargs)(**arguments)
def _wrap_for_task_http_headers(fn: Callable[..., Any]) -> Callable[..., Any]:
"""Wrap a function to restore HTTP headers in background task workers.
Uses functools.wraps so docket sees the original signature for dependency
resolution while the wrapper eagerly populates _task_http_headers before
the user's function runs.
"""
@functools.wraps(fn)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
task_info = get_task_context()
token = None
if task_info is not None and _task_http_headers.get() is None:
token = await _restore_task_http_headers(
task_info.session_id, task_info.task_id
)
try:
result = fn(*args, **kwargs)
if inspect.isawaitable(result):
result = await result
return result
finally:
if token is not None:
_task_http_headers.reset(token)
return wrapper
@overload
def tool(fn: F) -> F: ...
@overload