diff --git a/tests/integration/observability/test_passthrough_upstream_error_visibility.py b/tests/integration/observability/test_passthrough_upstream_error_visibility.py index 0d86440045b..83560f863ed 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_visibility.py +++ b/tests/integration/observability/test_passthrough_upstream_error_visibility.py @@ -1,16 +1,20 @@ import gzip import json +import threading from hashlib import sha256 from pathlib import Path from typing import Final +from urllib.parse import parse_qsl, urlsplit import httpx import pytest import yaml +from google import genai +from google.genai import types from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.process import owned_proxy_process -from integration._support.wire import Reply, Request, wire_server +from integration._support.wire import Reply, Request, Wire, wire_server from openai import AsyncOpenAI, NotFoundError, OpenAI from pydantic import JsonValue @@ -265,9 +269,7 @@ def test_gemini_streaming_passthrough_request_tags_reach_spend_log(gateway: Gate raw_tags: Final = row["request_tags"] request_tags: Final = json.loads(raw_tags) if isinstance(raw_tags, str) else raw_tags policy_tags: Final = [ - tag - for tag in request_tags - if not (isinstance(tag, str) and tag.startswith("User-Agent: ")) + tag for tag in request_tags if not (isinstance(tag, str) and tag.startswith("User-Agent: ")) ] assert policy_tags == ["gemini-js-sdk", "pass-through-endpoint"], row assert row["call_type"] == "pass_through_endpoint", row @@ -650,3 +652,359 @@ def test_budget_rejected_call_keeps_budget_normalized_error(gateway: Gateway, tm ) ) assert len(budget_rows) == 1, budget_rows + + +_GEMINI_REPLY: Final[dict[str, JsonValue]] = { + "candidates": [{"content": {"parts": [{"text": "scripted gemini"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, +} +_GEMINI_SSE: Final = (f"data: {json.dumps(_GEMINI_REPLY)}\r\n\r\n".encode(),) +_GEMINI_BODY: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": {"parts": [{"text": "be terse"}]}, + "generationConfig": {"temperature": 0.2, "maxOutputTokens": 32, "thinkingConfig": {"thinkingBudget": 0}}, +} +_GEMINI_SDK_BODY: Final[dict[str, JsonValue]] = { + **_GEMINI_BODY, + "systemInstruction": {"parts": [{"text": "be terse"}], "role": "user"}, +} + + +def _upstream_query(request: Request) -> tuple[str, list[tuple[str, str]]]: + split: Final = urlsplit(request.target) + return split.path, parse_qsl(split.query, keep_blank_values=True) + + +def test_gemini_passthrough_both_auth_spellings_swap_the_virtual_key_for_the_proxy_key( + gateway: Gateway, tmp_path: Path +) -> None: + def respond(request: Request) -> Reply: + if ":streamGenerateContent" in request.target: + return Reply(chunks=_GEMINI_SSE, content_type="text/event-stream") + return Reply(body=json.dumps(_GEMINI_REPLY).encode()) + + path: Final = tmp_path / "gemini-auth-spellings.yaml" + model_path: Final = "/v1beta/models/gemini-2.5-flash" + with wire_server(respond) as wire: + _gemini_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + key: Final = scenario.key() + sent: Final[list[httpx.Request]] = [] + client: Final = genai.Client( + api_key=key, + http_options=types.HttpOptions( + base_url=f"{str(candidate.client.base_url).rstrip('/')}/gemini", + client_args={"event_hooks": {"request": [sent.append]}, "timeout": 30, "trust_env": False}, + ), + ) + generated: Final = client.models.generate_content( + model="gemini-2.5-flash", + contents="hi", + config=types.GenerateContentConfig( + system_instruction="be terse", + temperature=0.2, + max_output_tokens=32, + thinking_config=types.ThinkingConfig(thinking_budget=0), + ), + ) + streamed_text: Final = "".join( + chunk.text or "" + for chunk in client.models.generate_content_stream( + model="gemini-2.5-flash", + contents="hi", + config=types.GenerateContentConfig( + system_instruction="be terse", + temperature=0.2, + max_output_tokens=32, + thinking_config=types.ThinkingConfig(thinking_budget=0), + ), + ) + ) + assert generated.text == "scripted gemini", generated + assert streamed_text == "scripted gemini", streamed_text + assert len(sent) == 2, sent + assert all(request.headers.get("x-goog-api-key") == key for request in sent), sent + assert parse_qsl(sent[1].url.query.decode(), keep_blank_values=True) == [("alt", "sse")], sent[1].url + query_spelling: Final = candidate.client.post( + f"/gemini{model_path}:generateContent", params={"key": key}, json=_GEMINI_BODY + ) + assert query_spelling.status_code == 200, query_spelling.text + assert query_spelling.json() == _GEMINI_REPLY, query_spelling.text + with candidate.client.stream( + "POST", + f"/gemini{model_path}:streamGenerateContent", + params={"key": key, "alt": "sse"}, + json=_GEMINI_BODY, + ) as streamed_response: + streamed: Final = streamed_response.read() + assert streamed_response.status_code == 200, streamed + assert streamed == b"".join(_GEMINI_SSE), streamed + received: Final = wire.drain() + assert [(request.method, *_upstream_query(request)) for request in received] == [ + ("POST", f"{model_path}:generateContent", [("key", "scripted")]), + ("POST", f"{model_path}:streamGenerateContent", [("alt", "sse"), ("key", "scripted")]), + ("POST", f"{model_path}:generateContent", [("key", "scripted")]), + ("POST", f"{model_path}:streamGenerateContent", [("key", "scripted"), ("alt", "sse")]), + ], received + for upstream, sdk_request in zip(received[:2], sent, strict=True): + sdk_body: Final = json.loads(sdk_request.content) + assert json.loads(upstream.body) == sdk_body, upstream.body + assert sdk_body == _GEMINI_SDK_BODY, sdk_request.content + for request in received[2:]: + assert json.loads(request.body) == _GEMINI_BODY, request.body + for request in received: + assert {name: value for name, value in request.headers.items() if key in value} == {}, ( + request.headers + ) + assert "x-goog-api-key" not in request.headers, request.headers + assert key not in request.target, request.target + + +_CHAT_REPLY: Final[dict[str, JsonValue]] = { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "scripted chat"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 3, "total_tokens": 8}, +} +_RESPONSE_OBJECT: Final[dict[str, JsonValue]] = { + "id": "resp_scripted", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": { + "input_tokens": 5, + "output_tokens": 3, + "total_tokens": 8, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, +} +_RESPONSE_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "response.created", + "sequence_number": 0, + "response": {**_RESPONSE_OBJECT, "status": "in_progress", "output": []}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_scripted", + "output_index": 0, + "content_index": 0, + "delta": "scripted response", + "logprobs": [], + }, + {"type": "response.completed", "sequence_number": 2, "response": _RESPONSE_OBJECT}, +) +_RESPONSE_SSE: Final = tuple( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _RESPONSE_EVENTS +) +_ASSISTANTS_PAGE: Final[dict[str, JsonValue]] = { + "object": "list", + "data": [ + { + "id": "asst_scripted", + "object": "assistant", + "created_at": 1, + "name": "scripted", + "description": None, + "model": "gpt-4o-mini", + "instructions": None, + "tools": [], + "metadata": {}, + } + ], + "first_id": "asst_scripted", + "last_id": "asst_scripted", + "has_more": False, +} + + +_RESPONSE_STREAM_GATE: Final = threading.Event() +_RESPONSE_STREAM_GATE.set() + + +def _openai_peer(request: Request) -> Reply: + if request.target.startswith("/v1/assistants"): + return Reply(body=json.dumps(_ASSISTANTS_PAGE).encode()) + if request.target == "/v1/chat/completions": + return Reply(body=json.dumps(_CHAT_REPLY).encode()) + if json.loads(request.body).get("stream") is True: + return Reply(chunks=_RESPONSE_SSE, content_type="text/event-stream", gate_after_first=_RESPONSE_STREAM_GATE) + return Reply(body=json.dumps(_RESPONSE_OBJECT).encode()) + + +def _openai_sdk(candidate: Gateway, key: str, prefix: str, sent: list[httpx.Request]) -> OpenAI: + return OpenAI( + api_key=key, + base_url=f"{str(candidate.client.base_url).rstrip('/')}{prefix}/v1", + max_retries=0, + http_client=httpx.Client(timeout=30, trust_env=False, event_hooks={"request": [sent.append]}), + ) + + +def _assert_openai_upstream_auth(request: Request, key: str) -> None: + assert request.headers["authorization"] == "Bearer scripted", request.headers + assert {name: value for name, value in request.headers.items() if key in value} == {}, request.headers + assert key.encode() not in request.body, request.body + + +def _assert_openai_sdk_request( + wire: Wire, + sent: list[httpx.Request], + key: str, + method: str, + target: str, + expected_body: dict[str, JsonValue] | None, +) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [(method, target)], received + assert len(sent) == 1, sent + client_request: Final = sent[0] + upstream: Final = received[0] + _assert_openai_upstream_auth(upstream, key) + if expected_body is None: + assert client_request.content == b"", client_request.content + assert upstream.body == b"", upstream.body + else: + assert json.loads(client_request.content) == expected_body, client_request.content + assert json.loads(upstream.body) == expected_body, upstream.body + return upstream + + +_OPENAI_CHAT_BODY: Final[dict[str, JsonValue]] = { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "user": "end-user-1", + "store": True, + "temperature": 0.2, +} +_OPENAI_RESPONSES_BODY: Final[dict[str, JsonValue]] = { + "model": "gpt-4o-mini", + "input": "hi", + "store": False, + "user": "end-user-1", + "temperature": 0.2, +} +_OPENAI_RESPONSES_STREAM_BODY: Final[dict[str, JsonValue]] = { + **_OPENAI_RESPONSES_BODY, + "stream": True, +} + + +@pytest.mark.parametrize( + ("base_path", "surface"), + [ + pytest.param("/openai", "chat", id="openai-chat"), + pytest.param("/openai_passthrough", "chat", id="openai-passthrough-chat"), + pytest.param("/openai_passthrough", "responses", id="openai-passthrough-responses"), + pytest.param("/openai_passthrough", "responses-stream", id="openai-passthrough-responses-stream"), + pytest.param("/openai", "assistants", id="openai-assistants"), + pytest.param("/openai_passthrough", "assistants", id="openai-passthrough-assistants"), + ], +) +def test_openai_passthrough_sdk_success_on_both_aliases_forwards_the_native_request( + gateway: Gateway, tmp_path: Path, base_path: str, surface: str +) -> None: + path: Final = tmp_path / "openai-success.yaml" + with wire_server(_openai_peer) as wire: + _openai_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + key: Final = scenario.key() + sent: Final[list[httpx.Request]] = [] + with _openai_sdk(candidate, key, base_path, sent) as sdk: + if surface == "chat": + chat: Final = sdk.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + user="end-user-1", + store=True, + temperature=0.2, + ) + assert chat.model_dump(exclude_unset=True) == _CHAT_REPLY, chat + _assert_openai_sdk_request(wire, sent, key, "POST", "/v1/chat/completions", _OPENAI_CHAT_BODY) + elif surface == "responses": + response: Final = sdk.responses.create( + model="gpt-4o-mini", + input="hi", + store=False, + user="end-user-1", + temperature=0.2, + ) + assert response.output_text == "scripted response", response + assert response.id == "resp_scripted", response + _assert_openai_sdk_request(wire, sent, key, "POST", "/v1/responses", _OPENAI_RESPONSES_BODY) + elif surface == "responses-stream": + _RESPONSE_STREAM_GATE.clear() + try: + stream: Final = iter( + sdk.responses.create( + model="gpt-4o-mini", + input="hi", + store=False, + user="end-user-1", + temperature=0.2, + stream=True, + ) + ) + first_event: Final = next(stream) + finally: + _RESPONSE_STREAM_GATE.set() + assert first_event.type == _RESPONSE_EVENTS[0]["type"], first_event + events: Final = (first_event, *stream) + streamed_text: Final = "".join( + event.delta for event in events if event.type == "response.output_text.delta" + ) + final: Final = next(event.response for event in events if event.type == "response.completed") + assert [event.type for event in events] == [event["type"] for event in _RESPONSE_EVENTS], events + assert streamed_text == "scripted response", events + assert final.output_text == "scripted response", final + _assert_openai_sdk_request( + wire, sent, key, "POST", "/v1/responses", _OPENAI_RESPONSES_STREAM_BODY + ) + else: + assistants: Final = sdk.beta.assistants.list(limit=2, order="desc") + assert [assistant.id for assistant in assistants.data] == ["asst_scripted"], assistants + upstream: Final = _assert_openai_sdk_request( + wire, + sent, + key, + "GET", + "/v1/assistants?limit=2&order=desc", + None, + ) + assert upstream.headers["openai-beta"] == "assistants=v2", upstream.headers + + +def test_openai_passthrough_responses_metadata_reaches_upstream(gateway: Gateway, tmp_path: Path) -> None: + pytest.skip("BUG: /openai_passthrough/v1/responses drops the native metadata field before forwarding upstream") + path: Final = tmp_path / "openai-metadata.yaml" + with wire_server(_openai_peer) as wire: + _openai_config(path, wire.url) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + with candidate.scenario() as scenario: + sent: Final[list[httpx.Request]] = [] + with _openai_sdk(candidate, scenario.key(), "/openai_passthrough", sent) as sdk: + sdk.responses.create(model="gpt-4o-mini", input="hi", metadata={"session": "scripted"}) + received: Final = wire.drain() + assert [json.loads(request.body) for request in received] == [json.loads(sent[0].content)], received diff --git a/tests/integration/providers/test_anthropic_passthrough_wire.py b/tests/integration/providers/test_anthropic_passthrough_wire.py new file mode 100644 index 00000000000..a791eb9c2a2 --- /dev/null +++ b/tests/integration/providers/test_anthropic_passthrough_wire.py @@ -0,0 +1,332 @@ +import json +import threading +from collections.abc import Iterator +from itertools import accumulate, dropwhile +from pathlib import Path +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(240) + +_API_KEY: Final = "scripted-anthropic-key" +_OAUTH_TOKEN: Final = "sk-ant-oat01-scripted-proxy-token" +_MODEL: Final = "claude-sonnet-4-5" +_TEXT: Final = "scripted passthrough answer" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_SYSTEM: Final[list[JsonValue]] = [{"type": "text", "text": "You are a terse assistant"}] +_TOOLS: Final[list[JsonValue]] = [ + { + "name": "get_weather", + "description": "Look up the weather", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } +] +_THINKING: Final[dict[str, JsonValue]] = {"type": "enabled", "budget_tokens": 1024} +_MESSAGE: Final[dict[str, JsonValue]] = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [{"type": "text", "text": _TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 7, "output_tokens": 3}, +} +_SSE_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**_MESSAGE, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _TEXT}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, +) +_SSE_FRAMES: Final = tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _SSE_EVENTS) +_CLAUDE_CODE_STREAM_GATE: Final = threading.Event() + + +def _respond(request: Request) -> Reply: + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + gate: Final = _CLAUDE_CODE_STREAM_GATE if request.headers.get("x-app") == "cli" else None + return Reply(chunks=_SSE_FRAMES, content_type="text/event-stream", gate_after_first=gate) + return Reply(body=json.dumps(_MESSAGE).encode()) + + +@pytest.fixture(scope="module") +def rig() -> Iterator[Gateway]: + with gateway_from_environment() as gateway: + yield gateway + + +@pytest.fixture(scope="module") +def wire() -> Iterator[Wire]: + with wire_server(_respond) as served: + yield served + + +@pytest.fixture(autouse=True) +def _drained(wire: Wire) -> None: + wire.drain() + + +def _config(directory: Path) -> Path: + path: Final = directory / "anthropic-passthrough.yaml" + path.write_text( + "model_list: []\n" + "general_settings:\n" + " master_key: os.environ/LITELLM_MASTER_KEY\n" + " database_url: os.environ/DATABASE_URL\n" + " store_model_in_db: true\n" + "router_settings:\n" + " disable_cooldowns: true\n" + ) + return path + + +@pytest.fixture(scope="module") +def api_key_proxy(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("anthropic-passthrough-api-key") + overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _API_KEY} + with owned_proxy(rig, directory, overrides, config=_config(directory), workers=2) as owned: + yield owned + + +@pytest.fixture(scope="module") +def oauth_proxy(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("anthropic-passthrough-oauth") + overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": _OAUTH_TOKEN} + with owned_proxy(rig, directory, overrides, config=_config(directory), workers=2) as owned: + yield owned + + +@pytest.fixture(scope="module") +def credentialless_proxy(rig: Gateway, wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("anthropic-passthrough-credentialless") + overrides: Final = {"ANTHROPIC_API_BASE": wire.url, "ANTHROPIC_API_KEY": "", "ANTHROPIC_AUTH_TOKEN": ""} + with owned_proxy(rig, directory, overrides, config=_config(directory), workers=2) as owned: + yield owned + + +def _sdk(proxy: Gateway, key: str, sent: list[httpx.Request]) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=key, + base_url=f"{str(proxy.client.base_url).rstrip('/')}/anthropic", + max_retries=0, + http_client=httpx.Client(timeout=30, trust_env=False, event_hooks={"request": [sent.append]}), + ) + + +def _only_upstream_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", "/v1/messages")], received + return received[0] + + +def _assert_virtual_key_absent(upstream: Request, key: str) -> None: + leaked: Final = {name: value for name, value in upstream.headers.items() if key in value} + assert leaked == {}, upstream.headers + assert key.encode() not in upstream.body, upstream.body + + +def _assert_sdk_call_reached_upstream_intact(upstream: Request, sent: httpx.Request, key: str) -> None: + assert _JSON_OBJECT.validate_json(upstream.body) == _JSON_OBJECT.validate_json(sent.content), upstream.body + assert upstream.headers["x-api-key"] == _API_KEY, upstream.headers + assert "authorization" not in upstream.headers, upstream.headers + assert upstream.headers["anthropic-version"] == sent.headers["anthropic-version"], upstream.headers + assert upstream.headers["anthropic-beta"] == "interleaved-thinking-2025-05-14", upstream.headers + _assert_virtual_key_absent(upstream, key) + + +@pytest.mark.parametrize("stream", [False, True], ids=["json", "sse"]) +def test_anthropic_sdk_call_reaches_upstream_with_full_body_and_proxy_key( + api_key_proxy: Gateway, wire: Wire, stream: bool +) -> None: + with api_key_proxy.scenario() as scenario: + key: Final = scenario.key() + sent: Final[list[httpx.Request]] = [] + with _sdk(api_key_proxy, key, sent) as sdk: + arguments: Final = { + "model": _MODEL, + "max_tokens": 2048, + "system": _SYSTEM, + "tools": _TOOLS, + "thinking": _THINKING, + "messages": [{"role": "user", "content": "weather in Paris?"}], + "extra_headers": {"anthropic-beta": "interleaved-thinking-2025-05-14"}, + } + if stream: + with sdk.messages.stream(**arguments) as events: + message = events.get_final_message() + else: + message = sdk.messages.create(**arguments) + assert message.model_dump(exclude_none=True) == anthropic.types.Message.model_validate(_MESSAGE).model_dump( + exclude_none=True + ), message + upstream: Final = _only_upstream_request(wire) + assert len(sent) == 1, sent + _assert_sdk_call_reached_upstream_intact(upstream, sent[0], key) + + +@pytest.mark.parametrize("auth_style", ["api-key", "oauth"]) +def test_anthropic_proxy_credentials_override_caller_credentials( + api_key_proxy: Gateway, oauth_proxy: Gateway, wire: Wire, auth_style: str +) -> None: + proxy: Final = api_key_proxy if auth_style == "api-key" else oauth_proxy + caller_api_key: Final = f"caller-anthropic-key-{auth_style}" + caller_authorization: Final = f"Bearer caller-anthropic-token-{auth_style}" + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + } + with ( + proxy.scenario() as scenario, + httpx.Client(base_url=str(proxy.client.base_url), timeout=30, trust_env=False) as client, + ): + key: Final = scenario.key() + response: Final = client.post( + "/anthropic/v1/messages", + json=body, + headers={ + "x-litellm-api-key": key, + "x-api-key": caller_api_key, + "Authorization": caller_authorization, + }, + ) + assert response.status_code == 200, response.text + message: Final = anthropic.types.Message.model_validate_json(response.content) + assert message.model_dump(exclude_none=True) == anthropic.types.Message.model_validate(_MESSAGE).model_dump( + exclude_none=True + ), response.text + upstream: Final = _only_upstream_request(wire) + assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body + assert { + name: value + for name, value in upstream.headers.items() + if caller_api_key in value or caller_authorization in value + } == {}, upstream.headers + if auth_style == "api-key": + assert upstream.headers["x-api-key"] == _API_KEY, upstream.headers + assert "authorization" not in upstream.headers, upstream.headers + else: + assert upstream.headers["authorization"] == f"Bearer {_OAUTH_TOKEN}", upstream.headers + assert "x-api-key" not in upstream.headers, upstream.headers + _assert_virtual_key_absent(upstream, key) + + +def test_claude_code_bearer_stream_merges_caller_beta_with_the_credential_beta( + oauth_proxy: Gateway, wire: Wire +) -> None: + _CLAUDE_CODE_STREAM_GATE.clear() + with oauth_proxy.scenario() as scenario: + key: Final = scenario.key() + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_tokens": 2048, + "stream": True, + "system": _SYSTEM, + "tools": _TOOLS, + "thinking": _THINKING, + "messages": [{"role": "user", "content": "weather in Paris?"}], + } + with oauth_proxy.client.stream( + "POST", + "/anthropic/v1/messages", + json=body, + headers={ + "Authorization": f"Bearer {key}", + "anthropic-version": "2023-06-01", + "anthropic-beta": "claude-code-20250219,interleaved-thinking-2025-05-14", + "x-app": "cli", + }, + ) as response: + raw: Final = iter(response.iter_raw()) + try: + assert response.status_code == 200, response.status_code + assert response.headers.get("content-type", "").startswith("text/event-stream"), response.headers + first_frame: Final = next( + dropwhile( + lambda buffered: b"\n\n" not in buffered, + accumulate(raw, lambda buffered, chunk: buffered + chunk, initial=b""), + ) + ) + assert first_frame == _SSE_FRAMES[0], first_frame + finally: + _CLAUDE_CODE_STREAM_GATE.set() + streamed: Final = first_frame + b"".join(raw) + assert streamed == b"".join(_SSE_FRAMES), streamed + _CLAUDE_CODE_STREAM_GATE.clear() + upstream: Final = _only_upstream_request(wire) + assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body + assert upstream.headers["authorization"] == f"Bearer {_OAUTH_TOKEN}", upstream.headers + assert "x-api-key" not in upstream.headers, upstream.headers + assert upstream.headers["anthropic-version"] == "2023-06-01", upstream.headers + assert upstream.headers["anthropic-beta"] == ( + "claude-code-20250219,interleaved-thinking-2025-05-14,oauth-2025-04-20" + ), upstream.headers + assert upstream.headers["x-app"] == "cli", upstream.headers + _assert_virtual_key_absent(upstream, key) + + +@pytest.mark.parametrize( + ("credential_header", "credential"), + [("authorization", "Bearer caller-anthropic-oauth-token"), ("x-api-key", "caller-anthropic-api-key")], + ids=["caller-oauth-token", "caller-api-key"], +) +def test_credentialless_proxy_forwards_the_caller_anthropic_credential_and_never_the_virtual_key( + credentialless_proxy: Gateway, wire: Wire, credential_header: str, credential: str +) -> None: + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + } + with ( + credentialless_proxy.scenario() as scenario, + httpx.Client(base_url=str(credentialless_proxy.client.base_url), timeout=30, trust_env=False) as client, + ): + key: Final = scenario.key() + response: Final = client.post( + "/anthropic/v1/messages", json=body, headers={"x-litellm-api-key": key, credential_header: credential} + ) + assert response.status_code == 200, response.text + message: Final = anthropic.types.Message.model_validate_json(response.content) + assert message.model_dump(exclude_none=True) == anthropic.types.Message.model_validate(_MESSAGE).model_dump( + exclude_none=True + ), response.text + upstream: Final = _only_upstream_request(wire) + assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body + assert upstream.headers[credential_header] == credential, upstream.headers + assert {"authorization", "x-api-key"} & set(upstream.headers) == {credential_header}, upstream.headers + _assert_virtual_key_absent(upstream, key) + + refused: Final = client.post("/anthropic/v1/messages", json=body, headers={"x-litellm-api-key": key}) + assert refused.status_code == 401, refused.text + assert "The LiteLLM virtual key is not forwarded to Anthropic" in refused.text, refused.text + assert wire.drain() == () + + +def test_anthropic_sdk_metadata_user_id_reaches_upstream(api_key_proxy: Gateway, wire: Wire) -> None: + pytest.skip("BUG: /anthropic/v1/messages drops the native metadata.user_id field before forwarding upstream") + with api_key_proxy.scenario() as scenario: + key: Final = scenario.key() + sent: Final[list[httpx.Request]] = [] + with _sdk(api_key_proxy, key, sent) as sdk: + sdk.messages.create( + model=_MODEL, + max_tokens=64, + metadata={"user_id": "user_scripted_session_1"}, + messages=[{"role": "user", "content": "hi"}], + ) + upstream: Final = _only_upstream_request(wire) + assert _JSON_OBJECT.validate_json(upstream.body) == _JSON_OBJECT.validate_json(sent[0].content), upstream.body diff --git a/tests/integration/providers/test_bedrock_passthrough_stream_wire.py b/tests/integration/providers/test_bedrock_passthrough_stream_wire.py index bb9bbc65f30..74d3ac7d2ef 100644 --- a/tests/integration/providers/test_bedrock_passthrough_stream_wire.py +++ b/tests/integration/providers/test_bedrock_passthrough_stream_wire.py @@ -1,10 +1,22 @@ +import base64 import json +import threading +import uuid +from collections.abc import Callable, Iterable +from pathlib import Path from typing import Final +from urllib.parse import quote, unquote +import boto3 import pytest +import yaml +from botocore.awsrequest import AWSPreparedRequest +from botocore.config import Config from integration._support.client import Gateway +from integration._support.process import owned_proxy from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue _MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" _EVENT_STREAM: Final = "application/vnd.amazon.eventstream" @@ -40,3 +52,299 @@ def test_bedrock_passthrough_converse_stream_response_carries_event_stream_conte assert len(wire.drain()) == 1, response.text assert response.headers.get("content-type") == _EVENT_STREAM, dict(response.headers) assert response.content == _STREAM_BYTES, response.text + + +_ACCOUNT: Final = "123456789012" +_CONVERSE_BODY: Final[dict[str, JsonValue]] = { + "messages": [{"role": "user", "content": [{"text": "full body"}]}], + "system": [{"text": "be terse"}], + "inferenceConfig": {"maxTokens": 64, "temperature": 0.2, "stopSequences": ["END"]}, + "toolConfig": { + "tools": [ + { + "toolSpec": { + "name": "lookup", + "description": "look a value up", + "inputSchema": {"json": {"type": "object", "properties": {"q": {"type": "string"}}}}, + } + } + ], + "toolChoice": {"auto": {}}, + }, +} +_INVOKE_BODY: Final[dict[str, JsonValue]] = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 64, + "system": "be terse", + "messages": [{"role": "user", "content": "full body"}], +} +_CONVERSE_REPLY: Final[dict[str, JsonValue]] = { + "output": {"message": {"role": "assistant", "content": [{"text": "bedrock converse"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 2, "totalTokens": 7}, + "metrics": {"latencyMs": 12}, +} +_INVOKE_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": "bedrock invoke"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 2}, +} + + +_CONVERSE_STREAM_FRAMES: Final = tuple(_aws_event_frame(kind, payload, "sc", "u") for kind, payload in _EVENTS) +_INVOKE_STREAM_FRAMES: Final = ( + _aws_event_frame( + "chunk", + {"bytes": base64.b64encode(json.dumps(_INVOKE_REPLY).encode()).decode()}, + "sc", + "u", + ), +) + + +_GATED_STREAM_TARGETS: Final = (f"/model/{_MODEL_ID}/converse-stream", "-direct-") +_BEDROCK_STREAM_GATE: Final = threading.Event() +_BEDROCK_STREAM_GATE.set() + + +def _full_body_peer(request: Request) -> Reply: + if request.target.endswith("/converse-stream"): + gated: Final = any(marker in unquote(request.target) for marker in _GATED_STREAM_TARGETS) + return Reply( + chunks=_CONVERSE_STREAM_FRAMES, + content_type=_EVENT_STREAM, + gate_after_first=_BEDROCK_STREAM_GATE if gated else None, + ) + if request.target.endswith("/invoke-with-response-stream"): + return Reply(chunks=_INVOKE_STREAM_FRAMES, content_type=_EVENT_STREAM) + if request.target.endswith("/converse"): + return Reply(body=json.dumps(_CONVERSE_REPLY).encode()) + return Reply(body=json.dumps(_INVOKE_REPLY).encode()) + + +def _boto3_runtime(gateway: Gateway, key: str): # noqa: ANN202 - boto3 clients are untyped + client: Final = boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(gateway.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + aws_access_key_id="AKIACALLERSIDE", + aws_secret_access_key="caller-side-secret", + config=Config(retries={"max_attempts": 1}), + ) + + def bearer(request: AWSPreparedRequest, **_: object) -> None: + request.headers["Authorization"] = f"Bearer {key}" + + client.meta.events.register("before-send.bedrock-runtime.*", bearer) + return client + + +def _events_read_before_the_upstream_finishes(open_stream: Callable[[], Iterable[object]]) -> tuple[object, ...]: + _BEDROCK_STREAM_GATE.clear() + events: Final = iter(open_stream()) + try: + first: Final = next(events) + finally: + _BEDROCK_STREAM_GATE.set() + return (first, *events) + + +def _assert_signed_with_the_deployment_key(request: Request, key: str) -> None: + authorization: Final = request.headers["authorization"] + assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIASCRIPTEDPROVIDER/"), request.headers + assert "/us-east-1/bedrock/aws4_request" in authorization, request.headers + assert {name: value for name, value in request.headers.items() if key in value} == {}, request.headers + assert key not in request.target, request.target + + +def test_boto3_converse_invoke_and_stream_methods_forward_full_bodies_for_deployments_ids_and_profile_arns( + gateway: Gateway, tmp_path: Path +) -> None: + suffix: Final = uuid.uuid4().hex[:12] + raw_id: Final = f"anthropic.claude-scripted-{suffix}-v1:0" + profile_arn: Final = f"arn:aws:bedrock:us-east-1:{_ACCOUNT}:application-inference-profile/{suffix}" + router_name: Final = f"bedrock-router-{suffix}" + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with wire_server(_full_body_peer) as wire: + deployment_params: Final = { + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + } + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + { + "model_name": router_name, + "litellm_params": {"model": f"bedrock/{_MODEL_ID}", **deployment_params}, + }, + { + "model_name": raw_id, + "litellm_params": {"model": f"bedrock/{raw_id}", **deployment_params}, + }, + { + "model_name": profile_arn, + "litellm_params": {"model": f"bedrock/{profile_arn}", **deployment_params}, + }, + ], + } + config_path: Final = tmp_path / "bedrock-passthrough.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(gateway, tmp_path, {}, config=config_path, workers=1) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + runtime: Final = _boto3_runtime(candidate, key) + router_converse: Final = runtime.converse(modelId=router_name, **_CONVERSE_BODY) + profile_converse: Final = runtime.converse(modelId=profile_arn, **_CONVERSE_BODY) + raw_invoke: Final = runtime.invoke_model(modelId=raw_id, body=json.dumps(_INVOKE_BODY)) + profile_invoke: Final = runtime.invoke_model(modelId=profile_arn, body=json.dumps(_INVOKE_BODY)) + raw_converse_stream: Final = runtime.converse_stream(modelId=raw_id, **_CONVERSE_BODY) + profile_converse_stream: Final = runtime.converse_stream(modelId=profile_arn, **_CONVERSE_BODY) + raw_invoke_stream: Final = runtime.invoke_model_with_response_stream( + modelId=raw_id, body=json.dumps(_INVOKE_BODY) + ) + profile_invoke_stream: Final = runtime.invoke_model_with_response_stream( + modelId=profile_arn, body=json.dumps(_INVOKE_BODY) + ) + router_converse_stream: Final = _events_read_before_the_upstream_finishes( + lambda: runtime.converse_stream(modelId=router_name, **_CONVERSE_BODY)["stream"] + ) + router_invoke_stream: Final = runtime.invoke_model_with_response_stream( + modelId=router_name, body=json.dumps(_INVOKE_BODY) + ) + assert tuple( + {name: reply[name] for name in _CONVERSE_REPLY} for reply in (router_converse, profile_converse) + ) == (_CONVERSE_REPLY, _CONVERSE_REPLY), (router_converse, profile_converse) + assert tuple(json.loads(response["body"].read()) for response in (raw_invoke, profile_invoke)) == ( + _INVOKE_REPLY, + _INVOKE_REPLY, + ), (raw_invoke, profile_invoke) + expected_converse_events: Final = [{kind: payload} for kind, payload in _EVENTS] + assert [list(response["stream"]) for response in (raw_converse_stream, profile_converse_stream)] == [ + expected_converse_events, + expected_converse_events, + ], (raw_converse_stream, profile_converse_stream) + assert [ + [json.loads(event["chunk"]["bytes"]) for event in response["body"]] + for response in (raw_invoke_stream, profile_invoke_stream) + ] == [[_INVOKE_REPLY], [_INVOKE_REPLY]], (raw_invoke_stream, profile_invoke_stream) + assert list(router_converse_stream) == expected_converse_events, router_converse_stream + assert [json.loads(event["chunk"]["bytes"]) for event in router_invoke_stream["body"]] == [_INVOKE_REPLY] + + received: Final = wire.drain() + encoded_arn: Final = quote(profile_arn, safe=":") + assert [(request.method, request.target) for request in received] == [ + ("POST", f"/model/{_MODEL_ID}/converse"), + ("POST", f"/model/{encoded_arn}/converse"), + ("POST", f"/model/{raw_id}/invoke"), + ("POST", f"/model/{encoded_arn}/invoke"), + ("POST", f"/model/{raw_id}/converse-stream"), + ("POST", f"/model/{encoded_arn}/converse-stream"), + ("POST", f"/model/{raw_id}/invoke-with-response-stream"), + ("POST", f"/model/{encoded_arn}/invoke-with-response-stream"), + ("POST", f"/model/{_MODEL_ID}/converse-stream"), + ("POST", f"/model/{_MODEL_ID}/invoke-with-response-stream"), + ], received + assert tuple(json.loads(request.body) for request in received) == ( + _CONVERSE_BODY, + _CONVERSE_BODY, + _INVOKE_BODY, + _INVOKE_BODY, + _CONVERSE_BODY, + _CONVERSE_BODY, + _INVOKE_BODY, + _INVOKE_BODY, + _CONVERSE_BODY, + _INVOKE_BODY, + ), tuple(request.body for request in received) + for request in received: + _assert_signed_with_the_deployment_key(request, key) + + +def test_boto3_direct_raw_model_ids_forward_full_bodies(gateway: Gateway, tmp_path: Path) -> None: + direct_raw_id: Final = f"anthropic.claude-direct-{uuid.uuid4().hex[:12]}-v1:0" + with wire_server(_full_body_peer) as wire: + config_path: Final = tmp_path / "bedrock-direct.yaml" + config_path.write_text( + yaml.safe_dump( + { + "model_list": [{"model_name": "bedrock/*", "litellm_params": {"model": "bedrock/*"}}], + "general_settings": {"supported_db_objects": []}, + } + ) + ) + aws_environment: Final = { + "AWS_ACCESS_KEY_ID": "AKIASCRIPTEDPROVIDER", + "AWS_SECRET_ACCESS_KEY": "scripted-secret", + "AWS_REGION_NAME": "us-east-1", + "AWS_BEDROCK_RUNTIME_ENDPOINT": wire.url, + } + with ( + owned_proxy( + gateway, + tmp_path, + aws_environment, + config=config_path, + remove_environment=( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_SECURITY_TOKEN", + "AWS_BEARER_TOKEN_BEDROCK", + "AWS_PROFILE", + "AWS_DEFAULT_PROFILE", + "AWS_SHARED_CREDENTIALS_FILE", + "AWS_CONFIG_FILE", + "AWS_WEB_IDENTITY_TOKEN_FILE", + "AWS_ROLE_ARN", + "AWS_CONTAINER_CREDENTIALS_FULL_URI", + "AWS_CONTAINER_CREDENTIALS_RELATIVE_URI", + "AWS_REGION", + "AWS_DEFAULT_REGION", + ), + workers=1, + ) as candidate, + candidate.scenario() as scenario, + ): + key: Final = scenario.key() + runtime: Final = _boto3_runtime(candidate, key) + direct_converse: Final = runtime.converse(modelId=direct_raw_id, **_CONVERSE_BODY) + direct_invoke: Final = runtime.invoke_model(modelId=direct_raw_id, body=json.dumps(_INVOKE_BODY)) + direct_converse_stream: Final = _events_read_before_the_upstream_finishes( + lambda: runtime.converse_stream(modelId=direct_raw_id, **_CONVERSE_BODY)["stream"] + ) + direct_invoke_stream: Final = runtime.invoke_model_with_response_stream( + modelId=direct_raw_id, body=json.dumps(_INVOKE_BODY) + ) + assert {name: direct_converse[name] for name in _CONVERSE_REPLY} == _CONVERSE_REPLY, direct_converse + assert json.loads(direct_invoke["body"].read()) == _INVOKE_REPLY, direct_invoke + assert list(direct_converse_stream) == [{kind: payload} for kind, payload in _EVENTS], ( + direct_converse_stream + ) + assert [json.loads(event["chunk"]["bytes"]) for event in direct_invoke_stream["body"]] == [_INVOKE_REPLY] + + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [ + ("POST", f"/model/{direct_raw_id}/converse"), + ("POST", f"/model/{direct_raw_id}/invoke"), + ("POST", f"/model/{direct_raw_id}/converse-stream"), + ("POST", f"/model/{direct_raw_id}/invoke-with-response-stream"), + ], received + assert tuple(json.loads(request.body) for request in received) == ( + _CONVERSE_BODY, + _INVOKE_BODY, + _CONVERSE_BODY, + _INVOKE_BODY, + ), tuple(request.body for request in received) + for request in received: + _assert_signed_with_the_deployment_key(request, key) diff --git a/tests/integration/providers/test_openai_passthrough_multipart_wire.py b/tests/integration/providers/test_openai_passthrough_multipart_wire.py new file mode 100644 index 00000000000..50c6fe3ac0e --- /dev/null +++ b/tests/integration/providers/test_openai_passthrough_multipart_wire.py @@ -0,0 +1,276 @@ +import io +import json +import threading +from collections.abc import Iterator +from email.parser import BytesParser +from email.policy import HTTP +from itertools import accumulate, dropwhile +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from openai import OpenAI +from pydantic import BaseModel + +pytestmark = pytest.mark.timeout(240) + +_PNG_A: Final = b"\x89PNG\r\n\x1a\nscripted-image-a" +_PNG_B: Final = b"\x89PNG\r\n\x1a\nscripted-image-b" +_MASK: Final = b"\x89PNG\r\n\x1a\nscripted-mask" +_AUDIO: Final = b"RIFF\x24\x00\x00\x00WAVEscripted-audio" +_JSONL: Final = b'{"custom_id":"1","method":"POST","url":"/v1/chat/completions","body":{}}\n' +_EDITED: Final = "c2NyaXB0ZWQtZWRpdA==" +_TRANSCRIPT: Final = "scripted transcript" +_IMAGE_STREAM_FRAMES: Final = ( + b'event: image_generation.partial_image\ndata: {"type":"image_generation.partial_image"}\n\n', + b'event: image_generation.completed\ndata: {"type":"image_generation.completed"}\n\n', +) +_IMAGE_STREAM_GATE: Final = threading.Event() +_FILE_OBJECT: Final = { + "id": "file-scripted", + "object": "file", + "bytes": len(_JSONL), + "created_at": 1, + "filename": "batch.jsonl", + "purpose": "batch", + "status": "processed", +} + +Part = tuple[str, str | None, str | None, bytes] + + +class _Image(BaseModel): + b64_json: str + + +class _ImageResponse(BaseModel): + created: int + data: tuple[_Image, ...] + + +def _respond(request: Request) -> Reply: + if request.target == "/v1/audio/transcriptions": + return Reply(body=_TRANSCRIPT.encode(), content_type="text/plain") + if request.target == "/v1/files": + return Reply(body=json.dumps(_FILE_OBJECT).encode()) + if request.target == "/v1/images/edits": + parts: Final = _parts(request.headers["content-type"], request.body) + if ("stream", None, None, b"true") in parts: + return Reply( + chunks=_IMAGE_STREAM_FRAMES, + content_type="text/event-stream", + gate_after_first=_IMAGE_STREAM_GATE, + ) + return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": _EDITED}]}).encode()) + + +@pytest.fixture(scope="module") +def wire() -> Iterator[Wire]: + with wire_server(_respond) as served: + yield served + + +@pytest.fixture(autouse=True) +def _drained(wire: Wire) -> None: + wire.drain() + + +@pytest.fixture(scope="module") +def proxy(wire: Wire, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("openai-passthrough-multipart") + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["environment_variables"] = {"OPENAI_API_BASE": wire.url, "OPENAI_API_KEY": "scripted"} + path: Final = directory / "openai-multipart.yaml" + path.write_text(yaml.safe_dump(config)) + with gateway_from_environment() as rig, owned_proxy(rig, directory, {}, config=path, workers=2) as owned: + yield owned + + +def _parts(content_type: str, body: bytes) -> tuple[Part, ...]: + envelope: Final = f"content-type: {content_type}\r\n\r\n".encode() + body + parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope) + assert parsed.is_multipart(), content_type + return tuple( + ( + str(part.get_param("name", header="content-disposition")), + part.get_filename(), + part.get_content_type() if part.get_filename() is not None else None, + part.get_payload(decode=True), + ) + for part in parsed.iter_parts() + ) + + +def _only_upstream(wire: Wire, target: str) -> Request: + received: Final = wire.drain() + assert [(request.method, request.target) for request in received] == [("POST", target)], received + upstream: Final = received[0] + assert upstream.headers["authorization"] == "Bearer scripted", upstream.headers + return upstream + + +def _sdk(proxy: Gateway, key: str, prefix: str) -> OpenAI: + return OpenAI( + api_key=key, + base_url=f"{str(proxy.client.base_url).rstrip('/')}{prefix}/v1", + max_retries=0, + http_client=httpx.Client(timeout=30, trust_env=False), + ) + + +def _image_edit_files( + stream: bool, image_field: str, mask_field: str +) -> list[tuple[str, tuple[str | None, str | bytes] | tuple[str, bytes, str]]]: + return [ + ("model", (None, "gpt-image-1")), + ("prompt", (None, "make them match")), + ("n", (None, "2")), + ("stream", (None, "true" if stream else "false")), + (image_field, ("a.png", _PNG_A, "image/png")), + (image_field, ("b.png", _PNG_B, "image/png")), + (mask_field, ("m.png", _MASK, "image/png")), + ] + + +def _expected_image_edit_parts(stream: bool, image_field: str, mask_field: str) -> tuple[Part, ...]: + return ( + ("model", None, None, b"gpt-image-1"), + ("prompt", None, None, b"make them match"), + ("n", None, None, b"2"), + ("stream", None, None, b"true" if stream else b"false"), + (image_field, "a.png", "image/png", _PNG_A), + (image_field, "b.png", "image/png", _PNG_B), + (mask_field, "m.png", "image/png", _MASK), + ) + + +@pytest.mark.parametrize( + ("image_field", "mask_field"), + [("image[]", "mask"), ("image", "mask"), ("image[]", "mask[]")], + ids=["bracketed-images", "plain-repeated-images", "bracketed-mask"], +) +def test_curl_image_edit_forwards_repeated_image_parts_mask_and_string_typed_fields( + proxy: Gateway, wire: Wire, image_field: str, mask_field: str +) -> None: + with proxy.scenario() as scenario: + response: Final = proxy.client.post( + "/openai/v1/images/edits", + headers={"Authorization": f"Bearer {scenario.key()}"}, + files=_image_edit_files(False, image_field, mask_field), + ) + assert response.status_code == 200, response.text + assert _ImageResponse.model_validate_json(response.content).data == (_Image(b64_json=_EDITED),), response.text + upstream: Final = _only_upstream(wire, "/v1/images/edits") + assert _parts(upstream.headers["content-type"], upstream.body) == _expected_image_edit_parts( + False, image_field, mask_field + ), upstream.body + + +def test_curl_streaming_image_edit_relays_each_sse_frame_as_the_upstream_sends_it(proxy: Gateway, wire: Wire) -> None: + pytest.skip("BUG: multipart image edits with stream=true do not relay the SSE response incrementally") + _IMAGE_STREAM_GATE.clear() + with ( + proxy.scenario() as scenario, + proxy.client.stream( + "POST", + "/openai/v1/images/edits", + headers={"Authorization": f"Bearer {scenario.key()}"}, + files=_image_edit_files(True, "image[]", "mask"), + ) as response, + ): + raw: Final = iter(response.iter_raw()) + try: + assert response.status_code == 200, response.status_code + assert response.headers.get("content-type") == "text/event-stream", response.headers + first_frame: Final = next( + dropwhile( + lambda buffered: b"\n\n" not in buffered, + accumulate(raw, lambda buffered, chunk: buffered + chunk, initial=b""), + ) + ) + assert first_frame == _IMAGE_STREAM_FRAMES[0], first_frame + finally: + _IMAGE_STREAM_GATE.set() + streamed: Final = first_frame + b"".join(raw) + assert streamed == b"".join(_IMAGE_STREAM_FRAMES), streamed + upstream: Final = _only_upstream(wire, "/v1/images/edits") + assert _parts(upstream.headers["content-type"], upstream.body) == _expected_image_edit_parts( + True, "image[]", "mask" + ) + + +def test_sdk_image_edit_with_two_images_reaches_upstream_as_two_file_parts(proxy: Gateway, wire: Wire) -> None: + with proxy.scenario() as scenario, _sdk(proxy, scenario.key(), "/openai") as sdk: + edited: Final = sdk.images.edit( + model="gpt-image-1", + image=[("a.png", io.BytesIO(_PNG_A), "image/png"), ("b.png", io.BytesIO(_PNG_B), "image/png")], + mask=("m.png", io.BytesIO(_MASK), "image/png"), + prompt="make them match", + n=2, + ) + assert [image.b64_json for image in edited.data or []] == [_EDITED], edited + upstream: Final = _only_upstream(wire, "/v1/images/edits") + assert _parts(upstream.headers["content-type"], upstream.body) == ( + ("prompt", None, None, b"make them match"), + ("model", None, None, b"gpt-image-1"), + ("n", None, None, b"2"), + ("image[]", "a.png", "image/png", _PNG_A), + ("image[]", "b.png", "image/png", _PNG_B), + ("mask", "m.png", "image/png", _MASK), + ), upstream.body + + +def test_sdk_transcription_and_file_upload_forward_the_file_and_its_text_fields(proxy: Gateway, wire: Wire) -> None: + with proxy.scenario() as scenario: + key: Final = scenario.key() + with _sdk(proxy, key, "/openai") as sdk: + transcript: Final = sdk.audio.transcriptions.create( + model="whisper-1", file=("clip.wav", io.BytesIO(_AUDIO), "audio/wav"), response_format="text" + ) + assert transcript == _TRANSCRIPT, transcript + transcription: Final = _only_upstream(wire, "/v1/audio/transcriptions") + assert _parts(transcription.headers["content-type"], transcription.body) == ( + ("model", None, None, b"whisper-1"), + ("response_format", None, None, b"text"), + ("file", "clip.wav", "audio/wav", _AUDIO), + ), transcription.body + with _sdk(proxy, key, "/openai_passthrough") as sdk: + uploaded: Final = sdk.files.create( + file=("batch.jsonl", io.BytesIO(_JSONL), "application/jsonl"), purpose="batch" + ) + assert uploaded.model_dump(exclude_unset=True) == _FILE_OBJECT, uploaded + upload: Final = _only_upstream(wire, "/v1/files") + assert _parts(upload.headers["content-type"], upload.body) == ( + ("purpose", None, None, b"batch"), + ("file", "batch.jsonl", "application/jsonl", _JSONL), + ), upload.body + + +def test_file_less_multipart_form_stays_multipart_with_every_repeated_field(proxy: Gateway, wire: Wire) -> None: + with proxy.scenario() as scenario: + response: Final = proxy.client.post( + "/openai/v1/images/edits", + headers={"Authorization": f"Bearer {scenario.key()}"}, + files=[ + ("model", (None, "gpt-image-1")), + ("prompt", (None, "no file")), + ("image_url[]", (None, "https://example.com/a.png")), + ("image_url[]", (None, "https://example.com/b.png")), + ], + ) + assert response.status_code == 200, response.text + parsed: Final = _ImageResponse.model_validate_json(response.content) + assert parsed == _ImageResponse(created=1, data=(_Image(b64_json=_EDITED),)), response.text + upstream: Final = _only_upstream(wire, "/v1/images/edits") + assert upstream.headers["content-type"].startswith("multipart/form-data; boundary="), upstream.headers + assert _parts(upstream.headers["content-type"], upstream.body) == ( + ("model", None, None, b"gpt-image-1"), + ("prompt", None, None, b"no file"), + ("image_url[]", None, None, b"https://example.com/a.png"), + ("image_url[]", None, None, b"https://example.com/b.png"), + ), upstream.body diff --git a/tests/integration/providers/test_vertex_passthrough_wire.py b/tests/integration/providers/test_vertex_passthrough_wire.py new file mode 100644 index 00000000000..eb9471bc948 --- /dev/null +++ b/tests/integration/providers/test_vertex_passthrough_wire.py @@ -0,0 +1,516 @@ +import contextlib +import datetime +import json +import socket +import socketserver +import ssl +import threading +from collections.abc import Generator, Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from itertools import accumulate, dropwhile +from pathlib import Path +from queue import Empty, SimpleQueue +from typing import Final + +import anthropic +import httpx +import pytest +import yaml +from cryptography import x509 +from cryptography.hazmat.primitives import hashes, serialization +from cryptography.hazmat.primitives.asymmetric import ec +from cryptography.x509.oid import NameOID +from google import genai +from google.genai import types +from google.oauth2.credentials import Credentials +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(240) + +_PROJECT: Final = "scripted-project" +_TOKEN: Final = "scripted-vertex-token" +_CALLER_GOOGLE_TOKEN: Final = "caller-google-token" +_GEMINI_ROUTER_NAME: Final = "vertex-gemini-router" +_REGIONAL: Final = "us-central1-aiplatform.googleapis.com:443" +_GLOBAL: Final = "aiplatform.googleapis.com:443" +_BYO: Final = "us-east5-aiplatform.googleapis.com:443" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_GEMINI_REPLY: Final[dict[str, JsonValue]] = { + "candidates": [{"content": {"parts": [{"text": "scripted vertex"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, +} +_GEMINI_SSE: Final = (f"data: {json.dumps(_GEMINI_REPLY)}\r\n\r\n".encode(),) +_GEMINI_BODY: Final[dict[str, JsonValue]] = { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "systemInstruction": {"parts": [{"text": "be terse"}]}, + "generationConfig": {"temperature": 0.2, "maxOutputTokens": 32}, +} +_GEMINI_SDK_BODY: Final[dict[str, JsonValue]] = { + **_GEMINI_BODY, + "systemInstruction": {"parts": [{"text": "be terse"}], "role": "user"}, +} +_CLAUDE_REPLY: Final[dict[str, JsonValue]] = { + "id": "msg_scripted", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "scripted claude"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, +} +_CLAUDE_EVENTS: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "message_start", "message": {**_CLAUDE_REPLY, "content": [], "stop_reason": None}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "scripted claude"}}, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 2}, + }, + {"type": "message_stop"}, +) +_CLAUDE_SSE: Final = tuple( + f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in _CLAUDE_EVENTS +) + + +_VERTEX_STREAM_GATE: Final = threading.Event() +_VERTEX_STREAM_GATE.set() + + +@dataclass(frozen=True, slots=True) +class _Tunnel: + url: str + authorities: SimpleQueue[str] + + def drain(self) -> list[str]: + drained: Final[list[str]] = [] + with contextlib.suppress(Empty): + while True: + drained.append(self.authorities.get_nowait()) + return drained + + +def _tls_context(directory: Path) -> ssl.SSLContext: + key: Final = ec.generate_private_key(ec.SECP256R1()) + now: Final = datetime.datetime.now(datetime.UTC) + name: Final = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "aiplatform.googleapis.com")]) + certificate: Final = ( + x509.CertificateBuilder() + .subject_name(name) + .issuer_name(name) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - datetime.timedelta(days=1)) + .not_valid_after(now + datetime.timedelta(days=1)) + .sign(key, hashes.SHA256()) + ) + certificate_file: Final = directory / "vertex.pem" + key_file: Final = directory / "vertex.key" + certificate_file.write_bytes(certificate.public_bytes(serialization.Encoding.PEM)) + key_file.write_bytes( + key.private_bytes(serialization.Encoding.PEM, serialization.PrivateFormat.PKCS8, serialization.NoEncryption()) + ) + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + context.load_cert_chain(certificate_file, key_file) + return context + + +def _pipe(source: socket.socket, sink: socket.socket) -> None: + with contextlib.suppress(OSError): + for chunk in iter(lambda: source.recv(65536), b""): + sink.sendall(chunk) + with contextlib.suppress(OSError): + sink.shutdown(socket.SHUT_WR) + + +@contextmanager +def _vertex_tunnel(destination: Wire) -> Generator[_Tunnel, None, None]: + authorities: Final[SimpleQueue[str]] = SimpleQueue() + destination_port: Final = int(destination.url.rsplit(":", 1)[1]) + + class Tunnel(socketserver.StreamRequestHandler): + rbufsize = 0 + request: socket.socket + + def handle(self) -> None: + authority: Final = self.rfile.readline().decode().split()[1] + while self.rfile.readline() not in (b"\r\n", b""): + pass + authorities.put(authority) + if authority not in (_REGIONAL, _GLOBAL, _BYO): + self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n") + return + self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n") + self.request.settimeout(10) + with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream: + outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream)) + outbound.start() + _pipe(upstream, self.request) + outbound.join(timeout=12) + + with socketserver.ThreadingTCPServer(("127.0.0.1", 0), Tunnel) as server: + thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05}) + thread.start() + try: + yield _Tunnel(f"http://127.0.0.1:{server.server_address[1]}", authorities) + finally: + server.shutdown() + thread.join(timeout=6) + + +def _vertex_peer(request: Request) -> Reply: + if "enerateContent" in request.target: + if ":streamGenerateContent" in request.target: + return Reply(chunks=_GEMINI_SSE, content_type="text/event-stream") + return Reply(body=json.dumps(_GEMINI_REPLY).encode()) + if ":streamRawPredict" in request.target or request.target.endswith(":rawPredict?alt=sse"): + return Reply(chunks=_CLAUDE_SSE, content_type="text/event-stream", gate_after_first=_VERTEX_STREAM_GATE) + return Reply(body=json.dumps(_CLAUDE_REPLY).encode()) + + +def _token_peer(request: Request) -> Reply: + assert request.target == "/_oauth/token", request.target + return Reply(body=json.dumps({"access_token": _TOKEN, "expires_in": 3600, "token_type": "Bearer"}).encode()) + + +@dataclass(frozen=True, slots=True) +class _Rig: + proxy: Gateway + vertex: Wire + tokens: Wire + tunnel: _Tunnel + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("vertex-passthrough") + with ( + wire_server(_vertex_peer, tls=_tls_context(directory)) as vertex, + wire_server(_token_peer) as tokens, + _vertex_tunnel(vertex) as tunnel, + ): + credentials: Final = service_account_json(_PROJECT, tokens.url) + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["model_list"] = [ + *config.get("model_list", []), + { + "model_name": _GEMINI_ROUTER_NAME, + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": _PROJECT, + "vertex_location": "us-central1", + "vertex_credentials": credentials, + "use_in_pass_through": True, + }, + }, + { + "model_name": "claude-sonnet-4-5", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-4-5", + "vertex_project": _PROJECT, + "vertex_location": "global", + "vertex_credentials": credentials, + "use_in_pass_through": True, + }, + }, + ] + path: Final = directory / "vertex-passthrough.yaml" + path.write_text(yaml.safe_dump(config)) + environment: Final = { + "HTTPS_PROXY": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "AIOHTTP_TRUST_ENV": "True", + "SSL_VERIFY": "False", + } + with ( + gateway_from_environment() as base, + owned_proxy(base, directory, environment, config=path, workers=2) as proxy, + ): + yield _Rig(proxy, vertex, tokens, tunnel) + + +@pytest.fixture(autouse=True) +def _drained(rig: _Rig) -> None: + rig.vertex.drain() + rig.tunnel.drain() + + +def _assert_minted_token_and_no_virtual_key(upstream: Request, key: str) -> None: + assert upstream.headers["authorization"] == f"Bearer {_TOKEN}", upstream.headers + assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers + assert key not in upstream.target, upstream.target + + +@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"]) +def test_google_genai_vertex_calls_reach_the_router_deployment_with_a_minted_token(rig: _Rig, alias: str) -> None: + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + sent: Final[list[httpx.Request]] = [] + client: Final = genai.Client( + vertexai=True, + project=_PROJECT, + location="us-central1", + credentials=Credentials(token=key), + http_options=types.HttpOptions( + base_url=f"{str(rig.proxy.client.base_url).rstrip('/')}{alias}", + client_args={"event_hooks": {"request": [sent.append]}, "timeout": 30, "trust_env": False}, + ), + ) + generated: Final = client.models.generate_content( + model=_GEMINI_ROUTER_NAME, + contents="hi", + config=types.GenerateContentConfig( + system_instruction="be terse", + temperature=0.2, + max_output_tokens=32, + ), + ) + streamed: Final = "".join( + chunk.text or "" + for chunk in client.models.generate_content_stream( + model=_GEMINI_ROUTER_NAME, + contents="hi", + config=types.GenerateContentConfig( + system_instruction="be terse", + temperature=0.2, + max_output_tokens=32, + ), + ) + ) + assert generated.text == "scripted vertex", generated + assert streamed == "scripted vertex", streamed + received: Final = rig.vertex.drain() + model_path: Final = f"/v1beta1/projects/{_PROJECT}/locations/us-central1/publishers/google/models" + assert [(request.method, request.target) for request in received] == [ + ("POST", f"{model_path}/gemini-2.5-flash:generateContent"), + ("POST", f"{model_path}/gemini-2.5-flash:streamGenerateContent?alt=sse"), + ], received + assert rig.tunnel.drain() == [_REGIONAL, _REGIONAL] + assert len(sent) == len(received), sent + for upstream, sdk_request in zip(received, sent, strict=True): + sdk_body: Final = _JSON_OBJECT.validate_json(sdk_request.content) + assert _JSON_OBJECT.validate_json(upstream.body) == sdk_body, upstream.body + assert sdk_body == _GEMINI_SDK_BODY, upstream.body + _assert_minted_token_and_no_virtual_key(upstream, key) + + +def test_vertex_abbreviated_beta_route_uses_router_project_and_location(rig: _Rig) -> None: + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + response: Final = rig.proxy.client.post( + f"/vertex-ai/v1beta1/publishers/google/models/{_GEMINI_ROUTER_NAME}:generateContent", + json=_GEMINI_BODY, + headers={"Authorization": f"Bearer {key}"}, + ) + assert response.status_code == 200, response.text + assert response.json() == _GEMINI_REPLY, response.text + received: Final = rig.vertex.drain() + assert [(request.method, request.target) for request in received] == [ + ( + "POST", + f"/v1beta1/projects/{_PROJECT}/locations/us-central1/publishers/google/models/" + "gemini-2.5-flash:generateContent", + ) + ], received + assert rig.tunnel.drain() == [_REGIONAL] + assert _JSON_OBJECT.validate_json(received[0].body) == _GEMINI_BODY, received[0].body + _assert_minted_token_and_no_virtual_key(received[0], key) + + +@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"]) +def test_vertex_rest_spellings_mint_token_and_add_sse_query_upstream(rig: _Rig, alias: str) -> None: + model_path: Final = f"/v1/projects/{_PROJECT}/locations/us-central1/publishers/google/models" + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + headers: Final = {"Authorization": f"Bearer {key}"} + generated: Final = rig.proxy.client.post( + f"{alias}{model_path}/{_GEMINI_ROUTER_NAME}:generateContent", json=_GEMINI_BODY, headers=headers + ) + assert generated.status_code == 200, generated.text + assert generated.json() == _GEMINI_REPLY, generated.text + with rig.proxy.client.stream( + "POST", + f"{alias}{model_path}/{_GEMINI_ROUTER_NAME}:streamGenerateContent", + json=_GEMINI_BODY, + headers=headers, + ) as streamed_response: + streamed: Final = streamed_response.read() + assert streamed_response.status_code == 200, streamed + assert streamed == b"".join(_GEMINI_SSE), streamed + received: Final = rig.vertex.drain() + assert [(request.method, request.target) for request in received] == [ + ("POST", f"{model_path}/gemini-2.5-flash:generateContent"), + ("POST", f"{model_path}/gemini-2.5-flash:streamGenerateContent?alt=sse"), + ], received + assert rig.tunnel.drain() == [_REGIONAL, _REGIONAL] + for upstream in received: + assert _JSON_OBJECT.validate_json(upstream.body) == _GEMINI_BODY, upstream.body + _assert_minted_token_and_no_virtual_key(upstream, key) + + +@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"]) +@pytest.mark.parametrize("stream", [False, True], ids=["rawPredict", "streamRawPredict"]) +def test_anthropic_vertex_sdk_partner_call_on_the_global_location_reaches_upstream_intact( + rig: _Rig, alias: str, stream: bool +) -> None: + sent: Final[list[httpx.Request]] = [] + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + client: Final = anthropic.AnthropicVertex( + base_url=f"{str(rig.proxy.client.base_url).rstrip('/')}{alias}/v1", + region="global", + project_id=_PROJECT, + access_token=key, + max_retries=0, + http_client=httpx.Client(timeout=30, trust_env=False, event_hooks={"request": [sent.append]}), + ) + arguments: Final[dict[str, JsonValue]] = { + "model": "claude-sonnet-4-5", + "max_tokens": 64, + "system": "be terse", + "messages": [{"role": "user", "content": "hi"}], + } + if stream: + with client.messages.stream(**arguments) as events: # type: ignore[arg-type] + text = events.get_final_text() + else: + text = "".join( + block.text + for block in client.messages.create(**arguments).content # type: ignore[arg-type] + if block.type == "text" + ) + client.close() + assert text == "scripted claude" + action: Final = "streamRawPredict" if stream else "rawPredict" + model_path: Final = f"/v1/projects/{_PROJECT}/locations/global/publishers/anthropic/models/claude-sonnet-4-5" + received: Final = rig.vertex.drain() + assert [(request.method, request.target) for request in received] == [ + ("POST", f"{model_path}:{action}" + ("?alt=sse" if stream else "")) + ], received + assert rig.tunnel.drain() == [_GLOBAL] + upstream: Final = received[0] + assert _JSON_OBJECT.validate_json(upstream.body) == _JSON_OBJECT.validate_json(sent[0].content), upstream.body + _assert_minted_token_and_no_virtual_key(upstream, key) + + +@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"]) +@pytest.mark.parametrize("action", ["rawPredict", "streamRawPredict"]) +def test_vertex_partner_stream_body_adds_the_sse_query_and_relays_each_frame_before_the_upstream_finishes( + rig: _Rig, alias: str, action: str +) -> None: + model_path: Final = f"/v1/projects/{_PROJECT}/locations/global/publishers/anthropic/models/claude-sonnet-4-5" + body: Final[dict[str, JsonValue]] = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 64, + "stream": True, + "messages": [{"role": "user", "content": "hi"}], + } + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + _VERTEX_STREAM_GATE.clear() + with rig.proxy.client.stream( + "POST", f"{alias}{model_path}:{action}", json=body, headers={"Authorization": f"Bearer {key}"} + ) as response: + raw: Final = iter(response.iter_raw()) + try: + assert response.status_code == 200, response.status_code + first_frame: Final = next( + dropwhile( + lambda buffered: b"\n\n" not in buffered, + accumulate(raw, lambda buffered, chunk: buffered + chunk, initial=b""), + ), + b"", + ) + assert first_frame == _CLAUDE_SSE[0], first_frame + finally: + _VERTEX_STREAM_GATE.set() + streamed: Final = first_frame + b"".join(raw) + assert streamed == b"".join(_CLAUDE_SSE), streamed + received: Final = rig.vertex.drain() + assert [(request.method, request.target) for request in received] == [ + ("POST", f"{model_path}:{action}?alt=sse") + ], received + assert rig.tunnel.drain() == [_GLOBAL] + assert _JSON_OBJECT.validate_json(received[0].body) == body, received[0].body + _assert_minted_token_and_no_virtual_key(received[0], key) + + +def test_credential_less_vertex_project_forwards_the_caller_google_token_and_strips_the_virtual_key( + rig: _Rig, +) -> None: + model_path: Final = "/v1/projects/byo-project/locations/us-east5/publishers/anthropic/models/claude-byo" + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + body: Final[dict[str, JsonValue]] = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + } + response: Final = rig.proxy.client.post( + f"/vertex_ai{model_path}:rawPredict", + json=body, + headers={"Authorization": f"Bearer {key}", "x-goog-api-key": _CALLER_GOOGLE_TOKEN}, + ) + assert response.status_code == 200, response.text + assert response.json() == _CLAUDE_REPLY, response.text + received: Final = rig.vertex.drain() + assert [(request.method, request.target) for request in received] == [("POST", f"{model_path}:rawPredict")], ( + received + ) + assert rig.tunnel.drain() == [_BYO] + upstream: Final = received[0] + assert upstream.headers["x-goog-api-key"] == _CALLER_GOOGLE_TOKEN, upstream.headers + assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers + assert "authorization" not in upstream.headers, upstream.headers + assert "x-litellm-api-key" not in upstream.headers, upstream.headers + assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body + + +@pytest.mark.parametrize("alias", ["/vertex_ai", "/vertex-ai"]) +def test_credential_less_vertex_project_forwards_the_caller_bearer_token_when_the_virtual_key_rides_in_x_litellm_api_key( + rig: _Rig, alias: str +) -> None: + model_path: Final = "/v1/projects/byo-project/locations/us-east5/publishers/anthropic/models/claude-byo" + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + second_key: Final = scenario.key() + body: Final[dict[str, JsonValue]] = { + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 64, + "messages": [{"role": "user", "content": "hi"}], + } + response: Final = rig.proxy.client.post( + f"{alias}{model_path}:rawPredict", + json=body, + headers={ + "x-litellm-api-key": f"Bearer {key}", + "Authorization": f"Bearer {_CALLER_GOOGLE_TOKEN}", + "api-key": second_key, + "x-caller-echo": key, + "x-caller-trace": "trace-1", + }, + ) + assert response.status_code == 200, response.text + assert response.json() == _CLAUDE_REPLY, response.text + received: Final = rig.vertex.drain() + assert [(request.method, request.target) for request in received] == [("POST", f"{model_path}:rawPredict")], ( + received + ) + assert rig.tunnel.drain() == [_BYO] + upstream: Final = received[0] + assert upstream.headers["authorization"] == f"Bearer {_CALLER_GOOGLE_TOKEN}", upstream.headers + assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers + assert {name: value for name, value in upstream.headers.items() if second_key in value} == {}, upstream.headers + assert "x-litellm-api-key" not in upstream.headers, upstream.headers + assert "x-goog-api-key" not in upstream.headers, upstream.headers + assert upstream.headers["x-caller-trace"] == "trace-1", upstream.headers + assert key not in upstream.target, upstream.target + assert _JSON_OBJECT.validate_json(upstream.body) == body, upstream.body diff --git a/tests/integration/spend/test_passthrough_request_tags.py b/tests/integration/spend/test_passthrough_request_tags.py index e6f5e15e161..f3c0b66b20d 100644 --- a/tests/integration/spend/test_passthrough_request_tags.py +++ b/tests/integration/spend/test_passthrough_request_tags.py @@ -1,9 +1,12 @@ import json import uuid from collections.abc import Callable, Mapping +from email.parser import BytesParser +from email.policy import HTTP from hashlib import sha256 from pathlib import Path from typing import Final +from urllib.parse import parse_qsl import pytest from integration._support.client import Gateway, JsonValue, Scenario, eventually, object_value @@ -419,3 +422,362 @@ def test_anthropic_passthrough_spend_row_carries_key_team_project_tags_and_spend row ) assert _spend_logs_metadata(row) == {"cost_center": marker, "team_field": marker}, row + + +_CONTRACT_TOKEN_ENV: Final = "INTEGRATION_PASSTHROUGH_CONTRACT_TOKEN" +_CONTRACT_TOKEN: Final = "scripted-configured-endpoint-token" +_ENDPOINTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +def _received_reply(request: Request) -> Reply: + return Reply(body=json.dumps({"received": request.target}).encode()) + + +def _create_endpoint(gateway: Gateway, scenario: Scenario, body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + response: Final = gateway.request("POST", "/config/pass_through_endpoint", dict(body)) + assert response.status_code == 200, response.text + created: Final = _ENDPOINTS.validate_python(response.json()["endpoints"]) + assert len(created) == 1, response.text + endpoint_id: Final = str(created[0]["id"]) + scenario.cleanups.callback( + lambda: gateway.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}) + ) + return created[0] + + +def _readback(gateway: Gateway, endpoint_id: str) -> list[dict[str, JsonValue]]: + response: Final = gateway.request("GET", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}) + assert response.status_code == 200, response.text + return _ENDPOINTS.validate_python(response.json()["endpoints"]) + + +def _stored_endpoint_ids() -> list[JsonValue]: + rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("general_settings",)) + assert len(rows) == 1, rows + settings: Final = object_value(rows[0]["param_value"]) + stored: Final = _ENDPOINTS.validate_python(settings.get("pass_through_endpoints") or []) + return [endpoint["id"] for endpoint in stored] + + +def _stored_endpoint(endpoint_id: str) -> dict[str, JsonValue]: + rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("general_settings",)) + assert len(rows) == 1, rows + settings: Final = object_value(rows[0]["param_value"]) + stored: Final = _ENDPOINTS.validate_python(settings.get("pass_through_endpoints") or []) + matches: Final = [endpoint for endpoint in stored if endpoint.get("id") == endpoint_id] + assert len(matches) == 1, matches + return matches[0] + + +def _query(request: Request) -> tuple[str, tuple[tuple[str, str], ...]]: + path, _, query = request.target.partition("?") + return path, tuple(parse_qsl(query, keep_blank_values=True)) + + +def _multipart_parts(content_type: str, body: bytes) -> list[tuple[str, str | None, str | None, bytes]]: + parsed: Final = BytesParser(policy=HTTP).parsebytes(f"content-type: {content_type}\r\n\r\n".encode() + body) + assert parsed.is_multipart(), content_type + return [ + ( + str(part.get_param("name", header="content-disposition")), + part.get_filename(), + part.get("content-type"), + part.get_payload(decode=True), + ) + for part in parsed.iter_parts() + ] + + +def test_configured_endpoint_forwards_json_multipart_and_query_and_an_update_keeps_omitted_fields( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + path: Final = f"/integration-contract-{marker}" + subpath_root: Final = f"/integration-contract-subpath-{marker}" + with ( + wire_server(_received_reply) as wire, + owned_proxy(gateway, tmp_path, {_CONTRACT_TOKEN_ENV: _CONTRACT_TOKEN}) as candidate, + candidate.scenario() as scenario, + ): + created: Final = _create_endpoint( + candidate, + scenario, + { + "path": path, + "target": f"{wire.url}/base", + "headers": {"Authorization": f"Bearer os.environ/{_CONTRACT_TOKEN_ENV}", "x-static": "static-v1"}, + "methods": ["POST"], + "default_query_params": {"api-version": "2024-01-01"}, + "auth": False, + }, + ) + _create_endpoint( + candidate, + scenario, + {"path": subpath_root, "target": f"{wire.url}/nested", "include_subpath": True, "auth": True}, + ) + endpoint_id: Final = str(created["id"]) + body: Final[dict[str, JsonValue]] = {"input": marker, "options": {"trace": marker, "stream": False}, "n": 2} + + sent_json: Final = candidate.client.post(path, params={"trace": "1"}, json=body) + assert sent_json.status_code == 200, sent_json.text + assert sent_json.json() == {"received": "/base?api-version=2024-01-01&trace=1"}, sent_json.text + overridden: Final = candidate.client.post(path, params={"api-version": "2025-02-02"}, json=body) + assert overridden.status_code == 200, overridden.text + form: Final = candidate.client.post( + path, + files=[ + ("tag", (None, "first")), + ("tag", (None, "second")), + ("file", ("one.bin", b"\x00one", "application/octet-stream")), + ("file", ("two.bin", b"\x00two", "application/octet-stream")), + ], + ) + assert form.status_code == 200, form.text + refused_verb: Final = candidate.client.get(path) + assert refused_verb.status_code == 405, refused_verb.text + nested: Final = candidate.request("POST", f"{subpath_root}/sub/route", body, params={"q": "1"}) + assert nested.status_code == 200, nested.text + before: Final = wire.drain() + assert [(request.method, *_query(request)) for request in before] == [ + ("POST", "/base", (("api-version", "2024-01-01"), ("trace", "1"))), + ("POST", "/base", (("api-version", "2025-02-02"),)), + ("POST", "/base", (("api-version", "2024-01-01"),)), + ("POST", "/nested/sub/route", (("q", "1"),)), + ], before + for request in before[:3]: + assert request.headers["authorization"] == f"Bearer {_CONTRACT_TOKEN}", request.headers + assert request.headers["x-static"] == "static-v1", request.headers + assert [json.loads(request.body) for request in (before[0], before[1], before[3])] == [body] * 3, before + assert _multipart_parts(before[2].headers["content-type"], before[2].body) == [ + ("tag", None, None, b"first"), + ("tag", None, None, b"second"), + ("file", "one.bin", "application/octet-stream", b"\x00one"), + ("file", "two.bin", "application/octet-stream", b"\x00two"), + ], before[2].body + + retargeted: Final = candidate.request( + "POST", f"/config/pass_through_endpoint/{endpoint_id}", {"path": path, "target": f"{wire.url}/v2"} + ) + assert retargeted.status_code == 200, retargeted.text + expected: Final = {**created, "target": f"{wire.url}/v2"} + assert _ENDPOINTS.validate_python(retargeted.json()["endpoints"]) == [expected], retargeted.text + retargeted_readback: Final = _readback(candidate, endpoint_id) + assert retargeted_readback == [expected] + assert _stored_endpoint(endpoint_id) == { + key: value for key, value in retargeted_readback[0].items() if key != "is_from_config" + } + after_retarget: Final = candidate.client.post(path, json=body) + assert after_retarget.status_code == 200, after_retarget.text + assert after_retarget.json() == {"received": "/v2?api-version=2024-01-01"}, after_retarget.text + retargeted_request: Final = wire.drain() + assert [(request.method, request.target) for request in retargeted_request] == [ + ("POST", "/v2?api-version=2024-01-01") + ], retargeted_request + assert retargeted_request[0].headers["authorization"] == f"Bearer {_CONTRACT_TOKEN}", retargeted_request + + locked: Final = candidate.request( + "POST", + f"/config/pass_through_endpoint/{endpoint_id}", + {"path": path, "target": f"{wire.url}/v2", "auth": True, "headers": {"x-static": "static-v2"}}, + ) + assert locked.status_code == 200, locked.text + locked_expected: Final = {**expected, "auth": True, "headers": {"x-static": "static-v2"}} + locked_readback: Final = _readback(candidate, endpoint_id) + assert locked_readback == [locked_expected] + assert _stored_endpoint(endpoint_id) == { + key: value for key, value in locked_readback[0].items() if key != "is_from_config" + } + anonymous: Final = candidate.client.post(path, json=body) + assert anonymous.status_code == 401, anonymous.text + assert wire.drain() == () + keyed: Final = candidate.request("POST", path, body) + assert keyed.status_code == 200, keyed.text + locked_request: Final = wire.drain() + assert [(request.method, request.target) for request in locked_request] == [ + ("POST", "/v2?api-version=2024-01-01") + ], locked_request + assert locked_request[0].headers["x-static"] == "static-v2", locked_request[0].headers + assert "authorization" not in locked_request[0].headers, locked_request[0].headers + refused_after_update: Final = candidate.request("GET", path) + assert refused_after_update.status_code == 405, refused_after_update.text + keyed_form: Final = candidate.client.post( + path, + headers={"Authorization": f"Bearer {candidate.key}"}, + files=[ + ("tag", (None, "first")), + ("tag", (None, "second")), + ("file", ("one.bin", b"\x00one", "application/octet-stream")), + ("file", ("two.bin", b"\x00two", "application/octet-stream")), + ], + ) + assert keyed_form.status_code == 200, keyed_form.text + locked_form: Final = wire.drain() + assert [(request.method, request.target) for request in locked_form] == [ + ("POST", "/v2?api-version=2024-01-01") + ], locked_form + assert _multipart_parts(locked_form[0].headers["content-type"], locked_form[0].body) == [ + ("tag", None, None, b"first"), + ("tag", None, None, b"second"), + ("file", "one.bin", "application/octet-stream", b"\x00one"), + ("file", "two.bin", "application/octet-stream", b"\x00two"), + ], locked_form[0].body + + +def test_renaming_a_subpath_endpoint_moves_both_routes_keeps_stored_fields_and_never_forwards_caller_headers( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + path: Final = f"/integration-rename-{marker}" + renamed_path: Final = f"/integration-renamed-{marker}" + body: Final[dict[str, JsonValue]] = {"input": marker} + with ( + wire_server(_received_reply) as wire, + owned_proxy(gateway, tmp_path, {}, workers=1) as candidate, + candidate.scenario() as scenario, + ): + created: Final = _create_endpoint( + candidate, + scenario, + { + "path": path, + "target": f"{wire.url}/base", + "include_subpath": True, + "auth": True, + "headers": {"x-static": "v1"}, + "methods": ["POST"], + }, + ) + endpoint_id: Final = str(created["id"]) + expected: Final = { + "id": endpoint_id, + "path": path, + "target": f"{wire.url}/base", + "headers": {"x-static": "v1"}, + "default_query_params": {}, + "include_subpath": True, + "cost_per_request": 0.0, + "timeout": None, + "auth": True, + "guardrails": None, + "is_from_config": False, + "methods": ["POST"], + } + assert created == expected, created + assert _readback(candidate, endpoint_id) == [expected] + key: Final = scenario.key(allowed_passthrough_routes=[path, renamed_path]) + caller_headers: Final = {"x-caller": f"caller-{marker}"} + first: Final = candidate.request("POST", path, body, key=key, headers=caller_headers) + assert first.status_code == 200, first.text + assert first.json() == {"received": "/base"}, first.text + second: Final = candidate.request("POST", f"{path}/sub", body, key=key, headers=caller_headers) + assert second.status_code == 200, second.text + assert second.json() == {"received": "/base/sub"}, second.text + initial_requests: Final = wire.drain() + assert [(request.method, request.target) for request in initial_requests] == [ + ("POST", "/base"), + ("POST", "/base/sub"), + ], initial_requests + + updated: Final = candidate.request( + "POST", + f"/config/pass_through_endpoint/{endpoint_id}", + {"path": renamed_path, "target": f"{wire.url}/base"}, + ) + assert updated.status_code == 200, updated.text + updated_expected: Final = {**expected, "path": renamed_path} + assert _ENDPOINTS.validate_python(updated.json()["endpoints"]) == [updated_expected], updated.text + updated_readback: Final = _readback(candidate, endpoint_id) + assert updated_readback == [updated_expected] + assert _stored_endpoint(endpoint_id) == { + key: value for key, value in updated_readback[0].items() if key != "is_from_config" + } + + old_path: Final = candidate.request("POST", path, body, key=key, headers=caller_headers) + old_subpath: Final = candidate.request("POST", f"{path}/sub", body, key=key, headers=caller_headers) + assert old_path.status_code == 404, old_path.text + assert old_subpath.status_code == 404, old_subpath.text + assert wire.drain() == () + + renamed: Final = candidate.request("POST", renamed_path, body, key=key, headers=caller_headers) + assert renamed.status_code == 200, renamed.text + assert renamed.json() == {"received": "/base"}, renamed.text + renamed_subpath: Final = candidate.request("POST", f"{renamed_path}/sub", body, key=key, headers=caller_headers) + assert renamed_subpath.status_code == 200, renamed_subpath.text + assert renamed_subpath.json() == {"received": "/base/sub"}, renamed_subpath.text + refused_verbs: Final = tuple( + candidate.request("GET", renamed_path, key=key, headers=caller_headers).status_code for _ in range(4) + ) + assert refused_verbs == (405, 405, 405, 405), refused_verbs + anonymous: Final = candidate.client.post(f"{renamed_path}/sub", json=body, headers=caller_headers) + assert anonymous.status_code == 401, anonymous.text + renamed_requests: Final = wire.drain() + assert [(request.method, request.target) for request in (*initial_requests, *renamed_requests)] == [ + ("POST", "/base"), + ("POST", "/base/sub"), + ("POST", "/base"), + ("POST", "/base/sub"), + ], renamed_requests + for upstream in (*initial_requests, *renamed_requests): + assert upstream.headers["x-static"] == "v1", upstream.headers + assert "x-caller" not in upstream.headers, upstream.headers + assert {name: value for name, value in upstream.headers.items() if key in value} == {}, upstream.headers + assert key not in upstream.target, upstream.target + assert wire.drain() == () + + +def test_configured_endpoint_without_auth_serves_its_subpaths_without_a_key(gateway: Gateway) -> None: + pytest.skip("BUG: an auth=false include_subpath pass-through endpoint answers 401 on every subpath request") + marker: Final = uuid.uuid4().hex + path: Final = f"/integration-open-subpath-{marker}" + with wire_server(_received_reply) as wire, gateway.scenario() as scenario: + _create_endpoint( + gateway, scenario, {"path": path, "target": f"{wire.url}/open", "include_subpath": True, "auth": False} + ) + response: Final = gateway.client.post(f"{path}/sub", json={"input": marker}) + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/open/sub")] + + +def test_deleted_configured_endpoint_leaves_readback_storage_and_runtime_and_team_get_filters( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + with ( + wire_server(_received_reply) as wire, + owned_proxy(gateway, tmp_path, {}) as candidate, + candidate.scenario() as scenario, + ): + allowed: Final = _create_endpoint( + candidate, scenario, {"path": f"/integration-allowed-{marker}", "target": f"{wire.url}/allowed"} + ) + hidden: Final = _create_endpoint( + candidate, scenario, {"path": f"/integration-hidden-{marker}", "target": f"{wire.url}/hidden"} + ) + team: Final = scenario.team(metadata={"allowed_passthrough_routes": [allowed["path"]]}) + team_view: Final = candidate.request("GET", f"/config/pass_through_endpoint/team/{team}") + assert team_view.status_code == 200, team_view.text + assert [endpoint["path"] for endpoint in _ENDPOINTS.validate_python(team_view.json()["endpoints"])] == [ + allowed["path"] + ], team_view.text + + served: Final = candidate.request("POST", str(hidden["path"]), {"input": marker}) + assert served.status_code == 200, served.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/hidden")] + + deleted: Final = candidate.request( + "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(hidden["id"])} + ) + assert deleted.status_code == 200, deleted.text + assert _ENDPOINTS.validate_python(deleted.json()["endpoints"]) == [hidden], deleted.text + assert _readback(candidate, str(hidden["id"])) == [] + listed: Final = candidate.request("GET", "/config/pass_through_endpoint") + assert listed.status_code == 200, listed.text + listed_ids: Final = [endpoint["id"] for endpoint in _ENDPOINTS.validate_python(listed.json()["endpoints"])] + assert hidden["id"] not in listed_ids and allowed["id"] in listed_ids, listed.text + stored_ids: Final = _stored_endpoint_ids() + assert hidden["id"] not in stored_ids and allowed["id"] in stored_ids, stored_ids + + refused: Final = candidate.request("POST", str(hidden["path"]), {"input": marker}) + assert refused.status_code == 404, refused.text + assert wire.drain() == ()