claude-subagent-cache-proxy.py

claude-subagent-cache-proxy.py

#!/usr/bin/env python3
"""
Transparent Anthropic API proxy that fixes prompt-cache breakpoint placement
for our stage-1/stage-2 labeling subagents.

Problem: Claude Code sends STAGE1 (the big taxonomy prompt) and the per-item
request as ONE text block in the user message, and auto-caches the whole block.
Since the item varies per subagent, the cached unit is unique → each subagent
WRITES the 37k STAGE1 instead of READING a shared copy.

Fix: intercept requests whose user content contains our STAGE1 signature, split
the text block at the "<request>" boundary into two blocks, and put a
cache_control breakpoint on the STAGE1 half. Then STAGE1 becomes a stable shared
prefix that all subagents READ (~0.1x), and only the small item is fresh.

Everything else (main Opus loop, non-matching requests) passes through untouched.

Run:   python3 cache_proxy.py            # listens on 127.0.0.1:8787
Then relaunch Claude Code with:
       ANTHROPIC_BASE_URL=http://127.0.0.1:8787 claude
"""
import json, sys, http.server, socketserver, urllib.request, urllib.error

UPSTREAM = "https://api.anthropic.com"
HOST = "127.0.0.1"
PORT = 8787

# Explicit breakpoint markers WE inject in the prompt. The proxy caches
# everything up to </cache> and keeps the item after it fresh. No guessing.
OPEN = "<cache>"
CLOSE = "</cache>"

# Tools a schema-classifier subagent is actually allowed to call. Everything else
# Workflow hands it is dead weight we strip from OUR (<cache>-marked) requests.
KEEP_TOOLS = {"StructuredOutput"}

rewrites = {"seen": 0, "rewritten": 0}

# --- one-shot dump of system+tools from the first subagent request we see ---
# Remove after inspection. Writes once, then never again.
DUMPFILE = "/tmp/cache_proxy_dump.json"
_dumped = {"done": False}

def dump_wrapper(body):
    if _dumped["done"]:
        return
    _dumped["done"] = True
    tools = body.get("tools") or []

    def block_view(b):
        if not isinstance(b, dict):
            return {"type": "raw", "chars": len(str(b))}
        txt = b.get("text", "")
        return {"type": b.get("type"),
                "chars": len(json.dumps(b, ensure_ascii=False)),
                "cache_control": bool(b.get("cache_control")),
                "preview": (txt[:300] if isinstance(txt, str) else "")}

    msgs_view = []
    for m in (body.get("messages") or []):
        c = m.get("content")
        blocks = [{"type": "text", "text": c}] if isinstance(c, str) else (c or [])
        msgs_view.append({"role": m.get("role"),
                          "blocks": [block_view(b) for b in blocks]})

    summary = {
        "system": body.get("system"),
        "system_chars": len(json.dumps(body.get("system"), ensure_ascii=False)),
        "n_tools": len(tools),
        "tools": [{"name": t.get("name"),
                   "chars": len(json.dumps(t, ensure_ascii=False))} for t in tools],
        "tools_full": tools,
        "messages_view": msgs_view,
        "top_keys": sorted(body.keys()),
    }
    try:
        with open(DUMPFILE, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=1)
        sys.stderr.write(f"[proxy] DUMPED system+tools -> {DUMPFILE}\n")
    except Exception as e:
        sys.stderr.write(f"[proxy] dump failed: {e}\n")

def split_text(t: str):
    """Return (stage1_clean, item) if the marker is present, else None."""
    if CLOSE not in t:
        return None
    i = t.find(CLOSE)
    stage1 = t[:i].replace(OPEN, "")          # frozen cacheable prefix, tags stripped
    item = t[i + len(CLOSE):]                  # fresh, varying suffix
    return stage1, item

def rewrite_body(raw: bytes) -> bytes:
    """Inject a cache_control breakpoint at our <cache>..</cache> marker.

    Only touches a text block that STARTS with <cache> — that is uniquely the
    subagent prompt we build. The orchestrator's own history mentions the marker
    in prose / file dumps but never as the first chars of a content block, so
    those requests pass through untouched (otherwise we'd blow the 4-breakpoint
    cap on the main loop). We also clear any pre-existing cache_control so the
    total can never exceed Anthropic's limit of 4.
    """
    try:
        body = json.loads(raw)
    except Exception:
        return raw
    msgs = body.get("messages")
    if not isinstance(msgs, list):
        return raw

    # Is this one of OUR subagent requests? Require a block that starts with the
    # open marker; bail out (pass through) otherwise.
    def is_ours(blocks):
        if not isinstance(blocks, list):
            return False
        return any(isinstance(b, dict) and b.get("type") == "text"
                   and b.get("text", "").lstrip().startswith(OPEN)
                   and CLOSE in b.get("text", "") for b in blocks)

    hit = any(is_ours([{"type": "text", "text": m.get("content")}]
                       if isinstance(m.get("content"), str) else m.get("content"))
              for m in msgs)
    if not hit:
        return raw
    dump_wrapper(body)

    changed = False

    # Strip the tool balast. Workflow hands every subagent the full orchestrator
    # toolset (~65 tools, ~28k tokens) but a schema-classifier can only ever call
    # StructuredOutput. Drop the other 64 so we stop paying a cached read of tool
    # defs the subagent physically cannot invoke. Keep the schema tool untouched
    # (its input_schema carries the category enum — the actual contract).
    tools = body.get("tools")
    if isinstance(tools, list):
        kept = [t for t in tools if isinstance(t, dict) and t.get("name") in KEEP_TOOLS]
        if kept and len(kept) != len(tools):
            for t in kept:
                t.pop("cache_control", None)   # tiny now; no point caching
            body["tools"] = kept
            changed = True
    kept_msgs = []
    for m in msgs:
        c = m.get("content")
        blocks = [{"type": "text", "text": c}] if isinstance(c, str) else c
        if not isinstance(blocks, list):
            kept_msgs.append(m)
            continue
        newblocks = []
        for b in blocks:
            if isinstance(b, dict):
                b.pop("cache_control", None)        # clear existing breakpoints
            txt = b.get("text", "") if isinstance(b, dict) else ""
            # Drop the injected workspace context: Claude Code prepends the
            # skills list and the whole CLAUDE.md/memory header as
            # <system-reminder> text blocks. A classifier subagent never needs
            # them, yet they sit inside the cached prefix (~7.6k tokens/item).
            # Strip any text block that starts with the reminder marker; keep our
            # <cache> block and all non-text blocks (tool_result, images).
            if isinstance(b, dict) and b.get("type") == "text" \
                    and txt.lstrip().startswith("<system-reminder>"):
                changed = True
                continue
            starts = isinstance(b, dict) and b.get("type") == "text" \
                and txt.lstrip().startswith(OPEN) and CLOSE in txt
            sp = split_text(txt) if starts else None
            if sp:
                stage1, item = sp
                newblocks.append({"type": "text", "text": stage1,
                                  "cache_control": {"type": "ephemeral"}})
                if item.strip():
                    newblocks.append({"type": "text", "text": item})
                changed = True
            else:
                newblocks.append(b)
        # If stripping emptied the message (e.g. a role=system message whose
        # only block was the skills <system-reminder>), drop the whole message.
        # The API rejects a message with content == [], so leaving an empty
        # shell here is what caused the "must contain at least one block" 400.
        if not newblocks:
            changed = True
            continue
        m["content"] = newblocks
        kept_msgs.append(m)
    body["messages"] = kept_msgs
    if changed:
        rewrites["rewritten"] += 1
        # one-line audit of what we actually forward (tools kept, prefix sizes)
        try:
            to = body.get("tools") or []
            sys_c = len(json.dumps(body.get("system"), ensure_ascii=False))
            tool_c = sum(len(json.dumps(t, ensure_ascii=False)) for t in to)
            msg_c = len(json.dumps(body.get("messages"), ensure_ascii=False))
            with open("/tmp/cache_proxy_out.log", "a", encoding="utf-8") as f:
                f.write(f"tools_out={len(to)} tool_chars={tool_c} sys_chars={sys_c} "
                        f"msg_chars={msg_c} names={[t.get('name') for t in to]}\n")
        except Exception:
            pass
        return json.dumps(body, ensure_ascii=False).encode("utf-8")
    return raw

class Handler(http.server.BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    def _proxy(self):
        rewrites["seen"] += 1
        length = int(self.headers.get("Content-Length", 0) or 0)
        raw = self.rfile.read(length) if length else b""
        # NOTE: path is "/v1/messages?beta=true" — match the substring, not a
        # suffix, or the query string defeats the guard and we pass through raw.
        if raw and "/v1/messages" in self.path:
            raw = rewrite_body(raw)
        # rebuild upstream request, preserve auth/beta headers
        url = UPSTREAM + self.path
        req = urllib.request.Request(url, data=raw if raw else None, method=self.command)
        for k, v in self.headers.items():
            if k.lower() in ("host", "content-length", "connection", "accept-encoding"):
                continue
            req.add_header(k, v)
        if raw:
            req.add_header("Content-Length", str(len(raw)))
        try:
            up = urllib.request.urlopen(req, timeout=600)
        except urllib.error.HTTPError as e:
            up = e  # forward error responses (status + body) verbatim
        except Exception as e:
            self.send_response(502); self.end_headers()
            self.wfile.write(str(e).encode()); return
        # stream response back (works for SSE too)
        self.send_response(up.status if hasattr(up, "status") else up.code)
        hop = ("transfer-encoding", "content-length", "connection", "content-encoding")
        for k, v in up.headers.items():
            if k.lower() in hop:
                continue
            self.send_header(k, v)
        self.send_header("Connection", "close")
        self.end_headers()
        while True:
            chunk = up.read(8192)
            if not chunk:
                break
            try:
                self.wfile.write(chunk); self.wfile.flush()
            except BrokenPipeError:
                break

    do_POST = _proxy
    do_GET = _proxy

    def log_message(self, *a):
        sys.stderr.write(f"[proxy] seen={rewrites['seen']} rewritten={rewrites['rewritten']}  {self.command} {self.path}\n")

class Threaded(socketserver.ThreadingMixIn, http.server.HTTPServer):
    daemon_threads = True

    def handle_error(self, request, client_address):
        # The client (Claude Code) routinely drops a connection before we finish
        # streaming — cancelled/superseded requests. That surfaces as
        # BrokenPipe/ConnectionReset from deep in socketserver; it is expected
        # and harmless, so swallow it instead of dumping a traceback.
        exc = sys.exc_info()[1]
        if isinstance(exc, (BrokenPipeError, ConnectionResetError)):
            return
        super().handle_error(request, client_address)

if __name__ == "__main__":
    print(f"cache-proxy → {UPSTREAM}  listening http://{HOST}:{PORT}", file=sys.stderr)
    print(f"relaunch:  ANTHROPIC_BASE_URL=http://{HOST}:{PORT} claude", file=sys.stderr)
    Threaded((HOST, PORT), Handler).serve_forever()