diff --git a/tests/integration/messages_endpoint/_claude_code.py b/tests/integration/messages_endpoint/_claude_code.py new file mode 100644 index 00000000000..5688a7e9776 --- /dev/null +++ b/tests/integration/messages_endpoint/_claude_code.py @@ -0,0 +1,511 @@ +"""Shared Claude Code-shaped request builders and upstream stream fixtures for the /v1/messages contracts.""" + +import json +from collections.abc import Mapping +from typing import Final + +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +ANTHROPIC_API_KEY: Final = "synthetic-anthropic-key" +OPENAI_API_KEY: Final = "synthetic-openai-key" +OPENAI_BACKEND: Final = "gpt-5.4-mini" +SONNET: Final = "claude-sonnet-4-5" +FABLE: Final = "claude-fable-5-1" +OPUS: Final = "claude-opus-5-5" +CLI_BETA: Final = ( + "claude-code-20250219,interleaved-thinking-2025-05-14,thinking-token-count-2026-05-13," + "context-management-2025-06-27,prompt-caching-scope-2026-01-05" +) +FRONTIER_CLI_BETA: Final = ( + f"{CLI_BETA},mid-conversation-system-2026-04-07,per-turn-control-2026-07-01," + "mid-conversation-tool-changes-2026-07-01,effort-2025-11-24" +) +CACHE: Final = {"type": "ephemeral"} +CONTEXT_MANAGEMENT: Final = {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]} +THINKING_BUDGET: Final = {"budget_tokens": 31999, "type": "enabled", "display": "omitted"} +THINKING_ADAPTIVE: Final = {"type": "adaptive", "display": "omitted"} +METADATA_USER_ID: Final = json.dumps( + { + "device_id": "0" * 64, + "account_uuid": "", + "session_id": "00000000-0000-4000-8000-000000000000", + } +) + + +def schema(properties: JsonValue, required: tuple[str, ...]) -> dict[str, JsonValue]: + return { + "$schema": "https://json-schema.org/draft/2020-12/schema", + "type": "object", + "properties": properties, + "required": list(required), + "additionalProperties": False, + } + + +def field(description: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"description": description, **extra} + + +def tools() -> tuple[dict[str, JsonValue], ...]: + _MAX: Final = 9007199254740991 + return ( + { + "name": "Bash", + "description": "Executes a given bash command and returns its output.", + "input_schema": schema( + { + "command": field("The command to execute", type="string"), + "timeout": field("Optional timeout in milliseconds (max 600000)", type="number"), + "description": field( + "Clear, concise description of what this command does in active voice.", type="string" + ), + "run_in_background": field("Set to true to run this command in the background.", type="boolean"), + "dangerouslyDisableSandbox": field( + "Set this to true to dangerously override sandbox mode and run commands without sandboxing.", + type="boolean", + ), + }, + ("command",), + ), + }, + { + "name": "Read", + "description": "Reads a file from the local filesystem.", + "input_schema": schema( + { + "file_path": field("The absolute path to the file to read", type="string"), + "offset": field("The line number to start reading from.", type="integer", minimum=0, maximum=_MAX), + "limit": field("The number of lines to read.", type="integer", exclusiveMinimum=0, maximum=_MAX), + "pages": field('Page range for PDF files (e.g., "1-5", "3", "10-20").', type="string"), + }, + ("file_path",), + ), + }, + { + "name": "Edit", + "description": "Performs exact string replacements in files.", + "input_schema": schema( + { + "file_path": field("The absolute path to the file to modify", type="string"), + "old_string": field("The text to replace", type="string"), + "new_string": field( + "The text to replace it with (must be different from old_string)", type="string" + ), + "replace_all": field( + "Replace all occurrences of old_string (default false)", default=False, type="boolean" + ), + }, + ("file_path", "old_string", "new_string"), + ), + }, + { + "name": "Agent", + "description": "Launch a new agent to handle complex, multi-step tasks.", + "input_schema": schema( + { + "description": field("A short (3-5 word) description of the task", type="string"), + "prompt": field("The task for the agent to perform", type="string"), + "subagent_type": field("The type of specialized agent to use for this task", type="string"), + "model": field( + "Optional model override for this agent.", + type="string", + enum=["sonnet", "opus", "haiku", "fable"], + ), + "run_in_background": field( + "Agents run in the background by default; you will be notified when one completes.", + type="boolean", + ), + "isolation": field("Isolation mode.", type="string", enum=["worktree", "remote"]), + }, + ("description", "prompt"), + ), + }, + ) + + +def system_blocks() -> tuple[dict[str, JsonValue], ...]: + return ( + {"type": "text", "text": "x-anthropic-billing-header: cc_version=2.1.283.00; cc_entrypoint=sdk-cli;"}, + {"type": "text", "text": "Synthetic agent identity system prompt.", "cache_control": CACHE}, + {"type": "text", "text": "Synthetic interactive agent instructions.", "cache_control": CACHE}, + ) + + +def claude_code_request(cache_bust: str) -> dict[str, JsonValue]: + reminders: Final = ( + f"\n{cache_bust}\n", + "\nSynthetic model identity reminder.\n", + "\nSynthetic agent types reminder.\n", + "\nSynthetic skills reminder.\n", + "\n15000000 tokens left\n", + "\nSynthetic date reminder.\n", + "\nSynthetic attribution reminder.\n", + ) + return { + "model": "", + "system": list(system_blocks()), + "messages": [ + { + "role": "user", + "content": [ + *[{"type": "text", "text": reminder} for reminder in reminders], + {"type": "text", "text": "Reply with exactly the word PONG", "cache_control": CACHE}, + ], + } + ], + "tools": list(tools()), + "metadata": {"user_id": METADATA_USER_ID}, + "max_tokens": 32000, + "thinking": dict(THINKING_BUDGET), + "context_management": dict(CONTEXT_MANAGEMENT), + "stream": True, + } + + +def frontier_request( + cache_bust: str, + effort: str, + max_tokens: int, + prompt_text: str = "Reply with exactly the word PONG", + stream: bool = True, +) -> dict[str, JsonValue]: + return { + "model": "", + "system": list(system_blocks()), + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"\n{cache_bust}\n"}, + {"type": "text", "text": prompt_text}, + ], + }, + { + "role": "system", + "content": [ + { + "type": "text", + "text": "# Environment\nSynthetic environment block.", + "cache_control": CACHE, + } + ], + }, + ], + "tools": list(tools()), + "metadata": {"user_id": METADATA_USER_ID}, + "max_tokens": max_tokens, + "thinking": dict(THINKING_ADAPTIVE), + "context_management": dict(CONTEXT_MANAGEMENT), + "output_config": {"effort": effort}, + "stream": stream, + } + + +def tool_loop_turn2( + base: dict[str, JsonValue], + assistant_content: tuple[dict[str, JsonValue], ...], + tool_results: tuple[tuple[str, JsonValue], ...], +) -> dict[str, JsonValue]: + return { + **base, + "messages": [ + *base["messages"], + {"role": "assistant", "content": list(assistant_content)}, + { + "role": "user", + "content": [ + {"tool_use_id": tool_use_id, "type": "tool_result", "content": content} + for tool_use_id, content in tool_results + ], + }, + { + "role": "system", + "content": [ + { + "type": "text", + "text": "14999970 tokens left", + "cache_control": CACHE, + }, + { + "type": "text", + "text": "First privately list what you need next; then request every item that doesn't depend on another's result in this one response.", + }, + ], + }, + ], + } + + +def cli_headers(key: str, beta: str = CLI_BETA) -> dict[str, str]: + return { + "accept": "application/json", + "content-type": "application/json", + "user-agent": "claude-cli/2.1.283 (external, sdk-cli)", + "x-claude-code-session-id": "00000000-0000-4000-8000-000000000000", + "x-stainless-arch": "x64", + "x-stainless-lang": "js", + "x-stainless-os": "Linux", + "x-stainless-package-version": "0.112.1", + "x-stainless-retry-count": "0", + "x-stainless-runtime": "node", + "x-stainless-runtime-version": "v26.3.0", + "x-stainless-timeout": "600", + "anthropic-beta": beta, + "anthropic-dangerous-direct-browser-access": "true", + "anthropic-version": "2023-06-01", + "x-app": "cli", + "x-api-key": key, + } + + +def sse_frame(event: str, data: JsonValue) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + +def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]: + frames: Final = tuple(frame for frame in text.split("\n\n") if frame.strip()) + return tuple( + ( + next(line.removeprefix("event: ") for line in frame.splitlines() if line.startswith("event: ")), + json.loads(next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: "))), + ) + for frame in frames + ) + + +def _start_usage(usage: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {key: value for key, value in usage.items() if key != "output_tokens"} + + +def text_stream(identity: str, model: str, text: str, usage: dict[str, int]) -> tuple[bytes, ...]: + return ( + sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": _start_usage(usage), + }, + }, + ), + sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ), + sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": usage["output_tokens"]}, + }, + ), + sse_frame("message_stop", {"type": "message_stop"}), + ) + + +def tool_use_stream( + identity: str, + model: str, + thinking: str, + signature: str, + tool_calls: tuple[tuple[str, str, JsonValue], ...], + usage: dict[str, int], +) -> tuple[bytes, ...]: + frames: list[bytes] = [ + sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": _start_usage(usage), + }, + }, + ), + sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + ), + sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": thinking}}, + ), + sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": signature}}, + ), + sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + ] + for index, (tool_id, name, tool_input) in enumerate(tool_calls, start=1): + arguments: Final = json.dumps(tool_input) + frames += [ + sse_frame( + "content_block_start", + { + "type": "content_block_start", + "index": index, + "content_block": {"type": "tool_use", "id": tool_id, "name": name, "input": {}}, + }, + ), + sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "input_json_delta", "partial_json": arguments[: len(arguments) // 2]}, + }, + ), + sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": index, + "delta": {"type": "input_json_delta", "partial_json": arguments[len(arguments) // 2 :]}, + }, + ), + sse_frame("content_block_stop", {"type": "content_block_stop", "index": index}), + ] + frames += [ + sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": usage["output_tokens"]}, + }, + ), + sse_frame("message_stop", {"type": "message_stop"}), + ] + return tuple(frames) + + +def responses_completed( + identity: str, + model: str, + output_items: tuple[dict[str, JsonValue], ...], + usage: dict[str, int], + status: str = "completed", + incomplete_details: JsonValue = None, +) -> bytes: + return json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": status, + "incomplete_details": incomplete_details, + "model": model, + "output": list(output_items), + "usage": usage, + } + ).encode() + + +def responses_stream(identity: str, model: str, output_items: tuple[dict[str, JsonValue], ...]) -> tuple[bytes, ...]: + frames: list[bytes] = [ + sse_frame( + "response.created", + { + "type": "response.created", + "response": { + "id": f"resp_{identity}", + "object": "response", + "status": "in_progress", + "model": model, + "output": [], + }, + }, + ) + ] + for index, item in enumerate(output_items): + item_id: Final = str(item.get("id", f"item_{index}")) + frames.append( + sse_frame( + "response.output_item.added", + { + "type": "response.output_item.added", + "output_index": index, + "item": {**item, "content": []} if item.get("type") == "message" else item, + }, + ) + ) + if item.get("type") == "message": + text: Final = "".join(part.get("text", "") for part in item.get("content", ()) if isinstance(part, dict)) + frames.append( + sse_frame( + "response.output_text.delta", + {"type": "response.output_text.delta", "output_index": index, "item_id": item_id, "delta": text}, + ) + ) + if item.get("type") == "reasoning": + summary_text: Final = "".join( + str(part.get("text", "")) for part in item.get("summary", ()) if isinstance(part, dict) + ) + if summary_text: + frames.append( + sse_frame( + "response.reasoning_summary_text.delta", + { + "type": "response.reasoning_summary_text.delta", + "output_index": index, + "item_id": item_id, + "delta": summary_text, + }, + ) + ) + if item.get("type") == "function_call": + frames.append( + sse_frame( + "response.function_call_arguments.delta", + { + "type": "response.function_call_arguments.delta", + "output_index": index, + "item_id": item_id, + "delta": item.get("arguments", ""), + }, + ) + ) + frames.append( + sse_frame( + "response.output_item.done", + {"type": "response.output_item.done", "output_index": index, "item": item}, + ) + ) + frames.append( + sse_frame( + "response.completed", + { + "type": "response.completed", + "response": { + "id": f"resp_{identity}", + "object": "response", + "status": "completed", + "model": model, + "output": list(output_items), + "usage": {"input_tokens": 41, "output_tokens": 5, "total_tokens": 46}, + }, + }, + ) + ) + return tuple(frames) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_client_disconnect_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_client_disconnect_wire.py new file mode 100644 index 00000000000..0e7f00943c2 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_client_disconnect_wire.py @@ -0,0 +1,85 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def test_client_disconnect_mid_stream_still_bills_the_message(gateway: Gateway) -> None: + identity: Final = f"msg_dc_{uuid.uuid4().hex}" + request_body: Final = cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}") + + def respond(request: Request) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": cc.SONNET, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 1}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "PONG"}}, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 4}, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ), + pause_between_chunks=3.0, + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.SONNET}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + with gateway.client.stream( + "POST", + "/v1/messages", + params={"beta": "true"}, + json={**request_body, "model": model}, + headers={ + **cc.cli_headers(gateway.key), + "authorization": f"Bearer {gateway.key}", + }, + ) as response: + assert response.status_code == 200, response.status_code + first: Final = next(response.iter_text()) + assert "message_start" in first, first + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + return_last_on_timeout=True, + ) + assert rows and rows[0]["prompt_tokens"] == 12, rows + assert wire.disconnected.empty(), ( + "closing the client stream must not abort the upstream call before it finishes; " + f"wire recorded a disconnect on {wire.disconnected.get_nowait()}" + ) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_compaction_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_compaction_wire.py new file mode 100644 index 00000000000..a1e903ab5c6 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_compaction_wire.py @@ -0,0 +1,108 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_CONTEXT_MANAGEMENT: Final = { + "edits": [ + {"type": "clear_thinking_20251015", "keep": "all"}, + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}, + ] +} +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": ""} + + +def _compaction_stream(identity: str) -> tuple[bytes, ...]: + return ( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": cc.FABLE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 1}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": dict(_COMPACTION_BLOCK)}, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + "context_management": { + "applied_edits": [{"type": "compact_20260112", "compacted_at": "2026-09-26T00:00:00Z"}] + }, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ) + + +def test_compaction_edit_and_applied_edit_block_round_trip_through_anthropic(gateway: Gateway) -> None: + identity: Final = f"msg_cm_{uuid.uuid4().hex}" + request_body: Final = { + **cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000), + "context_management": _CONTEXT_MANAGEMENT, + } + turn3: Final = cc.tool_loop_turn2( + request_body, + ( + dict(_COMPACTION_BLOCK), + {"type": "text", "text": "continuing after compaction"}, + ), + (), + ) + seen: list[dict[str, JsonValue]] = [] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen.append(body) + expected: Final = {**(turn3 if len(seen) == 2 else request_body), "model": cc.FABLE} + assert body == expected, { + key: (expected.get(key), body.get(key)) + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + assert body["context_management"] == _CONTEXT_MANAGEMENT + if len(seen) == 1: + return Reply(content_type="text/event-stream", chunks=_compaction_stream(identity)) + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream("msg_cm_next", cc.FABLE, "OK", {"input_tokens": 20, "output_tokens": 2}), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response1: Final = gateway.request( + "POST", "/v1/messages", {**request_body, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response1.status_code == 200, response1.text + events: Final = cc.sse_events(response1.text) + assert events[1][1]["content_block"] == _COMPACTION_BLOCK, events[1] + deltas: Final = [data for event, data in events if event == "message_delta"] + assert len(deltas) == 1 and deltas[0].get("context_management") == { + "applied_edits": [{"type": "compact_20260112", "compacted_at": "2026-09-26T00:00:00Z"}] + }, events + response2: Final = gateway.request( + "POST", "/v1/messages", {**turn3, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + assert len(wire.drain()) == 2 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_count_tokens_wire.py new file mode 100644 index 00000000000..3a4bd61fc8e --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_count_tokens_wire.py @@ -0,0 +1,50 @@ +import uuid +from typing import Final + +import pytest + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def test_count_tokens_forwards_to_anthropic_and_bills_nothing(gateway: Gateway) -> None: + pytest.skip( + "BUG: /v1/messages/count_tokens on an anthropic deployment runs the internal token_counter " + "and never forwards to the provider" + ) + request_body: Final = { + key: value + for key, value in cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}").items() + if key not in ("stream", "max_tokens", "thinking", "output_config") + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages/count_tokens", request.target + body = cc.JSON_OBJECT.validate_json(request.body) + assert body == {**request_body, "model": cc.FABLE}, body + return Reply(body=b'{"input_tokens": 37}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + assert cc.JSON_OBJECT.validate_json(response.content) == {"input_tokens": 37} + assert len(wire.drain()) == 1 + call_id: Final = response.headers.get("x-litellm-call-id", "") + assert call_id, dict(response.headers) + leftover: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=20, + return_last_on_timeout=True, + ) + assert all(float(row["spend"]) == 0 for row in leftover), leftover diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_document_input_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_document_input_wire.py new file mode 100644 index 00000000000..7fa67549372 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_document_input_wire.py @@ -0,0 +1,132 @@ +import base64 +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_PDF_BYTES: Final = ( + b"%PDF-1.1\n" + b"1 0 obj<>endobj\n" + b"2 0 obj<>endobj\n" + b"3 0 obj<>endobj\n" + b"trailer<>\n%%EOF" +) +_DOC_BLOCK: Final = { + "type": "document", + "source": {"type": "base64", "data": base64.b64encode(_PDF_BYTES).decode(), "media_type": "application/pdf"}, + "citations": {"enabled": True}, +} + + +def _cited_stream(identity: str) -> tuple[bytes, ...]: + return ( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": cc.FABLE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 1}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "A page."}}, + ), + cc.sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "citations_delta", + "citation": { + "type": "page_location", + "document_index": 0, + "document_title": "dot.pdf", + "start_page_number": 1, + "end_page_number": 1, + "cited_text": "Page", + }, + }, + }, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 6}, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ) + + +def test_base64_pdf_document_with_citations_reaches_anthropic_identical(gateway: Gateway) -> None: + request_body: Final = cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000) + request_body["messages"] = [ + { + "role": "user", + "content": [ + dict(_DOC_BLOCK), + {"type": "text", "text": f"What is on page one? {uuid.uuid4().hex}"}, + ], + } + ] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = {**request_body, "model": cc.FABLE} + assert body == expected, { + key: (expected.get(key), body.get(key)) + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + return Reply(content_type="text/event-stream", chunks=_cited_stream(f"msg_doc_{uuid.uuid4().hex}")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + events: Final = cc.sse_events(response.text) + citations: Final = [ + data["delta"] for event, data in events if data.get("delta", {}).get("type") == "citations_delta" + ] + assert citations == [ + { + "type": "citations_delta", + "citation": { + "type": "page_location", + "document_index": 0, + "document_title": "dot.pdf", + "start_page_number": 1, + "end_page_number": 1, + "cited_text": "Page", + }, + } + ], citations + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_fallback_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_fallback_wire.py new file mode 100644 index 00000000000..0075acb6da8 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_fallback_wire.py @@ -0,0 +1,92 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def _error_529() -> Reply: + return Reply( + status=529, + body=json.dumps({"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}).encode(), + ) + + +def test_anthropic_overloaded_primary_falls_back_to_second_deployment(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = f"msg_fb_{uuid.uuid4().hex}" + request_body: Final = cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}") + + def respond_fallback(request: Request) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity, cc.SONNET, "PONG", {"input_tokens": 12, "output_tokens": 4}), + ) + + with ( + wire_server(lambda request: _error_529()) as primary, + wire_server(respond_fallback) as fallback, + ): + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + { + "model_name": "cc-primary", + "litellm_params": { + "model": f"anthropic/{cc.SONNET}", + "api_key": cc.ANTHROPIC_API_KEY, + "api_base": primary.url, + "model_info": {"id": "primary-cc"}, + }, + }, + { + "model_name": "cc-fallback-group", + "litellm_params": { + "model": f"anthropic/{cc.SONNET}", + "api_key": cc.ANTHROPIC_API_KEY, + "api_base": fallback.url, + "model_info": {"id": "fallback-cc"}, + }, + }, + ] + config["router_settings"] = { + "num_retries": 0, + "disable_cooldowns": True, + "fallbacks": [{"cc-primary": ["cc-fallback-group"]}], + } + path: Final = tmp_path / "fallbacks.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy( + gateway, tmp_path, {"REDIS_HOST": "127.0.0.1", "REDIS_PORT": "6379"}, config=path + ) as candidate: + with candidate.client.stream( + "POST", + "/v1/messages", + params={"beta": "true"}, + json={**request_body, "model": "cc-primary"}, + headers={**cc.cli_headers(candidate.key), "authorization": f"Bearer {candidate.key}"}, + ) as response: + assert response.status_code == 200, response.status_code + body: Final = "".join(response.iter_text()) + assert "message_stop" in body, body + assert "PONG" in body, body + deployments: Final = candidate.get("/model/info")["data"] + fallback_id: Final = next( + entry["model_info"]["id"] + for entry in deployments + if entry["litellm_params"]["api_base"] == fallback.url + ) + assert response.headers.get("x-litellm-model-id") == fallback_id, dict(response.headers) + assert len(primary.drain()) == 1 + assert len(fallback.drain()) == 1 + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda values: len(values) == 1, + seconds=70, + ) + assert len(rows) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_frontier_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_frontier_wire.py new file mode 100644 index 00000000000..a9a41b96751 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_frontier_wire.py @@ -0,0 +1,111 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + + +def _diff(expected: dict[str, JsonValue], body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + key: {"expected": expected.get(key), "upstream": body.get(key)} + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + + +def test_claude_code_adaptive_thinking_and_effort_reach_anthropic_intact(gateway: Gateway) -> None: + identity: Final = f"msg_fable_{uuid.uuid4().hex}" + request_body: Final = cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000) + cli_beta: Final = frozenset(cc.FRONTIER_CLI_BETA.split(",")) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == cc.ANTHROPIC_API_KEY + assert request.headers["anthropic-version"] == "2023-06-01" + upstream_beta: Final = request.headers.get("anthropic-beta", "") + assert cli_beta <= frozenset(upstream_beta.split(",")), upstream_beta + assert upstream_beta.split(",").count("effort-2025-11-24") == 1, upstream_beta + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = {**request_body, "model": cc.FABLE} + assert body == expected, _diff(expected, body) + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity, cc.FABLE, "PONG", {"input_tokens": 12, "output_tokens": 4}), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + events: Final = cc.sse_events(response.text) + assert [event for event, _ in events][-1] == "message_stop" + assert events[4][1]["delta"]["stop_reason"] == "end_turn" + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == 12 and rows[0]["completion_tokens"] == 4 + + +def test_claude_code_xhigh_effort_reaches_anthropic_and_charges_by_usage(gateway: Gateway) -> None: + identity: Final = f"msg_opus_{uuid.uuid4().hex}" + request_body: Final = cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "xhigh", 128000) + cli_beta: Final = frozenset(cc.FRONTIER_CLI_BETA.split(",")) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == cc.ANTHROPIC_API_KEY + upstream_beta: Final = request.headers.get("anthropic-beta", "") + assert cli_beta <= frozenset(upstream_beta.split(",")), upstream_beta + assert upstream_beta.split(",").count("effort-2025-11-24") == 1, upstream_beta + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = {**request_body, "model": cc.OPUS} + assert body == expected, _diff(expected, body) + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity, cc.OPUS, "PONG", {"input_tokens": 10, "output_tokens": 5}), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{cc.OPUS}", + api_base=wire.url, + api_key=cc.ANTHROPIC_API_KEY, + input_cost_per_token=1e-6, + output_cost_per_token=5e-6, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(rows[0]["spend"]) == pytest.approx(10 * 1e-6 + 5 * 5e-6) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_image_input_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_image_input_wire.py new file mode 100644 index 00000000000..0deed7db56a --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_image_input_wire.py @@ -0,0 +1,69 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_PNG_B64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAQAAAAECAIAAAAmkwkpAAAAEElEQVR4nGP4z8AARwzEcQCukw/x0F8jngAAAABJRU5ErkJggg==" +_IMAGE_BLOCK: Final = { + "type": "image", + "source": {"type": "base64", "data": _PNG_B64, "media_type": "image/png"}, +} + + +def test_tool_result_image_block_and_pasted_image_reach_anthropic_identical(gateway: Gateway) -> None: + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/dot.png and say what colour it is", + ) + with_image_result: Final = cc.tool_loop_turn2( + turn1, + ({"type": "tool_use", "id": "toolu_img", "name": "Read", "input": {"file_path": "/tmp/cc_probe/dot.png"}},), + (("toolu_img", [dict(_IMAGE_BLOCK)]),), + ) + pasted: Final = cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000) + pasted["messages"] = [ + { + "role": "user", + "content": [ + dict(_IMAGE_BLOCK), + {"type": "text", "text": f"What colour is this? {uuid.uuid4().hex}"}, + ], + } + ] + seen: list[dict[str, JsonValue]] = [] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen.append(body) + expected: Final = {**(pasted if len(seen) == 2 else with_image_result), "model": cc.FABLE} + assert body == expected, { + key: (expected.get(key), body.get(key)) + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream( + f"msg_img_{uuid.uuid4().hex}", cc.FABLE, "RED", {"input_tokens": 20, "output_tokens": 2} + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response1: Final = gateway.request( + "POST", "/v1/messages", {**with_image_result, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response1.status_code == 200, response1.text + response2: Final = gateway.request( + "POST", "/v1/messages", {**pasted, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + assert len(wire.drain()) == 2 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_interleaved_thinking_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_interleaved_thinking_wire.py new file mode 100644 index 00000000000..c6a46ca2806 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_interleaved_thinking_wire.py @@ -0,0 +1,167 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_TURN1_TOOL: Final = ("toolu_a", "Read", {"file_path": "/tmp/cc_probe/a.txt"}) +_TURN2_TOOL: Final = ("toolu_b", "Read", {"file_path": "/tmp/cc_probe/b.txt"}) + + +def _diff(expected: dict[str, JsonValue], body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + key: {"expected": expected.get(key), "upstream": body.get(key)} + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + + +def _turn2(base: dict[str, JsonValue]) -> dict[str, JsonValue]: + return cc.tool_loop_turn2( + base, + ( + {"type": "thinking", "thinking": "plan", "signature": "sig1"}, + {"type": "tool_use", "id": _TURN1_TOOL[0], "name": _TURN1_TOOL[1], "input": _TURN1_TOOL[2]}, + ), + ((_TURN1_TOOL[0], "ALPHA"),), + ) + + +def _turn3(turn2: dict[str, JsonValue]) -> dict[str, JsonValue]: + return cc.tool_loop_turn2( + turn2, + ( + {"type": "thinking", "thinking": "got A", "signature": "sig2"}, + {"type": "text", "text": "got A"}, + {"type": "tool_use", "id": _TURN2_TOOL[0], "name": _TURN2_TOOL[1], "input": _TURN2_TOOL[2]}, + ), + ((_TURN2_TOOL[0], "BRAVO"),), + ) + + +def _interleaved_stream(identity: str) -> tuple[bytes, ...]: + usage: Final = {"input_tokens": 20, "output_tokens": 12} + return ( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": cc.FABLE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": usage["input_tokens"], "output_tokens": 1}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "got A"}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "sig2"}}, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "got A"}}, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 1}), + cc.sse_frame( + "content_block_start", + { + "type": "content_block_start", + "index": 2, + "content_block": {"type": "tool_use", "id": _TURN2_TOOL[0], "name": _TURN2_TOOL[1], "input": {}}, + }, + ), + cc.sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 2, + "delta": {"type": "input_json_delta", "partial_json": '{"file_path": "/tmp/cc_pr'}, + }, + ), + cc.sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 2, + "delta": {"type": "input_json_delta", "partial_json": 'obe/b.txt"}'}, + }, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 2}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "tool_use", "stop_sequence": None}, + "usage": {"output_tokens": usage["output_tokens"]}, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ) + + +def test_interleaved_thinking_text_and_tool_use_history_reaches_anthropic_identical(gateway: Gateway) -> None: + identity: Final = f"msg_il_{uuid.uuid4().hex}" + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/a.txt then /tmp/cc_probe/b.txt one at a time and reply with both words", + ) + turn2: Final = _turn2(turn1) + turn3: Final = _turn3(turn2) + seen: list[dict[str, JsonValue]] = [] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + upstream_beta: Final = request.headers.get("anthropic-beta", "") + assert upstream_beta.split(",").count("interleaved-thinking-2025-05-14") == 1, upstream_beta + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen.append(body) + expected: Final = {**turn3, "model": cc.FABLE} if len(seen) == 2 else {**turn2, "model": cc.FABLE} + assert body == expected, _diff(expected, body) + if len(seen) == 1: + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream("msg_il_turn2", cc.FABLE, "got A", {"input_tokens": 20, "output_tokens": 4}), + ) + return Reply(content_type="text/event-stream", chunks=_interleaved_stream(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response2: Final = gateway.request( + "POST", "/v1/messages", {**turn2, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + response3: Final = gateway.request( + "POST", "/v1/messages", {**turn3, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response3.status_code == 200, response3.text + events: Final = cc.sse_events(response3.text) + started: Final = [ + (data["index"], data["content_block"]["type"]) for event, data in events if event == "content_block_start" + ] + assert started == [(0, "thinking"), (1, "text"), (2, "tool_use")], started + assert events[-1][0] == "message_stop" + assert len(wire.drain()) == 2 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_long_context_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_long_context_beta_wire.py new file mode 100644 index 00000000000..a198ec8500a --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_long_context_beta_wire.py @@ -0,0 +1,53 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + +_BETA_1M: Final = f"{cc.FRONTIER_CLI_BETA.replace(',effort-2025-11-24', ',context-1m-2025-08-07,effort-2025-11-24')}" + + +def test_1m_context_beta_forwarded_and_tiered_prompt_priced_above_200k(gateway: Gateway) -> None: + identity: Final = f"msg_1m_{uuid.uuid4().hex}" + request_body: Final = cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + upstream_beta: Final = request.headers.get("anthropic-beta", "") + assert upstream_beta.split(",").count("context-1m-2025-08-07") == 1, upstream_beta + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity, cc.FABLE, "PONG", {"input_tokens": 250000, "output_tokens": 100}), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{cc.FABLE}", + api_base=wire.url, + api_key=cc.ANTHROPIC_API_KEY, + input_cost_per_token=1e-6, + input_cost_per_token_above_200k_tokens=2e-6, + output_cost_per_token=5e-6, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, _BETA_1M), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(rows[0]["spend"]) == pytest.approx(250000 * 2e-6 + 100 * 5e-6), dict(rows[0]) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_model_switch_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_model_switch_wire.py new file mode 100644 index 00000000000..ccdd4b5bd06 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_model_switch_wire.py @@ -0,0 +1,86 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + + +def test_claude_code_mid_loop_model_switch_replays_history_byte_identical(gateway: Gateway) -> None: + identity1: Final = f"msg_sw1_{uuid.uuid4().hex}" + identity2: Final = f"msg_sw2_{uuid.uuid4().hex}" + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/hello.txt and reply with its single word", + ) + turn2: Final = cc.tool_loop_turn2( + turn1, + ( + {"type": "thinking", "thinking": "need to read the file", "signature": "sig_anthropic_1"}, + { + "type": "tool_use", + "id": "toolu_read_1", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/hello.txt"}, + }, + ), + (("toolu_read_1", "1\tPROBE\n2\t"),), + ) + seen: list[dict[str, JsonValue]] = [] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen.append(body) + if len(seen) == 1: + assert body == {**turn1, "model": cc.FABLE}, body.get("model") + return Reply( + content_type="text/event-stream", + chunks=cc.tool_use_stream( + identity1, + cc.FABLE, + "need to read the file", + "sig_anthropic_1", + (("toolu_read_1", "Read", {"file_path": "/tmp/cc_probe/hello.txt"}),), + {"input_tokens": 20, "output_tokens": 10}, + ), + ) + expected: Final = {**turn2, "model": cc.OPUS} + assert body == expected, { + key: (expected.get(key), body.get(key)) + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity2, cc.OPUS, "PROBE", {"input_tokens": 30, "output_tokens": 3}), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + fable: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + opus: Final = scenario.model(model=f"anthropic/{cc.OPUS}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response1: Final = gateway.request( + "POST", "/v1/messages", {**turn1, "model": fable}, params={"beta": "true"}, headers=headers + ) + assert response1.status_code == 200, response1.text + response2: Final = gateway.request( + "POST", "/v1/messages", {**turn2, "model": opus}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + assert len(wire.drain()) == 2 + assert seen[1]["model"] == cc.OPUS + rows: Final = eventually( + lambda: read_rows( + 'SELECT model FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity2,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["model"] == f"anthropic/{cc.OPUS}" diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_native_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_native_wire.py index 6d7fd461fa6..6391912d781 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_native_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_native_wire.py @@ -1,265 +1,50 @@ -import json import uuid from typing import Final 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, TypeAdapter +from integration.messages_endpoint import _claude_code as cc -_API_KEY: Final = "synthetic-anthropic-key" -_MODEL: Final = "claude-sonnet-4-5" -_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) -_CLI_BETA: Final = ( - "claude-code-20250219,interleaved-thinking-2025-05-14,thinking-token-count-2026-05-13," - "context-management-2025-06-27,prompt-caching-scope-2026-01-05" -) -_CACHE: Final = {"type": "ephemeral"} - - -def _schema(properties: JsonValue, required: tuple[str, ...]) -> dict[str, JsonValue]: - return { - "$schema": "https://json-schema.org/draft/2020-12/schema", - "type": "object", - "properties": properties, - "required": list(required), - "additionalProperties": False, - } - - -def _field(description: str, **extra: JsonValue) -> dict[str, JsonValue]: - return {"description": description, **extra} - - -def _tools() -> tuple[dict[str, JsonValue], ...]: - _MAX: Final = 9007199254740991 - return ( - { - "name": "Bash", - "description": "Executes a given bash command and returns its output.", - "input_schema": _schema( - { - "command": _field("The command to execute", type="string"), - "timeout": _field("Optional timeout in milliseconds (max 600000)", type="number"), - "description": _field( - "Clear, concise description of what this command does in active voice.", type="string" - ), - "run_in_background": _field("Set to true to run this command in the background.", type="boolean"), - "dangerouslyDisableSandbox": _field( - "Set this to true to dangerously override sandbox mode and run commands without sandboxing.", - type="boolean", - ), - }, - ("command",), - ), - }, - { - "name": "Read", - "description": "Reads a file from the local filesystem.", - "input_schema": _schema( - { - "file_path": _field("The absolute path to the file to read", type="string"), - "offset": _field("The line number to start reading from.", type="integer", minimum=0, maximum=_MAX), - "limit": _field("The number of lines to read.", type="integer", exclusiveMinimum=0, maximum=_MAX), - "pages": _field('Page range for PDF files (e.g., "1-5", "3", "10-20").', type="string"), - }, - ("file_path",), - ), - }, - { - "name": "Edit", - "description": "Performs exact string replacements in files.", - "input_schema": _schema( - { - "file_path": _field("The absolute path to the file to modify", type="string"), - "old_string": _field("The text to replace", type="string"), - "new_string": _field( - "The text to replace it with (must be different from old_string)", type="string" - ), - "replace_all": _field( - "Replace all occurrences of old_string (default false)", default=False, type="boolean" - ), - }, - ("file_path", "old_string", "new_string"), - ), - }, - { - "name": "Agent", - "description": "Launch a new agent to handle complex, multi-step tasks.", - "input_schema": _schema( - { - "description": _field("A short (3-5 word) description of the task", type="string"), - "prompt": _field("The task for the agent to perform", type="string"), - "subagent_type": _field("The type of specialized agent to use for this task", type="string"), - "model": _field( - "Optional model override for this agent.", - type="string", - enum=["sonnet", "opus", "haiku", "fable"], - ), - "run_in_background": _field( - "Agents run in the background by default; you will be notified when one completes.", - type="boolean", - ), - "isolation": _field("Isolation mode.", type="string", enum=["worktree", "remote"]), - }, - ("description", "prompt"), - ), - }, - ) - - -def _claude_code_request(cache_bust: str) -> dict[str, JsonValue]: - reminders: Final = ( - f"\n{cache_bust}\n", - "\nSynthetic model identity reminder.\n", - "\nSynthetic agent types reminder.\n", - "\nSynthetic skills reminder.\n", - "\n15000000 tokens left\n", - "\nSynthetic date reminder.\n", - "\nSynthetic attribution reminder.\n", - ) - return { - "model": "", - "system": [ - {"type": "text", "text": "Synthetic billing header block from a Claude Code request."}, - {"type": "text", "text": "Synthetic agent identity system prompt.", "cache_control": _CACHE}, - {"type": "text", "text": "Synthetic interactive agent instructions.", "cache_control": _CACHE}, - ], - "messages": [ - { - "role": "user", - "content": [ - *[{"type": "text", "text": reminder} for reminder in reminders], - { - "type": "text", - "text": "Reply with exactly the word PONG", - "cache_control": _CACHE, - }, - ], - } - ], - "tools": list(_tools()), - "metadata": { - "user_id": json.dumps( - { - "device_id": "0" * 64, - "account_uuid": "", - "session_id": "00000000-0000-4000-8000-000000000000", - } - ) - }, - "max_tokens": 32000, - "thinking": {"budget_tokens": 31999, "type": "enabled", "display": "omitted"}, - "context_management": {"edits": [{"type": "clear_thinking_20251015", "keep": "all"}]}, - "stream": True, - } - - -def _sse_frame(event: str, data: JsonValue) -> bytes: - return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() - - -def _message_stream(identity: str) -> tuple[bytes, ...]: - return ( - _sse_frame( - "message_start", - { - "type": "message_start", - "message": { - "id": identity, - "type": "message", - "role": "assistant", - "model": _MODEL, - "content": [], - "stop_reason": None, - "stop_sequence": None, - "usage": {"input_tokens": 12, "output_tokens": 1}, - }, - }, - ), - _sse_frame( - "content_block_start", - {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, - ), - _sse_frame( - "content_block_delta", - {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "PONG"}}, - ), - _sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), - _sse_frame( - "message_delta", - { - "type": "message_delta", - "delta": {"stop_reason": "end_turn", "stop_sequence": None}, - "usage": {"output_tokens": 4}, - }, - ), - _sse_frame("message_stop", {"type": "message_stop"}), - ) - - -def sse_events(text: str) -> tuple[tuple[str, dict[str, object]], ...]: - frames: Final = tuple(frame for frame in text.split("\n\n") if frame.strip()) - return tuple( - ( - next(line.removeprefix("event: ") for line in frame.splitlines() if line.startswith("event: ")), - json.loads(next(line.removeprefix("data: ") for line in frame.splitlines() if line.startswith("data: "))), - ) - for frame in frames - ) +_MODEL: Final = cc.SONNET def test_claude_code_streaming_request_reaches_anthropic_intact_and_streams_back(gateway: Gateway) -> None: identity: Final = f"msg_cc_{uuid.uuid4().hex}" - request_body: Final = _claude_code_request(f"cache-bust-{uuid.uuid4().hex}") - cli_beta: Final = frozenset(_CLI_BETA.split(",")) + request_body: Final = cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}") + cli_beta: Final = frozenset(cc.CLI_BETA.split(",")) def respond(request: Request) -> Reply: assert request.method == "POST" assert request.target == "/v1/messages", request.target - assert request.headers["x-api-key"] == _API_KEY + assert request.headers["x-api-key"] == cc.ANTHROPIC_API_KEY assert request.headers["anthropic-version"] == "2023-06-01" upstream_beta: Final = frozenset(request.headers.get("anthropic-beta", "").split(",")) assert cli_beta <= upstream_beta, request.headers.get("anthropic-beta") - body: Final = _JSON_OBJECT.validate_json(request.body) + body: Final = cc.JSON_OBJECT.validate_json(request.body) expected: Final = {**request_body, "model": _MODEL} assert body == expected, { key: (expected.get(key), body.get(key)) for key in expected.keys() | body.keys() if expected.get(key) != body.get(key) } - return Reply(content_type="text/event-stream", chunks=_message_stream(identity)) + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity, _MODEL, "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=_API_KEY) + 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", {**request_body, "model": model}, params={"beta": "true"}, - headers={ - "accept": "application/json", - "content-type": "application/json", - "user-agent": "claude-cli/2.1.283 (external, sdk-cli)", - "x-claude-code-session-id": "00000000-0000-4000-8000-000000000000", - "x-stainless-arch": "x64", - "x-stainless-lang": "js", - "x-stainless-os": "Linux", - "x-stainless-package-version": "0.112.1", - "x-stainless-retry-count": "0", - "x-stainless-runtime": "node", - "x-stainless-runtime-version": "v26.3.0", - "x-stainless-timeout": "600", - "anthropic-beta": _CLI_BETA, - "anthropic-dangerous-direct-browser-access": "true", - "anthropic-version": "2023-06-01", - "x-app": "cli", - "x-api-key": gateway.key, - }, + headers=cc.cli_headers(gateway.key), ) assert response.status_code == 200, response.text assert response.headers["content-type"].startswith("text/event-stream"), dict(response.headers) - events: Final = sse_events(response.text) + events: Final = cc.sse_events(response.text) assert [event for event, _ in events] == [ "message_start", "content_block_start", diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_prompt_cache_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_prompt_cache_wire.py new file mode 100644 index 00000000000..54bbe2a0704 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_prompt_cache_wire.py @@ -0,0 +1,86 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_USAGE: Final = { + "input_tokens": 10, + "cache_read_input_tokens": 3000, + "cache_creation_input_tokens": 200, + "output_tokens": 5, +} + + +def test_claude_code_cached_turn_charges_cache_read_and_creation_rates(gateway: Gateway) -> None: + identity: Final = f"msg_pc_{uuid.uuid4().hex}" + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/hello.txt and reply with its single word", + ) + turn2: Final = cc.tool_loop_turn2( + turn1, + ( + {"type": "thinking", "thinking": "need to read the file", "signature": "sig_anthropic_1"}, + { + "type": "tool_use", + "id": "toolu_read_1", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/hello.txt"}, + }, + ), + (("toolu_read_1", "1\tPROBE\n2\t"),), + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = {**turn2, "model": cc.FABLE} + assert body == expected, { + key: (expected.get(key), body.get(key)) + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + return Reply(content_type="text/event-stream", chunks=cc.text_stream(identity, cc.FABLE, "PROBE", _USAGE)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{cc.FABLE}", + api_base=wire.url, + api_key=cc.ANTHROPIC_API_KEY, + input_cost_per_token=1e-6, + output_cost_per_token=5e-6, + cache_read_input_token_cost=1e-7, + cache_creation_input_token_cost=1.25e-6, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**turn2, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + events: Final = cc.sse_events(response.text) + usage: Final = events[0][1]["message"]["usage"] + assert usage["cache_read_input_tokens"] == 3000, usage + assert usage["cache_creation_input_tokens"] == 200, usage + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(rows[0]["spend"]) == pytest.approx(10 * 1e-6 + 3000 * 1e-7 + 200 * 1.25e-6 + 5 * 5e-6), dict( + rows[0] + ) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_tool_loop_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_tool_loop_wire.py new file mode 100644 index 00000000000..f386a206451 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_tool_loop_wire.py @@ -0,0 +1,165 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_THINKING: Final = "need to read the file" +_SIGNATURE: Final = "sig_probe_1" +_USAGE: Final = {"input_tokens": 20, "output_tokens": 10} + + +def _diff(expected: dict[str, JsonValue], body: dict[str, JsonValue]) -> dict[str, JsonValue]: + return { + key: {"expected": expected.get(key), "upstream": body.get(key)} + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + + +def test_claude_code_tool_loop_round_trips_thinking_tool_use_and_tool_result(gateway: Gateway) -> None: + identity1: Final = f"msg_tl1_{uuid.uuid4().hex}" + identity2: Final = f"msg_tl2_{uuid.uuid4().hex}" + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/hello.txt and reply with its single word", + ) + calls: Final = (("toolu_read_1", "Read", {"file_path": "/tmp/cc_probe/hello.txt"}),) + turn2: Final = cc.tool_loop_turn2( + turn1, + ( + {"type": "thinking", "thinking": _THINKING, "signature": _SIGNATURE}, + { + "type": "tool_use", + "id": "toolu_read_1", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/hello.txt"}, + }, + ), + (("toolu_read_1", "1\tPROBE\n2\t"),), + ) + seen: list[dict[str, JsonValue]] = [] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen.append(body) + expected: Final = {**turn2, "model": cc.FABLE} if len(seen) == 2 else {**turn1, "model": cc.FABLE} + assert body == expected, _diff(expected, body) + if len(seen) == 1: + return Reply( + content_type="text/event-stream", + chunks=cc.tool_use_stream(identity1, cc.FABLE, _THINKING, _SIGNATURE, calls, _USAGE), + ) + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(identity2, cc.FABLE, "PROBE", {"input_tokens": 30, "output_tokens": 3}), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response1: Final = gateway.request( + "POST", "/v1/messages", {**turn1, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response1.status_code == 200, response1.text + events: Final = cc.sse_events(response1.text) + assert [ + ( + event, + data.get("delta", {}).get( + "type", data.get("content_block", {}).get("type", data.get("delta", {}).get("stop_reason")) + ), + ) + for event, data in events + ] == [ + ("message_start", None), + ("content_block_start", "thinking"), + ("content_block_delta", "thinking_delta"), + ("content_block_delta", "signature_delta"), + ("content_block_stop", None), + ("content_block_start", "tool_use"), + ("content_block_delta", "input_json_delta"), + ("content_block_delta", "input_json_delta"), + ("content_block_stop", None), + ("message_delta", "tool_use"), + ("message_stop", None), + ] + assert events[5][1]["content_block"]["id"] == "toolu_read_1" + assert events[5][1]["content_block"]["name"] == "Read" + partial: Final = events[6][1]["delta"]["partial_json"] + events[7][1]["delta"]["partial_json"] + assert partial == '{"file_path": "/tmp/cc_probe/hello.txt"}' + response2: Final = gateway.request( + "POST", "/v1/messages", {**turn2, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + events2: Final = cc.sse_events(response2.text) + assert events2[2][1]["delta"] == {"type": "text_delta", "text": "PROBE"} + assert events2[4][1]["delta"]["stop_reason"] == "end_turn" + assert len(wire.drain()) == 2 + rows: Final = eventually( + lambda: read_rows( + 'SELECT prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity2,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == 30 + + +def test_claude_code_parallel_tool_results_reach_anthropic_in_client_order(gateway: Gateway) -> None: + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/hello.txt and /tmp/cc_probe/world.txt and reply with both words", + ) + turn2: Final = cc.tool_loop_turn2( + turn1, + ( + {"type": "thinking", "thinking": _THINKING, "signature": _SIGNATURE}, + { + "type": "tool_use", + "id": "toolu_read_1", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/hello.txt"}, + }, + { + "type": "tool_use", + "id": "toolu_read_2", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/world.txt"}, + }, + ), + (("toolu_read_2", "1\tPROBE2\n2\t"), ("toolu_read_1", "1\tPROBE\n2\t")), + ) + + def respond(request: Request) -> Reply: + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = {**turn2, "model": cc.FABLE} + assert body == expected, _diff(expected, body) + results: Final = [block for block in body["messages"][3]["content"] if block["type"] == "tool_result"] + assert [block["tool_use_id"] for block in results] == ["toolu_read_2", "toolu_read_1"] + return Reply( + content_type="text/event-stream", + chunks=cc.text_stream(f"msg_mt_{uuid.uuid4().hex}", cc.FABLE, "PROBE PROBE2", _USAGE), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**turn2, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_upstream_errors_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_upstream_errors_wire.py new file mode 100644 index 00000000000..17cfaee9dd2 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_upstream_errors_wire.py @@ -0,0 +1,136 @@ +import json +import uuid +from typing import Final + +import pytest + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def _error_body(error_type: str, message: str) -> bytes: + return json.dumps({"type": "error", "error": {"type": error_type, "message": message}}).encode() + + +def _assert_upstream_error_status_passthrough(gateway: Gateway, status: int, error_type: str) -> None: + request_body: Final = {**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": False} + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body = cc.JSON_OBJECT.validate_json(request.body) + assert body["stream"] is False + return Reply(status=status, body=_error_body(error_type, "Upstream rejected")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{cc.SONNET}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY, num_retries=0 + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == status, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + assert payload["type"] == "error", payload + assert payload["error"]["type"] == error_type, payload + assert len(wire.drain()) == 1 + call_id: Final = response.headers.get("x-litellm-call-id", "") + assert call_id, dict(response.headers) + leftover: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)), + lambda values: len(values) == 1, + seconds=20, + return_last_on_timeout=True, + ) + assert all(float(row["spend"]) == 0 for row in leftover), leftover + + +def test_anthropic_529_overloaded_error_passes_through_with_client_status(gateway: Gateway) -> None: + pytest.skip( + "BUG: upstream 529 overloaded_error is re-raised through exception_type as InternalServerError " + "and reaches the client as 500 api_error" + ) + _assert_upstream_error_status_passthrough(gateway, 529, "overloaded_error") + + +def test_anthropic_429_rate_limit_error_passes_through_with_client_status(gateway: Gateway) -> None: + _assert_upstream_error_status_passthrough(gateway, 429, "rate_limit_error") + + +def test_anthropic_stream_stop_reason_max_tokens_and_refusal_reach_client(gateway: Gateway) -> None: + for stop_reason in ("max_tokens", "refusal"): + _assert_stream_stop_reason_reaches_client(gateway, stop_reason) + + +def _assert_stream_stop_reason_reaches_client(gateway: Gateway, stop_reason: str) -> None: + identity: Final = f"msg_stop_{uuid.uuid4().hex}" + request_body: Final = cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}") + + def respond(request: Request) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": cc.SONNET, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 1}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "PAR"}}, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": stop_reason, "stop_sequence": None}, + "usage": {"output_tokens": 32000}, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{cc.SONNET}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + events: Final = cc.sse_events(response.text) + assert events[2][1]["delta"]["text"] == "PAR" + assert events[4][1]["delta"]["stop_reason"] == stop_reason + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda values: len(values) == 1, + seconds=20, + return_last_on_timeout=True, + ) + assert isinstance(rows, list) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_web_search_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_web_search_wire.py new file mode 100644 index 00000000000..a275b2e3305 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_claude_code_web_search_wire.py @@ -0,0 +1,184 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +WEB_SEARCH_TOOL: Final = { + "name": "WebSearch", + "description": "Search the web. Returns result blocks with titles and URLs.", + "input_schema": cc.schema( + { + "query": cc.field("The search query to use", type="string", minLength=2), + "allowed_domains": cc.field( + "Only include search results from these domains", type="array", items={"type": "string"} + ), + "blocked_domains": cc.field( + "Never include search results from these domains", type="array", items={"type": "string"} + ), + }, + ("query",), + ), +} + + +def _web_search_stream(identity: str) -> tuple[bytes, ...]: + return ( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": cc.FABLE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 1, "server_tool_use": {"web_search_requests": 1}}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}, + }, + ), + cc.sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "input_json_delta", "partial_json": '{"query": "current LiteLLM version"}'}, + }, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "content_block_start", + { + "type": "content_block_start", + "index": 1, + "content_block": { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "title": "litellm releases", + "url": "https://example.com/litellm", + "page_age": None, + "encrypted_content": "enc_ws_1", + } + ], + }, + }, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 1}), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "1.104.0"}}, + ), + cc.sse_frame( + "content_block_delta", + { + "type": "content_block_delta", + "index": 2, + "delta": { + "type": "citations_delta", + "citation": { + "type": "web_search_result_location", + "url": "https://example.com/litellm", + "title": "litellm releases", + "cited_text": "version 1.104.0", + "encrypted_index": "eidx_1", + }, + }, + }, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 2}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 15, "server_tool_use": {"web_search_requests": 1}}, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ) + + +def test_claude_code_web_search_tool_passthrough_and_cited_response(gateway: Gateway) -> None: + identity: Final = f"msg_ws_{uuid.uuid4().hex}" + request_body: Final = { + **cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Use web search to find the current LiteLLM version and answer in one word", + ), + } + request_body["tools"] = [*request_body["tools"], WEB_SEARCH_TOOL] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/v1/messages", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = {**request_body, "model": cc.FABLE} + assert body == expected, { + key: (expected.get(key), body.get(key)) + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + assert body["tools"][-1] == WEB_SEARCH_TOOL + return Reply(content_type="text/event-stream", chunks=_web_search_stream(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{cc.FABLE}", + api_base=wire.url, + api_key=cc.ANTHROPIC_API_KEY, + input_cost_per_token=1e-6, + output_cost_per_token=5e-6, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + events: Final = cc.sse_events(response.text) + started: Final = [ + (data["index"], data["content_block"]["type"]) for event, data in events if event == "content_block_start" + ] + assert started == [(0, "server_tool_use"), (1, "web_search_tool_result"), (2, "text")], started + citations: Final = [ + data["delta"] for event, data in events if data.get("delta", {}).get("type") == "citations_delta" + ] + assert len(citations) == 1 and citations[0]["citation"]["url"] == "https://example.com/litellm", citations + start_usage: Final = events[0][1]["message"]["usage"] + assert start_usage["server_tool_use"]["web_search_requests"] == 1, start_usage + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + token_cost: Final = 20 * 1e-6 + 15 * 5e-6 + assert float(rows[0]["spend"]) >= token_cost, dict(rows[0]) diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_compaction_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_compaction_bridge_wire.py new file mode 100644 index 00000000000..2e1d4d431ff --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_compaction_bridge_wire.py @@ -0,0 +1,107 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def test_compact_edit_maps_to_responses_context_management(gateway: Gateway) -> None: + request_body: Final = { + **cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000), + "context_management": { + "edits": [ + {"type": "clear_thinking_20251015", "keep": "all"}, + {"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}, + ] + }, + "stream": False, + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body.get("context_management") == [{"type": "compaction", "compact_threshold": 150000}], body.get( + "context_management" + ) + return Reply( + body=cc.responses_completed( + "cm", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + }, + ), + {"input_tokens": 41, "output_tokens": 3, "total_tokens": 44}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + assert payload["content"] == [{"type": "text", "text": "OK"}], payload["content"] + assert len(wire.drain()) == 1 + + +def test_compaction_output_item_reaches_client_as_compaction_block(gateway: Gateway) -> None: + pytest.skip( + "BUG: the responses bridge drops compaction output items in translate_response " + "(transformation.py handles only message/reasoning/function_call), so the client loses " + "the compaction block entirely" + ) + request_body: Final = { + **cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000), + "context_management": { + "edits": [{"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 150000}}] + }, + "stream": False, + } + + def respond(request: Request) -> Reply: + return Reply( + body=cc.responses_completed( + "cm", + cc.OPENAI_BACKEND, + ( + {"type": "compaction", "content": ""}, + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "OK", "annotations": []}], + }, + ), + {"input_tokens": 41, "output_tokens": 3, "total_tokens": 44}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + compaction_blocks: Final = [block for block in payload["content"] if block.get("type") == "compaction"] + assert compaction_blocks == [{"type": "compaction", "content": ""}], payload["content"] + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_count_tokens_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_count_tokens_bridge_wire.py new file mode 100644 index 00000000000..b03e06ff867 --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_count_tokens_bridge_wire.py @@ -0,0 +1,39 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def test_count_tokens_on_openai_deployment_returns_token_count(gateway: Gateway) -> None: + request_body: Final = { + key: value + for key, value in cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}").items() + if key not in ("stream", "max_tokens", "thinking", "output_config") + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses/input_tokens", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body["model"] == cc.OPENAI_BACKEND, body + assert body["instructions"] == str(request_body["system"]), body["instructions"] + assert len(body["input"]) == 1 and body["input"][0]["role"] == "user", body["input"] + assert "cache-bust-" in body["input"][0]["content"], body["input"] + assert body["tools"] == request_body["tools"], body["tools"] + return Reply(body=b'{"object": "response.input_tokens", "input_tokens": 37}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + assert payload == {"input_tokens": 37}, payload + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_document_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_document_bridge_wire.py new file mode 100644 index 00000000000..121d007858e --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_document_bridge_wire.py @@ -0,0 +1,73 @@ +import base64 +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + +_PDF_BYTES: Final = ( + b"%PDF-1.1\n" + b"1 0 obj<>endobj\n" + b"2 0 obj<>endobj\n" + b"3 0 obj<>endobj\n" + b"trailer<>\n%%EOF" +) +_PDF_B64: Final = base64.b64encode(_PDF_BYTES).decode() + + +def test_pdf_document_block_maps_to_input_file_on_bridge(gateway: Gateway) -> None: + request_body: Final = {**cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000), "stream": False} + doc_text: Final = f"What is on page one? {uuid.uuid4().hex}" + request_body["messages"] = [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": {"type": "base64", "data": _PDF_B64, "media_type": "application/pdf"}, + "title": "dot.pdf", + }, + {"type": "text", "text": doc_text}, + ], + } + ] + + def respond(request: Request) -> Reply: + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + user_msg: Final = body["input"][0] + assert user_msg["content"][0] == { + "type": "input_file", + "filename": "dot.pdf", + "file_data": f"data:application/pdf;base64,{_PDF_B64}", + }, user_msg + assert user_msg["content"][1] == {"type": "input_text", "text": doc_text}, user_msg + return Reply( + body=cc.responses_completed( + "doc", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Page one.", "annotations": []}], + }, + ), + {"input_tokens": 41, "output_tokens": 3, "total_tokens": 44}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_errors_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_errors_bridge_wire.py new file mode 100644 index 00000000000..32a7ce3bd6e --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_errors_bridge_wire.py @@ -0,0 +1,77 @@ +import json +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc + + +def test_openai_429_error_comes_back_in_anthropic_shape(gateway: Gateway) -> None: + request_body: Final = {**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": False} + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + return Reply( + status=429, + body=json.dumps( + {"error": {"type": "rate_limit_error", "message": "Rate limit reached", "code": "rate_limit_exceeded"}} + ).encode(), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY, num_retries=0 + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 429, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + assert payload["type"] == "error", payload + assert payload["error"]["type"] == "rate_limit_error", payload + assert len(wire.drain()) == 1 + + +def test_incomplete_responses_completion_maps_to_max_tokens(gateway: Gateway) -> None: + request_body: Final = {**cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}"), "stream": False} + + def respond(request: Request) -> Reply: + assert request.target == "/responses", request.target + return Reply( + body=cc.responses_completed( + "inc", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "PAR", "annotations": []}], + }, + ), + {"input_tokens": 41, "output_tokens": 5, "total_tokens": 46}, + status="incomplete", + incomplete_details={"reason": "max_output_tokens"}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + assert payload["stop_reason"] == "max_tokens", payload + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_frontier_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_frontier_bridge_wire.py new file mode 100644 index 00000000000..355004c1f16 --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_frontier_bridge_wire.py @@ -0,0 +1,149 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_INSTRUCTIONS: Final = "\n".join(block["text"] for block in cc.system_blocks()) +_OUTPUT_ITEMS: Final = ( + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "short plan"}], + "encrypted_content": "enc_1", + }, + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "PONG", "annotations": []}], + }, +) + + +def _expected_responses_body(request_body: dict[str, JsonValue], effort: str) -> dict[str, JsonValue]: + user_blocks: Final = request_body["messages"][0]["content"] + expected_tools: Final = tuple( + { + "type": "function", + "name": tool["name"], + "strict": False, + "description": tool["description"], + "parameters": tool["input_schema"], + } + for tool in request_body["tools"] + ) + input_items: Final = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": block["text"]} for block in user_blocks], + } + ] + for message in request_body["messages"][1:]: + input_items.append( + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": block["text"]} + for block in message["content"] + if block.get("type") == "text" + ], + } + ) + return { + "model": cc.OPENAI_BACKEND, + "input": input_items, + "include": ["reasoning.encrypted_content"], + "instructions": _INSTRUCTIONS, + "max_output_tokens": request_body["max_tokens"], + "tools": list(expected_tools), + "reasoning": {"effort": effort}, + "stream": True, + "user": cc.METADATA_USER_ID[:64], + "prompt_cache_key": "00000000-0000-4000-8000-000000000000", + } + + +def _assert_client_events(text: str) -> None: + events: Final = cc.sse_events(text) + assert [event for event, _ in events] == [ + "message_start", + "content_block_start", + "content_block_delta", + "content_block_delta", + "content_block_stop", + "content_block_start", + "content_block_delta", + "content_block_stop", + "message_delta", + "message_stop", + ], [event for event, _ in events] + assert events[1][1]["content_block"]["type"] == "thinking" + assert events[2][1]["delta"] == {"type": "thinking_delta", "thinking": "short plan"} + assert events[3][1]["delta"]["type"] == "signature_delta" + assert events[6][1]["delta"] == {"type": "text_delta", "text": "PONG"} + assert events[8][1]["delta"]["stop_reason"] == "end_turn" + + +def test_claude_code_frontier_body_becomes_reasoning_effort_on_responses_bridge(gateway: Gateway) -> None: + request_body: Final = cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + assert request.headers["authorization"] == f"Bearer {cc.OPENAI_API_KEY}" + body: Final = cc.JSON_OBJECT.validate_json(request.body) + expected: Final = _expected_responses_body(request_body, "high") + assert body == expected, { + key: {"expected": expected.get(key), "upstream": body.get(key)} + for key in expected.keys() | body.keys() + if expected.get(key) != body.get(key) + } + return Reply( + content_type="text/event-stream", chunks=cc.responses_stream("bridge1", cc.OPENAI_BACKEND, _OUTPUT_ITEMS) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + _assert_client_events(response.text) + assert len(wire.drain()) == 1 + + +def test_claude_code_legacy_thinking_budget_maps_to_reasoning_effort_on_bridge(gateway: Gateway) -> None: + request_body: Final = cc.claude_code_request(f"cache-bust-{uuid.uuid4().hex}") + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body["reasoning"] == {"effort": "high"}, body.get("reasoning") + assert body["model"] == cc.OPENAI_BACKEND + return Reply( + content_type="text/event-stream", chunks=cc.responses_stream("bridge2", cc.OPENAI_BACKEND, _OUTPUT_ITEMS) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + _assert_client_events(response.text) + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_image_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_image_bridge_wire.py new file mode 100644 index 00000000000..8abdde6f2cd --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_image_bridge_wire.py @@ -0,0 +1,112 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_PNG_B64: Final = "iVBORw0KGgoAAAANSUhEUgAAAAQAAAAECAIAAAAmkwkpAAAAEElEQVR4nGP4z8AARwzEcQCukw/x0F8jngAAAABJRU5ErkJggg==" +_IMAGE_BLOCK: Final = { + "type": "image", + "source": {"type": "base64", "data": _PNG_B64, "media_type": "image/png"}, +} +_DATA_URL: Final = f"data:image/png;base64,{_PNG_B64}" + + +def _respond_ok(request: Request) -> Reply: + return Reply( + body=cc.responses_completed( + "img", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "RED", "annotations": []}], + }, + ), + {"input_tokens": 41, "output_tokens": 3, "total_tokens": 44}, + ) + ) + + +def test_tool_result_image_maps_to_input_image_on_bridge(gateway: Gateway) -> None: + turn1: Final = { + **cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/dot.png and say what colour it is", + ), + "stream": False, + } + turn2: Final = cc.tool_loop_turn2( + turn1, + ({"type": "tool_use", "id": "call_img", "name": "Read", "input": {"file_path": "/tmp/cc_probe/dot.png"}},), + (("call_img", [dict(_IMAGE_BLOCK)]),), + ) + + def respond(request: Request) -> Reply: + body: Final = cc.JSON_OBJECT.validate_json(request.body) + outputs: Final = [ + item for item in body["input"] if isinstance(item, dict) and item.get("type") == "function_call_output" + ] + assert len(outputs) == 1 and outputs[0]["call_id"] == "call_img", outputs + image_messages: Final = [ + item + for item in body["input"] + if isinstance(item, dict) + and item.get("type") == "message" + and item.get("role") == "user" + and any(isinstance(part, dict) and part.get("type") == "input_image" for part in item.get("content", ())) + ] + assert image_messages, body["input"] + image_parts: Final = [ + part + for part in image_messages[0]["content"] + if isinstance(part, dict) and part.get("type") == "input_image" + ] + assert image_parts[0]["image_url"] == _DATA_URL, image_parts + return _respond_ok(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**turn2, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 + + +def test_pasted_image_maps_to_input_image_on_bridge(gateway: Gateway) -> None: + request_body: Final = {**cc.frontier_request(f"cache-bust-{uuid.uuid4().hex}", "high", 64000), "stream": False} + pasted_text: Final = f"What colour is this? {uuid.uuid4().hex}" + request_body["messages"] = [ + {"role": "user", "content": [dict(_IMAGE_BLOCK), {"type": "text", "text": pasted_text}]} + ] + + def respond(request: Request) -> Reply: + body: Final = cc.JSON_OBJECT.validate_json(request.body) + user_msg: Final = body["input"][0] + assert user_msg["content"][0] == {"type": "input_image", "image_url": _DATA_URL}, user_msg + assert user_msg["content"][1] == {"type": "input_text", "text": pasted_text}, user_msg + return _respond_ok(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_interleaved_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_interleaved_bridge_wire.py new file mode 100644 index 00000000000..4944c399a09 --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_interleaved_bridge_wire.py @@ -0,0 +1,128 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + + +def _expected_input(turn1: dict[str, JsonValue]) -> list[JsonValue]: + return [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": block["text"]} for block in turn1["messages"][0]["content"]], + }, + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": block["text"]} for block in turn1["messages"][1]["content"]], + }, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "plan"}]}, + { + "type": "function_call", + "call_id": "toolu_a", + "name": "Read", + "arguments": '{"file_path": "/tmp/cc_probe/a.txt"}', + }, + {"type": "function_call_output", "call_id": "toolu_a", "output": "ALPHA"}, + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "14999970 tokens left"}, + { + "type": "input_text", + "text": "First privately list what you need next; then request every item that doesn't depend on another's result in this one response.", + }, + ], + }, + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "got A"}]}, + { + "type": "function_call", + "call_id": "toolu_b", + "name": "Read", + "arguments": '{"file_path": "/tmp/cc_probe/b.txt"}', + }, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "got A"}], + }, + {"type": "function_call_output", "call_id": "toolu_b", "output": "BRAVO"}, + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "14999970 tokens left"}, + { + "type": "input_text", + "text": "First privately list what you need next; then request every item that doesn't depend on another's result in this one response.", + }, + ], + }, + ] + + +def test_bridge_replays_interleaved_history_in_order(gateway: Gateway) -> None: + turn1: Final = { + **cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/a.txt then /tmp/cc_probe/b.txt one at a time", + ), + "stream": False, + } + turn2: Final = cc.tool_loop_turn2( + turn1, + ( + {"type": "thinking", "thinking": "plan", "signature": "sig_anthropic_1"}, + {"type": "tool_use", "id": "toolu_a", "name": "Read", "input": {"file_path": "/tmp/cc_probe/a.txt"}}, + ), + (("toolu_a", "ALPHA"),), + ) + turn3: Final = cc.tool_loop_turn2( + turn2, + ( + {"type": "thinking", "thinking": "got A", "signature": "sig_anthropic_2"}, + {"type": "text", "text": "got A"}, + {"type": "tool_use", "id": "toolu_b", "name": "Read", "input": {"file_path": "/tmp/cc_probe/b.txt"}}, + ), + (("toolu_b", "BRAVO"),), + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body["input"] == _expected_input(turn1), body["input"] + return Reply( + body=cc.responses_completed( + "il", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ALPHA BRAVO", "annotations": []}], + }, + ), + {"input_tokens": 50, "output_tokens": 4, "total_tokens": 54}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**turn3, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_model_switch_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_model_switch_bridge_wire.py new file mode 100644 index 00000000000..67f9a41300d --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_model_switch_bridge_wire.py @@ -0,0 +1,95 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + + +def test_anthropic_signed_thinking_in_history_crosses_to_responses_bridge(gateway: Gateway) -> None: + turn1: Final = cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/hello.txt and reply with its single word", + ) + turn2: Final = cc.tool_loop_turn2( + turn1, + ( + {"type": "thinking", "thinking": "need to read the file", "signature": "sig_anthropic_1"}, + { + "type": "tool_use", + "id": "toolu_read_1", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/hello.txt"}, + }, + ), + (("toolu_read_1", "1\tPROBE\n2\t"),), + ) + bridge_input_box: list[JsonValue] = [] + + def respond_anthropic(request: Request) -> Reply: + assert request.target == "/v1/messages", request.target + return Reply( + content_type="text/event-stream", + chunks=cc.tool_use_stream( + f"msg_sw_{uuid.uuid4().hex}", + cc.FABLE, + "need to read the file", + "sig_anthropic_1", + (("toolu_read_1", "Read", {"file_path": "/tmp/cc_probe/hello.txt"}),), + {"input_tokens": 20, "output_tokens": 10}, + ), + ) + + def respond_openai(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + bridge_input_box.append(body.get("input")) + reasoning_items: Final = [ + item for item in body["input"] if isinstance(item, dict) and item.get("type") == "reasoning" + ] + assert reasoning_items == [ + {"type": "reasoning", "summary": [{"type": "summary_text", "text": "need to read the file"}]} + ], reasoning_items + return Reply( + content_type="text/event-stream", + chunks=cc.responses_stream( + "sw", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "PROBE", "annotations": []}], + }, + ), + ), + ) + + with ( + wire_server(respond_anthropic) as wire_a, + wire_server(respond_openai) as wire_b, + gateway.scenario() as scenario, + ): + fable: Final = scenario.model(model=f"anthropic/{cc.FABLE}", api_base=wire_a.url, api_key=cc.ANTHROPIC_API_KEY) + openai_alias: Final = scenario.model( + model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire_b.url, api_key=cc.OPENAI_API_KEY + ) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response1: Final = gateway.request( + "POST", "/v1/messages", {**turn1, "model": fable}, params={"beta": "true"}, headers=headers + ) + assert response1.status_code == 200, response1.text + response2: Final = gateway.request( + "POST", "/v1/messages", {**turn2, "model": openai_alias}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + events: Final = cc.sse_events(response2.text) + assert events[-1][0] == "message_stop" + assert len(wire_a.drain()) == 1 + assert len(wire_b.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_tool_loop_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_tool_loop_bridge_wire.py new file mode 100644 index 00000000000..fa33c98048a --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_tool_loop_bridge_wire.py @@ -0,0 +1,223 @@ +import json +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_INSTRUCTIONS: Final = "\n".join(block["text"] for block in cc.system_blocks()) + + +def _turn1() -> dict[str, JsonValue]: + return { + **cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Read /tmp/cc_probe/hello.txt and reply with its single word", + ), + "stream": False, + } + + +def _upstream_items(calls: tuple[tuple[str, str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return ( + { + "type": "reasoning", + "id": "rs_1", + "summary": [{"type": "summary_text", "text": "short plan"}], + "encrypted_content": "enc_1", + }, + *( + { + "type": "function_call", + "id": f"fc_{i}", + "call_id": call_id, + "name": name, + "arguments": json.dumps(tool_input), + "status": "completed", + } + for i, (call_id, name, tool_input) in enumerate(calls, start=1) + ), + ) + + +def _expected_turn2_input( + turn1: dict[str, JsonValue], + assistant_content: tuple[dict[str, JsonValue], ...], + tool_results: tuple[tuple[str, JsonValue], ...], +) -> list[JsonValue]: + items: Final = [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": block["text"]} for block in turn1["messages"][0]["content"]], + }, + { + "type": "message", + "role": "system", + "content": [{"type": "input_text", "text": block["text"]} for block in turn1["messages"][1]["content"]], + }, + { + "type": "reasoning", + "summary": [{"type": "summary_text", "text": "short plan"}], + "encrypted_content": "enc_1", + }, + ] + items += [ + { + "type": "function_call", + "call_id": block["id"], + "name": block["name"], + "arguments": json.dumps(block["input"]), + } + for block in assistant_content + if block.get("type") == "tool_use" + ] + items += [ + {"type": "function_call_output", "call_id": tool_use_id, "output": content} + for tool_use_id, content in tool_results + ] + items.append( + { + "type": "message", + "role": "system", + "content": [ + {"type": "input_text", "text": "14999970 tokens left"}, + { + "type": "input_text", + "text": "First privately list what you need next; then request every item that doesn't depend on another's result in this one response.", + }, + ], + } + ) + return items + + +def test_responses_bridge_replays_reasoning_and_tool_call_on_turn_two(gateway: Gateway) -> None: + turn1: Final = _turn1() + calls: Final = (("call_1", "Read", {"file_path": "/tmp/cc_probe/hello.txt"}),) + tool_results: Final = (("call_1", "1\tPROBE\n2\t"),) + seen: list[dict[str, JsonValue]] = [] + signature_box: list[str] = [] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen.append(body) + if len(seen) == 1: + return Reply( + body=cc.responses_completed( + "tl1", + cc.OPENAI_BACKEND, + _upstream_items(calls), + {"input_tokens": 41, "output_tokens": 5, "total_tokens": 46}, + ) + ) + assistant_content: Final = ( + {"type": "thinking", "thinking": "short plan", "signature": signature_box[0]}, + {"type": "tool_use", "id": "call_1", "name": "Read", "input": {"file_path": "/tmp/cc_probe/hello.txt"}}, + ) + assert body["input"] == _expected_turn2_input(turn1, assistant_content, tool_results), body["input"] + return Reply( + body=cc.responses_completed( + "tl2", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "PROBE", "annotations": []}], + }, + ), + {"input_tokens": 50, "output_tokens": 3, "total_tokens": 53}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + headers: Final = cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA) + response1: Final = gateway.request( + "POST", "/v1/messages", {**turn1, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response1.status_code == 200, response1.text + payload: Final = cc.JSON_OBJECT.validate_json(response1.content) + assert payload["stop_reason"] == "tool_use", payload + assert payload["content"][0]["type"] == "thinking", payload["content"] + signature: Final = payload["content"][0].get("signature") + assert signature, payload["content"][0] + signature_box.append(signature) + assert payload["content"][1] == { + "type": "tool_use", + "id": "call_1", + "name": "Read", + "input": {"file_path": "/tmp/cc_probe/hello.txt"}, + }, payload["content"] + assistant_content: Final = ( + {"type": "thinking", "thinking": "short plan", "signature": signature}, + {"type": "tool_use", "id": "call_1", "name": "Read", "input": {"file_path": "/tmp/cc_probe/hello.txt"}}, + ) + turn2: Final = cc.tool_loop_turn2(turn1, assistant_content, tool_results) + response2: Final = gateway.request( + "POST", "/v1/messages", {**turn2, "model": model}, params={"beta": "true"}, headers=headers + ) + assert response2.status_code == 200, response2.text + payload2: Final = cc.JSON_OBJECT.validate_json(response2.content) + assert payload2["stop_reason"] == "end_turn", payload2 + assert len(wire.drain()) == 2 + + +def test_responses_bridge_replays_parallel_tool_calls_in_order(gateway: Gateway) -> None: + turn1: Final = _turn1() + calls: Final = ( + ("call_1", "Read", {"file_path": "/tmp/cc_probe/hello.txt"}), + ("call_2", "Read", {"file_path": "/tmp/cc_probe/world.txt"}), + ) + + def respond(request: Request) -> Reply: + body: Final = cc.JSON_OBJECT.validate_json(request.body) + seen_items: Final = body["input"] + outputs: Final = [ + item for item in seen_items if isinstance(item, dict) and item.get("type") == "function_call_output" + ] + assert [item["call_id"] for item in outputs] == ["call_1", "call_2"], outputs + return Reply( + body=cc.responses_completed( + "mt", + cc.OPENAI_BACKEND, + ( + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "PROBE PROBE2", "annotations": []}], + }, + ), + {"input_tokens": 50, "output_tokens": 4, "total_tokens": 54}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + assistant_content: Final = tuple( + {"type": "tool_use", "id": call_id, "name": name, "input": tool_input} + for call_id, name, tool_input in calls + ) + turn2: Final = cc.tool_loop_turn2( + turn1, assistant_content, (("call_1", "1\tPROBE\n2\t"), ("call_2", "1\tPROBE2\n2\t")) + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**turn2, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/responses_bridge/test_claude_code_web_search_bridge_wire.py b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_web_search_bridge_wire.py new file mode 100644 index 00000000000..e06e1b8d403 --- /dev/null +++ b/tests/integration/messages_endpoint/responses_bridge/test_claude_code_web_search_bridge_wire.py @@ -0,0 +1,95 @@ +import uuid +from typing import Final + +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.messages_endpoint import _claude_code as cc +from pydantic import JsonValue + +_WEB_SEARCH_TOOL: Final = { + "name": "WebSearch", + "description": "Search the web. Returns result blocks with titles and URLs.", + "input_schema": cc.schema( + { + "query": cc.field("The search query to use", type="string", minLength=2), + "allowed_domains": cc.field( + "Only include search results from these domains", type="array", items={"type": "string"} + ), + "blocked_domains": cc.field( + "Never include search results from these domains", type="array", items={"type": "string"} + ), + }, + ("query",), + ), +} + + +def test_web_search_tool_and_cited_output_on_responses_bridge(gateway: Gateway) -> None: + request_body: Final = { + **cc.frontier_request( + f"cache-bust-{uuid.uuid4().hex}", + "high", + 64000, + prompt_text="Use web search to find the current LiteLLM version and answer in one word", + ), + "stream": False, + } + request_body["tools"] = [*request_body["tools"], _WEB_SEARCH_TOOL] + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses", request.target + body: Final = cc.JSON_OBJECT.validate_json(request.body) + tools: Final = body["tools"] + assert tools[-1] == { + "type": "function", + "name": "WebSearch", + "strict": False, + "description": _WEB_SEARCH_TOOL["description"], + "parameters": _WEB_SEARCH_TOOL["input_schema"], + }, tools[-1] + return Reply( + body=cc.responses_completed( + "ws", + cc.OPENAI_BACKEND, + ( + {"type": "web_search_call", "id": "ws_1", "status": "completed"}, + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "1.104.0", + "annotations": [ + { + "type": "url_citation", + "url": "https://example.com/litellm", + "title": "litellm releases", + "start_index": 0, + "end_index": 7, + } + ], + } + ], + }, + ), + {"input_tokens": 41, "output_tokens": 5, "total_tokens": 46}, + ) + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{cc.OPENAI_BACKEND}", api_base=wire.url, api_key=cc.OPENAI_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {**request_body, "model": model}, + params={"beta": "true"}, + headers=cc.cli_headers(gateway.key, cc.FRONTIER_CLI_BETA), + ) + assert response.status_code == 200, response.text + payload: Final = cc.JSON_OBJECT.validate_json(response.content) + assert payload["content"] == [{"type": "text", "text": "1.104.0"}], payload["content"] + assert len(wire.drain()) == 1