diff --git a/tests/integration/compatibility/test_a2a_wire_versions.py b/tests/integration/compatibility/test_a2a_wire_versions.py index 7a828ba2487..2823911ace3 100644 --- a/tests/integration/compatibility/test_a2a_wire_versions.py +++ b/tests/integration/compatibility/test_a2a_wire_versions.py @@ -3,7 +3,6 @@ import uuid from typing import Final import pytest - from integration._support.client import Gateway from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server @@ -115,3 +114,103 @@ def test_a2a_versions_and_legacy_casing_preserve_real_wire_and_response(gateway: actual: Final = wire.drain() assert len(tuple(item for item in actual if item.method == "POST")) == 1 assert any(item.method == "GET" for item in actual) + + +@pytest.mark.covers("compatibility.a2a.versioned_card_path_agent_is_reached_with_bearer_and_blocking_send") +def test_agent_serving_its_card_only_at_versioned_path_is_reached_with_bearer_and_answers(gateway: Gateway) -> None: + marker: Final = "foundry" + uuid.uuid4().hex + bearer: Final = "Bearer synthetic-entra-" + marker + + def upstream(request: Request) -> Reply: + assert request.headers.get("authorization") == bearer, request.headers + if request.method == "GET": + if request.target != "/agentCard/v1.0": + return Reply(status=404, body=json.dumps({"error": "not found"}).encode()) + return Reply( + body=json.dumps( + { + "protocolVersion": "0.3", + "name": marker, + "description": "Synthetic prompt agent", + "version": "1.0.0", + "url": wire.url + "/", + "capabilities": {"streaming": False}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + ).encode() + ) + assert request.method == "POST" and request.target == "/", request.target + body: Final = json.loads(request.body) + assert body["jsonrpc"] == "2.0" and body["method"] == "message/send", body + message: Final = body["params"]["message"] + assert message["kind"] == "message" and message["role"] == "user", message + assert message["parts"] == [{"kind": "text", "text": "synthetic ping"}], message + return Reply( + body=json.dumps( + { + "jsonrpc": "2.0", + "id": body["id"], + "result": { + "kind": "message", + "role": "agent", + "messageId": marker + "-out", + "parts": [{"kind": "text", "text": "synthetic pong"}], + }, + } + ).encode() + ) + + with wire_server(upstream) as wire, gateway.scenario() as scenario: + card: Final = { + "protocolVersion": "0.3", + "name": marker, + "description": "Synthetic prompt agent", + "version": "1.0.0", + "url": wire.url + "/", + "capabilities": {"streaming": False}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + created: Final = gateway.request( + "POST", + "/v1/agents", + {"agent_name": marker, "agent_card_params": card, "static_headers": {"Authorization": bearer}}, + ) + assert created.status_code == 200, created.text + identity: Final = created.json()["agent_id"] + + def cleanup() -> None: + deleted: Final = gateway.request("DELETE", f"/v1/agents/{identity}") + assert deleted.status_code == 200, deleted.text + assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (identity,)) == [] + + scenario.cleanups.callback(cleanup) + response: Final = gateway.request( + "POST", + f"/a2a/{identity}", + { + "jsonrpc": "2.0", + "id": marker, + "method": "message/send", + "params": { + "message": { + "kind": "message", + "role": "user", + "messageId": marker + "-in", + "parts": [{"kind": "text", "text": "synthetic ping"}], + } + }, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["jsonrpc"] == "2.0" and body["id"] == marker and "error" not in body, response.text + assert body["result"]["kind"] == "message", response.text + assert body["result"]["messageId"] == marker + "-out", response.text + assert body["result"]["parts"] == [{"kind": "text", "text": "synthetic pong"}], response.text + actual: Final = wire.drain() + assert tuple(item.target for item in actual if item.method == "GET")[-1] == "/agentCard/v1.0", actual + assert tuple(item.target for item in actual if item.method == "POST") == ("/",), actual diff --git a/tests/integration/providers/test_anthropic_advisor_wire.py b/tests/integration/providers/test_anthropic_advisor_wire.py index 77fa27cd2a9..b2f44d8f155 100644 --- a/tests/integration/providers/test_anthropic_advisor_wire.py +++ b/tests/integration/providers/test_anthropic_advisor_wire.py @@ -1,28 +1,35 @@ import json import uuid +from pathlib import Path from typing import Final import pytest +import yaml from integration._support.client import Gateway +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server _ADVISOR_KEY: Final = "synthetic-advisor-key" +_PROXY_ANTHROPIC_KEY: Final = "sk-proxy-owned-anthropic-secret" _QUESTION: Final = "which index should this query use" _ADVICE: Final = "use the composite index on (tenant_id, created_at)" _FINAL_ANSWER: Final = "done, the composite index is the right one" -_ADVISOR_CALL_MESSAGE: Final = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": "advisor-call", - "type": "function", - "function": {"name": "advisor", "arguments": json.dumps({"question": _QUESTION})}, - } - ], -} +def _advisor_call_message(question: str) -> dict[str, object]: + return { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "advisor-call", + "type": "function", + "function": {"name": "advisor", "arguments": json.dumps({"question": question})}, + } + ], + } + + _FINAL_MESSAGE: Final = {"role": "assistant", "content": _FINAL_ANSWER} @@ -41,7 +48,7 @@ def _chat_completion(identity: str, message: dict[str, object], finish_reason: s ) -def _executor_reply(body: dict[str, object], identity: str) -> Reply: +def _executor_reply(body: dict[str, object], identity: str, question: str) -> Reply: messages: Final = body["messages"] assert isinstance(messages, list) if any(message.get("role") == "tool" for message in messages): @@ -50,7 +57,7 @@ def _executor_reply(body: dict[str, object], identity: str) -> Reply: tools: Final = body["tools"] assert isinstance(tools, list) assert tools[0]["function"]["name"] == "advisor" - return _chat_completion(identity, _ADVISOR_CALL_MESSAGE, "tool_calls") + return _chat_completion(identity, _advisor_call_message(question), "tool_calls") @pytest.mark.covers("providers.anthropic_messages_advisor.sub_call_uses_the_configured_advisor_deployment") @@ -58,18 +65,20 @@ def test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_ gateway: Gateway, ) -> None: identity: Final = "advisor-wire-" + uuid.uuid4().hex + migration: Final = "please plan the migration " + identity + question: Final = _QUESTION + " " + identity def respond(request: Request) -> Reply: body: Final = json.loads(request.body) if request.target == "/v1/chat/completions": assert request.headers["authorization"] == "Bearer integration-provider-key" - return _executor_reply(body, identity) + return _executor_reply(body, identity, question) assert request.target == "/v1/messages" assert request.headers["x-api-key"] == _ADVISOR_KEY assert body["model"] == "claude-opus-4-1-20250805" assert body["messages"] == [ - {"role": "user", "content": "please plan the migration"}, - {"role": "user", "content": _QUESTION}, + {"role": "user", "content": migration}, + {"role": "user", "content": question}, ] assert "tools" not in body return Reply( @@ -88,7 +97,7 @@ def test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_ ) with wire_server(respond) as wire, gateway.scenario() as scenario: - executor: Final = scenario.model(model="hosted_vllm/llama-3.3-70b", api_base=wire.url + "/v1") + executor: Final = scenario.model(model="hosted_vllm/gpt-4o-mini", api_base=wire.url + "/v1") advisor: Final = scenario.model( model="anthropic/claude-opus-4-1-20250805", api_base=wire.url, api_key=_ADVISOR_KEY ) @@ -98,7 +107,7 @@ def test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_ { "model": executor, "max_tokens": 64, - "messages": [{"role": "user", "content": "please plan the migration"}], + "messages": [{"role": "user", "content": migration}], "tools": [{"type": "advisor_20260301", "name": "advisor", "model": advisor}], }, ) @@ -111,3 +120,82 @@ def test_advisor_sub_call_reaches_the_router_deployment_with_its_key_instead_of_ "/v1/messages", "/v1/chat/completions", ] + + +def _advice_reply(identity: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": f"msg-{identity}", + "type": "message", + "role": "assistant", + "model": "claude-opus-4-1-20250805", + "content": [{"type": "text", "text": _ADVICE}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + } + ).encode() + ) + + +@pytest.mark.covers("providers.anthropic_messages_advisor.caller_api_base_without_api_key_never_receives_the_proxy_key") +def test_advisor_api_base_without_api_key_is_rejected_before_the_proxy_anthropic_key_reaches_the_caller_host( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "advisor-leak-" + uuid.uuid4().hex + question: Final = _QUESTION + " " + identity + + def executor(request: Request) -> Reply: + assert request.target == "/v1/chat/completions", request.target + return _executor_reply(json.loads(request.body), identity, question) + + def caller_host(request: Request) -> Reply: + return _advice_reply(identity) + + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"]["allow_client_side_credentials"] = True + path: Final = tmp_path / "client-side-credentials.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + wire_server(executor) as executor_wire, + wire_server(caller_host) as caller_wire, + owned_proxy(gateway, tmp_path, {"ANTHROPIC_API_KEY": _PROXY_ANTHROPIC_KEY}, config=path) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model="hosted_vllm/gpt-4o-mini", api_base=executor_wire.url + "/v1") + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": "please plan the migration"}], + "tools": [ + { + "type": "advisor_20260301", + "name": "advisor", + "model": "anthropic/claude-opus-4-1-20250805", + "api_base": caller_wire.url, + } + ], + }, + ) + received: Final = caller_wire.drain() + assert [ + (request.target, request.headers.get("x-api-key"), json.loads(request.body)["messages"]) + for request in received + ] == [], response.text + assert response.is_error, response.text + assert response.json() == { + "type": "error", + "error": { + "type": "api_error", + "message": ( + "advisor tool definition sets 'api_base' without 'api_key'. A caller-supplied api_base is only " + "honored alongside a caller-supplied api_key, so the proxy's own credentials are never sent to a " + "caller-chosen destination." + ), + }, + }, response.text + assert executor_wire.drain() == (), response.text diff --git a/tests/integration/providers/test_anthropic_legacy_thinking_budget_wire.py b/tests/integration/providers/test_anthropic_legacy_thinking_budget_wire.py new file mode 100644 index 00000000000..242e5c7ec5a --- /dev/null +++ b/tests/integration/providers/test_anthropic_legacy_thinking_budget_wire.py @@ -0,0 +1,77 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +_MODEL: Final = "claude-sonnet-4-6" +_KEY: Final = "synthetic-anthropic-key" +_THINKING: Final = {"type": "enabled", "budget_tokens": 8000} +_TOOL: Final = { + "name": "read_file", + "description": "read a file", + "input_schema": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}, +} +_NEXT_CALL: Final = {"type": "tool_use", "id": "call-2", "name": "read_file", "input": {"path": "schema.prisma"}} + + +def _tool_loop_history(identity: str) -> tuple[dict[str, object], ...]: + return ( + {"role": "user", "content": f"open the config for {identity}"}, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "call-1", "name": "read_file", "input": {"path": "config.yaml"}}], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call-1", "content": "model_list: []"}]}, + ) + + +def _tool_use_reply(identity: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": f"msg-{identity}", + "type": "message", + "role": "assistant", + "model": _MODEL, + "content": [_NEXT_CALL], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 40, "output_tokens": 12}, + } + ).encode() + ) + + +@pytest.mark.covers("providers.anthropic_messages.claude_4_6_legacy_thinking_budget_reaches_the_wire_unchanged") +def test_claude_4_6_thinking_budget_tokens_on_messages_is_forwarded_instead_of_rewritten_to_adaptive( + gateway: Gateway, +) -> None: + identity: Final = "legacy-thinking-" + uuid.uuid4().hex + history: Final = _tool_loop_history(identity) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == _KEY + body: Final = json.loads(request.body) + assert body["thinking"] == _THINKING, body + assert "output_config" not in body, body + assert body["max_tokens"] == 32768, body + assert body["messages"] == list(history), body + assert body["tools"] == [_TOOL], body + return _tool_use_reply(identity) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 32768, "thinking": _THINKING, "messages": history, "tools": [_TOOL]}, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["content"] == [_NEXT_CALL], response.text + assert body["stop_reason"] == "tool_use", response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_anthropic_messages_fireworks_stop_wire.py b/tests/integration/providers/test_anthropic_messages_fireworks_stop_wire.py new file mode 100644 index 00000000000..adec5784aa8 --- /dev/null +++ b/tests/integration/providers/test_anthropic_messages_fireworks_stop_wire.py @@ -0,0 +1,65 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "accounts/fireworks/models/glm-5p3" +_API_KEY: Final = "synthetic-fireworks-key" +_STOP: Final = "" +_ANSWER: Final = "allow" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@pytest.mark.covers( + "providers.anthropic_messages_adapter.stop_sequences_and_disabled_thinking_reach_openai_compatible_provider_as_stop_and_reasoning_effort" +) +def test_messages_stop_sequences_to_fireworks_are_sent_as_stop_not_stop_sequences(gateway: Gateway) -> None: + prompt: Final = "classify this tool call " + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert "stop_sequences" not in body, body + assert body["stop"] == [_STOP], body + assert body["reasoning_effort"] == "none", body + assert body["model"] == _MODEL, body + assert body["messages"] == [{"role": "user", "content": prompt}], body + return Reply( + body=json.dumps( + { + "id": "fw-classifier", + "object": "chat.completion", + "created": 1, + "model": _MODEL, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _ANSWER}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 6, "total_tokens": 15}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"fireworks_ai/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": prompt}], + "stop_sequences": [_STOP], + "thinking": {"type": "disabled"}, + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["content"] == [{"type": "text", "text": _ANSWER}], response.text + assert payload["stop_reason"] == "end_turn", response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")] diff --git a/tests/integration/providers/test_anthropic_messages_openai_bridge_wire.py b/tests/integration/providers/test_anthropic_messages_openai_bridge_wire.py new file mode 100644 index 00000000000..72eba1d89a5 --- /dev/null +++ b/tests/integration/providers/test_anthropic_messages_openai_bridge_wire.py @@ -0,0 +1,82 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gpt-5.4-mini" +_API_KEY: Final = "synthetic-openai-key" +_CORRECTION: Final = "Stop refactoring the parser and only fix the failing test instead." +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _responses_reply(identity: str, content: str) -> bytes: + return json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": _BACKEND, + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": content, "annotations": []}], + } + ], + "usage": {"input_tokens": 41, "output_tokens": 5, "total_tokens": 46}, + } + ).encode() + + +@pytest.mark.covers("providers.anthropic_messages_openai_bridge.midturn_system_correction_reaches_the_wire") +def test_midturn_system_correction_is_forwarded_to_openai_responses(gateway: Gateway) -> None: + identity: Final = f"openai-midturn-system-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == "/responses" + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _BACKEND + assert body["instructions"] == "You are a coding agent." + assert body["input"] == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Fix the failing test."}]}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "I will start by refactoring the parser."}], + }, + {"type": "message", "role": "system", "content": [{"type": "input_text", "text": _CORRECTION}]}, + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Continue."}]}, + ], body + return Reply(body=_responses_reply(identity, "Understood.")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "system": "You are a coding agent.", + "messages": [ + {"role": "user", "content": "Fix the failing test."}, + {"role": "assistant", "content": "I will start by refactoring the parser."}, + {"role": "system", "content": _CORRECTION}, + {"role": "user", "content": "Continue."}, + ], + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["content"] == [{"type": "text", "text": "Understood."}], response.text + assert payload["stop_reason"] == "end_turn", response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/integration/providers/test_anthropic_messages_openai_tools_wire.py b/tests/integration/providers/test_anthropic_messages_openai_tools_wire.py new file mode 100644 index 00000000000..605fa45e17b --- /dev/null +++ b/tests/integration/providers/test_anthropic_messages_openai_tools_wire.py @@ -0,0 +1,92 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +_BACKEND: Final = "gpt-5.4-mini" +_API_KEY: Final = "synthetic-openai-key" +_TOOL_SCHEMA: Final = { + "type": "object", + "properties": { + "city": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + "include_forecast": {"type": "boolean"}, + }, + "required": ["city"], +} + + +@pytest.mark.covers("providers.anthropic_messages_bridge.optional_tool_properties_stay_optional_on_the_wire") +def test_messages_tool_with_optional_properties_reaches_openai_responses_non_strict(gateway: Gateway) -> None: + identity: Final = f"messages-optional-tool-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + body: Final = json.loads(request.body) + assert body["model"] == _BACKEND, body + assert body["tools"] == [ + { + "type": "function", + "name": "get_weather", + "strict": False, + "description": "Current weather for a city", + "parameters": _TOOL_SCHEMA, + } + ], body["tools"] + return Reply( + body=json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": _BACKEND, + "output": [ + { + "type": "function_call", + "id": f"fc_{identity}", + "call_id": f"call_{identity}", + "name": "get_weather", + "arguments": json.dumps({"city": "Paris"}), + "status": "completed", + } + ], + "usage": {"input_tokens": 30, "output_tokens": 9, "total_tokens": 39}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": "What is the weather in Paris?"}], + "tools": [ + { + "name": "get_weather", + "description": "Current weather for a city", + "input_schema": _TOOL_SCHEMA, + } + ], + }, + ) + assert response.status_code == 200, response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] + body: Final = response.json() + assert body["stop_reason"] == "tool_use", response.text + assert body["content"] == [ + { + "type": "tool_use", + "id": f"call_{identity}", + "name": "get_weather", + "input": {"city": "Paris"}, + } + ], response.text diff --git a/tests/integration/providers/test_anthropic_messages_timeout_wire.py b/tests/integration/providers/test_anthropic_messages_timeout_wire.py new file mode 100644 index 00000000000..76c29ad7763 --- /dev/null +++ b/tests/integration/providers/test_anthropic_messages_timeout_wire.py @@ -0,0 +1,58 @@ +import json +import time +import uuid +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway +from integration._support.wire import Reply, Request, wire_server + +_UPSTREAM_STALL_SECONDS: Final = 4.0 +_CONFIGURED_TIMEOUT_SECONDS: Final = 1.0 + + +@pytest.mark.covers("providers.anthropic_messages.configured_timeout_aborts_stalled_upstream") +def test_messages_endpoint_honors_configured_timeout_against_stalled_upstream(gateway: Gateway) -> None: + prompt: Final = "stall-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages" + assert request.headers["x-api-key"] == "synthetic-anthropic-key" + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["model"] == "claude-sonnet-4-5-20250929" + assert body["messages"] == [{"role": "user", "content": prompt}] + assert body["max_tokens"] == 16 + assert "timeout" not in body + time.sleep(_UPSTREAM_STALL_SECONDS) + return Reply( + body=json.dumps( + { + "id": "msg_stalled", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "too late"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 2}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", + api_base=wire.url, + api_key="synthetic-anthropic-key", + timeout=_CONFIGURED_TIMEOUT_SECONDS, + ) + started: Final = time.monotonic() + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + ) + elapsed: Final = time.monotonic() - started + assert response.status_code == 408, response.text + assert elapsed < _UPSTREAM_STALL_SECONDS, f"timed out only after {elapsed:.2f}s: {response.text}" + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_anthropic_thinking_signature_retry_wire.py b/tests/integration/providers/test_anthropic_thinking_signature_retry_wire.py new file mode 100644 index 00000000000..e414d8f0d11 --- /dev/null +++ b/tests/integration/providers/test_anthropic_thinking_signature_retry_wire.py @@ -0,0 +1,95 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "claude-sonnet-4-5-20250929" +KEY: Final = "synthetic-anthropic-key" +SIGNATURE_ERROR: Final = json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "messages.2.content.0.thinking.signature.str: Input should be a valid string", + }, + } +).encode() +TOOLS: Final = ({"name": "lookup", "input_schema": {"type": "object", "properties": {"key": {"type": "string"}}}},) + + +def _history_with_unsigned_thinking(identity: str) -> tuple[dict[str, object], ...]: + return ( + {"role": "user", "content": [{"type": "text", "text": f"first question {identity}"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "first answer"}]}, + {"role": "user", "content": [{"type": "text", "text": "second question"}]}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "replayed from another provider", "signature": None}, + {"type": "tool_use", "id": "call-1", "name": "lookup", "input": {"key": "value"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call-1", "content": "found"}]}, + ) + + +@pytest.mark.covers("providers.anthropic_messages.missing_thinking_signature_400_retries_without_thinking_blocks") +def test_missing_thinking_signature_400_retries_once_without_thinking_blocks_and_returns_200( + gateway: Gateway, +) -> None: + identity: Final = "thinking-signature-" + uuid.uuid4().hex + history: Final = _history_with_unsigned_thinking(identity) + tool_use_only_turn: Final = { + "role": "assistant", + "content": [{"type": "tool_use", "id": "call-1", "name": "lookup", "input": {"key": "value"}}], + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages" + assert request.headers["x-api-key"] == KEY + body: Final = json.loads(request.body) + assert body["model"] == MODEL + assert body["tools"] == list(TOOLS), body + if body["messages"][3]["content"][0]["type"] == "thinking": + assert body["messages"] == list(history), body + assert body["thinking"] == {"type": "enabled", "budget_tokens": 1024}, body + return Reply(status=400, body=SIGNATURE_ERROR) + assert body["messages"] == [*history[:3], tool_use_only_turn, history[4]], body + assert "thinking" not in body, body + return Reply( + body=json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": MODEL, + "content": [{"type": "text", "text": "recovered without thinking history"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 30, "output_tokens": 6}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{MODEL}", api_base=wire.url, api_key=KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "tools": list(TOOLS), + "messages": list(history), + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["id"] == identity, response.text + assert body["content"] == [{"type": "text", "text": "recovered without thinking history"}], response.text + assert body["stop_reason"] == "end_turn", response.text + assert [request.target for request in wire.drain()] == ["/v1/messages", "/v1/messages"] diff --git a/tests/integration/providers/test_anthropic_wire.py b/tests/integration/providers/test_anthropic_wire.py index 27895440712..7e9c5be227a 100644 --- a/tests/integration/providers/test_anthropic_wire.py +++ b/tests/integration/providers/test_anthropic_wire.py @@ -1,4 +1,5 @@ import json +import time import uuid from typing import Final @@ -8,10 +9,17 @@ from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server -@pytest.mark.covers("other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates") +@pytest.mark.covers( + "other.provider_wire.anthropic.tool_history_system_cache_and_internal_fields", + "quota_management.spend_tracking.cache_tokens.disjoint_classes_use_explicit_rates", +) def test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contracts(gateway: Gateway) -> None: identity: Final = "anthropic-wire-" + uuid.uuid4().hex - tool_schema: Final = {"type": "object", "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}}, "required": ["x", "y"]} + tool_schema: Final = { + "type": "object", + "properties": {"x": {"type": "integer"}, "y": {"type": "integer"}}, + "required": ["x", "y"], + } def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == "/v1/messages" @@ -21,27 +29,80 @@ def test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contra assert body["system"] == [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}] assert body["tools"][0]["name"] == "add" and body["tools"][0]["input_schema"] == tool_schema assert body["max_tokens"] == 16 - assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "rpm", "tpm"}.intersection(body) + assert not {"timeout", "stream_chunk_size", "litellm_params", "litellm_metadata", "rpm", "tpm"}.intersection( + body + ) messages: Final = body["messages"] assert [message["role"] for message in messages] == ["user", "assistant", "user"] assert messages[0]["content"] == [{"type": "text", "text": "first"}] - assert messages[1]["content"] == [{"type": "tool_use", "id": "history-call", "name": "add", "input": {"x": 1, "y": 2}}] - assert messages[2]["content"] == [{"type": "tool_result", "tool_use_id": "history-call", "content": "3"}, {"type": "text", "text": "next"}] - return Reply(body=json.dumps({"id": identity, "type": "message", "role": "assistant", "model": "claude-sonnet-4-5-20250929", "content": [{"type": "tool_use", "id": "next-call", "name": "add", "input": {"x": 3, "y": 4}}], "stop_reason": "tool_use", "stop_sequence": None, "usage": {"input_tokens": 10, "output_tokens": 4, "cache_read_input_tokens": 5, "cache_creation_input_tokens": 7}}).encode()) + assert messages[1]["content"] == [ + {"type": "tool_use", "id": "history-call", "name": "add", "input": {"x": 1, "y": 2}} + ] + assert messages[2]["content"] == [ + {"type": "tool_result", "tool_use_id": "history-call", "content": "3"}, + {"type": "text", "text": "next"}, + ] + return Reply( + body=json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "tool_use", "id": "next-call", "name": "add", "input": {"x": 3, "y": 4}}], + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": { + "input_tokens": 10, + "output_tokens": 4, + "cache_read_input_tokens": 5, + "cache_creation_input_tokens": 7, + }, + } + ).encode() + ) with wire_server(respond) as wire, gateway.scenario() as scenario: - model: Final = scenario.model(model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key", input_cost_per_token=0.001, output_cost_per_token=0.002, cache_read_input_token_cost=0.0001, cache_creation_input_token_cost=0.002) - response: Final = gateway.request("POST", "/v1/chat/completions", { - "model": model, "max_tokens": 16, "timeout": 5, - "messages": [ - {"role": "system", "content": [{"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}}]}, - {"role": "user", "content": "first"}, - {"role": "assistant", "tool_calls": [{"id": "history-call", "type": "function", "function": {"name": "add", "arguments": '{"x":1,"y":2}'}}]}, - {"role": "tool", "tool_call_id": "history-call", "content": "3"}, - {"role": "user", "content": "next"}, - ], - "tools": [{"type": "function", "function": {"name": "add", "parameters": tool_schema}}], - }) + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", + api_base=wire.url, + api_key="synthetic-anthropic-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + cache_read_input_token_cost=0.0001, + cache_creation_input_token_cost=0.002, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 16, + "timeout": 5, + "messages": [ + { + "role": "system", + "content": [ + {"type": "text", "text": "synthetic policy", "cache_control": {"type": "ephemeral"}} + ], + }, + {"role": "user", "content": "first"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "history-call", + "type": "function", + "function": {"name": "add", "arguments": '{"x":1,"y":2}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "history-call", "content": "3"}, + {"role": "user", "content": "next"}, + ], + "tools": [{"type": "function", "function": {"name": "add", "parameters": tool_schema}}], + }, + ) assert response.status_code == 200, response.text body: Final = response.json() assert body["id"].startswith("chatcmpl-") @@ -51,7 +112,14 @@ def test_anthropic_tool_history_and_cache_tokens_keep_wire_and_accounting_contra assert json.loads(tool["function"]["arguments"]) == {"x": 3, "y": 4} assert body["usage"]["prompt_tokens"] == 22 and body["usage"]["completion_tokens"] == 4 assert len(wire.drain()) == 1 - rows: Final = eventually(lambda: read_rows('SELECT spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (body["id"],)), lambda values: len(values) == 1, seconds=70) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (body["id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) assert float(rows[0]["spend"]) == pytest.approx(10 * 0.001 + 5 * 0.0001 + 7 * 0.002 + 4 * 0.002) assert rows[0]["prompt_tokens"] == 22 and rows[0]["completion_tokens"] == 4 metadata: Final = rows[0]["metadata"] @@ -81,3 +149,58 @@ def test_anthropic_bare_string_content_item_is_rejected_as_client_error_before_t ) assert response.status_code == 400, response.text assert wire.drain() == () + + +@pytest.mark.covers("other.provider_wire.anthropic.messages_request_timeout_reaches_transport") +def test_anthropic_messages_slow_upstream_is_cut_off_at_the_deployment_request_timeout(gateway: Gateway) -> None: + identity: Final = "anthropic-timeout-" + uuid.uuid4().hex + prompt: Final = f"slow answer {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages" + assert request.headers["x-api-key"] == "synthetic-anthropic-key" + body: Final = json.loads(request.body) + assert body["model"] == "claude-sonnet-4-5-20250929" + assert body["max_tokens"] == 16 + assert body["messages"] == [{"role": "user", "content": prompt}] + assert not { + "timeout", + "request_timeout", + "stream_chunk_size", + "litellm_params", + "litellm_metadata", + "rpm", + "tpm", + }.intersection(body) + time.sleep(1.5) + return Reply( + body=json.dumps( + { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": "late"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-5-20250929", + api_base=wire.url, + api_key="synthetic-anthropic-key", + request_timeout=0.3, + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]}, + headers={"anthropic-version": "2023-06-01"}, + ) + assert response.status_code == 408, response.text + assert "Timeout" in response.json()["error"]["message"], response.text + assert eventually(wire.drain, lambda requests: len(requests) == 1, seconds=5, return_last_on_timeout=True) diff --git a/tests/integration/providers/test_bedrock_converse_client_metadata_wire.py b/tests/integration/providers/test_bedrock_converse_client_metadata_wire.py new file mode 100644 index 00000000000..92d9ebd4a0c --- /dev/null +++ b/tests/integration/providers/test_bedrock_converse_client_metadata_wire.py @@ -0,0 +1,43 @@ +import json +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from integration.providers.test_bedrock_auth_wire import MODEL, RESPONSE, TOKEN + +ANTHROPIC_BETA: Final = ["interleaved-thinking-2025-05-14"] +CLIENT_METADATA: Final = {"originator": "codex_cli_rs", "version": "0.1.0", "session_id": "synthetic-session"} + + +def converse_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse" + body: Final = json.loads(request.body) + assert body["additionalModelRequestFields"] == {"anthropic_beta": ANTHROPIC_BETA}, body + return Reply(body=RESPONSE) + + +@pytest.mark.covers("providers.bedrock_converse.client_metadata_is_not_forwarded_in_additional_model_request_fields") +def test_client_metadata_is_dropped_from_converse_body_while_anthropic_beta_is_kept(gateway: Gateway) -> None: + with wire_server(converse_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "synthetic codex request"}], + "max_tokens": 16, + "anthropic_beta": ANTHROPIC_BETA, + "client_metadata": CLIENT_METADATA, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control" + assert len(wire.drain()) == 1, response.text diff --git a/tests/integration/providers/test_bedrock_deepseek_reasoning_wire.py b/tests/integration/providers/test_bedrock_deepseek_reasoning_wire.py new file mode 100644 index 00000000000..f7391da54d9 --- /dev/null +++ b/tests/integration/providers/test_bedrock_deepseek_reasoning_wire.py @@ -0,0 +1,93 @@ +import json +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +R1_MODEL: Final = "bedrock/converse/us.deepseek.r1-v1:0" +V3_MODEL: Final = "bedrock/converse/deepseek.v3.2" +TOKEN: Final = "synthetic-bedrock-bearer" +RESPONSE: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "deepseek reasoning wire control"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 5, "totalTokens": 14}, + "metrics": {"latencyMs": 1}, + } +).encode() + + +def r1_converse_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/model/us.deepseek.r1-v1%3A0/converse", request.target + assert request.headers["authorization"] == f"Bearer {TOKEN}" + body: Final = json.loads(request.body) + assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic r1 request"}]}] + assert body["inferenceConfig"] == {"maxTokens": 16}, body + assert body.get("additionalModelRequestFields") is None, body + return Reply(body=RESPONSE) + + +def v3_converse_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/model/deepseek.v3.2/converse", request.target + assert request.headers["authorization"] == f"Bearer {TOKEN}" + body: Final = json.loads(request.body) + assert body["messages"] == [{"role": "user", "content": [{"text": "synthetic v3 request"}]}] + assert body["inferenceConfig"] == {"maxTokens": 16}, body + assert body["additionalModelRequestFields"] == {"reasoning_effort": "high"}, body + return Reply(body=RESPONSE) + + +@pytest.mark.covers("providers.bedrock_converse.deepseek_r1_drops_thinking_and_reasoning_effort_before_provider") +def test_deepseek_r1_thinking_and_reasoning_effort_are_dropped_instead_of_leaking_into_converse( + gateway: Gateway, +) -> None: + with wire_server(r1_converse_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=R1_MODEL, + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + drop_params=True, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "synthetic r1 request"}], + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "reasoning_effort": "high", + "max_tokens": 16, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "deepseek reasoning wire control", response.text + assert response.json()["usage"]["total_tokens"] == 14, response.text + assert len(wire.drain()) == 1 + + +@pytest.mark.covers("providers.bedrock_converse.deepseek_v3_reasoning_effort_reaches_provider_raw") +def test_deepseek_v3_reasoning_effort_reaches_converse_raw_instead_of_as_anthropic_thinking( + gateway: Gateway, +) -> None: + with wire_server(v3_converse_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=V3_MODEL, api_key=TOKEN, aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "synthetic v3 request"}], + "reasoning_effort": "high", + "max_tokens": 16, + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "deepseek reasoning wire control", response.text + assert response.json()["usage"]["total_tokens"] == 14, response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_bedrock_invoke_cache_usage_wire.py b/tests/integration/providers/test_bedrock_invoke_cache_usage_wire.py new file mode 100644 index 00000000000..2638c3a2c8d --- /dev/null +++ b/tests/integration/providers/test_bedrock_invoke_cache_usage_wire.py @@ -0,0 +1,86 @@ +import json +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 + +MODEL_ID: Final = "us.amazon.nova-pro-v1:0" +TOKEN: Final = "synthetic-bedrock-bearer" +PROMPT: Final = "summarize the cached policy" +INPUT_TOKENS: Final = 11 +OUTPUT_TOKENS: Final = 4 +CACHE_READ_TOKENS: Final = 900 +CACHE_WRITE_TOKENS: Final = 300 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +CACHE_READ_RATE: Final = 0.0001 +CACHE_WRITE_RATE: Final = 0.0015 +RESPONSE: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "cached policy summary"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": INPUT_TOKENS, + "outputTokens": OUTPUT_TOKENS, + "totalTokens": INPUT_TOKENS + OUTPUT_TOKENS, + "cacheReadInputTokenCount": CACHE_READ_TOKENS, + "cacheWriteInputTokenCount": CACHE_WRITE_TOKENS, + }, + } +).encode() + + +def nova_invoke_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == f"/model/{MODEL_ID}/invoke", request.target + assert request.headers["authorization"] == f"Bearer {TOKEN}" + body: Final = json.loads(request.body) + assert body["messages"] == [{"role": "user", "content": [{"text": PROMPT}]}], body + return Reply(body=RESPONSE) + + +@pytest.mark.covers("providers.bedrock_invoke.count_suffixed_cache_usage_fields_are_reported_and_charged") +def test_nova_invoke_count_suffixed_cache_usage_fields_are_reported_and_charged(gateway: Gateway) -> None: + with wire_server(nova_invoke_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{MODEL_ID}", + api_key=TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + cache_read_input_token_cost=CACHE_READ_RATE, + cache_creation_input_token_cost=CACHE_WRITE_RATE, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 32, "messages": [{"role": "user", "content": PROMPT}]}, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["choices"][0]["message"]["content"] == "cached policy summary", response.text + usage: Final = body["usage"] + assert usage["prompt_tokens"] == INPUT_TOKENS + CACHE_READ_TOKENS + CACHE_WRITE_TOKENS, response.text + assert usage["completion_tokens"] == OUTPUT_TOKENS, response.text + assert usage["prompt_tokens_details"]["cached_tokens"] == CACHE_READ_TOKENS, response.text + assert usage["cache_read_input_tokens"] == CACHE_READ_TOKENS, response.text + assert usage["cache_creation_input_tokens"] == CACHE_WRITE_TOKENS, response.text + expected_cost: Final = ( + INPUT_TOKENS * INPUT_RATE + + CACHE_READ_TOKENS * CACHE_READ_RATE + + CACHE_WRITE_TOKENS * CACHE_WRITE_RATE + + OUTPUT_TOKENS * OUTPUT_RATE + ) + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_cost), response.text + assert len(wire.drain()) == 1 + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (body["id"],) + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert float(rows[0]["spend"]) == pytest.approx(expected_cost), rows + assert rows[0]["prompt_tokens"] == INPUT_TOKENS + CACHE_READ_TOKENS + CACHE_WRITE_TOKENS, rows diff --git a/tests/integration/providers/test_bedrock_knowledge_base_user_context_wire.py b/tests/integration/providers/test_bedrock_knowledge_base_user_context_wire.py new file mode 100644 index 00000000000..540b384d694 --- /dev/null +++ b/tests/integration/providers/test_bedrock_knowledge_base_user_context_wire.py @@ -0,0 +1,81 @@ +import json +import uuid +from collections.abc import Callable +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +ACCESS_KEY: Final = "AKIAINTEGRATION000003" +USER_CONTEXT: Final = {"userId": "reader@example.com"} +QUERY: Final = "synthetic knowledge base question" +RETRIEVE_RESPONSE: Final = json.dumps( + { + "retrievalResults": [ + { + "content": {"text": "permitted document text"}, + "score": 0.87, + "metadata": { + "x-amz-bedrock-kb-source-uri": "s3://synthetic-bucket/permitted.pdf", + "x-amz-bedrock-kb-chunk-id": "chunk-1", + }, + } + ] + } +).encode() + + +def retrieve_peer(knowledge_base_id: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == f"/knowledgebases/{knowledge_base_id}/retrieve", ( + request.target + ) + assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={ACCESS_KEY}/") + assert json.loads(request.body) == { + "retrievalQuery": {"text": QUERY}, + "retrievalConfiguration": {"vectorSearchConfiguration": {"numberOfResults": 3}}, + "userContext": USER_CONTEXT, + }, request.body + return Reply(body=RETRIEVE_RESPONSE) + + return respond + + +@pytest.mark.covers("providers.bedrock_knowledge_base.search_forwards_user_context_to_retrieve") +def test_vector_store_search_user_context_reaches_bedrock_retrieve_body(gateway: Gateway) -> None: + knowledge_base_id: Final = f"KB{uuid.uuid4().hex[:8].upper()}" + with wire_server(retrieve_peer(knowledge_base_id)) as wire, gateway.scenario() as scenario: + gateway.post( + "/vector_store/new", + { + "vector_store_id": knowledge_base_id, + "custom_llm_provider": "bedrock", + "litellm_params": { + "aws_region_name": "us-east-1", + "aws_access_key_id": ACCESS_KEY, + "aws_secret_access_key": "synthetic-knowledge-base-secret-key", + "aws_bedrock_runtime_endpoint": wire.url, + }, + }, + ) + scenario.cleanups.callback(gateway.post, "/vector_store/delete", {"vector_store_id": knowledge_base_id}) + response: Final = gateway.request( + "POST", + f"/v1/vector_stores/{knowledge_base_id}/search", + {"query": QUERY, "max_num_results": 3, "userContext": USER_CONTEXT}, + ) + assert response.status_code == 200, response.text + assert response.json()["data"] == [ + { + "score": 0.87, + "content": [{"text": "permitted document text", "type": "text"}], + "file_id": "s3://synthetic-bucket/permitted.pdf", + "filename": "permitted.pdf", + "attributes": { + "x-amz-bedrock-kb-source-uri": "s3://synthetic-bucket/permitted.pdf", + "x-amz-bedrock-kb-chunk-id": "chunk-1", + }, + } + ], response.text + assert len(wire.drain()) == 1 diff --git a/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py b/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py index bb7961160dc..aa66e82475b 100644 --- a/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_codex_input_wire.py @@ -86,3 +86,56 @@ def test_codex_agent_message_context_compaction_and_local_shell_call_reach_mantl forwarded: Final = wire.drain() assert len(forwarded) == 1, forwarded assert JSON_OBJECT.validate_json(forwarded[0].body)["input"] == expected_input, forwarded[0].body + + +SHELL_TOOL: Final[JsonValue] = { + "type": "function", + "name": "shell", + "description": "run a shell command", + "parameters": {"type": "object", "properties": {"command": {"type": "string"}}, "required": ["command"]}, +} +APPLY_PATCH_TOOL: Final[JsonValue] = { + "type": "function", + "name": "apply_patch", + "description": "apply a diff", + "parameters": {"type": "object", "properties": {"patch": {"type": "string"}}, "required": ["patch"]}, +} + + +@pytest.mark.covers("providers.bedrock_mantle.codex_additional_tools_input_item_is_hoisted_to_top_level_tools") +def test_codex_additional_tools_input_item_reaches_mantle_as_top_level_tools(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + expected_input: Final[list[JsonValue]] = [user_turn(f"hoist tools {marker}")] + expected_tools: Final[list[JsonValue]] = [SHELL_TOOL, APPLY_PATCH_TOOL] + + def mantle_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + assert request.headers["authorization"] == f"Bearer {TOKEN}" + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["model"] == "openai.gpt-5.6-sol", body + assert body["input"] == expected_input, body["input"] + assert body["tools"] == expected_tools, body + return Reply(body=RESPONSE) + + with wire_server(mantle_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=MODEL, api_key=TOKEN, api_base=wire.url, aws_region_name="us-east-2") + response: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "input": [ + {"type": "additional_tools", "role": "developer", "tools": [APPLY_PATCH_TOOL]}, + user_turn(f"hoist tools {marker}"), + ], + "tools": [SHELL_TOOL], + "store": False, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["output"][0]["content"][0]["text"] == "mantle wire control", response.text + forwarded: Final = wire.drain() + assert len(forwarded) == 1, forwarded + forwarded_body: Final = JSON_OBJECT.validate_json(forwarded[0].body) + assert forwarded_body["input"] == expected_input, forwarded[0].body + assert forwarded_body["tools"] == expected_tools, forwarded[0].body diff --git a/tests/integration/providers/test_bedrock_mantle_responses_wire.py b/tests/integration/providers/test_bedrock_mantle_responses_wire.py index 9bc6f83f8e4..3bd83b5019b 100644 --- a/tests/integration/providers/test_bedrock_mantle_responses_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_responses_wire.py @@ -104,3 +104,43 @@ def test_codex_agent_message_compaction_and_local_shell_items_are_rewritten_for_ } ], response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/openai/v1/responses")] + + +_MANTLE_MIN_MAX_OUTPUT_TOKENS: Final = 16 + + +def _mantle_peer_expecting_max_output_tokens(marker: str, expected: int) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["max_output_tokens"] == expected, request.body.decode() + assert body["input"] == f"clamp probe {marker}", request.body.decode() + return Reply(body=_RESPONSE) + + return respond + + +@pytest.mark.covers("providers.bedrock_mantle.max_output_tokens_below_minimum_is_clamped_to_16_on_the_wire") +def test_max_output_tokens_below_mantle_minimum_is_raised_to_16_before_reaching_mantle(gateway: Gateway) -> None: + marker: Final = uuid4().hex + peer: Final = _mantle_peer_expecting_max_output_tokens(marker, _MANTLE_MIN_MAX_OUTPUT_TOKENS) + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_TOKEN, aws_region_name="us-east-1") + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": f"clamp probe {marker}", "max_output_tokens": 5, "stream": False}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["status"] == "completed", response.text + assert payload["output"] == [ + { + **_OUTPUT_MESSAGE, + "phase": None, + "content": [ + {"type": "output_text", "text": "mantle wire control", "annotations": [], "logprobs": None} + ], + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/openai/v1/responses")] diff --git a/tests/integration/providers/test_bedrock_mantle_wire.py b/tests/integration/providers/test_bedrock_mantle_wire.py index 48cd0d770aa..ec32fe5a578 100644 --- a/tests/integration/providers/test_bedrock_mantle_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_wire.py @@ -1,5 +1,7 @@ import json +from collections.abc import Callable from typing import Final +from uuid import uuid4 import pytest from integration._support.client import Gateway @@ -51,3 +53,142 @@ def test_bedrock_mantle_context_overflow_returns_400_saying_prompt_is_too_long(g assert isinstance(message, str), response.text assert f"prompt is too long: {_PROMPT_TOKENS} tokens > {_MODEL_MAXIMUM} maximum" in message, response.text assert [(request.method, request.target) for request in wire.drain()] == [("POST", _RESPONSES_PATH)] + + +_ACCESS_KEY: Final = "AKIAINTEGRATION000003" +_SIGV4_PROMPT: Final = "synthetic sigv4 bridge control" +_SIGV4_RESPONSE: Final = json.dumps( + { + "id": "resp_synthetic_mantle_sigv4", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": _BACKEND, + "output": [ + { + "type": "message", + "id": "msg_synthetic_mantle_sigv4", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "mantle sigv4 wire control", "annotations": []}], + } + ], + "usage": {"input_tokens": 21, "output_tokens": 4, "total_tokens": 25}, + } +).encode() + + +def _sigv4_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == _RESPONSES_PATH, request.target + assert request.headers["authorization"].startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), dict( + request.headers + ) + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _BACKEND, body + assert _SIGV4_PROMPT in json.dumps(body["input"]), body + return Reply(body=_SIGV4_RESPONSE) + + +@pytest.mark.covers("providers.bedrock_mantle.chat_bridge_keeps_deployment_aws_credentials_for_sigv4") +def test_chat_completions_bridge_signs_mantle_responses_request_with_deployment_aws_keys(gateway: Gateway) -> None: + with wire_server(_sigv4_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock_mantle/{_BACKEND}", + api_base=wire.url, + api_key=None, + aws_access_key_id=_ACCESS_KEY, + aws_secret_access_key="synthetic-secret-key-for-testing", + aws_region_name="us-east-1", + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": _SIGV4_PROMPT}]}, + ) + assert response.status_code == 200, response.text + body: Final = _JSON_OBJECT.validate_json(response.content) + choices: Final = body["choices"] + assert isinstance(choices, list) and len(choices) == 1, response.text + choice: Final = choices[0] + assert isinstance(choice, dict), response.text + assert choice["message"] == {"role": "assistant", "content": "mantle sigv4 wire control"}, response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", _RESPONSES_PATH)] + + +_CLAUDE_BACKEND: Final = "anthropic.claude-sonnet-5-v1:0" +_MESSAGES_PATH: Final = "/anthropic/v1/messages" +_STREAM_EVENTS: Final = ( + ( + "message_start", + { + "message": { + "id": "msg_mantle_stream", + "type": "message", + "role": "assistant", + "model": _CLAUDE_BACKEND, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 11, "output_tokens": 1}, + } + }, + ), + ("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": "mantle "}}), + ("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": "stream control"}}), + ("content_block_stop", {"index": 0}), + ("message_delta", {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 4}}), + ("message_stop", {}), +) +_STREAM_FRAMES: Final = tuple( + f"event: {kind}\ndata: {json.dumps({'type': kind, **payload})}\n\n".encode() for kind, payload in _STREAM_EVENTS +) + + +def _streaming_messages_peer(prompt: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == _MESSAGES_PATH + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _CLAUDE_BACKEND, body + assert body["stream"] is True, body + assert body["messages"] == [{"role": "user", "content": prompt}], body + return Reply(content_type="text/event-stream", chunks=_STREAM_FRAMES) + + return respond + + +@pytest.mark.covers("providers.bedrock_mantle.messages_stream_sends_stream_true_and_relays_sse_events") +def test_bedrock_mantle_messages_stream_relays_anthropic_sse_instead_of_failing_on_event_stream_decode( + gateway: Gateway, +) -> None: + prompt: Final = f"synthetic mantle stream control {uuid4().hex}" + with wire_server(_streaming_messages_peer(prompt)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock_mantle/{_CLAUDE_BACKEND}", api_base=wire.url, api_key=_API_KEY, aws_region_name="us-east-1" + ) + with gateway.client.stream( + "POST", + "/v1/messages", + json={ + "model": model, + "max_tokens": 64, + "stream": True, + "messages": [{"role": "user", "content": prompt}], + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read().decode() + assert response.headers["content-type"].startswith("text/event-stream"), dict(response.headers) + events: Final = tuple( + _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.iter_lines() + if line.startswith("data: ") + ) + assert tuple(event["type"] for event in events) == tuple(kind for kind, _ in _STREAM_EVENTS), events + assert ( + "".join(str(event["delta"]["text"]) for event in events if event["type"] == "content_block_delta") + == "mantle stream control" + ), events + assert [(request.method, request.target) for request in wire.drain()] == [("POST", _MESSAGES_PATH)] diff --git a/tests/integration/providers/test_bedrock_marengo_embed_3_wire.py b/tests/integration/providers/test_bedrock_marengo_embed_3_wire.py new file mode 100644 index 00000000000..9cdedef6ef6 --- /dev/null +++ b/tests/integration/providers/test_bedrock_marengo_embed_3_wire.py @@ -0,0 +1,35 @@ +import json +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0" +TOKEN: Final = "synthetic-bedrock-bearer" +INPUT: Final = "hello world" +VECTOR: Final = [0.1, 0.2, 0.3] +RESPONSE: Final = json.dumps({"data": [{"embedding": VECTOR}]}).encode() + + +def marengo_3_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target == "/model/us.twelvelabs.marengo-embed-3-0-v1%3A0/invoke", request.target + assert request.headers["authorization"] == f"Bearer {TOKEN}" + assert json.loads(request.body) == {"inputType": "text", "text": {"inputText": INPUT}}, request.body + return Reply(body=RESPONSE) + + +@pytest.mark.covers("providers.bedrock_embedding.marengo_3_text_input_reaches_bedrock_nested_under_input_type") +def test_marengo_3_text_embedding_nests_input_text_under_input_type(gateway: Gateway) -> None: + with wire_server(marengo_3_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_key=TOKEN, + api_base=wire.url, + aws_region_name="us-east-1", + ) + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": INPUT}) + assert response.status_code == 200, response.text + assert response.json()["data"] == [{"object": "embedding", "index": 0, "embedding": VECTOR}], response.text + assert len(wire.drain()) == 1, "the embedding request never reached Bedrock" diff --git a/tests/integration/providers/test_bedrock_role_configuration.py b/tests/integration/providers/test_bedrock_role_configuration.py index ac8edbdfde0..857535e5e35 100644 --- a/tests/integration/providers/test_bedrock_role_configuration.py +++ b/tests/integration/providers/test_bedrock_role_configuration.py @@ -7,7 +7,6 @@ from urllib.parse import parse_qs import pytest import yaml - from integration._support.client import Gateway from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -30,7 +29,10 @@ def test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock(gatew assert parameters["RoleArn"] == [role] assert parameters["RoleSessionName"][0] in {"integration-yaml-session", "integration-db-session"} result = f"{assumed_key}synthetic-assumed-secret-key-for-testing{assumed_token}2035-01-01T00:00:00Zarn:aws:sts::123456789012:assumed-role/integration/sessionintegration:session0" - return Reply(content_type="text/xml", body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}synthetic-sts-request'.encode()) + return Reply( + content_type="text/xml", + body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}synthetic-sts-request'.encode(), + ) def bedrock(request: Request) -> Reply: assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse" @@ -41,8 +43,11 @@ def test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock(gatew with wire_server(sts) as authority, wire_server(bedrock) as provider: parameters: Final = { - "model": MODEL, "aws_region_name": "us-east-1", "aws_role_name": "os.environ/INTEGRATION_ROLE_ARN", - "aws_session_name": "integration-yaml-session", "aws_bedrock_runtime_endpoint": provider.url, + "model": MODEL, + "aws_region_name": "us-east-1", + "aws_role_name": "os.environ/INTEGRATION_ROLE_ARN", + "aws_session_name": "integration-yaml-session", + "aws_bedrock_runtime_endpoint": provider.url, "aws_sts_endpoint": authority.url, } alias: Final = "integration-role-yaml-" + uuid.uuid4().hex @@ -53,23 +58,144 @@ def test_role_reference_from_db_and_yaml_reaches_real_sts_http_and_bedrock(gatew empty: Final = tmp_path / "empty-aws-config" empty.write_text("") overrides: Final = { - "INTEGRATION_ROLE_ARN": role, "AWS_ACCESS_KEY_ID": "AKIAINTEGRATION000001", "AWS_SECRET_ACCESS_KEY": "synthetic-source-secret-key-for-testing", - "AWS_CONFIG_FILE": str(empty), "AWS_SHARED_CREDENTIALS_FILE": str(empty), "AWS_EC2_METADATA_DISABLED": "true", - "AWS_ENDPOINT_URL_STS": authority.url, "AWS_DEFAULT_REGION": "us-east-1", "LITELLM_RUST": "false", + "INTEGRATION_ROLE_ARN": role, + "AWS_ACCESS_KEY_ID": "AKIAINTEGRATION000001", + "AWS_SECRET_ACCESS_KEY": "synthetic-source-secret-key-for-testing", + "AWS_CONFIG_FILE": str(empty), + "AWS_SHARED_CREDENTIALS_FILE": str(empty), + "AWS_EC2_METADATA_DISABLED": "true", + "AWS_ENDPOINT_URL_STS": authority.url, + "AWS_DEFAULT_REGION": "us-east-1", + "LITELLM_RUST": "false", } - with owned_proxy(gateway, tmp_path, overrides, config=path, remove_environment=tuple(name for name in os.environ if name.startswith("AWS_"))) as candidate, candidate.scenario() as scenario: - database_model: Final = scenario.model(**{**parameters, "api_key": None, "aws_session_name": "integration-db-session"}) + with ( + owned_proxy( + gateway, + tmp_path, + overrides, + config=path, + remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")), + ) as candidate, + candidate.scenario() as scenario, + ): + database_model: Final = scenario.model( + **{**parameters, "api_key": None, "aws_session_name": "integration-db-session"} + ) for generation in range(2): for model in (alias, database_model): - response: Final = candidate.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "synthetic role request"}], "cache": {"no-cache": True}}) + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "synthetic role request"}], + "cache": {"no-cache": True}, + }, + ) assert response.status_code == 200, response.text assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control" assert response.json()["usage"]["total_tokens"] == 15 assert len(provider.drain()) == 1 if generation == 0: - target: Final = next(entry for entry in candidate.get("/model/info")["data"] if entry["model_name"] == database_model) - response: Final = candidate.request("PATCH", f"/model/{target['model_info']['id']}/update", {"model_info": {"description": "role reload"}}) + target: Final = next( + entry for entry in candidate.get("/model/info")["data"] if entry["model_name"] == database_model + ) + response: Final = candidate.request( + "PATCH", + f"/model/{target['model_info']['id']}/update", + {"model_info": {"description": "role reload"}}, + ) assert response.status_code == 200, response.text - assumed: Final = tuple(parse_qs(request.body.decode()) for request in authority.drain() if parse_qs(request.body.decode())["Action"] == ["AssumeRole"]) - assert {entry["RoleSessionName"][0] for entry in assumed} == {"integration-yaml-session", "integration-db-session"} + assumed: Final = tuple( + parse_qs(request.body.decode()) + for request in authority.drain() + if parse_qs(request.body.decode())["Action"] == ["AssumeRole"] + ) + assert {entry["RoleSessionName"][0] for entry in assumed} == { + "integration-yaml-session", + "integration-db-session", + } assert all(entry["RoleArn"] == [role] for entry in assumed) + + +@pytest.mark.covers("providers.bedrock_assume_role.repeat_requests_reuse_cached_sts_session_per_session_name") +def test_repeat_requests_under_one_session_name_assume_role_once_per_session_name( + gateway: Gateway, tmp_path: Path +) -> None: + role: Final = "arn:aws:iam::123456789012:role/integration-" + uuid.uuid4().hex + assumed_key: Final = "ASIAINTEGRATION000002" + assumed_token: Final = "synthetic-cached-session-token" + first_session: Final = "integration-attributed-user-a-" + uuid.uuid4().hex[:8] + second_session: Final = "integration-attributed-user-b-" + uuid.uuid4().hex[:8] + + def sts(request: Request) -> Reply: + parameters: Final = parse_qs(request.body.decode()) + action: Final = parameters["Action"][0] + assert request.method == "POST" and action in {"GetCallerIdentity", "AssumeRole"} + if action == "GetCallerIdentity": + result = "arn:aws:iam::123456789012:user/integration-sourceintegration-source123456789012" + else: + assert parameters["RoleArn"] == [role] + result = f"{assumed_key}synthetic-assumed-secret-key-for-testing{assumed_token}2035-01-01T00:00:00Zarn:aws:sts::123456789012:assumed-role/integration/sessionintegration:session0" + return Reply( + content_type="text/xml", + body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}synthetic-sts-request'.encode(), + ) + + def bedrock(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/model/anthropic.claude-3-haiku-20240307-v1%3A0/converse" + assert f"Credential={assumed_key}/" in request.headers["authorization"] + assert request.headers["x-amz-security-token"] == assumed_token + return Reply(body=RESPONSE) + + with wire_server(sts) as authority, wire_server(bedrock) as provider: + empty: Final = tmp_path / "empty-aws-config" + empty.write_text("") + overrides: Final = { + "AWS_ACCESS_KEY_ID": "AKIAINTEGRATION000002", + "AWS_SECRET_ACCESS_KEY": "synthetic-source-secret-key-for-testing", + "AWS_CONFIG_FILE": str(empty), + "AWS_SHARED_CREDENTIALS_FILE": str(empty), + "AWS_EC2_METADATA_DISABLED": "true", + "AWS_ENDPOINT_URL_STS": authority.url, + "AWS_DEFAULT_REGION": "us-east-1", + "LITELLM_RUST": "false", + } + with ( + owned_proxy( + gateway, + tmp_path, + overrides, + remove_environment=tuple(name for name in os.environ if name.startswith("AWS_")), + ) as candidate, + candidate.scenario() as scenario, + ): + parameters: Final = { + "model": MODEL, + "api_key": None, + "aws_region_name": "us-east-1", + "aws_role_name": role, + "aws_bedrock_runtime_endpoint": provider.url, + "aws_sts_endpoint": authority.url, + } + first_model: Final = scenario.model(**{**parameters, "aws_session_name": first_session}) + second_model: Final = scenario.model(**{**parameters, "aws_session_name": second_session}) + for model in (first_model, first_model, second_model, second_model): + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "synthetic cached role request"}], + "cache": {"no-cache": True}, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["choices"][0]["message"]["content"] == "bedrock wire control" + assert len(provider.drain()) == 1 + assumed: Final = tuple( + parse_qs(request.body.decode()) + for request in authority.drain() + if parse_qs(request.body.decode())["Action"] == ["AssumeRole"] + ) + assert tuple(entry["RoleSessionName"][0] for entry in assumed) == (first_session, second_session), assumed diff --git a/tests/integration/providers/test_bedrock_thinking_tokens_wire.py b/tests/integration/providers/test_bedrock_thinking_tokens_wire.py index 074adeb41f6..19d7b1e291c 100644 --- a/tests/integration/providers/test_bedrock_thinking_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_thinking_tokens_wire.py @@ -1,4 +1,5 @@ import json +import uuid from typing import Final import pytest @@ -34,13 +35,13 @@ _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) _JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) -def redacted_thinking_peer(request: Request) -> Reply: +def redacted_thinking_peer(request: Request, prompts: tuple[str, str]) -> Reply: assert request.method == "POST" and request.target == "/model/global.anthropic.claude-opus-4-8/converse" assert request.headers["authorization"] == f"Bearer {TOKEN}" body: Final = json.loads(request.body) assert body["messages"] in ( - [{"role": "user", "content": [{"text": PROMPT}]}], - [{"role": "user", "content": [{"text": RESPONSES_PROMPT}]}], + [{"role": "user", "content": [{"text": prompts[0]}]}], + [{"role": "user", "content": [{"text": prompts[1]}]}], ), body assert body["additionalModelRequestFields"]["thinking"]["type"] == "adaptive", body return Reply(body=RESPONSE) @@ -48,7 +49,9 @@ def redacted_thinking_peer(request: Request) -> Reply: @pytest.mark.covers("other.provider_wire.bedrock.hidden_thinking_tokens_are_not_reported_as_text") def test_bedrock_redacted_thinking_is_not_reported_as_zero_reasoning_tokens(gateway: Gateway) -> None: - with wire_server(redacted_thinking_peer) as wire, gateway.scenario() as scenario: + identity: Final = " " + uuid.uuid4().hex + prompts: Final = (PROMPT + identity, RESPONSES_PROMPT + identity) + with wire_server(lambda request: redacted_thinking_peer(request, prompts)) as wire, gateway.scenario() as scenario: model: Final = scenario.model( model=MODEL, api_key=TOKEN, aws_region_name="us-east-1", aws_bedrock_runtime_endpoint=wire.url ) @@ -57,7 +60,7 @@ def test_bedrock_redacted_thinking_is_not_reported_as_zero_reasoning_tokens(gate "/v1/chat/completions", { "model": model, - "messages": [{"role": "user", "content": PROMPT}], + "messages": [{"role": "user", "content": prompts[0]}], "max_tokens": 4000, "reasoning_effort": "max", }, @@ -76,7 +79,7 @@ def test_bedrock_redacted_thinking_is_not_reported_as_zero_reasoning_tokens(gate responses: Final = gateway.request( "POST", "/v1/responses", - {"model": model, "input": RESPONSES_PROMPT, "max_output_tokens": 4000, "reasoning": {"effort": "max"}}, + {"model": model, "input": prompts[1], "max_output_tokens": 4000, "reasoning": {"effort": "max"}}, ) assert responses.status_code == 200, responses.text responses_body: Final = _JSON_OBJECT.validate_json(responses.content) diff --git a/tests/integration/providers/test_fireworks_ai_session_affinity_wire.py b/tests/integration/providers/test_fireworks_ai_session_affinity_wire.py new file mode 100644 index 00000000000..9c8490a4591 --- /dev/null +++ b/tests/integration/providers/test_fireworks_ai_session_affinity_wire.py @@ -0,0 +1,69 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_MODEL: Final = "accounts/fireworks/models/kimi-k3" +_API_KEY: Final = "synthetic-fireworks-key" +_PROMPT: Final = "keep this conversation on one replica" +_SESSION_ID: Final = "conversation-affinity-6220" +_CACHED_TOKENS: Final = 7 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _cached_reply(request: Request, identity: str) -> Reply: + assert request.method == "POST" + assert request.target == "/chat/completions" + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _MODEL, body + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": _MODEL, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "pinned"}, "finish_reason": "stop"} + ], + "usage": { + "prompt_tokens": 12, + "completion_tokens": 1, + "total_tokens": 13, + "prompt_tokens_details": {"cached_tokens": _CACHED_TOKENS}, + }, + } + ).encode() + ) + + +@pytest.mark.covers("other.provider_wire.fireworks_ai.session_id_sent_as_affinity_header_and_cached_tokens_logged") +def test_fireworks_session_id_sends_affinity_header_and_logs_cache_read_tokens(gateway: Gateway) -> None: + identity: Final = f"fw-session-affinity-{uuid.uuid4().hex}" + with wire_server(lambda request: _cached_reply(request, identity)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"fireworks_ai/{_MODEL}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}]}, + headers={"x-litellm-session-id": _SESSION_ID}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["id"] == identity, response.text + requests: Final = wire.drain() + assert [(request.method, request.target) for request in requests] == [("POST", "/chat/completions")] + assert requests[0].headers.get("x-session-affinity") == _SESSION_ID, requests[0].headers + rows: Final = eventually( + lambda: read_rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda values: len(values) == 1, + seconds=70, + ) + usage_values: Final = object_value(object_value(rows[0]["metadata"])["additional_usage_values"]) + assert usage_values.get("cache_read_input_tokens") == _CACHED_TOKENS, rows[0]["metadata"] diff --git a/tests/integration/providers/test_gemini_messages_cache_control_wire.py b/tests/integration/providers/test_gemini_messages_cache_control_wire.py new file mode 100644 index 00000000000..71073b3dd5e --- /dev/null +++ b/tests/integration/providers/test_gemini_messages_cache_control_wire.py @@ -0,0 +1,86 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gemini-2.5-flash" +_API_KEY: Final = "synthetic-gemini-key" +_CACHE_NAME: Final = "cachedContents/synthetic-cache" +_CACHED_POLICY: Final = " ".join(f"policy clause {index} applies" for index in range(600)) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _generate_content_reply(text: str) -> bytes: + return json.dumps( + { + "candidates": [ + {"content": {"parts": [{"text": text}], "role": "model"}, "finishReason": "STOP", "index": 0} + ], + "usageMetadata": { + "promptTokenCount": 1300, + "candidatesTokenCount": 5, + "totalTokenCount": 1305, + "cachedContentTokenCount": 1290, + }, + "modelVersion": _BACKEND, + } + ).encode() + + +@pytest.mark.covers("other.provider_wire.gemini.messages_cache_control_creates_cached_content_with_anthropic_ttl") +def test_gemini_messages_cache_control_creates_cached_content_and_generates_from_it(gateway: Gateway) -> None: + identity: Final = f"gemini-messages-cache-{uuid.uuid4().hex}" + user_prompt: Final = f"Summarize the policy. Request {identity}." + + def respond(request: Request) -> Reply: + assert request.headers["x-goog-api-key"] == _API_KEY, request.headers + if request.method == "GET": + assert request.target == f"/models/{_BACKEND}:cachedContents", request.target + return Reply(body=b"{}") + assert request.method == "POST", request.method + body: Final = _JSON_OBJECT.validate_json(request.body) + if request.target == f"/models/{_BACKEND}:cachedContents": + assert isinstance(body["displayName"], str) and body["displayName"], body + assert body == { + "contents": [{"role": "user", "parts": [{"text": "."}]}], + "model": f"models/{_BACKEND}", + "displayName": body["displayName"], + "ttl": "300s", + "system_instruction": {"parts": [{"text": _CACHED_POLICY}]}, + "tools": None, + } + return Reply(body=json.dumps({"name": _CACHE_NAME, "model": f"models/{_BACKEND}"}).encode()) + assert request.target == f"/models/{_BACKEND}:generateContent", request.target + assert body == { + "contents": [{"role": "user", "parts": [{"text": user_prompt}]}], + "generationConfig": {"max_output_tokens": 32}, + "cachedContent": _CACHE_NAME, + } + return Reply(body=_generate_content_reply("The policy applies.")) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"gemini/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 32, + "system": [ + {"type": "text", "text": _CACHED_POLICY, "cache_control": {"type": "ephemeral", "ttl": "5m"}} + ], + "messages": [{"role": "user", "content": user_prompt}], + }, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["content"] == [{"type": "text", "text": "The policy applies."}], response.text + assert [(request.method, request.target) for request in wire.drain()] == [ + ("GET", f"/models/{_BACKEND}:cachedContents"), + ("POST", f"/models/{_BACKEND}:cachedContents"), + ("POST", f"/models/{_BACKEND}:generateContent"), + ] diff --git a/tests/integration/providers/test_nvidia_nim_ranking_wire.py b/tests/integration/providers/test_nvidia_nim_ranking_wire.py new file mode 100644 index 00000000000..9ed7a4eb48e --- /dev/null +++ b/tests/integration/providers/test_nvidia_nim_ranking_wire.py @@ -0,0 +1,47 @@ +import json +from typing import Final + +import pytest +from integration._support.client import JSON_OBJECT, Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "nvidia_nim/ranking/nvidia/llama-3.2-nv-rerankqa-1b-v2" +QUERY: Final = "which passage shows the gateway diagram" +IMAGE_PASSAGE: Final = "data:image/png;base64,aW50ZWdyYXRpb24tc3ludGhldGljLWltYWdl" +TEXT_PASSAGE: Final = "the gateway proxies rerank calls" +RESPONSE: Final = json.dumps( + {"rankings": [{"index": 0, "logit": 0.82}, {"index": 1, "logit": -1.4}], "usage": {"total_tokens": 11}} +).encode() + + +def ranking_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/ranking", request.target + assert request.headers["authorization"] == "Bearer integration-provider-key" + body: Final = JSON_OBJECT.validate_json(request.body) + assert body == { + "model": "nvidia/llama-3.2-nv-rerankqa-1b-v2", + "query": {"text": QUERY}, + "passages": [{"image": IMAGE_PASSAGE}, {"text": TEXT_PASSAGE}], + }, body + return Reply(body=RESPONSE) + + +@pytest.mark.covers( + "providers.nvidia_nim_ranking.image_passages_reach_ranking_without_top_k_and_top_n_is_applied_locally" +) +def test_nvidia_nim_ranking_keeps_image_passages_and_applies_top_n_without_sending_top_k(gateway: Gateway) -> None: + with wire_server(ranking_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=MODEL, api_base=wire.url, model_info={"mode": "rerank"}) + response: Final = gateway.request( + "POST", + "/v1/rerank", + { + "model": model, + "query": QUERY, + "documents": [{"image": IMAGE_PASSAGE}, {"text": TEXT_PASSAGE}], + "top_n": 1, + }, + ) + assert response.status_code == 200, response.text + assert response.json()["results"] == [{"index": 0, "relevance_score": 0.82}], response.text + assert len(wire.drain()) == 1, "Expected exactly one provider ranking call" diff --git a/tests/integration/providers/test_rerank_latency_headers_wire.py b/tests/integration/providers/test_rerank_latency_headers_wire.py new file mode 100644 index 00000000000..62624e0480a --- /dev/null +++ b/tests/integration/providers/test_rerank_latency_headers_wire.py @@ -0,0 +1,50 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +MODEL: Final = "cohere/synthetic-rerank-model-without-pricing" +QUERY: Final = "which document mentions the gateway" +DOCUMENTS: Final = ("the gateway proxies rerank calls", "unrelated synthetic text") +RESPONSE: Final = json.dumps( + { + "id": "synthetic-rerank-id", + "results": [{"index": 0, "relevance_score": 0.91}, {"index": 1, "relevance_score": 0.03}], + "meta": {"api_version": {"version": "2"}, "billed_units": {"search_units": 1}}, + } +).encode() + + +def rerank_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target.endswith("/rerank"), request.target + body: Final = json.loads(request.body) + assert body["query"] == QUERY and body["documents"] == list(DOCUMENTS), request.body + return Reply(body=RESPONSE) + + +@pytest.mark.covers("providers.rerank.response_carries_latency_and_cost_headers") +def test_rerank_response_carries_call_id_latency_and_cost_headers_like_chat_completions(gateway: Gateway) -> None: + with wire_server(rerank_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_key="synthetic-cohere-key", + api_base=wire.url, + model_info={"mode": "rerank"}, + ) + response: Final = gateway.request( + "POST", "/v1/rerank", {"model": model, "query": QUERY, "documents": list(DOCUMENTS), "top_n": 2} + ) + assert response.status_code == 200, response.text + assert [(result["index"], result["relevance_score"]) for result in response.json()["results"]] == [ + (0, 0.91), + (1, 0.03), + ], response.text + assert len(wire.drain()) == 1, "Expected exactly one provider rerank call" + assert response.headers["x-litellm-model-group"] == model, response.text + assert uuid.UUID(response.headers["x-litellm-call-id"]).version == 4, response.headers + assert float(response.headers["x-litellm-response-cost"]) == 0.0, response.headers + assert float(response.headers["x-litellm-response-duration-ms"]) > 0, response.headers + assert float(response.headers["x-litellm-overhead-duration-ms"]) >= 0, response.headers diff --git a/tests/integration/providers/test_responses_bridge_incomplete.py b/tests/integration/providers/test_responses_bridge_incomplete.py index 3352ca5775e..e700d17ea88 100644 --- a/tests/integration/providers/test_responses_bridge_incomplete.py +++ b/tests/integration/providers/test_responses_bridge_incomplete.py @@ -62,3 +62,129 @@ def test_chat_over_responses_deployment_returns_length_when_output_tokens_run_ou assert body["choices"][0]["message"]["role"] == "assistant", response.text assert body["usage"]["prompt_tokens"] == 12 and body["usage"]["completion_tokens"] == 16, response.text assert body["usage"]["total_tokens"] == 28, response.text + + +@pytest.mark.covers("other.provider_wire.responses_bridge.sub_minimum_max_tokens_clamped_to_provider_floor") +def test_messages_over_responses_deployment_with_max_tokens_1_is_clamped_to_16_instead_of_400(gateway: Gateway) -> None: + identity: Final = "responses-clamp-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + body: Final = json.loads(request.body) + assert body["model"] == "gpt-5.4" + if body["max_output_tokens"] < 16: + return Reply( + status=400, + body=json.dumps( + { + "error": { + "message": "Invalid 'max_output_tokens': integer below minimum value. Expected a value >= 16, but got 1 instead.", + "type": "invalid_request_error", + "param": "max_output_tokens", + "code": "integer_below_min_value", + } + } + ).encode(), + ) + assert body["max_output_tokens"] == 16 + assert body["input"] == [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": f"warmup probe {identity}"}], + } + ] + return Reply( + body=json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-5.4", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 12, "output_tokens": 1, "total_tokens": 13}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-5.4", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 1, + "messages": [{"role": "user", "content": f"warmup probe {identity}"}], + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert len(wire.drain()) == 1 + assert body["role"] == "assistant", response.text + assert body["content"] == [{"type": "text", "text": "ok"}], response.text + assert body["stop_reason"] == "end_turn", response.text + + +@pytest.mark.covers("providers.responses_bridge.sub_minimum_max_tokens_is_raised_to_the_openai_floor") +def test_messages_over_responses_deployment_with_max_tokens_one_reaches_openai_as_sixteen(gateway: Gateway) -> None: + identity: Final = "responses-min-tokens-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + body: Final = json.loads(request.body) + assert body["model"] == "gpt-5.6-sol" + assert body["max_output_tokens"] == 16, body + return Reply( + body=json.dumps( + { + "id": f"resp_{identity}", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-5.6-sol", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 9, "output_tokens": 1, "total_tokens": 10}, + } + ).encode() + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/responses/gpt-5.6-sol", api_base=wire.url, api_key="synthetic-openai-key" + ) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 1, + "messages": [{"role": "user", "content": f"warmup {identity}"}], + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert len(wire.drain()) == 1 + assert body["content"] == [{"type": "text", "text": "ok"}], response.text + assert body["usage"]["input_tokens"] == 9 and body["usage"]["output_tokens"] == 1, response.text diff --git a/tests/integration/providers/test_responses_bridge_namespace_tools.py b/tests/integration/providers/test_responses_bridge_namespace_tools.py new file mode 100644 index 00000000000..746dac03ced --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_namespace_tools.py @@ -0,0 +1,159 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +JSON_LIST: Final = TypeAdapter(list[dict[str, JsonValue]]) +NAMESPACE: Final = "mcp__everything" +TOOL_NAME: Final = "get_sum" +FLATTENED_NAME: Final = f"{NAMESPACE}__{TOOL_NAME}" +CALL_ID: Final = "call_synthetic_get_sum" +ARGUMENTS: Final = json.dumps({"a": 2, "b": 3}) +PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "required": ["a", "b"], + "properties": {"a": {"type": "number"}, "b": {"type": "number"}}, +} +NAMESPACE_TOOL: Final[dict[str, JsonValue]] = { + "type": "namespace", + "name": NAMESPACE, + "description": "Tools exposed by the everything MCP server", + "tools": [ + { + "type": "function", + "name": TOOL_NAME, + "description": "Adds two numbers", + "strict": False, + "parameters": PARAMETERS, + } + ], +} +EXPECTED_CHAT_TOOLS: Final[list[JsonValue]] = [ + { + "type": "function", + "function": { + "name": FLATTENED_NAME, + "description": "Tools exposed by the everything MCP server\n\nAdds two numbers", + "parameters": PARAMETERS, + "strict": False, + }, + } +] + + +def tool_call_completion(marker: str) -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "created": 1789788253, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": CALL_ID, + "type": "function", + "function": {"name": FLATTENED_NAME, "arguments": ARGUMENTS}, + } + ], + }, + } + ], + "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}, + } + ).encode() + + +def text_completion(marker: str) -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{marker}-final", + "object": "chat.completion", + "created": 1789788254, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "message": {"role": "assistant", "content": "The sum is 5"}, + } + ], + "usage": {"prompt_tokens": 40, "completion_tokens": 5, "total_tokens": 45}, + } + ).encode() + + +@pytest.mark.covers("other.provider_wire.responses_bridge.codex_namespace_tools_reach_chat_upstream_and_round_trip") +def test_codex_namespace_tool_is_flattened_for_chat_upstream_and_restored_in_responses_output( + gateway: Gateway, +) -> None: + marker: Final = uuid.uuid4().hex + prompt: Final = f"add 2 and 3 {marker}" + + def chat_peer(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions", request.target + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["tools"] == EXPECTED_CHAT_TOOLS, body + messages: Final = JSON_LIST.validate_python(body["messages"]) + if len(messages) == 1: + return Reply(body=tool_call_completion(marker)) + assert messages[1]["role"] == "assistant", messages + history_calls: Final = JSON_LIST.validate_python(messages[1]["tool_calls"]) + assert [(call["id"], call["function"]) for call in history_calls] == [ + (CALL_ID, {"name": FLATTENED_NAME, "arguments": ARGUMENTS}) + ], messages + assert messages[2] == {"role": "tool", "tool_call_id": CALL_ID, "content": "5"}, messages + return Reply(body=text_completion(marker)) + + with wire_server(chat_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="deepseek/gpt-4o-mini", api_base=wire.url + "/v1") + first: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": prompt, "tools": [NAMESPACE_TOOL], "store": False}, + ) + assert first.status_code == 200, first.text + first_output: Final = JSON_LIST.validate_python(JSON_OBJECT.validate_json(first.content)["output"]) + calls: Final = tuple(item for item in first_output if item["type"] == "function_call") + assert len(calls) == 1, first.text + assert calls[0]["name"] == TOOL_NAME, first.text + assert calls[0]["namespace"] == NAMESPACE, first.text + assert calls[0]["call_id"] == CALL_ID, first.text + assert calls[0]["arguments"] == ARGUMENTS, first.text + + second: Final = gateway.request( + "POST", + "/v1/responses", + { + "model": model, + "input": [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": prompt}]}, + { + "type": "function_call", + "call_id": CALL_ID, + "name": TOOL_NAME, + "namespace": NAMESPACE, + "arguments": ARGUMENTS, + }, + {"type": "function_call_output", "call_id": CALL_ID, "output": "5"}, + ], + "tools": [NAMESPACE_TOOL], + "store": False, + }, + ) + assert second.status_code == 200, second.text + second_output: Final = JSON_LIST.validate_python(JSON_OBJECT.validate_json(second.content)["output"]) + assert [item["type"] for item in second_output] == ["message"], second.text + assert JSON_LIST.validate_python(second_output[0]["content"])[0]["text"] == "The sum is 5", second.text + assert len(wire.drain()) == 2 diff --git a/tests/integration/providers/test_responses_bridge_stream_options.py b/tests/integration/providers/test_responses_bridge_stream_options.py new file mode 100644 index 00000000000..a0efc8d47d0 --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_stream_options.py @@ -0,0 +1,98 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + + +def _responses_stream(identity: str, text: str) -> tuple[bytes, ...]: + completed: Final = { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.3-codex", + "output": [ + { + "type": "message", + "id": f"msg_{identity}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + ], + "usage": { + "input_tokens": 11, + "output_tokens": 4, + "total_tokens": 15, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } + events: Final = ( + {"type": "response.created", "response": {**completed, "status": "in_progress", "output": [], "usage": None}}, + { + "type": "response.output_text.delta", + "item_id": f"msg_{identity}", + "output_index": 0, + "content_index": 0, + "delta": text, + }, + {"type": "response.completed", "response": completed}, + ) + return tuple(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode() for event in events) + + +@pytest.mark.covers("providers.responses_bridge.always_include_stream_usage_keeps_include_usage_off_the_responses_wire") +def test_messages_stream_with_always_include_stream_usage_omits_include_usage_from_responses_request( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "responses-stream-options-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + return Reply(content_type="text/event-stream", chunks=_responses_stream(identity, "usage control")) + + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"].update({"always_include_stream_usage": True}) + path: Final = tmp_path / "always_include_stream_usage.yaml" + path.write_text(yaml.safe_dump(config)) + with ( + wire_server(respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=path) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-5.3-codex", api_base=wire.url, api_key="synthetic-openai-key") + response: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "stream": True, + "messages": [{"role": "user", "content": f"count the usage {identity}"}], + }, + ) + assert response.status_code == 200, response.text + assert "event: message_stop" in response.text, response.text + requests: Final = wire.drain() + assert len(requests) == 1, response.text + assert json.loads(requests[0].body) == { + "model": "gpt-5.3-codex", + "input": [ + { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": f"count the usage {identity}"}], + } + ], + "include": ["reasoning.encrypted_content"], + "max_output_tokens": 64, + "stream": True, + }, response.text diff --git a/tests/integration/providers/test_responses_client_header_forwarding_wire.py b/tests/integration/providers/test_responses_client_header_forwarding_wire.py new file mode 100644 index 00000000000..50557cd1727 --- /dev/null +++ b/tests/integration/providers/test_responses_client_header_forwarding_wire.py @@ -0,0 +1,87 @@ +import json +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "gpt-5.4-mini" +_API_KEY: Final = "synthetic-openai-key" +_CLIENT_HEADER: Final = "x-my-new-header" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_OUTPUT_MESSAGE: Final[dict[str, JsonValue]] = { + "type": "message", + "id": "msg_forwarded", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "header wire control", "annotations": []}], +} +_RESPONSE: Final = json.dumps( + { + "id": "resp_forwarded", + "object": "response", + "status": "completed", + "created_at": 1700000000, + "model": _BACKEND, + "output": [_OUTPUT_MESSAGE], + "usage": { + "input_tokens": 9, + "output_tokens": 3, + "total_tokens": 12, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + } +).encode() + + +def _forwarding_config(directory: Path) -> Path: + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["general_settings"]["forward_client_headers_to_llm_api"] = True + path: Final = directory / "forwarding.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@pytest.mark.covers("providers.responses_api.forwarded_client_headers_reach_the_provider") +def test_client_x_header_is_forwarded_to_the_provider_on_responses(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = f"hello-from-client-{uuid4().hex}" + prompt: Final = f"forward my header {marker}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert request.headers["authorization"] == f"Bearer {_API_KEY}" + assert request.headers.get(_CLIENT_HEADER) == marker, dict(request.headers) + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _BACKEND and body["input"] == prompt, request.body + return Reply(body=_RESPONSE) + + with ( + wire_server(respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_forwarding_config(tmp_path)) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"openai/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = candidate.request( + "POST", + "/v1/responses", + {"model": model, "input": prompt, "stream": False}, + headers={_CLIENT_HEADER: marker}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["output"] == [ + { + **_OUTPUT_MESSAGE, + "phase": None, + "content": [ + {"type": "output_text", "text": "header wire control", "annotations": [], "logprobs": None} + ], + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/responses")] diff --git a/tests/integration/providers/test_sagemaker_chat_wire.py b/tests/integration/providers/test_sagemaker_chat_wire.py new file mode 100644 index 00000000000..346f4e59e0f --- /dev/null +++ b/tests/integration/providers/test_sagemaker_chat_wire.py @@ -0,0 +1,90 @@ +import json +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_ENDPOINT: Final = "integration-vllm-endpoint" +_INFERENCE_COMPONENT: Final = "integration-vllm-component" +_SERVED_MODEL: Final = "integration-org/served-chat-model" +_ACCESS_KEY: Final = "AKIAINTEGRATION000003" +_PROMPT: Final = "synthetic inference component request" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _completion(identity: str) -> bytes: + return json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": _SERVED_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "sagemaker wire control"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}, + } + ).encode() + + +@pytest.mark.covers( + "providers.sagemaker_chat_wire.inference_component_header_is_signed_and_hf_model_name_is_the_body_model" +) +def test_sagemaker_chat_signs_the_inference_component_header_and_sends_hf_model_name_as_the_body_model( + gateway: Gateway, +) -> None: + identity: Final = f"sagemaker-chat-{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + assert request.method == "POST", request + assert request.target == "/", request + assert request.headers["x-amzn-sagemaker-inference-component"] == _INFERENCE_COMPONENT, dict(request.headers) + authorization: Final = request.headers["authorization"] + assert authorization.startswith(f"AWS4-HMAC-SHA256 Credential={_ACCESS_KEY}/"), authorization + signed_headers: Final = next(part for part in authorization.split(", ") if part.startswith("SignedHeaders=")) + assert "x-amzn-sagemaker-inference-component" in signed_headers.removeprefix("SignedHeaders=").split(";"), ( + authorization + ) + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["model"] == _SERVED_MODEL, body + assert body["messages"] == [{"role": "user", "content": _PROMPT}], body + assert body["max_tokens"] == 16, body + assert "hf_model_name" not in body, body + return Reply(body=_completion(identity)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"sagemaker_chat/{_ENDPOINT}", + api_key=None, + api_base=None, + model_id=_INFERENCE_COMPONENT, + hf_model_name=_SERVED_MODEL, + aws_access_key_id=_ACCESS_KEY, + aws_secret_access_key="synthetic-secret-key-for-testing", + aws_region_name="us-east-1", + sagemaker_base_url=wire.url, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": _PROMPT}], "max_tokens": 16}, + ) + assert response.status_code == 200, response.text + payload: Final = _JSON_OBJECT.validate_json(response.content) + assert payload["id"] == identity, response.text + assert payload["choices"] == [ + { + "finish_reason": "stop", + "index": 0, + "message": {"role": "assistant", "content": "sagemaker wire control"}, + "provider_specific_fields": {}, + } + ], response.text + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/")], response.text diff --git a/tests/integration/providers/test_vertex_batch_output_info_wire.py b/tests/integration/providers/test_vertex_batch_output_info_wire.py new file mode 100644 index 00000000000..a7ac896076c --- /dev/null +++ b/tests/integration/providers/test_vertex_batch_output_info_wire.py @@ -0,0 +1,124 @@ +import base64 +import functools +import json +from typing import Final + +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +PROJECT: Final = "cc-scripted-project" +LOCATION: Final = "us-central1" +MODEL: Final = "vertex_ai/gemini-2.5-flash" +VERTEX_MODEL_RESOURCE: Final = "publishers/google/models/gemini-2.5-flash" +BUCKET: Final = "integration-batch-bucket" +INPUT_FILE_ID: Final = f"gs://{BUCKET}/litellm-vertex-files/{VERTEX_MODEL_RESOURCE}/input.jsonl" +OUTPUT_PREFIX: Final = INPUT_FILE_ID.rsplit("/", 1)[0] +JOB_NAME: Final = f"projects/{PROJECT}/locations/{LOCATION}/batchPredictionJobs/7412345678901234567" +JOB_ID: Final = JOB_NAME.rsplit("/", 1)[-1] +EXPECTED_VERTEX_BODY: Final = { + "inputConfig": {"gcsSource": {"uris": [INPUT_FILE_ID]}, "instancesFormat": "jsonl"}, + "outputConfig": {"predictionsFormat": "jsonl", "gcsDestination": {"outputUriPrefix": OUTPUT_PREFIX}}, + "model": VERTEX_MODEL_RESOURCE, +} +VERTEX_REPLY: Final = { + "name": JOB_NAME, + "displayName": "litellm-vertex-batch-scripted", + "model": VERTEX_MODEL_RESOURCE, + "inputConfig": {"gcsSource": {"uris": [INPUT_FILE_ID]}, "instancesFormat": "jsonl"}, + "outputConfig": {"predictionsFormat": "jsonl", "gcsDestination": {"outputUriPrefix": OUTPUT_PREFIX}}, + "outputInfo": None, + "state": "JOB_STATE_PENDING", + "createTime": "2026-07-24T20:00:00.000000Z", + "updateTime": "2026-07-24T20:00:00.000000Z", +} + + +@functools.cache +def _vertex_private_key_pem() -> str: + return ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + + +def _vertex_service_account_json(url: str) -> str: + return json.dumps( + { + "type": "service_account", + "project_id": PROJECT, + "private_key_id": "scripted", + "private_key": _vertex_private_key_pem(), + "client_email": f"scripted@{PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{url}/_oauth/authorize", + "token_uri": f"{url}/_oauth/token", + } + ) + + +def _encoded(raw: str, model: str, prefix: str) -> str: + return prefix + base64.urlsafe_b64encode(f"litellm:{raw};model,{model}".encode()).decode().rstrip("=") + + +def vertex_peer(request: Request) -> Reply: + assert request.method == "POST", request.method + assert request.target == f"/v1/projects/{PROJECT}/locations/{LOCATION}/batchPredictionJobs", request.target + assert request.headers["authorization"] == "Bearer scripted-token" + assert request.headers["content-type"] == "application/json; charset=utf-8" + body: Final = json.loads(request.body) + display_name: Final = body.pop("displayName") + assert isinstance(display_name, str) and display_name.startswith("litellm-vertex-batch-"), display_name + assert body == EXPECTED_VERTEX_BODY, body + return Reply(body=json.dumps(VERTEX_REPLY).encode()) + + +@pytest.mark.covers("other.provider_wire.vertex_ai.batch_create_with_null_output_info_returns_batch_instead_of_500") +def test_vertex_batch_create_survives_explicit_null_output_info(gateway: Gateway) -> None: + with wire_server(vertex_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=MODEL, + api_key=None, + api_base=wire.url, + vertex_project=PROJECT, + vertex_location=LOCATION, + vertex_credentials=_vertex_service_account_json(gateway.upstream_url), + ) + response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": INPUT_FILE_ID, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert ( + body["id"], + body["object"], + body["status"], + body["input_file_id"], + body["output_file_id"], + body["error_file_id"], + body["completion_window"], + ) == ( + _encoded(JOB_ID, model, "batch_"), + "batch", + "validating", + _encoded(INPUT_FILE_ID, model, "file-"), + _encoded(f"{OUTPUT_PREFIX}/predictions.jsonl", model, "file-"), + None, + "24h", + ), response.text + requests: Final = wire.drain() + assert len(requests) == 1, f"Expected exactly one Vertex POST, saw {[request.target for request in requests]}" diff --git a/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py b/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py new file mode 100644 index 00000000000..449c5c9c105 --- /dev/null +++ b/tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py @@ -0,0 +1,138 @@ +import json +import time +from typing import Final + +import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +_BACKEND: Final = "gemini-3.7-flash" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_MODEL_PATH: Final = f"/v1/projects/{_PROJECT}/locations/{_LOCATION}/publishers/google/models/{_BACKEND}" +_PROMPT: Final = "Write a very long numbered list." +_PART_COUNT: Final = 8000 +_LINES_PER_FRAGMENT: Final = 64 +_STREAM_BUDGET_SECONDS: Final = 10.0 +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +class _Delta(BaseModel): + model_config = ConfigDict(extra="ignore") + content: str | None = None + + +class _Choice(BaseModel): + model_config = ConfigDict(extra="ignore") + delta: _Delta + finish_reason: str | None = None + + +class _Chunk(BaseModel): + model_config = ConfigDict(extra="ignore") + choices: tuple[_Choice, ...] + + +def _service_account_json(token_url: str) -> str: + private_key: Final = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + return json.dumps( + { + "type": "service_account", + "project_id": _PROJECT, + "private_key_id": "scripted", + "private_key": private_key, + "client_email": f"scripted@{_PROJECT}.iam.gserviceaccount.com", + "client_id": "0", + "auth_uri": f"{token_url}/_oauth/authorize", + "token_uri": f"{token_url}/_oauth/token", + } + ) + + +def _expected_text() -> str: + return "".join(f"{index}. item\n" for index in range(_PART_COUNT)) + + +def _gemini_response_fragments() -> tuple[bytes, ...]: + document: Final = json.dumps( + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": f"{index}. item\n"} for index in range(_PART_COUNT)], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 40000, "totalTokenCount": 40009}, + "modelVersion": _BACKEND, + }, + indent=2, + ) + lines: Final = document.split("\n") + fragments: Final = tuple( + "\n".join(lines[start : start + _LINES_PER_FRAGMENT]).encode() + b"\n" + for start in range(0, len(lines), _LINES_PER_FRAGMENT) + ) + return (b"data: " + fragments[0], *fragments[1:], b"\n") + + +@pytest.mark.covers("providers.vertex_gemini.fragmented_stream_json_is_parsed_once_and_stays_live") +def test_vertex_gemini_stream_split_across_many_fragments_completes_without_stalling(gateway: Gateway) -> None: + fragments: Final = _gemini_response_fragments() + + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target == f"{_MODEL_PATH}:streamGenerateContent?alt=sse" + assert request.headers["authorization"] == "Bearer scripted-token" + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["contents"] == [{"role": "user", "parts": [{"text": _PROMPT}]}] + assert body["generationConfig"] == {"temperature": 0.0} + return Reply(content_type="text/event-stream", chunks=fragments) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"vertex_ai/{_BACKEND}", + api_base=f"{wire.url}{_MODEL_PATH}", + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=_service_account_json(gateway.upstream_url.rstrip("/")), + ) + started: Final = time.monotonic() + with gateway.client.stream( + "POST", + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": True, + "temperature": 0.0, + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + timeout=_STREAM_BUDGET_SECONDS, + ) as response: + assert response.status_code == 200, response.read() + lines: Final = tuple(line for line in response.iter_lines() if line.startswith("data: ")) + elapsed: Final = time.monotonic() - started + assert elapsed < _STREAM_BUDGET_SECONDS, f"stream took {elapsed:.1f}s for {len(fragments)} fragments" + assert lines[-1] == "data: [DONE]", lines[-3:] + chunks: Final = tuple(_Chunk.model_validate_json(line.removeprefix("data: ")) for line in lines[:-1]) + choices: Final = tuple(choice for chunk in chunks for choice in chunk.choices) + assert "".join(choice.delta.content or "" for choice in choices) == _expected_text() + assert tuple(choice.finish_reason for choice in choices if choice.finish_reason) == ("stop",) + assert [(request.method, request.target) for request in wire.drain()] == [ + ("POST", f"{_MODEL_PATH}:streamGenerateContent?alt=sse") + ] diff --git a/tests/integration/providers/test_websearch_interception_wire.py b/tests/integration/providers/test_websearch_interception_wire.py index cba4f174233..a6f098cf64f 100644 --- a/tests/integration/providers/test_websearch_interception_wire.py +++ b/tests/integration/providers/test_websearch_interception_wire.py @@ -183,6 +183,8 @@ from integration._support.client import Gateway, eventually _QUERY: Final = "integration capped search" _TEXT_BLOCK: Final = {"type": "text", "text": "searching once more"} _NOT_INTERCEPTED: Final = "native tool reached the provider" +_FINAL_BLOCK: Final = {"type": "text", "text": "answered from the stored backend"} +_OWNED_RESULT_TEXT: Final = "Title: Owned result\nURL: https://owned.invalid/a\nSnippet: owned snippet" _SEARCH_RESULT_BLOCK: Final = { "type": "web_search_result", "url": "https://owned.invalid/a", @@ -297,3 +299,111 @@ def test_capped_websearch_interception_loop_ends_turn_instead_of_exposing_intern assert content[2] == _TEXT_BLOCK, response.text targets: Final = tuple((request.method, urlsplit(request.target).path) for request in wire.drain()) assert targets[-3:] == (("POST", "/v1/messages"), ("GET", "/search"), ("POST", "/v1/messages")), targets + + +@pytest.mark.covers("other.provider_wire.anthropic.websearch_interception_uses_database_search_tool_backend") +def test_database_created_search_tool_backend_receives_the_intercepted_query_over_a_same_named_config_tool( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "websearch-db-" + uuid.uuid4().hex + tool_name: Final = "integration-db-searxng-" + uuid.uuid4().hex + searched: Final = threading.Event() + + def respond(request: Request) -> Reply: + parts: Final = urlsplit(request.target) + if request.method == "GET" and parts.path == "/database/search": + assert parse_qs(parts.query)["q"] == [_QUERY], request.target + searched.set() + return Reply( + body=json.dumps( + { + "results": [ + {"title": "Owned result", "url": "https://owned.invalid/a", "content": "owned snippet"} + ] + } + ).encode() + ) + assert request.method == "POST" and parts.path == "/v1/messages", request.target + body: Final = json.loads(request.body) + if any(tool.get("type") == "web_search_20250305" for tool in body["tools"]): + return _anthropic_reply(identity, [{"type": "text", "text": _NOT_INTERCEPTED}], "end_turn") + assert [tool["name"] for tool in body["tools"]] == ["litellm_web_search"], body["tools"] + results: Final = [ + block + for message in body["messages"] + if isinstance(message["content"], list) + for block in message["content"] + if block["type"] == "tool_result" + ] + if not results: + return _anthropic_reply(identity, [_TEXT_BLOCK, _search_tool_use(identity)], "tool_use") + assert results == [{"type": "tool_result", "tool_use_id": identity, "content": _OWNED_RESULT_TEXT}], results + return _anthropic_reply(identity, [_FINAL_BLOCK], "end_turn") + + def send(candidate: Gateway, model: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": identity + " attempt " + uuid.uuid4().hex}], + "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 3}], + }, + ) + + def searched_through_proxy(response: httpx.Response) -> bool: + return searched.is_set() and _NOT_INTERCEPTED not in response.text + + with wire_server(respond) as wire, gateway.scenario() as scenario: + created: Final = gateway.post( + "/search_tools", + { + "search_tool": { + "search_tool_name": tool_name, + "litellm_params": {"search_provider": "searxng", "api_base": wire.url + "/database"}, + } + }, + ) + scenario.cleanups.callback(gateway.request, "DELETE", f"/search_tools/{created['search_tool_id']}") + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["search_tools"] = [ + { + "search_tool_name": tool_name, + "litellm_params": {"search_provider": "searxng", "api_base": wire.url + "/config"}, + } + ] + config["litellm_settings"].update( + { + "callbacks": ["websearch_interception"], + "websearch_interception_params": { + "enabled": True, + "enabled_providers": ["anthropic"], + "search_tool_name": tool_name, + }, + } + ) + path: Final = tmp_path / "websearch-db.yaml" + path.write_text(yaml.safe_dump(config)) + environment: Final = {"ANTHROPIC_API_BASE": wire.url} + with owned_proxy(gateway, tmp_path, environment, config=path) as candidate, candidate.scenario() as models: + model: Final = models.model( + model="anthropic/claude-sonnet-4-5-20250929", api_base=wire.url, api_key="synthetic-anthropic-key" + ) + response: Final = eventually(lambda: send(candidate, model), searched_through_proxy, seconds=40) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["stop_reason"] == "end_turn", response.text + assert body["content"][-1] == _FINAL_BLOCK, response.text + found: Final = [ + (result["url"], result["title"]) + for block in body["content"] + if block["type"] == "web_search_tool_result" + for result in block["content"] + ] + assert found == [("https://owned.invalid/a", "Owned result")], response.text + assert "litellm_web_search" not in response.text, response.text + targets: Final = tuple((request.method, urlsplit(request.target).path) for request in wire.drain()) + assert targets[-3:] == (("POST", "/v1/messages"), ("GET", "/database/search"), ("POST", "/v1/messages")), ( + targets + ) diff --git a/tests/integration/routing/test_advisor_failure_cooldown.py b/tests/integration/routing/test_advisor_failure_cooldown.py new file mode 100644 index 00000000000..41618e29109 --- /dev/null +++ b/tests/integration/routing/test_advisor_failure_cooldown.py @@ -0,0 +1,101 @@ +import json +import uuid +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +_ADVISOR_KEY: Final = "synthetic-advisor-key" +_QUESTION: Final = "which index should this query use" +_PROXY_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" + + +def _executor_reply(body: dict[str, object], identity: str) -> Reply: + tools: Final = body.get("tools") + message: Final = ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "advisor-call", + "type": "function", + "function": {"name": "advisor", "arguments": json.dumps({"question": _QUESTION})}, + } + ], + } + if isinstance(tools, list) + else {"role": "assistant", "content": "served without an advisor"} + ) + return Reply( + body=json.dumps( + { + "id": f"chatcmpl-{identity}-{uuid.uuid4().hex[:8]}", + "object": "chat.completion", + "created": 1, + "model": "llama-3.3-70b-versatile", + "choices": [{"index": 0, "message": message, "finish_reason": "tool_calls" if tools else "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, + } + ).encode() + ) + + +def _cooldowns_enabled_config(directory: Path) -> Path: + loaded: Final = yaml.safe_load(_PROXY_CONFIG.read_text()) + path: Final = directory / "cooldowns_enabled.yaml" + path.write_text(yaml.safe_dump({**loaded, "router_settings": {"num_retries": 0}})) + return path + + +@pytest.mark.covers("routing.cooldown.advisor_sub_call_failure_does_not_cool_down_the_executor_deployment") +def test_advisor_sub_call_401_leaves_the_executor_deployment_serving_the_next_request( + gateway: Gateway, tmp_path: Path +) -> None: + identity: Final = "advisor-cooldown-" + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + if request.target == "/v1/chat/completions": + return _executor_reply(json.loads(request.body), identity) + assert request.target == "/v1/messages" + assert request.headers["x-api-key"] == _ADVISOR_KEY + return Reply( + status=401, + body=json.dumps( + {"type": "error", "error": {"type": "authentication_error", "message": "invalid x-api-key"}} + ).encode(), + ) + + with ( + wire_server(respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=_cooldowns_enabled_config(tmp_path)) as candidate, + candidate.scenario() as scenario, + ): + executor: Final = scenario.model(model="hosted_vllm/gpt-4o-mini", api_base=wire.url + "/v1") + advisor: Final = scenario.model( + model="anthropic/claude-opus-4-1-20250805", api_base=wire.url, api_key=_ADVISOR_KEY + ) + advised: Final = candidate.request( + "POST", + "/v1/messages", + { + "model": executor, + "max_tokens": 64, + "messages": [{"role": "user", "content": identity}], + "tools": [{"type": "advisor_20260301", "name": "advisor", "model": advisor}], + }, + ) + assert advised.status_code == 401, advised.text + assert [request.target for request in wire.drain()] == ["/v1/chat/completions", "/v1/messages"] + unrelated: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": executor, "messages": [{"role": "user", "content": identity + " unrelated"}]}, + ) + assert unrelated.status_code == 200, unrelated.text + assert unrelated.json()["choices"][0]["message"]["content"] == "served without an advisor", unrelated.text + assert [request.target for request in wire.drain()] == ["/v1/chat/completions"] diff --git a/tests/integration/routing/test_key_tpm_reservation.py b/tests/integration/routing/test_key_tpm_reservation.py new file mode 100644 index 00000000000..8d0e06679cd --- /dev/null +++ b/tests/integration/routing/test_key_tpm_reservation.py @@ -0,0 +1,59 @@ +import json +import time +import uuid +from collections import Counter +from concurrent.futures import ThreadPoolExecutor +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +KEY_TPM_LIMIT: Final = 100 +MAX_TOKENS: Final = 80 +CONCURRENT_REQUESTS: Final = 10 +PROVIDER_HOLD_SECONDS: Final = 2.0 +UPSTREAM_REPLY: Final = json.dumps( + { + "id": "chatcmpl_tpm_reservation", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "reserved"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40}, + } +).encode() + + +@pytest.mark.covers("quota_management.key_tpm_limit.concurrent_requests_reserve_tokens_before_provider_call") +def test_concurrent_requests_over_key_tpm_are_rejected_before_reaching_provider(gateway: Gateway) -> None: + probe: Final = "tpm reservation probe " + uuid.uuid4().hex[:8] + messages: Final[list[JsonValue]] = [{"role": "user", "content": probe}] + + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/v1/chat/completions") + assert json.loads(request.body) == {"model": "gpt-4o-mini", "max_tokens": MAX_TOKENS, "messages": messages} + time.sleep(PROVIDER_HOLD_SECONDS) + return Reply(body=UPSTREAM_REPLY) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + key: Final = scenario.key(tpm_limit=KEY_TPM_LIMIT) + body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": MAX_TOKENS, + "messages": messages, + } + + def send(_: int) -> httpx.Response: + return gateway.request("POST", "/v1/chat/completions", body, key=key) + + with ThreadPoolExecutor(max_workers=CONCURRENT_REQUESTS) as pool: + responses: Final = tuple(pool.map(send, range(CONCURRENT_REQUESTS))) + statuses: Final = Counter(response.status_code for response in responses) + assert statuses == Counter({200: 1, 429: CONCURRENT_REQUESTS - 1}), tuple( + response.text for response in responses + ) + assert tuple(json.loads(request.body)["messages"] for request in wire.drain()) == (messages,) diff --git a/tests/integration/routing/test_priority_model_tpm_enforcement.py b/tests/integration/routing/test_priority_model_tpm_enforcement.py new file mode 100644 index 00000000000..c3d0446c1e1 --- /dev/null +++ b/tests/integration/routing/test_priority_model_tpm_enforcement.py @@ -0,0 +1,116 @@ +import json +import uuid +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +OPENAI_MODEL: Final = "gpt-4.1-mini" +PROMPT_TOKENS: Final = 30 +COMPLETION_TOKENS: Final = 10 +MODEL_TPM: Final = PROMPT_TOKENS + COMPLETION_TOKENS +PREMIUM_SHARE: Final = 0.5 +UPSTREAM_REPLY: Final = json.dumps( + { + "id": "chatcmpl_model_tpm_enforcement", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "model tpm control"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + } +).encode() + + +@pytest.mark.covers("other.routing.priority_rate_limits.tpm_only_model_rejects_priority_traffic_at_capacity") +def test_tpm_only_model_returns_429_to_priority_key_once_recorded_tokens_reach_model_tpm( + gateway: Gateway, tmp_path: Path +) -> None: + probe: Final = "model tpm probe " + uuid.uuid4().hex + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions" + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + body: Final = json.loads(request.body) + assert body["messages"][0]["content"].startswith(probe), body + assert body == { + "model": OPENAI_MODEL, + "messages": [{"role": "user", "content": body["messages"][0]["content"]}], + "max_tokens": 16, + } + return Reply(body=UPSTREAM_REPLY) + + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["litellm_settings"] = { + **configuration["litellm_settings"], + "callbacks": ["dynamic_rate_limiter_v3"], + "priority_reservation": {"premium": PREMIUM_SHARE}, + } + path: Final = tmp_path / "priority.yaml" + path.write_text(yaml.safe_dump(configuration)) + with ( + wire_server(respond) as wire, + owned_proxy(gateway, tmp_path, {}, config=path) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + tpm=MODEL_TPM, + ) + key: Final = scenario.key(metadata={"priority": "premium"}) + responses: Final[SimpleQueue[httpx.Response]] = SimpleQueue() + + def attempt() -> httpx.Response: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"{probe} {uuid.uuid4().hex}"}], + }, + key=key, + ) + responses.put(response) + return response + + first: Final = attempt() + assert first.status_code == 200, first.text + assert first.json()["usage"]["total_tokens"] == MODEL_TPM, first.text + blocked: Final = eventually(attempt, lambda response: response.status_code == 429, seconds=30) + served: Final = tuple(responses.get_nowait() for _ in range(responses.qsize())) + assert all(response.status_code == 200 for response in served[:-1]), [r.status_code for r in served] + assert len(wire.drain()) == len(served) - 1 + assert blocked.headers["x-litellm-priority"] == "premium", blocked.headers + assert blocked.headers["rate_limit_type"] == "tokens", blocked.headers + detail: Final = ( + f"Model capacity reached for {model}. Priority: premium, Rate limit type: tokens, " + f"Model TPM: {MODEL_TPM}, Model RPM: not configured, Remaining: 0" + ) + assert blocked.json() == { + "error": { + "message": detail, + "type": "throttling_error", + "param": None, + "code": "429", + "provider_specific_fields": {"error": detail}, + } + }, blocked.text diff --git a/tests/integration/routing/test_priority_rate_limit_headers.py b/tests/integration/routing/test_priority_rate_limit_headers.py index 2d21f3ba8b2..bd92a362885 100644 --- a/tests/integration/routing/test_priority_rate_limit_headers.py +++ b/tests/integration/routing/test_priority_rate_limit_headers.py @@ -5,7 +5,7 @@ from typing import Final import pytest import yaml -from integration._support.client import Gateway +from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server @@ -27,6 +27,27 @@ UPSTREAM_REPLY: Final = json.dumps( ).encode() +CHAT_MODEL: Final = "gpt-5.6" +MAX_COMPLETION_TOKENS: Final = 64 + + +def _chat_frames(identity: str, text: str) -> tuple[bytes, ...]: + events: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {"choices": [], "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}}, + ) + frames: Final = tuple( + b"data: " + + json.dumps( + {"id": identity, "object": "chat.completion.chunk", "created": 1, "model": CHAT_MODEL, **event} + ).encode() + + b"\n\n" + for event in events + ) + return (*frames, b"data: [DONE]\n\n") + + @pytest.mark.covers("other.routing.priority_rate_limits.v1_messages_success_exposes_v3_priority_headers") def test_non_streaming_v1_messages_success_carries_v3_priority_rate_limit_headers( gateway: Gateway, tmp_path: Path @@ -86,3 +107,89 @@ def test_non_streaming_v1_messages_success_carries_v3_priority_rate_limit_header } observed: Final = {name: response.headers.get(name) for name in expected} assert observed == expected, response.headers + + +@pytest.mark.covers("other.routing.priority_rate_limits.streaming_success_logs_v3_remaining_values_for_callbacks") +def test_streaming_chat_completion_success_logs_v3_rate_limit_remaining_values_for_callbacks( + gateway: Gateway, tmp_path: Path +) -> None: + probe: Final = "streaming remaining probe " + uuid.uuid4().hex + sink_secret: Final = "synthetic-sink-secret-" + uuid.uuid4().hex + + def provider(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions", request.target + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + assert json.loads(request.body) == { + "model": CHAT_MODEL, + "messages": [{"role": "user", "content": probe}], + "max_completion_tokens": MAX_COMPLETION_TOKENS, + "stream": True, + "stream_options": {"include_usage": True}, + }, request.body + return Reply(content_type="text/event-stream", chunks=_chat_frames("chatcmpl_" + probe[-8:], "streamed")) + + def sink(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {sink_secret}" + return Reply() + + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["litellm_settings"] = { + **configuration["litellm_settings"], + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + } + path: Final = tmp_path / "per_key_streaming.yaml" + path.write_text(yaml.safe_dump(configuration)) + with ( + wire_server(provider) as wire, + wire_server(sink) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url, "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}"}, + config=path, + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=f"openai/{CHAT_MODEL}", + api_base=wire.url + "/v1", + api_key="synthetic-openai-key", + ) + key: Final = scenario.key(model_rpm_limit={model: MODEL_RPM}, model_tpm_limit={model: MODEL_TPM}) + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": probe}], + "max_completion_tokens": MAX_COMPLETION_TOKENS, + "stream": True, + "stream_options": {"include_usage": True}, + }, + key=key, + ) + assert response.status_code == 200, response.text + assert '"content":"streamed"' in response.text, response.text + assert len(wire.drain()) == 1 + batches: Final[ + list[Request] + ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered() -> tuple[dict, ...]: + batches.extend(endpoint.drain()) + return tuple( + event for batch in batches for event in json.loads(batch.body) if event.get("model_group") == model + ) + + events: Final = eventually(delivered, lambda values: len(values) == 1, seconds=10) + assert (events[0]["status"], events[0]["stream"]) == ("success", True), json.dumps(events[0]) + additional_headers: Final = events[0]["hidden_params"]["additional_headers"] or {} + observed: Final = {name: value for name, value in additional_headers.items() if name.startswith("x-ratelimit-")} + remaining_tokens: Final = observed.get("x-ratelimit-model_per_key-remaining-tokens") + assert isinstance(remaining_tokens, int) and 0 < remaining_tokens <= MODEL_TPM, json.dumps(observed) + assert {name: value for name, value in observed.items() if not name.endswith("-remaining-tokens")} == { + "x-ratelimit-model_per_key-limit-requests": MODEL_RPM, + "x-ratelimit-model_per_key-remaining-requests": MODEL_RPM - 1, + "x-ratelimit-model_per_key-limit-tokens": MODEL_TPM, + }, json.dumps(events[0]["hidden_params"]) diff --git a/tests/integration/routing/test_team_model_tpm_limit.py b/tests/integration/routing/test_team_model_tpm_limit.py new file mode 100644 index 00000000000..741c41c9285 --- /dev/null +++ b/tests/integration/routing/test_team_model_tpm_limit.py @@ -0,0 +1,89 @@ +import json +import threading +import uuid +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually +from integration._support.wire import Reply, Request, wire_server + +PROVIDER_MODEL: Final = "gpt-4o-mini" +TEAM_MODEL_TPM: Final = 100 +MAX_TOKENS: Final = 60 +CONCURRENT_REQUESTS: Final = 3 +UPSTREAM_REPLY: Final = json.dumps( + { + "id": "chatcmpl-team-tpm-control", + "object": "chat.completion", + "created": 1, + "model": PROVIDER_MODEL, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "team tpm control"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, + } +).encode() + + +@pytest.mark.covers("routing.team_model_tpm.concurrent_requests_over_the_limit_are_rejected_before_the_provider_call") +def test_concurrent_team_model_tpm_requests_reserve_tokens_before_reaching_the_provider(gateway: Gateway) -> None: + probe: Final = "team tpm probe " + uuid.uuid4().hex + release: Final = threading.Event() + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/chat/completions" + assert request.headers["authorization"] == "Bearer synthetic-team-tpm-key" + body: Final = json.loads(request.body) + content: Final = body["messages"][0]["content"] + assert body == { + "model": PROVIDER_MODEL, + "messages": [{"role": "user", "content": content}], + "max_tokens": MAX_TOKENS, + } + assert content.startswith(probe), content + release.wait(timeout=10) + return Reply(body=UPSTREAM_REPLY) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{PROVIDER_MODEL}", + api_base=wire.url + "/v1", + api_key="synthetic-team-tpm-key", + ) + team: Final = scenario.team(metadata={"model_tpm_limit": {model: TEAM_MODEL_TPM}}) + key: Final = scenario.key(team_id=team) + + def send(index: int) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": MAX_TOKENS, + "messages": [{"role": "user", "content": f"{probe} {index}"}], + }, + key=key, + ) + + with ThreadPoolExecutor(max_workers=CONCURRENT_REQUESTS) as pool: + futures: Final[tuple[Future[httpx.Response], ...]] = tuple( + pool.submit(send, index) for index in range(CONCURRENT_REQUESTS) + ) + eventually( + lambda: sum(future.done() for future in futures) + wire.received.qsize(), + lambda settled: settled >= CONCURRENT_REQUESTS, + seconds=10, + ) + release.set() + responses: Final = tuple(future.result(timeout=15) for future in futures) + statuses: Final = tuple(sorted(response.status_code for response in responses)) + assert statuses == (200, 429, 429), tuple(response.text for response in responses) + assert len(wire.drain()) == 1, statuses + served: Final = next(response for response in responses if response.status_code == 200) + assert served.json()["usage"] == {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14}, served.text + for rejected in (response for response in responses if response.status_code == 429): + error: Final = rejected.json()["error"] + assert (error["type"], error["code"], error["param"]) == ("throttling_error", "429", None), rejected.text + assert f"Limit type: tokens. Current limit: {TEAM_MODEL_TPM}," in error["message"], rejected.text diff --git a/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py new file mode 100644 index 00000000000..1ec1f486bed --- /dev/null +++ b/tests/integration/sdk/test_aiohttp_session_rebuild_wire.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +import textwrap +import threading +from collections.abc import Iterator +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +CONFIGURED_KEEPALIVE_SECONDS: Final = 1 +IDLE_SECONDS: Final = 2 +RESPONSES: Final = TypeAdapter(list[dict[str, JsonValue]]) + +REBUILT_SESSION_EXCHANGE: Final = textwrap.dedent( + """ + import asyncio, json, sys + from aiohttp import ClientSession + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + async def main(base_url: str, idle_seconds: float) -> None: + shared = ClientSession() + handler = AsyncHTTPHandler(shared_session=shared) + await shared.close() + first = await handler.post(f"{base_url}/embeddings", json={"input": "warm-up"}) + await asyncio.sleep(idle_seconds) + second = await handler.post(f"{base_url}/embeddings", json={"input": "warm-up"}) + print(json.dumps([first.json(), second.json()])) + await handler.close() + + asyncio.run(main(sys.argv[1], float(sys.argv[2]))) + """ +) + + +class _ConnectionCountingPeer(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, address: tuple[str, int]) -> None: + super().__init__(address, _ConnectionHandler) + self.lock = threading.Lock() + self.connections = 0 + + def next_connection(self) -> int: + with self.lock: + self.connections += 1 + return self.connections + + +class _ConnectionHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + server: _ConnectionCountingPeer + + def setup(self) -> None: + super().setup() + self.connection_number = self.server.next_connection() + + def do_POST(self) -> None: + self.rfile.read(int(self.headers["Content-Length"])) + body: Final = json.dumps({"connection": self.connection_number}).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + +@pytest.fixture +def connection_counting_peer() -> Iterator[str]: + server: Final = _ConnectionCountingPeer(("127.0.0.1", 0)) + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_address[1]}" + server.shutdown() + server.server_close() + thread.join(timeout=10) + + +def _rebuilt_session_exchange(base_url: str) -> list[dict[str, JsonValue]]: + completed: Final = subprocess.run( + [sys.executable, "-P", "-c", REBUILT_SESSION_EXCHANGE, base_url, str(IDLE_SECONDS)], + env={ + **os.environ, + "AIOHTTP_KEEPALIVE_TIMEOUT": str(CONFIGURED_KEEPALIVE_SECONDS), + "AIOHTTP_SO_KEEPALIVE": "true", + }, + capture_output=True, + text=True, + timeout=60, + check=False, + ) + assert completed.returncode == 0, completed.stderr + return RESPONSES.validate_json(completed.stdout) + + +@pytest.mark.covers("sdk.aiohttp_transport.rebuilt_shared_session_keeps_configured_keepalive_timeout") +def test_rebuilt_shared_session_drops_idle_connection_after_configured_keepalive_timeout( + connection_counting_peer: str, +) -> None: + observed: Final = _rebuilt_session_exchange(connection_counting_peer) + assert observed == [{"connection": 1}, {"connection": 2}], observed diff --git a/tests/integration/streaming/test_file_content_streaming.py b/tests/integration/streaming/test_file_content_streaming.py new file mode 100644 index 00000000000..9ec7665307b --- /dev/null +++ b/tests/integration/streaming/test_file_content_streaming.py @@ -0,0 +1,48 @@ +import threading +import uuid +from collections.abc import Callable +from typing import Final + +import pytest +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server + +STREAM_CHUNK_BYTES: Final = 1024 * 1024 +HEAD: Final = b"h" * STREAM_CHUNK_BYTES +TAIL: Final = b'{"custom_id": "tail", "response": {"status_code": 200}}\n' + + +def _file_content_gated_after_head(gate: threading.Event) -> Callable[[Request], Reply]: + def respond(_request: Request) -> Reply: + return Reply(content_type="application/octet-stream", chunks=(HEAD, TAIL), gate_after_first=gate) + + return respond + + +@pytest.mark.covers("streaming.file_content.body_reaches_client_before_upstream_finishes_sending") +def test_file_content_streams_the_first_megabyte_to_the_client_before_the_upstream_sends_the_rest( + gateway: Gateway, +) -> None: + file_id: Final = "file-" + uuid.uuid4().hex + gate: Final = threading.Event() + with gateway.scenario() as scenario, wire_server(_file_content_gated_after_head(gate)) as wire: + model: Final = scenario.model(api_base=wire.url + "/v1") + with gateway.client.stream( + "GET", + f"/v1/files/{file_id}/content", + params={"model": model}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + assert response.status_code == 200, response.read().decode() + chunks: Final = response.iter_bytes(chunk_size=STREAM_CHUNK_BYTES) + head: Final = next(chunks) + assert head == HEAD, f"First {len(head)} bytes differ from the upstream head before the gate was released" + gate.set() + rest: Final = b"".join(chunks) + assert rest == TAIL, rest + requests: Final = wire.drain() + assert len(requests) == 1, requests + assert requests[0].method == "GET", requests[0] + assert requests[0].target == f"/v1/files/{file_id}/content", requests[0].target + assert requests[0].headers["authorization"] == "Bearer integration-provider-key", requests[0].headers + assert requests[0].body == b"", requests[0].body diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py index 7466a9454b8..bd89b869ef2 100644 --- a/tests/integration/streaming/test_stream_contracts.py +++ b/tests/integration/streaming/test_stream_contracts.py @@ -261,6 +261,87 @@ def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(ga ) +def reasoning_first_stream(identity: str) -> tuple[bytes, ...]: + usage: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [], + "usage": {"prompt_tokens": 11, "completion_tokens": 6, "total_tokens": 17}, + } + return ( + frame(identity, {"role": "assistant", "content": None, "reasoning_content": "Let me "}), + frame(identity, {"content": None, "reasoning_content": "think."}), + frame(identity, {"content": "Hello "}), + frame(identity, {"content": "there"}), + frame(identity, {}, finish="stop"), + b"data: " + json.dumps(usage).encode() + b"\n\n", + b"data: [DONE]\n\n", + ) + + +@pytest.mark.covers("streaming.messages_bridge.reasoning_content_only_chunks_open_a_thinking_block_first") +def test_messages_stream_opens_thinking_block_at_index_zero_for_reasoning_content_only_chunks( + gateway: Gateway, +) -> None: + identity: Final = "messages-reasoning-first-" + uuid.uuid4().hex + with ( + wire_server( + lambda request: Reply(content_type="text/event-stream", chunks=reasoning_first_stream(identity)) + ) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model="hosted_vllm/reasoning-model", api_base=wire.url + "/v1") + with gateway.client.stream( + "POST", + "/v1/messages", + json={ + "model": model, + "max_tokens": 64, + "stream": True, + "messages": [{"role": "user", "content": identity}], + }, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + assert response.headers["content-type"].startswith("text/event-stream"), text + events: Final = tuple(json.loads(line) for line in sse_data_lines(text)) + blocks: Final = tuple( + (event["index"], event.get("content_block") or event["delta"]) + for event in events + if event["type"] in ("content_block_start", "content_block_delta") + ) + assert blocks == ( + (0, {"type": "thinking", "thinking": "", "signature": ""}), + (0, {"type": "thinking_delta", "thinking": "Let me "}), + (0, {"type": "thinking_delta", "thinking": "think."}), + (1, {"type": "text", "text": ""}), + (1, {"type": "text_delta", "text": "Hello "}), + (1, {"type": "text_delta", "text": "there"}), + ), text + assert tuple(event["type"] 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_delta", + "content_block_stop", + "message_delta", + "message_stop", + ), text + message_delta: Final = next(event for event in events if event["type"] == "message_delta") + assert message_delta["usage"] == {"input_tokens": 11, "output_tokens": 6}, text + requests: Final = wire.drain() + assert len(requests) == 1 + outbound: Final = json.loads(requests[0].body) + assert outbound["stream"] is True and outbound["messages"] == [{"role": "user", "content": identity}], outbound + + @pytest.mark.covers("other.streaming.responses_bridge.empty_choices_chunks_complete_stream") def test_responses_stream_completes_through_empty_choices_metadata_and_usage_chunks(gateway: Gateway) -> None: identity: Final = "responses-empty-choices-" + uuid.uuid4().hex @@ -445,9 +526,7 @@ def test_primary_stream_with_empty_first_chunk_then_disconnect_falls_back_and_bi abort_after=2, ) ) as primary, - wire_server( - lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity)) - ) as fallback, + wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as fallback, ): config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["model_list"] = [