trl-token-reduction 0.1.0__py3-none-any.whl

plugin/__init__.py

@@ -0,0 +1,2 @@
+"""Claude Code plugin: expose the AST retrieval engine so the agent pulls only
+the relevant code slices instead of grepping + dumping whole files."""

plugin/claude-code/.mcp.json

@@ -0,0 +1,16 @@
+{
+  "mcpServers": {
+    "trl-retrieve": {
+      "command": "python",
+      "args": [
+        "-m",
+        "plugin.mcp_server"
+      ],
+      "cwd": "/ABSOLUTE/PATH/TO/trl-token-reduction",
+      "env": {
+        "PYTHONPATH": "/ABSOLUTE/PATH/TO/trl-token-reduction",
+        "TRL_REPO": "."
+      }
+    }
+  }
+}

plugin/claude-code/SKILL.md

@@ -0,0 +1,7 @@
+# trl-retrieve (code retrieval)
+
+When you need to understand this codebase, DO NOT grep and read whole files.
+Call the MCP tool **`retrieve_code`** (server `trl-retrieve`) with a natural-language
+question — it returns the exact relevant source slices (function/class/method) at a
+fraction of the tokens. Use **`explain_symbol`** to fetch the full source of a named
+symbol. Only fall back to reading whole files if retrieval returns nothing relevant.

plugin/cli.py

@@ -0,0 +1,31 @@
+"""Bash-callable retrieval for Claude Code (and for humans).
+
+  python -m plugin.cli "how does auth work?" [--repo DIR] [--budget 1200] [--k 8]
+
+Prints only the relevant code slices (file:line + source) -- a drop-in for
+`grep -r` + reading whole files, at a fraction of the tokens."""
+import argparse, os, sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from plugin.index_store import get_index
+from trl.retrieval import retrieve
+from trl.util import count_tokens
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("query")
+    ap.add_argument("--repo", default=None)   # None -> resolve $TRL_REPO/$CLAUDE_PROJECT_DIR/git-root/cwd
+    ap.add_argument("--budget", type=int, default=1200)
+    ap.add_argument("--k", type=int, default=8)
+    a = ap.parse_args()
+    idx = get_index(a.repo)   # index_store resolves + caches at <repo>/.trl/index.json
+    r = retrieve(idx, a.query, token_budget=a.budget, k=a.k)
+    if not r["symbols"]:
+        print("(no relevant symbols found)"); return
+    hdr = (f"# {len(r['symbols'])} slices, {r['tokens']} tokens "
+           f"(indexed {len(idx['files'])} files, {len(idx['symbols'])} symbols)")
+    print(hdr); print(r["context"])
+
+
+if __name__ == "__main__":
+    main()

plugin/codex/config.toml

@@ -0,0 +1,15 @@
+# Codex retrieval plugin — merge this into ~/.codex/config.toml (global) or
+# .codex/config.toml (project-scoped; trusted projects only).
+#
+# Same STDIO MCP server as the Claude Code plugin — the protocol is client-neutral,
+# so nothing about the server changes. Edit the two paths if the project moves.
+# Set TRL_REPO to the repo you want Codex to retrieve from ("." = the cwd Codex runs in).
+
+[mcp_servers.trl-retrieve]
+command = "python"
+args = ["-m", "plugin.mcp_server"]
+env = { PYTHONPATH = "/ABSOLUTE/PATH/TO/trl-token-reduction", TRL_REPO = "." }
+
+# Optional — route Codex's API-key calls through the token-reduction proxy too
+# (applies to API-key mode, not ChatGPT-subscription mode):
+#   openai_base_url = "http://localhost:8899/v1"

plugin/index_store.py

@@ -0,0 +1,71 @@
+"""Shared: resolve the user's project at RUNTIME and build-or-load a persistent,
+incremental AST index for it.
+
+Why runtime resolution: a Claude Code plugin MCP server's cwd is the PLUGIN root
+(not the user's project), and ${...} interpolation inside a plugin's .mcp.json is
+unreliable. So we never trust the manifest to pass the path -- we resolve it here.
+"""
+import os, sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from trl.retrieval import build_index, save_index, load_index
+
+INDEX_DIR = ".trl"
+INDEX_NAME = "index.json"
+
+
+def _git_root(start):
+    """Nearest ancestor of `start` containing a .git dir, else None."""
+    d = os.path.abspath(start)
+    while True:
+        if os.path.isdir(os.path.join(d, ".git")):
+            return d
+        parent = os.path.dirname(d)
+        if parent == d:
+            return None
+        d = parent
+
+
+def _resolve_repo(repo=None):
+    """Resolve the target project dir at RUNTIME. Order:
+    explicit arg -> $TRL_REPO -> $CLAUDE_PROJECT_DIR -> nearest .git above cwd -> cwd.
+    Always returns a path."""
+    if repo:
+        return os.path.abspath(repo)
+    for env in ("TRL_REPO", "CLAUDE_PROJECT_DIR"):
+        v = os.environ.get(env)
+        if v:
+            return os.path.abspath(v)
+    g = _git_root(os.getcwd())
+    return g if g else os.getcwd()
+
+
+def _has_explicit_repo(repo=None):
+    """True only for a RELIABLE project signal (explicit arg or env). git-root/cwd are
+    NOT reliable for a plugin MCP server, whose cwd is the plugin root -- callers that
+    run there must gate on this to avoid indexing the wrong tree."""
+    return bool(repo or os.environ.get("TRL_REPO") or os.environ.get("CLAUDE_PROJECT_DIR"))
+
+
+def index_path(repo):
+    return os.path.join(repo, INDEX_DIR, INDEX_NAME)
+
+
+def get_index(repo=None):
+    """Up-to-date index for the resolved repo, cached at <repo>/.trl/index.json so
+    `/trl-index` and `retrieve_code` agree regardless of cwd. Incremental: only
+    changed files are re-parsed."""
+    repo = _resolve_repo(repo)
+    cache = index_path(repo)
+    prev = None
+    if os.path.exists(cache):
+        try:
+            prev = load_index(cache)
+        except Exception:
+            prev = None
+    idx = build_index(repo, prev=prev)
+    try:
+        os.makedirs(os.path.dirname(cache), exist_ok=True)
+        save_index(idx, cache)
+    except Exception:
+        pass
+    return idx

plugin/mcp_server.py

@@ -0,0 +1,143 @@
+"""MCP server exposing the retrieval engine to Claude Code (or any MCP client).
+
+The target project is resolved AT RUNTIME (see plugin.index_store._resolve_repo):
+$TRL_REPO -> $CLAUDE_PROJECT_DIR -> nearest .git above cwd -> cwd. We do NOT rely on
+manifest ${...} interpolation. Because an MCP server's cwd is the PLUGIN root, if no
+reliable signal is present the tools return a "run /trl-index or set TRL_REPO" hint
+instead of silently indexing the plugin's own tree. Both tools accept repo= override.
+
+Tools:
+  retrieve_code(query, repo=?)   -> relevant slices instead of whole files
+  explain_symbol(name, repo=?)   -> exact source of a named function/class/method
+"""
+import os, sys
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from plugin.index_store import get_index, _resolve_repo, _has_explicit_repo
+from trl.retrieval import retrieve
+from plugin.retrieval_policy import (RetrievalGuard, DEFAULT_BUDGET, DEFAULT_K,
+                                     ADAPTIVE_CEIL, IDLE_RESET_S)
+
+try:
+    from mcp.server.fastmcp import FastMCP
+except Exception:
+    FastMCP = None
+
+_PLUGIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+_UNRESOLVED = ("TRL: couldn't resolve your project -- this MCP server runs from the "
+               "plugin directory. Run /trl-index in your project, set TRL_REPO, or "
+               "call this tool with repo='/path/to/your/project'.")
+
+
+def _target(repo=None):
+    """Resolved project path, or None if we can't CONFIDENTLY resolve it (no explicit
+    arg/env and resolution falls back to the plugin root)."""
+    if _has_explicit_repo(repo):
+        return _resolve_repo(repo)
+    resolved = _resolve_repo(repo)
+    if os.path.abspath(resolved) == _PLUGIN_ROOT:
+        return None
+    return resolved
+
+
+def _log_savings(kind, query, target, slice_tokens, symbols):
+    """Best-effort adoption + savings log (opt-in via env TRL_SAVINGS_LOG). Records the
+    slices actually returned vs the whole-file counterfactual (the files those slices came
+    from), so cumulative savings sum across sessions. NEVER raises -- must not break retrieval."""
+    path = os.environ.get("TRL_SAVINGS_LOG")
+    if path and path.strip().lower() == "off":
+        return
+    if not path:
+        # Deterministic default: log into the repo we are actually retrieving from, so
+        # tracking never depends on launcher env / cwd (fixes the split/undercounted logs).
+        try:
+            if target and os.path.isdir(target):
+                path = os.path.join(target, ".trl", "savings.jsonl")
+            else:
+                return
+        except Exception:
+            return
+    try:
+        import json, time
+        from trl.util import count_tokens
+        files = {s.file for s in symbols}
+        whole = 0
+        for f in files:
+            try:
+                whole += count_tokens(open(f, encoding="utf-8", errors="ignore").read())
+            except Exception:
+                pass
+        rec = {"ts": time.time(), "tool": kind, "repo": target, "query": (query or "")[:120],
+               "slice_tokens": int(slice_tokens), "wholefile_tokens": whole,
+               "saved": max(0, whole - int(slice_tokens)), "n_slices": len(symbols),
+               "n_files": len(files)}
+        os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
+        with open(path, "a", encoding="utf-8") as fh:
+            fh.write(json.dumps(rec) + "\n")
+    except Exception:
+        pass
+
+
+if FastMCP is not None:
+    mcp = FastMCP("trl-retrieve")
+    _GUARD = RetrievalGuard(idle_reset_s=IDLE_RESET_S)  # over-retrieval guard, resets on idle
+
+    @mcp.tool()
+    def retrieve_code(query: str = "", budget: int = DEFAULT_BUDGET, k: int = DEFAULT_K,
+                      repo: str = "", token_budget: int = 0, limit: int = 0, path: str = "",
+                      project: str = "") -> str:
+        """Return the most relevant code slices for a question, instead of reading whole
+        files. Make ONE good, specific query then ANSWER from the result -- do not re-run
+        variations of the same question. Args: query (REQUIRED), k (max slices, default 5),
+        budget (token budget, default 700), repo (optional project path). Aliases
+        token_budget/limit/path/project accepted."""
+        query = (query or "").strip()
+        if not query:
+            return ("retrieve_code needs a 'query' string, e.g. "
+                    "retrieve_code(query=\"how does the stash update\").")
+        # Over-retrieval guard: nudge/stop a runaway loop before it re-bills huge history.
+        blocked, prefix = _GUARD.check(query)
+        if blocked:
+            return prefix
+        # Clamp to lean caps so a large model-supplied budget/k can't defeat the fix.
+        budget = min(token_budget or budget or DEFAULT_BUDGET, DEFAULT_BUDGET)
+        k = min(limit or k or DEFAULT_K, DEFAULT_K)
+        target = _target(repo or path or project or None)
+        if target is None:
+            return _UNRESOLVED
+        r = retrieve(get_index(target), query, token_budget=budget, k=k, adaptive_ceil=ADAPTIVE_CEIL)
+        _log_savings("retrieve_code", query, target, r["tokens"], r["symbols"])
+        return prefix + (r["context"] or "(no relevant symbols found)")
+
+    @mcp.tool()
+    def explain_symbol(name: str = "", repo: str = "", path: str = "", project: str = "") -> str:
+        """Return the exact source of a function/class/method by name. Use sparingly, after
+        retrieve_code, for one specific symbol. Args: name (REQUIRED), repo (optional project
+        path). Aliases path/project accepted."""
+        name = (name or "").strip()
+        if not name:
+            return "explain_symbol needs a 'name', e.g. explain_symbol(name=\"StashService\")."
+        blocked, prefix = _GUARD.check("symbol:" + name)
+        if blocked:
+            return prefix
+        target = _target(repo or path or project or None)
+        if target is None:
+            return _UNRESOLVED
+        idx = get_index(target)
+        hits = [s for s in idx["symbols"] if s.name == name]
+        if not hits:
+            return f"(no symbol named {name})"
+        picked = hits[:3]
+        from trl.util import count_tokens
+        _log_savings("explain_symbol", name, target, sum(count_tokens(s.source) for s in picked), picked)
+        return prefix + "\n\n".join(f"# {s.file}:{s.start_line}-{s.end_line} ({s.kind})\n{s.source}"
+                           for s in picked)
+
+    def main():
+        mcp.run()
+else:
+    def main():
+        print("install the MCP SDK:  pip install mcp", file=sys.stderr); sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

plugin/retrieval_policy.py

@@ -0,0 +1,109 @@
+"""Single source of truth for the retrieval POLICY, shared by the shipping MCP
+(plugin/mcp_server.py) and the paired billed-token harness
+(validate/paired_billed_bench.py) so the benchmark measures exactly what ships.
+
+Background: the trl-retrieve MCP is quality-neutral and cheaper per retrieval, BUT
+had an over-retrieval failure mode -- an agent could loop retrieve_code 25+ times and
+the re-billed slice history ballooned to 280-380k input tokens, flipping the total
+cost above a plain grep+read agent. This module fixes that with two levers, applied
+identically in the tool and in the harness:
+
+  Lever 1 -- leaner slices: smaller k + token budget + a lower adaptive ceiling, so
+  each result is small and per-call history growth is bounded. (retrieve() defaults to
+  an adaptive ceiling of 6000 tokens/call; we cap the retrieval path far below that.)
+
+  Lever 2 -- a per-session repeat-call guard: after a few calls, nudge the agent to
+  ANSWER instead of retrieving more; after a hard cap, stop returning slices entirely
+  so the loop cannot run away. Exact-duplicate queries are short-circuited immediately.
+
+Positioning guardrail (do NOT regress): this makes the tool cheaper AND quality-neutral
+vs a shell agent post-fix. It does NOT restore any blanket "93% billed savings" claim;
+93% was content reduction vs whole-file reads, not billed savings in an agent loop.
+"""
+
+# --- Lever 1: leaner slices (was budget=1200, k=8, adaptive ceiling=6000) ---
+DEFAULT_BUDGET = 700     # token budget per retrieve_code call
+DEFAULT_K = 5            # max slices per call
+ADAPTIVE_CEIL = 1500     # hard cap on per-call returned context (overrides retrieve()'s 6000)
+
+# --- Lever 2: repeat-call guard, per session / per agent loop ---
+SOFT_NUDGE_AFTER = 3     # calls 4.. get a "you likely have enough, answer now" prefix
+HARD_STOP_AFTER = 6      # calls 7.. return only a stop-and-answer nudge (no new slices)
+IDLE_RESET_S = 90        # (shipping MCP only) auto-reset the counter after this much idle,
+                         # so a NEW user turn starts fresh while an intra-turn runaway
+                         # loop (back-to-back calls) still hits the cap.
+
+RETRIEVE_DESC = (
+    "Return the most relevant code slices for a question, instead of reading whole "
+    "files. Make ONE good, specific query and then ANSWER from the result -- do NOT "
+    "re-run variations of the same question; you rarely need more than a couple of "
+    "retrievals. Args: query (the question, REQUIRED string), k (max slices, default 5), "
+    "budget (token budget, default 700).")
+EXPLAIN_DESC = (
+    "Return the exact source of a function/class/method/table by name. Use sparingly, "
+    "after retrieve_code, when you need one specific named symbol in full. Args: name (REQUIRED).")
+
+_SOFT_MSG = ("NOTE: you have already retrieved code {n} times this turn -- you very likely "
+             "have enough context now. Prefer to ANSWER the question. Only retrieve again if "
+             "a specific, named thing is still missing.\n\n")
+_HARD_MSG = ("You have retrieved code {n} times this turn -- that is enough context. Stop "
+             "retrieving and ANSWER the user's question now from what you already have. "
+             "(Further retrieval is paused for this turn to prevent runaway token cost.)")
+_DUP_MSG = ("You already ran this exact query this turn and the result has not changed. "
+            "Answer now from the context you already retrieved instead of re-retrieving.")
+
+
+def _norm(query):
+    return " ".join((query or "").lower().split())
+
+
+class RetrievalGuard:
+    """Per-session counter shared by both retrieval tools. One instance lives for the
+    life of ONE agent loop: the MCP server process (shipping), or one run_session in the
+    harness (which builds a fresh guard per session).
+
+    Usage: call check(query) BEFORE retrieving.
+      returns (blocked, prefix)
+        blocked=True  -> return `prefix` as the WHOLE tool result; do NOT retrieve
+                         (no slices are added to history -- this is what caps the loop).
+        blocked=False -> prepend `prefix` (possibly "") to the real slices.
+
+    idle_reset_s: if set (shipping MCP), the counter auto-resets when more than that many
+    seconds elapse between calls, so distinct user turns don't share one budget while a
+    back-to-back runaway loop still trips the cap. Left None in the harness for determinism.
+    """
+
+    def __init__(self, idle_reset_s=None, clock=None):
+        self.idle_reset_s = idle_reset_s
+        self._clock = clock
+        self._last = None
+        self.n = 0
+        self._seen = set()
+
+    def _now(self):
+        if self._clock is not None:
+            return self._clock()
+        import time
+        return time.monotonic()
+
+    def reset(self):
+        self.n = 0
+        self._seen = set()
+
+    def check(self, query):
+        if self.idle_reset_s is not None:
+            now = self._now()
+            if self._last is not None and (now - self._last) > self.idle_reset_s:
+                self.reset()
+            self._last = now
+        norm = _norm(query)
+        if norm and norm in self._seen:
+            return True, _DUP_MSG              # exact repeat: don't count, don't retrieve
+        if self.n >= HARD_STOP_AFTER:
+            return True, _HARD_MSG.format(n=self.n)   # cap reached: stop, no slices
+        self.n += 1
+        if norm:
+            self._seen.add(norm)
+        if self.n > SOFT_NUDGE_AFTER:
+            return False, _SOFT_MSG.format(n=self.n)   # nudge, but still retrieve
+        return False, ""

proxy/__init__.py

@@ -0,0 +1,5 @@
+"""v0 product skin: an OpenAI-compatible drop-in proxy. Point your client's
+base_url at this server and it applies the token-reduction levers to every
+request before forwarding upstream -- zero code change for the caller."""
+from .transform import transform_chat_request
+__all__ = ["transform_chat_request", "transform_anthropic_request"]

proxy/compress_endpoint.py

@@ -0,0 +1,49 @@
+"""Local /compress endpoint for the composer-compression browser extension.
+
+No upstream call. Takes the user's pasted bulk context, runs the SAME
+compression + fact-guard (or the text retriever) the rest of the engine uses,
+and returns the shorter text plus token counts and the numeric facts the guard
+guarantees survived -- so the extension can show a trustworthy pre-send preview.
+
+Isolated + pure so it unit-tests offline (like transform.py)."""
+import re
+
+from trl.message import Message, TOOL_RESULT
+from trl.compress import compress_request
+from trl.util import count_tokens
+
+_NUM = re.compile(r"-?\d[\d,]*")
+
+
+def handle_compress(req: dict, engine) -> dict:
+    text = (req.get("text") or "").strip()
+    if not text:
+        return {"error": "no text"}
+    mode = req.get("mode", "compress")
+    question = (req.get("question") or "").strip()
+    budget = int(req.get("budget", 1200))
+    before = count_tokens(text)
+
+    if mode == "retrieve":
+        # Doc-slice mode: keep only the passages relevant to `question`.
+        from trl.retrieval import build_text_index, retrieve_text
+        idx = build_text_index({"pasted": text})
+        r = retrieve_text(idx, question or text, token_budget=budget, rerank=False)
+        out = r.get("context", "") or text
+    else:
+        # Compression mode: summarize the tail, then the deterministic fact-guard
+        # re-injects any dropped number. Short text (<200 chars) passes through.
+        msg = Message("tool", TOOL_RESULT, text)
+        new_msgs, _ = compress_request([msg], engine.mode, engine.local)
+        out = "\n".join(m.content for m in new_msgs).strip() or text
+
+    after = count_tokens(out)
+    preserved = sorted({n.replace(",", "") for n in _NUM.findall(text)}, key=len)
+    return {
+        "compressed": out,
+        "tokens_before": before,
+        "tokens_after": after,
+        "saved_pct": round(100 * (1 - after / before), 1) if before else 0.0,
+        "preserved_facts": preserved,
+        "mode": mode,
+    }

proxy/server.py

@@ -0,0 +1,179 @@
+"""Drop-in OpenAI-compatible proxy. Point your client's base_url here; every
+/v1/chat/completions request gets the token-reduction levers applied, then is
+forwarded upstream with YOUR Authorization header (the proxy never stores keys).
+
+Run:  python -m proxy.server            # listens on :8899, forwards to OpenAI
+Then: client base_url = http://localhost:8899/v1   (key unchanged)
+
+Stdlib only (http.server + urllib) so it has zero extra deps. Adds response
+header  X-TRL-Tokens-Saved  so you can see the reduction per request."""
+from __future__ import annotations
+
+import json
+import os
+import sys
+import urllib.request
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from trl import Engine
+from trl.util import load_config
+from proxy.transform import transform_chat_request, transform_anthropic_request
+from proxy.compress_endpoint import handle_compress
+
+_DEFAULT_CFG = {
+    "arms": {"treatment": {"native_prompt_cache": True, "compress_history": True,
+                            "compress_tool_outputs": True, "compression_mode": "safe"}},
+    "local_model": {"provider": "none"},
+    "retrieval": {"enabled": True, "token_budget": 800, "k": 8, "expand_call_graph": True},
+}
+
+
+def _load_cfg():
+    # Prefer an explicit path, then ./config.yaml, then a repo-relative copy;
+    # fall back to a sane built-in default so the pip-installed CLI runs anywhere.
+    import os as _os
+    for cand in (_os.environ.get("TRL_CONFIG"), "config.yaml",
+                 _os.path.join(_os.path.dirname(__file__), "..", "config.yaml")):
+        if cand and _os.path.exists(cand):
+            try:
+                return load_config(cand)
+            except Exception:
+                pass
+    return _DEFAULT_CFG
+
+
+_CFG = _load_cfg()
+_ENGINE = Engine(_CFG)
+_UPSTREAM = os.environ.get("TRL_UPSTREAM", "https://api.openai.com")
+
+
+_MAX_BODY = 10 * 1024 * 1024        # 10 MB request-body cap (reject before read)
+_ALLOWED_ORIGINS = {"https://claude.ai"}
+_CORS_BASE = {
+    "Access-Control-Allow-Methods": "POST, OPTIONS",
+    "Access-Control-Allow-Headers": "Content-Type",
+    "Access-Control-Allow-Private-Network": "true",
+}
+
+
+class Handler(BaseHTTPRequestHandler):
+    def _cors(self):
+        # Scope CORS: the browser extension's origin is chrome-extension://<id>
+        # (random per install, so allow the scheme) plus claude.ai. Arbitrary
+        # websites get NO Access-Control-Allow-Origin, so a browser refuses to let
+        # them read responses from the local proxy. Non-browser SDK clients ignore
+        # CORS entirely, so this doesn't affect the base_url-swap proxy usage.
+        h = dict(_CORS_BASE)
+        origin = self.headers.get("Origin", "")
+        if origin.startswith("chrome-extension://") or origin in _ALLOWED_ORIGINS:
+            h["Access-Control-Allow-Origin"] = origin
+            h["Vary"] = "Origin"
+        return h
+
+    def do_OPTIONS(self):
+        self._send(204, b"", self._cors())
+
+    def do_GET(self):
+        # Lightweight liveness probe so clients (e.g. the browser extension) can
+        # tell "engine running" from "engine down" and show a useful message.
+        if self.path.rstrip("/").endswith("/health"):
+            return self._send(200, b'{"status":"ok","service":"trl-proxy"}',
+                              {"Content-Type": "application/json", **self._cors()})
+        return self._send(404, b'{"error":"not found"}',
+                          {"Content-Type": "application/json", **self._cors()})
+
+    def _send(self, code, body: bytes, headers=None):
+        self.send_response(code)
+        for k, v in (headers or {}).items():
+            self.send_header(k, v)
+        self.end_headers()
+        self.wfile.write(body)
+
+    def do_POST(self):
+        try:
+            n = int(self.headers.get("Content-Length") or 0)
+        except (TypeError, ValueError):
+            return self._send(411, b'{"error":"missing or invalid Content-Length"}',
+                              {"Content-Type": "application/json", **self._cors()})
+        if n > _MAX_BODY:                       # reject oversized bodies BEFORE read
+            return self._send(413, b'{"error":"request body too large"}',
+                              {"Content-Type": "application/json", **self._cors()})
+        raw = self.rfile.read(n)
+        path = self.path.rstrip("/")
+        if path.endswith("/compress"):
+            try:
+                result = handle_compress(json.loads(raw), _ENGINE)
+            except Exception as e:
+                return self._send(400, json.dumps({"error": str(e)}).encode(), self._cors())
+            return self._send(200, json.dumps(result).encode(),
+                              {"Content-Type": "application/json", **self._cors()})
+        is_anthropic = path.endswith("/messages")
+        is_openai = path.endswith("/chat/completions")
+        if not (is_anthropic or is_openai):
+            return self._send(404, b'{"error":"only /v1/chat/completions or /v1/messages"}',
+                              {"Content-Type": "application/json"})
+        try:
+            req = json.loads(raw)
+            if is_anthropic:
+                new_req, meta = transform_anthropic_request(req, _ENGINE)
+            else:
+                new_req, meta = transform_chat_request(req, _ENGINE)
+        except Exception as e:
+            return self._send(400, json.dumps({"error": str(e)}).encode())
+
+        fwd_headers = {"Content-Type": "application/json"}
+        if self.headers.get("Authorization"):
+            fwd_headers["Authorization"] = self.headers.get("Authorization")
+        for h in ("x-api-key", "anthropic-version", "anthropic-beta", "OpenAI-Organization"):
+            if self.headers.get(h):
+                fwd_headers[h] = self.headers.get(h)
+        up = urllib.request.Request(_UPSTREAM + self.path,
+                                    data=json.dumps(new_req).encode(), headers=fwd_headers)
+        saved_hdrs = {"X-TRL-Tokens-Saved": str(meta["tokens_saved"]),
+                      "X-TRL-Tokens-Before": str(meta["tokens_before"]),
+                      "X-TRL-Tokens-After": str(meta["tokens_after"])}
+        streaming = False
+        try:
+            with urllib.request.urlopen(up, timeout=300) as r:
+                if new_req.get("stream"):
+                    # stream the upstream SSE response through, unchanged
+                    streaming = True
+                    self.send_response(r.status)
+                    self.send_header("Content-Type",
+                                     r.headers.get("Content-Type", "text/event-stream"))
+                    for k, v in saved_hdrs.items():
+                        self.send_header(k, v)
+                    self.end_headers()
+                    while True:
+                        chunk = r.read(2048)
+                        if not chunk:
+                            break
+                        self.wfile.write(chunk); self.wfile.flush()
+                else:
+                    body = r.read()
+                    self._send(r.status, body, {"Content-Type": "application/json", **saved_hdrs})
+        except urllib.error.HTTPError as e:
+            self._send(e.code, e.read())
+        except Exception as e:
+            if streaming:
+                # headers already sent mid-stream; a second status line would
+                # corrupt the response -- just drop the connection.
+                self.close_connection = True
+                return
+            self._send(502, json.dumps({"error": f"upstream: {e}"}).encode())
+
+    def log_message(self, *a):   # quiet
+        pass
+
+
+def main(port=None):
+    if port is None:   # honor TRL_PORT for the pip-installed `trl-proxy` script too
+        port = int(os.environ.get("TRL_PORT", "8899"))
+    print(f"TRL proxy on http://localhost:{port}  ->  {_UPSTREAM}")
+    print("point your client base_url at  http://localhost:%d/v1  (key unchanged)" % port)
+    ThreadingHTTPServer(("127.0.0.1", port), Handler).serve_forever()
+
+
+if __name__ == "__main__":
+    main()