diff --git a/tests/integration/_support/claude_code.py b/tests/integration/_support/claude_code.py index 7bb05ce941e..9d80832f776 100644 --- a/tests/integration/_support/claude_code.py +++ b/tests/integration/_support/claude_code.py @@ -31,6 +31,8 @@ CLAUDE_CODE_REASONING_BETAS: Final = ( "interleaved-thinking-2025-05-14", "thinking-token-count-2026-05-13", ) +CLAUDE_CODE_CACHING_BETAS: Final = ("extended-cache-ttl-2025-04-11", "prompt-caching-scope-2026-01-05") +CACHING_CLI_BETA: Final = f"{CLI_BETA},extended-cache-ttl-2025-04-11" REASONING_FIELDS: Final = ("thinking", "output_config", "reasoning_effort", "temperature") _OUTPUT_USAGE_KEYS: Final = frozenset({"output_tokens", "output_tokens_details"}) METADATA_USER_ID: Final = json.dumps( @@ -731,6 +733,114 @@ def forwarded(sent: Mapping[str, JsonValue], request: Request) -> Forwarded: ) +@dataclass(frozen=True, slots=True) +class CacheForwarded: + breakpoints: dict[str, JsonValue] + other_changes: dict[str, JsonValue] + caching_betas: tuple[str, ...] + + +def _marked(prefix: str, blocks: JsonValue) -> tuple[tuple[str, JsonValue], ...]: + if not isinstance(blocks, list): + return () + return tuple( + (f"{prefix}[{index}]", block["cache_control"]) + for index, block in enumerate(blocks) + if isinstance(block, dict) and "cache_control" in block + ) + + +def _marked_message(index: int, message: JsonValue) -> tuple[tuple[str, JsonValue], ...]: + return _marked(f"messages[{index}].content", message.get("content")) if isinstance(message, dict) else () + + +def _marked_messages(messages: JsonValue) -> tuple[tuple[str, JsonValue], ...]: + if not isinstance(messages, list): + return () + return tuple(chain.from_iterable(_marked_message(index, message) for index, message in enumerate(messages))) + + +def cache_breakpoints(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + top_level: Final = (("cache_control", body["cache_control"]),) if "cache_control" in body else () + return dict( + ( + *top_level, + *_marked("system", body.get("system")), + *_marked("tools", body.get("tools")), + *_marked_messages(body.get("messages")), + ) + ) + + +def _unmarked_block(block: JsonValue) -> JsonValue: + if not isinstance(block, dict): + return block + return {key: value for key, value in block.items() if key != "cache_control"} + + +def _unmarked_blocks(blocks: JsonValue) -> JsonValue: + return [_unmarked_block(block) for block in blocks] if isinstance(blocks, list) else blocks + + +def _unmarked_message(message: JsonValue) -> JsonValue: + if not isinstance(message, dict) or "content" not in message: + return message + return {**message, "content": _unmarked_blocks(message["content"])} + + +def _unmarked_value(key: str, value: JsonValue) -> JsonValue: + match key, value: + case ("system" | "tools", _): + return _unmarked_blocks(value) + case ("messages", list()): + return [_unmarked_message(message) for message in value] + case _: + return value + + +def with_block_breakpoint( + body: Mapping[str, JsonValue], field: str, index: int, control: Mapping[str, JsonValue] +) -> dict[str, JsonValue]: + blocks: Final = body[field] + assert isinstance(blocks, list), body + return { + **body, + field: [ + {**block, "cache_control": dict(control)} if position == index and isinstance(block, dict) else block + for position, block in enumerate(blocks) + ], + } + + +def without_cache_breakpoints(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {key: _unmarked_value(key, value) for key, value in body.items() if key != "cache_control"} + + +def caching_betas(anthropic_beta: str) -> tuple[str, ...]: + return tuple(sorted(beta for beta in anthropic_beta.split(",") if beta in CLAUDE_CODE_CACHING_BETAS)) + + +def _without_model(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {key: value for key, value in without_cache_breakpoints(body).items() if key != "model"} + + +def cache_forwarded(sent: Mapping[str, JsonValue], request: Request) -> CacheForwarded: + body: Final = JSON_OBJECT.validate_json(request.body) + return CacheForwarded( + breakpoints=cache_breakpoints(body), + other_changes=body_diff(_without_model(sent), _without_model(body)), + caching_betas=caching_betas(request.headers.get("anthropic-beta", "")), + ) + + +def streamed_start_usage(stream: str) -> JsonValue: + return next( + JSON_OBJECT.validate_python(data["message"]).get("usage") + for event, data in _client_events(stream) + if event == "message_start" + ) + + def _appended(block: Mapping[str, JsonValue], key: str, text: str) -> dict[str, JsonValue]: return {**block, key: f"{block.get(key) or ''}{text}"} diff --git a/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_breakpoint_injection_wire.py b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_breakpoint_injection_wire.py new file mode 100644 index 00000000000..bca8e67b68c --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_breakpoint_injection_wire.py @@ -0,0 +1,175 @@ +import uuid +from collections.abc import Callable, Mapping +from typing import Final + +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL: Final = "claude-sonnet-4-5" +_FIVE_MINUTES: Final = {"type": "ephemeral"} +_SUBAGENT_BILLING: Final = ( + "x-anthropic-billing-header: cc_version=2.1.283.00; cc_entrypoint=sdk-cli; cc_is_subagent=true;" +) +_MAIN_AGENT_BILLING: Final = "x-anthropic-billing-header: cc_version=2.1.283.00; cc_entrypoint=sdk-cli;" + + +def _unmarked_claude_code_turn() -> dict[str, JsonValue]: + return cc.without_cache_breakpoints({**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": False}) + + +def _one_shot_turn(billing: str) -> dict[str, JsonValue]: + return { + "system": [{"type": "text", "text": billing}], + "messages": [{"role": "user", "content": [{"type": "text", "text": f"Summarize {uuid.uuid4().hex}"}]}], + "max_tokens": 512, + "stream": False, + } + + +def _forward( + gateway: Gateway, + sent: Mapping[str, JsonValue], + key_fields: Mapping[str, JsonValue], + deployment: Mapping[str, JsonValue], +) -> tuple[cc.CacheForwarded, dict[str, JsonValue], dict[str, JsonValue]]: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + return Reply( + body=cc.message_reply( + identity, _MODEL, ({"type": "text", "text": "PONG"},), {"input_tokens": 12, "output_tokens": 4} + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY, **deployment + ) + key: Final = scenario.key(**key_fields) + response: Final = gateway.request( + "POST", "/v1/messages", {**sent, "model": model}, key=key, headers=cc.cli_headers(key, cc.CACHING_CLI_BETA) + ) + assert response.status_code == 200, response.text + upstream: Final = wire.drain() + rows: Final = eventually( + lambda: read_rows( + "SELECT model_id, metadata->>'litellm_gateway_injected_cache' AS injected_for " + 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert len(upstream) == 1, upstream + return cc.cache_forwarded(sent, upstream[0]), cc.JSON_OBJECT.validate_json(upstream[0].body), rows[0] + + +def test_prompt_caching_key_adds_breakpoints_to_the_system_prompt_and_last_message(gateway: Gateway) -> None: + sent: Final = _unmarked_claude_code_turn() + forwarded, _, spend_row = _forward(gateway, sent, {"enable_prompt_caching": True}, {}) + assert cc.cache_breakpoints(sent) == {} + assert forwarded.breakpoints == {"system[2]": _FIVE_MINUTES, "messages[0].content[7]": _FIVE_MINUTES} + assert forwarded.other_changes == {} + assert forwarded.caching_betas == ("extended-cache-ttl-2025-04-11", "prompt-caching-scope-2026-01-05") + assert spend_row["injected_for"] == spend_row["model_id"], spend_row + + +def test_prompt_caching_key_turns_a_string_system_prompt_into_a_cached_block(gateway: Gateway) -> None: + sent: Final = {**_unmarked_claude_code_turn(), "system": "Synthetic agent identity system prompt."} + _, body, _ = _forward(gateway, sent, {"enable_prompt_caching": True}, {}) + assert body["system"] == [ + {"type": "text", "text": "Synthetic agent identity system prompt.", "cache_control": _FIVE_MINUTES} + ] + + +def test_key_without_prompt_caching_adds_no_breakpoints(gateway: Gateway) -> None: + sent: Final = _unmarked_claude_code_turn() + forwarded, _, spend_row = _forward(gateway, sent, {}, {}) + assert forwarded.breakpoints == {} + assert forwarded.other_changes == {} + assert spend_row["injected_for"] is None, spend_row + + +def test_prompt_caching_key_adds_nothing_when_the_client_marked_any_breakpoint(gateway: Gateway) -> None: + sent: Final = cc.with_block_breakpoint(_unmarked_claude_code_turn(), "tools", 23, _FIVE_MINUTES) + forwarded, _, spend_row = _forward(gateway, sent, {"enable_prompt_caching": True}, {}) + assert forwarded.breakpoints == {"tools[23]": _FIVE_MINUTES} + assert forwarded.other_changes == {} + assert spend_row["injected_for"] is None, spend_row + + +@pytest.mark.parametrize( + ("billing", "received"), + ( + pytest.param(_SUBAGENT_BILLING, {}, id="subagent-gets-no-breakpoints"), + pytest.param( + _MAIN_AGENT_BILLING, + {"system[0]": _FIVE_MINUTES, "messages[0].content[0]": _FIVE_MINUTES}, + id="main-agent-gets-breakpoints", + ), + ), +) +def test_prompt_caching_key_skips_one_shot_claude_code_subagent_requests( + gateway: Gateway, billing: str, received: dict[str, JsonValue] +) -> None: + sent: Final = _one_shot_turn(billing) + forwarded, _, _ = _forward(gateway, sent, {"enable_prompt_caching": True}, {}) + assert forwarded.breakpoints == received + assert forwarded.other_changes == {} + + +def _three_client_breakpoints() -> dict[str, JsonValue]: + system_marked: Final = cc.with_block_breakpoint(_unmarked_claude_code_turn(), "system", 1, _FIVE_MINUTES) + return {**cc.with_block_breakpoint(system_marked, "tools", 23, _FIVE_MINUTES), "cache_control": dict(_FIVE_MINUTES)} + + +def _four_client_breakpoints() -> dict[str, JsonValue]: + return cc.with_block_breakpoint(_three_client_breakpoints(), "system", 2, _FIVE_MINUTES) + + +@pytest.mark.parametrize( + ("sent", "received"), + ( + pytest.param( + _three_client_breakpoints, + { + "cache_control": _FIVE_MINUTES, + "system[1]": _FIVE_MINUTES, + "tools[23]": _FIVE_MINUTES, + "messages[0].content[7]": _FIVE_MINUTES, + }, + id="three-client-breakpoints-plus-last-message", + ), + pytest.param( + _four_client_breakpoints, + { + "cache_control": _FIVE_MINUTES, + "system[1]": _FIVE_MINUTES, + "system[2]": _FIVE_MINUTES, + "tools[23]": _FIVE_MINUTES, + }, + id="four-client-breakpoints-leave-no-room", + ), + ), +) +def test_configured_injection_points_stop_at_anthropics_four_breakpoint_limit( + gateway: Gateway, sent: Callable[[], dict[str, JsonValue]], received: dict[str, JsonValue] +) -> None: + client_body: Final = sent() + forwarded, _, _ = _forward( + gateway, + client_body, + {}, + { + "cache_control_injection_points": [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1}, + ] + }, + ) + assert forwarded.breakpoints == received + assert forwarded.other_changes == {} diff --git a/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_breakpoint_passthrough_wire.py b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_breakpoint_passthrough_wire.py new file mode 100644 index 00000000000..144d209073c --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_breakpoint_passthrough_wire.py @@ -0,0 +1,86 @@ +import uuid +from collections.abc import Callable +from typing import Final + +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL: Final = "claude-sonnet-4-5" +_FIVE_MINUTES: Final = {"type": "ephemeral"} +_ONE_HOUR: Final = {"type": "ephemeral", "ttl": "1h"} +_ONE_HOUR_GLOBAL: Final = {"type": "ephemeral", "ttl": "1h", "scope": "global"} +_CLAUDE_CODE_BREAKPOINTS: Final = { + "system[1]": _FIVE_MINUTES, + "system[2]": _FIVE_MINUTES, + "messages[0].content[7]": _FIVE_MINUTES, +} + + +def _claude_code_default(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return body + + +def _one_hour_global_system_prompt(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return cc.with_block_breakpoint(body, "system", 2, _ONE_HOUR_GLOBAL) + + +def _one_hour_tool_breakpoint(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return cc.with_block_breakpoint(body, "tools", 23, _ONE_HOUR) + + +def _top_level_automatic_caching(body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {**body, "cache_control": dict(_FIVE_MINUTES)} + + +@pytest.mark.parametrize( + ("mark", "received"), + ( + pytest.param(_claude_code_default, _CLAUDE_CODE_BREAKPOINTS, id="claude-code-default"), + pytest.param( + _one_hour_global_system_prompt, + {**_CLAUDE_CODE_BREAKPOINTS, "system[2]": _ONE_HOUR_GLOBAL}, + id="1h-global-scope-on-system-prompt", + ), + pytest.param( + _one_hour_tool_breakpoint, + {**_CLAUDE_CODE_BREAKPOINTS, "tools[23]": _ONE_HOUR}, + id="1h-on-last-tool", + ), + pytest.param( + _top_level_automatic_caching, + {"cache_control": _FIVE_MINUTES, **_CLAUDE_CODE_BREAKPOINTS}, + id="top-level-automatic-caching", + ), + ), +) +def test_client_cache_breakpoints_reach_anthropic_unchanged( + gateway: Gateway, mark: Callable[[dict[str, JsonValue]], dict[str, JsonValue]], received: dict[str, JsonValue] +) -> None: + sent: Final = mark({**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": False}) + + def respond(request: Request) -> Reply: + return Reply( + body=cc.message_reply( + f"msg_{uuid.uuid4().hex}", + _MODEL, + ({"type": "text", "text": "PONG"},), + {"input_tokens": 12, "output_tokens": 4}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request( + "POST", "/v1/messages", {**sent, "model": model}, headers=cc.cli_headers(gateway.key, cc.CACHING_CLI_BETA) + ) + assert response.status_code == 200, response.text + upstream: Final = wire.drain() + assert len(upstream) == 1, upstream + forwarded: Final = cc.cache_forwarded(sent, upstream[0]) + assert cc.cache_breakpoints(sent) == received + assert forwarded.breakpoints == received + assert forwarded.other_changes == {} + assert forwarded.caching_betas == ("extended-cache-ttl-2025-04-11", "prompt-caching-scope-2026-01-05") diff --git a/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_token_pricing_wire.py b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_token_pricing_wire.py new file mode 100644 index 00000000000..ccbe7491b4d --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_token_pricing_wire.py @@ -0,0 +1,80 @@ +import uuid +from typing import Final + +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "claude-sonnet-4-5" +_INPUT_RATE: Final = 1e-6 +_OUTPUT_RATE: Final = 2e-6 +_CACHE_READ_RATE: Final = 1e-7 +_CACHE_WRITE_5M_RATE: Final = 1.25e-6 +_CACHE_WRITE_1H_RATE: Final = 2e-6 +_CONTENT: Final = ({"type": "text", "text": "PONG"},) +_USAGE: Final = { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, +} + + +def _reply(identity: str, stream: bool) -> Reply: + if stream: + return Reply(chunks=cc.message_stream(identity, _MODEL, _CONTENT, _USAGE), content_type="text/event-stream") + return Reply(body=cc.message_reply(identity, _MODEL, _CONTENT, _USAGE)) + + +@pytest.mark.parametrize("stream", (pytest.param(False, id="non-streamed"), pytest.param(True, id="streamed"))) +def test_cache_reads_and_5m_and_1h_cache_writes_are_each_billed_at_their_own_rate( + gateway: Gateway, stream: bool +) -> None: + identity: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + return _reply(identity, stream) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MODEL}", + api_base=wire.url, + api_key=cc.ANTHROPIC_API_KEY, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_5M_RATE, + cache_creation_input_token_cost_above_1hr=_CACHE_WRITE_1H_RATE, + ) + body: Final = {**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": stream, "model": model} + response: Final = gateway.request("POST", "/v1/messages", body) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + "SELECT spend, prompt_tokens, " + "(metadata->'cost_breakdown'->>'cache_read_cost')::float AS cache_read_cost, " + "(metadata->'cost_breakdown'->>'cache_creation_cost')::float AS cache_creation_cost " + 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + input_tokens, output_tokens, cache_read, cache_write_5m, cache_write_1h = 100, 50, 400, 100, 200 + row: Final = rows[0] + assert row["prompt_tokens"] == input_tokens + cache_read + cache_write_5m + cache_write_1h, row + assert row["cache_read_cost"] == pytest.approx(cache_read * _CACHE_READ_RATE), row + assert row["cache_creation_cost"] == pytest.approx( + cache_write_5m * _CACHE_WRITE_5M_RATE + cache_write_1h * _CACHE_WRITE_1H_RATE + ), row + assert float(row["spend"]) == pytest.approx( + input_tokens * _INPUT_RATE + + output_tokens * _OUTPUT_RATE + + cache_read * _CACHE_READ_RATE + + cache_write_5m * _CACHE_WRITE_5M_RATE + + cache_write_1h * _CACHE_WRITE_1H_RATE + ), row diff --git a/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_usage_response_wire.py b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_usage_response_wire.py new file mode 100644 index 00000000000..afd950c6af4 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/caching/test_anthropic_cache_usage_response_wire.py @@ -0,0 +1,64 @@ +import uuid +from typing import Final + +from integration._support import claude_code as cc +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL: Final = "claude-sonnet-4-5" +_ANTHROPIC_USAGE: Final = { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, +} + + +def _client_body(stream: bool) -> dict[str, JsonValue]: + return {**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": stream} + + +def test_streamed_cache_read_and_write_tokens_reach_the_client_unchanged(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + return Reply( + chunks=cc.message_stream( + f"msg_{uuid.uuid4().hex}", _MODEL, ({"type": "text", "text": "PONG"},), _ANTHROPIC_USAGE + ), + content_type="text/event-stream", + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request("POST", "/v1/messages", {**_client_body(stream=True), "model": model}) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert cc.streamed_start_usage(response.text) == { + "input_tokens": 100, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + }, response.text + + +def test_non_streamed_cache_read_and_write_tokens_reach_the_client_unchanged(gateway: Gateway) -> None: + def respond(request: Request) -> Reply: + return Reply( + body=cc.message_reply( + f"msg_{uuid.uuid4().hex}", _MODEL, ({"type": "text", "text": "PONG"},), _ANTHROPIC_USAGE + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request("POST", "/v1/messages", {**_client_body(stream=False), "model": model}) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + assert response.json()["usage"] == { + "input_tokens": 100, + "output_tokens": 50, + "cache_read_input_tokens": 400, + "cache_creation_input_tokens": 300, + "cache_creation": {"ephemeral_5m_input_tokens": 100, "ephemeral_1h_input_tokens": 200}, + }, response.text