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:
MattMaximo
2026-09-01 23:00:14 -04:00
committed by Teknium
parent eb4d77c2a8
commit 6f625f7381
3 changed files with 87 additions and 1 deletions

View File

@@ -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)

View File

@@ -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

View File

@@ -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.