proxy_server.py

"""
proxy_server.py - Local HTTP proxy that intercepts Anthropic API calls and
captures token usage data from responses.
"""

from __future__ import annotations

import http.server
import json
import ssl
import threading
import time
import urllib.error
import urllib.request
from http.client import HTTPResponse
from io import BytesIO


# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

ANTHROPIC_HOST = "api.anthropic.com"
ANTHROPIC_BASE_URL = f"https://{ANTHROPIC_HOST}"
DEFAULT_PORT = 7834

# Headers that must not be forwarded upstream verbatim (hop-by-hop).
_HOP_BY_HOP_HEADERS = frozenset(
    [
        "connection",
        "keep-alive",
        "proxy-authenticate",
        "proxy-authorization",
        "te",
        "trailers",
        "transfer-encoding",
        "upgrade",
        # Host is rewritten to the upstream host.
        "host",
    ]
)

# CORS headers for browser-facing /track and /stats endpoints.
_CORS_HEADERS = {
    "Access-Control-Allow-Origin": "*",
    "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
    "Access-Control-Allow-Headers": "Content-Type, X-Session-Id",
}


# ---------------------------------------------------------------------------
# Helper – create an SSL context that validates api.anthropic.com
# ---------------------------------------------------------------------------

def _make_ssl_context():
    ctx = ssl.create_default_context()
    return ctx


# ---------------------------------------------------------------------------
# Request handler
# ---------------------------------------------------------------------------

class _ProxyHandler(http.server.BaseHTTPRequestHandler):
    """HTTP request handler that proxies to api.anthropic.com."""

    # Suppress the default "127.0.0.1 - - [date] METHOD path HTTP/1.x" logs.
    def log_message(self, format, *args):  # noqa: A002
        pass

    # ------------------------------------------------------------------
    # Routing
    # ------------------------------------------------------------------

    def do_OPTIONS(self):
        """Handle CORS preflight for /track and /stats."""
        self._send_cors_preflight()

    def do_GET(self):
        if self.path == "/stats":
            self._handle_stats()
        else:
            self._proxy_request()

    def do_POST(self):
        if self.path == "/track":
            self._handle_track()
        else:
            self._proxy_request()

    # Catch-all: forward any other HTTP verb to upstream.
    def do_PUT(self):
        self._proxy_request()

    def do_PATCH(self):
        self._proxy_request()

    def do_DELETE(self):
        self._proxy_request()

    def do_HEAD(self):
        self._proxy_request()

    # ------------------------------------------------------------------
    # /stats endpoint
    # ------------------------------------------------------------------

    def _handle_stats(self):
        storage = self.server.token_proxy.storage
        if storage is None:
            data = {"error": "No storage configured"}
            status = 503
        else:
            data = storage.get_all_totals()
            status = 200

        body = json.dumps(data).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        for k, v in _CORS_HEADERS.items():
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(body)

    # ------------------------------------------------------------------
    # /track endpoint
    # ------------------------------------------------------------------

    def _handle_track(self):
        try:
            length = int(self.headers.get("Content-Length", 0))
            raw = self.rfile.read(length)
            payload = json.loads(raw)
        except Exception as exc:
            self._send_json_error(400, f"Bad request: {exc}")
            return

        session_id = payload.get("session_id") or self._make_session_id()
        model = payload.get("model", "unknown")
        input_tokens = int(payload.get("input_tokens", 0))
        output_tokens = int(payload.get("output_tokens", 0))
        cache_creation = int(payload.get("cache_creation_input_tokens", 0))
        cache_read = int(payload.get("cache_read_input_tokens", 0))

        self.server.token_proxy._record(
            session_id, model, input_tokens, output_tokens, cache_creation, cache_read
        )

        body = json.dumps({"status": "ok"}).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        for k, v in _CORS_HEADERS.items():
            self.send_header(k, v)
        self.end_headers()
        self.wfile.write(body)

    # ------------------------------------------------------------------
    # Proxy logic
    # ------------------------------------------------------------------

    def _proxy_request(self):
        proxy = self.server.token_proxy
        session_id = self.headers.get("X-Session-Id") or self._make_session_id()

        # Build the upstream URL.
        upstream_url = ANTHROPIC_BASE_URL + self.path

        # Read request body (if any).
        content_length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(content_length) if content_length > 0 else None

        # Build upstream headers (drop hop-by-hop, set correct Host).
        upstream_headers = {}
        for name, value in self.headers.items():
            if name.lower() not in _HOP_BY_HOP_HEADERS:
                upstream_headers[name] = value
        upstream_headers["Host"] = ANTHROPIC_HOST

        req = urllib.request.Request(
            upstream_url,
            data=body,
            headers=upstream_headers,
            method=self.command,
        )

        ssl_ctx = _make_ssl_context()

        try:
            response: HTTPResponse = urllib.request.urlopen(req, context=ssl_ctx)
        except urllib.error.HTTPError as exc:
            # HTTPError is itself a valid response object – pass it through.
            response = exc
        except urllib.error.URLError as exc:
            self._send_json_error(502, f"Upstream connection failed: {exc.reason}")
            return
        except Exception as exc:
            self._send_json_error(502, f"Proxy error: {exc}")
            return

        content_type = response.headers.get("Content-Type", "")
        is_streaming = "text/event-stream" in content_type

        # Forward status line.
        self.send_response(response.status)

        # Forward upstream response headers (drop hop-by-hop).
        for name, value in response.headers.items():
            if name.lower() not in _HOP_BY_HOP_HEADERS:
                self.send_header(name, value)

        if is_streaming:
            # For SSE we stream chunks; Transfer-Encoding is chunked implicitly.
            self.send_header("Transfer-Encoding", "chunked")
        self.end_headers()

        if is_streaming:
            self._stream_response(response, session_id)
        else:
            self._buffer_response(response, session_id)

    # ------------------------------------------------------------------
    # Streaming response handler
    # ------------------------------------------------------------------

    def _stream_response(self, response, session_id):
        """Forward SSE stream while extracting token usage events."""
        proxy = self.server.token_proxy

        input_tokens = 0
        output_tokens = 0
        cache_creation = 0
        cache_read = 0
        model = "unknown"
        accumulated_line = b""

        try:
            while True:
                chunk = response.read(4096)
                if not chunk:
                    break

                # Forward raw chunk to client as an HTTP chunked-transfer chunk.
                chunk_size = f"{len(chunk):X}\r\n".encode()
                self.wfile.write(chunk_size)
                self.wfile.write(chunk)
                self.wfile.write(b"\r\n")
                self.wfile.flush()

                # Parse SSE lines from the chunk to extract usage.
                accumulated_line += chunk
                lines = accumulated_line.split(b"\n")
                # Keep the last (possibly incomplete) piece for next iteration.
                accumulated_line = lines[-1]

                for line in lines[:-1]:
                    line = line.rstrip(b"\r")
                    if not line.startswith(b"data: "):
                        continue
                    data_str = line[6:].decode("utf-8", errors="replace").strip()
                    if data_str == "[DONE]":
                        continue
                    try:
                        event_data = json.loads(data_str)
                    except json.JSONDecodeError:
                        continue

                    event_type = event_data.get("type", "")

                    if event_type == "message_start":
                        msg = event_data.get("message", {})
                        if not model or model == "unknown":
                            model = msg.get("model", "unknown")
                        usage = msg.get("usage", {})
                        input_tokens = usage.get("input_tokens", input_tokens)
                        cache_creation = usage.get(
                            "cache_creation_input_tokens", cache_creation
                        )
                        cache_read = usage.get("cache_read_input_tokens", cache_read)

                    elif event_type == "message_delta":
                        usage = event_data.get("usage", {})
                        output_tokens = usage.get("output_tokens", output_tokens)

            # Terminating zero-length chunk for chunked transfer encoding.
            self.wfile.write(b"0\r\n\r\n")
            self.wfile.flush()

        except (BrokenPipeError, ConnectionResetError):
            # Client disconnected; still record whatever usage we captured.
            pass
        except Exception:
            pass
        finally:
            response.close()
            if input_tokens or output_tokens:
                proxy._record(
                    session_id, model, input_tokens, output_tokens,
                    cache_creation, cache_read
                )

    # ------------------------------------------------------------------
    # Non-streaming (buffered) response handler
    # ------------------------------------------------------------------

    def _buffer_response(self, response, session_id):
        """Read full response, extract usage, forward to client."""
        proxy = self.server.token_proxy
        try:
            body = response.read()
        except Exception:
            body = b""
        finally:
            response.close()

        try:
            self.wfile.write(body)
            self.wfile.flush()
        except (BrokenPipeError, ConnectionResetError):
            pass

        # Parse usage from JSON response body.
        try:
            data = json.loads(body)
            usage = data.get("usage", {})
            if usage:
                model = data.get("model", "unknown")
                input_tokens = usage.get("input_tokens", 0)
                output_tokens = usage.get("output_tokens", 0)
                cache_creation = usage.get("cache_creation_input_tokens", 0)
                cache_read = usage.get("cache_read_input_tokens", 0)
                proxy._record(
                    session_id, model, input_tokens, output_tokens,
                    cache_creation, cache_read
                )
        except (json.JSONDecodeError, AttributeError):
            pass

    # ------------------------------------------------------------------
    # Utilities
    # ------------------------------------------------------------------

    def _make_session_id(self):
        """Generate a session ID from client address and current time."""
        addr = self.client_address[0] if self.client_address else "unknown"
        ts = int(time.time() * 1000)
        return f"{addr}-{ts}"

    def _send_cors_preflight(self):
        self.send_response(204)
        for k, v in _CORS_HEADERS.items():
            self.send_header(k, v)
        self.end_headers()

    def _send_json_error(self, status, message):
        body = json.dumps({"error": message}).encode("utf-8")
        try:
            self.send_response(status)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
        except (BrokenPipeError, ConnectionResetError):
            pass


# ---------------------------------------------------------------------------
# Main proxy class
# ---------------------------------------------------------------------------

class TokenCounterProxy:
    """
    Local HTTP proxy that intercepts Anthropic API requests, records token
    usage to a TokenStorage instance, and forwards all traffic transparently.
    """

    def __init__(self, port: int = DEFAULT_PORT):
        self.port = port
        self.storage = None
        self._on_update = None
        self._server: http.server.HTTPServer | None = None
        self._thread: threading.Thread | None = None
        self._lock = threading.Lock()

    # ------------------------------------------------------------------
    # Configuration
    # ------------------------------------------------------------------

    def set_storage(self, storage):
        """Attach a TokenStorage instance for persisting usage records."""
        self.storage = storage

    def set_on_update(self, callback):
        """Register a callable invoked after every usage record is written."""
        self._on_update = callback

    # ------------------------------------------------------------------
    # Server lifecycle
    # ------------------------------------------------------------------

    def start(self):
        """Start the proxy in a background daemon thread."""
        server = http.server.HTTPServer(("127.0.0.1", self.port), _ProxyHandler)
        # Allow address reuse to avoid "Address already in use" on quick restarts.
        server.allow_reuse_address = True
        # Attach a back-reference so the handler can reach this instance.
        server.token_proxy = self
        self._server = server

        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        self._thread = thread

    def stop(self):
        """Shut down the server and wait for the thread to exit."""
        if self._server is not None:
            self._server.shutdown()
            self._server.server_close()
            self._server = None
        if self._thread is not None:
            self._thread.join(timeout=5)
            self._thread = None

    # ------------------------------------------------------------------
    # Internal: record usage
    # ------------------------------------------------------------------

    def _record(
        self,
        session_id: str,
        model: str,
        input_tokens: int,
        output_tokens: int,
        cache_creation: int = 0,
        cache_read: int = 0,
    ):
        """Thread-safe write to storage + fire update callback."""
        with self._lock:
            if self.storage is not None:
                self.storage.record_usage(
                    session_id=session_id,
                    model=model,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    cache_creation=cache_creation,
                    cache_read=cache_read,
                )
            if self._on_update is not None:
                try:
                    self._on_update()
                except Exception:
                    pass


# ---------------------------------------------------------------------------
# Standalone entry-point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys

    port = DEFAULT_PORT
    if len(sys.argv) > 1:
        try:
            port = int(sys.argv[1])
        except ValueError:
            print(f"Invalid port '{sys.argv[1]}', using default {DEFAULT_PORT}.")

    from storage import TokenStorage

    storage = TokenStorage()
    proxy = TokenCounterProxy(port=port)
    proxy.set_storage(storage)

    def _on_update():
        totals = storage.get_all_totals()
        print(
            f"  [usage] input={totals['total_input']}  output={totals['total_output']}"
            f"  cache_creation={totals['total_cache_creation']}"
            f"  cache_read={totals['total_cache_read']}"
            f"  requests={totals['request_count']}"
        )

    proxy.set_on_update(_on_update)
    proxy.start()

    print(f"Token-counter proxy listening on http://127.0.0.1:{port}")
    print("Configure your Anthropic client base URL to point here.")
    print("Press Ctrl-C to stop.\n")

    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nShutting down…")
        proxy.stop()