hermes: split AI usage exporter engines

This commit is contained in:
Hermes Agent 2026-08-17 15:33:56 +00:00
parent fcc8f650c2
commit ed43dbaa7f
9 changed files with 903 additions and 365 deletions

View File

@ -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"

View File

@ -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

View File

@ -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

View File

@ -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:

View File

@ -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

View File

@ -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)

View File

@ -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(

View File

@ -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()

View File

@ -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"]}