165 lines
8.1 KiB
Python
165 lines
8.1 KiB
Python
"""Trace Codex HTTP boundaries without logging credentials or prompt contents."""
|
|
import argparse
|
|
from collections import Counter
|
|
import http.server
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import signal
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
import tomllib
|
|
import urllib.error
|
|
import urllib.request
|
|
from codex_http_transport import ModelTransport
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--route-prefix", default="")
|
|
parser.add_argument("--disable-compression", action="store_true")
|
|
parser.add_argument("--direct", action="store_true", help="compare direct CLI transport with the same upstream URL")
|
|
parser.add_argument("--production", action="store_true", help="exercise the scheduler's bounded-retry transport")
|
|
parser.add_argument("--timeout", type=float, default=180)
|
|
parser.add_argument("--report", type=Path)
|
|
args = parser.parse_args()
|
|
config_root = Path(os.environ.get("CODEX_HOME", str(Path.home() / ".codex")))
|
|
config = tomllib.loads((config_root / "config.toml").read_text())
|
|
provider = config["model_provider"]
|
|
upstream = config["model_providers"][provider]["base_url"].rstrip("/") + args.route_prefix.rstrip("/")
|
|
started = time.monotonic()
|
|
output_lock = threading.Lock()
|
|
requests = []
|
|
if args.report:
|
|
args.report.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
def event(event_type, **values):
|
|
with output_lock:
|
|
line = json.dumps({"event": event_type, "seconds": round(time.monotonic() - started, 3), **values})
|
|
print(line, flush=True)
|
|
if args.report:
|
|
with args.report.open("a", encoding="utf-8") as stream:
|
|
stream.write(line + "\n")
|
|
|
|
class Proxy(http.server.BaseHTTPRequestHandler):
|
|
protocol_version = "HTTP/1.1"
|
|
|
|
def log_message(self, *_):
|
|
pass
|
|
|
|
def do_POST(self):
|
|
if self.path != "/responses":
|
|
self.send_error(404)
|
|
return
|
|
length = int(self.headers.get("Content-Length", 0))
|
|
if not 0 < length <= 16 * 1024 * 1024:
|
|
self.send_error(413)
|
|
return
|
|
body = self.rfile.read(length)
|
|
number = len(requests) + 1
|
|
requests.append(number)
|
|
metadata = {}
|
|
if not self.headers.get("Content-Encoding"):
|
|
payload = json.loads(body)
|
|
metadata = {"model": payload.get("model"), "stream": payload.get("stream"),
|
|
"toolTypes": dict(Counter(row.get("type") for row in payload.get("tools", []))),
|
|
"inputTypes": dict(Counter(row.get("type") for row in payload.get("input", []) if isinstance(row, dict))),
|
|
"previousResponse": bool(payload.get("previous_response_id")),
|
|
"additionalTools": [{"keys": list(row), "fieldTypes": {k: type(v).__name__ for k, v in row.items()},
|
|
"count": len(row.get("tools", [])),
|
|
"types": dict(Counter(tool.get("type") for tool in row.get("tools", [])))}
|
|
for row in payload.get("input", []) if isinstance(row, dict) and row.get("type") == "additional_tools"]}
|
|
event("request", number=number, bytes=length, encoding=self.headers.get("Content-Encoding"), **metadata)
|
|
headers = {k: v for k, v in self.headers.items() if k.lower() not in {"host", "connection", "content-length", "transfer-encoding"}}
|
|
request = urllib.request.Request(upstream + "/responses", data=body, headers=headers)
|
|
sent_headers = False
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=60) as response:
|
|
event("upstream_headers", number=number, status=response.status)
|
|
self.send_response(response.status)
|
|
self.send_header("Content-Type", response.headers.get("Content-Type", "text/event-stream"))
|
|
self.send_header("Connection", "close")
|
|
self.end_headers()
|
|
sent_headers = True
|
|
count = 0
|
|
terminal = None
|
|
for line in response:
|
|
if count == 0:
|
|
event("first_body", number=number)
|
|
count += len(line)
|
|
if line.startswith(b"data:"):
|
|
try:
|
|
item = json.loads(line[5:])
|
|
kind = item.get("type")
|
|
if kind in {"response.created", "response.completed", "response.failed", "error"}:
|
|
event("upstream_event", number=number, kind=kind)
|
|
if kind in {"response.completed", "response.failed", "error"}:
|
|
terminal = kind
|
|
except ValueError:
|
|
pass
|
|
self.wfile.write(line)
|
|
self.wfile.flush()
|
|
event("upstream_end", number=number, bytes=count, terminal=terminal)
|
|
except Exception as error:
|
|
event("upstream_error", number=number, errorType=type(error).__name__, status=getattr(error, "code", None))
|
|
if not sent_headers:
|
|
try:
|
|
self.send_error(502, "Diagnostic upstream request failed")
|
|
except OSError:
|
|
pass
|
|
self.close_connection = True
|
|
|
|
server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Proxy)
|
|
server.daemon_threads = True
|
|
threading.Thread(target=server.serve_forever, daemon=True).start()
|
|
production = ModelTransport(provider, upstream, log_path=Path(__file__).resolve().parents[1] / f"docs/current/continuation-runs/transport-probe-{os.getpid()}.jsonl") if args.production else None
|
|
url = production.url if production else upstream if args.direct else f"http://127.0.0.1:{server.server_port}"
|
|
command = ["codex", "--ask-for-approval", "never", "--sandbox", "danger-full-access", "exec", "--ephemeral", "--json",
|
|
"--skip-git-repo-check", "-C", "/tmp", "-c", 'model_reasoning_effort="low"',
|
|
"-c", f"model_providers.{provider}.base_url={json.dumps(url)}",
|
|
"Run pwd exactly once, then report its actual output. Do not edit files or run other commands."]
|
|
if args.disable_compression or production:
|
|
command[1:1] = ["--disable", "enable_request_compression"]
|
|
child = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
|
|
text=True, start_new_session=True)
|
|
commands = []
|
|
completed = []
|
|
|
|
def collect():
|
|
for line in child.stdout:
|
|
try:
|
|
item = json.loads(line)
|
|
except ValueError:
|
|
continue
|
|
detail = item.get("item", {})
|
|
event("cli", kind=item.get("type"), itemType=detail.get("type"), exitCode=detail.get("exit_code"))
|
|
if item.get("type") == "item.completed" and detail.get("type") == "command_execution" and detail.get("exit_code") == 0:
|
|
commands.append(True)
|
|
if item.get("type") == "turn.completed":
|
|
completed.append(True)
|
|
|
|
reader = threading.Thread(target=collect, daemon=True)
|
|
reader.start()
|
|
try:
|
|
code = child.wait(timeout=args.timeout)
|
|
except subprocess.TimeoutExpired:
|
|
os.killpg(child.pid, signal.SIGTERM)
|
|
try:
|
|
child.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
os.killpg(child.pid, signal.SIGKILL)
|
|
child.wait()
|
|
code = 124
|
|
reader.join(timeout=2)
|
|
server.shutdown()
|
|
server.server_close()
|
|
if production:
|
|
production.close()
|
|
event("result", exitCode=code, requests=production.requests if production else len(requests), commands=len(commands), completed=bool(completed))
|
|
return 0 if code == 0 and commands and completed else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|