diff --git a/services/hermes/agent-deployment.yaml b/services/hermes/agent-deployment.yaml index 845eb6f9..526c5c55 100644 --- a/services/hermes/agent-deployment.yaml +++ b/services/hermes/agent-deployment.yaml @@ -25,7 +25,7 @@ spec: ai.bstein.dev/execution: Hermes Kanban with durable direct Codex and Claude Code CLI workers ai.bstein.dev/model-policy: Jetson-assisted AUTO routing, low through xhigh, cross-provider fallback ai.bstein.dev/placement: rpi5 preferred; Jetson deferred until state storage is available - ai.bstein.dev/config-rev: "20260816-auth-health-quarantine-v6" + ai.bstein.dev/config-rev: "20260816-ai-usage-poller-v1" prometheus.io/scrape: "true" prometheus.io/path: /metrics prometheus.io/port: "9010" diff --git a/services/hermes/scripts/ai_usage_exporter.py b/services/hermes/scripts/ai_usage_exporter.py index 178c2c54..16be50a6 100644 --- a/services/hermes/scripts/ai_usage_exporter.py +++ b/services/hermes/scripts/ai_usage_exporter.py @@ -48,6 +48,9 @@ METRIC_HELP = { "atlas_ai_quota_used_percent": "Used percentage in a first-party coding CLI quota window.", "atlas_ai_quota_window_duration_seconds": "Nominal duration of a coding CLI quota window.", } +CODEX_CLEANUP_TIMEOUT_SECONDS = 5 +POLL_STARTUP_GRACE_SECONDS = 45 +POLL_PROGRESS_BUDGET_SECONDS = 60 @dataclass(frozen=True) @@ -228,6 +231,63 @@ def parse_claude_payload(payload: dict[str, Any]) -> list[Sample]: return samples +def _close_stream(stream: Any) -> bool: + """Close a subprocess pipe without allowing cleanup errors to escape.""" + if stream is None: + return True + try: + stream.close() + except Exception: + return False + return True + + +def _cleanup_codex_process( + process: subprocess.Popen[str], selector: selectors.BaseSelector | None +) -> bool: + """Stop and reap a Codex child within two fixed waits, then close its pipes.""" + cleanup_ok = True + if selector is not None: + try: + selector.close() + except Exception: + cleanup_ok = False + + # Closing the request pipe first also gives a responsive app-server an EOF. + cleanup_ok = _close_stream(getattr(process, "stdin", None)) and cleanup_ok + try: + process.terminate() + except ProcessLookupError: + pass + except Exception: + cleanup_ok = False + + needs_kill = False + try: + process.wait(timeout=CODEX_CLEANUP_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + needs_kill = True + except Exception: + cleanup_ok = False + needs_kill = True + + if needs_kill: + try: + process.kill() + except ProcessLookupError: + pass + except Exception: + cleanup_ok = False + try: + process.wait(timeout=CODEX_CLEANUP_TIMEOUT_SECONDS) + except Exception: + cleanup_ok = False + + cleanup_ok = _close_stream(getattr(process, "stdout", None)) and cleanup_ok + cleanup_ok = _close_stream(getattr(process, "stderr", None)) and cleanup_ok + return cleanup_ok + + def query_codex(timeout: float = 20) -> tuple[dict[str, Any], dict[str, Any]]: """Read Codex account quota and usage through its structured app-server protocol.""" process = subprocess.Popen( @@ -256,6 +316,9 @@ def query_codex(timeout: float = 20) -> tuple[dict[str, Any], dict[str, Any]]: {"id": 2, "method": "account/rateLimits/read", "params": None}, {"id": 3, "method": "account/usage/read", "params": None}, ) + selector: selectors.BaseSelector | None = None + rate_result: dict[str, Any] + usage_result: dict[str, Any] try: if process.stdin is None or process.stdout is None: raise RuntimeError("Codex app-server pipes are unavailable") @@ -278,14 +341,13 @@ def query_codex(timeout: float = 20) -> tuple[dict[str, Any], dict[str, Any]]: response = responses.get(response_id) if not response or "error" in response or not isinstance(response.get("result"), dict): raise RuntimeError(f"Codex app-server response {response_id} failed") - return responses[2]["result"], responses[3]["result"] + rate_result = responses[2]["result"] + usage_result = responses[3]["result"] finally: - process.terminate() - try: - process.wait(timeout=5) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=5) + cleanup_ok = _cleanup_codex_process(process, selector) + if not cleanup_ok: + raise RuntimeError("Codex app-server cleanup failed") + return rate_result, usage_result def query_claude() -> dict[str, Any]: @@ -322,6 +384,22 @@ class Collector: self._lock = threading.Lock() self._providers = {name: ProviderState() for name in ("openai", "anthropic")} + def record_failure( + self, + provider: str, + *, + started: float | None = None, + monotonic_started: float | None = None, + ) -> None: + """Record a failed attempt while retaining the provider's last good samples.""" + attempted_at = time.time() if started is None else started + duration_started = time.monotonic() if monotonic_started is None else monotonic_started + with self._lock: + state = self._providers[provider] + state.last_attempt = attempted_at + state.duration = max(0, time.monotonic() - duration_started) + state.fetch_success = False + def refresh_provider(self, provider: str) -> None: """Refresh one provider while retaining the last good values on failure.""" started = time.time() @@ -333,19 +411,21 @@ class Collector: samples = parse_claude_payload(query_claude()) else: raise ValueError("unknown provider") - success = True - except (OSError, RuntimeError, ValueError, json.JSONDecodeError) as error: + except Exception as error: print(f"{provider} quota collection deferred: {type(error).__name__}", flush=True) - samples = [] - success = False + self.record_failure( + provider, + started=started, + monotonic_started=monotonic_started, + ) + return with self._lock: state = self._providers[provider] state.last_attempt = started state.duration = time.monotonic() - monotonic_started - state.fetch_success = success - if success: - state.samples = samples - state.last_success = time.time() + state.fetch_success = True + state.samples = samples + state.last_success = time.time() def render(self) -> bytes: """Render the current provider states in Prometheus text format.""" @@ -381,38 +461,149 @@ class Collector: return ("\n".join(lines) + "\n").encode("utf-8") +class PollingEngine: + """Run isolated provider polls and track bounded forward progress.""" + + def __init__( + self, + collector: Collector, + *, + interval: float, + startup_grace: float = POLL_STARTUP_GRACE_SECONDS, + progress_timeout: float | None = None, + clock: Any = time.monotonic, + ) -> None: + self.collector = collector + self.interval = interval + self.startup_grace = startup_grace + self.progress_timeout = ( + interval + POLL_PROGRESS_BUDGET_SECONDS + if progress_timeout is None + else progress_timeout + ) + self._clock = clock + self._started_at = clock() + self._last_progress: float | None = None + self._progress_lock = threading.Lock() + self._thread: threading.Thread | None = None + + def _mark_progress(self) -> None: + """Record completion of one bounded provider attempt.""" + with self._progress_lock: + self._last_progress = self._clock() + + def poll_once(self) -> None: + """Refresh every provider even if another provider fails unexpectedly.""" + for provider in ("openai", "anthropic"): + started = time.time() + monotonic_started = time.monotonic() + try: + self.collector.refresh_provider(provider) + except Exception as error: + print( + f"{provider} quota collection isolated: {type(error).__name__}", + flush=True, + ) + try: + self.collector.record_failure( + provider, + started=started, + monotonic_started=monotonic_started, + ) + except Exception as record_error: + print( + f"{provider} quota failure accounting deferred: " + f"{type(record_error).__name__}", + flush=True, + ) + finally: + self._mark_progress() + + def run(self) -> None: + """Poll forever without allowing a cycle-level exception to stop the thread.""" + while True: + try: + self.poll_once() + except Exception as error: + print(f"quota polling cycle deferred: {type(error).__name__}", flush=True) + time.sleep(self.interval) + + def start(self) -> None: + """Start the daemon poller exactly once.""" + if self._thread is not None: + return + self._thread = threading.Thread( + target=self.run, + name="ai-usage-poller", + daemon=True, + ) + self._thread.start() + + def is_healthy(self) -> bool: + """Report thread liveness and progress, independent of provider success.""" + thread = self._thread + if thread is None or not thread.is_alive(): + return False + now = self._clock() + with self._progress_lock: + last_progress = self._last_progress + if last_progress is None: + return now - self._started_at <= self.startup_grace + return now - last_progress <= self.progress_timeout + + +def make_handler(collector: Collector, poller: PollingEngine) -> type[BaseHTTPRequestHandler]: + """Build an HTTP handler bound to one collector and polling engine.""" + + class Handler(BaseHTTPRequestHandler): + def _respond(self, status: int, payload: bytes, content_type: str) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def do_GET(self) -> None: # noqa: N802 + try: + if self.path == "/metrics": + self._respond( + 200, + collector.render(), + "text/plain; version=0.0.4", + ) + elif self.path == "/healthz": + healthy = poller.is_healthy() + self._respond( + 200 if healthy else 503, + b"ok\n" if healthy else b"poller unhealthy\n", + "text/plain; charset=utf-8", + ) + else: + self.send_error(404) + except (BrokenPipeError, ConnectionResetError): + return + + def log_message(self, _format: str, *_args: object) -> None: + return + + return Handler + + +class Server(ThreadingHTTPServer): + """Threaded metrics server that does not retain disconnected clients.""" + + daemon_threads = True + + def main() -> int: """Poll quota APIs and serve only sanitized metrics and health endpoints.""" collector = Collector() interval = max(60, int(os.environ.get("ATLAS_AI_USAGE_INTERVAL_SECONDS", "300"))) - def polling_loop() -> None: - while True: - for provider in ("openai", "anthropic"): - collector.refresh_provider(provider) - time.sleep(interval) - - class Handler(BaseHTTPRequestHandler): - def do_GET(self) -> None: # noqa: N802 - if self.path == "/metrics": - payload = collector.render() - self.send_response(200) - self.send_header("Content-Type", "text/plain; version=0.0.4") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) - elif self.path == "/healthz": - self.send_response(200) - self.end_headers() - else: - self.send_error(404) - - def log_message(self, _format: str, *_args: object) -> None: - return - - threading.Thread(target=polling_loop, daemon=True).start() + poller = PollingEngine(collector, interval=interval) + poller.start() port = int(os.environ.get("ATLAS_AI_USAGE_PORT", "9010")) - server = ThreadingHTTPServer(("0.0.0.0", port), Handler) + server = Server(("0.0.0.0", port), make_handler(collector, poller)) server.serve_forever() return 0 diff --git a/testing/tests/test_hermes_ai_usage_exporter.py b/testing/tests/test_hermes_ai_usage_exporter.py index 9b78b7f3..9fe54d72 100644 --- a/testing/tests/test_hermes_ai_usage_exporter.py +++ b/testing/tests/test_hermes_ai_usage_exporter.py @@ -4,9 +4,12 @@ from __future__ import annotations import importlib.util import json +import subprocess import sys from datetime import date +from io import BytesIO, StringIO from pathlib import Path +from types import SimpleNamespace ROOT = Path(__file__).resolve().parents[2] @@ -173,3 +176,218 @@ for response_id, result in responses.items(): "account/usage/read", ] assert all("/status" not in json.dumps(request) for request in requests) + + +def test_codex_cleanup_is_bounded_when_process_cannot_be_reaped(monkeypatch): + mod = load_module() + processes = [] + selectors = [] + + class StubbornProcess: + def __init__(self, *_args, **_kwargs): + self.stdin = StringIO() + self.stdout = StringIO( + "\n".join( + json.dumps({"id": response_id, "result": result}) + for response_id, result in ( + (1, {}), + (2, {"rateLimits": {}}), + (3, {"summary": {}, "dailyUsageBuckets": []}), + ) + ) + + "\n" + ) + self.terminate_calls = 0 + self.kill_calls = 0 + self.wait_timeouts = [] + processes.append(self) + + def terminate(self): + self.terminate_calls += 1 + + def kill(self): + self.kill_calls += 1 + + def wait(self, timeout): + self.wait_timeouts.append(timeout) + raise subprocess.TimeoutExpired("codex-mock", timeout) + + class TrackingSelector: + def __init__(self): + self.fileobj = None + self.closed = False + selectors.append(self) + + def register(self, fileobj, _events): + self.fileobj = fileobj + + def select(self, _timeout): + return [(SimpleNamespace(fileobj=self.fileobj), None)] + + def close(self): + self.closed = True + + monkeypatch.setattr(mod.subprocess, "Popen", StubbornProcess) + monkeypatch.setattr(mod.selectors, "DefaultSelector", TrackingSelector) + + try: + mod.query_codex(timeout=1) + except Exception as error: # The cleanup result must be sanitized and non-TimeoutExpired. + assert type(error) is RuntimeError + assert str(error) == "Codex app-server cleanup failed" + else: + raise AssertionError("unreaped Codex process was reported as successful") + + process = processes[0] + assert process.terminate_calls == 1 + assert process.kill_calls == 1 + assert process.wait_timeouts == [5, 5] + assert process.stdin.closed + assert process.stdout.closed + assert selectors[0].closed + + +def test_unexpected_provider_failure_is_sanitized_and_preserves_last_good(monkeypatch, capsys): + mod = load_module() + collector = mod.Collector() + previous = mod.Sample( + "atlas_ai_quota_used_percent", + {"provider": "openai", "limit": "codex", "window": "seven_day"}, + 42, + ) + collector._providers["openai"].samples = [previous] + secret = "provider-body-must-not-appear" + + class UnexpectedProviderFailure(Exception): + pass + + def fail(): + raise UnexpectedProviderFailure(secret) + + monkeypatch.setattr(mod, "query_codex", fail) + collector.refresh_provider("openai") + + output = capsys.readouterr().out + state = collector._providers["openai"] + assert state.samples == [previous] + assert state.fetch_success is False + assert state.last_attempt > 0 + assert "UnexpectedProviderFailure" in output + assert secret not in output + + +def test_polling_engine_isolates_providers_and_future_cycles(monkeypatch, capsys): + mod = load_module() + collector = mod.Collector() + calls = [] + secret = "cleanup-detail-must-not-appear" + + class UnexpectedCleanupFailure(Exception): + pass + + def refresh(provider): + calls.append(provider) + if provider == "openai" and calls.count("openai") == 1: + raise UnexpectedCleanupFailure(secret) + + monkeypatch.setattr(collector, "refresh_provider", refresh) + engine = mod.PollingEngine(collector, interval=60) + + engine.poll_once() + engine.poll_once() + + output = capsys.readouterr().out + assert calls == ["openai", "anthropic", "openai", "anthropic"] + assert collector._providers["openai"].fetch_success is False + assert collector._providers["openai"].last_attempt > 0 + assert "UnexpectedCleanupFailure" in output + assert secret not in output + + +def test_health_tracks_poller_liveness_and_progress_not_provider_success(): + mod = load_module() + now = [100.0] + collector = mod.Collector() + collector._providers["openai"].fetch_success = False + engine = mod.PollingEngine( + collector, + interval=60, + startup_grace=10, + progress_timeout=75, + clock=lambda: now[0], + ) + + class ThreadState: + alive = True + + def is_alive(self): + return self.alive + + thread = ThreadState() + engine._thread = thread + assert engine.is_healthy() + + now[0] = 111.0 + assert not engine.is_healthy() + + engine._mark_progress() + assert engine.is_healthy() + now[0] = 187.0 + assert not engine.is_healthy() + + thread.alive = False + now[0] = 100.0 + assert not engine.is_healthy() + + +def test_health_endpoint_uses_poller_state_and_ignores_provider_failure(): + mod = load_module() + collector = mod.Collector() + + class Poller: + healthy = False + + def is_healthy(self): + return self.healthy + + poller = Poller() + handler_class = mod.make_handler(collector, poller) + handler = object.__new__(handler_class) + handler.path = "/healthz" + statuses = [] + handler.send_response = statuses.append + handler.send_header = lambda _name, _value: None + handler.end_headers = lambda: None + handler.wfile = BytesIO() + + handler.do_GET() + assert statuses == [503] + + poller.healthy = True + statuses.clear() + collector._providers["openai"].fetch_success = False + handler.do_GET() + assert statuses == [200] + + +def test_metrics_response_tolerates_disconnected_client(): + mod = load_module() + collector = mod.Collector() + + class HealthyPoller: + def is_healthy(self): + return True + + class Disconnected(BytesIO): + def write(self, _body): + raise BrokenPipeError + + handler_class = mod.make_handler(collector, HealthyPoller()) + handler = object.__new__(handler_class) + handler.path = "/metrics" + handler.send_response = lambda _status: None + handler.send_header = lambda _name, _value: None + handler.end_headers = lambda: None + handler.wfile = Disconnected() + + handler.do_GET()