You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
116 lines
4.6 KiB
116 lines
4.6 KiB
"""Recording reverse proxy for the Anthropic API, run in its own container.
|
|
|
|
Claude Code is pointed at it with ANTHROPIC_BASE_URL. Every request body is
|
|
recorded (decompressed) before being forwarded unchanged, so a test can check
|
|
what actually left the machine. Responses are streamed back untouched.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import gzip
|
|
import json
|
|
import os
|
|
import zlib
|
|
|
|
import aiohttp
|
|
from aiohttp import web
|
|
|
|
DEFAULT_UPSTREAM = "https://api.anthropic.com"
|
|
|
|
# Hop-by-hop headers (RFC 9110) plus the ones aiohttp recomputes.
|
|
_SKIP_REQUEST_HEADERS = {"host", "content-length", "connection", "keep-alive", "transfer-encoding", "upgrade"}
|
|
_SKIP_RESPONSE_HEADERS = {"content-length", "connection", "keep-alive", "transfer-encoding", "upgrade"}
|
|
|
|
|
|
def _decode(body: bytes, encoding: str) -> str:
|
|
encoding = encoding.lower().strip()
|
|
try:
|
|
if encoding == "gzip":
|
|
body = gzip.decompress(body)
|
|
elif encoding == "deflate":
|
|
body = zlib.decompress(body)
|
|
elif encoding == "zstd":
|
|
import zstandard
|
|
|
|
body = zstandard.ZstdDecompressor().decompressobj().decompress(body)
|
|
elif encoding == "br":
|
|
import brotli # optional
|
|
|
|
body = brotli.decompress(body)
|
|
except Exception: # keep the raw bytes: a failed decode must not hide data
|
|
pass
|
|
return body.decode("utf-8", errors="replace")
|
|
|
|
|
|
class RecordingProxy:
|
|
"""Forwards every request to `upstream` and appends its body to `record_file`."""
|
|
|
|
def __init__(self, upstream: str, auth: str | None, record_file: str) -> None:
|
|
self.upstream = upstream
|
|
# "oauth:<token>" or "apikey:<key>": replaces whatever credential the client
|
|
# sent, so the sandboxed client only ever holds a dummy one.
|
|
self.auth = auth
|
|
self.record_file = record_file
|
|
self._session: aiohttp.ClientSession | None = None
|
|
|
|
async def serve(self, host: str, port: int) -> None:
|
|
self._session = aiohttp.ClientSession(auto_decompress=False, trust_env=True,
|
|
timeout=aiohttp.ClientTimeout(total=None, sock_read=600))
|
|
app = web.Application(client_max_size=256 * 1024 * 1024)
|
|
app.router.add_route("*", "/{tail:.*}", self._handle)
|
|
runner = web.AppRunner(app, access_log=None)
|
|
await runner.setup()
|
|
await web.TCPSite(runner, host, port).start()
|
|
print(f"listening on {host}:{port}", flush=True)
|
|
await asyncio.Event().wait()
|
|
|
|
async def _handle(self, request: web.Request) -> web.StreamResponse:
|
|
assert self._session is not None
|
|
body = await request.read()
|
|
with open(self.record_file, "a") as f:
|
|
f.write(json.dumps({"method": request.method, "path": request.path_qs,
|
|
"body": _decode(body, request.headers.get("Content-Encoding", ""))}) + "\n")
|
|
|
|
headers = {k: v for k, v in request.headers.items() if k.lower() not in _SKIP_REQUEST_HEADERS}
|
|
if self.auth:
|
|
headers = _inject_auth(headers, self.auth)
|
|
url = self.upstream.rstrip("/") + request.path_qs
|
|
async with self._session.request(request.method, url, headers=headers, data=body) as upstream:
|
|
response = web.StreamResponse(status=upstream.status, reason=upstream.reason)
|
|
for k, v in upstream.headers.items():
|
|
if k.lower() not in _SKIP_RESPONSE_HEADERS:
|
|
response.headers.add(k, v)
|
|
await response.prepare(request)
|
|
async for chunk in upstream.content.iter_any():
|
|
await response.write(chunk)
|
|
await response.write_eof()
|
|
return response
|
|
|
|
|
|
def _inject_auth(headers: dict[str, str], auth: str) -> dict[str, str]:
|
|
kind, _, secret = auth.strip().partition(":")
|
|
out = {k: v for k, v in headers.items() if k.lower() not in ("authorization", "x-api-key")}
|
|
if kind == "oauth":
|
|
out["Authorization"] = f"Bearer {secret}"
|
|
elif kind == "apikey":
|
|
out["x-api-key"] = secret
|
|
else:
|
|
raise ValueError(f"unknown auth kind {kind!r}")
|
|
return out
|
|
|
|
|
|
def main() -> None:
|
|
"""Container entry point: listen on 0.0.0.0:8080, record to /records."""
|
|
auth_file = os.environ.get("BENCH_AUTH_FILE", "/run/secrets/bench-auth")
|
|
auth = open(auth_file).read() if os.path.exists(auth_file) else None
|
|
proxy = RecordingProxy(os.environ.get("BENCH_UPSTREAM", DEFAULT_UPSTREAM), auth,
|
|
os.environ.get("BENCH_RECORD_FILE", "/records/wire.jsonl"))
|
|
try:
|
|
asyncio.run(proxy.serve("0.0.0.0", int(os.environ.get("BENCH_PORT", "8080"))))
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|