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