fix(tools): propagate caller contextvars in DaemonThreadPoolExecutor.submit
Some bundled CPython runtime builds strip stdlib ThreadPoolExecutor's copy_context() propagation, so work submitted to the daemon pool runs in a bare context. Under the multiplexed gateway this dropped the profile secret scope in pool workers: the context-compression timeout fence resolved auxiliary provider keys (SURPLUS_API_KEY) with UnscopedSecretError, silently degrading LLM compression to lossy deterministic summaries and driving re-read loops in affected sessions. Restore stdlib semantics in submit() by snapshotting the caller's context and running the callable inside it (a no-op re-application on runtimes that already propagate). Mirrors the gateway's _run_in_executor_with_context pattern. Tests: daemon pool worker sees caller contextvars; scoped get_secret works in a daemon-pool worker under multiplex while scoped misses still fail closed (no env leak).
This commit is contained in:
@@ -347,3 +347,34 @@ class TestRelayRoutingStampGlobals:
|
||||
ss.set_multiplex_active(False)
|
||||
for name in self.AUTH_VARS:
|
||||
assert not ss._is_global_env(name), name
|
||||
|
||||
|
||||
class TestSecretScopeAcrossExecutorThreads:
|
||||
"""Multiplexed profile state must reach pool workers (see #95119).
|
||||
|
||||
The context-compression timeout fence runs auxiliary LLM calls in a
|
||||
daemon thread pool. Bundled CPython runtime builds omit
|
||||
``ThreadPoolExecutor``'s context propagation, so the profile secret
|
||||
scope was absent in the worker and ``get_secret`` failed closed with
|
||||
``UnscopedSecretError``, silently degrading compression to lossy
|
||||
deterministic summaries. ``DaemonThreadPoolExecutor.submit`` restores
|
||||
stdlib context semantics; these tests lock that in.
|
||||
"""
|
||||
|
||||
def test_scoped_read_works_in_daemon_pool_worker(self, monkeypatch):
|
||||
from tools.daemon_pool import DaemonThreadPoolExecutor
|
||||
|
||||
monkeypatch.setenv("SURPLUS_API_KEY", "env-key")
|
||||
ss.set_multiplex_active(True)
|
||||
token = ss.set_secret_scope({"SURPLUS_API_KEY": "scope-key"})
|
||||
pool = DaemonThreadPoolExecutor(max_workers=1)
|
||||
try:
|
||||
# The scope (authoritative under multiplex) must reach the worker.
|
||||
seen = pool.submit(ss.get_secret, "SURPLUS_API_KEY").result(timeout=10)
|
||||
assert seen == "scope-key"
|
||||
# A scoped miss must still not borrow the (cross-profile) env value.
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "env-leak")
|
||||
assert pool.submit(ss.get_secret, "OPENAI_API_KEY").result(timeout=10) is None
|
||||
finally:
|
||||
pool.shutdown(wait=True)
|
||||
ss.reset_secret_scope(token)
|
||||
|
||||
@@ -69,6 +69,30 @@ def test_wedged_worker_does_not_block_interpreter_exit():
|
||||
assert "main-done" in proc.stdout
|
||||
|
||||
|
||||
def test_submit_propagates_caller_contextvars():
|
||||
"""Pool workers inherit contextvars set in the submitting context.
|
||||
|
||||
Stdlib ThreadPoolExecutor snapshots the caller's context with
|
||||
``copy_context()``; some bundled CPython runtime builds strip that, so
|
||||
the daemon pool restores it explicitly. Without the fix this returns
|
||||
the default because the worker runs in a bare context.
|
||||
"""
|
||||
from contextvars import ContextVar
|
||||
|
||||
var = ContextVar("daemon_pool_test_var", default="unset")
|
||||
|
||||
pool = DaemonThreadPoolExecutor(max_workers=1)
|
||||
try:
|
||||
token = var.set("hello")
|
||||
try:
|
||||
seen = pool.submit(var.get).result(timeout=10)
|
||||
finally:
|
||||
var.reset(token)
|
||||
assert seen == "hello"
|
||||
finally:
|
||||
pool.shutdown(wait=True)
|
||||
|
||||
|
||||
def _repo_root():
|
||||
import pathlib
|
||||
|
||||
|
||||
@@ -16,7 +16,15 @@ exit hook insists on joining.
|
||||
- the interpreter's non-daemon thread join at shutdown skips them.
|
||||
|
||||
Semantics are otherwise identical (initializer/initargs, work queue,
|
||||
idle-thread reuse). Use it for any pool whose work is best-effort or
|
||||
idle-thread reuse) and, since #95119, so is context propagation:
|
||||
``submit`` snapshots the submitting context with ``copy_context()`` and
|
||||
runs each work item inside it, matching stdlib ``ThreadPoolExecutor``.
|
||||
That matters because some bundled CPython runtime builds omit stdlib's
|
||||
context propagation entirely, which silently drops contextvar-based state
|
||||
(profile secret scope, HERMES_HOME override) in pool workers — e.g. the
|
||||
context-compression timeout fence resolved auxiliary provider keys with
|
||||
``UnscopedSecretError`` under the multiplexed gateway. Use it for any
|
||||
pool whose work is best-effort or
|
||||
independently interruptible and must never hold the process open:
|
||||
concurrent tool execution, background memory sync, catalog fan-out,
|
||||
subagent timeout wrappers. Do NOT use it for work that must complete
|
||||
@@ -30,6 +38,7 @@ import threading
|
||||
import weakref
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from concurrent.futures.thread import _worker
|
||||
from contextvars import copy_context
|
||||
|
||||
__all__ = ["DaemonThreadPoolExecutor"]
|
||||
|
||||
@@ -37,6 +46,28 @@ __all__ = ["DaemonThreadPoolExecutor"]
|
||||
class DaemonThreadPoolExecutor(ThreadPoolExecutor):
|
||||
"""ThreadPoolExecutor variant whose workers do not block process exit."""
|
||||
|
||||
def submit(self, fn, /, *args, **kwargs):
|
||||
"""Submit a callable, propagating the caller's contextvars.
|
||||
|
||||
Stdlib ``ThreadPoolExecutor`` snapshots the submitting context with
|
||||
``copy_context()`` and runs each work item inside it, so pool
|
||||
workers inherit contextvar state such as the multiplexed profile
|
||||
secret scope. Some bundled CPython runtime builds strip that
|
||||
propagation from the stdlib executor (their ``_WorkItem.run`` calls
|
||||
the callable directly), which broke auxiliary LLM key resolution
|
||||
from the context-compression timeout fence with
|
||||
``UnscopedSecretError``. Restore the stdlib behavior explicitly so
|
||||
the daemon pool behaves identically on every runtime; on runtimes
|
||||
that already propagate, the inner ``ctx.run`` re-applies the same
|
||||
immutable context and is a no-op.
|
||||
"""
|
||||
ctx = copy_context()
|
||||
|
||||
def _run_with_context(*call_args, **call_kwargs):
|
||||
return ctx.run(fn, *call_args, **call_kwargs)
|
||||
|
||||
return super().submit(_run_with_context, *args, **kwargs)
|
||||
|
||||
def _adjust_thread_count(self) -> None:
|
||||
# Mirrors CPython's implementation (3.8–3.13) with two changes:
|
||||
# daemon=True and no _threads_queues registration.
|
||||
|
||||
Reference in New Issue
Block a user