scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

test(integration): regression tests for July provider translation, routing and streaming bugs (#42693)

Browse source 
* test(integration): optional Anthropic tool properties stay optional on the OpenAI Responses wire (Pylon #6619)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Bedrock InvokeModel count-suffixed cache usage fields are reported and charged (Pylon #6708)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): anthropic messages honors the deployment request timeout (Pylon #6505)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): drop client_metadata before the Bedrock Converse body reaches the provider (Pylon #6645)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): repeat Bedrock requests under one session name assume the role once (Pylon #6681)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): messages stream keeps include_usage off the Responses wire with always_include_stream_usage (Pylon #6466)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): clamp sub-16 max_tokens to the Responses API floor instead of 400 (Pylon #6539)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): forwarded client x- headers reach the provider on /v1/responses (Pylon #6565)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Anthropic messages stop_sequences reach OpenAI-compatible providers as stop (Pylon #6536)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Codex namespace tools reach a chat upstream flattened and round-trip through /v1/responses (Pylon #6409)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): keep Claude 4.6 legacy thinking budget_tokens on /v1/messages (Pylon #6727)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): nvidia nim ranking keeps image passages and applies top_n without sending top_k (Pylon #6401)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): tpm-only model rejects priority traffic once recorded tokens reach the model tpm (Pylon #6344)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): reasoning-only chunks open an Anthropic thinking block at index zero on /v1/messages streams (Pylon #6337)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): file content streams to the client before the upstream finishes sending (Pylon #6315)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): agent whose card lives only at agentCard/v1.0 is reached with bearer auth (Pylon #6249)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): vertex batch create returns a batch when outputInfo is null (Pylon #6374)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): fireworks session id is sent as x-session-affinity and cached tokens land in spend log metadata (Pylon #6220)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): advisor sub-call failure does not cool down the executor deployment (Pylon #6212)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Gemini /v1/messages cache_control creates cachedContent with Anthropic ttl (Pylon #6221)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock_mantle max_output_tokens below 16 is clamped before reaching Mantle (Pylon #6262)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): missing thinking signature 400 on /v1/messages retries without thinking blocks (Pylon #6222)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): format Gemini messages cache_control wire test (Pylon #6221)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): rebuilt shared aiohttp session keeps the configured keepalive timeout (Pylon #6387)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): sagemaker_chat signs the inference component header and sends hf_model_name as the body model (Pylon #6187)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Bedrock Converse DeepSeek drops Anthropic thinking and sends V3 reasoning_effort raw (Pylon #6149)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): concurrent team model TPM requests are reserved before the provider call (Pylon #6075)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): /v1/messages honors the configured timeout against a stalled upstream (Pylon #6025)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Codex additional_tools input items reach Bedrock Mantle as top-level tools (Pylon #6012)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): advisor api_base without api_key never sends the proxy Anthropic key to the caller host (Pylon #6226)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): rerank responses carry call id, latency and cost headers (Pylon #5981)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): parse the outbound Anthropic body with the typed JSON adapter (Pylon #6025)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Bedrock Knowledge Base search forwards userContext to the Retrieve body (Pylon #5991)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): Marengo 3.0 text embeddings reach Bedrock nested under inputType (Pylon #5949)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): sub-16 max_tokens over a responses deployment reaches OpenAI as 16 (Pylon #6008)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): midturn system correction reaches the OpenAI Responses wire via /v1/messages (Pylon #6449)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): concurrent requests over a key tpm limit are rejected before reaching the provider (Pylon #5737)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): chat to responses bridge keeps deployment AWS credentials for Bedrock Mantle SigV4 (Pylon #5870)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): vertex gemini stream split across many fragments completes without stalling the proxy (Pylon #5838)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): bedrock mantle /v1/messages stream keeps stream true and relays SSE events (Pylon #5596)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): large chat payloads are released from worker memory after the request ends (Pylon #5920)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): streaming success logs v3 rate limit remaining values for callbacks (Pylon #5767)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): a database created search tool backs Anthropic web search interception (Pylon #5669)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): register july provider regression contracts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): make july provider regression tests deterministic under cache and worker sharing

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): drop order-fragile worker memory probe pending a real retention regression check

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): apply ruff import sorting and formatting

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): drop stale contract entry and pass question to advisor executor

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(integration): use tiktoken-backed executor model in advisor tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: kerry <kerry@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-23 09:51:06 -07:00 • committed by GitHub
parent 21530d887b
commit a286ebf42e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
39 changed files with 3434 additions and 62 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

101
tests/integration/compatibility/test_a2a_wire_versions.py
View file

@ -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

124
tests/integration/providers/test_anthropic_advisor_wire.py
View file

@ -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

77
tests/integration/providers/test_anthropic_legacy_thinking_budget_wire.py Normal file
View file

@ -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

65
tests/integration/providers/test_anthropic_messages_fireworks_stop_wire.py Normal file
View file

@ -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 = "</block>"
_ANSWER: Final = "<verdict>allow</verdict>"
_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")]

82
tests/integration/providers/test_anthropic_messages_openai_bridge_wire.py Normal file
View file

@ -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")]

92
tests/integration/providers/test_anthropic_messages_openai_tools_wire.py Normal file
View file

@ -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

58
tests/integration/providers/test_anthropic_messages_timeout_wire.py Normal file
View file

@ -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

95
tests/integration/providers/test_anthropic_thinking_signature_retry_wire.py Normal file
View file

@ -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"]

161
tests/integration/providers/test_anthropic_wire.py
View file

@ -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)

43
tests/integration/providers/test_bedrock_converse_client_metadata_wire.py Normal file
View file

@ -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

93
tests/integration/providers/test_bedrock_deepseek_reasoning_wire.py Normal file
View file

@ -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

86
tests/integration/providers/test_bedrock_invoke_cache_usage_wire.py Normal file
View file

@ -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

81
tests/integration/providers/test_bedrock_knowledge_base_user_context_wire.py Normal file
View file

@ -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

53
tests/integration/providers/test_bedrock_mantle_codex_input_wire.py
View file

@ -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

40
tests/integration/providers/test_bedrock_mantle_responses_wire.py
View file

@ -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")]

141
tests/integration/providers/test_bedrock_mantle_wire.py
View file

@ -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)]

35
tests/integration/providers/test_bedrock_marengo_embed_3_wire.py Normal file
View file

@ -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"

154
tests/integration/providers/test_bedrock_role_configuration.py
View file

@ -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"<AssumeRoleResult><Credentials><AccessKeyId>{assumed_key}</AccessKeyId><SecretAccessKey>synthetic-assumed-secret-key-for-testing</SecretAccessKey><SessionToken>{assumed_token}</SessionToken><Expiration>2035-01-01T00:00:00Z</Expiration></Credentials><AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/integration/session</Arn><AssumedRoleId>integration:session</AssumedRoleId></AssumedRoleUser><PackedPolicySize>0</PackedPolicySize></AssumeRoleResult>"
return Reply(content_type="text/xml", body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}<ResponseMetadata><RequestId>synthetic-sts-request</RequestId></ResponseMetadata></{action}Response>'.encode())
return Reply(
content_type="text/xml",
body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}<ResponseMetadata><RequestId>synthetic-sts-request</RequestId></ResponseMetadata></{action}Response>'.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 = "<GetCallerIdentityResult><Arn>arn:aws:iam::123456789012:user/integration-source</Arn><UserId>integration-source</UserId><Account>123456789012</Account></GetCallerIdentityResult>"
else:
assert parameters["RoleArn"] == [role]
result = f"<AssumeRoleResult><Credentials><AccessKeyId>{assumed_key}</AccessKeyId><SecretAccessKey>synthetic-assumed-secret-key-for-testing</SecretAccessKey><SessionToken>{assumed_token}</SessionToken><Expiration>2035-01-01T00:00:00Z</Expiration></Credentials><AssumedRoleUser><Arn>arn:aws:sts::123456789012:assumed-role/integration/session</Arn><AssumedRoleId>integration:session</AssumedRoleId></AssumedRoleUser><PackedPolicySize>0</PackedPolicySize></AssumeRoleResult>"
return Reply(
content_type="text/xml",
body=f'<{action}Response xmlns="https://sts.amazonaws.com/doc/2011-06-15/">{result}<ResponseMetadata><RequestId>synthetic-sts-request</RequestId></ResponseMetadata></{action}Response>'.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

15
tests/integration/providers/test_bedrock_thinking_tokens_wire.py
View file

@ -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)

69
tests/integration/providers/test_fireworks_ai_session_affinity_wire.py Normal file
View file

@ -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"]

86
tests/integration/providers/test_gemini_messages_cache_control_wire.py Normal file
View file

@ -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"),
]

47
tests/integration/providers/test_nvidia_nim_ranking_wire.py Normal file
View file

@ -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"

50
tests/integration/providers/test_rerank_latency_headers_wire.py Normal file
View file

@ -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

126
tests/integration/providers/test_responses_bridge_incomplete.py
View file

@ -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

159
tests/integration/providers/test_responses_bridge_namespace_tools.py Normal file
View file

@ -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

98
tests/integration/providers/test_responses_bridge_stream_options.py Normal file
View file

@ -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

87
tests/integration/providers/test_responses_client_header_forwarding_wire.py Normal file
View file

@ -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")]

90
tests/integration/providers/test_sagemaker_chat_wire.py Normal file
View file

@ -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

124
tests/integration/providers/test_vertex_batch_output_info_wire.py Normal file
View file

@ -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]}"

138
tests/integration/providers/test_vertex_gemini_fragmented_stream_wire.py Normal file
View file

@ -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")
]

110
tests/integration/providers/test_websearch_interception_wire.py
View file

@ -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
)

101
tests/integration/routing/test_advisor_failure_cooldown.py Normal file
View file

@ -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"]

59
tests/integration/routing/test_key_tpm_reservation.py Normal file
View file

@ -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,)

116
tests/integration/routing/test_priority_model_tpm_enforcement.py Normal file
View file

@ -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

109
tests/integration/routing/test_priority_rate_limit_headers.py
View file

@ -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"])

89
tests/integration/routing/test_team_model_tpm_limit.py Normal file
View file

@ -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

109
tests/integration/sdk/test_aiohttp_session_rebuild_wire.py Normal file
View file

@ -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

48
tests/integration/streaming/test_file_content_streaming.py Normal file
View file

@ -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

85
tests/integration/streaming/test_stream_contracts.py
View file

@ -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"] = [