hermes: split AI usage exporter engines
This commit is contained in:
parent
fcc8f650c2
commit
ed43dbaa7f
@ -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"
|
||||
|
||||
@ -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
|
||||
|
||||
138
services/hermes/scripts/ai_usage_codex.py
Normal file
138
services/hermes/scripts/ai_usage_codex.py
Normal 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
|
||||
@ -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:
|
||||
|
||||
49
services/hermes/scripts/ai_usage_http.py
Normal file
49
services/hermes/scripts/ai_usage_http.py
Normal 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
|
||||
255
testing/tests/test_hermes_ai_usage_codex.py
Normal file
255
testing/tests/test_hermes_ai_usage_codex.py
Normal 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)
|
||||
@ -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(
|
||||
|
||||
275
testing/tests/test_hermes_ai_usage_exporter_coverage.py
Normal file
275
testing/tests/test_hermes_ai_usage_exporter_coverage.py
Normal 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()
|
||||
@ -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"]}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user