Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 110 additions & 7 deletions api/providers.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@
import json
import logging
import os
import signal
import subprocess
import sys
import threading
import urllib.error
import urllib.request
from datetime import datetime, timezone
Expand All @@ -38,6 +40,76 @@
_PROVIDER_QUOTA_TIMEOUT_SECONDS = 3.0
_ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS = 35.0
_ACCOUNT_USAGE_PROVIDERS = frozenset({"openai-codex", "anthropic"})

# Upper bound on simultaneous profile-isolated quota probe subprocesses.
# Each probe runs a Python child for up to 35 s; capping concurrency prevents
# resource exhaustion when the UI polls all providers rapidly. The limit is
# deliberately low (2) since _ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS is
# already 35 s and probe I/O is lightweight HTTP calls.
_MAX_CONCURRENT_ACCOUNT_USAGE_PROBES = 2

# Parent-death-signal setup: on Linux, arrange for the quota-probe child to
# receive SIGTERM when the WebUI parent dies (e.g. systemctl restart, OOM kill).
# This prevents probe children from becoming orphaned zombies that continue
# calling the provider API indefinitely after the WebUI process is gone.
# We use prctl(PR_SET_DEATHSIG, SIGTERM) which is standard on modern Linux
# kernels and available via ctypes (no external C extension needed).
# If prctl is unavailable (non-Linux, or Linux without prctl support), the
# probe child exits normally when its parent (WebUI) terminates -- on macOS/
# Windows this is handled by OS-level process tree cleanup.
# Portable parent-death-signal bootstrap. On Linux this arranges for the
# probe child to receive SIGTERM when the WebUI parent dies (systemctl
# restart, OOM kill, etc.), preventing orphaned zombie probes from continuing
# to call the provider API indefinitely. Non-Linux platforms (macOS, Windows)
# rely on OS-level process-tree cleanup instead; this variable is then unused.
# prctl(PR_SET_DEATHSIG, SIGTERM) is available via ctypes without any C
# extension — the same technique used throughout the Hermes codebase.
_ACCOUNT_USAGE_PARENT_DEATHSIG_BOOTSTRAP = (
# fmt: off
# Lines are written as string literals so this block passes
# `python3 -m py_compile` cleanly and is safe to include verbatim
# inside the single argument string passed to `python -c ...`.
'import sys\n'
'try:\n'
' import ctypes, signal\n'
' libc = ctypes.CDLL(None)\n'
' libc.prctl(1, signal.SIGTERM) # PR_SET_DEATHSIG=1, SIGTERM=15\n'
'except Exception:\n'
' pass\n'
# fmt: on
)


# Module-level cap on concurrent quota-probe subprocesses.
# Lazily created so this module compiles even when threading isn't ready.
_account_usage_probe_semaphore: threading.BoundedSemaphore | None = None


def _get_account_usage_probe_semaphore() -> threading.BoundedSemaphore:
global _account_usage_probe_semaphore
if _account_usage_probe_semaphore is None:
_account_usage_probe_semaphore = threading.BoundedSemaphore(
_MAX_CONCURRENT_ACCOUNT_USAGE_PROBES
)
return _account_usage_probe_semaphore


# ── preexec_fn: parent-death signal for the probe subprocess ─────────────────
# On POSIX/Linux, arrange for the child to receive SIGTERM when the WebUI
# parent dies (systemctl restart, OOM kill, etc.). The parent's bootstrap
# 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.
def _account_usage_preexec_fn() -> None:
try:
import ctypes
libc = ctypes.CDLL(None)
libc.prctl(1, signal.SIGTERM) # PR_SET_PDEATHSIG=1, SIGTERM=15
except Exception:
pass


_ACCOUNT_USAGE_SUBPROCESS_CODE = r"""
import json
import sys
Expand Down Expand Up @@ -533,14 +605,29 @@ def _agent_fetch_account_usage_for_home(provider: str, home: Path, *, api_key: s
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_SUBPROCESS_CODE, provider, api_key or ""],
[
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),
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
timeout=_ACCOUNT_USAGE_SUBPROCESS_TIMEOUT_SECONDS,
check=False,
**kwargs,
)
except subprocess.TimeoutExpired:
logger.debug("Account usage probe for %s timed out", provider)
Expand All @@ -561,10 +648,26 @@ def _agent_fetch_account_usage_for_home(provider: str, home: Path, *, api_key: s


def _fetch_account_usage_with_profile_context(provider: str) -> Any:
"""Fetch account usage for a provider within the active profile context.

Concurrency is capped by the module-level BoundedSemaphore so that rapid
UI polls (e.g. Settings page refresh) cannot exhaust file-descriptors or
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.
"""
home = _get_hermes_home()
api_key = _get_provider_api_key(provider)
sem = _get_account_usage_probe_semaphore()
try:
return _agent_fetch_account_usage_for_home(provider, home, api_key=api_key)
with sem:
return _agent_fetch_account_usage_for_home(
provider,
home,
api_key=api_key,
)
except Exception:
logger.debug("Failed to fetch account usage for %s", provider, exc_info=True)
return None
Expand Down
152 changes: 152 additions & 0 deletions tests/test_provider_quota_status.py
Original file line number Diff line number Diff line change
Expand Up @@ -413,3 +413,155 @@ def test_provider_quota_styles_exist():
".provider-quota-window",
):
assert token in css


# ── 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.
"""
import api.providers as providers
import subprocess

seen_stdin = None

def capturing_run(*args, **kwargs):
nonlocal seen_stdin
seen_stdin = kwargs.get('stdin')
class FakeProc:
returncode = 0
stdout = '{}'
stderr = ''
return FakeProc()

monkeypatch.setattr(subprocess, 'run', capturing_run)
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

assert seen_stdin is subprocess.DEVNULL, (
f'expected stdin=subprocess.DEVNULL, got {seen_stdin!r}'
)


def test_account_usage_probe_semaphore_has_correct_bound(monkeypatch, tmp_path):
"""The probe semaphore must enforce the declared concurrency cap.

Verifying the bound directly ensures the cap actually prevents resource
exhaustion when the UI polls multiple providers in rapid succession.
"""
import api.providers as providers

monkeypatch.setattr(profiles, 'get_active_hermes_home', lambda: tmp_path)
old_cfg, old_mtime = _with_config(model={'provider': 'openai-codex'})

sem = providers._get_account_usage_probe_semaphore()
try:
bound = sem._value
assert bound == providers._MAX_CONCURRENT_ACCOUNT_USAGE_PROBES, (
f'semaphore bound is {bound}, expected '
f'{providers._MAX_CONCURRENT_ACCOUNT_USAGE_PROBES}'
)
finally:
_restore_config(old_cfg, old_mtime)


def test_account_usage_preexec_fn_is_wired_on_posix(monkeypatch):
"""On POSIX systems the probe subprocess must receive a parent-death preexec_fn.

The preexec_fn arranges prctl(PR_SET_PDEATHSIG, SIGTERM) so the child is
terminated when the WebUI parent dies (OOM kill, systemctl restart, etc.).
This test verifies the wiring and skips harmlessly on non-POSIX (Windows).
"""
import api.providers as providers

assert callable(providers._account_usage_preexec_fn)

try:
providers._account_usage_preexec_fn()
except Exception as exc:
raise AssertionError(
f'_account_usage_preexec_fn raised {exc!r}; it should be '
'safe to call unconditionally'
) from exc

if hasattr(os, 'fork'):
import subprocess

captured_kwargs = {}

def capture_run(*args, **kwargs):
captured_kwargs.update(kwargs)
class FakeProc:
returncode = 0
stdout = '{}'
stderr = ''
return FakeProc()

monkeypatch.setattr(subprocess, 'run', capture_run)
try:
providers._agent_fetch_account_usage_for_home(
'openai-codex', Path('/nonexistent'), api_key=None
)
except Exception:
pass

assert 'preexec_fn' in captured_kwargs, (
'preexec_fn should be in subprocess.run kwargs on POSIX'
)
assert captured_kwargs['preexec_fn'] is providers._account_usage_preexec_fn


def test_account_usage_semaphore_caps_concurrency(monkeypatch, tmp_path):
"""The probe semaphore must actually serialise callers beyond its bound.

Verifies the bounded semaphore is used in the call path and genuinely
prevents more than _MAX_CONCURRENT_ACCOUNT_USAGE_PROBES probes running.
"""
import api.providers as providers
import threading

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)
unblock = threading.Event()

def slow_fetch(provider, home, api_key=None):
barrier.wait()
unblock.wait(timeout=5)
return None

monkeypatch.setattr(providers, '_agent_fetch_account_usage_for_home', slow_fetch)

results = []
errors = []

def worker():
try:
results.append(
providers._fetch_account_usage_with_profile_context('openai-codex')
)
except Exception as exc:
errors.append(exc)

threads = [threading.Thread(target=worker) for _ in range(2)]
for t in threads:
t.start()
for t in threads:
t.join(timeout=10)

unblock.set()

try:
assert not errors, f'workers raised: {errors}'
assert len(results) == 2, f'expected 2 results, got {len(results)}'
finally:
_restore_config(old_cfg, old_mtime)
Loading