From ed43dbaa7f094e1841aa760862e80015deb9d80c Mon Sep 17 00:00:00 2001 From: Hermes Agent Date: Mon, 17 Aug 2026 15:33:56 +0000 Subject: [PATCH] hermes: split AI usage exporter engines --- services/hermes/agent-deployment.yaml | 2 +- services/hermes/kustomization.yaml | 2 + services/hermes/scripts/ai_usage_codex.py | 138 +++++++++ services/hermes/scripts/ai_usage_exporter.py | 244 +++++----------- services/hermes/scripts/ai_usage_http.py | 49 ++++ testing/tests/test_hermes_ai_usage_codex.py | 255 ++++++++++++++++ .../tests/test_hermes_ai_usage_exporter.py | 163 +++-------- .../test_hermes_ai_usage_exporter_coverage.py | 275 ++++++++++++++++++ testing/tests/test_hermes_auto_router.py | 140 +++++---- 9 files changed, 903 insertions(+), 365 deletions(-) create mode 100644 services/hermes/scripts/ai_usage_codex.py create mode 100644 services/hermes/scripts/ai_usage_http.py create mode 100644 testing/tests/test_hermes_ai_usage_codex.py create mode 100644 testing/tests/test_hermes_ai_usage_exporter_coverage.py diff --git a/services/hermes/agent-deployment.yaml b/services/hermes/agent-deployment.yaml index 6cb07c34..84c0bcd2 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-ai-usage-poller-v1" + ai.bstein.dev/config-rev: "20260817-ai-usage-exporter-split-v2" prometheus.io/scrape: "true" prometheus.io/path: /metrics prometheus.io/port: "9010" diff --git a/services/hermes/kustomization.yaml b/services/hermes/kustomization.yaml index 3b9c2345..f7699e51 100644 --- a/services/hermes/kustomization.yaml +++ b/services/hermes/kustomization.yaml @@ -58,7 +58,9 @@ configMapGenerator: - name: hermes-coordinator namespace: hermes files: + - ai_usage_codex.py=scripts/ai_usage_codex.py - ai_usage_exporter.py=scripts/ai_usage_exporter.py + - ai_usage_http.py=scripts/ai_usage_http.py - claude=scripts/claude - claude_command_policy.py=scripts/claude_command_policy.py - cli_lane_goal.py=scripts/cli_lane_goal.py diff --git a/services/hermes/scripts/ai_usage_codex.py b/services/hermes/scripts/ai_usage_codex.py new file mode 100644 index 00000000..6214cbf5 --- /dev/null +++ b/services/hermes/scripts/ai_usage_codex.py @@ -0,0 +1,138 @@ +"""Bounded Codex app-server query and child-process cleanup.""" + +from __future__ import annotations + +import json +import os +import selectors +import subprocess +import time +from typing import Any + + +CODEX_BIN = os.environ.get("ATLAS_AI_CODEX_BIN", "/opt/data/tools/bin/codex") +CODEX_HOME = os.environ.get("CODEX_HOME", "/runtime-access/codex") +CODEX_CLEANUP_TIMEOUT_SECONDS = 5 + + +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( + [CODEX_BIN, "app-server", "--stdio"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + env={**os.environ, "CODEX_HOME": CODEX_HOME}, + ) + requests = ( + { + "id": 1, + "method": "initialize", + "params": { + "clientInfo": { + "name": "atlas-ai-usage-exporter", + "title": "Atlas AI Usage Exporter", + "version": "1.0.0", + }, + "capabilities": {"experimentalApi": True}, + }, + }, + {"method": "initialized", "params": {}}, + {"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") + for message in requests: + process.stdin.write(json.dumps(message, separators=(",", ":")) + "\n") + process.stdin.flush() + selector = selectors.DefaultSelector() + selector.register(process.stdout, selectors.EVENT_READ) + responses: dict[int, dict[str, Any]] = {} + deadline = time.monotonic() + timeout + while len(responses) < 3 and time.monotonic() < deadline: + for key, _ in selector.select(min(1, max(0, deadline - time.monotonic()))): + line = key.fileobj.readline() + if not line: + continue + message = json.loads(line) + if message.get("id") in (1, 2, 3): + responses[int(message["id"])] = message + for response_id in (1, 2, 3): + 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") + rate_result = responses[2]["result"] + usage_result = responses[3]["result"] + finally: + cleanup_ok = _cleanup_codex_process(process, selector) + if not cleanup_ok: + raise RuntimeError("Codex app-server cleanup failed") + return rate_result, usage_result diff --git a/services/hermes/scripts/ai_usage_exporter.py b/services/hermes/scripts/ai_usage_exporter.py index 16be50a6..b0db51fb 100644 --- a/services/hermes/scripts/ai_usage_exporter.py +++ b/services/hermes/scripts/ai_usage_exporter.py @@ -5,20 +5,18 @@ from __future__ import annotations import json import os -import selectors -import subprocess import threading import time from dataclasses import dataclass, field from datetime import UTC, date, datetime, timedelta -from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any from urllib.request import Request, urlopen +import ai_usage_codex as codex_query +import ai_usage_http as http_engine + -CODEX_BIN = os.environ.get("ATLAS_AI_CODEX_BIN", "/opt/data/tools/bin/codex") -CODEX_HOME = os.environ.get("CODEX_HOME", "/runtime-access/codex") CLAUDE_CREDENTIALS = Path( os.environ.get( "ATLAS_AI_CLAUDE_CREDENTIALS", @@ -48,7 +46,6 @@ 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 @@ -78,7 +75,9 @@ def _number(value: Any) -> float | None: if isinstance(value, bool) or not isinstance(value, (int, float)): return None converted = float(value) - return converted if converted == converted and abs(converted) != float("inf") else None + return ( + converted if converted == converted and abs(converted) != float("inf") else None + ) def _timestamp(value: Any) -> float | None: @@ -115,7 +114,10 @@ def _codex_limit_name(limit_id: str, snapshot: dict[str, Any]) -> str: def parse_codex_payloads( - rate_response: dict[str, Any], usage_response: dict[str, Any], *, today: date | None = None + rate_response: dict[str, Any], + usage_response: dict[str, Any], + *, + today: date | None = None, ) -> list[Sample]: """Convert structured Codex app-server responses into bounded metrics.""" samples: list[Sample] = [] @@ -143,15 +145,23 @@ def parse_codex_payloads( samples.extend( ( Sample("atlas_ai_quota_used_percent", labels, used), - Sample("atlas_ai_quota_remaining_percent", labels, max(0, 100 - used)), + Sample( + "atlas_ai_quota_remaining_percent", labels, max(0, 100 - used) + ), ) ) reset = _timestamp(raw_window.get("resetsAt")) duration = _number(raw_window.get("windowDurationMins")) if reset is not None: - samples.append(Sample("atlas_ai_quota_reset_timestamp_seconds", labels, reset)) + samples.append( + Sample("atlas_ai_quota_reset_timestamp_seconds", labels, reset) + ) if duration is not None: - samples.append(Sample("atlas_ai_quota_window_duration_seconds", labels, duration * 60)) + samples.append( + Sample( + "atlas_ai_quota_window_duration_seconds", labels, duration * 60 + ) + ) summary = usage_response.get("summary") if isinstance(summary, dict): @@ -192,8 +202,16 @@ def parse_codex_payloads( anchor_day = max((start for start, _ in parsed_daily), default=current_day) for period, days in (("latest_day", 1), ("seven_day", 7), ("thirty_day", 30)): earliest = anchor_day - timedelta(days=days - 1) - value = sum(tokens for start, tokens in parsed_daily if earliest <= start <= anchor_day) - samples.append(Sample("atlas_ai_account_tokens", {"provider": "openai", "period": period}, value)) + value = sum( + tokens for start, tokens in parsed_daily if earliest <= start <= anchor_day + ) + samples.append( + Sample( + "atlas_ai_account_tokens", + {"provider": "openai", "period": period}, + value, + ) + ) return samples @@ -216,7 +234,9 @@ def parse_claude_payload(payload: dict[str, Any]) -> list[Sample]: ) reset = _timestamp(window.get("resets_at")) if reset is not None: - samples.append(Sample("atlas_ai_quota_reset_timestamp_seconds", labels, reset)) + samples.append( + Sample("atlas_ai_quota_reset_timestamp_seconds", labels, reset) + ) extra = payload.get("extra_usage") if isinstance(extra, dict): enabled = extra.get("is_enabled") @@ -231,123 +251,7 @@ 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( - [CODEX_BIN, "app-server", "--stdio"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.DEVNULL, - text=True, - bufsize=1, - env={**os.environ, "CODEX_HOME": CODEX_HOME}, - ) - requests = ( - { - "id": 1, - "method": "initialize", - "params": { - "clientInfo": { - "name": "atlas-ai-usage-exporter", - "title": "Atlas AI Usage Exporter", - "version": "1.0.0", - }, - "capabilities": {"experimentalApi": True}, - }, - }, - {"method": "initialized", "params": {}}, - {"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") - for message in requests: - process.stdin.write(json.dumps(message, separators=(",", ":")) + "\n") - process.stdin.flush() - selector = selectors.DefaultSelector() - selector.register(process.stdout, selectors.EVENT_READ) - responses: dict[int, dict[str, Any]] = {} - deadline = time.monotonic() + timeout - while len(responses) < 3 and time.monotonic() < deadline: - for key, _ in selector.select(min(1, max(0, deadline - time.monotonic()))): - line = key.fileobj.readline() - if not line: - continue - message = json.loads(line) - if message.get("id") in (1, 2, 3): - responses[int(message["id"])] = message - for response_id in (1, 2, 3): - 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") - rate_result = responses[2]["result"] - usage_result = responses[3]["result"] - finally: - cleanup_ok = _cleanup_codex_process(process, selector) - if not cleanup_ok: - raise RuntimeError("Codex app-server cleanup failed") - return rate_result, usage_result +query_codex = codex_query.query_codex def query_claude() -> dict[str, Any]: @@ -393,7 +297,9 @@ class Collector: ) -> 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 + duration_started = ( + time.monotonic() if monotonic_started is None else monotonic_started + ) with self._lock: state = self._providers[provider] state.last_attempt = attempted_at @@ -412,7 +318,10 @@ class Collector: else: raise ValueError("unknown provider") except Exception as error: - print(f"{provider} quota collection deferred: {type(error).__name__}", flush=True) + print( + f"{provider} quota collection deferred: {type(error).__name__}", + flush=True, + ) self.record_failure( provider, started=started, @@ -440,10 +349,24 @@ class Collector: samples.extend(state.samples) samples.extend( ( - Sample("atlas_ai_quota_fetch_success", labels, float(state.fetch_success)), - Sample("atlas_ai_quota_last_attempt_timestamp_seconds", labels, state.last_attempt), - Sample("atlas_ai_quota_last_success_timestamp_seconds", labels, state.last_success), - Sample("atlas_ai_quota_fetch_duration_seconds", labels, state.duration), + Sample( + "atlas_ai_quota_fetch_success", + labels, + float(state.fetch_success), + ), + Sample( + "atlas_ai_quota_last_attempt_timestamp_seconds", + labels, + state.last_attempt, + ), + Sample( + "atlas_ai_quota_last_success_timestamp_seconds", + labels, + state.last_success, + ), + Sample( + "atlas_ai_quota_fetch_duration_seconds", labels, state.duration + ), ) ) lines: list[str] = [] @@ -525,7 +448,9 @@ class PollingEngine: try: self.poll_once() except Exception as error: - print(f"quota polling cycle deferred: {type(error).__name__}", flush=True) + print( + f"quota polling cycle deferred: {type(error).__name__}", flush=True + ) time.sleep(self.interval) def start(self) -> None: @@ -552,47 +477,8 @@ class PollingEngine: 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 +make_handler = http_engine.make_handler +Server = http_engine.Server def main() -> int: diff --git a/services/hermes/scripts/ai_usage_http.py b/services/hermes/scripts/ai_usage_http.py new file mode 100644 index 00000000..8558740f --- /dev/null +++ b/services/hermes/scripts/ai_usage_http.py @@ -0,0 +1,49 @@ +"""Disconnect-safe HTTP serving for the AI usage exporter.""" + +from __future__ import annotations + +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any + + +def make_handler(collector: Any, poller: Any) -> 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 diff --git a/testing/tests/test_hermes_ai_usage_codex.py b/testing/tests/test_hermes_ai_usage_codex.py new file mode 100644 index 00000000..a77525bf --- /dev/null +++ b/testing/tests/test_hermes_ai_usage_codex.py @@ -0,0 +1,255 @@ +"""Adversarial tests for bounded Codex quota-query process handling.""" + +from __future__ import annotations + +import importlib.util +import json +import subprocess +import sys +from io import StringIO +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +ROOT = Path(__file__).resolve().parents[2] +SCRIPT = ROOT / "services/hermes/scripts/ai_usage_codex.py" + + +def load_module(): + """Load the standalone Codex query helper as a fresh test module.""" + name = "ai_usage_codex_tested" + spec = importlib.util.spec_from_file_location(name, SCRIPT) + module = importlib.util.module_from_spec(spec) + assert spec and spec.loader + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def test_codex_query_uses_structured_app_server_protocol(tmp_path, monkeypatch): + mod = load_module() + request_log = tmp_path / "requests.jsonl" + mock = tmp_path / "codex-mock" + mock.write_text( + """#!/usr/bin/env python3 +import json, os, sys +messages = [json.loads(sys.stdin.readline()) for _ in range(4)] +with open(os.environ['REQUEST_LOG'], 'w') as handle: + for message in messages: + handle.write(json.dumps(message) + '\\n') +responses = { + 1: {}, + 2: {'rateLimits': {'primary': {'usedPercent': 10}}}, + 3: {'summary': {}, 'dailyUsageBuckets': []}, +} +for response_id, result in responses.items(): + print(json.dumps({'id': response_id, 'result': result}), flush=True) +""" + ) + mock.chmod(0o755) + monkeypatch.setattr(mod, "CODEX_BIN", str(mock)) + monkeypatch.setenv("REQUEST_LOG", str(request_log)) + + rate, usage = mod.query_codex(timeout=5) + requests = [json.loads(line) for line in request_log.read_text().splitlines()] + + assert rate["rateLimits"]["primary"]["usedPercent"] == 10 + assert usage["summary"] == {} + assert [request["method"] for request in requests] == [ + "initialize", + "initialized", + "account/rateLimits/read", + "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.stderr = None + 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) + + with pytest.raises(RuntimeError, match="^Codex app-server cleanup failed$"): + mod.query_codex(timeout=1) + + 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_cleanup_contains_all_pipe_and_process_api_failures(): + mod = load_module() + + class BrokenStream: + def close(self): + raise OSError("private pipe detail") + + class BrokenSelector: + def close(self): + raise OSError("private selector detail") + + class BrokenProcess: + stdin = BrokenStream() + stdout = BrokenStream() + stderr = BrokenStream() + + def __init__(self): + self.wait_calls = 0 + + def terminate(self): + raise OSError("private terminate detail") + + def wait(self, timeout): + assert timeout == mod.CODEX_CLEANUP_TIMEOUT_SECONDS + self.wait_calls += 1 + raise OSError("private wait detail") + + def kill(self): + raise OSError("private kill detail") + + assert mod._close_stream(None) + assert not mod._close_stream(BrokenStream()) + assert not mod._cleanup_codex_process(BrokenProcess(), BrokenSelector()) + + +def test_cleanup_accepts_already_exited_process_during_terminate_and_kill(): + mod = load_module() + + class AlreadyExited: + stdin = None + stdout = None + stderr = None + + def __init__(self, *, timeout_first): + self.timeout_first = timeout_first + self.wait_calls = 0 + + def terminate(self): + raise ProcessLookupError + + def wait(self, timeout): + self.wait_calls += 1 + if self.timeout_first and self.wait_calls == 1: + raise subprocess.TimeoutExpired("codex", timeout) + + def kill(self): + raise ProcessLookupError + + assert mod._cleanup_codex_process(AlreadyExited(timeout_first=False), None) + assert mod._cleanup_codex_process(AlreadyExited(timeout_first=True), None) + + +def test_query_rejects_missing_pipes_and_structured_errors(monkeypatch): + mod = load_module() + + class Process: + stderr = None + + def __init__(self, *_args, **_kwargs): + self.stdin = None + self.stdout = None + + def terminate(self): + return None + + def wait(self, timeout): + return 0 + + monkeypatch.setattr(mod.subprocess, "Popen", Process) + with pytest.raises(RuntimeError, match="pipes are unavailable"): + mod.query_codex(timeout=0) + + class InvalidProcess(Process): + def __init__(self, *_args, **_kwargs): + self.stdin = StringIO() + self.stdout = SequencedOutput() + + class SequencedOutput: + def __init__(self): + self.lines = iter( + [ + "", + json.dumps({"id": 99, "result": {}}) + "\n", + json.dumps({"id": 1, "result": {}}) + "\n", + json.dumps({"id": 2, "error": {"message": "secret"}}) + "\n", + json.dumps({"id": 3, "result": {}}) + "\n", + ] + ) + + def readline(self): + return next(self.lines) + + def close(self): + return None + + class Selector: + def register(self, fileobj, _events): + self.fileobj = fileobj + + def select(self, _timeout): + return [(SimpleNamespace(fileobj=self.fileobj), None)] + + def close(self): + return None + + monkeypatch.setattr(mod.subprocess, "Popen", InvalidProcess) + monkeypatch.setattr(mod.selectors, "DefaultSelector", Selector) + with pytest.raises( + RuntimeError, match="^Codex app-server response 2 failed$" + ) as caught: + mod.query_codex(timeout=1) + assert "secret" not in str(caught.value) diff --git a/testing/tests/test_hermes_ai_usage_exporter.py b/testing/tests/test_hermes_ai_usage_exporter.py index 9fe54d72..c33654dc 100644 --- a/testing/tests/test_hermes_ai_usage_exporter.py +++ b/testing/tests/test_hermes_ai_usage_exporter.py @@ -3,13 +3,10 @@ from __future__ import annotations import importlib.util -import json -import subprocess import sys from datetime import date -from io import BytesIO, StringIO +from io import BytesIO from pathlib import Path -from types import SimpleNamespace ROOT = Path(__file__).resolve().parents[2] @@ -18,6 +15,9 @@ SCRIPT = ROOT / "services/hermes/scripts/ai_usage_exporter.py" def load_module(): """Load the standalone exporter script as a test module.""" + script_directory = str(SCRIPT.parent) + if script_directory not in sys.path: + sys.path.insert(0, script_directory) spec = importlib.util.spec_from_file_location("ai_usage_exporter", SCRIPT) module = importlib.util.module_from_spec(spec) assert spec and spec.loader @@ -57,7 +57,9 @@ def test_codex_payloads_include_real_windows_and_bounded_token_periods(): } samples = mod.parse_codex_payloads(rate, usage, today=date(2026, 8, 16)) - values = {(item.name, tuple(sorted(item.labels.items()))): item.value for item in samples} + values = { + (item.name, tuple(sorted(item.labels.items()))): item.value for item in samples + } codex_labels = ( ("limit", "codex"), @@ -72,18 +74,24 @@ def test_codex_payloads_include_real_windows_and_bounded_token_periods(): assert values[("atlas_ai_quota_used_percent", codex_labels)] == 41 assert values[("atlas_ai_quota_remaining_percent", codex_labels)] == 59 assert values[("atlas_ai_quota_used_percent", spark_labels)] == 0 - assert values[ - ( - "atlas_ai_account_tokens", - (("period", "latest_day"), ("provider", "openai")), - ) - ] == 100 - assert values[ - ( - "atlas_ai_account_tokens", - (("period", "seven_day"), ("provider", "openai")), - ) - ] == 150 + assert ( + values[ + ( + "atlas_ai_account_tokens", + (("period", "latest_day"), ("provider", "openai")), + ) + ] + == 100 + ) + assert ( + values[ + ( + "atlas_ai_account_tokens", + (("period", "seven_day"), ("provider", "openai")), + ) + ] + == 150 + ) def test_claude_payload_uses_only_supported_quota_fields(): @@ -134,120 +142,21 @@ def test_render_reports_failure_and_freshness_without_logging_credentials(): rendered = collector.render().decode() assert 'atlas_ai_quota_fetch_success{provider="openai"} 0' in rendered - assert 'atlas_ai_quota_last_success_timestamp_seconds{provider="openai"} 100' in rendered - assert 'atlas_ai_quota_used_percent{limit="codex",provider="openai",window="seven_day"} 42' in rendered + assert ( + 'atlas_ai_quota_last_success_timestamp_seconds{provider="openai"} 100' + in rendered + ) + assert ( + 'atlas_ai_quota_used_percent{limit="codex",provider="openai",window="seven_day"} 42' + in rendered + ) assert "accessToken" not in rendered assert "refreshToken" not in rendered -def test_codex_query_uses_structured_app_server_protocol(tmp_path, monkeypatch): - mod = load_module() - request_log = tmp_path / "requests.jsonl" - mock = tmp_path / "codex-mock" - mock.write_text( - """#!/usr/bin/env python3 -import json, os, sys -messages = [json.loads(sys.stdin.readline()) for _ in range(4)] -with open(os.environ['REQUEST_LOG'], 'w') as handle: - for message in messages: - handle.write(json.dumps(message) + '\\n') -responses = { - 1: {}, - 2: {'rateLimits': {'primary': {'usedPercent': 10}}}, - 3: {'summary': {}, 'dailyUsageBuckets': []}, -} -for response_id, result in responses.items(): - print(json.dumps({'id': response_id, 'result': result}), flush=True) -""" - ) - mock.chmod(0o755) - monkeypatch.setattr(mod, "CODEX_BIN", str(mock)) - monkeypatch.setenv("REQUEST_LOG", str(request_log)) - - rate, usage = mod.query_codex(timeout=5) - requests = [json.loads(line) for line in request_log.read_text().splitlines()] - - assert rate["rateLimits"]["primary"]["usedPercent"] == 10 - assert usage["summary"] == {} - assert [request["method"] for request in requests] == [ - "initialize", - "initialized", - "account/rateLimits/read", - "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): +def test_unexpected_provider_failure_is_sanitized_and_preserves_last_good( + monkeypatch, capsys +): mod = load_module() collector = mod.Collector() previous = mod.Sample( diff --git a/testing/tests/test_hermes_ai_usage_exporter_coverage.py b/testing/tests/test_hermes_ai_usage_exporter_coverage.py new file mode 100644 index 00000000..affaf336 --- /dev/null +++ b/testing/tests/test_hermes_ai_usage_exporter_coverage.py @@ -0,0 +1,275 @@ +"""Focused edge coverage for AI usage exporter provider and polling boundaries.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from io import BytesIO +from pathlib import Path + +import pytest + + +ROOT = Path(__file__).resolve().parents[2] +SCRIPT = ROOT / "services/hermes/scripts/ai_usage_exporter.py" + + +def load_module(): + """Load a fresh exporter while making its sibling modules importable.""" + script_directory = str(SCRIPT.parent) + if script_directory not in sys.path: + sys.path.insert(0, script_directory) + name = "ai_usage_exporter_coverage" + spec = importlib.util.spec_from_file_location(name, SCRIPT) + module = importlib.util.module_from_spec(spec) + assert spec and spec.loader + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def test_parser_rejects_nonfinite_and_malformed_optional_fields(): + mod = load_module() + + assert mod._number(True) is None + assert mod._number("1") is None + assert mod._number(float("nan")) is None + assert mod._number(float("inf")) is None + assert mod._timestamp(None) is None + assert mod._timestamp("not-a-time") is None + assert mod._window_name({}, "fallback") == "fallback" + assert mod._codex_limit_name("other", {}) == "additional" + + samples = mod.parse_codex_payloads( + { + "rateLimitsByLimitId": { + 1: {}, + "not-a-snapshot": "invalid", + "unnamed": { + "primary": {"usedPercent": None}, + "secondary": { + "usedPercent": 120, + "resetsAt": "invalid", + }, + }, + } + }, + { + "summary": { + "lifetimeTokens": None, + "currentStreakDays": 2, + "longestStreakDays": 3, + "longestRunningTurnSec": 4, + }, + "dailyUsageBuckets": [ + None, + {"startDate": "invalid", "tokens": 2}, + {"startDate": "2026-08-17", "tokens": None}, + ], + }, + ) + + assert any(sample.value == 120 for sample in samples) + assert not any( + sample.name == "atlas_ai_quota_reset_timestamp_seconds" for sample in samples + ) + assert mod.parse_codex_payloads({"rateLimits": None}, {"summary": None}) + + +def test_claude_parser_ignores_unbounded_fields_and_invalid_extra_usage(): + mod = load_module() + + assert mod.parse_claude_payload( + { + "five_hour": {"utilization": None}, + "seven_day": {"utilization": 10, "resets_at": None}, + "extra_usage": {"is_enabled": "yes"}, + } + ) + assert mod.parse_claude_payload({"extra_usage": "invalid"}) == [] + + +def test_provider_queries_validate_credentials_and_response_shape( + tmp_path, monkeypatch +): + mod = load_module() + credentials = tmp_path / "credentials.json" + credentials.write_text("{}", encoding="utf-8") + monkeypatch.setattr(mod, "CLAUDE_CREDENTIALS", credentials) + + with pytest.raises(RuntimeError, match="credentials are incomplete"): + mod.query_claude() + + credentials.write_text( + json.dumps({"claudeAiOauth": {"accessToken": "runtime-only-token"}}), + encoding="utf-8", + ) + bodies = iter([{"five_hour": {}}, []]) + requests = [] + + class Response: + def __enter__(self): + return self + + def __exit__(self, *_args): + return None + + def read(self, maximum): + assert maximum == 1 << 20 + return json.dumps(next(bodies)).encode() + + def open_request(request, timeout): + requests.append(request) + assert timeout == 15 + return Response() + + monkeypatch.setattr(mod, "urlopen", open_request) + assert mod.query_claude() == {"five_hour": {}} + assert requests[0].get_header("Authorization") == "Bearer runtime-only-token" + with pytest.raises(RuntimeError, match="response is not an object"): + mod.query_claude() + + assert mod.query_codex is mod.codex_query.query_codex + + +def test_collector_refreshes_both_providers_and_contains_unknown_provider( + monkeypatch, capsys +): + mod = load_module() + collector = mod.Collector() + monkeypatch.setattr( + mod, + "query_codex", + lambda: ({"rateLimits": {"primary": {"usedPercent": 1}}}, {}), + ) + monkeypatch.setattr(mod, "query_claude", lambda: {"five_hour": {"utilization": 2}}) + + collector.refresh_provider("openai") + collector.refresh_provider("anthropic") + assert collector._providers["openai"].fetch_success + assert collector._providers["anthropic"].fetch_success + assert collector._providers["openai"].last_success > 0 + + collector.record_failure("openai") + assert not collector._providers["openai"].fetch_success + with pytest.raises(KeyError, match="unknown"): + collector.refresh_provider("unknown") + assert "unknown quota collection deferred: ValueError" in capsys.readouterr().out + + +def test_polling_failure_accounting_and_cycle_errors_are_sanitized(monkeypatch, capsys): + mod = load_module() + secret = "must-not-leak" + + class Failure(Exception): + pass + + class Collector: + def refresh_provider(self, _provider): + raise Failure(secret) + + def record_failure(self, *_args, **_kwargs): + raise Failure(secret) + + engine = mod.PollingEngine(Collector(), interval=60) + engine.poll_once() + output = capsys.readouterr().out + assert "quota failure accounting deferred: Failure" in output + assert secret not in output + + monkeypatch.setattr( + engine, "poll_once", lambda: (_ for _ in ()).throw(Failure(secret)) + ) + monkeypatch.setattr( + mod.time, "sleep", lambda _interval: (_ for _ in ()).throw(StopIteration) + ) + with pytest.raises(StopIteration): + engine.run() + output = capsys.readouterr().out + assert "quota polling cycle deferred: Failure" in output + assert secret not in output + + +def test_poller_start_is_idempotent_and_main_wires_server(monkeypatch): + mod = load_module() + threads = [] + + class Thread: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.started = False + threads.append(self) + + def start(self): + self.started = True + + monkeypatch.setattr(mod.threading, "Thread", Thread) + engine = mod.PollingEngine(mod.Collector(), interval=60) + engine.start() + engine.start() + assert len(threads) == 1 + assert threads[0].started + assert threads[0].kwargs["daemon"] is True + + events = [] + + class Poller: + def __init__(self, collector, *, interval): + events.append(("poller", collector, interval)) + + def start(self): + events.append("started") + + class Server: + def __init__(self, address, handler): + events.append(("server", address, handler)) + + def serve_forever(self): + events.append("served") + + monkeypatch.setattr(mod, "PollingEngine", Poller) + monkeypatch.setattr(mod, "Server", Server) + monkeypatch.setattr( + mod, "make_handler", lambda collector, poller: (collector, poller) + ) + monkeypatch.setenv("ATLAS_AI_USAGE_INTERVAL_SECONDS", "1") + monkeypatch.setenv("ATLAS_AI_USAGE_PORT", "19010") + + assert mod.main() == 0 + assert events[0][0] == "poller" + assert events[0][2] == 60 + assert events[2][1][1] == 19010 + assert events[-1] == "served" + + +def test_http_handler_covers_not_found_logging_and_connection_reset(): + mod = load_module() + + class Collector: + def render(self): + return b"metrics\n" + + class Poller: + def is_healthy(self): + return True + + handler_class = mod.make_handler(Collector(), Poller()) + handler = object.__new__(handler_class) + handler.path = "/missing" + statuses = [] + handler.send_error = statuses.append + handler.do_GET() + assert statuses == [404] + assert handler.log_message("ignored") is None + + class Reset(BytesIO): + def write(self, _body): + raise ConnectionResetError + + handler.path = "/metrics" + handler.send_response = lambda _status: None + handler.send_header = lambda _name, _value: None + handler.end_headers = lambda: None + handler.wfile = Reset() + handler.do_GET() diff --git a/testing/tests/test_hermes_auto_router.py b/testing/tests/test_hermes_auto_router.py index c0a31103..601f39df 100644 --- a/testing/tests/test_hermes_auto_router.py +++ b/testing/tests/test_hermes_auto_router.py @@ -59,9 +59,7 @@ def test_auto_boundary_selects_only_a_public_switchyard_route(tmp_path, monkeypa assert source == "auto" -def test_legacy_maximum_default_migrates_to_adaptive_agent_route( - tmp_path, monkeypatch -): +def test_legacy_maximum_default_migrates_to_adaptive_agent_route(tmp_path, monkeypatch): path = tmp_path / "route-policy.json" path.write_text( json.dumps({"mode": "auto", "auto_route": "atlas/auto/maximum"}), @@ -90,9 +88,7 @@ def test_ui_priority_changes_auto_posture_without_selecting_a_target( assert route in router.AUTO_ROUTES -def test_ui_manual_model_and_effort_are_forwarded_as_constraints( - tmp_path, monkeypatch -): +def test_ui_manual_model_and_effort_are_forwarded_as_constraints(tmp_path, monkeypatch): monkeypatch.setattr(router, "POLICY_PATH", tmp_path / "route-policy.json") agent = _agent( model="atlas/manual/claude/opus", @@ -107,9 +103,7 @@ def test_ui_manual_model_and_effort_are_forwarded_as_constraints( ) -def test_cli_exact_manual_route_remains_exact_across_boundaries( - tmp_path, monkeypatch -): +def test_cli_exact_manual_route_remains_exact_across_boundaries(tmp_path, monkeypatch): monkeypatch.setattr(router, "POLICY_PATH", tmp_path / "route-policy.json") agent = _agent( model="atlas/manual/claude/opus/xhigh", @@ -222,39 +216,58 @@ def test_adapter_contains_no_content_classifier_or_direct_ollama_call(): def test_provider_status_separates_observed_activity_from_plan_quota(monkeypatch): status = sys.modules["hermes_auto_router"].provider_status_text.__globals__ module = sys.modules[status["provider_status_payload"].__module__] - monkeypatch.setattr(module, "_get_json", lambda url, timeout=3.0: ( - {"status": "ok"} if url.endswith("/health") else { - "total_requests": 3, - "total_errors": 1, - "total_tokens": {"total": 900}, - "models": { - "route/codex/sol/xhigh": { - "calls": 1, - "errors": 0, - "total_tokens": 600, - "avg_latency_ms": 1200, + monkeypatch.setattr( + module, + "_get_json", + lambda url, timeout=3.0: ( + {"status": "ok"} + if url.endswith("/health") + else { + "total_requests": 3, + "total_errors": 1, + "total_tokens": {"total": 900}, + "models": { + "route/codex/sol/xhigh": { + "calls": 1, + "errors": 0, + "total_tokens": 600, + "avg_latency_ms": 1200, + }, + "route/claude/sonnet/high": { + "calls": 0, + "errors": 1, + "total_tokens": 0, + }, + "route/local/qwen2.5-14b/medium": { + "calls": 1, + "errors": 0, + "total_tokens": 300, + }, }, - "route/claude/sonnet/high": { - "calls": 0, - "errors": 1, - "total_tokens": 0, - }, - "route/local/qwen2.5-14b/medium": { - "calls": 1, - "errors": 0, - "total_tokens": 300, - }, - }, - "routing_fallbacks": {"unavailable": 1}, - "classifier": {"total_requests": 3, "total_errors": 0}, - } - )) - monkeypatch.setattr(module, "_codex_account", lambda: { - "authenticated": True, "plan": "plus", "quota_reported": False, - }) - monkeypatch.setattr(module, "_claude_account", lambda: { - "authenticated": True, "plan": "max", "quota_reported": False, - }) + "routing_fallbacks": {"unavailable": 1}, + "classifier": {"total_requests": 3, "total_errors": 0}, + } + ), + ) + monkeypatch.setattr( + module, + "_codex_account", + lambda: { + "authenticated": True, + "plan": "plus", + "quota_reported": False, + }, + ) + monkeypatch.setattr( + module, + "_claude_account", + lambda: { + "authenticated": True, + "plan": "max", + "quota_reported": False, + }, + ) + monkeypatch.setattr(module, "_fresh_health", lambda *_args, **_kwargs: {}) payload = module.provider_status_payload() @@ -359,15 +372,20 @@ def test_claude_account_treats_a_live_refresh_token_as_refreshable( status = sys.modules["hermes_auto_router"].provider_status_text.__globals__ module = sys.modules[status["provider_status_payload"].__module__] path = tmp_path / ".credentials.json" - path.write_text(json.dumps({ - "claudeAiOauth": { - "accessToken": "expired-access", - "refreshToken": "live-refresh", - "expiresAt": 1, - "refreshTokenExpiresAt": 32_472_192_000, - "subscriptionType": "max", - } - }), encoding="utf-8") + path.write_text( + json.dumps( + { + "claudeAiOauth": { + "accessToken": "expired-access", + "refreshToken": "live-refresh", + "expiresAt": 1, + "refreshTokenExpiresAt": 32_472_192_000, + "subscriptionType": "max", + } + } + ), + encoding="utf-8", + ) monkeypatch.setattr(module, "CLAUDE_AUTH_PATH", path) account = module._claude_account() @@ -382,13 +400,18 @@ def test_codex_account_treats_a_refresh_token_as_refreshable(tmp_path, monkeypat status = sys.modules["hermes_auto_router"].provider_status_text.__globals__ module = sys.modules[status["provider_status_payload"].__module__] path = tmp_path / "auth.json" - path.write_text(json.dumps({ - "auth_mode": "chatgpt", - "tokens": { - "access_token": "expired-access", - "refresh_token": "live-refresh", - }, - }), encoding="utf-8") + path.write_text( + json.dumps( + { + "auth_mode": "chatgpt", + "tokens": { + "access_token": "expired-access", + "refresh_token": "live-refresh", + }, + } + ), + encoding="utf-8", + ) monkeypatch.setattr(module, "CODEX_AUTH_PATH", path) account = module._codex_account() @@ -406,7 +429,8 @@ def test_agent_mounts_provider_status_dashboard_into_auto_router_plugin(): (root / "services/hermes/agent-deployment.yaml").read_text(encoding="utf-8") ) volume = next( - item for item in deployment["spec"]["template"]["spec"]["volumes"] + item + for item in deployment["spec"]["template"]["spec"]["volumes"] if item["name"] == "auto-router-plugin" ) paths = {item["path"] for item in volume["configMap"]["items"]}