144 lines
4.8 KiB
Python
144 lines
4.8 KiB
Python
"""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
|
|
CODEX_QUERY_TIMEOUT_SECONDS = min(
|
|
60.0,
|
|
max(20.0, float(os.environ.get("ATLAS_AI_CODEX_QUERY_TIMEOUT_SECONDS", "45"))),
|
|
)
|
|
|
|
|
|
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 | None = None) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
"""Read Codex account quota and usage through its structured app-server protocol."""
|
|
timeout = CODEX_QUERY_TIMEOUT_SECONDS if timeout is None else timeout
|
|
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
|