diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index 6c2d11db..10bcbd7f 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -42,6 +42,10 @@ from ..payment.price import sats_usd_price from ..payment.usage import NormalizedUsage, normalize_usage from ..wallet import recieve_token, send_token from . import messages_dispatch +from .cache_breakpoints import ( + inject_anthropic_cache_breakpoints, + is_explicit_cache_model, +) from .count_tokens import count_tokens_locally from .litellm_routing import detect_litellm_prefix @@ -445,6 +449,19 @@ class BaseUpstreamProvider: return body + def _upstream_accepts_cache_control(self) -> bool: + """True when this upstream accepts explicit ``cache_control`` markers. + + Only OpenRouter (documents Anthropic + Alibaba explicit caching) and the + native Anthropic API accept the markers. Stamping them toward an + automatic-cache or non-supporting upstream risks a 400, so injection is + confined to these. Base URL is also checked so an OpenRouter endpoint + configured through the generic provider is still recognised. + """ + if self.provider_type in ("openrouter", "anthropic"): + return True + return "openrouter.ai" in (self.base_url or "") + def prepare_request_body( self, body: bytes | None, model_obj: Model ) -> bytes | None: @@ -514,6 +531,27 @@ class BaseUpstreamProvider: data["stream_options"] = merged changed = True + # Explicit-cache models (Anthropic Claude, Alibaba Qwen / deepseek-v3.2) + # cache nothing without ``cache_control`` markers in the body. Clients + # that don't recognise a routstr URL as one of these never send them, so + # caching silently never engages over routstr even though it works + # against OpenRouter directly. Stamp the standard breakpoints so caching + # works by default, deferring to any client-set markers. Gated to + # upstreams that accept the markers (OpenRouter / Anthropic) so they + # never leak to an automatic-cache provider that would reject them. + if ( + "messages" in data + and isinstance(data.get("messages"), list) + and self._upstream_accepts_cache_control() + and is_explicit_cache_model( + model_obj.id, + model_obj.forwarded_model_id, + model_obj.canonical_slug, + ) + ): + if inject_anthropic_cache_breakpoints(data): + changed = True + if changed: return json.dumps(data).encode() return body diff --git a/routstr/upstream/cache_breakpoints.py b/routstr/upstream/cache_breakpoints.py new file mode 100644 index 00000000..6979819b --- /dev/null +++ b/routstr/upstream/cache_breakpoints.py @@ -0,0 +1,157 @@ +"""Inject explicit prompt-cache breakpoints into OpenAI-shaped requests. + +Some upstreams cache *explicitly*: the request must carry +``cache_control: {"type": "ephemeral"}`` markers on the content blocks that +should be cached. Two model families use this identical wire format: + +* **Anthropic Claude** — direct or via OpenRouter's ``anthropic/*`` models. +* **Alibaba's explicit-cache models on OpenRouter** — ``qwen/qwen3-max``, + ``qwen/qwen-plus``, ``qwen/qwen3.6-plus``, ``qwen/qwen3-coder-plus``, + ``qwen/qwen3-coder-flash`` and ``deepseek/deepseek-v3.2`` — which OpenRouter + documents as using "the same syntax as Anthropic explicit caching". + +Every other provider routstr proxies (OpenAI, Azure, xAI/Grok, Groq, Moonshot, +default DeepSeek, Gemini implicit, Fireworks) caches *automatically* and needs +no markers — they are left untouched. + +A client that doesn't know it is talking to one of these models *through* +routstr (e.g. an OpenAI-compatible coding agent pointed at a routstr URL) never +emits the markers — it only adds them when it recognises the provider as +OpenRouter. So caching silently never engages over routstr even though the same +client caches fine talking to OpenRouter directly. + +This module restores caching by stamping the standard breakpoints onto the +forwarded body — the system prompt, the last tool, and the last conversation +message (the format allows up to four; we use three, matching the common +agent convention) — but only when the client supplied none of its own, so +explicit client control always wins. The caller is responsible for only +applying this toward an upstream that accepts the markers (OpenRouter / +Anthropic), so they never leak to a provider that would reject them. +""" + +from __future__ import annotations + +from typing import Any + +# The single ephemeral marker stamped onto each chosen breakpoint. A 5-minute +# TTL (the default for ``ephemeral``) — deliberately not the 1h tier, which +# carries a higher cache-write premium and should stay opt-in. +EPHEMERAL_CACHE_CONTROL: dict[str, str] = {"type": "ephemeral"} + +# Alibaba's explicit-cache models on OpenRouter. Matched as substrings of the +# model id (any spelling routstr carries). Snapshot endpoints that OpenRouter +# documents as *not* supporting explicit caching (e.g. ``qwen3.5-plus-02-15``) +# are different families and deliberately absent from this list. +_ALIBABA_EXPLICIT_CACHE_SLUGS: tuple[str, ...] = ( + "qwen3-max", + "qwen-plus", + "qwen3.6-plus", + "qwen3-coder-plus", + "qwen3-coder-flash", + "deepseek-v3.2", +) + + +def is_explicit_cache_model(model_id: str | None, *fallbacks: str | None) -> bool: + """True when the target model uses the explicit ``cache_control`` dialect. + + Covers the Claude family (broadly — every Claude model supports it) and + Alibaba's documented explicit-cache models, across the id spellings routstr + carries: the OpenRouter id (``anthropic/claude-...``, ``qwen/qwen3-max``), + the bare upstream id, and any forwarded/canonical alias. + """ + for candidate in (model_id, *fallbacks): + if not candidate: + continue + lowered = candidate.lower() + if "claude" in lowered or "anthropic/" in lowered: + return True + if any(slug in lowered for slug in _ALIBABA_EXPLICIT_CACHE_SLUGS): + return True + return False + + +def _has_cache_control(obj: Any) -> bool: + """Recursively detect any client-supplied ``cache_control`` marker.""" + if isinstance(obj, dict): + if "cache_control" in obj: + return True + return any(_has_cache_control(v) for v in obj.values()) + if isinstance(obj, list): + return any(_has_cache_control(v) for v in obj) + return False + + +def body_has_cache_control(data: dict) -> bool: + """True when the request already carries cache_control on messages/tools.""" + return _has_cache_control(data.get("messages")) or _has_cache_control( + data.get("tools") + ) + + +def _stamp_text_content(message: dict) -> bool: + """Add the ephemeral marker to a message's last text block. + + A string content is promoted to the array form Anthropic requires for + cache markers; an existing array gets the marker on its last text part. + Returns True when a marker was placed. + """ + content = message.get("content") + if isinstance(content, str): + if not content: + return False + message["content"] = [ + { + "type": "text", + "text": content, + "cache_control": dict(EPHEMERAL_CACHE_CONTROL), + } + ] + return True + if isinstance(content, list): + for part in reversed(content): + if isinstance(part, dict) and part.get("type") == "text": + part["cache_control"] = dict(EPHEMERAL_CACHE_CONTROL) + return True + return False + + +def _stamp_system_prompt(messages: list) -> None: + for message in messages: + if isinstance(message, dict) and message.get("role") in ( + "system", + "developer", + ): + _stamp_text_content(message) + return + + +def _stamp_last_tool(tools: Any) -> None: + if isinstance(tools, list) and tools and isinstance(tools[-1], dict): + tools[-1]["cache_control"] = dict(EPHEMERAL_CACHE_CONTROL) + + +def _stamp_last_conversation_message(messages: list) -> None: + for message in reversed(messages): + if isinstance(message, dict) and message.get("role") in ("user", "assistant"): + if _stamp_text_content(message): + return + + +def inject_anthropic_cache_breakpoints(data: dict) -> bool: + """Stamp ephemeral cache breakpoints onto an OpenAI-shaped chat body. + + Mutates ``data`` in place (the established convention in + ``prepare_request_body``) and returns True when anything changed. No-ops + when the body isn't chat-shaped or the client already set cache_control. + """ + messages = data.get("messages") + if not isinstance(messages, list) or not messages: + return False + if body_has_cache_control(data): + return False + + _stamp_system_prompt(messages) + _stamp_last_tool(data.get("tools")) + _stamp_last_conversation_message(messages) + return True diff --git a/tests/unit/test_cache_breakpoints.py b/tests/unit/test_cache_breakpoints.py new file mode 100644 index 00000000..641623b1 --- /dev/null +++ b/tests/unit/test_cache_breakpoints.py @@ -0,0 +1,189 @@ +"""Tests for Anthropic cache-breakpoint injection on forwarded requests. + +Anthropic prompt caching is explicit; a client that doesn't recognise a routstr +URL as Anthropic-backed never sends ``cache_control`` markers, so caching never +engages over routstr. ``prepare_request_body`` must stamp the standard +breakpoints for Anthropic-family models while always deferring to client-set +markers and never touching automatic-cache providers. +""" + +import json +import os + +import pytest + +os.environ.setdefault("UPSTREAM_BASE_URL", "http://test") +os.environ.setdefault("UPSTREAM_API_KEY", "test") +os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to") + +from routstr.upstream import GenericUpstreamProvider +from routstr.upstream.cache_breakpoints import ( + body_has_cache_control, + inject_anthropic_cache_breakpoints, + is_explicit_cache_model, +) + + +def _chat_body() -> dict: + return { + "model": "anthropic/claude-sonnet-4.5", + "stream": True, + "messages": [ + {"role": "system", "content": "You are concise."}, + {"role": "user", "content": "Hello"}, + ], + "tools": [ + {"type": "function", "function": {"name": "a"}}, + {"type": "function", "function": {"name": "b"}}, + ], + } + + +# --------------------------------------------------------------------------- +# Model detection +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "model_id,expected", + [ + ("anthropic/claude-sonnet-4.5", True), + ("claude-haiku-4-5-20251001", True), + # Alibaba explicit-cache models share Anthropic's wire format + ("qwen/qwen3-max", True), + ("qwen/qwen3-coder-plus", True), + ("deepseek/deepseek-v3.2", True), + # Automatic-cache providers need no markers + ("openai/gpt-4o", False), + ("google/gemini-2.5-flash", False), + ("deepseek/deepseek-chat", False), + ("qwen/qwen3.5-plus-02-15", False), # snapshot, no explicit caching + (None, False), + ], +) +def test_is_explicit_cache_model(model_id: str | None, expected: bool) -> None: + assert is_explicit_cache_model(model_id) is expected + + +def test_is_explicit_cache_model_uses_fallbacks() -> None: + # routstr id is opaque but a forwarded/canonical alias reveals the family. + assert is_explicit_cache_model("model-xyz", None, "anthropic/claude-opus-4.1") + + +# --------------------------------------------------------------------------- +# Breakpoint placement +# --------------------------------------------------------------------------- + + +def test_injects_three_breakpoints() -> None: + data = _chat_body() + assert inject_anthropic_cache_breakpoints(data) is True + + # system prompt promoted to array form with a marker + system = data["messages"][0]["content"] + assert system == [ + { + "type": "text", + "text": "You are concise.", + "cache_control": {"type": "ephemeral"}, + } + ] + # last tool marked + assert data["tools"][-1]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in data["tools"][0] + # last user message marked + user = data["messages"][1]["content"] + assert user[-1]["cache_control"] == {"type": "ephemeral"} + + +def test_defers_to_client_supplied_cache_control() -> None: + data = _chat_body() + data["messages"][1]["content"] = [ + {"type": "text", "text": "Hello", "cache_control": {"type": "ephemeral"}} + ] + assert body_has_cache_control(data) is True + # No additional stamping when the client already controls caching. + assert inject_anthropic_cache_breakpoints(data) is False + assert "cache_control" not in data["tools"][-1] + + +def test_marks_last_text_part_of_array_content() -> None: + data = _chat_body() + data["messages"][1]["content"] = [ + {"type": "text", "text": "first"}, + {"type": "image_url", "image_url": {"url": "x"}}, + {"type": "text", "text": "last"}, + ] + inject_anthropic_cache_breakpoints(data) + parts = data["messages"][1]["content"] + assert parts[2]["cache_control"] == {"type": "ephemeral"} + assert "cache_control" not in parts[0] + + +def test_noop_without_messages() -> None: + assert inject_anthropic_cache_breakpoints({"prompt": "x"}) is False + + +# --------------------------------------------------------------------------- +# prepare_request_body integration +# --------------------------------------------------------------------------- + + +def _model(model_id: str): # type: ignore[no-untyped-def] + from routstr.payment.models import Architecture, Model, Pricing + + return Model( + id=model_id, + name=model_id, + created=0, + description="", + context_length=200000, + architecture=Architecture( + modality="text->text", + input_modalities=["text"], + output_modalities=["text"], + tokenizer="Claude", + instruct_type=None, + ), + pricing=Pricing(prompt=0.0, completion=0.0), + ) + + +def _openrouter_provider() -> "GenericUpstreamProvider": + # OpenRouter endpoint via the generic provider — recognised by base URL. + return GenericUpstreamProvider(base_url="https://openrouter.ai/api/v1") + + +@pytest.mark.parametrize( + "model_id", ["anthropic/claude-sonnet-4.5", "qwen/qwen3-max", "deepseek/deepseek-v3.2"] +) +def test_prepare_request_body_injects_for_explicit_models(model_id: str) -> None: + provider = _openrouter_provider() + body = json.dumps(_chat_body()).encode() + out = provider.prepare_request_body(body, _model(model_id)) + assert out is not None + data = json.loads(out) + assert body_has_cache_control(data) is True + assert data["tools"][-1]["cache_control"] == {"type": "ephemeral"} + + +def test_prepare_request_body_skips_for_automatic_provider_model() -> None: + provider = _openrouter_provider() + body = json.dumps(_chat_body()).encode() + out = provider.prepare_request_body(body, _model("openai/gpt-4o")) + assert out is not None + data = json.loads(out) + assert body_has_cache_control(data) is False + + +def test_prepare_request_body_skips_when_upstream_rejects_markers() -> None: + # Claude id but a non-OpenRouter/Anthropic upstream → must NOT inject, + # since the markers could be rejected by an upstream that doesn't accept them. + from routstr.upstream import GenericUpstreamProvider + + provider = GenericUpstreamProvider(base_url="https://some-gateway.example/v1") + body = json.dumps(_chat_body()).encode() + out = provider.prepare_request_body(body, _model("anthropic/claude-sonnet-4.5")) + assert out is not None + data = json.loads(out) + assert body_has_cache_control(data) is False