scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-17 04:34:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 172

test(integration): cover health_routing endpoint contracts (#44744)

Browse source 
* test(integration): cover health_routing endpoint contracts

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

* test(integration): tighten health_routing contracts per review

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

* test(integration): close health_routing review gaps

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

* style(integration): format observed routing fallback cells

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

* test(integration): close health_routing final review findings

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

* test(integration): resolve health_routing review findings

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

---------

Co-authored-by: mateo <mateo@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-10-14 06:34:22 +00:00 • committed by GitHub
parent 0f0c1bc80b
commit e226016aee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 2944 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

267
tests/integration/management/test_health_routing.py Normal file
View file

@ -0,0 +1,267 @@
import hashlib
import uuid
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, eventually
from integration._support.database import scratch_database
from integration._support.database_relay import database_relay, held_statement_relay
from integration._support.process import owned_proxy
_READINESS_PROBE_QUERY: Final = b"SELECT 1"
def _health_routing_config(directory: Path, general_settings: dict[str, object]) -> Path:
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["general_settings"] = {**config.get("general_settings", {}), **general_settings}
path: Final = directory / "health-routing.yaml"
path.write_text(yaml.safe_dump(config))
return path
def _plaintext_database_url(database_url: str) -> str:
parts: Final = urlsplit(database_url)
parameters: Final = tuple(
(name, value) for name, value in parse_qsl(parts.query, keep_blank_values=True) if name != "sslmode"
) + (("sslmode", "disable"),)
return urlunsplit(parts._replace(query=urlencode(parameters)))
def _assert_liveness_and_readiness(gateway: Gateway, *, status_code: int, body: dict[str, str]) -> None:
for path in ("/health/liveliness", "/health/liveness"):
response: Final = gateway.client.get(path)
assert response.status_code == status_code, response.text
assert response.json() == ("I'm alive!" if status_code == 200 else body), response.text
readiness: Final = gateway.client.get("/health/readiness")
assert readiness.status_code == status_code, readiness.text
assert readiness.json() == body, readiness.text
@pytest.mark.timeout(180)
def test_public_readiness_reports_database_outage_and_allow_unavailable_setting(
gateway: Gateway, tmp_path: Path
) -> None:
with scratch_database() as database_url:
with database_relay(database_url, _READINESS_PROBE_QUERY) as (relay, relayed_url):
with owned_proxy(
gateway,
tmp_path,
{
"DATABASE_URL": _plaintext_database_url(relayed_url),
"PRISMA_HEALTH_WATCHDOG_ENABLED": "false",
},
workers=1,
) as candidate:
ready_before_outage: Final = candidate.client.get("/health/readiness")
assert ready_before_outage.status_code == 200, ready_before_outage.text
assert ready_before_outage.json() == {"status": "healthy", "db": "connected"}, ready_before_outage.text
relay.arm()
disconnected_after_outage: Final = eventually(
lambda: candidate.client.get("/health/readiness"),
lambda response: (
relay.tripped.is_set() and response.json() == {"status": "healthy", "db": "disconnected"}
),
seconds=45,
return_last_on_timeout=True,
)
assert relay.tripped.is_set(), "readiness database probe did not reach the relay"
assert disconnected_after_outage.status_code == 503, disconnected_after_outage.text
assert disconnected_after_outage.json() == {
"status": "healthy",
"db": "disconnected",
}, disconnected_after_outage.text
allow_unavailable_config: Final = _health_routing_config(
tmp_path,
{"allow_requests_on_db_unavailable": True},
)
with database_relay(database_url, _READINESS_PROBE_QUERY) as (relay, relayed_url):
with owned_proxy(
gateway,
tmp_path,
{
"DATABASE_URL": _plaintext_database_url(relayed_url),
"PRISMA_HEALTH_WATCHDOG_ENABLED": "false",
},
config=allow_unavailable_config,
workers=1,
) as candidate:
ready_before_outage_with_allow: Final = candidate.client.get("/health/readiness")
assert ready_before_outage_with_allow.status_code == 200, ready_before_outage_with_allow.text
assert ready_before_outage_with_allow.json() == {"status": "healthy", "db": "connected"}, (
ready_before_outage_with_allow.text
)
relay.arm()
disconnected_with_allow: Final = eventually(
lambda: candidate.client.get("/health/readiness"),
lambda response: (
relay.tripped.is_set() and response.json() == {"status": "healthy", "db": "disconnected"}
),
seconds=45,
return_last_on_timeout=True,
)
assert relay.tripped.is_set(), "allow-unavailable readiness probe did not reach the relay"
assert disconnected_with_allow.status_code == 200, disconnected_with_allow.text
assert disconnected_with_allow.json() == {
"status": "healthy",
"db": "disconnected",
}, disconnected_with_allow.text
@pytest.mark.timeout(300)
def test_public_readiness_reports_stalled_database_lookups(gateway: Gateway, tmp_path: Path) -> None:
with scratch_database() as database_url:
for allow_unavailable in (False, True):
unknown_key = "sk-" + uuid.uuid4().hex
lookup_hash = hashlib.sha256(unknown_key.encode()).hexdigest().encode()
config = (
_health_routing_config(tmp_path, {"allow_requests_on_db_unavailable": True})
if allow_unavailable
else None
)
with held_statement_relay(database_url, lookup_hash) as (relay, relayed_url):
overrides = {
"DATABASE_URL": _plaintext_database_url(relayed_url),
"PRISMA_HEALTH_WATCHDOG_ENABLED": "false",
"PROXY_DB_LOOKUP_DEADLINE_SECONDS": "1",
"PROXY_DB_LOOKUP_STALL_WINDOW_SECONDS": "15",
}
with owned_proxy(
gateway,
tmp_path,
overrides,
config=config,
workers=1,
) as candidate:
expected_request_status = 400 if allow_unavailable else 503
unknown_model_message = (
"/chat/completions: Invalid model name passed in model=unconfigured-stalled-readiness-model. "
"Call `/v1/models` to view available models for your key."
)
expected_request_body = (
{
"error": {
"message": unknown_model_message,
"type": "invalid_request_error",
"param": None,
"code": "400",
"provider_specific_fields": {"error": unknown_model_message},
}
}
if allow_unavailable
else {
"error": {
"message": (
"Service Unavailable, the authentication database is temporarily unreachable. "
"Please retry shortly."
),
"type": "no_db_connection",
"param": "None",
"code": "503",
}
}
)
with ThreadPoolExecutor(max_workers=1) as executor:
pending = executor.submit(
candidate.request,
"POST",
"/v1/chat/completions",
{
"model": "unconfigured-stalled-readiness-model",
"messages": [{"role": "user", "content": "stalled lookup control"}],
},
key=unknown_key,
)
assert relay.held.wait(30), "unknown-key lookup did not reach the held statement relay"
expected_status = 200 if allow_unavailable else 503
stalled = eventually(
lambda: candidate.client.get("/health/readiness"),
lambda response, expected_status=expected_status: (
response.status_code == expected_status
and response.json() == {"status": "healthy", "db": "stalled"}
),
seconds=5,
return_last_on_timeout=True,
)
assert stalled.status_code == expected_status, stalled.text
assert stalled.json() == {"status": "healthy", "db": "stalled"}, stalled.text
relay.release()
recovered = eventually(
lambda: candidate.client.get("/health/readiness"),
lambda response: (
response.status_code == 200
and response.json() == {"status": "healthy", "db": "connected"}
),
seconds=30,
return_last_on_timeout=True,
)
assert recovered.status_code == 200, recovered.text
assert recovered.json() == {"status": "healthy", "db": "connected"}, recovered.text
rejected = pending.result(timeout=60)
assert (rejected.status_code, JSON_OBJECT.validate_json(rejected.content)) == (
expected_request_status,
expected_request_body,
), rejected.text
@pytest.mark.parametrize("token_source", ("config", "environment"))
def test_drain_endpoint_requires_token_and_changes_health_state_only_after_authorization(
token_source: str, gateway: Gateway, tmp_path: Path
) -> None:
disabled: Final = gateway.client.get("/health/drain")
assert disabled.status_code == 404, disabled.text
assert disabled.json() == {"detail": "Not Found"}, disabled.text
token: Final = "synthetic-health-drain-token"
general_settings: Final = {
"enable_drain_endpoint": True,
**({"drain_endpoint_token": token} if token_source == "config" else {}),
}
config: Final = _health_routing_config(tmp_path, general_settings)
overrides: Final = {"DRAIN_ENDPOINT_TOKEN": token} if token_source == "environment" else {}
remove_environment: Final = ("DRAIN_ENDPOINT_TOKEN",) if token_source == "config" else ()
with owned_proxy(
gateway,
tmp_path,
overrides,
config=config,
workers=1,
remove_environment=remove_environment,
) as candidate:
_assert_liveness_and_readiness(
candidate,
status_code=200,
body={"status": "healthy", "db": "connected"},
)
missing: Final = candidate.client.get("/health/drain")
assert missing.status_code == 401, missing.text
assert missing.json() == {"detail": "Invalid or missing X-Drain-Token"}, missing.text
wrong: Final = candidate.client.get(
"/health/drain",
headers={"X-Drain-Token": "synthetic-wrong-token"},
)
assert wrong.status_code == 401, wrong.text
assert wrong.json() == {"detail": "Invalid or missing X-Drain-Token"}, wrong.text
_assert_liveness_and_readiness(
candidate,
status_code=200,
body={"status": "healthy", "db": "connected"},
)
drained: Final = candidate.client.get(
"/health/drain",
headers={"X-Drain-Token": token},
)
assert drained.status_code == 200, drained.text
assert drained.json() == {"status": "drained", "drained_requests": 0}, drained.text
for path in ("/health/liveliness", "/health/liveness", "/health/readiness"):
response: Final = candidate.client.get(path)
assert response.status_code == 503, response.text
assert response.json() == {"status": "shutting_down"}, response.text

1092
tests/integration/management/test_model_health_check.py
View file

File diff suppressed because it is too large Load diff

636
tests/integration/routing/test_auto_router_management.py Normal file
View file

@ -0,0 +1,636 @@
import hashlib
import json
import uuid
from collections.abc import Callable
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
from integration._support.database import read_rows, scratch_database
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, Wire, wire_server
from pydantic import JsonValue, TypeAdapter
JSON_ARRAY: Final = TypeAdapter(list[JsonValue])
def _assistant_content(response: httpx.Response) -> str:
body: Final = JSON_OBJECT.validate_json(response.content)
choices: Final = JSON_ARRAY.validate_python(body["choices"])
assert len(choices) == 1, response.text
choice: Final = object_value(choices[0])
message: Final = object_value(choice["message"])
return string_value(message["content"])
def _tier_peer(model: str, prompts: tuple[str, ...]) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
assert request.method == "POST", request.method
if request.target == "/v1/responses":
return Reply(
body=json.dumps(
{
"id": "resp-" + uuid.uuid4().hex[:8],
"object": "response",
"created_at": 1700000000,
"status": "completed",
"model": model,
"output": [
{
"id": "msg-" + uuid.uuid4().hex[:8],
"type": "message",
"status": "completed",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": "synthetic tier response",
"annotations": [],
}
],
}
],
"usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12},
}
).encode()
)
assert request.target == "/v1/chat/completions", request.target
body: Final = JSON_OBJECT.validate_json(request.body)
expected_bodies: Final = tuple(
{"model": model, "messages": [{"role": "user", "content": prompt}]} for prompt in prompts
)
assert body in expected_bodies, request.body.decode()
return Reply(
body=json.dumps(
{
"id": "chatcmpl-" + uuid.uuid4().hex[:8],
"object": "chat.completion",
"created": 1700000000,
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "synthetic tier response"},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12},
}
).encode()
)
return respond
def _deployment(
alias: str,
model: str,
wire: Wire,
input_cost_per_token: float,
output_cost_per_token: float,
) -> dict[str, JsonValue]:
return {
"model_name": alias,
"litellm_params": {
"model": f"openai/{model}",
"api_base": f"{wire.url}/v1",
"api_key": "synthetic-auto-router-key",
},
"model_info": {
"id": f"deployment-{alias}",
"input_cost_per_token": input_cost_per_token,
"output_cost_per_token": output_cost_per_token,
},
}
def _session_response(
gateway: Gateway,
key: str,
session_id: str,
database_url: str | None = None,
stored_session_id: str | None = None,
seconds: float = 70,
return_last_on_timeout: bool = False,
) -> httpx.Response:
if database_url is None or stored_session_id is None:
return eventually(
lambda: gateway.request(
"GET",
"/auto_router/session",
params={"session_id": session_id},
key=key,
),
lambda response: response.status_code == 200,
seconds=seconds,
return_last_on_timeout=return_last_on_timeout,
)
key_hash: Final = hashlib.sha256(key.encode()).hexdigest()
def read_response_state() -> tuple[httpx.Response, bool]:
rows: Final = read_rows(
'SELECT session_id FROM "LiteLLM_AutoRouterSession" WHERE api_key = %s AND session_id = %s',
(key_hash, stored_session_id),
database_url=database_url,
)
response: Final = gateway.request(
"GET",
"/auto_router/session",
params={"session_id": session_id},
key=key,
)
return response, bool(rows)
response_state: Final = eventually(
read_response_state,
lambda state: state[0].status_code == 200 or state[1],
seconds=seconds,
return_last_on_timeout=return_last_on_timeout,
)
return response_state[0]
def _assert_rollup_row(
api_key: str,
database_url: str,
router_name: str,
expected_session_id: str,
) -> None:
key_hash: Final = hashlib.sha256(api_key.encode()).hexdigest()
rows: Final = read_rows(
"SELECT session_id, router_name, router_type, last_model, turns, spend, saved_spend, "
"savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, "
"savings_estimated_baseline_models, baseline_models, tier_turns "
'FROM "LiteLLM_AutoRouterSession" WHERE api_key = %s AND session_id = %s AND router_name = %s',
(key_hash, expected_session_id, router_name),
database_url=database_url,
)
assert len(rows) == 1, rows
assert rows[0] == {
"session_id": expected_session_id,
"router_name": router_name,
"router_type": "complexity",
"last_model": "openai/integration-simple",
"turns": 1,
"spend": pytest.approx(0.000015),
"saved_spend": pytest.approx(0.000135),
"savings_estimated_turns": 1,
"savings_estimated_actual_spend": pytest.approx(0.000015),
"savings_estimated_saved_spend": pytest.approx(0.000135),
"savings_estimated_baseline_models": {"openai/integration-baseline": 1},
"baseline_models": {"openai/integration-baseline": 1},
"tier_turns": {"SIMPLE": 1},
}
def _rollup_row(api_key: str, database_url: str, router_name: str, session_id: str) -> dict[str, object]:
key_hash: Final = hashlib.sha256(api_key.encode()).hexdigest()
rows: Final = read_rows(
"SELECT session_id, router_name, router_type, last_model, turns, spend, saved_spend, "
"savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, "
"savings_estimated_baseline_models, baseline_models, tier_turns "
'FROM "LiteLLM_AutoRouterSession" WHERE api_key = %s AND session_id = %s AND router_name = %s',
(key_hash, session_id, router_name),
database_url=database_url,
)
assert len(rows) == 1, rows
return rows[0]
def _session_contract(session_id: str, router_name: str) -> dict[str, object]:
return {
"session_id": session_id,
"router_name": router_name,
"router_type": "complexity",
"turns": 1,
"last_model": "openai/integration-simple",
"spend": pytest.approx(0.000015),
"savings_estimated_turns": 1,
"savings_estimated_actual_spend": pytest.approx(0.000015),
"saved_spend": pytest.approx(0.000135),
"baseline_spend": pytest.approx(0.00015),
"savings_estimated_baseline_spend": pytest.approx(0.00015),
"baseline_model": "openai/integration-baseline",
"baseline_models": {"openai/integration-baseline": 1},
}
@pytest.mark.timeout(600)
def test_auto_router_session_tracks_key_scoped_spend_and_bounds_long_session_ids(
gateway: Gateway,
tmp_path: Path,
) -> None:
normal_session_id: Final = str(uuid.uuid4())
messages_session_id: Final = str(uuid.uuid4())
responses_session_id: Final = str(uuid.uuid4())
codex_session_id: Final = str(uuid.uuid4())
codex_thread_id: Final = str(uuid.uuid4())
opencode_session_id: Final = str(uuid.uuid4())
thread_only_session_id: Final = str(uuid.uuid4())
long_session_id: Final = "s" * 300
router_name: Final = "auto-router-" + uuid.uuid4().hex[:8]
normal_prompt: Final = f"integration-auto-simple {normal_session_id}"
messages_prompt: Final = f"integration-auto-simple {messages_session_id}"
messages_second_prompt: Final = f"integration-auto-medium {messages_session_id}"
responses_prompt: Final = f"integration-auto-simple {responses_session_id}"
codex_prompt: Final = f"integration-auto-simple {codex_session_id}"
opencode_prompt: Final = f"integration-auto-simple {opencode_session_id}"
thread_only_prompt: Final = f"integration-auto-simple {thread_only_session_id}"
long_prompt: Final = f"integration-auto-simple {long_session_id}"
messages_upstream_body: Final = {
"model": "integration-simple",
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": messages_prompt}],
}
],
"include": ["reasoning.encrypted_content"],
"max_output_tokens": 32,
}
messages_second_upstream_body: Final = {
"model": "integration-medium",
"input": [
{
"type": "message",
"role": "user",
"content": [{"type": "input_text", "text": messages_second_prompt}],
}
],
"include": ["reasoning.encrypted_content"],
"max_output_tokens": 32,
}
responses_upstream_body: Final = {
"model": "integration-simple",
"input": responses_prompt,
}
with (
scratch_database() as database_url,
wire_server(
_tier_peer(
"integration-simple",
(normal_prompt, codex_prompt, opencode_prompt, thread_only_prompt, long_prompt),
)
) as simple,
wire_server(_tier_peer("integration-medium", ())) as medium,
wire_server(_tier_peer("integration-complex", (normal_prompt,))) as complex_tier,
wire_server(_tier_peer("integration-baseline", (normal_prompt,))) as baseline,
):
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["model_list"] = [
_deployment("integration-simple", "integration-simple", simple, 0.000001, 0.000002),
_deployment("integration-medium", "integration-medium", medium, 0.000002, 0.000004),
_deployment("integration-complex", "integration-complex", complex_tier, 0.000004, 0.000008),
_deployment("integration-baseline", "integration-baseline", baseline, 0.00001, 0.00002),
{
"model_name": router_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "heuristic",
"keyword_tier_rules": [
{"keywords": ["integration-auto-simple"], "tier": "SIMPLE"},
{"keywords": ["integration-auto-medium"], "tier": "MEDIUM"},
],
"tiers": {
"SIMPLE": "integration-simple",
"MEDIUM": "integration-medium",
"COMPLEX": "integration-complex",
"REASONING": "integration-baseline",
},
},
},
},
]
config["router_settings"] = {**config["router_settings"], "num_retries": 0, "disable_cooldowns": True}
config_path: Final = tmp_path / "auto_router_management.yaml"
config_path.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {"DATABASE_URL": database_url}, config=config_path, workers=1) as candidate:
with candidate.scenario() as scenario:
key: Final = scenario.key()
other_key: Final = scenario.key()
response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": router_name, "messages": [{"role": "user", "content": normal_prompt}]},
key=key,
headers={"x-claude-code-session-id": normal_session_id},
)
assert response.status_code == 200, response.text
assert _assistant_content(response) == "synthetic tier response", response.text
session: Final = _session_response(candidate, key, normal_session_id)
assert JSON_OBJECT.validate_json(session.content) == {
"session_id": normal_session_id,
"router_name": router_name,
"router_type": "complexity",
"turns": 1,
"last_model": "openai/integration-simple",
"spend": pytest.approx(0.000015),
"savings_estimated_turns": 1,
"savings_estimated_actual_spend": pytest.approx(0.000015),
"saved_spend": pytest.approx(0.000135),
"baseline_spend": pytest.approx(0.00015),
"savings_estimated_baseline_spend": pytest.approx(0.00015),
"baseline_model": "openai/integration-baseline",
"baseline_models": {"openai/integration-baseline": 1},
}, session.text
_assert_rollup_row(key, database_url, router_name, normal_session_id)
hidden: Final = candidate.request(
"GET",
"/auto_router/session",
params={"session_id": normal_session_id},
key=other_key,
)
assert hidden.status_code == 404, hidden.text
assert JSON_OBJECT.validate_json(hidden.content) == {
"detail": f"No auto-routed turns recorded for session {normal_session_id!r} under this key"
}, hidden.text
messages_response: Final = candidate.client.post(
"/v1/messages",
json={
"model": router_name,
"max_tokens": 32,
"messages": [{"role": "user", "content": messages_prompt}],
},
headers={
"x-api-key": key,
"anthropic-version": "2023-06-01",
"x-claude-code-session-id": messages_session_id,
},
)
assert messages_response.status_code == 200, messages_response.text
assert JSON_OBJECT.validate_json(messages_response.content)["content"] == [
{"type": "text", "text": "synthetic tier response"}
], messages_response.text
messages_second_response: Final = candidate.client.post(
"/v1/messages",
json={
"model": router_name,
"max_tokens": 32,
"messages": [{"role": "user", "content": messages_second_prompt}],
},
headers={
"x-api-key": key,
"anthropic-version": "2023-06-01",
"x-claude-code-session-id": messages_session_id,
},
)
assert messages_second_response.status_code == 200, messages_second_response.text
assert JSON_OBJECT.validate_json(messages_second_response.content)["content"] == [
{"type": "text", "text": "synthetic tier response"}
], messages_second_response.text
messages_session: Final = _session_response(candidate, key, messages_session_id)
assert JSON_OBJECT.validate_json(messages_session.content) == {
"session_id": messages_session_id,
"router_name": router_name,
"router_type": "complexity",
"turns": 2,
"last_model": "openai/integration-medium",
"spend": pytest.approx(0.000045),
"savings_estimated_turns": 2,
"savings_estimated_actual_spend": pytest.approx(0.000045),
"saved_spend": pytest.approx(0.000255),
"baseline_spend": pytest.approx(0.0003),
"savings_estimated_baseline_spend": pytest.approx(0.0003),
"baseline_model": "openai/integration-baseline",
"baseline_models": {"openai/integration-baseline": 2},
}, messages_session.text
assert _rollup_row(key, database_url, router_name, messages_session_id) == {
"session_id": messages_session_id,
"router_name": router_name,
"router_type": "complexity",
"last_model": "openai/integration-medium",
"turns": 2,
"spend": pytest.approx(0.000045),
"saved_spend": pytest.approx(0.000255),
"savings_estimated_turns": 2,
"savings_estimated_actual_spend": pytest.approx(0.000045),
"savings_estimated_saved_spend": pytest.approx(0.000255),
"savings_estimated_baseline_models": {"openai/integration-baseline": 2},
"baseline_models": {"openai/integration-baseline": 2},
"tier_turns": {"SIMPLE": 1, "MEDIUM": 1},
}
responses_response: Final = candidate.client.post(
"/v1/responses",
json={"model": router_name, "input": responses_prompt},
headers={
"Authorization": f"Bearer {key}",
"user-agent": "codex_cli_rs/0.50.0",
"session-id": responses_session_id,
"thread-id": str(uuid.uuid4()),
},
)
assert responses_response.status_code == 200, responses_response.text
responses_output: Final = JSON_ARRAY.validate_python(
JSON_OBJECT.validate_json(responses_response.content)["output"]
)
assert len(responses_output) == 1, responses_response.text
assert object_value(responses_output[0])["content"] == [
{
"type": "output_text",
"text": "synthetic tier response",
"annotations": [],
"logprobs": None,
}
], responses_response.text
responses_session: Final = _session_response(candidate, key, responses_session_id)
assert JSON_OBJECT.validate_json(responses_session.content) == _session_contract(
responses_session_id, router_name
), responses_session.text
_assert_rollup_row(key, database_url, router_name, responses_session_id)
codex_response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": router_name, "messages": [{"role": "user", "content": codex_prompt}]},
key=key,
headers={
"user-agent": "codex_cli_rs/0.50.0",
"session-id": codex_session_id,
"thread-id": codex_thread_id,
},
)
assert codex_response.status_code == 200, codex_response.text
assert _assistant_content(codex_response) == "synthetic tier response", codex_response.text
codex_session: Final = _session_response(candidate, key, codex_session_id)
assert JSON_OBJECT.validate_json(codex_session.content) == _session_contract(
codex_session_id, router_name
), codex_session.text
_assert_rollup_row(key, database_url, router_name, codex_session_id)
codex_thread: Final = candidate.request(
"GET",
"/auto_router/session",
params={"session_id": codex_thread_id},
key=key,
)
assert codex_thread.status_code == 404, codex_thread.text
assert JSON_OBJECT.validate_json(codex_thread.content) == {
"detail": f"No auto-routed turns recorded for session {codex_thread_id!r} under this key"
}, codex_thread.text
codex_hidden: Final = candidate.request(
"GET",
"/auto_router/session",
params={"session_id": codex_session_id},
key=other_key,
)
assert codex_hidden.status_code == 404, codex_hidden.text
assert JSON_OBJECT.validate_json(codex_hidden.content) == {
"detail": f"No auto-routed turns recorded for session {codex_session_id!r} under this key"
}, codex_hidden.text
opencode_response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": router_name, "messages": [{"role": "user", "content": opencode_prompt}]},
key=key,
headers={
"x-session-id": opencode_session_id,
"user-agent": "opencode/0.15.0",
},
)
assert opencode_response.status_code == 200, opencode_response.text
assert _assistant_content(opencode_response) == "synthetic tier response", opencode_response.text
opencode_session: Final = _session_response(
candidate,
key,
opencode_session_id,
return_last_on_timeout=True,
)
assert JSON_OBJECT.validate_json(opencode_session.content) == _session_contract(
opencode_session_id, router_name
), opencode_session.text
_assert_rollup_row(key, database_url, router_name, opencode_session_id)
opencode_hidden: Final = candidate.request(
"GET",
"/auto_router/session",
params={"session_id": opencode_session_id},
key=other_key,
)
assert opencode_hidden.status_code == 404, opencode_hidden.text
assert JSON_OBJECT.validate_json(opencode_hidden.content) == {
"detail": f"No auto-routed turns recorded for session {opencode_session_id!r} under this key"
}, opencode_hidden.text
thread_only_response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": router_name, "messages": [{"role": "user", "content": thread_only_prompt}]},
key=key,
headers={
"user-agent": "codex_cli_rs/0.50.0",
"thread-id": thread_only_session_id,
},
)
assert thread_only_response.status_code == 200, thread_only_response.text
assert _assistant_content(thread_only_response) == "synthetic tier response", thread_only_response.text
thread_only_session: Final = _session_response(candidate, key, thread_only_session_id)
assert JSON_OBJECT.validate_json(thread_only_session.content) == _session_contract(
thread_only_session_id, router_name
), thread_only_session.text
_assert_rollup_row(key, database_url, router_name, thread_only_session_id)
thread_only_hidden: Final = candidate.request(
"GET",
"/auto_router/session",
params={"session_id": thread_only_session_id},
key=other_key,
)
assert thread_only_hidden.status_code == 404, thread_only_hidden.text
assert JSON_OBJECT.validate_json(thread_only_hidden.content) == {
"detail": f"No auto-routed turns recorded for session {thread_only_session_id!r} under this key"
}, thread_only_hidden.text
long_response: Final = candidate.request(
"POST",
"/v1/chat/completions",
{"model": router_name, "messages": [{"role": "user", "content": long_prompt}]},
key=key,
headers={"x-litellm-session-id": long_session_id},
)
assert long_response.status_code == 200, long_response.text
assert _assistant_content(long_response) == "synthetic tier response", long_response.text
bounded_session_id: Final = "sha256:" + hashlib.sha256(long_session_id.encode()).hexdigest()
long_session: Final = _session_response(
candidate,
key,
long_session_id,
database_url,
bounded_session_id,
seconds=70,
return_last_on_timeout=True,
)
assert JSON_OBJECT.validate_json(long_session.content) == {
"session_id": long_session_id,
"router_name": router_name,
"router_type": "complexity",
"turns": 1,
"last_model": "openai/integration-simple",
"spend": pytest.approx(0.000015),
"savings_estimated_turns": 1,
"savings_estimated_actual_spend": pytest.approx(0.000015),
"saved_spend": pytest.approx(0.000135),
"baseline_spend": pytest.approx(0.00015),
"savings_estimated_baseline_spend": pytest.approx(0.00015),
"baseline_model": "openai/integration-baseline",
"baseline_models": {"openai/integration-baseline": 1},
}, long_session.text
_assert_rollup_row(key, database_url, router_name, bounded_session_id)
simple_calls: Final = tuple(request for request in simple.drain() if request.method == "POST")
assert tuple(
(
request.method,
request.target,
json.loads(request.body),
)
for request in simple_calls
) == (
(
"POST",
"/v1/chat/completions",
{"model": "integration-simple", "messages": [{"role": "user", "content": normal_prompt}]},
),
("POST", "/v1/responses", messages_upstream_body),
("POST", "/v1/responses", responses_upstream_body),
(
"POST",
"/v1/chat/completions",
{"model": "integration-simple", "messages": [{"role": "user", "content": codex_prompt}]},
),
(
"POST",
"/v1/chat/completions",
{"model": "integration-simple", "messages": [{"role": "user", "content": opencode_prompt}]},
),
(
"POST",
"/v1/chat/completions",
{"model": "integration-simple", "messages": [{"role": "user", "content": thread_only_prompt}]},
),
(
"POST",
"/v1/chat/completions",
{"model": "integration-simple", "messages": [{"role": "user", "content": long_prompt}]},
),
), response.text
medium_calls: Final = tuple(request for request in medium.drain() if request.method == "POST")
assert tuple((request.method, request.target, json.loads(request.body)) for request in medium_calls) == (
("POST", "/v1/responses", messages_second_upstream_body),
), response.text
assert tuple(request for request in complex_tier.drain() if request.method == "POST") == (), response.text
assert tuple(request for request in baseline.drain() if request.method == "POST") == (), response.text

588
tests/integration/routing/test_observed_routing.py
View file

@ -1,19 +1,28 @@
import json
import uuid
from collections.abc import Callable
from pathlib import Path
from queue import SimpleQueue
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, object_value
from integration._support.client import JSON_OBJECT, Gateway, object_value
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
from pydantic import JsonValue
@pytest.mark.covers("other.routing.retries.several_attempts_reach_success_without_hidden_retries", "other.routing.errors.nonretryable_and_exhausted_failures_remain_errors")
@pytest.mark.covers(
"other.routing.retries.several_attempts_reach_success_without_hidden_retries",
"other.routing.errors.nonretryable_and_exhausted_failures_remain_errors",
)
def test_retry_counts_and_public_errors_match_actual_provider_attempts(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
with (
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
gateway.scenario() as scenario,
):
original: Final = object_value(gateway.get("/router/settings")["current_values"])["num_retries"]
provider_model: Final = "errors-" + uuid.uuid4().hex
model: Final = scenario.model(model=f"openai/{provider_model}", input_cost_per_token=0, output_cost_per_token=0)
@ -25,12 +34,23 @@ def test_retry_counts_and_public_errors_match_actual_provider_attempts(gateway:
scenario.cleanups.callback(remove)
try:
for index, (retries, statuses, status, attempts) in enumerate(((2, [500, 500, 200], 200, 3), (2, [400, 200], 400, 1), (1, [429, 429, 200], 429, 2), (1, [500, 500, 200], 500, 2))):
for index, (retries, statuses, status, attempts) in enumerate(
(
(2, [500, 500, 200], 200, 3),
(2, [400, 200], 400, 1),
(1, [429, 429, 200], 429, 2),
(1, [500, 500, 200], 500, 2),
)
):
gateway.post("/config/update", {"router_settings": {"num_retries": retries}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == retries
upstream.post(f"/__scripts/{provider_model}", json={"statuses": statuses}).raise_for_status()
upstream.get("/__observations").raise_for_status()
response: Final = gateway.request("POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"{provider_model} {index}"}]})
response: Final = gateway.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"{provider_model} {index}"}]},
)
assert response.status_code == status, response.text
requests: Final = upstream.get("/__observations").json()["requests"]
assert len(requests) == attempts
@ -42,9 +62,12 @@ def test_retry_counts_and_public_errors_match_actual_provider_attempts(gateway:
error: Final = response.json()["error"]
assert isinstance(error["message"], str) and "Controlled provider failure" in error["message"]
assert str(error["code"]) == str(status)
assert error["type"] == {400: "invalid_request_error", 429: "throttling_error", 500: "internal_server_error"}[status]
assert (
error["type"]
== {400: "invalid_request_error", 429: "throttling_error", 500: "internal_server_error"}[status]
)
assert error["param"] is None
assert "Traceback" not in response.text and "File \"" not in response.text
assert "Traceback" not in response.text and 'File "' not in response.text
finally:
gateway.post("/config/update", {"router_settings": {"num_retries": original}})
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == original
@ -58,22 +81,74 @@ def test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity
model: Final = json.loads(request.body)["model"]
assert model in {"primary-wire", "fallback-wire", "unrelated-wire"}
if model == "primary-wire":
return Reply(status=500, body=b'{"error":{"message":"synthetic primary unavailable","type":"api_error","code":"500"}}')
return Reply(body=json.dumps({"id": "response-" + model, "object": "chat.completion", "created": 1, "model": model, "choices": [{"index": 0, "message": {"role": "assistant", "content": "served " + model}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15}}).encode())
return Reply(
status=500,
body=b'{"error":{"message":"synthetic primary unavailable","type":"api_error","code":"500"}}',
)
return Reply(
body=json.dumps(
{
"id": "response-" + model,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "served " + model},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
).encode()
)
with wire_server(respond) as wire:
path: Final = tmp_path / "fallback.yaml"
path.write_text(yaml.safe_dump({"model_list": [{"model_name": alias, "litellm_params": {"model": "openai/" + upstream, "api_key": "synthetic-routing-key", "api_base": wire.url + "/v1"}} for alias, upstream in (("primary", "primary-wire"), ("fallback", "fallback-wire"), ("unrelated", "unrelated-wire"))], "router_settings": {"num_retries": 0, "disable_cooldowns": True, "fallbacks": [{"primary": ["fallback"]}]}}))
path.write_text(
yaml.safe_dump(
{
"model_list": [
{
"model_name": alias,
"litellm_params": {
"model": "openai/" + upstream,
"api_key": "synthetic-routing-key",
"api_base": wire.url + "/v1",
},
}
for alias, upstream in (
("primary", "primary-wire"),
("fallback", "fallback-wire"),
("unrelated", "unrelated-wire"),
)
],
"router_settings": {
"num_retries": 0,
"disable_cooldowns": True,
"fallbacks": [{"primary": ["fallback"]}],
},
}
)
)
loaded: Final = yaml.safe_load(path.read_text())
router: Final = Router(model_list=loaded["model_list"], **loaded["router_settings"])
try:
result: Final = router.completion(model="primary", messages=[{"role": "user", "content": "fallback control"}])
result: Final = router.completion(
model="primary", messages=[{"role": "user", "content": "fallback control"}]
)
assert result.id == "response-fallback-wire"
assert result.choices[0].message.content == "served fallback-wire"
assert result.choices[0].finish_reason == "stop"
assert result.usage.prompt_tokens == 11 and result.usage.completion_tokens == 4
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("primary-wire", "fallback-wire")
control: Final = router.completion(model="unrelated", messages=[{"role": "user", "content": "independent route"}])
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == (
"primary-wire",
"fallback-wire",
)
control: Final = router.completion(
model="unrelated", messages=[{"role": "user", "content": "independent route"}]
)
assert control.id == "response-unrelated-wire"
assert tuple(json.loads(request.body)["model"] for request in wire.drain()) == ("unrelated-wire",)
finally:
@ -82,17 +157,494 @@ def test_loaded_fallback_selects_expected_deployment_and_keeps_response_identity
@pytest.mark.covers("other.routing.alias_update.persisted_target_changes_only_selected_route")
def test_saved_deployment_target_update_changes_wire_and_preserves_control(gateway: Gateway) -> None:
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
with (
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
gateway.scenario() as scenario,
):
prefix: Final = "target-" + uuid.uuid4().hex
model: Final = scenario.model(model="openai/" + prefix + "-first", input_cost_per_token=0, output_cost_per_token=0)
other: Final = scenario.model(model="openai/" + prefix + "-control", input_cost_per_token=0, output_cost_per_token=0)
model: Final = scenario.model(
model="openai/" + prefix + "-first", input_cost_per_token=0, output_cost_per_token=0
)
other: Final = scenario.model(
model="openai/" + prefix + "-control", input_cost_per_token=0, output_cost_per_token=0
)
target: Final = next(entry for entry in gateway.get("/model/info")["data"] if entry["model_name"] == model)
for generation, suffix in enumerate(("first", "second")):
if generation:
response: Final = gateway.request("PATCH", f"/model/{target['model_info']['id']}/update", {"litellm_params": {"model": "openai/" + prefix + "-second"}})
response: Final = gateway.request(
"PATCH",
f"/model/{target['model_info']['id']}/update",
{"litellm_params": {"model": "openai/" + prefix + "-second"}},
)
assert response.status_code == 200, response.text
upstream.get("/__observations").raise_for_status()
for alias in (model, other):
assert gateway.chat(alias, text=f"{prefix} generation {generation}")["usage"]["total_tokens"] == 40
requests: Final = upstream.get("/__observations").json()["requests"]
assert [request["body"]["model"] for request in requests] == [prefix + "-" + suffix, prefix + "-control"]
@pytest.mark.parametrize(
(
"case",
"primary_error",
"body_flag",
"expected_sequence",
"expected_model",
"expected_status",
"expected_attempted_fallbacks",
"stream_request",
"key_metadata",
),
(
("context", "context", None, ("primary-up", "big-up"), "big-up", 200, "1", False, False),
("content-policy", "content-policy", None, ("primary-up", "safe-up"), "safe-up", 200, "1", False, False),
("server-error", "server-error", None, ("primary-up", "general-up"), "general-up", 200, "1", False, False),
("disabled-context", "context", "disable_fallbacks", ("primary-up",), "primary-up", 400, None, False, False),
(
"disabled-server-error",
"server-error",
"disable_fallbacks",
("primary-up",),
"primary-up",
500,
None,
False,
False,
),
(
"disabled-mid-stream",
"mid-stream",
"disable_fallbacks",
("primary-up",),
"primary-up",
200,
"0",
True,
False,
),
("key-metadata-disabled", "context", None, ("primary-up",), "primary-up", 400, None, False, True),
("stream-context", "context", None, ("primary-up", "big-up"), "big-up", 200, "1", True, False),
(
"stream-mid-stream",
"mid-stream",
None,
("primary-up", "general-up"),
"general-up",
200,
"0",
True,
False,
),
("mock-context", None, "mock_testing_context_fallbacks", ("big-up",), "big-up", 200, "1", False, False),
(
"mock-content-policy",
None,
"mock_testing_content_policy_fallbacks",
("safe-up",),
"safe-up",
200,
"1",
False,
False,
),
("mock-server-error", None, "mock_testing_fallbacks", ("general-up",), "general-up", 200, "1", False, False),
("ordered-context", "context", None, ("primary-up", "big-up"), "big-up", 200, "1", False, False),
(
"ordered-content-policy",
"content-policy",
None,
("primary-up", "safe-up"),
"safe-up",
200,
"1",
False,
False,
),
(
"ordered-server-error",
"server-error",
None,
("primary-up", "primary-tier2-up"),
"primary-tier2-up",
200,
"1",
False,
False,
),
("body-context", "context", "body_fallbacks", ("primary-up", "big-up"), "big-up", 200, "1", False, False),
(
"body-content-policy",
"content-policy",
"body_fallbacks",
("primary-up", "safe-up"),
"safe-up",
200,
"1",
False,
False,
),
),
)
def test_health_routing_fallbacks_preserve_error_specific_targets_and_request_controls(
gateway: Gateway,
tmp_path: Path,
case: str,
primary_error: str | None,
body_flag: str | None,
expected_sequence: tuple[str, ...],
expected_model: str,
expected_status: int,
expected_attempted_fallbacks: str | None,
stream_request: bool,
key_metadata: bool,
) -> None:
prompt: Final = f"health-routing-{case}-{uuid.uuid4().hex}"
call_order: Final = SimpleQueue[str]()
def respond(model: str) -> Reply:
if stream_request:
served_chunk: Final = (
b"data: "
+ json.dumps(
{
"id": "response-" + model,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "served " + model},
"finish_reason": "stop",
}
],
}
).encode()
+ b"\n\n"
)
usage_chunk: Final = (
b"data: "
+ json.dumps(
{
"id": "response-" + model,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
).encode()
+ b"\n\n"
)
return Reply(
content_type="text/event-stream",
chunks=(served_chunk, usage_chunk, b"data: [DONE]\n\n"),
)
return Reply(
body=json.dumps(
{
"id": "response-" + model,
"object": "chat.completion",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "served " + model},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}
).encode()
)
ordered: Final = case.startswith("ordered-")
body_fallbacks: Final = case.startswith("body-")
def handler(model: str) -> Callable[[Request], Reply]:
def handle(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
actual_body: Final = json.loads(request.body)
assert request.target == "/v1/chat/completions", request.target
assert actual_body == {
"model": model,
"messages": [{"role": "user", "content": prompt}],
**({"stream": True, "stream_options": {"include_usage": True}} if stream_request else {}),
}, request.body.decode()
call_order.put(model)
if model == "primary-up" and primary_error == "context":
return Reply(
status=400,
body=b'{"error":{"message":"This model\'s maximum context length is 4096 tokens","type":"invalid_request_error","code":"context_length_exceeded"}}',
)
if model == "primary-up" and primary_error == "content-policy":
return Reply(
status=400,
body=b'{"error":{"message":"synthetic content policy rejection","type":"invalid_request_error","code":"content_policy_violation"}}',
)
if model == "primary-up" and primary_error == "server-error":
return Reply(
status=500,
body=b'{"error":{"message":"synthetic upstream failure","type":"api_error","code":"500"}}',
)
if model == "primary-up" and primary_error == "mid-stream":
empty_first: Final = (
b"data: "
+ json.dumps(
{
"id": "response-" + model,
"object": "chat.completion.chunk",
"created": 1,
"model": model,
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": ""},
"finish_reason": None,
}
],
}
).encode()
+ b"\n\n"
)
return Reply(
content_type="text/event-stream",
chunks=(empty_first, b":" + b"x" * 4_000_000 + b"\n\n", empty_first),
abort_after=2,
)
return respond(model)
return handle
with (
wire_server(handler("primary-up")) as primary,
wire_server(handler("primary-tier2-up")) as primary_tier2,
wire_server(handler("big-up")) as big,
wire_server(handler("safe-up")) as safe,
wire_server(handler("general-up")) as general,
):
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["model_list"] = [
*(
(
{
"model_name": "primary",
"litellm_params": {
"model": "openai/primary-up",
"api_base": primary.url + "/v1",
"api_key": "synthetic-health-routing-key",
"order": 1,
},
},
{
"model_name": "primary",
"litellm_params": {
"model": "openai/primary-tier2-up",
"api_base": primary_tier2.url + "/v1",
"api_key": "synthetic-health-routing-key",
"order": 2,
},
},
)
if ordered
else (
{
"model_name": "primary",
"litellm_params": {
"model": "openai/primary-up",
"api_base": primary.url + "/v1",
"api_key": "synthetic-health-routing-key",
},
},
)
),
*(
{
"model_name": alias,
"litellm_params": {
"model": f"openai/{upstream_model}",
"api_base": wire.url + "/v1",
"api_key": "synthetic-health-routing-key",
},
}
for alias, upstream_model, wire in (
("big", "big-up", big),
("safe", "safe-up", safe),
("general", "general-up", general),
)
),
]
config["router_settings"] = {
"num_retries": 0,
"disable_cooldowns": True,
**(
{}
if body_fallbacks
else {
"context_window_fallbacks": [{"primary": ["big"]}],
"content_policy_fallbacks": [{"primary": ["safe"]}],
}
),
"fallbacks": [{"primary": ["general"]}],
}
config["general_settings"] = {
**config.get("general_settings", {}),
"dangerously_allow_mock_testing_request_params": True,
}
path: Final = tmp_path / f"health-routing-{case}.yaml"
path.write_text(yaml.safe_dump(config))
with gateway.scenario() as scenario, owned_proxy(gateway, tmp_path, {}, config=path) as candidate:
request_key: Final = scenario.key(metadata={"disable_fallbacks": True}) if key_metadata else candidate.key
request_body: Final[dict[str, JsonValue]] = {
"model": "primary",
"messages": [{"role": "user", "content": prompt}],
**({"stream": True} if stream_request else {}),
**(
{body_flag: True}
if body_flag
in (
"disable_fallbacks",
"mock_testing_context_fallbacks",
"mock_testing_content_policy_fallbacks",
"mock_testing_fallbacks",
)
else {}
),
**(
{
"context_window_fallbacks": [{"primary": ["big"]}],
"content_policy_fallbacks": [{"primary": ["safe"]}],
}
if body_fallbacks
else {}
),
}
if stream_request:
with candidate.client.stream(
"POST",
"/v1/chat/completions",
json=request_body,
headers={"Authorization": f"Bearer {request_key}"},
) as response:
response_text: Final = response.read().decode()
response_status: Final = response.status_code
response_headers: Final = response.headers
assert response_status == expected_status, response_text
assert response_headers.get("x-litellm-attempted-fallbacks") == expected_attempted_fallbacks, (
response_text
)
assert tuple(call_order.get_nowait() for _ in range(call_order.qsize())) == expected_sequence, (
response_text
)
event_lines: Final = tuple(
line.removeprefix("data:") for line in response_text.splitlines() if line.startswith("data:")
)
events: Final = tuple(
JSON_OBJECT.validate_json(event) for event in event_lines if event.strip() != "[DONE]"
)
if expected_status == 200 and not (primary_error == "mid-stream" and body_flag == "disable_fallbacks"):
assert tuple({**event, "created": 0} if "created" in event else event for event in events) == (
*(
(
{
"id": "response-primary-up",
"object": "chat.completion.chunk",
"created": 0,
"model": "primary",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}],
},
)
if primary_error == "mid-stream"
else ()
),
{
"id": "response-" + expected_model,
"object": "chat.completion.chunk",
"created": 0,
"model": expected_model.removesuffix("-up"),
"choices": [
{
"index": 0,
"delta": {"role": "assistant", "content": "served " + expected_model},
}
],
},
{
"id": "response-" + expected_model,
"object": "chat.completion.chunk",
"created": 0,
"model": expected_model.removesuffix("-up"),
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
},
), response_text
return
assert len(events) == 2, response_text
first_event: Final = events[0]
assert ({**first_event, "created": 0} if "created" in first_event else first_event) == {
"id": "response-primary-up",
"object": "chat.completion.chunk",
"created": 0,
"model": "primary",
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}],
}, response_text
assert {key: events[1]["error"][key] for key in events[1]["error"] if key != "message"} == {
"type": None,
"param": None,
"code": "500",
}, response_text
message: Final = events[1]["error"]["message"]
assert isinstance(message, str), response_text
assert message.startswith("litellm.APIConnectionError: "), response_text
else:
non_stream_response: Final = candidate.request(
"POST", "/v1/chat/completions", request_body, key=request_key
)
assert non_stream_response.status_code == expected_status, non_stream_response.text
assert (
non_stream_response.headers.get("x-litellm-attempted-fallbacks") == expected_attempted_fallbacks
), non_stream_response.text
if expected_status == 400:
assert JSON_OBJECT.validate_json(non_stream_response.content) == {
"error": {
"message": (
"litellm.ContextWindowExceededError: litellm.BadRequestError: "
"ContextWindowExceededError: OpenAIException - This model's maximum context length "
"is 4096 tokens"
),
"type": "invalid_request_error",
"param": None,
"code": "400",
}
}, non_stream_response.text
elif expected_status == 500:
assert JSON_OBJECT.validate_json(non_stream_response.content) == {
"error": {
"message": "litellm.InternalServerError: InternalServerError: OpenAIException - "
"synthetic upstream failure",
"type": "internal_server_error",
"param": None,
"code": "500",
}
}, non_stream_response.text
else:
assert JSON_OBJECT.validate_json(non_stream_response.content) == {
"id": "response-" + expected_model,
"object": "chat.completion",
"created": 1,
"model": expected_model,
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "served " + expected_model},
"finish_reason": "stop",
"provider_specific_fields": {},
}
],
"usage": {"prompt_tokens": 11, "completion_tokens": 4, "total_tokens": 15},
}, non_stream_response.text
assert tuple(call_order.get_nowait() for _ in range(call_order.qsize())) == expected_sequence, (
non_stream_response.text
)

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

@ -2,18 +2,23 @@ import asyncio
import json
import threading
import uuid
from collections.abc import Callable
from pathlib import Path
from queue import SimpleQueue
from typing import Final
import pytest
import yaml
from hypothesis import Phase, example, given, settings
from hypothesis import strategies as st
from integration._support.client import Gateway, eventually
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.wire import Reply, wire_server
from integration._support.wire import Reply, Request, wire_server
from openai import OpenAI
from pydantic import JsonValue, TypeAdapter
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
def frame(identity: str, delta: dict, *, finish: str | None = None) -> bytes:
@ -165,6 +170,13 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat
model: Final = scenario.model(
api_base=wire.url + "/v1", input_cost_per_token=0.001, output_cost_per_token=0.002
)
eventually(
lambda: tuple(
object_value(entry).get("model_name") for entry in gateway.get("/model/info")["data"]
),
lambda model_names, model=model: model in model_names,
seconds=60,
)
with OpenAI(
api_key=gateway.key, base_url=str(gateway.client.base_url), timeout=5, max_retries=0
) as client:
@ -434,6 +446,100 @@ def sse_data_lines(text: str) -> tuple[str, ...]:
return tuple(line.removeprefix("data: ") for line in text.splitlines() if line.startswith("data: "))
def _native_anthropic_message(model: str, identity: str, text: str) -> dict[str, object]:
return {
"id": identity,
"type": "message",
"role": "assistant",
"model": model,
"content": [{"type": "text", "text": text}],
"stop_reason": "end_turn",
"stop_sequence": None,
"usage": {"input_tokens": 11, "output_tokens": 4},
}
def _native_responses_response(model: str, identity: str, text: str) -> dict[str, object]:
return {
"id": identity,
"object": "response",
"created_at": 1,
"status": "completed",
"model": model,
"output": [
{
"id": "msg_" + identity,
"type": "message",
"status": "completed",
"role": "assistant",
"content": [{"type": "output_text", "text": text, "annotations": []}],
}
],
"usage": {"input_tokens": 11, "output_tokens": 4, "total_tokens": 15},
}
def _native_sse_event(event: dict[str, object]) -> bytes:
return f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
def _native_anthropic_stream(model: str, identity: str, text: str) -> tuple[bytes, ...]:
return (
_native_sse_event(
{
"type": "message_start",
"message": {
**_native_anthropic_message(model, identity, ""),
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 11, "output_tokens": 0},
},
}
),
_native_sse_event({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
_native_sse_event(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": text},
}
),
_native_sse_event({"type": "content_block_stop", "index": 0}),
_native_sse_event(
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 4},
}
),
_native_sse_event({"type": "message_stop"}),
)
def _native_responses_stream(model: str, identity: str, text: str) -> tuple[bytes, ...]:
response: Final = _native_responses_response(model, identity, text)
return (
_native_sse_event(
{
"type": "response.created",
"sequence_number": 0,
"response": {**response, "status": "in_progress", "output": []},
}
),
_native_sse_event(
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "msg_" + identity,
"output_index": 0,
"content_index": 0,
"delta": text,
}
),
_native_sse_event({"type": "response.completed", "sequence_number": 2, "response": response}),
)
@pytest.mark.covers("other.streaming.usage.provider_cost_object_completes_stream_and_bills_total_cost")
def test_perplexity_stream_with_cost_breakdown_object_completes_and_bills_total_cost(gateway: Gateway) -> None:
identity: Final = "stream-cost-object-" + uuid.uuid4().hex
@ -594,6 +700,280 @@ def test_primary_stream_with_empty_first_chunk_then_disconnect_falls_back_and_bi
assert float(rows[0]["spend"]) == pytest.approx(0.019), rows
@pytest.mark.parametrize(
("route", "failure_mode"),
(
("messages", "non-streaming"),
("messages", "before-first-byte"),
("messages", "mid-stream"),
("messages", "request-body-fallbacks"),
("responses", "non-streaming"),
("responses", "before-first-byte"),
("responses", "mid-stream"),
("responses", "request-body-fallbacks"),
),
)
def test_native_messages_and_responses_routes_fall_back_and_bill_the_fallback(
gateway: Gateway,
tmp_path: Path,
route: str,
failure_mode: str,
) -> None:
identity: Final = f"native-fallback-{route}-{failure_mode}-{uuid.uuid4().hex}"
is_messages: Final = route == "messages"
path: Final = "/v1/messages" if is_messages else "/v1/responses"
primary_model: Final = "claude-sonnet-4-5-20250929" if is_messages else "gpt-5.6"
fallback_model: Final = "claude-opus-4-5-20251101" if is_messages else "gpt-5.6"
stream: Final = failure_mode in ("before-first-byte", "mid-stream")
client_body: Final = (
{
"model": "primary",
"max_tokens": 12,
"messages": [{"role": "user", "content": identity}],
"stream": stream,
**({"fallbacks": ["fallback"]} if failure_mode == "request-body-fallbacks" else {}),
}
if is_messages
else {
"model": "primary",
"input": identity,
"stream": stream,
**({"fallbacks": ["fallback"]} if failure_mode == "request-body-fallbacks" else {}),
}
)
primary_body: Final = (
{
"model": primary_model,
"max_tokens": 12,
"messages": [{"role": "user", "content": identity}],
"stream": stream,
}
if is_messages
else {"model": primary_model, "input": identity, "stream": stream}
)
fallback_body: Final = {**primary_body, "model": fallback_model}
call_order: Final[SimpleQueue[str]] = SimpleQueue()
received: Final[SimpleQueue[tuple[str, Request]]] = SimpleQueue()
fallback_text: Final = f"fallback response {identity}"
primary_text: Final = f"primary response {identity}"
def response_for(model: str, text: str) -> Reply:
if is_messages:
if stream:
return Reply(content_type="text/event-stream", chunks=_native_anthropic_stream(model, identity, text))
return Reply(body=json.dumps(_native_anthropic_message(model, identity, text)).encode())
if stream:
return Reply(content_type="text/event-stream", chunks=_native_responses_stream(model, identity, text))
return Reply(body=json.dumps(_native_responses_response(model, identity, text)).encode())
def upstream(name: str, model: str) -> Callable[[Request], Reply]:
def handle(request: Request) -> Reply:
if request.method == "GET":
return Reply(body=b'{"object":"list","data":[]}')
assert request.method == "POST" and request.target == path, (request.method, request.target)
expected_body: Final = primary_body if name == "primary" else fallback_body
assert json.loads(request.body) == expected_body, request.body.decode()
if is_messages:
assert request.headers.get("x-api-key") == "synthetic-native-fallback-key", request.headers
assert request.headers.get("anthropic-version") == "2023-06-01", request.headers
assert "authorization" not in request.headers, request.headers
else:
assert request.headers.get("authorization") == "Bearer synthetic-native-fallback-key", request.headers
call_order.put(name)
received.put((name, request))
if name == "primary" and failure_mode == "mid-stream":
primary_events: Final = (
_native_anthropic_stream(primary_model, identity + "-primary", primary_text)
if is_messages
else _native_responses_stream(primary_model, identity + "-primary", primary_text)
)
first_event: Final = primary_events[0]
second_event: Final = (
primary_events[1]
if is_messages
else _native_sse_event(
{
"type": "response.in_progress",
"sequence_number": 1,
"response": {
**_native_responses_response(primary_model, identity + "-primary", primary_text),
"status": "in_progress",
"output": [],
},
}
)
)
return Reply(
content_type="text/event-stream",
chunks=(first_event, second_event, first_event),
abort_after=2,
pause_between_chunks=0.05,
)
if name == "primary" and failure_mode in (
"non-streaming",
"before-first-byte",
"request-body-fallbacks",
):
return Reply(
status=500,
body=b'{"error":{"message":"synthetic primary failure","type":"server_error","code":"internal_error"}}',
)
return response_for(model, fallback_text if name == "fallback" else primary_text)
return handle
with (
wire_server(upstream("primary", primary_model)) as primary,
wire_server(upstream("fallback", fallback_model)) as fallback,
):
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
config["model_list"] = [
{
"model_name": model_group,
"litellm_params": {
"model": f"{'anthropic' if is_messages else 'openai'}/{model}",
"api_key": "synthetic-native-fallback-key",
"api_base": server.url if is_messages else server.url + "/v1",
"input_cost_per_token": 0.001,
"output_cost_per_token": 0.002,
},
}
for model_group, model, server in (
("primary", primary_model, primary),
("fallback", fallback_model, fallback),
)
]
config["router_settings"] = {
"num_retries": 0,
"disable_cooldowns": True,
"fallbacks": [] if failure_mode == "request-body-fallbacks" else [{"primary": ["fallback"]}],
}
proxy_config: Final = tmp_path / f"native-{route}-{failure_mode}.yaml"
proxy_config.write_text(yaml.safe_dump(config))
with owned_proxy(gateway, tmp_path, {}, config=proxy_config) as candidate:
request_headers: Final = {
"x-litellm-call-id": identity,
**(
{"x-api-key": candidate.key, "anthropic-version": "2023-06-01"}
if is_messages
else {"Authorization": f"Bearer {candidate.key}"}
),
}
with candidate.client.stream("POST", path, json=client_body, headers=request_headers) as response:
response_text: Final = response.read().decode()
assert response.status_code == 200, response_text
assert response.headers.get("x-litellm-attempted-fallbacks") == "1", response_text
assert tuple(call_order.get_nowait() for _ in range(call_order.qsize())) == ("primary", "fallback")
assert tuple(
(name, request.method, request.target, json.loads(request.body))
for name, request in (received.get_nowait() for _ in range(received.qsize()))
) == (("primary", "POST", path, primary_body), ("fallback", "POST", path, fallback_body))
if stream:
event_lines: Final = tuple(line for line in sse_data_lines(response_text) if line != "[DONE]")
events: Final = tuple(JSON_OBJECT.validate_json(line) for line in event_lines)
event_types: Final = tuple(event["type"] for event in events)
expected_types: Final = (
(
"message_start",
"content_block_start",
"content_block_delta",
"content_block_stop",
"message_delta",
"message_stop",
)
if is_messages
else ("response.created", "response.output_text.delta", "response.completed")
)
assert event_types.count("message_start" if is_messages else "response.created") == 1, response_text
assert event_types.count("message_stop" if is_messages else "response.completed") == 1, response_text
assert event_types[-1] == ("message_stop" if is_messages else "response.completed"), response_text
assert event_types == expected_types, response_text
output_text: Final = (
"".join(
string_value(object_value(event["delta"])["text"])
for event in events
if event["type"] == "content_block_delta"
)
if is_messages
else "".join(
string_value(event["delta"])
for event in events
if event["type"] == "response.output_text.delta"
)
)
assert output_text == fallback_text, response_text
else:
response_body: Final = JSON_OBJECT.validate_json(response.content)
expected_response: Final = (
_native_anthropic_message(fallback_model, identity, fallback_text)
if is_messages
else {
"id": response_body["id"],
"created_at": 1,
"error": None,
"incomplete_details": None,
"instructions": None,
"metadata": None,
"model": fallback_model,
"object": "response",
"output": [
{
"id": "msg_" + identity,
"content": [
{
"annotations": [],
"text": fallback_text,
"type": "output_text",
"logprobs": None,
}
],
"role": "assistant",
"status": "completed",
"type": "message",
"phase": None,
}
],
"parallel_tool_calls": None,
"temperature": None,
"tool_choice": None,
"tools": None,
"top_p": None,
"max_output_tokens": None,
"previous_response_id": None,
"reasoning": None,
"status": "completed",
"text": None,
"truncation": None,
"usage": {
"input_tokens": 11,
"input_tokens_details": None,
"output_tokens": 4,
"output_tokens_details": None,
"total_tokens": 15,
"cost": None,
},
"user": None,
"store": None,
}
)
assert response_body == expected_response, response_text
if not is_messages:
assert response_body["id"].startswith("resp_"), response_text
rows: Final = eventually(
lambda: read_rows(
"SELECT model, model_group, spend, prompt_tokens, completion_tokens, status "
'FROM "LiteLLM_SpendLogs" WHERE request_id=%s OR litellm_call_id=%s',
(identity, identity),
),
lambda values: len(values) == 1,
seconds=70,
)
assert rows[0]["model_group"] == "fallback", rows
assert rows[0]["status"] == "success", rows
assert (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == (11, 4), rows
assert float(rows[0]["spend"]) == pytest.approx(0.019), rows
@pytest.mark.covers("other.streaming.failure.truncated_transport_raises_and_control_recovers")
def test_truncated_http_stream_is_an_error_and_next_stream_succeeds() -> None:
import litellm