From ce9adc5e2ce5a030946cc4ea4567eddef19cc341 Mon Sep 17 00:00:00 2001 From: nesquena-hermes Date: Sun, 7 Jun 2026 12:07:13 -0700 Subject: [PATCH] =?UTF-8?q?Release=20v0.51.318=20=E2=80=94=20Release=20KH?= =?UTF-8?q?=20(Phase=203=20light:=20warm=20account-usage=20probe=20worker?= =?UTF-8?q?=20pool,=20#3722)=20(#3792)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase-3-LOW backend refactor. #3722 (@rodboev, #1912): per-probe subprocess.run -> warm worker pool for codex/anthropic quota probes; all hardening retained + idle reaper + credential invalidation + fallback. Codex SAFE, Opus SHIP, real-thread self-verify clean, CI 11/11. Follow-up #3787. Co-authored-by: rodboev --- CHANGELOG.md | 5 + api/providers.py | 380 ++++++++++++++++---- tests/test_agent_source_dependency_audit.py | 11 +- tests/test_provider_quota_status.py | 303 ++++++++++++++-- 4 files changed, 601 insertions(+), 98 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b0375b25..11a24067 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,11 @@ ## [Unreleased] +## [v0.51.318] — 2026-06-07 — Release KH (Phase 3 light — warm account-usage probe worker pool) + +### Changed +- **Account-usage quota probes now reuse a warm worker pool instead of spawning a fresh interpreter per probe.** Quota checks for the providers that need an isolated subprocess previously started a new Python interpreter — and re-imported the agent account-usage stack — on every uncached probe, which is slow when the quota UI is opened across several profiles/providers. A profile-home-keyed pool of persistent warm workers (with a 5-minute idle shutdown) eliminates the repeated spawn/import overhead while keeping every existing safety property: the concurrency semaphore cap, the LRU result cache, the parent-death-signal hardening, and subprocess env scrubbing. Workers are invalidated when provider credentials change so a child process can't retain stale keys, contention falls back to a one-shot probe, and all workers are cleaned up at process exit. (#3722 / #1912, @rodboev) + ## [v0.51.317] — 2026-06-07 — Release KG (Phase 3 light — align CSP enforcement with report policy) ### Fixed diff --git a/api/providers.py b/api/providers.py index 12e6f580..680f6f70 100644 --- a/api/providers.py +++ b/api/providers.py @@ -7,6 +7,7 @@ multi-provider support). from __future__ import annotations +import atexit import base64 import hashlib import json @@ -75,6 +76,7 @@ _PROVIDER_QUOTA_TIMEOUT_SECONDS = 3.0 _ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS = 35.0 _ACCOUNT_USAGE_CACHE_TTL_SECONDS = 45.0 _ACCOUNT_USAGE_CACHE_MAX_ENTRIES = 64 +_ACCOUNT_USAGE_WORKER_IDLE_SECONDS = 5 * 60 _ACCOUNT_USAGE_PROVIDERS = frozenset({"openai-codex", "anthropic"}) # Upper bound on simultaneous profile-isolated quota probe subprocesses. @@ -129,6 +131,8 @@ _account_usage_probe_semaphore: threading.BoundedSemaphore | None = None # represented as non-None snapshots and remain cacheable. _account_usage_status_cache: dict[tuple[str, str, str], tuple[float, Any]] = {} _account_usage_status_cache_lock = threading.Lock() +_account_usage_worker_pool: dict[str, "_AccountUsageProbeWorker"] = {} +_account_usage_worker_pool_lock = threading.Lock() def _get_account_usage_probe_semaphore() -> threading.BoundedSemaphore: @@ -146,7 +150,7 @@ def _get_account_usage_probe_semaphore() -> threading.BoundedSemaphore: # code (_ACCOUNT_USAGE_PARENT_DEATHSIG_BOOTSTRAP) also covers the grandchild # fork inside the child, but this preexec_fn handles the direct child-process # case. Returns None on non-POSIX or when prctl is unavailable so that -# subprocess.run() works on Windows/macOS without changes. +# subprocess startup works on Windows/macOS without changes. def _account_usage_preexec_fn() -> None: try: import ctypes @@ -159,6 +163,7 @@ def _account_usage_preexec_fn() -> None: _ACCOUNT_USAGE_SUBPROCESS_CODE = r""" import base64 import json +import os import sys from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone @@ -627,17 +632,48 @@ def _fetch_codex_account_usage_from_pool(): return None -provider = sys.argv[1] -api_key = sys.argv[2] or None -try: - snapshot = fetch_account_usage(provider, api_key=api_key) -except Exception: - snapshot = None -if str(provider or "").strip().lower() == "openai-codex": - pool_snapshot = _fetch_codex_account_usage_from_pool() - if isinstance(getattr(pool_snapshot, "pool", None), dict): - snapshot = pool_snapshot -print(json.dumps(_snapshot_payload(snapshot))) +def _fetch_snapshot(provider, api_key, env_var=None): + previous = os.environ.get(env_var) if env_var else None + had_previous = bool(env_var and env_var in os.environ) + if env_var and api_key: + os.environ[env_var] = api_key + try: + try: + snapshot = fetch_account_usage(provider, api_key=api_key) + except Exception: + snapshot = None + if str(provider or "").strip().lower() == "openai-codex": + pool_snapshot = _fetch_codex_account_usage_from_pool() + if isinstance(getattr(pool_snapshot, "pool", None), dict): + snapshot = pool_snapshot + return _snapshot_payload(snapshot) + finally: + if env_var and api_key: + if had_previous: + os.environ[env_var] = previous + else: + os.environ.pop(env_var, None) + + +def _run_worker(): + for raw_line in sys.stdin: + try: + request = json.loads(raw_line) + provider = request.get("provider") + api_key = request.get("api_key") or None + env_var = request.get("env_var") or None + payload = _fetch_snapshot(provider, api_key, env_var=env_var) + except Exception: + payload = None + print(json.dumps(payload), flush=True) + + +if len(sys.argv) > 1 and sys.argv[1] == "--worker": + _run_worker() +else: + provider = sys.argv[1] + api_key = sys.argv[2] or None + print(json.dumps(_fetch_snapshot(provider, api_key)), flush=True) """ # SECTION: Provider ↔ env var mapping @@ -1186,6 +1222,267 @@ def _account_usage_payload_to_snapshot(payload: Any) -> Any: ) +class _AccountUsageProbeWorker: + def __init__(self, home: Path): + self.home = Path(home) + self.last_used = time.monotonic() + self._lock = threading.RLock() + self._proc: subprocess.Popen[str] | None = None + + def close(self) -> None: + with self._lock: + proc = self._proc + self._proc = None + self._close_process(proc) + + @staticmethod + def _close_process(proc: subprocess.Popen[str] | None) -> None: + if proc is None: + return + for stream_name in ("stdin", "stdout"): + stream = getattr(proc, stream_name, None) + try: + if stream is not None: + stream.close() + except Exception: + pass + try: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=1.0) + except Exception: + proc.kill() + except Exception: + pass + + def fetch(self, provider: str, *, api_key: str | None = None) -> Any: + if not self._lock.acquire(blocking=False): + return _fetch_account_usage_once_for_home(provider, self.home, api_key=api_key) + try: + return self._fetch_locked(provider, api_key=api_key) + finally: + self._lock.release() + + def _fetch_locked(self, provider: str, *, api_key: str | None = None) -> Any: + self.last_used = time.monotonic() + proc = self._ensure_process(provider) + if proc is None or proc.stdin is None or proc.stdout is None: + return None + + request = json.dumps({ + "provider": provider, + "api_key": api_key or "", + "env_var": _provider_env_var_for((provider or "").strip().lower()), + }) + "\n" + result: dict[str, Any] = {} + + def round_trip() -> None: + try: + proc.stdin.write(request) + proc.stdin.flush() + result["line"] = proc.stdout.readline() + except Exception as exc: + result["error"] = exc + + thread = threading.Thread(target=round_trip, daemon=True) + thread.start() + thread.join(_ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS) + self.last_used = time.monotonic() + if thread.is_alive(): + self.close() + thread.join(timeout=1.0) + logger.debug("Account usage worker for %s timed out", provider) + return None + if result.get("error") is not None: + exc = result["error"] + self.close() + logger.debug( + "Account usage worker for %s failed", + provider, + exc_info=(type(exc), exc, exc.__traceback__), + ) + return None + + line = str(result.get("line") or "").strip() + if not line: + self.close() + logger.debug("Account usage worker for %s exited before responding", provider) + return None + try: + payload = json.loads(line) + except json.JSONDecodeError: + self.close() + logger.debug("Account usage worker for %s returned invalid JSON", provider) + return None + return _account_usage_payload_to_snapshot(payload) + + def _ensure_process(self, provider: str) -> subprocess.Popen[str] | None: + if self._proc is not None and self._proc.poll() is None: + return self._proc + old_proc = self._proc + self._proc = None + self._close_process(old_proc) + try: + from api.config import PYTHON_EXE + except Exception: + PYTHON_EXE = sys.executable or "python3" + + kwargs: dict[str, Any] = { + "stdin": subprocess.PIPE, + "stdout": subprocess.PIPE, + "stderr": subprocess.DEVNULL, + "text": True, + "bufsize": 1, + } + if hasattr(os, "fork"): # POSIX + kwargs["preexec_fn"] = _account_usage_preexec_fn + + try: + self._proc = subprocess.Popen( + [ + PYTHON_EXE, + "-c", + _ACCOUNT_USAGE_PARENT_DEATHSIG_BOOTSTRAP + _ACCOUNT_USAGE_SUBPROCESS_CODE, + "--worker", + ], + env=_account_usage_subprocess_env(self.home, provider, None), + **kwargs, + ) + except Exception: + self._proc = None + logger.debug("Account usage worker for %s failed to launch", provider, exc_info=True) + return self._proc + + +def _launch_account_usage_worker_process( + home: Path, + provider: str, + *, + stdin: Any = subprocess.PIPE, + stdout: Any = subprocess.PIPE, +) -> subprocess.Popen[str] | None: + try: + from api.config import PYTHON_EXE + except Exception: + PYTHON_EXE = sys.executable or "python3" + + kwargs: dict[str, Any] = { + "stdin": stdin, + "stdout": stdout, + "stderr": subprocess.DEVNULL, + "text": True, + "bufsize": 1, + } + if hasattr(os, "fork"): # POSIX + kwargs["preexec_fn"] = _account_usage_preexec_fn + + try: + return subprocess.Popen( + [ + PYTHON_EXE, + "-c", + _ACCOUNT_USAGE_PARENT_DEATHSIG_BOOTSTRAP + _ACCOUNT_USAGE_SUBPROCESS_CODE, + "--worker", + ], + env=_account_usage_subprocess_env(home, provider, None), + **kwargs, + ) + except Exception: + logger.debug("Account usage worker for %s failed to launch", provider, exc_info=True) + return None + + +def _fetch_account_usage_once_for_home(provider: str, home: Path, *, api_key: str | None = None) -> Any: + proc = _launch_account_usage_worker_process(Path(home), provider) + if proc is None or proc.stdin is None or proc.stdout is None: + _AccountUsageProbeWorker._close_process(proc) + return None + request = json.dumps({ + "provider": provider, + "api_key": api_key or "", + "env_var": _provider_env_var_for((provider or "").strip().lower()), + }) + "\n" + try: + stdout, _stderr = proc.communicate(request, timeout=_ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + _AccountUsageProbeWorker._close_process(proc) + return None + except Exception: + _AccountUsageProbeWorker._close_process(proc) + return None + try: + line = str(stdout or "").splitlines()[0] + payload = json.loads(line.strip()) + except json.JSONDecodeError: + return None + except IndexError: + return None + return _account_usage_payload_to_snapshot(payload) + + +def _get_account_usage_probe_worker(home: Path) -> _AccountUsageProbeWorker: + key = str(Path(home)) + with _account_usage_worker_pool_lock: + worker = _account_usage_worker_pool.get(key) + if worker is None: + worker = _AccountUsageProbeWorker(Path(home)) + _account_usage_worker_pool[key] = worker + return worker + + +def _cleanup_account_usage_probe_workers( + *, + now: float | None = None, + idle_seconds: float = _ACCOUNT_USAGE_WORKER_IDLE_SECONDS, +) -> None: + cutoff = time.monotonic() if now is None else now + stale: list[tuple[str, _AccountUsageProbeWorker]] = [] + with _account_usage_worker_pool_lock: + for key, worker in list(_account_usage_worker_pool.items()): + if worker._lock.acquire(blocking=False): + try: + proc = worker._proc + is_dead = proc is None or proc.poll() is not None + if is_dead or cutoff - worker.last_used >= idle_seconds: + stale.append((key, worker)) + _account_usage_worker_pool.pop(key, None) + finally: + worker._lock.release() + for _key, worker in stale: + worker.close() + + +def _close_account_usage_probe_workers() -> None: + with _account_usage_worker_pool_lock: + workers = list(_account_usage_worker_pool.values()) + _account_usage_worker_pool.clear() + _close_account_usage_probe_worker_list(workers) + + +def _close_account_usage_probe_worker_list(workers: list[_AccountUsageProbeWorker]) -> None: + for worker in workers: + worker.close() + + +def _close_account_usage_probe_workers_async() -> None: + with _account_usage_worker_pool_lock: + workers = list(_account_usage_worker_pool.values()) + _account_usage_worker_pool.clear() + if not workers: + return + thread = threading.Thread( + target=_close_account_usage_probe_worker_list, + args=(workers,), + daemon=True, + name="account-usage-worker-close", + ) + thread.start() + + +atexit.register(_close_account_usage_probe_workers) + + def _account_usage_cache_key(provider: str, home: Path, api_key: str | None) -> tuple[str, str, str]: key_fingerprint = "" if api_key: @@ -1211,10 +1508,11 @@ def invalidate_account_usage_status_cache(provider_id: str | None = None) -> Non with _account_usage_status_cache_lock: if not normalized: _account_usage_status_cache.clear() - return - for key in list(_account_usage_status_cache): - if key[0] == normalized: - _account_usage_status_cache.pop(key, None) + else: + for key in list(_account_usage_status_cache): + if key[0] == normalized: + _account_usage_status_cache.pop(key, None) + _close_account_usage_probe_workers_async() def _set_cached_account_usage( @@ -1246,51 +1544,11 @@ def _set_cached_account_usage( def _agent_fetch_account_usage_for_home(provider: str, home: Path, *, api_key: str | None = None) -> Any: try: - from api.config import PYTHON_EXE + _cleanup_account_usage_probe_workers() + return _get_account_usage_probe_worker(home).fetch(provider, api_key=api_key) except Exception: - PYTHON_EXE = sys.executable or "python3" - - try: - # On POSIX (Linux/macOS), wire parent-death signal so the child dies - # cleanly if the WebUI parent terminates. preexec_fn is not safe on - # Windows, where OS-level process-tree cleanup handles child orphans. - kwargs: dict[str, Any] = { - "stdin": subprocess.DEVNULL, - "stdout": subprocess.PIPE, - "stderr": subprocess.PIPE, - "text": True, - "timeout": _ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS, - "check": False, - } - if hasattr(os, "fork"): # POSIX - kwargs["preexec_fn"] = _account_usage_preexec_fn - - proc = subprocess.run( - [ - PYTHON_EXE, "-c", - _ACCOUNT_USAGE_PARENT_DEATHSIG_BOOTSTRAP + _ACCOUNT_USAGE_SUBPROCESS_CODE, - provider, - api_key or "", - ], - env=_account_usage_subprocess_env(home, provider, api_key), - **kwargs, - ) - except subprocess.TimeoutExpired: - logger.debug("Account usage probe for %s timed out", provider) + logger.debug("Account usage probe for %s failed", provider, exc_info=True) return None - except Exception: - logger.debug("Account usage probe for %s failed to launch", provider, exc_info=True) - return None - - if proc.returncode != 0: - logger.debug("Account usage probe for %s exited with status %s", provider, proc.returncode) - return None - try: - payload = json.loads((proc.stdout or "").strip() or "null") - except json.JSONDecodeError: - logger.debug("Account usage probe for %s returned invalid JSON", provider) - return None - return _account_usage_payload_to_snapshot(payload) def _fetch_account_usage_with_profile_context(provider: str, *, refresh: bool = False) -> Any: @@ -1301,8 +1559,7 @@ def _fetch_account_usage_with_profile_context(provider: str, *, refresh: bool = memory by spawning more than _MAX_CONCURRENT_ACCOUNT_USAGE_PROBES probe subprocesses simultaneously. Each probe runs up to 35 s. - A warm worker-pool (reuse of persistent subprocess handles) is a natural - follow-up if this first slice proves insufficient in production. + Warm per-profile worker processes handle the actual probe requests. """ home = _get_hermes_home() api_key = _get_provider_api_key(provider) @@ -2122,6 +2379,7 @@ def set_provider_key(provider_id: str, api_key: str | None) -> dict[str, Any]: # Using invalidate_models_cache() instead of reload_config() to avoid # disrupting active streaming sessions that may be reading config.cfg. invalidate_models_cache() + invalidate_account_usage_status_cache(provider_id) return { "ok": True, diff --git a/tests/test_agent_source_dependency_audit.py b/tests/test_agent_source_dependency_audit.py index a9f74e55..9f56f5fb 100644 --- a/tests/test_agent_source_dependency_audit.py +++ b/tests/test_agent_source_dependency_audit.py @@ -145,8 +145,15 @@ def test_audit_embedded_worker_import_line_anchors_are_source_lines(): if finding["path"] == "api/providers.py" and finding["anchor"] == "agent.account_usage" ) - assert account_usage["line"] == 168 - assert account_usage["text"] == "from agent.account_usage import fetch_account_usage" + expected_text = "from agent.account_usage import fetch_account_usage" + assert account_usage["text"] == expected_text + # The reported line number must be an ACCURATE source anchor — i.e. that line + # in the real file actually contains the import. Asserting the exact number is + # brittle (any edit above it shifts it); asserting the line CONTENT at the + # reported line is the real invariant ("line anchors are source lines"). + source_lines = (REPO / account_usage["path"]).read_text(encoding="utf-8").splitlines() + assert account_usage["line"] >= 1 + assert source_lines[account_usage["line"] - 1].strip() == expected_text def test_audit_reports_runtime_state_and_provider_imports(): diff --git a/tests/test_provider_quota_status.py b/tests/test_provider_quota_status.py index a70e0182..b1a6a0be 100644 --- a/tests/test_provider_quota_status.py +++ b/tests/test_provider_quota_status.py @@ -1307,37 +1307,252 @@ def test_provider_quota_styles_exist(): # ── Regression tests for #1912 ──────────────────────────────────────────────── -def test_account_usage_subprocess_uses_devnull_stdin(monkeypatch): - """Account-usage probe subprocess must receive stdin=DEVNULL. - DEVNULL prevents the child from inheriting any pipe that could block or - leak data. This is a defence-in-depth measure beyond the parent-death - signal; it is tested separately to make the invariant explicit. - """ +class _FakeAccountUsageWorkerProcess: + def __init__(self, payloads=None): + self.payloads = list(payloads or []) + self.requests = [] + self.returncode = None + self.terminated = False + self.killed = False + self.stdin = self._Stdin(self) + self.stdout = self._Stdout(self) + + class _Stdin: + def __init__(self, proc): + self.proc = proc + self.closed = False + + def write(self, value): + self.proc.requests.append(json.loads(value)) + + def flush(self): + pass + + def close(self): + self.closed = True + + class _Stdout: + def __init__(self, proc): + self.proc = proc + self.closed = False + + def readline(self): + request = self.proc.requests[-1] + if self.proc.payloads: + payload = self.proc.payloads.pop(0) + else: + payload = { + "provider": request["provider"], + "source": "usage_api", + "title": "Account limits", + "windows": [], + "details": [], + "available": True, + "unavailable_reason": None, + "fetched_at": "2030-03-17T12:30:00Z", + } + return json.dumps(payload) + "\n" + + def close(self): + self.closed = True + + def poll(self): + return self.returncode + + def terminate(self): + self.terminated = True + self.returncode = -15 + + def wait(self, timeout=None): + return self.returncode + + def kill(self): + self.killed = True + self.returncode = -9 + + +def test_account_usage_worker_reuses_process_for_same_home(monkeypatch, tmp_path): + import api.providers as providers + import subprocess + + launched = [] + + def fake_popen(*args, **kwargs): + proc = _FakeAccountUsageWorkerProcess() + launched.append((args, kwargs, proc)) + return proc + + monkeypatch.setattr(subprocess, "Popen", fake_popen) + providers._close_account_usage_probe_workers() + try: + first = providers._agent_fetch_account_usage_for_home("openai-codex", tmp_path) + second = providers._agent_fetch_account_usage_for_home("anthropic", tmp_path, api_key="sk-test") + finally: + providers._close_account_usage_probe_workers() + + assert len(launched) == 1 + assert first.provider == "openai-codex" + assert second.provider == "anthropic" + assert launched[0][2].requests == [ + {"provider": "openai-codex", "api_key": "", "env_var": None}, + {"provider": "anthropic", "api_key": "sk-test", "env_var": "ANTHROPIC_API_KEY"}, + ] + + +def test_account_usage_worker_pool_is_keyed_by_home(monkeypatch, tmp_path): + import api.providers as providers + import subprocess + + launched = [] + + def fake_popen(*args, **kwargs): + proc = _FakeAccountUsageWorkerProcess() + launched.append((args, kwargs, proc)) + return proc + + home_a = tmp_path / "a" + home_b = tmp_path / "b" + monkeypatch.setattr(subprocess, "Popen", fake_popen) + providers._close_account_usage_probe_workers() + try: + providers._agent_fetch_account_usage_for_home("openai-codex", home_a) + providers._agent_fetch_account_usage_for_home("openai-codex", home_b) + finally: + providers._close_account_usage_probe_workers() + + assert len(launched) == 2 + + +def test_account_usage_worker_idle_cleanup_closes_stale_process(monkeypatch, tmp_path): + import api.providers as providers + import subprocess + + launched = [] + + def fake_popen(*args, **kwargs): + proc = _FakeAccountUsageWorkerProcess() + launched.append((args, kwargs, proc)) + return proc + + monkeypatch.setattr(subprocess, "Popen", fake_popen) + providers._close_account_usage_probe_workers() + try: + providers._agent_fetch_account_usage_for_home("openai-codex", tmp_path) + worker = providers._account_usage_worker_pool[str(tmp_path)] + providers._cleanup_account_usage_probe_workers( + now=worker.last_used + providers._ACCOUNT_USAGE_WORKER_IDLE_SECONDS + 1 + ) + providers._agent_fetch_account_usage_for_home("openai-codex", tmp_path) + finally: + providers._close_account_usage_probe_workers() + + assert len(launched) == 2 + assert launched[0][2].terminated is True + + +def test_busy_account_usage_worker_uses_one_shot_fallback(monkeypatch, tmp_path): + import api.providers as providers + + worker = providers._AccountUsageProbeWorker(tmp_path) + calls = [] + + def fake_one_shot(provider, home, *, api_key=None): + calls.append((provider, Path(home), api_key)) + return SimpleNamespace(provider=provider, source="usage_api", windows=(), details=(), available=True) + + monkeypatch.setattr(providers, "_fetch_account_usage_once_for_home", fake_one_shot) + locked = threading.Event() + release = threading.Event() + + def hold_lock(): + with worker._lock: + locked.set() + release.wait(timeout=5) + + holder = threading.Thread(target=hold_lock) + holder.start() + assert locked.wait(timeout=5) + try: + snapshot = worker.fetch("anthropic", api_key="sk-test") + finally: + release.set() + holder.join(timeout=5) + worker.close() + + assert snapshot.provider == "anthropic" + assert calls == [("anthropic", tmp_path, "sk-test")] + + +def test_account_usage_cleanup_removes_null_proc_worker(monkeypatch, tmp_path): + import api.providers as providers + import subprocess + + launched = [] + + def fake_popen(*args, **kwargs): + proc = _FakeAccountUsageWorkerProcess() + launched.append(proc) + return proc + + monkeypatch.setattr(subprocess, "Popen", fake_popen) + providers._close_account_usage_probe_workers() + try: + providers._agent_fetch_account_usage_for_home("openai-codex", tmp_path) + worker = providers._account_usage_worker_pool[str(tmp_path)] + worker.close() + providers._cleanup_account_usage_probe_workers() + assert str(tmp_path) not in providers._account_usage_worker_pool + finally: + providers._close_account_usage_probe_workers() + + assert len(launched) == 1 + + +def test_provider_key_mutation_invalidates_warm_account_usage_workers(monkeypatch, tmp_path): + import api.providers as providers + + invalidated = [] + + monkeypatch.setattr(providers, "_get_hermes_home", lambda: tmp_path) + monkeypatch.setattr(providers, "invalidate_models_cache", lambda: None) + monkeypatch.setattr( + providers, + "invalidate_account_usage_status_cache", + lambda provider_id=None: invalidated.append(provider_id), + ) + + updated = providers.set_provider_key("anthropic", "sk-test-quota-worker") + removed = providers.set_provider_key("anthropic", None) + + assert updated["ok"] is True + assert removed["ok"] is True + assert invalidated == ["anthropic", "anthropic"] + + +def test_account_usage_worker_uses_controlled_pipe_stdin(monkeypatch): + """Account-usage probe workers must not inherit process stdin.""" import api.providers as providers import subprocess seen_stdin = None - def capturing_run(*args, **kwargs): + def capturing_popen(*args, **kwargs): nonlocal seen_stdin seen_stdin = kwargs.get('stdin') - class FakeProc: - returncode = 0 - stdout = '{}' - stderr = '' - return FakeProc() + return _FakeAccountUsageWorkerProcess() - monkeypatch.setattr(subprocess, 'run', capturing_run) + monkeypatch.setattr(subprocess, 'Popen', capturing_popen) + providers._close_account_usage_probe_workers() try: providers._agent_fetch_account_usage_for_home( 'openai-codex', Path('/nonexistent'), api_key=None ) - except Exception: - pass # errors are expected on a fake env; we only care about stdin + finally: + providers._close_account_usage_probe_workers() - assert seen_stdin is subprocess.DEVNULL, ( - f'expected stdin=subprocess.DEVNULL, got {seen_stdin!r}' + assert seen_stdin is subprocess.PIPE, ( + f'expected stdin=subprocess.PIPE, got {seen_stdin!r}' ) @@ -1387,24 +1602,21 @@ def test_account_usage_preexec_fn_is_wired_on_posix(monkeypatch): captured_kwargs = {} - def capture_run(*args, **kwargs): + def capture_popen(*args, **kwargs): captured_kwargs.update(kwargs) - class FakeProc: - returncode = 0 - stdout = '{}' - stderr = '' - return FakeProc() + return _FakeAccountUsageWorkerProcess() - monkeypatch.setattr(subprocess, 'run', capture_run) + monkeypatch.setattr(subprocess, 'Popen', capture_popen) + providers._close_account_usage_probe_workers() try: providers._agent_fetch_account_usage_for_home( 'openai-codex', Path('/nonexistent'), api_key=None ) - except Exception: - pass + finally: + providers._close_account_usage_probe_workers() assert 'preexec_fn' in captured_kwargs, ( - 'preexec_fn should be in subprocess.run kwargs on POSIX' + 'preexec_fn should be in subprocess.Popen kwargs on POSIX' ) assert captured_kwargs['preexec_fn'] is providers._account_usage_preexec_fn @@ -1421,12 +1633,27 @@ def test_account_usage_semaphore_caps_concurrency(monkeypatch, tmp_path): monkeypatch.setattr(profiles, 'get_active_hermes_home', lambda: tmp_path) old_cfg, old_mtime = _with_config(model={'provider': 'openai-codex'}) - barrier = threading.Barrier(2, timeout=2) + active = 0 + max_active = 0 + entered = 0 + first_two_entered = threading.Event() + third_entered = threading.Event() + lock = threading.Lock() unblock = threading.Event() def slow_fetch(provider, home, api_key=None): - barrier.wait() + nonlocal active, max_active, entered + with lock: + active += 1 + entered += 1 + max_active = max(max_active, active) + if entered == providers._MAX_CONCURRENT_ACCOUNT_USAGE_PROBES: + first_two_entered.set() + if entered > providers._MAX_CONCURRENT_ACCOUNT_USAGE_PROBES: + third_entered.set() unblock.wait(timeout=5) + with lock: + active -= 1 return None monkeypatch.setattr(providers, '_agent_fetch_account_usage_for_home', slow_fetch) @@ -1442,16 +1669,22 @@ def test_account_usage_semaphore_caps_concurrency(monkeypatch, tmp_path): except Exception as exc: errors.append(exc) - threads = [threading.Thread(target=worker) for _ in range(2)] + threads = [ + threading.Thread(target=worker) + for _ in range(providers._MAX_CONCURRENT_ACCOUNT_USAGE_PROBES + 1) + ] for t in threads: t.start() - for t in threads: - t.join(timeout=10) - - unblock.set() try: + assert first_two_entered.wait(timeout=2), 'first probe batch did not start' + assert not third_entered.wait(timeout=0.2), 'third probe bypassed semaphore' + unblock.set() + for t in threads: + t.join(timeout=10) assert not errors, f'workers raised: {errors}' - assert len(results) == 2, f'expected 2 results, got {len(results)}' + assert len(results) == len(threads), f'expected {len(threads)} results, got {len(results)}' + assert max_active <= providers._MAX_CONCURRENT_ACCOUNT_USAGE_PROBES finally: + unblock.set() _restore_config(old_cfg, old_mtime)