264 lines
7.8 KiB
Python
264 lines
7.8 KiB
Python
"""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_default_query_deadline_has_margin_for_live_app_server_jitter(monkeypatch):
|
|
"""The default must not sit on the live app-server's common 20s boundary."""
|
|
monkeypatch.setenv("ATLAS_AI_CODEX_QUERY_TIMEOUT_SECONDS", "45")
|
|
mod = load_module()
|
|
|
|
assert mod.CODEX_QUERY_TIMEOUT_SECONDS == 45
|
|
|
|
|
|
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)
|