fix(auxiliary): isolate runtime cache by live context

This commit is contained in:
dfein38347g 2026-07-17 07:52:50 -07:00 committed by Teknium
parent 9e1b1d7536
commit fdc6c32d7d
5 changed files with 394 additions and 74 deletions

View file

@ -41,6 +41,8 @@ Payment / credit exhaustion fallback:
"""
import contextlib
import contextvars
import inspect
import json
import logging
import os
@ -2207,7 +2209,7 @@ def _read_main_model() -> str:
that gate on "the active main model" (e.g. ``vision_analyze``'s native
fast path) see the live runtime, not the persisted config default.
"""
override = _RUNTIME_MAIN_MODEL
override = _runtime_main_value("model")
if isinstance(override, str) and override.strip():
return override.strip()
try:
@ -2234,7 +2236,7 @@ def _read_main_provider() -> str:
Runtime override: see ``_read_main_model`` same mechanism for the
provider half of the runtime tuple.
"""
override = _RUNTIME_MAIN_PROVIDER
override = _runtime_main_value("provider")
if isinstance(override, str) and override.strip():
return override.strip().lower()
try:
@ -2263,7 +2265,7 @@ def _read_main_api_key() -> str:
the main model's credentials instead of falling to ``no-key-required``
(issue #9318).
"""
override = _RUNTIME_MAIN_API_KEY
override = _runtime_main_value("api_key")
if isinstance(override, str) and override.strip():
return override.strip()
try:
@ -2284,7 +2286,7 @@ def _read_main_base_url() -> str:
Same override-then-config pattern as ``_read_main_api_key``.
"""
override = _RUNTIME_MAIN_BASE_URL
override = _runtime_main_value("base_url")
if isinstance(override, str) and override.strip():
return override.strip()
try:
@ -2320,13 +2322,55 @@ def _read_main_api_key_if_same_host(aux_base_url: str) -> str:
return _read_main_api_key()
# Process-local override set by AIAgent at session/turn start. Single-threaded
# per turn — no lock needed. Cleared by ``clear_runtime_main()``.
# Compatibility mirrors for older readers/tests. The authoritative value is
# the ContextVar below: gateway sessions can overlap in one process, so a
# process-global tuple is not safe as routing or cache-key input.
_RUNTIME_MAIN_PROVIDER: str = ""
_RUNTIME_MAIN_MODEL: str = ""
_RUNTIME_MAIN_BASE_URL: str = ""
_RUNTIME_MAIN_API_KEY: str = ""
_RUNTIME_MAIN_API_KEY: Any = ""
_RUNTIME_MAIN_API_MODE: str = ""
_RUNTIME_MAIN_AUTH_MODE: str = ""
_RUNTIME_MAIN_CONTEXT: contextvars.ContextVar[Optional[Dict[str, Any]]] = (
contextvars.ContextVar("auxiliary_runtime_main", default=None)
)
_RUNTIME_MAIN_COMPAT_SNAPSHOT: Tuple[Any, ...] = ("", "", "", "", "", "")
_RUNTIME_MAIN_COMPAT_LOCK = threading.Lock()
def _compat_runtime_main() -> Optional[Dict[str, Any]]:
"""Expose deliberately patched legacy globals in a single main context.
``set_runtime_main`` mirrors values into the old module attributes for
introspection, but those mirrors must never become runtime inputs. A direct
patch is recognized only when it differs from the mirrored snapshot and
only on the main thread, keeping concurrent session workers isolated.
"""
if threading.current_thread() is not threading.main_thread():
return None
values = (
_RUNTIME_MAIN_PROVIDER,
_RUNTIME_MAIN_MODEL,
_RUNTIME_MAIN_BASE_URL,
_RUNTIME_MAIN_API_KEY,
_RUNTIME_MAIN_API_MODE,
_RUNTIME_MAIN_AUTH_MODE,
)
if values == _RUNTIME_MAIN_COMPAT_SNAPSHOT:
return None
return dict(zip(_MAIN_RUNTIME_FIELDS, values))
def _runtime_main_value(field: str) -> Any:
"""Read one runtime field through context-local/controlled legacy state."""
runtime = _RUNTIME_MAIN_CONTEXT.get()
if runtime is None:
runtime = _compat_runtime_main()
if isinstance(runtime, dict):
value = runtime.get(field)
if value:
return value
return ""
def set_runtime_main(
@ -2334,38 +2378,61 @@ def set_runtime_main(
model: str,
*,
base_url: str = "",
api_key: str = "",
api_key: Any = "",
api_mode: str = "",
auth_mode: str = "",
) -> None:
"""Record the live runtime provider/model/credentials for the current AIAgent.
"""Record the current context's live main runtime for auxiliary routing.
Called by ``run_agent.AIAgent._sync_runtime_main_for_aux_routing`` (or
equivalent setter) at the top of each turn so that
``_read_main_provider`` / ``_read_main_model`` reflect CLI/gateway
overrides instead of the stale config.yaml default.
For ``custom:`` providers, ``base_url`` and ``api_key`` must also be
recorded so that ``_resolve_auto`` can construct a valid client in
Step 1 instead of falling through to the aggregator chain.
Context-local state prevents concurrent gateway sessions from overwriting
one another while retaining compatibility mirrors for legacy readers.
"""
global _RUNTIME_MAIN_PROVIDER, _RUNTIME_MAIN_MODEL
global _RUNTIME_MAIN_BASE_URL, _RUNTIME_MAIN_API_KEY, _RUNTIME_MAIN_API_MODE
_RUNTIME_MAIN_PROVIDER = (provider or "").strip().lower()
_RUNTIME_MAIN_MODEL = (model or "").strip()
_RUNTIME_MAIN_BASE_URL = (base_url or "").strip()
_RUNTIME_MAIN_API_KEY = api_key.strip() if isinstance(api_key, str) else ""
_RUNTIME_MAIN_API_MODE = (api_mode or "").strip()
global _RUNTIME_MAIN_AUTH_MODE, _RUNTIME_MAIN_COMPAT_SNAPSHOT
runtime = {
"provider": (provider or "").strip().lower(),
"model": (model or "").strip(),
"base_url": (base_url or "").strip(),
"api_key": (
api_key.strip()
if isinstance(api_key, str)
else api_key if callable(api_key) else ""
),
"api_mode": (api_mode or "").strip(),
"auth_mode": (auth_mode or "").strip().lower(),
}
# Publish authoritative context before updating locked compatibility
# mirrors; concurrent sessions never read those mirrors at runtime.
_RUNTIME_MAIN_CONTEXT.set(runtime)
with _RUNTIME_MAIN_COMPAT_LOCK:
(
_RUNTIME_MAIN_PROVIDER,
_RUNTIME_MAIN_MODEL,
_RUNTIME_MAIN_BASE_URL,
_RUNTIME_MAIN_API_KEY,
_RUNTIME_MAIN_API_MODE,
_RUNTIME_MAIN_AUTH_MODE,
) = (runtime[field] for field in _MAIN_RUNTIME_FIELDS)
_RUNTIME_MAIN_COMPAT_SNAPSHOT = tuple(
runtime[field] for field in _MAIN_RUNTIME_FIELDS
)
def clear_runtime_main() -> None:
"""Clear the runtime override (e.g. on session end)."""
"""Clear the runtime override in the current context."""
global _RUNTIME_MAIN_PROVIDER, _RUNTIME_MAIN_MODEL
global _RUNTIME_MAIN_BASE_URL, _RUNTIME_MAIN_API_KEY, _RUNTIME_MAIN_API_MODE
_RUNTIME_MAIN_PROVIDER = ""
_RUNTIME_MAIN_MODEL = ""
_RUNTIME_MAIN_BASE_URL = ""
_RUNTIME_MAIN_API_KEY = ""
_RUNTIME_MAIN_API_MODE = ""
global _RUNTIME_MAIN_AUTH_MODE, _RUNTIME_MAIN_COMPAT_SNAPSHOT
_RUNTIME_MAIN_CONTEXT.set(None)
with _RUNTIME_MAIN_COMPAT_LOCK:
_RUNTIME_MAIN_PROVIDER = ""
_RUNTIME_MAIN_MODEL = ""
_RUNTIME_MAIN_BASE_URL = ""
_RUNTIME_MAIN_API_KEY = ""
_RUNTIME_MAIN_API_MODE = ""
_RUNTIME_MAIN_AUTH_MODE = ""
_RUNTIME_MAIN_COMPAT_SNAPSHOT = ("", "", "", "", "", "")
def _resolve_custom_runtime() -> Tuple[Optional[str], Optional[str], Optional[str]]:
@ -2784,6 +2851,14 @@ def _normalize_main_runtime(main_runtime: Optional[Dict[str, Any]]) -> Dict[str,
surface as the main agent. The OpenAI SDK accepts ``Callable[[], str]``
for ``api_key`` and calls it before every request.
"""
if main_runtime is None:
# Context-local state is inherited by tool worker wrappers while
# remaining isolated across concurrent gateway sessions. Never fall
# back to compatibility mirrors here: another session may have written
# them most recently, which would leak its endpoint/key into this call.
main_runtime = _RUNTIME_MAIN_CONTEXT.get()
if main_runtime is None:
main_runtime = _compat_runtime_main()
if not isinstance(main_runtime, dict):
return {}
normalized: Dict[str, Any] = {}
@ -3311,13 +3386,7 @@ def _evict_cached_clients(provider: str) -> None:
for key in stale_keys:
client = _client_cache.get(key, (None, None, None))[0]
if client is not None:
_force_close_async_httpx(client)
try:
close_fn = getattr(client, "close", None)
if callable(close_fn):
close_fn()
except Exception:
pass
_close_cached_client(client)
_client_cache.pop(key, None)
@ -4308,17 +4377,6 @@ def _resolve_auto(
runtime_api_key = runtime.get("api_key", "")
runtime_api_mode = str(runtime.get("api_mode") or "")
# Fall back to process-local globals when main_runtime dict was not
# provided or was incomplete. ``set_runtime_main()`` now records
# base_url/api_key/api_mode alongside provider/model, so custom:
# providers get the full credential surface in Step 1 of the
# auto-detect chain.
if not runtime_base_url and _RUNTIME_MAIN_BASE_URL:
runtime_base_url = _RUNTIME_MAIN_BASE_URL
if not runtime_api_key and _RUNTIME_MAIN_API_KEY:
runtime_api_key = _RUNTIME_MAIN_API_KEY
if not runtime_api_mode and _RUNTIME_MAIN_API_MODE:
runtime_api_mode = _RUNTIME_MAIN_API_MODE
# ── Warn once if OPENAI_BASE_URL is set but config.yaml uses a named
# provider (not 'custom'). This catches the common "env poisoning"
@ -5602,10 +5660,15 @@ def resolve_vision_provider_client(
rpc_api_key = None
rpc_api_mode = resolved_api_mode
if main_provider == "custom" or main_provider.startswith("custom:"):
if _RUNTIME_MAIN_BASE_URL:
rpc_base_url = _RUNTIME_MAIN_BASE_URL
rpc_api_key = _RUNTIME_MAIN_API_KEY or None
rpc_api_mode = resolved_api_mode or _RUNTIME_MAIN_API_MODE or None
runtime_base_url = _runtime_main_value("base_url")
if runtime_base_url:
rpc_base_url = runtime_base_url
rpc_api_key = _runtime_main_value("api_key") or None
rpc_api_mode = (
resolved_api_mode
or _runtime_main_value("api_mode")
or None
)
else:
# No live runtime recorded (non-gateway caller): fall
# back to resolving the configured custom endpoint.
@ -5742,6 +5805,35 @@ _client_cache_lock = threading.Lock()
_CLIENT_CACHE_MAX_SIZE = 64 # safety belt — evict oldest when exceeded
class _CallableCacheDiscriminator:
"""Hash a credential callback by identity without exposing its state."""
__slots__ = ("_callback",)
def __init__(self, callback: Any) -> None:
# Retain the callback so its id cannot be reused while cached.
self._callback = callback
def __hash__(self) -> int:
return id(self._callback)
def __eq__(self, other: object) -> bool:
return (
isinstance(other, _CallableCacheDiscriminator)
and self._callback is other._callback
)
def __repr__(self) -> str:
return "<callable-api-key>"
def _runtime_cache_discriminator(field: str, value: Any) -> Any:
"""Return a hashable, secret-safe runtime cache-key component."""
if field == "api_key" and callable(value):
return _CallableCacheDiscriminator(value)
return value
def _client_cache_key(
provider: str,
*,
@ -5755,7 +5847,10 @@ def _client_cache_key(
model: Optional[str] = None,
) -> tuple:
runtime = _normalize_main_runtime(main_runtime)
runtime_key = tuple(runtime.get(field, "") for field in _MAIN_RUNTIME_FIELDS) if provider == "auto" else ()
runtime_key = tuple(
_runtime_cache_discriminator(field, runtime.get(field, ""))
for field in _MAIN_RUNTIME_FIELDS
) if provider == "auto" else ()
# `auto` can now resolve through task-specific or main fallback policy,
# so the task participates in the cache key. Non-auto providers keep the
# old cache shape because the explicit provider/model tuple is sufficient.
@ -5778,13 +5873,7 @@ def _store_cached_client(cache_key: tuple, client: Any, default_model: Optional[
with _client_cache_lock:
old_entry = _client_cache.get(cache_key)
if old_entry is not None and old_entry[0] is not client:
_force_close_async_httpx(old_entry[0])
try:
close_fn = getattr(old_entry[0], "close", None)
if callable(close_fn):
close_fn()
except Exception:
pass
_close_cached_client(old_entry[0])
_client_cache[cache_key] = (client, default_model, bound_loop)
@ -5886,30 +5975,31 @@ def _force_close_async_httpx(client: Any) -> None:
pass
def _close_cached_client(client: Any) -> None:
"""Apply the canonical best-effort close policy to one cached client."""
if client is None:
return
_force_close_async_httpx(client)
try:
close_fn = getattr(client, "close", None)
if callable(close_fn) and not inspect.iscoroutinefunction(close_fn):
close_fn()
except Exception:
pass
def shutdown_cached_clients() -> None:
"""Close all cached clients (sync and async) to prevent event-loop errors.
Call this during CLI shutdown, *before* the event loop is closed, to
avoid ``AsyncHttpxClientWrapper.__del__`` raising on a dead loop.
"""
import inspect
with _client_cache_lock:
for key, entry in list(_client_cache.items()):
client = entry[0]
if client is None:
continue
# Mark any async httpx transport as closed first (prevents __del__
# from scheduling aclose() on a dead event loop).
_force_close_async_httpx(client)
# Sync clients: close the httpx connection pool cleanly.
# Async clients: skip — we already neutered __del__ above.
try:
close_fn = getattr(client, "close", None)
if close_fn and not inspect.iscoroutinefunction(close_fn):
close_fn()
except Exception:
pass
_close_cached_client(client)
_client_cache.clear()
@ -6061,7 +6151,7 @@ def _get_cached_client(
# the oldest entries (FIFO — dict preserves insertion order).
while len(_client_cache) >= _CLIENT_CACHE_MAX_SIZE:
evict_key, evict_entry = next(iter(_client_cache.items()))
_force_close_async_httpx(evict_entry[0])
_close_cached_client(evict_entry[0])
del _client_cache[evict_key]
_client_cache[cache_key] = (client, default_model, bound_loop)
else:

View file

@ -267,9 +267,9 @@ def _resolve_inference_base_url(
) -> str:
"""Best-effort base URL for the active inference provider."""
try:
from agent.auxiliary_client import _RUNTIME_MAIN_BASE_URL
from agent.auxiliary_client import _runtime_main_value
runtime = str(_RUNTIME_MAIN_BASE_URL or "").strip()
runtime = str(_runtime_main_value("base_url") or "").strip()
if runtime:
return runtime
except Exception:

View file

@ -160,6 +160,7 @@ def build_turn_context(
base_url=getattr(agent, "base_url", "") or "",
api_key=getattr(agent, "api_key", "") or "",
api_mode=getattr(agent, "api_mode", "") or "",
auth_mode=getattr(agent, "auth_mode", "") or "",
)
except Exception:
pass

View file

@ -1199,6 +1199,7 @@ class AIAgent:
"base_url": getattr(self, "base_url", "") or "",
"api_key": getattr(self, "api_key", "") or "",
"api_mode": getattr(self, "api_mode", "") or "",
"auth_mode": getattr(self, "auth_mode", "") or "",
}
def _check_compression_model_feasibility(self) -> None:

View file

@ -0,0 +1,228 @@
"""Regression coverage for implicit live-runtime auxiliary cache keys.
#49151/#49156 is specifically the ``provider='auto'`` path where callers omit
``main_runtime`` after a mid-session model switch. This is distinct from
#56889, which isolates callers that pass different explicit ``model=`` values.
"""
from concurrent.futures import ThreadPoolExecutor
from threading import Barrier
from unittest.mock import MagicMock, patch
import pytest
import agent.auxiliary_client as aux
@pytest.fixture(autouse=True)
def _clean_aux_state():
aux.shutdown_cached_clients()
aux.clear_runtime_main()
yield
aux.shutdown_cached_clients()
aux.clear_runtime_main()
def _runtime(model: str, *, provider: str = "custom:llama-swap") -> dict:
return {
"provider": provider,
"model": model,
"base_url": "http://llama-swap.test/v1",
"api_key": "local-key",
"api_mode": "chat_completions",
"auth_mode": "api_key",
}
def test_implicit_auto_cache_rebuilds_after_runtime_model_switch():
"""A /model switch must not reuse the old implicit-auto cache entry."""
built = []
def fake_resolve(_provider, _model, _async_mode, *, main_runtime, **_kwargs):
client = MagicMock(name=f"client-{main_runtime['model']}")
built.append((client, dict(main_runtime)))
return client, main_runtime["model"]
with patch.object(aux, "resolve_provider_client", side_effect=fake_resolve):
aux.set_runtime_main(**_runtime("qwen35b-code"))
first_client, first_model = aux._get_cached_client("auto")
aux.set_runtime_main(**_runtime("qwen27b-code"))
second_client, second_model = aux._get_cached_client("auto")
assert first_model == "qwen35b-code"
assert second_model == "qwen27b-code"
assert second_client is not first_client
assert [runtime["model"] for _, runtime in built] == [
"qwen35b-code",
"qwen27b-code",
]
def test_implicit_runtime_cache_key_covers_full_connection_and_auth_surface():
"""Provider/endpoint/credential/wire/auth changes all isolate auto clients."""
base = _runtime("same-model")
variants = [
{**base, "provider": "custom:other"},
{**base, "base_url": "https://other.test/v1"},
{**base, "api_key": "other-key"},
{**base, "api_mode": "codex_responses"},
{**base, "auth_mode": "entra_id", "api_key": lambda: "token"},
]
aux.set_runtime_main(**base)
baseline = aux._client_cache_key("auto", async_mode=False)
keys = []
for variant in variants:
aux.set_runtime_main(**variant)
keys.append(aux._client_cache_key("auto", async_mode=False))
assert all(key != baseline for key in keys)
assert len(set(keys)) == len(keys)
def test_implicit_runtime_is_isolated_between_concurrent_session_contexts():
"""Concurrent gateway sessions must not read each other's live runtime."""
barrier = Barrier(2)
def session(model: str):
aux.set_runtime_main(**_runtime(model))
barrier.wait()
normalized = aux._normalize_main_runtime(None)
return normalized["model"], aux._client_cache_key("auto", async_mode=False)
with ThreadPoolExecutor(max_workers=2) as pool:
first = pool.submit(session, "session-a-model")
second = pool.submit(session, "session-b-model")
model_a, key_a = first.result()
model_b, key_b = second.result()
assert model_a == "session-a-model"
assert model_b == "session-b-model"
assert key_a != key_b
def test_context_without_runtime_does_not_fall_back_to_other_session_globals():
"""A fresh context must not inherit another session's compatibility mirrors."""
aux.set_runtime_main(**_runtime("other-session-model"))
def fresh_context():
return aux._normalize_main_runtime(None)
import contextvars
assert contextvars.Context().run(fresh_context) == {}
def test_legacy_patched_globals_are_visible_only_without_an_active_runtime():
"""Direct legacy patches work, but never override context-local session state."""
with patch.object(aux, "_RUNTIME_MAIN_PROVIDER", "custom:legacy"), patch.object(
aux, "_RUNTIME_MAIN_MODEL", "legacy-model"
), patch.object(
aux, "_RUNTIME_MAIN_BASE_URL", "https://legacy.test/v1"
):
assert aux._normalize_main_runtime(None)["model"] == "legacy-model"
aux.set_runtime_main(**_runtime("active-session-model"))
runtime = aux._normalize_main_runtime(None)
assert runtime["model"] == "active-session-model"
assert runtime["base_url"] == "http://llama-swap.test/v1"
def test_concurrent_vision_probes_use_each_sessions_endpoint_and_model():
"""Vision auto-routing must not mix custom endpoints across sessions."""
barrier = Barrier(2)
def fake_resolve(provider, model, **kwargs):
barrier.wait()
client = MagicMock()
client.probed_base_url = kwargs.get("explicit_base_url")
return client, model
def probe(model: str, base_url: str):
runtime = _runtime(model)
runtime["base_url"] = base_url
aux.set_runtime_main(**runtime)
provider, client, resolved_model = aux.resolve_vision_provider_client()
assert client is not None
return provider, resolved_model, client.probed_base_url
with patch.object(
aux, "_resolve_task_provider_model", return_value=("auto", None, None, None, None)
), patch.object(aux, "_main_model_supports_vision", return_value=True), patch.object(
aux, "resolve_provider_client", side_effect=fake_resolve
):
with ThreadPoolExecutor(max_workers=2) as pool:
first = pool.submit(probe, "vision-a", "https://a.test/v1")
second = pool.submit(probe, "vision-b", "https://b.test/v1")
result_a = first.result()
result_b = second.result()
assert result_a == ("custom:llama-swap", "vision-a", "https://a.test/v1")
assert result_b == ("custom:llama-swap", "vision-b", "https://b.test/v1")
def test_explicit_model_cache_isolation_remains_independent_of_runtime_key():
"""#56889 remains covered: explicit model values isolate non-auto clients."""
first = aux._client_cache_key(
"openrouter", async_mode=False, model="anthropic/claude-opus-4.8"
)
second = aux._client_cache_key(
"openrouter", async_mode=False, model="openai/gpt-5.5"
)
assert first != second
def test_unhashable_callable_runtime_api_keys_are_safe_secret_free_discriminators():
"""Callable token providers remain cacheable without leaking returned tokens."""
class TokenProvider(list):
def __init__(self, token: str):
super().__init__()
self.token = token
def __call__(self) -> str:
return self.token
first_provider = TokenProvider("first-super-secret-token")
second_provider = TokenProvider("second-super-secret-token")
first = aux._client_cache_key(
"auto", async_mode=False, main_runtime={**_runtime("same"), "api_key": first_provider}
)
second = aux._client_cache_key(
"auto", async_mode=False, main_runtime={**_runtime("same"), "api_key": second_provider}
)
hash(first)
hash(second)
assert first != second
rendered = repr((first, second))
assert "first-super-secret-token" not in rendered
assert "second-super-secret-token" not in rendered
def test_fifo_eviction_closes_oldest_sync_client_once_after_65_entries():
"""The bounded cache closes, rather than merely drops, its FIFO sync victim."""
clients = []
def fake_resolve(_provider, model, _async_mode, **_kwargs):
client = MagicMock(name=f"client-{model}")
clients.append(client)
return client, model
with patch.object(aux, "resolve_provider_client", side_effect=fake_resolve):
for index in range(65):
aux._get_cached_client("custom", model=f"model-{index}")
assert len(aux._client_cache) == 64
clients[0].close.assert_called_once_with()
for client in clients[1:]:
client.close.assert_not_called()
aux.shutdown_cached_clients()
clients[0].close.assert_called_once_with()
for client in clients[1:]:
client.close.assert_called_once_with()