diff --git a/config.yaml b/config.yaml index 8921bb4..16eed47 100644 --- a/config.yaml +++ b/config.yaml @@ -9,6 +9,9 @@ proxy: upstream: openai_url: "https://api.openai.com" anthropic_url: "https://api.anthropic.com" + # Used by Codex CLI when signed in with a ChatGPT subscription + # (Plus/Pro/Team) instead of an OpenAI API key. + chatgpt_backend_url: "https://chatgpt.com/backend-api/codex" # API keys: can also be set via environment variables # OPENAI_API_KEY and ANTHROPIC_API_KEY diff --git a/pyproject.toml b/pyproject.toml index ec4e640..fc970c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dependencies = [ "tiktoken>=0.5.0", "rich>=13.0.0", "pyyaml>=6.0", + "zstandard>=0.22", ] [project.optional-dependencies] diff --git a/tests/test_server.py b/tests/test_server.py index dcc46bb..17063b1 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -145,6 +145,76 @@ def test_messages_compression_with_system(self, mock_post, client: TestClient): content = forwarded_body["messages"][0]["content"] assert "..." in content # database.py should be skeletonized + @patch("httpx.AsyncClient.post") + def test_messages_preserves_oauth_authorization(self, mock_post, client: TestClient): + mock_response = MagicMock() + mock_response.content = json.dumps({"id": "msg-test", "content": []}).encode() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_post.return_value = mock_response + + payload = { + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Hi"}], + "stream": False, + } + + response = client.post( + "/v1/messages?beta=true", + json=payload, + headers={ + "Authorization": "Bearer oauth-token", + "anthropic-version": "2023-06-01", + "anthropic-beta": "oauth-2025-04-20", + }, + ) + + assert response.status_code == 200 + call_args = mock_post.call_args + assert call_args.args[0] == "https://fake-anthropic.com/v1/messages?beta=true" + forwarded_headers = call_args.kwargs["headers"] + assert forwarded_headers["Authorization"] == "Bearer oauth-token" + assert forwarded_headers["anthropic-version"] == "2023-06-01" + assert forwarded_headers["anthropic-beta"] == "oauth-2025-04-20" + assert "x-api-key" not in forwarded_headers + + @patch("httpx.AsyncClient.post") + def test_count_tokens_passes_through_anthropic_auth( + self, mock_post, client: TestClient + ): + mock_response = MagicMock() + mock_response.content = json.dumps({"input_tokens": 12}).encode() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_post.return_value = mock_response + + payload = { + "model": "claude-sonnet-4-20250514", + "messages": [{"role": "user", "content": "Count me"}], + } + + response = client.post( + "/v1/messages/count_tokens?beta=true", + json=payload, + headers={ + "Authorization": "Bearer oauth-token", + "anthropic-version": "2023-06-01", + "anthropic-beta": "oauth-2025-04-20", + }, + ) + + assert response.status_code == 200 + assert response.json() == {"input_tokens": 12} + call_args = mock_post.call_args + assert call_args.args[0] == ( + "https://fake-anthropic.com/v1/messages/count_tokens?beta=true" + ) + assert call_args.kwargs["json"] == payload + forwarded_headers = call_args.kwargs["headers"] + assert forwarded_headers["Authorization"] == "Bearer oauth-token" + assert forwarded_headers["anthropic-beta"] == "oauth-2025-04-20" + assert "x-api-key" not in forwarded_headers + class TestMetrics: def test_metrics_tracked(self, client: TestClient, metrics: SessionMetrics): diff --git a/tests/test_session_cache.py b/tests/test_session_cache.py index 8836341..e7a2314 100644 --- a/tests/test_session_cache.py +++ b/tests/test_session_cache.py @@ -239,6 +239,84 @@ def test_double_apply_does_not_add_extra_breakpoints(self, cache): markers = _count_cache_markers(body) assert markers == info1["breakpoints"] + def test_existing_breakpoints_consume_budget(self, cache): + msgs = [ + _msg("user", [{"type": "text", "text": f"turn {i}"}]) + for i in range(8) + ] + for msg in msgs[:3]: + msg["content"][0]["cache_control"] = {"type": "ephemeral"} + body = { + "system": "System prompt", + "messages": msgs, + "tools": [{"name": "Read"}], + } + + modified, info = cache.apply(body) + + assert _count_cache_markers(modified) <= MAX_BREAKPOINTS + assert info["breakpoints"] <= 1 + + def test_surplus_existing_breakpoints_are_trimmed(self, cache): + msgs = [ + _msg("user", [{"type": "text", "text": f"turn {i}"}]) + for i in range(5) + ] + body = { + "system": [ + {"type": "text", "text": "system", "cache_control": {"type": "ephemeral"}}, + ], + "messages": msgs, + "tools": [ + {"name": "Read", "cache_control": {"type": "ephemeral"}}, + ], + } + for msg in msgs[:3]: + msg["content"][0]["cache_control"] = {"type": "ephemeral"} + + modified, info = cache.apply(body) + + assert _count_cache_markers(modified) == MAX_BREAKPOINTS + assert info["breakpoints"] == 0 + + def test_injected_breakpoint_before_one_hour_ttl_is_promoted(self, cache): + body = { + "system": [ + {"type": "text", "text": "first"}, + { + "type": "text", + "text": "second", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + }, + ], + "messages": [_msg("user", f"turn {i}") for i in range(5)], + "tools": [{"name": "Read"}], + } + + modified, info = cache.apply(body) + + assert modified["tools"][0]["cache_control"]["ttl"] == "1h" + assert modified["system"][1]["cache_control"]["ttl"] == "1h" + assert _cache_ttls_are_valid(modified) + + def test_existing_short_request_ttls_are_normalized(self, cache): + body = { + "system": [ + {"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}}, + { + "type": "text", + "text": "b", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + }, + ], + "messages": [_msg("user", "short"), _msg("assistant", "ok")], + } + + modified, info = cache.apply(body) + + assert modified["system"][0]["cache_control"]["ttl"] == "1h" + assert _cache_ttls_are_valid(modified) + def _count_cache_markers(body: dict) -> int: count = 0 @@ -259,6 +337,20 @@ def _count_cache_markers(body: dict) -> int: return count +def _cache_ttls_are_valid(body: dict) -> bool: + seen_short = False + owners = SessionCache._iter_cache_control_owners(body) + for owner in owners: + cache_control = owner.get("cache_control") + ttl = cache_control.get("ttl") if isinstance(cache_control, dict) else None + is_one_hour = ttl == "1h" + if seen_short and is_one_hour: + return False + if not is_one_hour: + seen_short = True + return True + + # ────────────────────────────────────────────────────────── # Session tracking # ────────────────────────────────────────────────────────── diff --git a/token_tamer/config.py b/token_tamer/config.py index 3a19894..94c273d 100644 --- a/token_tamer/config.py +++ b/token_tamer/config.py @@ -25,6 +25,9 @@ class ProxyConfig: class UpstreamConfig: openai_url: str = "https://api.openai.com" anthropic_url: str = "https://api.anthropic.com" + # ChatGPT-subscription backend used by Codex CLI when the user is signed in + # with a ChatGPT Plus/Pro/Team account instead of an OpenAI API key. + chatgpt_backend_url: str = "https://chatgpt.com/backend-api/codex" @dataclass @@ -126,6 +129,9 @@ def load_config(config_path: Union[str, Path, None] = None) -> Config: config.upstream = UpstreamConfig( openai_url=raw["upstream"].get("openai_url", config.upstream.openai_url), anthropic_url=raw["upstream"].get("anthropic_url", config.upstream.anthropic_url), + chatgpt_backend_url=raw["upstream"].get( + "chatgpt_backend_url", config.upstream.chatgpt_backend_url + ), ) # API keys (env vars take priority) diff --git a/token_tamer/server.py b/token_tamer/server.py index 93b04a7..5d4d5d2 100644 --- a/token_tamer/server.py +++ b/token_tamer/server.py @@ -9,16 +9,29 @@ from __future__ import annotations import copy +import gzip +import io import json import logging import time +import zlib from typing import AsyncGenerator, List, Optional import httpx -from fastapi import FastAPI, Request, Response +from fastapi import FastAPI, HTTPException, Request, Response from fastapi.responses import StreamingResponse from starlette.middleware.base import BaseHTTPMiddleware +try: + import zstandard as _zstd +except ImportError: # pragma: no cover + _zstd = None + +try: + import brotli as _brotli +except ImportError: # pragma: no cover + _brotli = None + logger = logging.getLogger("token_tamer") from . import __version__ @@ -67,6 +80,186 @@ def _request_has_tools(body: dict) -> bool: return False +def _anthropic_forward_headers(headers: dict, config: Config) -> dict: + """Build upstream headers for Anthropic-compatible requests. + + Claude Code can authenticate with either API keys (`x-api-key`) or a + Claude.ai OAuth/subscription token (`Authorization: Bearer ...`). Preserve + whichever credential the client sent instead of forcing all requests into + the API-key header shape. + """ + forward_headers = { + "Content-Type": "application/json", + "anthropic-version": headers.get("anthropic-version", "2023-06-01"), + } + + auth_header = headers.get("authorization", "") + api_key = headers.get("x-api-key", "") + if auth_header: + forward_headers["Authorization"] = auth_header + elif api_key: + forward_headers["x-api-key"] = api_key + else: + configured_key = config.get_api_key("anthropic") + if configured_key: + forward_headers["x-api-key"] = configured_key + + # Forward Claude/Anthropic feature headers such as anthropic-beta. + for key, value in headers.items(): + if key.startswith("anthropic-") and key != "anthropic-version": + forward_headers[key] = value + + return forward_headers + + +def _with_query(url: str, request: Request) -> str: + """Append the original query string when proxying a route.""" + query = request.url.query + return f"{url}?{query}" if query else url + + +# Headers Codex CLI sends when authenticated with a ChatGPT subscription. +# `chatgpt-account-id` is the load-bearing one — its presence is how we detect +# subscription mode and how the upstream identifies which account to bill. +_CHATGPT_PASSTHROUGH_HEADERS = ( + "chatgpt-account-id", + "originator", + "session_id", + "version", +) + +# Hop-by-hop / connection-scoped headers we must NOT forward. Content-Length is +# stripped because we re-serialize the body; Accept-Encoding because we want +# raw bytes (especially for SSE streaming). +_HEADERS_TO_STRIP = { + "host", + "content-length", + "connection", + "keep-alive", + "transfer-encoding", + "upgrade", + "proxy-authorization", + "proxy-connection", + "te", + "trailer", + "accept-encoding", + # We decompress incoming bodies and forward as plain bytes, so the + # upstream must not be told the body is still compressed. + "content-encoding", + # Strip incoming Content-Type so we can set a canonical one without + # ending up with a duplicate header (httpx joins case-different keys + # like "content-type" and "Content-Type" into a single comma-separated + # value, which strict backends reject as "Unsupported content type"). + "content-type", +} + + +def _decode_body(raw: bytes, encoding: str) -> bytes: + """Decompress a request body according to its Content-Encoding header. + + Codex CLI sends `Content-Encoding: zstd` (and historically gzip), so we + must decompress before JSON-parsing. Unknown encodings raise; callers turn + that into a 415. + """ + enc = (encoding or "").lower().strip() + if not enc or enc == "identity": + return raw + if enc == "gzip": + return gzip.decompress(raw) + if enc == "deflate": + try: + return zlib.decompress(raw) + except zlib.error: + return zlib.decompress(raw, -zlib.MAX_WBITS) + if enc == "zstd": + if _zstd is None: + raise RuntimeError("zstandard package not installed") + # Codex CLI emits streaming zstd frames without a Content-Size header, + # so plain `decompress(raw)` fails with "could not determine content + # size". The streaming reader handles both framed and sized inputs. + dctx = _zstd.ZstdDecompressor() + with dctx.stream_reader(io.BytesIO(raw)) as reader: + return reader.read() + if enc == "br": + if _brotli is None: + raise RuntimeError("brotli package not installed") + return _brotli.decompress(raw) + raise ValueError(f"Unsupported Content-Encoding: {enc!r}") + + +async def _read_json(request: Request) -> dict: + """Read and JSON-parse a request body, transparently handling compression.""" + raw = await request.body() + encoding = request.headers.get("content-encoding", "") + try: + decoded = _decode_body(raw, encoding) + except Exception as e: + logger.warning(f"Failed to decode {encoding!r} body: {e}") + raise HTTPException(status_code=415, detail=f"Cannot decode body: {e}") + if not decoded: + return {} + return json.loads(decoded) + + +def _is_chatgpt_subscription_request(headers: dict) -> bool: + """Detect a Codex CLI request authenticated via ChatGPT subscription. + + Codex always sends `chatgpt-account-id` in this mode and never in API-key + mode, so we use it as the routing signal. + """ + return bool(headers.get("chatgpt-account-id")) + + +def _openai_upstream_base(headers: dict, config: Config) -> str: + """Pick the correct upstream base URL for an OpenAI-format request. + + ChatGPT-subscription Codex traffic must go to the ChatGPT backend, not the + standard OpenAI API; everything else (API-key Codex, ChatCompletions + clients) uses the configured OpenAI URL. + """ + if _is_chatgpt_subscription_request(headers): + return config.upstream.chatgpt_backend_url + return config.upstream.openai_url + + +def _openai_forward_headers(headers: dict, config: Config) -> dict: + """Build upstream headers for OpenAI-compatible requests. + + Forwards all client headers verbatim except hop-by-hop and connection-scoped + ones (see `_HEADERS_TO_STRIP`). Preserving `User-Agent`, `OpenAI-Beta`, + Codex-specific headers, etc. is required because the ChatGPT backend + fingerprints on them and will throttle requests that look unrecognized. + """ + forward_headers: dict = { + k: v for k, v in headers.items() if k.lower() not in _HEADERS_TO_STRIP + } + forward_headers["Content-Type"] = "application/json" + + if not headers.get("authorization"): + configured_key = config.get_api_key("openai") + if configured_key: + forward_headers["Authorization"] = f"Bearer {configured_key}" + + return forward_headers + + +def _openai_upstream_path(headers: dict, path: str) -> str: + """Adapt an incoming OpenAI-style path to the upstream's path scheme. + + The OpenAI public API uses `/v1/...` (e.g., `/v1/responses`). The ChatGPT + backend (`chatgpt.com/backend-api/codex`) drops the `/v1` and serves + `/responses` directly, so a request proxied as-is hits the wrong URL and + chatgpt.com returns 403. Strip the `/v1` prefix when routing to that + backend; leave OpenAI-direct traffic untouched. + """ + if _is_chatgpt_subscription_request(headers): + if path.startswith("/v1/"): + return path[3:] + if path.startswith("v1/"): + return path[2:] + return path + + def create_app( config: Config, metrics: SessionMetrics, @@ -142,17 +335,9 @@ async def health(): @app.post("/v1/chat/completions") async def openai_chat_completions(request: Request): """Intercept OpenAI Chat Completions, compress, and forward.""" - body = await request.json() + body = await _read_json(request) headers = dict(request.headers) - # Resolve API key - auth_header = headers.get("authorization", "") - api_key = "" - if auth_header.startswith("Bearer "): - api_key = auth_header[7:] - if not api_key: - api_key = config.get_api_key("openai") - model = body.get("model", "gpt-4o") messages = body.get("messages", []) is_streaming = body.get("stream", False) @@ -215,11 +400,12 @@ async def openai_chat_completions(request: Request): metrics.record_request(req_metrics, cost_saved) # ── Forward to upstream ── - upstream_url = f"{config.upstream.openai_url}/v1/chat/completions" - forward_headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - } + upstream_url = _with_query( + f"{_openai_upstream_base(headers, config)}" + f"{_openai_upstream_path(headers, '/v1/chat/completions')}", + request, + ) + forward_headers = _openai_forward_headers(headers, config) if is_streaming: return StreamingResponse( @@ -249,17 +435,15 @@ async def openai_chat_completions(request: Request): @app.post("/v1/completions") async def openai_completions(request: Request): """Pass-through for legacy completions (no compression for non-chat).""" - body = await request.json() + body = await _read_json(request) headers = dict(request.headers) - auth_header = headers.get("authorization", "") - api_key = auth_header[7:] if auth_header.startswith("Bearer ") else config.get_api_key("openai") - - upstream_url = f"{config.upstream.openai_url}/v1/completions" - forward_headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - } + upstream_url = _with_query( + f"{_openai_upstream_base(headers, config)}" + f"{_openai_upstream_path(headers, '/v1/completions')}", + request, + ) + forward_headers = _openai_forward_headers(headers, config) is_streaming = body.get("stream", False) if is_streaming: @@ -279,13 +463,15 @@ async def openai_completions(request: Request): async def openai_models(request: Request): """Pass-through for model listing.""" headers = dict(request.headers) - auth_header = headers.get("authorization", "") - api_key = auth_header[7:] if auth_header.startswith("Bearer ") else config.get_api_key("openai") - upstream_url = f"{config.upstream.openai_url}/v1/models" + upstream_url = _with_query( + f"{_openai_upstream_base(headers, config)}" + f"{_openai_upstream_path(headers, '/v1/models')}", + request, + ) response = await http_client.get( upstream_url, - headers={"Authorization": f"Bearer {api_key}"}, + headers=_openai_forward_headers(headers, config), ) return Response( content=response.content, @@ -300,13 +486,9 @@ async def openai_models(request: Request): @app.post("/v1/messages") async def anthropic_messages(request: Request): """Intercept Anthropic Messages API, compress, and forward.""" - body = await request.json() + body = await _read_json(request) headers = dict(request.headers) - # Anthropic uses x-api-key header - api_key = headers.get("x-api-key", "") or config.get_api_key("anthropic") - anthropic_version = headers.get("anthropic-version", "2023-06-01") - model = body.get("model", "claude-sonnet-4-20250514") messages = body.get("messages", []) is_streaming = body.get("stream", False) @@ -417,17 +599,8 @@ async def anthropic_messages(request: Request): metrics.record_request(req_metrics, cost_saved) # ── Forward to upstream ── - upstream_url = f"{config.upstream.anthropic_url}/v1/messages" - forward_headers = { - "Content-Type": "application/json", - "x-api-key": api_key, - "anthropic-version": anthropic_version, - } - - # Forward any additional anthropic headers - for key, value in headers.items(): - if key.startswith("anthropic-") and key != "anthropic-version": - forward_headers[key] = value + upstream_url = _with_query(f"{config.upstream.anthropic_url}/v1/messages", request) + forward_headers = _anthropic_forward_headers(headers, config) cache_headers = { "X-TokenTamer-Cache-Breakpoints": str(cache_info["breakpoints"]), @@ -460,10 +633,47 @@ async def anthropic_messages(request: Request): }, ) + @app.post("/v1/messages/count_tokens") + async def anthropic_count_tokens(request: Request): + """Pass through Anthropic token counting requests. + + Claude Code calls this endpoint when using Anthropic Messages format + gateways. Token counting must see the original request body, so do not + compress or mutate it here. + """ + body = await _read_json(request) + headers = dict(request.headers) + + upstream_url = _with_query( + f"{config.upstream.anthropic_url}/v1/messages/count_tokens", + request, + ) + response = await http_client.post( + upstream_url, + headers=_anthropic_forward_headers(headers, config), + json=body, + ) + return Response( + content=response.content, + status_code=response.status_code, + headers={"Content-Type": response.headers.get("content-type", "application/json")}, + ) + # ────────────────────────────────────────────────────────── # OpenAI Responses API (used by Codex CLI) # ────────────────────────────────────────────────────────── + @app.get("/v1/responses") + async def openai_responses_reject_upgrade(request: Request): + """Reject WebSocket upgrade attempts cleanly. + + Codex CLI tries `wss://.../v1/responses` first, falling back to plain + HTTPS POST on rejection. Without this handler the GET would fall to the + catch-all, which forwards to chatgpt.com — wasting a round-trip and + polluting logs with 403s from the upstream's branded error page. + """ + return Response(status_code=426, content="WebSocket not supported") + @app.post("/v1/responses") async def openai_responses(request: Request): """Intercept OpenAI Responses API (Codex CLI). @@ -472,12 +682,9 @@ async def openai_responses(request: Request): `messages`. We compress textual code blocks in user messages only, leaving tool calls and reasoning items untouched. """ - body = await request.json() + body = await _read_json(request) headers = dict(request.headers) - auth_header = headers.get("authorization", "") - api_key = auth_header[7:] if auth_header.startswith("Bearer ") else config.get_api_key("openai") - model = body.get("model", "gpt-4o") is_streaming = body.get("stream", False) raw_input = body.get("input", "") @@ -485,6 +692,7 @@ async def openai_responses(request: Request): # Normalize input → list of pseudo-messages for the analyzer. # The Responses API accepts either a string or a list of items. synthesized: List[dict] = [] + has_typed_parts = False if isinstance(raw_input, str): synthesized = [{"role": "user", "content": raw_input}] elif isinstance(raw_input, list): @@ -497,10 +705,21 @@ async def openai_responses(request: Request): "tool_call", "tool_result", "reasoning"): continue content = item.get("content", "") + if not isinstance(content, str): + has_typed_parts = True synthesized.append({"role": role, "content": content}) has_tools = _request_has_tools(body) - skip_compression = passthrough or (has_tools and not compress_with_tools) + # Codex sends `content` as a list of typed parts (e.g., + # [{"type": "input_text", "text": "..."}]). Our compressor only knows + # how to handle plain strings; replacing typed parts with a flat string + # breaks the wire shape and chatgpt.com 400s. Skip compression in that + # case until we add typed-part-aware rewriting. + skip_compression = ( + passthrough + or (has_tools and not compress_with_tools) + or has_typed_parts + ) compressed_input = raw_input analysis = None @@ -539,15 +758,12 @@ async def openai_responses(request: Request): compressed_cost = counter.estimate_cost(compressed_tokens, 0, pricing.input, pricing.output) metrics.record_request(req_metrics, original_cost - compressed_cost) - upstream_url = f"{config.upstream.openai_url}/v1/responses" - forward_headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {api_key}", - } - # Forward OpenAI-Beta and other custom headers - for k, v in headers.items(): - if k.lower().startswith("openai-"): - forward_headers[k] = v + upstream_url = _with_query( + f"{_openai_upstream_base(headers, config)}" + f"{_openai_upstream_path(headers, '/v1/responses')}", + request, + ) + forward_headers = _openai_forward_headers(headers, config) if is_streaming: return StreamingResponse( @@ -563,6 +779,11 @@ async def openai_responses(request: Request): response = await http_client.post( upstream_url, headers=forward_headers, json=body, ) + if response.status_code >= 400: + logger.warning( + "Upstream /v1/responses %s from %s — body=%r", + response.status_code, upstream_url, response.content[:1000], + ) return Response( content=response.content, status_code=response.status_code, @@ -578,14 +799,23 @@ async def openai_responses(request: Request): @app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"]) async def catch_all(request: Request, path: str): - """Pass through any unrecognized routes to the OpenAI upstream.""" + """Pass through unrecognized routes to the right OpenAI-side upstream. + + Codex CLI on a ChatGPT subscription occasionally hits auxiliary routes + (e.g. account info) on the ChatGPT backend; route those there based on + the same `chatgpt-account-id` signal we use for /v1/responses. + """ body = await request.body() headers = dict(request.headers) # Remove host header to avoid conflicts headers.pop("host", None) - upstream_url = f"{config.upstream.openai_url}/{path}" + upstream_url = _with_query( + f"{_openai_upstream_base(headers, config)}" + f"{_openai_upstream_path(headers, '/' + path)}", + request, + ) response = await http_client.request( method=request.method, @@ -622,6 +852,15 @@ async def _stream_proxy( json=body, timeout=httpx.Timeout(connect=10.0, read=300.0, write=30.0, pool=10.0), ) as response: + if response.status_code >= 400: + # Read the whole error body so we can log it for diagnosis. + error_body = await response.aread() + logger.warning( + "Upstream stream %s from %s — body=%r", + response.status_code, url, error_body[:1000], + ) + yield error_body + return async for chunk in response.aiter_bytes(): yield chunk except httpx.ReadTimeout: diff --git a/token_tamer/session_cache.py b/token_tamer/session_cache.py index 1c975c4..6f62adb 100644 --- a/token_tamer/session_cache.py +++ b/token_tamer/session_cache.py @@ -109,27 +109,41 @@ def apply(self, body: dict) -> Tuple[dict, Dict[str, Any]]: "prefix_end_index": -1, # messages[0..prefix_end_index] are cached } + existing_breakpoints = self._count_cache_controls(body) + if existing_breakpoints > MAX_BREAKPOINTS: + self._trim_cache_controls(body, MAX_BREAKPOINTS) + existing_breakpoints = MAX_BREAKPOINTS + if len(messages) < MIN_MESSAGES_TO_CACHE: + self._normalize_cache_control_ttls(body) return body, info breakpoint_count = 0 # ── 1. Cache the tools array (most stable thing in the request) ── - if self._mark_tools_for_caching(body): + if ( + existing_breakpoints + breakpoint_count < MAX_BREAKPOINTS + and self._mark_tools_for_caching(body) + ): breakpoint_count += 1 # ── 2. Cache the system prompt ── - if self._mark_system_for_caching(body): + if ( + existing_breakpoints + breakpoint_count < MAX_BREAKPOINTS + and self._mark_system_for_caching(body) + ): breakpoint_count += 1 # ── 3. Cache the conversation prefix (all but trailing turns) ── - remaining_budget = MAX_BREAKPOINTS - breakpoint_count + remaining_budget = MAX_BREAKPOINTS - existing_breakpoints - breakpoint_count if remaining_budget > 0: prefix_idx = self._mark_conversation_prefix(body, messages) if prefix_idx >= 0: breakpoint_count += 1 info["prefix_end_index"] = prefix_idx + self._normalize_cache_control_ttls(body) + # ── 4. Track the session ── session_id = self._fingerprint(messages) turn_count = self._record_session(session_id) @@ -147,6 +161,77 @@ def apply(self, body: dict) -> Tuple[dict, Dict[str, Any]]: return body, info + @staticmethod + def _iter_cache_control_owners(body: dict) -> List[dict]: + """Return dict objects that may own Anthropic cache_control markers. + + Anthropic accepts at most four cache breakpoints per request. Claude + Code may already send some, so TokenTamer must treat existing markers + as consuming the same global budget as markers it injects. + """ + owners: List[dict] = [] + + for tool in body.get("tools") or []: + if isinstance(tool, dict) and "cache_control" in tool: + owners.append(tool) + + system = body.get("system") + if isinstance(system, list): + for block in system: + if isinstance(block, dict) and "cache_control" in block: + owners.append(block) + + for message in body.get("messages") or []: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and "cache_control" in block: + owners.append(block) + + return owners + + @classmethod + def _count_cache_controls(cls, body: dict) -> int: + return len(cls._iter_cache_control_owners(body)) + + @classmethod + def _trim_cache_controls(cls, body: dict, limit: int = MAX_BREAKPOINTS) -> int: + """Remove surplus cache_control markers from the end of the request.""" + owners = cls._iter_cache_control_owners(body) + removed = 0 + for owner in reversed(owners[limit:]): + owner.pop("cache_control", None) + removed += 1 + return removed + + @classmethod + def _normalize_cache_control_ttls(cls, body: dict) -> None: + """Keep cache_control TTLs valid in Anthropic processing order. + + Anthropic processes cache breakpoints as tools, system, then messages. + A longer-lived 1h block cannot appear after a shorter/default 5m block. + If Claude Code already sent any 1h cache block, promote earlier cache + blocks to 1h so TokenTamer's injected defaults do not make the request + invalid. + """ + owners = cls._iter_cache_control_owners(body) + last_one_hour = -1 + for index, owner in enumerate(owners): + cache_control = owner.get("cache_control") + if isinstance(cache_control, dict) and cache_control.get("ttl") == "1h": + last_one_hour = index + + if last_one_hour < 0: + return + + for owner in owners[: last_one_hour + 1]: + cache_control = owner.get("cache_control") + if not isinstance(cache_control, dict): + continue + cache_control["ttl"] = "1h" + def reset(self) -> None: """Forget all tracked sessions. Useful for tests.""" with self._lock: