fix: propagate contextvars through thread_pool_exec (#16247)

## Problem

`thread_pool_exec()` dispatches work via `loop.run_in_executor()`, which
submits the callable with a plain `executor.submit(func, *args)` and
does **not** copy the caller's `contextvars.Context`. So a `ContextVar`
set in the async caller is not visible inside the function running in
the worker thread.

This differs from `asyncio.to_thread()`, which runs the callable inside
a copied context. `run_in_executor()` has never propagated context
(verified on Python 3.12 and 3.13) — so this is a pre-existing gap in
the helper, **not** a regression or a Python-version compatibility
issue.

Concretely, any code that sets a `ContextVar` in async code and reads it
inside a function dispatched via `thread_pool_exec` (request tracing,
per-task state, Langfuse trace propagation, etc.) silently loses that
context.

## Fix

Copy the current context before submitting and run the callable inside
it with `ctx.run()`, matching what `asyncio.to_thread()` does:

```python
async def thread_pool_exec(func, *args, **kwargs):
    loop = asyncio.get_running_loop()
    ctx = contextvars.copy_context()
    if kwargs:
        inner = functools.partial(func, *args, **kwargs)
        return await loop.run_in_executor(_thread_pool_executor(), ctx.run, inner)
    return await loop.run_in_executor(_thread_pool_executor(), ctx.run, func, *args)
```

This explicitly **adds** ContextVar propagation to the helper (it does
not restore any prior behavior). Backward-compatible.

## Tests

`TestThreadPoolExec` covers propagation, the kwargs path, per-call
isolation and the unset-default case.

> Note: the branch name still contains `python313` for historical
reasons; the change is unrelated to any Python version.
This commit is contained in:
VincentLambert
2026-06-23 09:17:42 +02:00
committed by GitHub
parent d8ee1ffaad
commit 11e14a8353
2 changed files with 74 additions and 4 deletions

View File

@@ -16,6 +16,7 @@
import asyncio
import base64
import contextvars
import functools
import hashlib
import logging
@@ -250,8 +251,13 @@ def _thread_pool_executor():
async def thread_pool_exec(func, *args, **kwargs):
# loop.run_in_executor() submits the callable without propagating the caller's
# contextvars (unlike asyncio.to_thread, which copies the context). Copy the
# current context and run the callable inside it so ContextVars set by the
# caller (e.g. tracing / per-request state) are visible in the worker thread.
loop = asyncio.get_running_loop()
ctx = contextvars.copy_context()
if kwargs:
func = functools.partial(func, *args, **kwargs)
return await loop.run_in_executor(_thread_pool_executor(), func)
return await loop.run_in_executor(_thread_pool_executor(), func, *args)
inner = functools.partial(func, *args, **kwargs)
return await loop.run_in_executor(_thread_pool_executor(), ctx.run, inner)
return await loop.run_in_executor(_thread_pool_executor(), ctx.run, func, *args)