From b9079755bb62453a447e8d6d132b4b1e128e69ce Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:35:43 +0000 Subject: [PATCH 01/23] fix(ci): restore loadscope on the mcp-interactions shard so its MCP subscription cases do not race across workers (#46783) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-unit.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index fd188deeef8..c11caa56be1 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -589,7 +589,7 @@ jobs: reruns: 0 timeout-minutes: 20 job-timeout-minutes: 60 - dist: load + dist: loadscope - shard: proxy-infra test-path: |- From 77c22a4c015c84ed78e33835d86c9b022bf591c9 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:49:26 -0700 Subject: [PATCH 02/23] feat(proxy): add agent_access_default_deny for least-privilege agent access (#43979) * feat(proxy): add require_key_agent_access_defined for least-privilege agent access Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): align A2A model visibility assertions with supported routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep A2A model entries on their existing routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): point the listing model at the harness upstream Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): coerce require_key_agent_access_defined like the settings validator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): inject agent access lookups instead of patching the handler Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): rename require_key_agent_access_defined to agent_access_default_deny Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: ryan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 7 + .../auth/agent_permission_handler.py | 51 ++- litellm/proxy/proxy_server.py | 1 + .../test_agent_least_privilege.py | 338 ++++++++++++++++++ .../auth/test_agent_permission_handler.py | 180 ++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 6 files changed, 576 insertions(+), 7 deletions(-) create mode 100644 tests/integration/authorization/test_agent_least_privilege.py diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c3bb76e0f28..e4a0a5242b4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3027,6 +3027,13 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): "(see GitHub issue #12905)." ), ) + agent_access_default_deny: bool = Field( + False, + description=( + "When true, a key with no agent grants on itself or its team can reach no agents instead of all agents. " + "Proxy admins are exempt." + ), + ) user_url_validation: bool | None = Field( None, description=( diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index f671e1f9098..48213625c4f 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -9,9 +9,10 @@ import asyncio from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, Protocol, TypeAlias from fastapi import HTTPException +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.proxy._experimental.mcp_server.ui_session_utils import build_effective_auth_contexts @@ -47,6 +48,29 @@ class RestrictedAgentAccess: AgentAccess: TypeAlias = UnrestrictedAgentAccess | RestrictedAgentAccess +class AgentAccessLookup(Protocol): + def __call__(self, user_api_key_auth: UserAPIKeyAuth | None, /, *, strict: bool) -> Awaitable[AgentAccess]: ... + + +_DEFAULT_DENY_FLAG: Final = TypeAdapter(bool) + + +def agent_access_default_deny() -> bool: + from litellm.proxy.proxy_server import general_settings_view + + try: + return _DEFAULT_DENY_FLAG.validate_python(general_settings_view().get("agent_access_default_deny", False)) + except ValidationError: + return True + + +def _is_proxy_admin(user_api_key_auth: UserAPIKeyAuth | None) -> bool: + return user_api_key_auth is not None and user_api_key_auth.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN.value, + ) + + def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry @@ -80,7 +104,7 @@ class AgentRequestHandler: - If team has restrictions and key has restrictions: use intersection - If team has restrictions and key has none: inherit from team - If team has no restrictions: use key restrictions - - If no restrictions: allow all agents + - If no restrictions: allow all agents unless agent_access_default_deny is enabled for a non-admin key """ @staticmethod @@ -89,15 +113,22 @@ class AgentRequestHandler: resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, *, strict: bool = False, + default_deny: Callable[[], bool] = agent_access_default_deny, + lookup_key_access: AgentAccessLookup | None = None, + lookup_team_access: AgentAccessLookup | None = None, ) -> AgentAccess: """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" if managed_agent_policy(user_api_key_auth) is not None: return await _managed_actor_agent_access(user_api_key_auth) key_team_access: Final = await AgentRequestHandler.resolve_key_team_agent_access( - user_api_key_auth, strict=strict + user_api_key_auth, + strict=strict, + lookup_key_access=lookup_key_access, + lookup_team_access=lookup_team_access, ) - if strict and isinstance(key_team_access, UnrestrictedAgentAccess): + deny_undefined: Final = strict or (not _is_proxy_admin(user_api_key_auth) and default_deny()) + if deny_undefined and isinstance(key_team_access, UnrestrictedAgentAccess): return RestrictedAgentAccess(frozenset()) caller_access: Final = await AgentRequestHandler.agent_caller_access(user_api_key_auth, strict=strict) own_access: Final = _intersect_agent_access(key_team_access, caller_access) @@ -122,12 +153,18 @@ class AgentRequestHandler: user_api_key_auth: UserAPIKeyAuth | None, *, strict: bool = False, + lookup_key_access: AgentAccessLookup | None = None, + lookup_team_access: AgentAccessLookup | None = None, ) -> AgentAccess: try: - key_access: Final = await AgentRequestHandler.get_allowed_agents_for_key(user_api_key_auth, strict=strict) - team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team( - user_api_key_auth, strict=strict + key_lookup: Final = ( + AgentRequestHandler.get_allowed_agents_for_key if lookup_key_access is None else lookup_key_access ) + team_lookup: Final = ( + AgentRequestHandler._get_allowed_agents_for_team if lookup_team_access is None else lookup_team_access + ) + key_access: Final = await key_lookup(user_api_key_auth, strict=strict) + team_access: Final = await team_lookup(user_api_key_auth, strict=strict) except Exception as e: if strict: raise HTTPException(503, "Agent invocation policy is unavailable") from e diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e5bceb25ef3..60842648d16 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -18824,6 +18824,7 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro "cancel_on_disconnect": "Boolean", "disable_auto_add_proxy_admin_to_teams": "Boolean", "apply_user_budget_to_team_keys": "Boolean", + "agent_access_default_deny": "Boolean", "user_api_key_cache_max_size": "Integer", "transcribe_media_buckets": "List", } diff --git a/tests/integration/authorization/test_agent_least_privilege.py b/tests/integration/authorization/test_agent_least_privilege.py new file mode 100644 index 00000000000..8a7f1596379 --- /dev/null +++ b/tests/integration/authorization/test_agent_least_privilege.py @@ -0,0 +1,338 @@ +import json +import uuid +from contextlib import ExitStack +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +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, Request, Wire, wire_server + + +def _write_access_config(directory: Path, *, default_deny: bool, model_name: str, api_base: str) -> Path: + source: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **source, + "model_list": [ + { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "integration-provider-key", + "api_base": api_base, + }, + } + ], + "general_settings": { + **source["general_settings"], + "agent_access_default_deny": default_deny, + }, + } + path: Final = directory / f"agent-access-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _peer_response(request: Request, wire: Wire, agent_name: str) -> Reply: + if request.method == "GET": + assert request.target in ("/.well-known/agent-card.json", "/.well-known/agent.json") + return Reply( + body=json.dumps( + { + "protocolVersion": "0.3", + "name": agent_name, + "description": "Synthetic permission peer", + "version": "1.0.0", + "url": wire.url + "/", + "capabilities": {"streaming": False}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } + ).encode() + ) + assert request.method == "POST" and request.target == "/" + body: Final = json.loads(request.body) + assert body["jsonrpc"] == "2.0" and body["method"] == "message/send" + message: Final = body["params"]["message"] + assert message["role"] == "user" + assert message["kind"] == "message" + assert message["parts"] in ( + [{"kind": "text", "text": "synthetic ping"}], + [{"kind": "text", "text": "user: synthetic ping"}], + ) + return Reply( + body=json.dumps( + { + "jsonrpc": "2.0", + "id": body["id"], + "result": { + "kind": "message", + "role": "agent", + "messageId": message["messageId"] + "-out", + "parts": [{"kind": "text", "text": "synthetic pong"}], + }, + } + ).encode() + ) + + +def _create_agent(candidate: Gateway, wire: Wire, agent_name: str, cleanups: ExitStack) -> str: + created: Final = candidate.post( + "/v1/agents", + { + "agent_name": agent_name, + "agent_card_params": { + "protocolVersion": "0.3", + "name": agent_name, + "description": "Synthetic permission peer", + "version": "1.0.0", + "url": wire.url + "/", + "capabilities": {"streaming": False}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + }, + "litellm_params": {"make_public": False}, + }, + ) + agent_id: Final = string_value(created["agent_id"]) + cleanups.callback(_delete_agent, candidate, agent_id) + return agent_id + + +def _delete_agent(candidate: Gateway, agent_id: str) -> None: + deleted: Final = candidate.request("DELETE", f"/v1/agents/{agent_id}") + assert deleted.status_code == 200, deleted.text + assert read_rows('SELECT agent_id FROM "LiteLLM_AgentsTable" WHERE agent_id=%s', (agent_id,)) == [] + + +def _card(candidate: Gateway, agent_id: str, key: str) -> httpx.Response: + return candidate.request( + "GET", + f"/a2a/{agent_id}/.well-known/agent-card.json", + key=key, + headers={"Connection": "close"}, + ) + + +def _send(candidate: Gateway, agent_id: str, key: str, marker: str) -> httpx.Response: + return candidate.request( + "POST", + f"/a2a/{agent_id}", + { + "jsonrpc": "2.0", + "id": marker, + "method": "message/send", + "params": { + "message": { + "kind": "message", + "role": "user", + "messageId": marker + "-in", + "parts": [{"kind": "text", "text": "synthetic ping"}], + } + }, + }, + key=key, + headers={"Connection": "close"}, + ) + + +def _chat_completion(candidate: Gateway, agent_name: str, key: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + { + "model": f"a2a/{agent_name}", + "messages": [{"role": "user", "content": "synthetic ping"}], + }, + key=key, + headers={"Connection": "close"}, + ) + + +def _model_names(candidate: Gateway, path: str, key: str) -> frozenset[str]: + response: Final = candidate.request("GET", path, key=key, headers={"Connection": "close"}) + assert response.status_code == 200, response.text + payload: Final = response.json() + raw_entries: Final = payload if isinstance(payload, list) else payload["data"] + entries: Final = tuple(object_value(entry) for entry in raw_entries) + if path == "/model_group/info": + return frozenset(string_value(entry["model_group"]) for entry in entries) + return frozenset(string_value(entry["model_name"]) for entry in entries) + + +def _listed_agent_ids(candidate: Gateway, key: str) -> frozenset[str]: + response: Final = candidate.request("GET", "/v1/agents", key=key, headers={"Connection": "close"}) + assert response.status_code == 200, response.text + return frozenset(string_value(object_value(agent)["agent_id"]) for agent in response.json()) + + +def _assert_successful_send(response: httpx.Response, marker: str) -> None: + 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"]["messageId"] == marker + "-in-out", response.text + assert body["result"]["parts"] == [{"kind": "text", "text": "synthetic pong"}], response.text + + +def _assert_denial(response: httpx.Response, agent_id: str) -> None: + assert response.status_code == 403, response.text + assert response.json() == { + "detail": f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access." + }, response.text + + +@pytest.mark.timeout(180) +def test_agent_access_default_deny_denies_ungranted_key_but_keeps_explicit_grants( + gateway: Gateway, tmp_path: Path +) -> None: + agent_name: Final = "least-privilege-" + uuid.uuid4().hex + v2_model: Final = "listing-" + agent_name + config: Final = _write_access_config( + tmp_path, default_deny=True, model_name=v2_model, api_base=gateway.upstream_url + "/v1" + ) + + def upstream(request: Request) -> Reply: + return _peer_response(request, wire, agent_name) + + with ( + wire_server(upstream) as wire, + owned_proxy(gateway, tmp_path, {}, config=config, workers=2) as candidate, + candidate.scenario() as scenario, + ): + agent_id: Final = _create_agent(candidate, wire, agent_name, scenario.cleanups) + ungranted_team: Final = scenario.team() + key_a: Final = scenario.key(team_id=ungranted_team) + key_b: Final = scenario.key(object_permission={"agents": [agent_id]}) + granted_team: Final = scenario.team(object_permission={"agents": [agent_id]}) + key_c: Final = scenario.key(team_id=granted_team) + wire.drain() + + denied_cards: Final = tuple(_card(candidate, agent_id, key_a) for _ in range(8)) + _assert_denial(denied_cards[0], agent_id) + assert all(response.json() == denied_cards[0].json() for response in denied_cards), [ + response.text for response in denied_cards + ] + + denied_sends: Final = tuple( + _send(candidate, agent_id, key_a, f"denied-{index}-{uuid.uuid4().hex}") for index in range(8) + ) + _assert_denial(denied_sends[0], agent_id) + assert all(response.json() == denied_sends[0].json() for response in denied_sends), [ + response.text for response in denied_sends + ] + denied_completion: Final = _chat_completion(candidate, agent_name, key_a) + assert denied_completion.status_code == 403, denied_completion.text + assert denied_completion.json() == { + "error": { + "message": f"Agent '{agent_name}' is not allowed for your key/team. Contact proxy admin for access.", + "type": "permission_error", + "param": None, + "code": "403", + } + }, denied_completion.text + denied_agents: Final = tuple( + candidate.request("GET", f"/v1/agents/{agent_id}", key=key_a, headers={"Connection": "close"}) + for _ in range(8) + ) + _assert_denial(denied_agents[0], agent_id) + assert all(response.json() == denied_agents[0].json() for response in denied_agents), [ + response.text for response in denied_agents + ] + assert v2_model in _model_names(candidate, "/v2/model/info", candidate.key) + assert _listed_agent_ids(candidate, key_a) == frozenset() + assert f"a2a/{agent_name}" not in _model_names(candidate, "/v2/model/info", key_a) + assert f"a2a/{agent_name}" not in _model_names(candidate, "/model_group/info", key_a) + assert tuple(item for item in wire.drain() if item.method == "POST") == () + + for key in (key_b, key_c, candidate.key): + cards: Final = tuple(_card(candidate, agent_id, key) for _ in range(8)) + assert all(response.status_code == 200 for response in cards), [response.text for response in cards] + completion: Final = _chat_completion(candidate, agent_name, key) + assert completion.status_code == 200, completion.text + assert completion.json()["choices"][0]["message"]["content"] == "synthetic pong", completion.text + wire.drain() + for index in range(8): + marker: Final = f"granted-{index}-{uuid.uuid4().hex}" + _assert_successful_send(_send(candidate, agent_id, key, marker), marker) + posts: Final = tuple(item for item in wire.drain() if item.method == "POST") + assert len(posts) == 1, posts + if key != candidate.key: + assert agent_id in _listed_agent_ids(candidate, key) + assert f"a2a/{agent_name}" in _model_names(candidate, "/v2/model/info", key) + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize("workers", [2, 4]) +def test_agent_access_default_deny_toggles_at_runtime_through_config_api( + gateway: Gateway, tmp_path: Path, workers: int +) -> None: + agent_name: Final = "runtime-agent-" + uuid.uuid4().hex + + def upstream(request: Request) -> Reply: + return _peer_response(request, wire, agent_name) + + with ( + wire_server(upstream) as wire, + owned_proxy( + gateway, + tmp_path, + {}, + config=Path("tests/integration/proxy_config.yaml"), + workers=workers, + ) as candidate, + candidate.scenario() as scenario, + ): + agent_id: Final = _create_agent(candidate, wire, agent_name, scenario.cleanups) + ungranted_key: Final = scenario.key() + wire.drain() + assert _card(candidate, agent_id, ungranted_key).status_code == 200 + + enabled: Final = candidate.request( + "POST", + "/config/field/update", + { + "config_type": "general_settings", + "field_name": "agent_access_default_deny", + "field_value": "true", + }, + ) + assert enabled.status_code == 200, enabled.text + config_fields: Final = candidate.request( + "GET", + "/config/list", + params={"config_type": "general_settings"}, + ) + assert config_fields.status_code == 200, config_fields.text + setting: Final = next( + object_value(field) + for field in config_fields.json() + if object_value(field)["field_name"] == "agent_access_default_deny" + ) + assert setting["field_type"] == "Boolean" and setting["field_value"] == "true", setting + eventually( + lambda: tuple(_card(candidate, agent_id, ungranted_key) for _ in range(8)), + lambda responses: all(response.status_code == 403 for response in responses), + seconds=40, + ) + + disabled: Final = candidate.request( + "POST", + "/config/field/update", + { + "config_type": "general_settings", + "field_name": "agent_access_default_deny", + "field_value": False, + }, + ) + assert disabled.status_code == 200, disabled.text + eventually( + lambda: tuple(_card(candidate, agent_id, ungranted_key) for _ in range(8)), + lambda responses: all(response.status_code == 200 for response in responses), + seconds=40, + ) diff --git a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py index b5fa1cbc53d..45d3d7fb22e 100644 --- a/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/unit/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -15,6 +15,7 @@ from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentAccess, + AgentAccessLookup, AgentRequestHandler, RestrictedAgentAccess, UnrestrictedAgentAccess, @@ -43,6 +44,15 @@ def _agent_id(registry: AgentRegistry, agent_name: str) -> str: return agent.agent_id +def _lookup(result: AgentAccess | Exception) -> AgentAccessLookup: + async def lookup(_user_api_key_auth: UserAPIKeyAuth | None, /, *, strict: bool) -> AgentAccess: + if isinstance(result, Exception): + raise result + return result + + return lookup + + async def _single_context(user_api_key_auth: UserAPIKeyAuth) -> list[UserAPIKeyAuth]: return [user_api_key_auth] @@ -133,6 +143,176 @@ class TestAgentRequestHandler: ) assert result == UnrestrictedAgentAccess() + async def test_agent_access_default_deny_is_opt_in(self): + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + key_lookup: Final = _lookup(UnrestrictedAgentAccess()) + team_lookup: Final = _lookup(UnrestrictedAgentAccess()) + + assert ( + await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: False, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) + == UnrestrictedAgentAccess() + ) + assert await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: True, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) == RestrictedAgentAccess(frozenset()) + assert await AgentRequestHandler.resolve_agent_access( + None, + default_deny=lambda: True, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) == RestrictedAgentAccess(frozenset()) + + async def test_explicit_key_and_team_grants_work_with_key_agent_access_required(self): + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + unrestricted_lookup: Final = _lookup(UnrestrictedAgentAccess()) + + assert await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: True, + lookup_key_access=unrestricted_lookup, + lookup_team_access=_lookup(RestrictedAgentAccess(frozenset({"team-agent"}))), + ) == RestrictedAgentAccess(frozenset({"team-agent"})) + + assert await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: True, + lookup_key_access=_lookup(RestrictedAgentAccess(frozenset({"key-agent"}))), + lookup_team_access=unrestricted_lookup, + ) == RestrictedAgentAccess(frozenset({"key-agent"})) + + async def test_proxy_admin_is_exempt_from_key_agent_access_requirement(self): + admin: Final = UserAPIKeyAuth( + api_key="admin-key", + user_id="admin", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + string_admin: Final = admin.model_copy(update={"user_role": LitellmUserRoles.PROXY_ADMIN.value}) + key_lookup: Final = _lookup(UnrestrictedAgentAccess()) + team_lookup: Final = _lookup(UnrestrictedAgentAccess()) + + assert ( + await AgentRequestHandler.resolve_agent_access( + admin, + default_deny=lambda: True, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) + == UnrestrictedAgentAccess() + ) + assert ( + await AgentRequestHandler.resolve_agent_access( + string_admin, + default_deny=lambda: True, + lookup_key_access=key_lookup, + lookup_team_access=team_lookup, + ) + == UnrestrictedAgentAccess() + ) + + async def test_key_agent_access_requirement_precedes_agent_access_group_ceiling(self): + agent_key: Final = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id="caller-agent", + ) + resolve, _ = self._ceiling_resolver(frozenset({"agent-beta"})) + unrestricted_lookup: Final = _lookup(UnrestrictedAgentAccess()) + + assert await AgentRequestHandler.resolve_agent_access( + agent_key, + resolve, + default_deny=lambda: True, + lookup_key_access=unrestricted_lookup, + lookup_team_access=unrestricted_lookup, + ) == RestrictedAgentAccess(frozenset()) + + async def test_key_agent_access_requirement_fails_closed_on_lookup_errors(self): + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", team_id="test-team") + unrestricted_lookup: Final = _lookup(UnrestrictedAgentAccess()) + key_error_lookup: Final = _lookup(Exception("DB Error")) + + assert await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: True, + lookup_key_access=key_error_lookup, + lookup_team_access=unrestricted_lookup, + ) == RestrictedAgentAccess(frozenset()) + assert ( + await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: False, + lookup_key_access=key_error_lookup, + lookup_team_access=unrestricted_lookup, + ) + == UnrestrictedAgentAccess() + ) + + team_error_lookup: Final = _lookup(Exception("DB Error")) + key_grant_lookup: Final = _lookup(RestrictedAgentAccess(frozenset({"key-agent"}))) + assert await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: True, + lookup_key_access=key_grant_lookup, + lookup_team_access=team_error_lookup, + ) == RestrictedAgentAccess(frozenset()) + assert ( + await AgentRequestHandler.resolve_agent_access( + auth, + default_deny=lambda: False, + lookup_key_access=key_grant_lookup, + lookup_team_access=team_error_lookup, + ) + == UnrestrictedAgentAccess() + ) + + async def test_agent_access_default_deny_reads_general_settings(self, monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + monkeypatch.setattr(proxy_server, "prisma_client", None) + + with patch.object(proxy_server, "general_settings", {"agent_access_default_deny": True}): + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", auth) is False + + with patch.object(proxy_server, "general_settings", {}): + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", auth) is True + + @pytest.mark.parametrize( + ("settings", "expected_allowed"), + [ + pytest.param({"agent_access_default_deny": True}, False, id="true-bool"), + pytest.param({"agent_access_default_deny": "true"}, False, id="true-lower-string"), + pytest.param({"agent_access_default_deny": "True"}, False, id="true-capitalized-string"), + pytest.param({"agent_access_default_deny": 1}, False, id="one"), + pytest.param({"agent_access_default_deny": False}, True, id="false-bool"), + pytest.param({"agent_access_default_deny": "false"}, True, id="false-string"), + pytest.param({"agent_access_default_deny": 0}, True, id="zero"), + pytest.param({}, True, id="missing"), + pytest.param({"agent_access_default_deny": "not-a-bool"}, False, id="invalid-string"), + ], + ) + async def test_agent_access_default_deny_coerces_general_settings( + self, + monkeypatch: pytest.MonkeyPatch, + settings: dict[str, object], + expected_allowed: bool, + ): + from litellm.proxy import proxy_server + + auth: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + monkeypatch.setattr(proxy_server, "prisma_client", None) + + with patch.object(proxy_server, "general_settings", settings): + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", auth) is expected_allowed + async def test_disjoint_key_and_team_grants_deny_every_agent(self): """LIT-5143: a key restricted to one agent inside a team restricted to another must reach nothing. The empty intersection used to read as "no restrictions", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7abd84c4ad7..6f6eae1f98e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29649,6 +29649,12 @@ export interface components { * @default 1 */ admission_queue_timeout_seconds: number; + /** + * Agent Access Default Deny + * @description When true, a key with no agent grants on itself or its team can reach no agents instead of all agents. Proxy admins are exempt. + * @default false + */ + agent_access_default_deny: boolean; /** * Alert To Webhook Url * @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}` From ebac9479078b26a721e52b3eb45d7a77716c4460 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 15 Oct 2026 16:50:02 -0700 Subject: [PATCH 03/23] fix(rust-bridge): honor the configured request_timeout on native Messages (#46749) The route host now reports get_configured_request_timeout() in MessagesSettings, and prepare uses the call timeout first, then that setting, then the 600s default. Before this, litellm_settings.request_timeout and REQUEST_TIMEOUT were ignored on the native path, so a long turn that works on Python timed out after 10 minutes Co-authored-by: Claude Opus 5.5 --- .../crates/inference-messages/src/prepare.rs | 37 ++++++++++++++++++- .../crates/inference-messages/src/types.rs | 14 ++++++- litellm/rust_bridge/messages/route_host.py | 3 ++ .../rust_bridge/messages/test_route_host.py | 24 +++++++++++- 4 files changed, 75 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/inference-messages/src/prepare.rs b/litellm-rust/crates/inference-messages/src/prepare.rs index 3ec756b65d0..4b7873b4808 100644 --- a/litellm-rust/crates/inference-messages/src/prepare.rs +++ b/litellm-rust/crates/inference-messages/src/prepare.rs @@ -124,7 +124,7 @@ fn prepare_provider_request( url, body: transformed, environment, - timeout, + timeout: timeout.or_else(|| shaping.settings.request_timeout()), api_key: api_key.map(SecretValue::new), }) } @@ -253,6 +253,41 @@ mod tests { ); } + #[rstest] + #[case::nothing_configured(None, None, None)] + #[case::configured_setting(None, Some(1800.0), Some(Duration::from_secs(1800)))] + #[case::call_timeout_wins( + Some(Duration::from_secs(30)), + Some(1800.0), + Some(Duration::from_secs(30)) + )] + #[case::call_timeout_only(Some(Duration::from_secs(30)), None, Some(Duration::from_secs(30)))] + #[case::unusable_setting_ignored(None, Some(-1.0), None)] + #[case::zero_setting_ignored(None, Some(0.0), None)] + fn timeout_falls_back_to_the_configured_request_timeout( + shaping: MessagesShaping, + #[case] call_timeout: Option, + #[case] request_timeout: Option, + #[case] expected: Option, + ) { + let prepared = prepare(MessagesCall { + body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})), + api_key: Some("sk-test".into()), + api_base: Some("https://anthropic.test".into()), + custom_llm_provider: Some("anthropic".into()), + litellm_params: Default::default(), + extra_headers: None, + provider_specific_header: None, + timeout: call_timeout, + shaping: MessagesShaping { + settings: MessagesSettings { request_timeout, ..shaping.settings }, + ..shaping + }, + }) + .unwrap(); + assert_eq!(prepared.timeout, expected); + } + fn prepared_body(fields: Value, shaping: MessagesShaping) -> Result { prepare(MessagesCall { body: body(fields), diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 669f7b283a1..21eac0574ed 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -60,6 +60,16 @@ pub struct MessagesSettings { pub reasoning_auto_summary: bool, #[serde(default)] pub additional_drop_params: Vec, + #[serde(default)] + pub request_timeout: Option, +} + +impl MessagesSettings { + pub fn request_timeout(&self) -> Option { + self.request_timeout + .and_then(|secs| Duration::try_from_secs_f64(secs).ok()) + .filter(|timeout| !timeout.is_zero()) + } } #[cfg(test)] @@ -120,13 +130,15 @@ mod tests { }, "drop_params": true, "reasoning_auto_summary": true, - "additional_drop_params": ["metadata.user_id", "thinking"] + "additional_drop_params": ["metadata.user_id", "thinking"], + "request_timeout": 1800.0 }), MessagesShaping { settings: MessagesSettings { drop_params: true, reasoning_auto_summary: true, additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()], + request_timeout: Some(1800.0), }, capabilities: MessagesModelCapabilities { supports_reasoning: true, diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/messages/route_host.py index 69e27e78223..5dfd3bfd530 100644 --- a/litellm/rust_bridge/messages/route_host.py +++ b/litellm/rust_bridge/messages/route_host.py @@ -9,6 +9,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout from litellm.llms.anthropic.pass_through.utils import is_reasoning_auto_summary_enabled from litellm.rust_bridge import failures from litellm.rust_bridge.public_call import optional_str @@ -22,6 +23,7 @@ class MessagesSettings: drop_params: bool reasoning_auto_summary: bool additional_drop_params: Sequence[str] + request_timeout: float | None def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: @@ -69,5 +71,6 @@ def settings(kwargs: Mapping[str, object]) -> dict[str, object]: drop_params=_drop_params(kwargs), reasoning_auto_summary=is_reasoning_auto_summary_enabled(), additional_drop_params=_additional_drop_params(kwargs), + request_timeout=get_configured_request_timeout(), ) ) diff --git a/tests/unit/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/messages/test_route_host.py index e721b9624d2..43a17e5e593 100644 --- a/tests/unit/rust_bridge/messages/test_route_host.py +++ b/tests/unit/rust_bridge/messages/test_route_host.py @@ -1,10 +1,12 @@ from types import MappingProxyType from typing import Final -from litellm.rust_bridge.messages.route_host import response import pytest + import litellm +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS from litellm.rust_bridge.messages import route_host +from litellm.rust_bridge.messages.route_host import response def test_response_is_a_detached_public_messages_dict() -> None: @@ -31,6 +33,8 @@ def test_response_is_a_detached_public_messages_dict() -> None: def test_settings_project_caller_configuration_without_resolving_a_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "drop_params", False) monkeypatch.setattr(litellm, "reasoning_auto_summary", True) + monkeypatch.setattr(litellm, "request_timeout", 1800) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) projected: Final = route_host.settings({"drop_params": "true", "additional_drop_params": ["metadata.user_id"]}) @@ -38,9 +42,27 @@ def test_settings_project_caller_configuration_without_resolving_a_model(monkeyp "drop_params": True, "reasoning_auto_summary": True, "additional_drop_params": ("metadata.user_id",), + "request_timeout": 1800.0, } +@pytest.mark.parametrize( + ("timeout", "explicitly_set", "expected"), + [ + (DEFAULT_REQUEST_TIMEOUT_SECONDS, False, None), + (DEFAULT_REQUEST_TIMEOUT_SECONDS, True, DEFAULT_REQUEST_TIMEOUT_SECONDS), + (300, False, 300.0), + ], +) +def test_request_timeout_carries_only_a_configured_global_timeout( + monkeypatch: pytest.MonkeyPatch, timeout: float, explicitly_set: bool, expected: float | None +) -> None: + monkeypatch.setattr(litellm, "request_timeout", timeout) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", explicitly_set) + + assert route_host.settings({})["request_timeout"] == expected + + @pytest.mark.parametrize( ("global_flag", "kwargs", "expected"), [ From 1e5d83f6ce202898f53b4c3219fead6bd1f1e960 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 15 Oct 2026 16:53:36 -0700 Subject: [PATCH 04/23] ci: run mcp interactions and redis cluster iam integration tests only in circleci (#46780) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .circleci/config.yml | 13 +++++++++++++ .github/workflows/test-unit.yml | 13 ------------- tests/integration/run.py | 2 -- 3 files changed, 13 insertions(+), 15 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index 90bab08235a..0d442e32b0a 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1293,6 +1293,19 @@ jobs: uv venv --python 3.12 .venv-mcp-peer uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" + - when: + condition: + equal: [sdk, << parameters.suite >>] + steps: + - run: + name: Install redis-server 7.2.16 + command: | + mkdir -p /tmp/redis-bin + cid="$(docker create redis:7.2.16@sha256:0637954999d01b7c9ce9167db2da50656e2590d3b884f1c600c5f63bb6e6773c)" + docker cp "${cid}:/usr/local/bin/redis-server" /tmp/redis-bin/redis-server + docker rm "$cid" + /tmp/redis-bin/redis-server --version | grep -q 'v=7.2.16' + echo 'export PATH="/tmp/redis-bin:$PATH"' >> "$BASH_ENV" - start_postgres: image: postgres:16@sha256:e17e86066e5ef83e0952a9347f5c792b7ece00972e2aa787a6986f471b3dd3d5 server_args: "-c max_connections=300 -c shared_preload_libraries=pg_stat_statements -c pg_stat_statements.track=all -c pg_stat_statements.max=20000" diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index c11caa56be1..0558875bdad 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -565,7 +565,6 @@ jobs: - shard: mcp-elicitation test-path: |- - --override-ini=pythonpath=tests tests/unit/experimental_mcp_client/test_mcp_client.py tests/unit/proxy/_experimental/mcp_server/test_capabilities.py tests/unit/proxy/_experimental/mcp_server/test_interactions.py @@ -580,17 +579,6 @@ jobs: job-timeout-minutes: 60 dist: loadscope - - shard: mcp-interactions - test-path: |- - --override-ini=pythonpath=tests - tests/integration/mcp/test_interactions.py - python-version: "3.12" - workers: 4 - reruns: 0 - timeout-minutes: 20 - job-timeout-minutes: 60 - dist: loadscope - - shard: proxy-infra test-path: |- tests/unit/proxy/db @@ -1463,7 +1451,6 @@ jobs: tests/unit/caching/test_redis_connection_pool.py \ tests/unit/caching/test_redis_cluster_cache.py \ tests/unit/caching/test_evicted_client_closer.py \ - tests/integration/sdk/test_redis_cluster_iam_auth.py \ --tb=short -vv -rs -p no:rerunfailures -p no:pytest-retry --timeout=120 --durations=20 \ --cov=./litellm --cov-report=xml:coverage.xml diff --git a/tests/integration/run.py b/tests/integration/run.py index d6ddd9d5e89..7ab8d59c02c 100644 --- a/tests/integration/run.py +++ b/tests/integration/run.py @@ -33,8 +33,6 @@ GITHUB_FILES: Final = frozenset( "tests/integration/database/image/test_image_bedrock_realtime_extra.py", "tests/integration/database/image/test_offline_image_migration.py", "tests/integration/database/image/test_ui_image_serves_offline.py", - "tests/integration/mcp/test_interactions.py", - "tests/integration/sdk/test_redis_cluster_iam_auth.py", } ) From b602209ca6b33ab27b8d2332866a7b20d452e7d8 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 15 Oct 2026 16:56:40 -0700 Subject: [PATCH 05/23] test(integration): hand every harness server a port from one per-worker block (#46761) * test(integration): hand every harness server a port from one per-worker block About 45 places in the integration harness picked a port by binding port 0, reading the number back, closing the socket and starting a server on that number later. Another process could take the port in between, which surfaced as "Address already in use" flakes. Earlier fixes retried around the race in single call sites without removing it. tests/integration/_support/ports.py now owns every port the harness picks ahead of a server. Each pytest process claims a 200-port block between 10000 and 32767 by taking an exclusive flock on a per-block lock file in the system temp directory, so parallel runs on one machine never share a block and the lock goes away with the process. The range is above the fixed service ports and below the Linux and macOS ephemeral ranges. Ports go out in order, are never reused in a run, skip anything an outside process already holds, and a process that exhausts its block claims the next free one. The three _free_port copies, worker_endpoints' port-0 binds, free_loopback_port, the owned proxy's lost-port-race retry and the owned Redis retry loop are gone. Tests that restart a server on the same port now allocate that port up front instead of reading it back from the first server. TQ012 in scripts/check_test_quality.py rejects a port-0 bind read back with getsockname anywhere under tests/integration. * ci(integration): record listening sockets and the ephemeral port range * Revert "ci(integration): record listening sockets and the ephemeral port range" This reverts commit d3fafc3da2a628ea4e1d2d806e121d50bdaceb3b. --- scripts/check_test_quality.py | 56 ++++++++- tests/integration/README.md | 2 + tests/integration/_support/database_relay.py | 18 +-- tests/integration/_support/otlp_sink.py | 19 ++-- tests/integration/_support/ports.py | 61 ++++++++++ tests/integration/_support/process.py | 77 +++++-------- tests/integration/_support/redis_process.py | 106 ++++++------------ tests/integration/_support/workers.py | 17 ++- .../test_key_guardrail_opt_out_chaos.py | 9 +- ...st_responses_compact_id_ownership_chaos.py | 6 +- .../test_a2a_protocol_version_override.py | 10 +- .../test_granian_admission_saturation.py | 12 +- .../test_missing_body_param_status.py | 6 +- .../database/test_harness_lifecycle.py | 8 +- .../test_v1_migration_error_logging.py | 6 +- ...test_vector_store_file_list_managed_ids.py | 6 +- tests/integration/mcp/test_proxy_mcp_e2e.py | 10 +- ...zure_content_safety_dispatch_resilience.py | 6 +- .../observability/test_guardrail_effects.py | 7 +- .../test_guardrail_stream_scope_chaos.py | 12 +- .../observability/test_langtrace_delivery.py | 4 +- .../test_logging_only_scope_chaos.py | 6 +- .../test_passthrough_upstream_error_chaos.py | 5 +- .../test_straiker_v3_platform.py | 7 +- .../_assistants_route_errors_support.py | 7 -- tests/integration/providers/_coralbricks.py | 7 -- .../providers/_vertex_batch_rows.py | 7 -- ...anthropic_count_tokens_system_lift_wire.py | 10 +- ...st_anthropic_refusal_stop_details_chaos.py | 12 +- .../test_assistants_route_errors_wire.py | 4 +- .../test_bedrock_kimi_k3_cache_point_wire.py | 10 +- ...st_bedrock_mantle_context_overflow_wire.py | 5 +- .../test_bedrock_mantle_count_tokens_wire.py | 12 +- ..._bedrock_runtime_chat_completions_chaos.py | 13 +-- ...che_control_tool_call_marks_owned_proxy.py | 10 +- ...test_chat_completions_tool_search_chaos.py | 10 +- ...t_chatgpt_responses_caller_stream_chaos.py | 10 +- .../test_cloudflare_clef_omni_media_chaos.py | 10 +- .../test_coralbricks_provider_wire.py | 6 +- ...est_decisions_box_deployment_flag_chaos.py | 6 +- .../providers/test_decisions_chaos.py | 6 +- .../providers/test_decisions_wire.py | 18 +-- .../test_gemini_tts_voice_mapping_wire.py | 10 +- ...test_idle_upstream_reconnect_chaos_wire.py | 5 +- ...bridge_message_cache_control_wire_chaos.py | 11 +- ...es_bridge_prompt_cache_breakpoint_chaos.py | 8 +- ...idge_prompt_cache_breakpoint_wire_chaos.py | 11 +- ...test_responses_bridge_tool_search_chaos.py | 10 +- ...responses_minted_reasoning_replay_chaos.py | 10 +- .../test_sagemaker_embedding_chaos.py | 10 +- ...vertex_anthropic_structured_output_wire.py | 9 +- ...test_vertex_batch_row_line_errors_chaos.py | 4 +- ...ertex_batch_row_line_errors_owned_proxy.py | 6 +- .../test_vertex_batch_row_line_errors_wire.py | 6 +- .../test_vertex_partner_count_tokens_wire.py | 6 +- .../test_voyage_mongodb_host_wire.py | 11 +- ...est_declared_provider_bare_model_routes.py | 10 +- ...allback_hop_foreign_encrypted_reasoning.py | 5 +- tests/integration/sdk/test_http2_wire.py | 7 +- .../integration/sdk/test_owned_redis_ports.py | 33 +----- .../test_faked_responses_stream_chaos.py | 15 +-- .../test_responses_bridge_stream_chaos.py | 12 +- .../test_responses_bridge_stream_errors.py | 13 +-- ...tream_open_error_provider_headers_proxy.py | 12 +- .../integration/test_oci_proxy_integration.py | 10 +- .../unit/integration_support/test_process.py | 44 -------- tests/unit/test_check_test_quality.py | 18 +++ 67 files changed, 360 insertions(+), 575 deletions(-) create mode 100644 tests/integration/_support/ports.py diff --git a/scripts/check_test_quality.py b/scripts/check_test_quality.py index 09b87039048..6e82fe8c229 100644 --- a/scripts/check_test_quality.py +++ b/scripts/check_test_quality.py @@ -88,6 +88,13 @@ TQ011 An integration-test assertion that compares an elapsed time against a nu is a literal when it is built only from numbers, or names a module-level number. Only `assert` statements are examined, so a polling deadline in a `while` loop is left alone. Same path scope as TQ010. +TQ012 A port read back from a port-0 bind: `sock.bind((host, 0))` in a function + that also calls `getsockname`. The number is free only until the socket + closes, so a server started on it later can lose it to another process and + fail with "Address already in use". Take the port from `allocate_port()` in + tests/integration/_support/ports.py, which hands each pytest worker its own + block of ports. Every file under tests/integration is examined, including + tests/integration/_support. Every rule is suppressible with `# test-quality-ok: ` on the reported line, following the repo's `*-ok: ` convention. A suppression without a @@ -172,7 +179,7 @@ SUBPROCESS_SPAWNS: Final = frozenset(("run", "Popen", "check_output", "check_cal INTERPRETER_ISOLATION_FLAGS: Final = frozenset(("-I", "-P")) RULE_CODES: Final = frozenset(( - "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", "TQ010", "TQ011", + "TQ000", "TQ001", "TQ002", "TQ003", "TQ004", "TQ005", "TQ006", "TQ007", "TQ009", "TQ010", "TQ011", "TQ012", )) INTEGRATION_ROOT: Final = ("tests", "integration") @@ -181,6 +188,8 @@ SIGNAL_SENDERS: Final = frozenset(("os.kill", "os.killpg")) PROCESS_CONTROL_SIGNALS: Final = frozenset(("SIGKILL", "SIGSTOP", "SIGCONT")) ZERO_ARGUMENT_PROCESS_CONTROLS: Final = frozenset(("kill", "terminate", "suspend", "resume")) SIGNAL_METHOD: Final = "send_signal" +BIND_METHOD: Final = "bind" +PORT_READBACK: Final = "getsockname" CLOCK_READERS: Final = frozenset(( "time.monotonic", "time.perf_counter", "time.time", "time.monotonic_ns", "time.perf_counter_ns", "time.time_ns", @@ -877,6 +886,50 @@ def iter_wall_clock_bound_violations(path: Path, tree: ast.Module) -> Iterator[V ) +def is_integration_tree_path(path: Path) -> bool: + parts: Final = path.parts + return any(parts[index:index + 2] == INTEGRATION_ROOT for index in range(len(parts) - 1)) + + +def _calls_in_scope(node: ast.AST, scope: ast.AST) -> Iterator[tuple[ast.Call, ast.AST]]: + for child in ast.iter_child_nodes(node): + if isinstance(child, ast.Call): + yield child, scope + inner: Final = child if isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef)) else scope + yield from _calls_in_scope(child, inner) + + +def _binds_port_zero(call: ast.Call) -> bool: + if not isinstance(call.func, ast.Attribute) or call.func.attr != BIND_METHOD or not call.args: + return False + address: Final = call.args[0] + return ( + isinstance(address, ast.Tuple) + and len(address.elts) == 2 + and isinstance(address.elts[1], ast.Constant) + and address.elts[1].value == 0 + ) + + +def _reads_back_port(scope: ast.AST) -> bool: + return any(isinstance(node, ast.Attribute) and node.attr == PORT_READBACK for node in ast.walk(scope)) + + +def iter_port_readback_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: + if not is_integration_tree_path(path): + return + for call, scope in _calls_in_scope(tree, tree): + if _binds_port_zero(call) and _reads_back_port(scope): + yield Violation( + path, + call.lineno, + "TQ012", + "a port read back from a port-0 bind is free only until the socket closes, so a server started on " + "it later can lose it to another process; take the port from allocate_port() in " + f"tests/integration/_support/ports.py (suppress: `# {SUPPRESSION_TOKEN}: `)", + ) + + def iter_conftest_inventory_violations(path: Path, tree: ast.Module) -> Iterator[Violation]: if path.name != CONFTEST_NAME: return @@ -916,6 +969,7 @@ def check_file(path: Path) -> tuple[Violation, ...]: *iter_child_interpreter_violations(path, tree), *iter_process_signal_violations(path, tree), *iter_wall_clock_bound_violations(path, tree), + *iter_port_readback_violations(path, tree), ) if violation.line not in skip ) diff --git a/tests/integration/README.md b/tests/integration/README.md index 478ad1cd787..33263043489 100644 --- a/tests/integration/README.md +++ b/tests/integration/README.md @@ -28,6 +28,8 @@ Owned proxies disable `.env` loading by default. A contract testing dotenv behav Owned proxy workers configure a 32-thread executor before application startup, preserving additional worker startup hooks. Held synchronous provider calls must not consume a CPU-dependent executor limit before the rest of a fault burst can start. The proxy entrypoint imports LiteLLM inside `main`, so spawned workers begin supervision before importing the application. Worker healthchecks have a five-second deadline by default, so uvicorn kills a worker that stays paused longer than that. Tests never signal proxy processes themselves. `OwnedProxy.kill_worker` SIGKILLs a worker and returns once its replacement answers readiness with its own pid, `replacing_worker` does the same for a kill that has to land mid-burst, and `suspended` pauses a worker or the supervisor, which only works on a proxy launched with `healthcheck=WorkerHealthcheck.LENIENT`. Helper processes that are not the proxy stop through `stop_child`, `stop_processes` and `stop_process_group` and pause through `paused`. `scripts/check_test_quality.py` rule TQ010 rejects a raw `os.kill`, SIGKILL / SIGSTOP / SIGCONT, `.kill()`, `.terminate()`, `.suspend()` or `.send_signal()` anywhere under `tests/integration` outside `_support`, and TQ011 rejects an assertion that caps an elapsed time with a fixed number. Bound timings with `_support/timing.py` instead, relative to a measured control or to the injected delay. `run.py` measures how long this box takes to import the proxy application before it starts pytest, records it in `boot-calibration.json` beside the results and exports it as `INTEGRATION_IMPORT_SECONDS`. Readiness deadlines, graceful stops and the default per-test timeout are multiples of that measurement (`_support/budget.py`), so a slower box gets longer budgets and a two-worker proxy gets twice the readiness deadline of a one-worker one. Readiness requires responsive worker identities that are still alive and a positively identified owned listener. Unreadable socket records cannot establish ownership. Crash tests own their client tasks with `TaskGroup`, release held peers in `finally` and verify that the intended worker dies. A test requiring traffic on both workers prepares that traffic explicitly instead of assuming OS load balancing, for example by sending each half while the other worker is suspended +Any port the harness picks before the server that listens on it has started comes from `allocate_port()` in `_support/ports.py`. That covers owned proxies, scripted upstreams, fake providers, OTLP sinks, database relays and Redis. A server that binds port 0 itself and keeps that socket, such as `wire_server` without a port, needs no allocation. Each pytest process claims a block of 200 ports between 10000 and 32767 by holding an exclusive lock file per block in the system temp directory, so two runs on one machine, such as parallel worktrees, never share a block, and the OS drops the lock when the process exits. The range sits above the fixed service ports and below the Linux and macOS ephemeral ranges, so a port-0 bind or an outbound connection never lands in a block. Ports inside a block go out in order and are never handed out twice in a run, a port that something outside the harness already holds is skipped, and a process that uses up its block claims the next free one. A test that restarts a server on the same port allocates that port up front instead of reading it back from the first server. TQ012 in `scripts/check_test_quality.py` rejects reading a port back from a port-0 bind anywhere under `tests/integration` + The CircleCI workflow starts its own database and Redis, restricts test-phase egress to its owned services and writes JUnit plus an executed-node manifest. Missing setup, failed cleanup or a selected test with neither a passed call nor a skip fail qualification. Skipped nodes are listed under `skipped` in `execution.json`, so the skip reasons double as the open bug list. `GITHUB_FILES` in `run.py` lists the files excluded from the `run_integration.sh` selection There is no hand-maintained list of test cases. A positional argument is a file of the group or a pytest node id inside one (`path::test[param]`), so one cell of a parametrized file can run alone. The runner fails only when pytest fails, when collection errors, or when a selected file collects zero tests. Older tests still carry `@pytest.mark.covers(...)` decorators; the marker stays registered so they collect, but the IDs are not checked against anything and new tests should not use it. The GitHub Actions coverage census reads the `GROUPS` literal in `run.py` and treats every `tests/integration//test_*.py` file in a scheduled group as owned by CircleCI diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 2b04aa369fc..2e9f2ec4a41 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -1,5 +1,4 @@ import asyncio -import socket import threading import time from collections.abc import Generator @@ -7,23 +6,16 @@ from contextlib import contextmanager from typing import Final from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit -from pydantic import TypeAdapter from tests.integration._support.async_server import LoopbackServer, close_writer +from tests.integration._support.ports import allocate_port -PORT: Final = TypeAdapter(int) OUTAGE_SECONDS: Final = 10.0 -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return PORT.validate_python(reserve.getsockname()[1]) - - class DatabaseRelay: def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None: - self.port: Final = _free_port() + self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger @@ -83,7 +75,7 @@ class DatabaseRelay: class HeldStatementRelay: def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None: - self.port: Final = _free_port() + self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger @@ -144,7 +136,7 @@ class TriggerScanner: class DroppedConnectionRelay: def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None: - self.port: Final = _free_port() + self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger @@ -231,7 +223,7 @@ class FrontendStatementScanner: class StatementCountingRelay: def __init__(self, upstream_host: str, upstream_port: int) -> None: - self.port: Final = _free_port() + self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._loop: Final = asyncio.new_event_loop() diff --git a/tests/integration/_support/otlp_sink.py b/tests/integration/_support/otlp_sink.py index 9c5f0150b42..561159fe197 100644 --- a/tests/integration/_support/otlp_sink.py +++ b/tests/integration/_support/otlp_sink.py @@ -13,7 +13,6 @@ import datetime import json import os import signal -import socket import ssl import subprocess import sys @@ -35,6 +34,8 @@ import yaml from pydantic import JsonValue, TypeAdapter from typing_extensions import ReadOnly, TypedDict +from integration._support.ports import allocate_port + INTERNAL_MARKERS: Final = ("gen_ai.operation.name", "mcp.method.name", "litellm.guardrail_name") @@ -421,12 +422,6 @@ class ConnectSink: ca_pem: str -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return int(reserve.getsockname()[1]) - - def _pid_reachable(url: str) -> bool: try: return httpx.get(f"{url}/__pid", trust_env=False, timeout=2).status_code == 200 @@ -439,7 +434,7 @@ def owned_sinks(directory: Path) -> Iterator[SpanSinks]: from integration._support.process import group_members, signal_group, stop_root_process directory.mkdir(parents=True, exist_ok=True) - ports: Final = tuple(_free_port() for _ in range(3)) + ports: Final = tuple(allocate_port() for _ in range(3)) root: Final = Path(__file__).resolve().parents[3] with ExitStack() as stack: processes: Final = tuple( @@ -509,8 +504,8 @@ def _await_sink(url: str) -> None: @contextmanager def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]: - http_port: Final = _free_port() - grpc_port: Final = _free_port() + http_port: Final = allocate_port() + grpc_port: Final = allocate_port() with _spawn_sink(directory, "otlp-grpc-sink.log", ["--port", str(http_port), "--grpc-port", str(grpc_port)]): control_url: Final = f"http://127.0.0.1:{http_port}" _await_sink(control_url) @@ -519,8 +514,8 @@ def owned_grpc_sink(directory: Path) -> Iterator[GrpcSink]: @contextmanager def owned_connect_sink(directory: Path) -> Iterator[ConnectSink]: - http_port: Final = _free_port() - tunnel_port: Final = _free_port() + http_port: Final = allocate_port() + tunnel_port: Final = allocate_port() ca_dir: Final = directory / "mitm" with _spawn_sink( directory, diff --git a/tests/integration/_support/ports.py b/tests/integration/_support/ports.py new file mode 100644 index 00000000000..ad480dfeb9b --- /dev/null +++ b/tests/integration/_support/ports.py @@ -0,0 +1,61 @@ +import fcntl +import os +import socket +import tempfile +import threading +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +FIRST_PORT: Final = 10_000 +LAST_PORT: Final = 32_767 +BLOCK_SIZE: Final = 200 +LOCK_DIRECTORY: Final = Path(tempfile.gettempdir()) / "litellm-integration-ports" + + +def _blocks() -> Iterator[range]: + return (range(start, start + BLOCK_SIZE) for start in range(FIRST_PORT, LAST_PORT + 2 - BLOCK_SIZE, BLOCK_SIZE)) + + +def _claim_until_exit(block: range, directory: Path) -> bool: + descriptor: Final = os.open(directory / f"{block.start}.lock", os.O_RDWR | os.O_CREAT, 0o600) + try: + fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + os.close(descriptor) + return False + return True + + +def _bindable(port: int) -> bool: + with socket.socket() as probe: + try: + probe.bind(("127.0.0.1", port)) + except OSError: + return False + return True + + +def _free_ports(directory: Path) -> Iterator[int]: + directory.mkdir(parents=True, exist_ok=True) + for block in _blocks(): + if _claim_until_exit(block, directory): + yield from filter(_bindable, block) + + +class PortAllocator: + def __init__(self, directory: Path) -> None: + self._directory: Final = directory + self._guard: Final = threading.Lock() + self._ports: Final = _free_ports(directory) + + def allocate(self) -> int: + with self._guard: + port: Final = next(self._ports, None) + assert port is not None, ( + f"Every integration port block in {FIRST_PORT}-{LAST_PORT} is claimed (locks in {self._directory})" + ) + return port + + +allocate_port: Final = PortAllocator(LOCK_DIRECTORY).allocate diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 75277ba8cfc..3c12414553d 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -1,10 +1,8 @@ -import errno import json import threading import os import re import signal -import socket import subprocess import sys import time @@ -33,6 +31,7 @@ from pydantic import TypeAdapter from tests.integration._support.budget import import_seconds, readiness_for, readiness_seconds from tests.integration._support.client import GATEWAY_LIMITS, Gateway +from tests.integration._support.ports import allocate_port DB_PUSH: Final = ("--use_prisma_db_push",) MIGRATE_DEPLOY: Final = () @@ -406,16 +405,6 @@ def _stop(process: subprocess.Popen[bytes]) -> None: raise -_PORT_ATTEMPTS: Final = 3 -_BIND_COLLISION: Final = os.strerror(errno.EADDRINUSE).lower() - - -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - @dataclass(frozen=True, slots=True) class _Launch: process: subprocess.Popen[bytes] @@ -458,7 +447,7 @@ def _stop_launch(launch: _Launch) -> None: def _launch(command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path) -> _Launch: - port: Final = _free_port() + port: Final = allocate_port() log_path: Final = output / f"owned-proxy-{uuid.uuid4().hex}.log" started: Final = time.monotonic() with log_path.open("w") as log: @@ -473,10 +462,6 @@ def _launch(command: tuple[str, ...], root: Path, environment: Mapping[str, str] return _Launch(process, port, log_path, started) -def _lost_port_race(exit_code: int | None, log: Path) -> bool: - return exit_code is not None and _BIND_COLLISION in log.read_text().lower() - - def _owns_listener(launch: _Launch) -> bool: for member in group_members(launch.process.pid): try: @@ -600,42 +585,36 @@ def _wait_until_ready(launch: _Launch, workers: int = 1) -> BootPhases | None: ) -def _launch_until_bound( - command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, attempts: int, workers: int +def _launch_until_ready( + command: tuple[str, ...], root: Path, environment: Mapping[str, str], output: Path, workers: int ) -> _Launch: launch: Final = _launch(command, root, environment, output) try: phases: Final = _wait_until_ready(launch, workers) - exit_code: Final = launch.process.poll() - assert exit_code is None or (attempts > 1 and _lost_port_race(exit_code, launch.log)), ( - "Owned proxy exited before readiness" - ) + assert launch.process.poll() is None, "Owned proxy exited before readiness" except BaseException: _stop_launch(launch) raise - if exit_code is None: - assert phases is not None - launch.log.with_suffix(".startup.json").write_text( - json.dumps( - { - "root_pid": launch.process.pid, - "readiness_seconds": time.monotonic() - launch.started, - "logical_cpus": os.cpu_count(), - "workers": workers, - "spawned_at": _wall_clock(launch.started), - "phases": { - "prisma_started": phases.prisma_started, - "prisma_finished": phases.prisma_finished, - "listening": phases.listening, - "first_ready": phases.first_ready, - }, - } - ) - + "\n" + assert phases is not None + launch.log.with_suffix(".startup.json").write_text( + json.dumps( + { + "root_pid": launch.process.pid, + "readiness_seconds": time.monotonic() - launch.started, + "logical_cpus": os.cpu_count(), + "workers": workers, + "spawned_at": _wall_clock(launch.started), + "phases": { + "prisma_started": phases.prisma_started, + "prisma_finished": phases.prisma_finished, + "listening": phases.listening, + "first_ready": phases.first_ready, + }, + } ) - return launch - _stop_launch(launch) - return _launch_until_bound(command, root, environment, output, attempts - 1, workers) + + "\n" + ) + return launch def _proxy_root() -> Path: @@ -785,7 +764,7 @@ def owned_proxy_process( *database_setup, *extra_arguments, ) - launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS, workers) + launch: Final = _launch_until_ready(command, root, environment, output, workers) process: Final = launch.process try: with httpx.Client( @@ -825,7 +804,7 @@ def configless_proxy_process( str(healthcheck_seconds(WorkerHealthcheck.STRICT)), *DB_PUSH, ) - launch: Final = _launch_until_bound(command, _proxy_root(), environment, output, _PORT_ATTEMPTS, workers) + launch: Final = _launch_until_ready(command, _proxy_root(), environment, output, workers) try: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS @@ -858,7 +837,7 @@ def owned_gateway_image( "--timeout-worker-healthcheck", str(int(graceful_stop_seconds())), ) - launch: Final = _launch_until_bound(command, root, environment, output, _PORT_ATTEMPTS, workers) + launch: Final = _launch_until_ready(command, root, environment, output, workers) try: with httpx.Client( base_url=f"http://127.0.0.1:{launch.port}", timeout=15, trust_env=False, limits=GATEWAY_LIMITS @@ -1040,7 +1019,7 @@ def owned_upstream( directory: Path, *, tls: bool = False, port: int | None = None, inherited: Mapping[str, str] = os.environ ) -> Generator[UpstreamSlot]: certificate: Final = self_signed_certificate(directory) if tls else None - slot: Final = UpstreamSlot(directory, port or _free_port(), _harness_root(), certificate, inherited) + slot: Final = UpstreamSlot(directory, port or allocate_port(), _harness_root(), certificate, inherited) try: slot.start() yield slot diff --git a/tests/integration/_support/redis_process.py b/tests/integration/_support/redis_process.py index f5e7f509eb5..1a325d73659 100644 --- a/tests/integration/_support/redis_process.py +++ b/tests/integration/_support/redis_process.py @@ -1,7 +1,6 @@ import os import shutil import signal -import socket import subprocess import time import uuid @@ -14,18 +13,7 @@ from typing import Final, TextIO from redis import Redis from redis.exceptions import ConnectionError as RedisConnectionError -PORT_ATTEMPTS: Final = 5 - - -@dataclass(frozen=True, slots=True) -class PortUnavailable: - reason: str - - -def free_loopback_port() -> int: - with socket.socket() as reservation: - reservation.bind(("127.0.0.1", 0)) - return reservation.getsockname()[1] +from tests.integration._support.ports import allocate_port @dataclass @@ -60,13 +48,8 @@ class OwnedRedis: return int(result.stdout) def start(self) -> None: - unavailable: Final = self.try_start() - assert unavailable is None, unavailable.reason - - def try_start(self) -> PortUnavailable | None: assert self.process is None - if self.serves(): - return PortUnavailable(f"Another Redis already serves {self.host}:{self.port} before the owned one starts") + assert not self.serves(), f"Another Redis already serves {self.host}:{self.port} before the owned one starts" process: Final = subprocess.Popen( self.command, stdout=self.log, stderr=subprocess.STDOUT, start_new_session=True ) @@ -77,7 +60,7 @@ class OwnedRedis: code: Final = process.poll() if code is not None: self.process = None - return PortUnavailable( + raise AssertionError( f"Owned Redis exited with code {code} before readiness on {self.host}:{self.port}" ) try: @@ -86,13 +69,13 @@ class OwnedRedis: expected: Final = self.owned_pid() if actual == expected: self.server_pid = actual - return None + return try: process.wait(timeout=max(deadline - time.monotonic(), 0)) except subprocess.TimeoutExpired as error: raise AssertionError("Redis readiness reached a different process") from error self.process = None - return PortUnavailable(f"A different Redis answered on {self.host}:{self.port} during startup") + raise AssertionError(f"A different Redis answered on {self.host}:{self.port} during startup") except RedisConnectionError: pass assert time.monotonic() < deadline, "Owned Redis readiness deadline exceeded" @@ -148,33 +131,10 @@ class OwnedRedis: self.signal(signal.SIGCONT) -@contextmanager -def _started( - build: Callable[[int], OwnedRedis], - choose_port: Callable[[], int], - attempts: int, - log_name: str, -) -> Iterator[OwnedRedis]: - server: Final = build(choose_port()) - try: - unavailable: Final = server.try_start() - if unavailable is None: - yield server - return - assert attempts > 1, ( - f"Owned Redis found no usable port after {PORT_ATTEMPTS} attempts: {unavailable.reason} (log {log_name})" - ) - with _started(build, choose_port, attempts - 1, log_name) as started: - yield started - finally: - if server.process is not None: - server.stop() - - @contextmanager def owned_redis( directory: Path, - choose_port: Callable[[], int] = free_loopback_port, + choose_port: Callable[[], int] = allocate_port, inherited: Mapping[str, str] = os.environ, ) -> Iterator[OwnedRedis]: binary: Final = shutil.which("redis-server") @@ -198,31 +158,31 @@ def owned_redis( output.mkdir(parents=True, exist_ok=True) directory.mkdir(parents=True, exist_ok=True) with (output / f"owned-redis-{uuid.uuid4().hex}.log").open("w") as log: - - def build(port: int) -> OwnedRedis: - pid_file: Final = ( - str(directory / f"owned-redis-{port}.pid") - if binary - else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid" - ) - command: Final = ( - *prefix, - "--port", - str(port), - "--set-proc-title", - "no", - "--pidfile", - pid_file, - "--bind", - "0.0.0.0" if not binary else "127.0.0.1", - "--protected-mode", - "no", - "--save", - "", - "--appendonly", - "no", - ) - return OwnedRedis(host, port, command, log, pid_file) - - with _started(build, choose_port, PORT_ATTEMPTS, log.name) as server: + port: Final = choose_port() + pid_file: Final = ( + str(directory / f"owned-redis-{port}.pid") if binary else f"/tmp/integration-redis-{uuid.uuid4().hex}.pid" + ) + command: Final = ( + *prefix, + "--port", + str(port), + "--set-proc-title", + "no", + "--pidfile", + pid_file, + "--bind", + "0.0.0.0" if not binary else "127.0.0.1", + "--protected-mode", + "no", + "--save", + "", + "--appendonly", + "no", + ) + server: Final = OwnedRedis(host, port, command, log, pid_file) + try: + server.start() yield server + finally: + if server.process is not None: + server.stop() diff --git a/tests/integration/_support/workers.py b/tests/integration/_support/workers.py index 33e086576c5..e75b0b3c6bb 100644 --- a/tests/integration/_support/workers.py +++ b/tests/integration/_support/workers.py @@ -1,7 +1,6 @@ from __future__ import annotations import os -import socket import time import uuid from collections.abc import Callable, Iterator, Mapping @@ -23,6 +22,7 @@ from redis.exceptions import TimeoutError as RedisTimeoutError from tests.integration._support.client import Gateway from tests.integration._support.database import scratch_database +from tests.integration._support.ports import allocate_port from tests.integration._support.redis_process import owned_redis from tests.integration._support.routing import WORKER_DATABASE_PREFIX @@ -78,15 +78,12 @@ def wait_for_services(upstream: str, proxies: tuple[Gateway, ...], cache: Redis) def worker_endpoints() -> Mapping[str, str]: - with socket.socket() as upstream, socket.socket() as provider: - upstream.bind(("127.0.0.1", 0)) - provider.bind(("127.0.0.1", 0)) - return MappingProxyType( - { - "INTEGRATION_UPSTREAM_URL": f"http://127.0.0.1:{upstream.getsockname()[1]}", - "INTEGRATION_PROVIDER_PORT": str(provider.getsockname()[1]), - } - ) + return MappingProxyType( + { + "INTEGRATION_UPSTREAM_URL": f"http://127.0.0.1:{allocate_port()}", + "INTEGRATION_PROVIDER_PORT": str(allocate_port()), + } + ) def _replica_urls(database: str) -> Mapping[str, str]: diff --git a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py index 78b226f619b..e136691aa2e 100644 --- a/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py +++ b/tests/integration/authorization/test_key_guardrail_opt_out_chaos.py @@ -1,5 +1,4 @@ import json -import socket import threading import time import uuid @@ -15,6 +14,7 @@ from integration._support.client import Gateway, eventually, string_value from integration._support.database import read_rows from integration._support.process import owned_proxy, owned_proxy_process from integration._support.timing import assert_cut_short, timed +from integration._support.ports import allocate_port from integration.authorization._guardrail_opt_out import ( chat, guardrail_config, @@ -34,14 +34,9 @@ class _GuardrailSink: self._action: Final = action self._server: ThreadingHTTPServer | None = None self._thread: threading.Thread | None = None - self._port: Final = self._claim_port() + self._port: Final = allocate_port() self.start() - def _claim_port(self) -> int: - with socket.socket() as probe: - probe.bind(("127.0.0.1", 0)) - return probe.getsockname()[1] - @property def url(self) -> str: return f"http://127.0.0.1:{self._port}" diff --git a/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py index 66af1fa3f26..d6ff66dc7d1 100644 --- a/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py +++ b/tests/integration/authorization/test_responses_compact_id_ownership_chaos.py @@ -15,6 +15,7 @@ import pytest from integration._support.client import Gateway, eventually from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Wire, wire_server +from integration._support.ports import allocate_port from integration.authorization._responses_compact import ( ALIASES, OTHER_TEAM_DETAIL, @@ -158,9 +159,9 @@ def _settled(target: _Target) -> bool: def test_upstream_outage_mid_burst_keeps_refusals_and_lands_each_owner_compact_once(gateway: Gateway) -> None: upstream: Final = CompactUpstream() created: Final = marker("mint") + port: Final = allocate_port() with gateway.scenario() as scenario: - with wire_server(upstream.respond) as first: - port: Final = urlsplit(first.url).port + with wire_server(upstream.respond, port=port) as first: actors: Final = tenants(scenario, first) target: Final = _Target(gateway, actors, mint(gateway, actors.owner, actors.model, created), vendor_id()) upstreamed(first) @@ -168,7 +169,6 @@ def test_upstream_outage_mid_burst_keeps_refusals_and_lands_each_owner_compact_o before_received: Final = upstreamed(first) during: Final = _send_all(target, _calls("during", 16)) alive: Final = gateway.request("GET", "/health/liveliness") - assert port is not None with wire_server(upstream.respond, port=port) as second: eventually(lambda: _settled(target), lambda settled: settled, seconds=60) after: Final = _send_all(target, _calls("after", 16)) diff --git a/tests/integration/compatibility/test_a2a_protocol_version_override.py b/tests/integration/compatibility/test_a2a_protocol_version_override.py index 299259714c4..cc36fd58e89 100644 --- a/tests/integration/compatibility/test_a2a_protocol_version_override.py +++ b/tests/integration/compatibility/test_a2a_protocol_version_override.py @@ -10,13 +10,13 @@ from __future__ import annotations import json import os -import socket import time from typing import Final import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows +from integration._support.ports import allocate_port from integration.compatibility._a2a_peers import ( HINT, KIND_ERROR, @@ -284,14 +284,8 @@ def test_registered_agent_persists_protocol_version(gateway: Gateway) -> None: _assert_legacy_wire(posts(agent.wire), agent.marker) # type: ignore[arg-type] -def _closed_local_port() -> int: - with socket.socket() as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def test_pinned_agent_survives_localhost_retry(gateway: Gateway) -> None: - dead: Final = f"http://localhost:{_closed_local_port()}/a2a/unreachable" + dead: Final = f"http://localhost:{allocate_port()}/a2a/unreachable" with peer("lgretry", "langgraph", card_url=dead) as agent, gateway.scenario() as scenario: identity: Final = register(gateway, scenario.cleanups.callback, agent.marker, agent.url, PINNED) _send_ok(gateway, identity, agent.marker) diff --git a/tests/integration/compatibility/test_granian_admission_saturation.py b/tests/integration/compatibility/test_granian_admission_saturation.py index 02c349b2c03..b8c1a2512f6 100644 --- a/tests/integration/compatibility/test_granian_admission_saturation.py +++ b/tests/integration/compatibility/test_granian_admission_saturation.py @@ -1,6 +1,5 @@ import asyncio import os -import socket import subprocess import sys import time @@ -12,22 +11,17 @@ import pytest from tests._master_key import MASTER_KEY from tests.integration._support.process import stop_child +from tests.integration._support.ports import allocate_port MAX_IN_FLIGHT: Final = 8 MAX_QUEUED: Final = 8 REQUESTS: Final = 200 -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: - listener.bind(("127.0.0.1", 0)) - return int(listener.getsockname()[1]) - - @pytest.mark.asyncio async def test_granian_admission_control_saturation(tmp_path: Path) -> None: - fake_port: Final = _free_port() - proxy_port: Final = _free_port() + fake_port: Final = allocate_port() + proxy_port: Final = allocate_port() fake_script: Final = Path(__file__).parents[2] / "_fake_openai_endpoint_server.py" config_path: Final = tmp_path / "saturation_config.yaml" proxy_env: Final = {**{k: v for k, v in os.environ.items() if k != "DATABASE_URL"}, "LITELLM_MASTER_KEY": MASTER_KEY} diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index 32338940931..42ee063ddc7 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import json import os -import socket import subprocess import sys import uuid @@ -24,6 +23,7 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy_process, paused, stop_child from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from openai import APIStatusError, AsyncOpenAI, BadRequestError, OpenAI from openai.types import Video from openai.types.eval_create_response import EvalCreateResponse as _Eval @@ -1504,9 +1504,7 @@ def _healthy(url: str) -> int: @contextmanager def _upstream(directory: Path) -> Iterator[tuple[subprocess.Popen[bytes], str]]: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - port: Final = int(reserve.getsockname()[1]) + port: Final = allocate_port() root: Final = Path(__file__).resolve().parents[2] url: Final = f"http://127.0.0.1:{port}" with (directory / "upstream.log").open("w") as log: diff --git a/tests/integration/database/test_harness_lifecycle.py b/tests/integration/database/test_harness_lifecycle.py index 032e00ba854..c8de39c0d38 100644 --- a/tests/integration/database/test_harness_lifecycle.py +++ b/tests/integration/database/test_harness_lifecycle.py @@ -44,6 +44,7 @@ from tests.integration._support.database_relay import ( from tests.integration._support.runtime import configure_executor from tests.integration._support.sharding import ShardExecution, shard_errors from tests.integration._support.wire import Reply, wire_server +from tests.integration._support.ports import allocate_port from tests.integration.run import _without_provider_api_keys @@ -143,10 +144,11 @@ async def _unused_peer(reader: asyncio.StreamReader, writer: asyncio.StreamWrite def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> None: + port: Final = allocate_port() with socket.socket() as occupied: - occupied.bind(("127.0.0.1", 0)) + occupied.bind(("127.0.0.1", port)) occupied.listen() - relay: Final = LoopbackServer(occupied.getsockname()[1], _unused_peer) + relay: Final = LoopbackServer(port, _unused_peer) with pytest.raises(OSError, match=os.strerror(errno.EADDRINUSE).lower()): relay.start() relay.stop() @@ -156,7 +158,7 @@ def test_relay_bind_failure_closes_its_loop_and_allows_repeated_cleanup() -> Non def test_upstream_readiness_failure_reaps_the_started_process(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: monkeypatch.setattr(process_support, "_UPSTREAM_READY_SECONDS", 0) children_before: Final = frozenset(child.pid for child in psutil.Process().children(recursive=True)) - slot: Final = process_support.UpstreamSlot(tmp_path, process_support._free_port(), process_support._harness_root()) + slot: Final = process_support.UpstreamSlot(tmp_path, allocate_port(), process_support._harness_root()) try: with pytest.raises(AssertionError, match="readiness deadline exceeded"): slot.start() diff --git a/tests/integration/database/test_v1_migration_error_logging.py b/tests/integration/database/test_v1_migration_error_logging.py index 1fe78fb68f3..3ca2179d983 100644 --- a/tests/integration/database/test_v1_migration_error_logging.py +++ b/tests/integration/database/test_v1_migration_error_logging.py @@ -25,7 +25,7 @@ from psycopg import sql from psycopg.types.json import Jsonb from tests.integration._support.client import eventually -from tests.integration._support.process import _free_port as free_port +from tests.integration._support.ports import allocate_port from tests.integration._support.process import stop_process_group REPO_ROOT: Final = Path(__file__).resolve().parents[3] @@ -94,7 +94,7 @@ def _unreachable_database_url(database_url: str) -> str: parsed.path.lstrip("/"), encode_password=False, ) - return _replace_port(url, free_port()) + return _replace_port(url, allocate_port()) @contextmanager @@ -388,7 +388,7 @@ def test_unreachable_database_recovers_after_postgres_forwarder_starts( port: Final = admin.port assert hostname is not None and port is not None target_host: Final = "127.0.0.1" if hostname == "localhost" else hostname - forwarding_port: Final = free_port() + forwarding_port: Final = allocate_port() forwarded_url: Final = _replace_port(database_url, forwarding_port, "127.0.0.1") command, environment = _migration_invocation(forwarded_url, tmp_path, {}) output_path: Final = _migration_log_path(request.node.name, tmp_path) diff --git a/tests/integration/management/test_vector_store_file_list_managed_ids.py b/tests/integration/management/test_vector_store_file_list_managed_ids.py index fed1b3eaa1f..d5765827969 100644 --- a/tests/integration/management/test_vector_store_file_list_managed_ids.py +++ b/tests/integration/management/test_vector_store_file_list_managed_ids.py @@ -22,6 +22,7 @@ from integration._support.client import Gateway, Scenario, eventually, object_va from integration._support.database import read_rows from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from openai import APIStatusError, AsyncOpenAI, OpenAI from pydantic import JsonValue, TypeAdapter @@ -1202,7 +1203,8 @@ def test_provider_outage_mid_burst_fails_loudly_and_mapping_resumes_after_recove provider_a: Final = _provider_file_id(bearer, "a.txt") respond: Final = _provider(store, _constant_listing(store, provider_a)) with gateway.scenario() as scenario: - with wire_server(respond) as wire: + port: Final = allocate_port() + with wire_server(respond, port=port) as wire: model: Final = scenario.model(api_base=wire.url + "/v1", api_key=bearer) _wait_until_every_worker_serves(gateway, model) wire.drain() @@ -1212,8 +1214,6 @@ def test_provider_outage_mid_burst_fails_loudly_and_mapping_resumes_after_recove served: Final = _burst(gateway, store, member.key, model, 40) assert [_ids(_listed(response)) for response in served] == [(managed_a,)] * 40 assert sum(1 for request in wire.drain() if request.method == "GET") == 40 - port: Final = urlsplit(wire.url).port - assert port is not None failed: Final = _burst(gateway, store, member.key, model, 20) assert [response.status_code for response in failed] == [500] * 20, [r.text for r in failed[:3]] for response in failed: diff --git a/tests/integration/mcp/test_proxy_mcp_e2e.py b/tests/integration/mcp/test_proxy_mcp_e2e.py index 959d2620f58..2f690463143 100644 --- a/tests/integration/mcp/test_proxy_mcp_e2e.py +++ b/tests/integration/mcp/test_proxy_mcp_e2e.py @@ -45,6 +45,7 @@ from tests.integration._support.process import stop_child from tests.integration._support.timing import assert_enforced_within from tests.integration._support.wire import Reply, Wire, wire_server from tests.integration._support.wire import Request as WireRequest +from tests.integration._support.ports import allocate_port CONFIG_TEMPLATE_PATH = Path("tests/integration/mcp/_config_mcp_e2e.yaml") MCP_SERVER_SCRIPT = Path("tests/integration/mcp/_math_server.py") @@ -294,8 +295,9 @@ def _start_proxy_server( sock: Final = socket_factory() startup.callback(sock.close) sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - sock.bind(("127.0.0.1", 0)) - host, port = sock.getsockname() + host: Final = "127.0.0.1" + port: Final = allocate_port() + sock.bind((host, port)) config = uvicorn.Config(proxy_app, host=host, port=port, log_level="warning", lifespan="off") server = uvicorn.Server(config) @@ -369,9 +371,7 @@ def test_proxy_startup_failure_closes_acquired_resources(tmp_path: Path) -> None @contextmanager def _math_http_server(offset: int) -> typing.Iterator[str]: host = "127.0.0.1" - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: - sock.bind((host, 0)) - _, port = sock.getsockname() + port: Final = allocate_port() with tempfile.TemporaryFile() as server_log: process = subprocess.Popen( diff --git a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py index d443ab9a2d0..14eb687f69f 100644 --- a/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py +++ b/tests/integration/observability/test_azure_content_safety_dispatch_resilience.py @@ -1,5 +1,4 @@ import concurrent.futures -import socket import threading import uuid from contextlib import ExitStack @@ -13,6 +12,7 @@ from integration._support.client import Gateway, Scenario, eventually from integration._support.database import read_rows from integration._support.redis_process import owned_redis from integration._support.wire import Wire, wire_server +from integration._support.ports import allocate_port from integration.observability.azure_dispatch_support import ( ATTACK_MARKER, AzureBehavior, @@ -251,9 +251,7 @@ def test_x1_azure_server_stop_restart_during_mixed_burst( with ExitStack() as resources: redis: Final = resources.enter_context(owned_redis(tmp_path)) provider: Final = resources.enter_context(wire_server(provider_handler)) - with socket.socket() as reservation: - reservation.bind(("127.0.0.1", 0)) - azure_port: Final = int(reservation.getsockname()[1]) + azure_port: Final = allocate_port() first_azure_lifetime: Final = ExitStack() try: azure: Final = first_azure_lifetime.enter_context(wire_server(azure_handler(), port=azure_port)) diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index dca9c2ec8e2..f267cff04f3 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -1,7 +1,6 @@ import asyncio import json import re -import socket import threading import uuid from collections.abc import Callable, Iterator, Mapping @@ -22,6 +21,7 @@ from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from openai import AsyncOpenAI, OpenAI from pydantic import BaseModel, JsonValue @@ -1285,10 +1285,7 @@ def _blocked_stream_events(text: str) -> tuple[dict[str, object], ...]: def _dead_api_base() -> str: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - port: Final = reserve.getsockname()[1] - return f"http://127.0.0.1:{port}/v1" + return f"http://127.0.0.1:{allocate_port()}/v1" @pytest.mark.covers("other.observability.guardrails.responses_pre_call_denial_openai_sdk_streams_typed_message") diff --git a/tests/integration/observability/test_guardrail_stream_scope_chaos.py b/tests/integration/observability/test_guardrail_stream_scope_chaos.py index 6aa705e52b3..3ada2fb7753 100644 --- a/tests/integration/observability/test_guardrail_stream_scope_chaos.py +++ b/tests/integration/observability/test_guardrail_stream_scope_chaos.py @@ -3,7 +3,6 @@ from __future__ import annotations import json import os import shutil -import socket import subprocess import uuid from collections.abc import Iterator, Mapping, Sequence @@ -23,6 +22,7 @@ from integration._support.client import Gateway, JsonValue, eventually from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import TypeAdapter ChaosEndpoint: TypeAlias = Literal["chat", "messages", "responses"] @@ -455,12 +455,6 @@ def _owned_proxy( yield owned -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return int(reserve.getsockname()[1]) - - def _assert_down_wave( plans: Sequence[CallPlan], responses: Sequence[httpx.Response], @@ -485,7 +479,7 @@ def _assert_down_wave( def test_h1_sink_outage_keeps_scope_isolated_through_recovery(rig: ChaosRig, tmp_path: Path) -> None: - port_a: Final = _free_port() + port_a: Final = allocate_port() started: Final = Event() unavailable: Final = Event() release: Final = Event() @@ -763,7 +757,7 @@ def _owned_postgres(directory: Path) -> Iterator[PostgresCluster]: ) assert docker_info.returncode == 0, docker_info.stderr directory.mkdir(parents=True, exist_ok=True) - port: Final = _free_port() + port: Final = allocate_port() container_name: Final = f"litellm-stream-scope-h5-{uuid.uuid4().hex}" password: Final = uuid.uuid4().hex created: Final = subprocess.run( diff --git a/tests/integration/observability/test_langtrace_delivery.py b/tests/integration/observability/test_langtrace_delivery.py index 93006cb6c6b..f81780f9f72 100644 --- a/tests/integration/observability/test_langtrace_delivery.py +++ b/tests/integration/observability/test_langtrace_delivery.py @@ -21,6 +21,7 @@ from integration._support.client import Gateway, eventually from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process from integration._support.timing import assert_not_delayed_by, timed from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from openai import AsyncOpenAI, OpenAI from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest from opentelemetry.proto.trace.v1.trace_pb2 import Span, Status @@ -595,8 +596,7 @@ def test_langtrace_two_workers_deliver_every_burst_span_exactly_once(gateway: Ga def test_langtrace_sink_outage_mid_burst_recovers_on_the_same_port(gateway: Gateway, tmp_path: Path) -> None: key: Final = "synthetic-langtrace-key-" + uuid.uuid4().hex - with wire_server(_accepted) as probe: - port: Final = int(probe.url.rsplit(":", 1)[1]) + port: Final = allocate_port() host: Final = f"http://127.0.0.1:{port}" with wire_server(_upstream) as provider: overrides: Final = {"LANGTRACE_API_KEY": key, "LANGTRACE_API_HOST": host, "OTEL_BSP_SCHEDULE_DELAY": "300"} diff --git a/tests/integration/observability/test_logging_only_scope_chaos.py b/tests/integration/observability/test_logging_only_scope_chaos.py index a65e56b930b..48e9930d948 100644 --- a/tests/integration/observability/test_logging_only_scope_chaos.py +++ b/tests/integration/observability/test_logging_only_scope_chaos.py @@ -1,6 +1,5 @@ from __future__ import annotations -import socket import threading import time import uuid @@ -46,6 +45,7 @@ from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy, owned_proxy_process from integration._support.timing import Timed, assert_not_delayed_by, timed from integration._support.wire import Reply, Request +from integration._support.ports import allocate_port from openai import APIConnectionError as OpenAIAPIConnectionError from pydantic import JsonValue @@ -60,9 +60,7 @@ def test_K1_policy_edge_restart_mid_burst_keeps_output_observation_fail_open(gat def policy(_request: Request) -> Reply: return Reply(body=b'{"action":"NONE"}') - with socket.socket() as reservation: - reservation.bind(("127.0.0.1", 0)) - policy_port: Final = reservation.getsockname()[1] + policy_port: Final = allocate_port() with gateway.scenario() as scenario: deployments: Final = _chaos_models(scenario, marker) models: Final = tuple(deployment.model_name for deployment in deployments) diff --git a/tests/integration/observability/test_passthrough_upstream_error_chaos.py b/tests/integration/observability/test_passthrough_upstream_error_chaos.py index b4a8524d3f8..68ad3d50a80 100644 --- a/tests/integration/observability/test_passthrough_upstream_error_chaos.py +++ b/tests/integration/observability/test_passthrough_upstream_error_chaos.py @@ -15,6 +15,7 @@ from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.process import WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue _GENERATE_CONTENT: Final[dict[str, JsonValue]] = {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]} @@ -120,8 +121,8 @@ async def _fire_burst( async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gateway: Gateway, tmp_path: Path) -> None: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) path: Final = tmp_path / "chaos-outage.yaml" - with wire_server(_chaos_reply) as wire: - port: Final = int(wire.url.rsplit(":", 1)[1]) + port: Final = allocate_port() + with wire_server(_chaos_reply, port=port) as wire: config["environment_variables"] = {"GEMINI_API_BASE": wire.url, "GEMINI_API_KEY": "scripted"} path.write_text(yaml.safe_dump(config)) with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: diff --git a/tests/integration/observability/test_straiker_v3_platform.py b/tests/integration/observability/test_straiker_v3_platform.py index a419d343535..c078cbb0c61 100644 --- a/tests/integration/observability/test_straiker_v3_platform.py +++ b/tests/integration/observability/test_straiker_v3_platform.py @@ -10,7 +10,6 @@ from __future__ import annotations import hashlib import itertools import json -import socket import threading import uuid from collections.abc import Callable, Iterator @@ -30,6 +29,7 @@ from integration._support.client import Gateway, eventually, gateway_from_enviro from integration._support.database import read_rows from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port V3_KEY: Final = "sk_agt_synthetic_integration_key" V1_KEY: Final = "synthetic-v1-collection-key" @@ -426,10 +426,7 @@ class Rig: @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: root: Final = tmp_path_factory.mktemp("straiker") - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - port: Final = reserve.getsockname()[1] - sink: Final = Sink(port) + sink: Final = Sink(allocate_port()) sink.start() with gateway_from_environment() as gateway, wire_server(_provider) as provider: config: Final = _rig_config(sink.url, root) diff --git a/tests/integration/providers/_assistants_route_errors_support.py b/tests/integration/providers/_assistants_route_errors_support.py index 0a622818c2c..0307a8c59a9 100644 --- a/tests/integration/providers/_assistants_route_errors_support.py +++ b/tests/integration/providers/_assistants_route_errors_support.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio import json import re -import socket import threading import uuid from collections.abc import AsyncIterator, Callable, Generator, Mapping @@ -84,12 +83,6 @@ def new_marker() -> str: return f"m{uuid.uuid4().hex}" -def free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - def proxy_path(route: Route, marker: str) -> str: return route.path.format(marker=marker) diff --git a/tests/integration/providers/_coralbricks.py b/tests/integration/providers/_coralbricks.py index 0cfbcd94d64..e6449df9591 100644 --- a/tests/integration/providers/_coralbricks.py +++ b/tests/integration/providers/_coralbricks.py @@ -2,7 +2,6 @@ from __future__ import annotations import json import re -import socket import uuid from collections.abc import Callable, Mapping from dataclasses import dataclass @@ -380,12 +379,6 @@ def raw_deployment(scenario: Scenario, model_name: str, litellm_params: Mapping[ return model_name -def closed_port() -> int: - with socket.socket() as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def provider_calls(wire: Wire) -> tuple[Request, ...]: return tuple(request for request in wire.drain() if not is_readiness_probe(request)) diff --git a/tests/integration/providers/_vertex_batch_rows.py b/tests/integration/providers/_vertex_batch_rows.py index 9d4f9c57e75..7e47bd198f3 100644 --- a/tests/integration/providers/_vertex_batch_rows.py +++ b/tests/integration/providers/_vertex_batch_rows.py @@ -1,5 +1,4 @@ import json -import socket import threading from collections.abc import Callable from email.message import Message @@ -132,12 +131,6 @@ def stored_custom_ids(request: Request) -> tuple[str, ...]: ) -def closed_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def multipart_file(request: Request) -> bytes: envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body parts: Final[tuple[Message, ...]] = tuple(BytesParser(policy=HTTP).parsebytes(envelope).iter_parts()) diff --git a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py index a4b5c15ad54..c3cd4c5b0e7 100644 --- a/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py +++ b/tests/integration/providers/test_anthropic_count_tokens_system_lift_wire.py @@ -1,6 +1,5 @@ import asyncio import json -import socket import threading import uuid from collections.abc import Callable, Iterator, Mapping @@ -18,6 +17,7 @@ import pytest from anthropic.types import MessageParam from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from integration.providers._count_tokens_system_lift import ( ANTHROPIC_VERSION, ASSISTANT, @@ -131,12 +131,6 @@ def _listening(deployment: _Deployment, count: Callable[[Request], Reply] = _cou return wire_server(_peer(deployment.provider, count), port=deployment.port) -def _reserved_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return _CLIENT_ADDRESS.validate_python(reserve.getsockname())[1] - - def _cmdline(process: psutil.Process) -> tuple[str, ...]: try: return _ARGUMENTS.validate_python(process.cmdline()) @@ -211,7 +205,7 @@ def _settled_on_every_worker(gateway: Gateway, body: Mapping[str, JsonValue]) -> def _deploy(gateway: Gateway, scenario: Scenario, provider: _Provider) -> _Deployment: - port: Final = _reserved_port() + port: Final = allocate_port() model: Final = scenario.model( model_info=None, model=f"{provider.prefix}/{_MODEL}", api_base=f"http://127.0.0.1:{port}", api_key=_API_KEY ) diff --git a/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py index dd4b7602b96..ad7cc94146c 100644 --- a/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py +++ b/tests/integration/providers/test_anthropic_refusal_stop_details_chaos.py @@ -33,6 +33,7 @@ from integration._support.client import Gateway, eventually, held from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.timing import assert_not_delayed_by from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port pytestmark: Final = [ pytest.mark.timeout(2 * graceful_stop_seconds() + 120), @@ -180,12 +181,6 @@ def _open_upstream_connections(pid: int, upstream: str) -> int: ) -def _port(wire: Wire) -> int: - port: Final = urlsplit(wire.url).port - assert port is not None, wire.url - return port - - def _accept(_request: Request) -> Reply: return Reply() @@ -200,8 +195,9 @@ async def _timed_chat(gateway: Gateway) -> float: async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: Gateway, tmp_path: Path) -> None: calls: Final = _calls(30) markers: Final = frozenset(call.marker for call in calls) + sink_port: Final = allocate_port() with wire_server(respond) as upstream, ExitStack() as first_sink: - sink: Final = first_sink.enter_context(wire_server(_accept)) + sink: Final = first_sink.enter_context(wire_server(_accept, port=sink_port)) first: Final = Ledger(sink) config: Final = _chaos_config(tmp_path, upstream, sink.url, retries=6) with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: @@ -213,7 +209,7 @@ async def test_sink_outage_mid_burst_delivers_every_refusal_event_once(gateway: first.events() first_sink.close() served_during: Final = await _burst(owned.gateway, calls[10:20]) - with wire_server(_accept, port=_port(sink)) as restarted: + with wire_server(_accept, port=sink_port) as restarted: second: Final = Ledger(restarted) served_after: Final = await _burst(owned.gateway, calls[20:]) for item in (*served_first, *served_during, *served_after): diff --git a/tests/integration/providers/test_assistants_route_errors_wire.py b/tests/integration/providers/test_assistants_route_errors_wire.py index 395a99a8afe..a59713f073b 100644 --- a/tests/integration/providers/test_assistants_route_errors_wire.py +++ b/tests/integration/providers/test_assistants_route_errors_wire.py @@ -14,6 +14,7 @@ import pytest from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.process import OwnedProxy, WorkerHealthcheck, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request +from integration._support.ports import allocate_port from integration.providers._assistants_route_errors_support import ( ASSISTANT_MODEL, BODY_ROUTES, @@ -39,7 +40,6 @@ from integration.providers._assistants_route_errors_support import ( assert_reached_upstream_once, assistants_wire, error_reply, - free_port, held_error_peer, held_one_by_one, held_upstream_connections, @@ -89,7 +89,7 @@ class _Scripted: @pytest.fixture(scope="module") def scripted(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Scripted]: directory: Final = tmp_path_factory.mktemp("assistants-route-errors-scripted") - port: Final = free_port() + port: Final = allocate_port() config: Final = openai_assistants_config(directory, port, _DEPLOYMENT_TIMEOUT_SECONDS) with ( gateway_from_environment() as shared, diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index db4e564b8e5..7d978bc0561 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -3,7 +3,6 @@ import base64 import json import os import re -import socket import threading import time import uuid @@ -27,6 +26,7 @@ from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.timing import assert_within_factor_of_control, timed from integration._support.upstream import _aws_event_frame from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with @@ -1406,17 +1406,11 @@ def _assert_rows_once(served: tuple[_Served, ...]) -> None: assert all(row["status"] == "success" for row in rows), rows -def _reserved_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - @pytest.mark.timeout(900) async def test_c01_a_peer_outage_mid_traffic_fails_cleanly_and_recovery_lands_every_id_once( gateway: Gateway, tmp_path: Path ) -> None: - port: Final = _reserved_port() + port: Final = allocate_port() prompt: Final = _prompt() with wire_server(_peer(), port=port) as first_peer: config: Final = _owned_config(first_peer, tmp_path, names=("kimi-us", "claude-control")) diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py index 51fc1aa283b..551bab45e73 100644 --- a/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_wire.py @@ -12,6 +12,7 @@ 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, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter _BACKEND: Final = "openai.gpt-5.6-luna" @@ -776,9 +777,9 @@ def test_messages_stream_overflow_logs_one_failure_row(gateway: Gateway) -> None def test_chaos_mantle_peer_outage_mid_burst_logs_every_call_once_and_keeps_the_proxy_alive(gateway: Gateway) -> None: + port: Final = allocate_port() with gateway.scenario() as scenario: - with wire_server(_chaos_peer) as wire: - port: Final = int(wire.url.rsplit(":", 1)[1]) + with wire_server(_chaos_peer, port=port) as wire: model: Final = scenario.model(model=_MODEL, api_base=wire.url, api_key=_API_KEY) before: Final = _fire(gateway, _burst_bodies(model, "before")) assert len(wire.drain()) == len(before) diff --git a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py index ece9d6be4b2..5179514b2b8 100644 --- a/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_count_tokens_wire.py @@ -1,7 +1,6 @@ import asyncio import json import re -import socket import threading import uuid from collections.abc import Callable, Iterator, Mapping, Sequence @@ -23,6 +22,7 @@ from integration._support.bedrock_runtime_peer import respond as runtime_generat from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from integration.providers._count_tokens_system_lift import ( ASSISTANT, FOLLOW_UP, @@ -181,19 +181,13 @@ def _mantle_environment(port: int) -> Mapping[str, str]: _INHERITED_BEARER: Final = ("AWS_BEARER_TOKEN_BEDROCK",) -def _reserved_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - def _closed_port_url() -> str: - return f"http://127.0.0.1:{_reserved_port()}" + return f"http://127.0.0.1:{allocate_port()}" @pytest.fixture(scope="module") def mantle_port() -> int: - return _reserved_port() + return allocate_port() @pytest.fixture(scope="module") diff --git a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py index 71b63786c27..3cadc1b5ccf 100644 --- a/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py +++ b/tests/integration/providers/test_bedrock_runtime_chat_completions_chaos.py @@ -11,7 +11,7 @@ import socket import threading import uuid from collections.abc import Callable, Iterator, Mapping -from contextlib import ExitStack, contextmanager +from contextlib import contextmanager from dataclasses import dataclass from multiprocessing.process import BaseProcess from multiprocessing.sharedctypes import Synchronized @@ -30,6 +30,7 @@ from integration._support.client import Gateway, eventually, gateway_from_enviro from integration._support.database import read_rows from integration._support.process import kill_stray_process, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter BEDROCK_MODEL: Final = "us.openai.gpt-5.6-sol" @@ -227,14 +228,6 @@ def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], ) -def _free_ports(count: int) -> tuple[int, ...]: - with ExitStack() as reserved: - sockets: Final = tuple(reserved.enter_context(socket.socket()) for _ in range(count)) - for reserve in sockets: - reserve.bind(("127.0.0.1", 0)) - return tuple(reserve.getsockname()[1] for reserve in sockets) - - def _accepts_connections(port: int) -> bool: try: with socket.create_connection(("127.0.0.1", port), timeout=0.2): @@ -283,7 +276,7 @@ def _chaos_config(endpoints: Mapping[str, str], directory: Path) -> Path: def chaos_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_ChaosProxy]: directory: Final = tmp_path_factory.mktemp("bedrock-gpt-chat-completions-chaos") burst, peer_killed, slow_peer = ( - _Deployment(f"bedrock-gpt-chat-completions-chaos-{uuid.uuid4().hex}", port) for port in _free_ports(3) + _Deployment(f"bedrock-gpt-chat-completions-chaos-{uuid.uuid4().hex}", allocate_port()) for _ in range(3) ) endpoints: Final = { deployment.model: f"http://127.0.0.1:{deployment.peer_port}" for deployment in (burst, peer_killed, slow_peer) diff --git a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py index 27e1a819e06..02f9e42ff31 100644 --- a/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py +++ b/tests/integration/providers/test_cache_control_tool_call_marks_owned_proxy.py @@ -1,6 +1,5 @@ import asyncio import re -import socket import threading from collections import Counter from collections.abc import Callable, Sequence @@ -17,6 +16,7 @@ from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from integration.providers._cache_control_marks_support import ( CITIES, POINTS, @@ -57,12 +57,6 @@ class _Sent: call_id: str -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return int(reserve.getsockname()[1]) - - def _surface_request(surface: str, marker: str) -> tuple[str, dict[str, JsonValue]]: if surface == "messages": return "/v1/messages", messages_body(_MODEL, marker, stream=False) @@ -150,7 +144,7 @@ def _single_spend_row(item: _Sent) -> None: async def test_capped_burst_rides_out_a_provider_outage_and_logs_every_request_once( gateway: Gateway, tmp_path: Path ) -> None: - port: Final = _free_port() + port: Final = allocate_port() config: Final = owned_config( tmp_path, [anthropic_deployment(_MODEL, f"http://127.0.0.1:{port}", cache_control_injection_points=POINTS)] ) diff --git a/tests/integration/providers/test_chat_completions_tool_search_chaos.py b/tests/integration/providers/test_chat_completions_tool_search_chaos.py index ac9cb214069..eda002047b0 100644 --- a/tests/integration/providers/test_chat_completions_tool_search_chaos.py +++ b/tests/integration/providers/test_chat_completions_tool_search_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import socket import threading import uuid from collections.abc import Callable, Mapping @@ -22,6 +21,7 @@ from integration._support.database import read_rows from integration._support.openai_wire import chat_reply from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from integration.providers.test_chat_completions_tool_search_wire import ( DEFERRED_NAMES, UNLISTED_CLAUDE, @@ -191,12 +191,6 @@ def _assert_every_call_forwarded(wire: Wire, calls: tuple[_Call, ...]) -> None: _assert_forwarded(request, kinds[marker]) -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def _identities(served: tuple[_Served, ...], kind: Kind) -> tuple[str, ...]: return tuple(response_id(item.answer) for item in served if item.call.kind == kind) @@ -236,7 +230,7 @@ def _spend_status(model: str, expected: int) -> dict[str, str]: async def test_a_vendor_outage_fails_each_call_cleanly_and_the_recovered_vendor_gets_the_regex_tool( gateway: Gateway, ) -> None: - port: Final = _free_port() + port: Final = allocate_port() while_down: Final = _calls(15, ("search",)) after: Final = _calls(15, ("search",)) with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py index 40b87ba1b8b..d2787d49bd0 100644 --- a/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py +++ b/tests/integration/providers/test_chatgpt_responses_caller_stream_chaos.py @@ -20,6 +20,7 @@ from integration._support.client import Gateway, eventually, gateway_from_enviro from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(240) @@ -62,17 +63,10 @@ class _Rig: return f"http://127.0.0.1:{self.port}" -def _free_port() -> int: - with wire_server(cv.CodexVendor().respond) as probe: - port: Final = urlsplit(probe.url).port - assert port is not None, probe.url - return port - - @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: directory: Final = tmp_path_factory.mktemp("chatgpt-caller-stream-chaos") - port: Final = _free_port() + port: Final = allocate_port() overrides: Final = {"CHATGPT_TOKEN_DIR": str(cv.login(directory)), "CHATGPT_API_BASE": f"http://127.0.0.1:{port}"} config: Final = cv.proxy_config(directory, probe=False) with gateway_from_environment() as gateway: diff --git a/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py b/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py index a1cb2a84d71..75a1d55f67a 100644 --- a/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py +++ b/tests/integration/providers/test_cloudflare_clef_omni_media_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import socket import threading import uuid from dataclasses import dataclass @@ -19,6 +18,7 @@ from integration._support.client import Gateway, eventually, object_value, strin from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter _API_KEY: Final = "synthetic-cloudflare-key" @@ -163,12 +163,6 @@ def _assert_each_billed_once( assert by_call[item.call_id]["status"] == "success", (item.call_id, by_call[item.call_id]) -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def _assert_failed_cleanly(served: _Served) -> None: assert served.status >= 500, (served.status, served.text) assert "answers" not in served.text, served.text @@ -180,7 +174,7 @@ def _assert_failed_cleanly(served: _Served) -> None: async def test_an_upstream_dropping_connections_mid_burst_then_stopped_and_restarted_lands_each_success_once( gateway: Gateway, ) -> None: - port: Final = _free_port() + port: Final = allocate_port() first_wave: Final = _markers(12) dropped: Final = frozenset(marker for index, marker in enumerate(first_wave) if index % 2 == 1) while_down: Final = _markers(6) diff --git a/tests/integration/providers/test_coralbricks_provider_wire.py b/tests/integration/providers/test_coralbricks_provider_wire.py index d8231df5151..ac2bab379f1 100644 --- a/tests/integration/providers/test_coralbricks_provider_wire.py +++ b/tests/integration/providers/test_coralbricks_provider_wire.py @@ -8,6 +8,7 @@ import pytest from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value from integration._support.responses_vendor import same_response from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from integration.providers._coralbricks import ( ANSWER, API_KEY, @@ -24,7 +25,6 @@ from integration.providers._coralbricks import ( catalog_rows, async_anthropic_client, async_openai_client, - closed_port, deployment, expected_spend, failing_provider, @@ -365,7 +365,7 @@ def test_unreachable_api_base_reports_the_connection_error_and_leaves_the_proxy_ with gateway.scenario() as scenario: control: Final = scenario.model() alias: Final = scenario.model( - model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{closed_port()}/v1", api_key=API_KEY + model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{allocate_port()}/v1", api_key=API_KEY ) response: Final = eventually( lambda: gateway.request( @@ -593,7 +593,7 @@ def test_fallback_from_a_refused_primary_lands_on_coralbricks(gateway: Gateway) with wire_server(provider(identity)) as wire, gateway.scenario() as scenario: alias: Final = deployment(scenario, wire) primary: Final = scenario.model( - model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{closed_port()}/v1", api_key=API_KEY + model=f"{PROVIDER}/{MODEL}", api_base=f"http://127.0.0.1:{allocate_port()}/v1", api_key=API_KEY ) wire.drain() response: Final = eventually( diff --git a/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py b/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py index 22c8b207998..d01cbde915e 100644 --- a/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py +++ b/tests/integration/providers/test_decisions_box_deployment_flag_chaos.py @@ -18,6 +18,7 @@ from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter _API_KEY: Final = "synthetic-decisions-key" @@ -180,10 +181,9 @@ async def test_a_flagged_deployment_recovers_box_answers_after_an_upstream_outag gateway: Gateway, ) -> None: base_url: Final = str(gateway.client.base_url) + port: Final = allocate_port() with gateway.scenario() as scenario: - with wire_server(_reply) as wire: - port: Final = urlsplit(wire.url).port - assert port is not None + with wire_server(_reply, port=port) as wire: model: Final = scenario.model( model=_backend(), api_base=wire.url, api_key=_API_KEY, model_info={"supports_box_questions": True} ) diff --git a/tests/integration/providers/test_decisions_chaos.py b/tests/integration/providers/test_decisions_chaos.py index d868e46e6a5..7bbbb09f638 100644 --- a/tests/integration/providers/test_decisions_chaos.py +++ b/tests/integration/providers/test_decisions_chaos.py @@ -18,6 +18,7 @@ from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter _BACKEND: Final = "pplx-decider-v1-27b" @@ -239,10 +240,9 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_serving_the_default_m @pytest.mark.timeout(180) async def test_upstream_outage_fails_its_calls_and_recovery_on_the_same_port_restores_them(gateway: Gateway) -> None: base_url: Final = str(gateway.client.base_url) + port: Final = allocate_port() with gateway.scenario() as scenario: - with wire_server(_reply) as wire: - port: Final = urlsplit(wire.url).port - assert port is not None + with wire_server(_reply, port=port) as wire: model: Final = scenario.model(model=f"perplexity/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) before: Final = await _burst(base_url, gateway.key, model, _calls(5, fail=False)) assert _health(gateway, model) == (1, 0) diff --git a/tests/integration/providers/test_decisions_wire.py b/tests/integration/providers/test_decisions_wire.py index 882bba1df51..87cd3105cc7 100644 --- a/tests/integration/providers/test_decisions_wire.py +++ b/tests/integration/providers/test_decisions_wire.py @@ -1,6 +1,5 @@ import json import math -import socket import uuid from collections.abc import Sequence from dataclasses import dataclass @@ -14,6 +13,7 @@ from integration._support.client import Gateway, Scenario, eventually, object_va from integration._support.database import read_rows from integration._support.process import owned_proxy_process from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario +from integration._support.ports import allocate_port from integration.cost_calculation.cost_tracking_case import JsonResponse from pydantic import JsonValue, TypeAdapter @@ -235,12 +235,6 @@ def _assert_connection_error(response: httpx.Response) -> None: assert _GENERIC_API_ERROR not in response.text, response.text -def _free_closed_port() -> int: - with socket.socket() as probe: - probe.bind(("127.0.0.1", 0)) - return probe.getsockname()[1] - - def _pass_through_config(directory: Path, pass_through_target: str, native_api_base: str) -> Path: base: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) config: Final = { @@ -431,8 +425,8 @@ def test_request_body_api_base_is_refused_like_chat_without_an_upstream_call(gat with gateway.scenario() as scenario: handle: Final = _register(scenario, _answer_body(_PERPLEXITY)) model: Final = _deployment(scenario, handle, _PERPLEXITY) - decisions: Final = _decide(gateway, model, api_base=f"http://127.0.0.1:{_free_closed_port()}") - chat: Final = _chat(gateway, model, api_base=f"http://127.0.0.1:{_free_closed_port()}") + decisions: Final = _decide(gateway, model, api_base=f"http://127.0.0.1:{allocate_port()}") + chat: Final = _chat(gateway, model, api_base=f"http://127.0.0.1:{allocate_port()}") assert 400 <= decisions.status_code < 500, decisions.text assert decisions.status_code == chat.status_code, (decisions.text, chat.text) assert _upstream_calls(gateway, handle) == [] @@ -542,7 +536,7 @@ def test_unreachable_upstream_fails_only_its_own_deployment(gateway: Gateway, pr handle: Final = _register(scenario, _answer_body(provider)) healthy: Final = _deployment(scenario, handle, provider) dead: Final = scenario.model( - model=provider.model, api_base=f"http://127.0.0.1:{_free_closed_port()}", api_key=provider.api_key + model=provider.model, api_base=f"http://127.0.0.1:{allocate_port()}", api_key=provider.api_key ) failed: Final = _decide(gateway, dead, num_retries=0) _assert_connection_error(failed) @@ -556,7 +550,7 @@ def test_an_unreachable_openrouter_deployment_reports_a_connection_error_on_chat gateway: Gateway, ) -> None: with gateway.scenario() as scenario: - dead_api_base: Final = f"http://127.0.0.1:{_free_closed_port()}" + dead_api_base: Final = f"http://127.0.0.1:{allocate_port()}" chat_model: Final = scenario.model(model=_OPENROUTER_CHAT_MODEL, api_base=dead_api_base, api_key=_API_KEY) decisions_model: Final = scenario.model(model=_OPENROUTER.model, api_base=dead_api_base, api_key=_API_KEY) chat: Final = _chat(gateway, chat_model, num_retries=0) @@ -572,7 +566,7 @@ def test_an_unreachable_openrouter_deployment_reports_a_connection_error_on_chat async def test_sdk_openrouter_connection_failures_raise_a_connection_error(gateway: Gateway) -> None: - dead_api_base: Final = f"http://127.0.0.1:{_free_closed_port()}" + dead_api_base: Final = f"http://127.0.0.1:{allocate_port()}" with pytest.raises(litellm.APIConnectionError): litellm.completion( model=_OPENROUTER_CHAT_MODEL, diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py index 1241c2a7f49..90ff32e01dd 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py @@ -11,7 +11,6 @@ import asyncio import base64 import dataclasses import json -import socket import threading import time from collections.abc import Callable, Mapping @@ -24,6 +23,7 @@ import openai import pytest from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from integration.providers._gemini_tts_voice_support import ( AUDIO_TOKENS, BACKEND, @@ -514,12 +514,6 @@ def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gatew assert voice_in(setups[0]) == NO_VOICE, setups[0] -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return int(reserve.getsockname()[1]) - - @dataclass(frozen=True, slots=True) class Fired: provider: str @@ -586,7 +580,7 @@ def _assert_served(items: tuple[Fired, ...], received: tuple[Request, ...]) -> N def test_peer_outage_mid_burst_ends_each_request_once_and_recovers_on_the_same_port(gateway: Gateway) -> None: - port: Final = _free_port() + port: Final = allocate_port() with gateway.scenario() as scenario: url: Final = f"http://127.0.0.1:{port}" models: Final = { diff --git a/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py b/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py index 727b7ff7f5b..a86c82bcb21 100644 --- a/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py +++ b/tests/integration/providers/test_idle_upstream_reconnect_chaos_wire.py @@ -8,6 +8,7 @@ from typing import Final import pytest from integration._support.client import Gateway from integration._support.idle_gateway import idle_gateway +from integration._support.ports import allocate_port from integration._support.process import graceful_stop_seconds from integration.providers._idle_upstream_support import ( DROP_AFTER_SECONDS, @@ -108,7 +109,9 @@ def test_upstream_restarted_on_the_same_port_with_connections_pooled_serves_a_bu before_restart: Final = stack.enter_context(ExitStack()) first: Final = MappingProxyType( { - row.name: before_restart.enter_context(idle_gateway(answer, drop_after=DROP_AFTER_SECONDS)) + row.name: before_restart.enter_context( + idle_gateway(answer, drop_after=DROP_AFTER_SECONDS, port=allocate_port()) + ) for row in RESTART_ROWS } ) diff --git a/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py b/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py index 74ea3272503..3de825ae4a6 100644 --- a/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py +++ b/tests/integration/providers/test_responses_bridge_message_cache_control_wire_chaos.py @@ -8,7 +8,6 @@ from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final, TypeAlias -from urllib.parse import urlsplit import httpx import pytest @@ -25,6 +24,7 @@ from integration._support.client import ( from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.openai_wire import answering_model_discovery, chat_reply from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from integration.providers._cache_control_marks_support import ( ANTHROPIC_MODEL, ASK, @@ -236,13 +236,6 @@ def _health(gateway: Gateway, model: str) -> Mapping[str, JsonValue]: return _JSON_OBJECT.validate_json(response.text) -def _free_port() -> int: - with wire_server(anthropic_peer) as probe: - port: Final = urlsplit(probe.url).port - assert port is not None, probe.url - return port - - def _deployment(scenario: Scenario, api_base: str) -> str: return scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=api_base, api_key=PROVIDER_KEY) @@ -262,7 +255,7 @@ async def test_mixed_burst_carries_the_system_mark_once_per_call(gateway: Gatewa async def test_peer_outage_fails_cleanly_and_the_restarted_peer_receives_the_marks( gateway: Gateway, spend: pcb.SpendLogs ) -> None: - port: Final = _free_port() + port: Final = allocate_port() while_down: Final = _calls(12) after: Final = _calls(12) with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py index e03af85d01a..5acf3b10e04 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_chaos.py @@ -8,7 +8,6 @@ from contextlib import ExitStack from pathlib import Path from queue import SimpleQueue from typing import Final -from urllib.parse import urlsplit import httpx import psutil @@ -17,6 +16,7 @@ import yaml from integration._support.client import Gateway, eventually from integration._support.process import owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from integration.providers._responses_bridge_prompt_cache_breakpoint import ( _Call, _JSON_OBJECT, @@ -86,8 +86,9 @@ async def test_worker_and_peer_outages_preserve_markers_and_recover( assert gate.wait(timeout=60), "The worker-kill burst was never released" return _responses_reply(request) + upstream_port: Final = allocate_port() with ExitStack() as peer_stack: - wire: Final = peer_stack.enter_context(wire_server(held)) + wire: Final = peer_stack.enter_context(wire_server(held, port=upstream_port)) config: Final = _chaos_config(wire, tmp_path) with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: try: @@ -113,9 +114,6 @@ async def test_worker_and_peer_outages_preserve_markers_and_recover( early_successful: Final = tuple(item for item in early_served if item.status == 200) for item in early_successful: _assert_spend_for_result(item, _CONFIG_MODEL) - upstream_port_value: Final = urlsplit(wire.url).port - assert upstream_port_value is not None - upstream_port: Final = upstream_port_value active_by_worker: Final = eventually( lambda: {pid: _open_upstream_connections(pid, upstream_port) for pid in workers}, lambda counts: sum(counts.values()) == len(calls) - len(early_calls), diff --git a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py index ee73d780170..1bc7853437b 100644 --- a/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py +++ b/tests/integration/providers/test_responses_bridge_prompt_cache_breakpoint_wire_chaos.py @@ -9,7 +9,6 @@ from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final, Literal, TypeAlias -from urllib.parse import urlsplit import httpx import pytest @@ -19,6 +18,7 @@ from integration._support import responses_vendor as rv from integration._support.client import Gateway, eventually, gateway_from_environment, string_value from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -213,13 +213,6 @@ def _health(gateway: Gateway, model: str) -> Mapping[str, JsonValue]: return rv.JSON_OBJECT.validate_json(response.text) -def _free_port() -> int: - with wire_server(pcb.respond) as probe: - port: Final = urlsplit(probe.url).port - assert port is not None, probe.url - return port - - def _chat(gateway: Gateway, model: str, marker: str, breakpoint: JsonValue) -> httpx.Response: return gateway.request("POST", "/v1/chat/completions", dict(_body(model, _Call("chat", False, marker), breakpoint))) @@ -285,7 +278,7 @@ async def test_mixed_burst_carries_every_marker_once(gateway: Gateway, spend: pc async def test_upstream_outage_fails_cleanly_and_the_restarted_upstream_serves_marked_calls( gateway: Gateway, spend: pcb.SpendLogs ) -> None: - port: Final = _free_port() + port: Final = allocate_port() while_down: Final = _calls(12, _ENDPOINTS) after: Final = _calls(12, _ENDPOINTS) with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py index bec1f5799eb..0101824054a 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py @@ -2,7 +2,6 @@ import asyncio import itertools import json import re -import socket import threading import uuid from collections.abc import Callable, Iterator, Mapping, Sequence @@ -28,6 +27,7 @@ from integration._support.openai_wire import answering_model_discovery from integration._support.database import read_rows from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue _CLAUDE: Final = cc.OPUS @@ -339,12 +339,6 @@ def _assert_each_lands_once( assert by_status[match] == "success", rows -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def _health_counts(gateway: Gateway, model: str) -> tuple[int, int]: response: Final = gateway.request("GET", f"/health?model={model}", None) assert response.status_code in (200, 503), response.text @@ -382,7 +376,7 @@ async def test_mixed_burst_carries_the_search_call_on_every_endpoint_and_lands_e async def test_vendor_outage_fails_each_call_cleanly_and_the_recovered_vendor_gets_the_tool_search( gateway: Gateway, ) -> None: - port: Final = _free_port() + port: Final = allocate_port() while_down: Final = _calls(15, ("responses", "chat", "messages")) after: Final = _calls(15, ("responses", "chat", "messages")) with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py b/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py index f003e5ffa63..16d37980efa 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_chaos.py @@ -1,7 +1,6 @@ import asyncio import json import re -import socket import threading import time import uuid @@ -25,6 +24,7 @@ from integration._support.database import read_rows from integration._support.process import OwnedProxy, owned_proxy_process from integration._support.tls import server_context, write_self_signed_cert from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue _GPT: Final = "gpt-5.6" @@ -226,12 +226,6 @@ def _assert_each_lands_once( assert by_status[match] == "success", rows -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def _health_counts(gateway: Gateway, model: str) -> tuple[int, int]: response: Final = gateway.request("GET", f"/health?model={model}", None) assert response.status_code in (200, 503), response.text @@ -249,7 +243,7 @@ def _marked(received: tuple[Request, ...]) -> dict[str, Request]: async def test_vendor_outage_fails_each_replay_cleanly_and_the_recovered_vendor_gets_them_without_minted_items( gateway: Gateway, ) -> None: - port: Final = _free_port() + port: Final = allocate_port() while_down: Final = _calls(15, ("responses", "chat", "messages")) after: Final = _calls(15, ("responses", "chat", "messages")) with gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_sagemaker_embedding_chaos.py b/tests/integration/providers/test_sagemaker_embedding_chaos.py index 680ad3bfd2d..8e984c81ff9 100644 --- a/tests/integration/providers/test_sagemaker_embedding_chaos.py +++ b/tests/integration/providers/test_sagemaker_embedding_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import socket import threading import uuid from collections.abc import Iterator @@ -29,6 +28,7 @@ from integration._support.sagemaker_embedding import ( respond, ) from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue pytestmark = pytest.mark.timeout(300) @@ -54,12 +54,6 @@ class _Served: call_id: str -def _free_port() -> int: - with socket.socket() as probe: - probe.bind(("127.0.0.1", 0)) - return probe.getsockname()[1] - - def _config(directory: Path) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["model_list"] = [openai_deployment(_MODEL, "openai", model_info={"mode": "embedding"})] @@ -80,7 +74,7 @@ def rig() -> Iterator[Gateway]: @pytest.fixture(scope="module") def peer_port() -> int: - return _free_port() + return allocate_port() @pytest.fixture(scope="module") diff --git a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py index 2d2098987a7..1809b8ce5f8 100644 --- a/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py +++ b/tests/integration/providers/test_vertex_anthropic_structured_output_wire.py @@ -28,6 +28,7 @@ from integration._support.process import OwnedProxy, graceful_stop_seconds, owne from integration._support.upstream import aws_event_stream_frame from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter _JSON_TOOL: Final = "json_tool_call" @@ -1734,9 +1735,9 @@ def test_vertex_chat_peer_error_reaches_the_caller_and_the_deployment_keeps_serv ) -> None: marker, marker_hex = _marker("wrapped") after, after_hex = _marker("wrapped") + port: Final = allocate_port() with gateway.scenario() as scenario: - with wire_server(_failing(status)) as failing: - port: Final = int(failing.url.rsplit(":", 1)[1]) + with wire_server(_failing(status), port=port) as failing: model: Final = _vertex_deployment(gateway, scenario, failing.url, _TOOL_MODEL) response: Final = gateway.request("POST", "/v1/chat/completions", _chat_body(model, marker)) assert response.status_code == status, response.text @@ -1830,8 +1831,8 @@ def test_upstream_outage_mid_run_recovers_with_envelopes_still_unwrapped(gateway assert observed == _payload(marker_hex), response.text _seen_once(_delivered(wire), tuple(marker_hex for _, marker_hex, _ in results)) - with wire_server(_peer) as wire: - port: Final = int(wire.url.rsplit(":", 1)[1]) + port: Final = allocate_port() + with wire_server(_peer, port=port) as wire: model: Final = _vertex_deployment(gateway, scenario, wire.url, _TOOL_MODEL) unwrapped(wire, wave()) outage: Final = wave() diff --git a/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py b/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py index 12e58ec6e41..b969a1dec8f 100644 --- a/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py +++ b/tests/integration/providers/test_vertex_batch_row_line_errors_chaos.py @@ -8,12 +8,12 @@ from typing import Final import httpx from integration._support.client import Gateway, eventually from integration._support.wire import Request, wire_server +from integration._support.ports import allocate_port from integration.providers._vertex_batch_rows import ( ROUTES, Route, chat_row, - closed_port, encoded, error_of, gcs_peer, @@ -124,7 +124,7 @@ def test_concurrent_mixed_burst_names_each_bad_files_own_line_and_lands_each_goo def test_bucket_outage_answers_500_without_blaming_a_row_and_recovers_on_restart(gateway: Gateway) -> None: - port: Final = closed_port() + port: Final = allocate_port() with gateway.scenario() as scenario, _burst_client(gateway) as client: model: Final = vertex_model(gateway, scenario, f"http://127.0.0.1:{port}") with wire_server(gcs_peer, port=port) as before: diff --git a/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py b/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py index f5df41c16b0..77c3c21ffa2 100644 --- a/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py +++ b/tests/integration/providers/test_vertex_batch_row_line_errors_owned_proxy.py @@ -17,12 +17,12 @@ import yaml from integration._support.client import Gateway, eventually, gateway_from_environment from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Request, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter from integration.providers._vertex_batch_rows import ( Route, chat_row, - closed_port, encoded, error_of, gcs_peer, @@ -97,8 +97,8 @@ def _config(gateway: Gateway, directory: Path, vertex_port: int, openai_port: in @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: directory: Final = tmp_path_factory.mktemp("vertex-batch-row-lines") - vertex_port: Final = closed_port() - openai_port: Final = closed_port() + vertex_port: Final = allocate_port() + openai_port: Final = allocate_port() with gateway_from_environment() as shared: config: Final = _config(shared, directory, vertex_port, openai_port) with owned_proxy_process(shared, directory, {}, config=config, workers=2) as owned: diff --git a/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py b/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py index e47dd42413d..f03719d9520 100644 --- a/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py +++ b/tests/integration/providers/test_vertex_batch_row_line_errors_wire.py @@ -7,6 +7,7 @@ import openai import pytest from integration._support.client import Gateway, Scenario from integration._support.wire import Reply, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue from integration.providers._vertex_batch_rows import ( @@ -16,7 +17,6 @@ from integration.providers._vertex_batch_rows import ( ROUTES, Route, chat_row, - closed_port, encoded, error_of, gcs_peer, @@ -264,7 +264,7 @@ def test_bucket_failures_keep_their_status_and_never_blame_a_row( gateway: Gateway, route: Route, reply: Reply | None, status: int ) -> None: with wire_server(lambda _request: reply or Reply()) as wire, gateway.scenario() as scenario: - url: Final = wire.url if reply is not None else f"http://127.0.0.1:{closed_port()}" + url: Final = wire.url if reply is not None else f"http://127.0.0.1:{allocate_port()}" model: Final = vertex_model(gateway, scenario, url) response: Final = upload(gateway, route, model, jsonl(encoded(chat_row("first", model, "one")))) assert response.status_code == status, response.text @@ -330,7 +330,7 @@ def _unreachable_anthropic_model(scenario: Scenario) -> str: return scenario.model( model="anthropic/claude-haiku-4-5", api_key="synthetic-anthropic-key", - api_base=f"http://127.0.0.1:{closed_port()}", + api_base=f"http://127.0.0.1:{allocate_port()}", num_retries=0, ) diff --git a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py index e181e38b1bf..4cf9a3159fe 100644 --- a/tests/integration/providers/test_vertex_partner_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_partner_count_tokens_wire.py @@ -1,6 +1,5 @@ import json import re -import socket import threading import uuid from collections.abc import Callable, Mapping, Sequence @@ -18,6 +17,7 @@ from integration._support.client import Gateway, Scenario, eventually from integration._support.process import owned_proxy_process from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue, TypeAdapter _BACKEND: Final = "claude-sonnet-4-6" @@ -224,9 +224,7 @@ def _local_count(gateway: Gateway, body: Mapping[str, JsonValue]) -> int: def _closed_port_url() -> str: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return f"http://127.0.0.1:{reserve.getsockname()[1]}" + return f"http://127.0.0.1:{allocate_port()}" def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: diff --git a/tests/integration/providers/test_voyage_mongodb_host_wire.py b/tests/integration/providers/test_voyage_mongodb_host_wire.py index 273020d48a8..62860c52568 100644 --- a/tests/integration/providers/test_voyage_mongodb_host_wire.py +++ b/tests/integration/providers/test_voyage_mongodb_host_wire.py @@ -2,7 +2,6 @@ import asyncio import json import os import re -import socket import uuid from collections.abc import Callable, Iterator from dataclasses import dataclass @@ -17,6 +16,7 @@ from integration._support.database import read_rows from integration._support.forward_proxy import ForwardProxy, refusing_forward_proxy from integration._support.process import owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from openai import APIStatusError, AsyncOpenAI, OpenAI from pydantic import JsonValue, TypeAdapter @@ -617,15 +617,8 @@ def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: return read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (call_id,)) -def _reserve_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - address: Final[Callable[[], tuple[str, int]]] = reserve.getsockname - return address()[1] - - async def test_forward_proxy_outage_mid_traffic_recovers_with_the_right_hosts(gateway: Gateway, tmp_path: Path) -> None: - port: Final = _reserve_port() + port: Final = allocate_port() config: Final = _write_config(tmp_path, "outage", _isolated_deployments(gateway.upstream_url)) with owned_proxy( gateway, diff --git a/tests/integration/routing/test_declared_provider_bare_model_routes.py b/tests/integration/routing/test_declared_provider_bare_model_routes.py index 054d4d88be4..0ee5f43c328 100644 --- a/tests/integration/routing/test_declared_provider_bare_model_routes.py +++ b/tests/integration/routing/test_declared_provider_bare_model_routes.py @@ -9,7 +9,6 @@ from __future__ import annotations import asyncio import json -import socket import time import uuid from collections.abc import Awaitable, Callable @@ -29,6 +28,7 @@ from websockets.exceptions import ConnectionClosed from tests.integration._support.client import JSON_OBJECT, Gateway, Scenario, object_value, string_value from tests.integration._support.upstream import delete_scenario, register_scenario from tests.integration._support.wire import Reply, Request, Wire, wire_server +from tests.integration._support.ports import allocate_port from tests.integration.cost_calculation.cost_tracking_case import RealtimeResponse _KEY: Final = "synthetic-declared-provider-key" @@ -497,14 +497,8 @@ def _marker(request: Request) -> str: return string_value(body["prompt"] if request.target == _IMAGES else body["input"]) -def _closed_port() -> int: - with socket.socket() as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - def test_burst_across_declared_and_dead_deployments_reaches_each_live_provider_once(gateway: Gateway) -> None: - dead_base: Final = f"http://127.0.0.1:{_closed_port()}/v1" + dead_base: Final = f"http://127.0.0.1:{allocate_port()}/v1" with wire_server(_burst_reply) as wire, gateway.scenario() as scenario: live: Final = tuple( (route, _deployment(scenario, wire, form, _bare())) diff --git a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py index 490c26226bd..f30c13993d4 100644 --- a/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py +++ b/tests/integration/routing/test_fallback_hop_foreign_encrypted_reasoning.py @@ -46,6 +46,7 @@ from integration._support.database import read_rows, scratch_database from integration._support.process import graceful_stop_seconds, owned_proxy_process from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, SseResponse, StoredResponse from pydantic import JsonValue @@ -1214,10 +1215,10 @@ def test_order_one_outage_mid_burst_hops_stripped_and_recovers_after_restart(gat order_two: Final = deployment( scenario, name, scenario_api_base(scenario, two_scenario, healthy_json()), order=2 ) - with wire_server(peer_reply) as peer: + port: Final = allocate_port() + with wire_server(peer_reply, port=port) as peer: order_one: Final = deployment(scenario, name, peer.url, order=1) target: Final = Target(name, frozenset({order_one, order_two})) - port: Final = int(peer.url.rsplit(":", 1)[-1]) warm: Final = marker() assert served_by_order_one(responses_answer(gateway, HTTPX, target, history(warm)), order_one), warm assert [json.loads(request.body)["input"] for request in peer.drain()] == [history(warm)] diff --git a/tests/integration/sdk/test_http2_wire.py b/tests/integration/sdk/test_http2_wire.py index 15bb366c7a2..4fcf590f0eb 100644 --- a/tests/integration/sdk/test_http2_wire.py +++ b/tests/integration/sdk/test_http2_wire.py @@ -10,7 +10,7 @@ import time from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path -from typing import Final, cast +from typing import Final import pytest from cryptography import x509 @@ -22,6 +22,7 @@ from hypercorn.asyncio import ( ) from hypercorn.config import Config from hypercorn.typing import ASGIReceiveCallable, ASGISendCallable, HTTPResponseBodyEvent, HTTPResponseStartEvent, Scope +from integration._support.ports import allocate_port STREAM_CHUNKS: Final = 3 @@ -102,9 +103,7 @@ async def _peer(scope: Scope, receive: ASGIReceiveCallable, send: ASGISendCallab @pytest.fixture(scope="module") def http2_tls_peer(tmp_path_factory: pytest.TempPathFactory) -> Iterator[str]: cert_file, key_file = _write_self_signed_cert(tmp_path_factory.mktemp("h2certs")) - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - port: Final = cast(int, sock.getsockname()[1]) + port: Final = allocate_port() shutdown: Final = threading.Event() def _serve() -> None: diff --git a/tests/integration/sdk/test_owned_redis_ports.py b/tests/integration/sdk/test_owned_redis_ports.py index 0212329b7b7..7e7dfa94852 100644 --- a/tests/integration/sdk/test_owned_redis_ports.py +++ b/tests/integration/sdk/test_owned_redis_ports.py @@ -1,12 +1,10 @@ from __future__ import annotations import re -from itertools import chain from pathlib import Path -from typing import Final import pytest -from integration._support.redis_process import PORT_ATTEMPTS, free_loopback_port, owned_redis +from integration._support.redis_process import owned_redis from redis import Redis @@ -15,20 +13,6 @@ def _redis_pid(host: str, port: int) -> int: return int(client.info("server")["process_id"]) -def test_owned_redis_moves_off_a_port_another_owned_redis_already_serves(tmp_path: Path) -> None: - with owned_redis(tmp_path / "first") as first: - candidates: Final = chain((first.port,), iter(free_loopback_port, None)) - with owned_redis(tmp_path / "second", lambda: next(candidates)) as second: - assert second.port != first.port - assert first.serves() - assert second.serves() - assert second.server_pid != first.server_pid - assert _redis_pid(second.host, second.port) == second.server_pid - assert _redis_pid(first.host, first.port) == first.server_pid - assert not second.serves() - assert first.serves() - - def test_owned_redis_restart_refuses_a_port_another_redis_took_while_it_was_down(tmp_path: Path) -> None: with owned_redis(tmp_path / "first") as first: first.stop() @@ -39,18 +23,3 @@ def test_owned_redis_restart_refuses_a_port_another_redis_took_while_it_was_down first.start() assert first.process is None assert _redis_pid(squatter.host, squatter.port) == squatter.server_pid - - -def test_owned_redis_gives_up_with_the_last_cause_after_a_bounded_number_of_ports(tmp_path: Path) -> None: - with owned_redis(tmp_path / "first") as first: - ports: Final = iter((first.port,) * PORT_ATTEMPTS) - with pytest.raises( - AssertionError, - match=( - f"after {PORT_ATTEMPTS} attempts: Another Redis already serves {re.escape(first.host)}:{first.port}" - ), - ): - with owned_redis(tmp_path / "second", lambda: next(ports)): - pytest.fail("Owned Redis unexpectedly reused a port already serving Redis") - assert next(ports, None) is None - assert first.serves() diff --git a/tests/integration/streaming/test_faked_responses_stream_chaos.py b/tests/integration/streaming/test_faked_responses_stream_chaos.py index f96538dfab6..28efe912157 100644 --- a/tests/integration/streaming/test_faked_responses_stream_chaos.py +++ b/tests/integration/streaming/test_faked_responses_stream_chaos.py @@ -1,6 +1,5 @@ import asyncio import json -import socket import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -27,6 +26,7 @@ from integration._support.client import ( from integration._support.database import read_rows from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from integration.streaming.faked_responses_case import ( ANTHROPIC_TOOLS, API_KEY, @@ -39,7 +39,7 @@ from integration.streaming.faked_responses_case import ( message_item, response_body, ) -from pydantic import JsonValue, TypeAdapter +from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -68,9 +68,6 @@ class _Served: call_id: str -_ADDRESS: Final = TypeAdapter(tuple[str, int]) - - @dataclass(frozen=True, slots=True) class _Rig: port: int @@ -81,12 +78,6 @@ class _Rig: return self.proxy.gateway -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return _ADDRESS.validate_python(reserve.getsockname())[1] - - def _deployment(name: str, model: str, port: int) -> dict[str, JsonValue]: return { "model_name": name, @@ -117,7 +108,7 @@ def _config(port: int, directory: Path) -> Path: @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: directory: Final = tmp_path_factory.mktemp("faked-stream-chaos") - port: Final = _free_port() + port: Final = allocate_port() with ( gateway_from_environment() as shared, owned_proxy_process(shared, directory, {}, config=_config(port, directory), workers=2) as owned, diff --git a/tests/integration/streaming/test_responses_bridge_stream_chaos.py b/tests/integration/streaming/test_responses_bridge_stream_chaos.py index ede5a824f52..754ee5d3981 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_chaos.py +++ b/tests/integration/streaming/test_responses_bridge_stream_chaos.py @@ -1,6 +1,5 @@ import asyncio import re -import socket import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -35,6 +34,7 @@ from integration._support.responses_stream import ( rate_limited_stream, ) from integration._support.wire import Reply, Request, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) @@ -94,12 +94,6 @@ class _Rig: return self.proxy.gateway -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - def _config(port: int, directory: Path) -> Path: stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) router_settings: Final = object_value(stock.get("router_settings") or {}) @@ -128,7 +122,7 @@ def _config(port: int, directory: Path) -> Path: @pytest.fixture(scope="module") def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: directory: Final = tmp_path_factory.mktemp("bridged-stream-chaos") - port: Final = _free_port() + port: Final = allocate_port() with ( gateway_from_environment() as shared, owned_proxy_process(shared, directory, {}, config=_config(port, directory), workers=2) as owned, @@ -382,7 +376,7 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_the_bridged async def test_proxy_sigterm_mid_burst_drains_the_spend_log_queue_and_the_restarted_proxy_serves( gateway: Gateway, tmp_path: Path ) -> None: - port: Final = _free_port() + port: Final = allocate_port() config: Final = _config(port, tmp_path) calls: Final = _calls(20, _CHAT_KINDS) held: Final = _Held(threading.Event(), SimpleQueue()) diff --git a/tests/integration/streaming/test_responses_bridge_stream_errors.py b/tests/integration/streaming/test_responses_bridge_stream_errors.py index d8d4a44cd68..3fba69b39b7 100644 --- a/tests/integration/streaming/test_responses_bridge_stream_errors.py +++ b/tests/integration/streaming/test_responses_bridge_stream_errors.py @@ -1,8 +1,6 @@ import json -import socket import uuid from collections.abc import Iterator, Mapping, Sequence -from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path from typing import Final @@ -47,6 +45,7 @@ from integration._support.responses_stream import ( sse_events, ) from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from pydantic import JsonValue _PRIMARY: Final = "bridged-primary" @@ -110,14 +109,6 @@ class FallbackProxy: spare_port: int -def _free_ports(count: int) -> tuple[int, ...]: - with ExitStack() as reserved: - sockets: Final = tuple(reserved.enter_context(socket.socket()) for _ in range(count)) - for reserve in sockets: - reserve.bind(("127.0.0.1", 0)) - return tuple(reserve.getsockname()[1] for reserve in sockets) - - def _fallback_config(directory: Path, primary_port: int, spare_port: int) -> Path: base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) deployments: Final = [ @@ -140,7 +131,7 @@ def _fallback_config(directory: Path, primary_port: int, spare_port: int) -> Pat @pytest.fixture(scope="module") def fallback_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[FallbackProxy]: directory: Final = tmp_path_factory.mktemp("bridged-fallbacks") - primary_port, spare_port = _free_ports(2) + primary_port, spare_port = allocate_port(), allocate_port() with ( gateway_from_environment() as shared, owned_proxy(shared, directory, {}, config=_fallback_config(directory, primary_port, spare_port)) as owned, diff --git a/tests/integration/streaming/test_stream_open_error_provider_headers_proxy.py b/tests/integration/streaming/test_stream_open_error_provider_headers_proxy.py index f614de3aeef..bd3da0a6728 100644 --- a/tests/integration/streaming/test_stream_open_error_provider_headers_proxy.py +++ b/tests/integration/streaming/test_stream_open_error_provider_headers_proxy.py @@ -1,7 +1,6 @@ import asyncio import json import re -import socket import threading import uuid from collections.abc import Iterator, Mapping, Sequence @@ -29,6 +28,7 @@ from integration._support.client import ( from integration._support.database import read_rows from integration._support.process import OwnedProxy, configless_proxy_process, graceful_stop_seconds from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.ports import allocate_port from openai.types.responses import ResponseFailedEvent from pydantic import JsonValue @@ -76,12 +76,6 @@ class _Served: call_id: str -def _free_port() -> int: - with socket.socket() as reserve: - reserve.bind(("127.0.0.1", 0)) - return reserve.getsockname()[1] - - @contextmanager def _configless(gateway: Gateway, directory: Path, port: int) -> Iterator[_Rig]: url: Final = f"http://127.0.0.1:{port}" @@ -99,7 +93,7 @@ def _configless(gateway: Gateway, directory: Path, port: int) -> Iterator[_Rig]: def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: with ( gateway_from_environment() as shared, - _configless(shared, tmp_path_factory.mktemp("configless"), _free_port()) as owned, + _configless(shared, tmp_path_factory.mktemp("configless"), allocate_port()) as owned, ): yield owned @@ -543,7 +537,7 @@ async def test_worker_sigkill_mid_burst_leaves_the_sibling_answering_each_stream async def test_proxy_sigterm_mid_burst_answers_each_in_flight_stream_and_the_restarted_proxy_serves( gateway: Gateway, tmp_path: Path ) -> None: - port: Final = _free_port() + port: Final = allocate_port() calls: Final = _stream_calls(20) held: Final = _held() with wire_server(held.respond, port=port) as wire: diff --git a/tests/integration/test_oci_proxy_integration.py b/tests/integration/test_oci_proxy_integration.py index d059a55b1f2..67f0cf53c93 100644 --- a/tests/integration/test_oci_proxy_integration.py +++ b/tests/integration/test_oci_proxy_integration.py @@ -32,7 +32,6 @@ from __future__ import annotations import json import os -import socket import subprocess import sys import time @@ -43,6 +42,7 @@ import httpx import pytest from tests._master_key import MASTER_KEY from tests.integration._support.process import stop_child +from tests.integration._support.ports import allocate_port # --------------------------------------------------------------------------- # Skip gate @@ -59,12 +59,6 @@ STARTUP_TIMEOUT_S = 90.0 REQUEST_TIMEOUT_S = 120.0 -def _free_port() -> int: - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("127.0.0.1", 0)) - return s.getsockname()[1] - - def _wait_for_health(base_url: str, proc: subprocess.Popen, deadline: float) -> None: """Poll /health/liveliness until the proxy answers or the deadline expires.""" while time.monotonic() < deadline: @@ -118,7 +112,7 @@ def _serve(config_path: str) -> Iterator[str]: env.pop("DATABASE_URL", None) env["STORE_MODEL_IN_DB"] = "False" - port = _free_port() + port = allocate_port() base_url = f"http://127.0.0.1:{port}" # Prefer the `litellm` console script that lives next to the active diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 1a6257e9193..26306a2d39a 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -1,10 +1,7 @@ from __future__ import annotations -import asyncio -import errno import importlib import os -import socket import subprocess import sys import threading @@ -37,7 +34,6 @@ from tests.integration._support.process import ( ) TESTS_DIR: Final = Path(__file__).resolve().parents[2] -UNRELATED_CRASH: Final = "Traceback (most recent call last):\nModuleNotFoundError: No module named 'litellm'\n" def test_worker_configuration_routes_stock_edges_to_its_own_servers_and_preserves_other_targets( @@ -111,46 +107,6 @@ def process_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType: return importlib.import_module("integration._support.process") -async def _refused_bind(port: int) -> OSError | None: - try: - await asyncio.get_running_loop().create_server(asyncio.Protocol, "127.0.0.1", port) - except OSError as refused: - return refused - return None - - -@pytest.fixture -def bind_error_line() -> Iterator[str]: - with socket.socket() as held: - held.bind(("127.0.0.1", 0)) - held.listen() - refused: Final = asyncio.run(_refused_bind(held.getsockname()[1])) - assert refused is not None and refused.errno == errno.EADDRINUSE - yield f"ERROR: {refused}\n" - - -def _written_log(directory: Path, text: str) -> Path: - log: Final = directory / "owned-proxy.log" - log.write_text(text) - return log - - -def test_lost_port_race_matches_the_bind_error_the_server_logs( - process_module: ModuleType, bind_error_line: str, tmp_path: Path -) -> None: - assert process_module._lost_port_race(1, _written_log(tmp_path, bind_error_line)) - - -def test_lost_port_race_ignores_an_exit_for_another_reason(process_module: ModuleType, tmp_path: Path) -> None: - assert not process_module._lost_port_race(1, _written_log(tmp_path, UNRELATED_CRASH)) - - -def test_lost_port_race_needs_the_process_to_have_exited( - process_module: ModuleType, bind_error_line: str, tmp_path: Path -) -> None: - assert not process_module._lost_port_race(None, _written_log(tmp_path, bind_error_line)) - - @pytest.mark.parametrize("failure", (psutil.AccessDenied, psutil.NoSuchProcess)) def test_an_unreadable_process_cannot_establish_listener_ownership( process_module: ModuleType, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, failure: type[psutil.Error] diff --git a/tests/unit/test_check_test_quality.py b/tests/unit/test_check_test_quality.py index a0b46cdcdaa..9d6b5cb4d88 100644 --- a/tests/unit/test_check_test_quality.py +++ b/tests/unit/test_check_test_quality.py @@ -640,6 +640,10 @@ _VIOLATING_SNIPPETS: Final = MappingProxyType( ), "TQ010": ("tests/integration/chaos/test_snippet.py", "import os, signal\n\nos.kill(1, signal.SIGKILL)\n"), "TQ011": ("tests/integration/chaos/test_snippet.py", "def test_fast():\n elapsed = run()\n assert elapsed < 2\n"), + "TQ012": ( + "tests/integration/_support/snippet.py", + "def port():\n probe.bind(('127.0.0.1', 0))\n return probe.getsockname()[1]\n", + ), } ) @@ -792,6 +796,20 @@ def test_lower_bounds_relative_bounds_counts_and_polling_deadlines_are_left_alon assert _chaos_codes(tmp_path, source) == [] +@pytest.mark.parametrize( + ("source", "expected"), + ( + ("def port():\n probe.bind(('127.0.0.1', 0))\n return probe.getsockname()[1]\n", ["TQ012"]), + ("def port():\n probe.bind((host, 0))\n _, number = probe.getsockname()\n return number\n", ["TQ012"]), + ("def serve(port):\n probe.bind(('127.0.0.1', port))\n return probe.getsockname()[1]\n", []), + ("def client():\n probe.bind(('127.0.0.1', 0))\n\n\ndef other():\n return probe.getsockname()\n", []), + ), +) +def test_a_port_read_back_from_a_port_zero_bind_is_flagged_anywhere_under_integration(tmp_path, source, expected): + assert _chaos_codes(tmp_path, source, "tests/integration/_support/snippet.py") == expected + assert _chaos_codes(tmp_path, source, "tests/unit/test_snippet.py") == [] + + def test_a_clock_bound_name_is_only_an_elapsed_time_in_the_function_that_binds_it(tmp_path): source = ( "import time\n\n\n" From 9d442b1300966d7c8873be87c50a4ab03ef91db0 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 15 Oct 2026 16:57:05 -0700 Subject: [PATCH 06/23] fix(terraform): stop the perpetual plan diff on litellm_mcp_server (#46781) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 1 + .../provider/litellm/resource_mcp_server.go | 1 + .../litellm/resource_mcp_server_crud.go | 17 +- .../litellm/resource_mcp_server_crud_test.go | 165 ++++++++++++++++++ 4 files changed, 181 insertions(+), 3 deletions(-) diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index e679329073b..4c45008623c 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -64,6 +64,7 @@ longer signal it. - **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead - **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs - **unified_access_group**: create now accepts any 2xx response instead of requiring exactly HTTP 200; `POST /v1/unified_access_group` legitimately returns 201, so creation previously succeeded on the proxy but failed in the provider, leaving the group out of state and forcing a `terraform import` to recover on the next apply's 409. Same fix and shape as the one already applied to `mcp_server`/`model`/`key`/`organization_member`; `handleResponse` (shared by several other resources) was the one status-check helper that fix didn't reach. The legacy `litellm_access_group` resource calls the unrelated `/access_group/new` endpoint, which already returns 200, so it was never affected +- **mcp_server**: A plan right after a clean apply of a minimal litellm_mcp_server no longer shows changes. alias is now computed so the proxy's default (the server name) is kept, spec_version keeps the configured value when the proxy returns none, and an mcp_info block the proxy fills with fields the provider does not model is no longer written to state ### Changed diff --git a/terraform/provider/litellm/resource_mcp_server.go b/terraform/provider/litellm/resource_mcp_server.go index f87a4dab90c..b891f5944be 100644 --- a/terraform/provider/litellm/resource_mcp_server.go +++ b/terraform/provider/litellm/resource_mcp_server.go @@ -24,6 +24,7 @@ func resourceLiteLLMMCPServer() *schema.Resource { "alias": { Type: schema.TypeString, Optional: true, + Computed: true, Description: "Alias for the MCP server", }, "description": { diff --git a/terraform/provider/litellm/resource_mcp_server_crud.go b/terraform/provider/litellm/resource_mcp_server_crud.go index 9d0a405392a..9bd91d10218 100644 --- a/terraform/provider/litellm/resource_mcp_server_crud.go +++ b/terraform/provider/litellm/resource_mcp_server_crud.go @@ -130,7 +130,9 @@ func updateSchemaFromResponse(d *schema.ResourceData, resp *MCPServerResponse) e d.Set("description", resp.Description) d.Set("url", resp.URL) d.Set("transport", resp.Transport) - d.Set("spec_version", resp.SpecVersion) + if resp.SpecVersion != "" { + d.Set("spec_version", resp.SpecVersion) + } d.Set("auth_type", resp.AuthType) d.Set("created_at", resp.CreatedAt) d.Set("created_by", resp.CreatedBy) @@ -151,8 +153,9 @@ func updateSchemaFromResponse(d *schema.ResourceData, resp *MCPServerResponse) e d.Set("args", resp.Args) } - // Set mcp_info - if resp.MCPInfo != nil { + if !mcpInfoHasModeledFields(resp.MCPInfo) { + d.Set("mcp_info", []interface{}{}) + } else { mcpInfoList := make([]map[string]interface{}, 1) mcpInfoMap := make(map[string]interface{}) @@ -180,6 +183,14 @@ func updateSchemaFromResponse(d *schema.ResourceData, resp *MCPServerResponse) e return nil } +func mcpInfoHasModeledFields(info *MCPInfo) bool { + return info != nil && + (info.ServerName != "" || + info.Description != "" || + info.LogoURL != "" || + info.MCPServerCostInfo != nil) +} + func resourceLiteLLMMCPServerCreate(d *schema.ResourceData, m interface{}) error { client, ok := m.(*Client) if !ok { diff --git a/terraform/provider/litellm/resource_mcp_server_crud_test.go b/terraform/provider/litellm/resource_mcp_server_crud_test.go index fe1eb522239..2fb2daab38e 100644 --- a/terraform/provider/litellm/resource_mcp_server_crud_test.go +++ b/terraform/provider/litellm/resource_mcp_server_crud_test.go @@ -1,6 +1,7 @@ package litellm import ( + "context" "encoding/json" "net/http" "net/http/httptest" @@ -184,3 +185,167 @@ func TestMCPServerSchemaWriteOnlyCredentialRules(t *testing.T) { }) } } + +func TestMCPServerReadProducesStablePlan(t *testing.T) { + const responseWithoutModeledMCPInfo = `{"alias":"lit9695_proof","approval_status":"active","args":[],"auth_type":"none","command":null,"created_at":"2026-10-15T22:48:21.530000Z","created_by":"default_user_id","credentials":null,"description":null,"env":{},"mcp_access_groups":[],"mcp_info":{"is_public":false,"is_public_explicit":false},"server_id":"dc577aa3-22ac-4fdf-a56b-f6c40e35c389","server_name":"lit9695_proof","status":"unknown","transport":"http","updated_at":"2026-10-15T22:48:21.530000Z","updated_by":"default_user_id","url":"https://mcp.example.com/mcp"}` + responseWithLogoURL := strings.Replace( + responseWithoutModeledMCPInfo, + `"mcp_info":{"is_public":false,"is_public_explicit":false}`, + `"mcp_info":{"is_public":false,"is_public_explicit":false,"logo_url":"https://x/logo.png"}`, + 1, + ) + if responseWithLogoURL == responseWithoutModeledMCPInfo { + t.Fatal("failed to add logo_url to MCP info response") + } + + cases := []struct { + name string + config map[string]interface{} + response string + wantEmptyDiff bool + wantChangedKey string + wantAlias string + wantSpecVersion string + wantMCPInfoLogoURL string + }{ + { + name: "minimal config", + config: map[string]interface{}{ + "server_name": "lit9695_proof", + "url": "https://mcp.example.com/mcp", + "transport": "http", + }, + response: responseWithoutModeledMCPInfo, + wantEmptyDiff: true, + wantAlias: "lit9695_proof", + wantSpecVersion: "2024-11-05", + }, + { + name: "configured spec version", + config: map[string]interface{}{ + "server_name": "lit9695_proof", + "url": "https://mcp.example.com/mcp", + "transport": "http", + "spec_version": "2025-03-26", + }, + response: responseWithoutModeledMCPInfo, + wantEmptyDiff: true, + wantSpecVersion: "2025-03-26", + }, + { + name: "configured alias drift remains visible", + config: map[string]interface{}{ + "server_name": "lit9695_proof", + "alias": "custom", + "url": "https://mcp.example.com/mcp", + "transport": "http", + }, + response: responseWithoutModeledMCPInfo, + wantChangedKey: "alias", + }, + { + name: "configured MCP info drift remains visible", + config: map[string]interface{}{ + "server_name": "lit9695_proof", + "url": "https://mcp.example.com/mcp", + "transport": "http", + "mcp_info": []interface{}{ + map[string]interface{}{"logo_url": "https://x/logo.png"}, + }, + }, + response: responseWithoutModeledMCPInfo, + wantChangedKey: "mcp_info", + }, + { + name: "modeled MCP info is retained", + config: map[string]interface{}{ + "server_name": "lit9695_proof", + "url": "https://mcp.example.com/mcp", + "transport": "http", + "mcp_info": []interface{}{ + map[string]interface{}{"logo_url": "https://x/logo.png"}, + }, + }, + response: responseWithLogoURL, + wantEmptyDiff: true, + wantMCPInfoLogoURL: "https://x/logo.png", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost && r.Method != http.MethodGet { + t.Errorf("unexpected request method: %s", r.Method) + w.WriteHeader(http.StatusNotFound) + return + } + if r.URL.Path != "/v1/mcp/server" && r.URL.Path != "/v1/mcp/server/dc577aa3-22ac-4fdf-a56b-f6c40e35c389" { + t.Errorf("unexpected request path: %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(tc.response)) + })) + defer server.Close() + + resource := resourceLiteLLMMCPServer() + d := schema.TestResourceDataRaw(t, resource.Schema, tc.config) + client := NewClient(server.URL, "test-key", true) + if err := resourceLiteLLMMCPServerCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + if err := resourceLiteLLMMCPServerRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + diff, err := resource.Diff(context.Background(), d.State(), terraform.NewResourceConfigRaw(tc.config), nil) + if err != nil { + t.Fatalf("plan diff failed: %v", err) + } + if tc.wantEmptyDiff { + if diff != nil && !diff.Empty() { + t.Fatalf("expected empty plan diff, got %v", diff.Attributes) + } + if diff != nil { + for key := range diff.Attributes { + if key == "alias" || key == "spec_version" || strings.HasPrefix(key, "mcp_info.") { + t.Fatalf("unexpected diff on %s", key) + } + } + } + } else { + if diff == nil || diff.Empty() { + t.Fatal("expected a plan diff") + } + hasExpectedDiff := false + for key := range diff.Attributes { + hasExpectedDiff = hasExpectedDiff || key == tc.wantChangedKey || + (tc.wantChangedKey == "mcp_info" && strings.HasPrefix(key, "mcp_info.")) + } + if !hasExpectedDiff { + t.Fatalf("expected a plan diff on %s, got %v", tc.wantChangedKey, diff.Attributes) + } + } + + if tc.wantAlias != "" && d.Get("alias") != tc.wantAlias { + t.Fatalf("alias = %v, want %s", d.Get("alias"), tc.wantAlias) + } + if tc.wantSpecVersion != "" && d.Get("spec_version") != tc.wantSpecVersion { + t.Fatalf("spec_version = %v, want %s", d.Get("spec_version"), tc.wantSpecVersion) + } + if tc.wantMCPInfoLogoURL != "" { + mcpInfo := d.Get("mcp_info").([]interface{}) + if len(mcpInfo) != 1 { + t.Fatalf("mcp_info has %d blocks, want 1", len(mcpInfo)) + } + gotLogoURL := mcpInfo[0].(map[string]interface{})["logo_url"] + if gotLogoURL != tc.wantMCPInfoLogoURL { + t.Fatalf("mcp_info logo_url = %v, want %s", gotLogoURL, tc.wantMCPInfoLogoURL) + } + } + }) + } +} From 5fd6379a7826611365ca34434367d8289cfc9bf7 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 15 Oct 2026 17:14:20 -0700 Subject: [PATCH 07/23] test(e2e): perf users start staggered with jitter, tool-call streams get a first token, small_bodies_5rps sends from the mixed case's callers (#46776) * test(e2e): perf streams take a reasoning or tool-call delta as the first token, tool-call replies had no time_to_first_token * test(e2e): perf users start at a random offset and jitter each interval by up to 25%, no more lockstep bursts at the planned rate * test(e2e): small_bodies_5rps sends from the mixed case's callers (two teams, JWT users, keys in and out of teams) through the pre and post guardrail --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/perf/cases/small_bodies_5rps.py | 39 +++++++++++---- tests/e2e/perf/setup/locustfile.py | 58 +++++++++++++++++++---- 2 files changed, 79 insertions(+), 18 deletions(-) diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index e118e453396..52fe41623e3 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -1,20 +1,23 @@ -"""Small bodies at a low rate on the small gateway: the shape of github issue #46276 (short chat -requests, master key and virtual keys, no JWT, no guardrails) at 5 rps on one 4-worker pod, far -under what the pod sustains, so nothing queues and gateway_added_seconds is the gateway's own -per-request overhead. A slower middleware chain or an extra Redis round trip per request shows up -directly instead of hiding behind queue wait.""" +"""Small bodies at a low rate on the small gateway: the request sizes of github issue #46276 (short +chat requests, half streamed) from the same callers as the mixed case (two teams, JWT users in and +out of a team, virtual keys, the pre and post guardrail) at 5 rps on one 4-worker pod, far under +what the pod sustains, so nothing queues and gateway_added_seconds is the gateway's own per-request +overhead. A slower middleware chain or an extra Redis round trip per request shows up directly +instead of hiding behind queue wait.""" from __future__ import annotations from typing import Final from setup.performance_test_case import ( + Guardrail, Limit, MockLatency, PerformanceTestCase, RequestShape, Team, Traffic, + User, VirtualKey, ) @@ -22,16 +25,32 @@ TINY: Final = RequestShape("tiny", input_tokens=(2_500, 4_000)) CASE: Final = PerformanceTestCase( name="small_bodies_5rps", - teams=(Team("team-a", models=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock")),), - users=(), + teams=( + Team("team-a", models=("gpt-6-luna", "claude-sonnet-5"), max_budget=5000.0, rpm_limit=6000), + Team("team-b", models=("claude-sonnet-5-bedrock", "claude-sonnet-5"), max_budget=5000.0), + ), + users=( + User("alice", team="team-a"), + User("bob", team="team-a"), + User("carol", team="team-a"), + User("dave", team="team-b"), + User("erin", team="team-b"), + User("frank", team="team-b"), + User("grace"), + User("heidi"), + User("ivan"), + User("judy", max_budget=1000.0), + ), virtual_keys=( - VirtualKey("key-a1", team="team-a"), - VirtualKey("key-a2", team="team-a"), + VirtualKey("key-a1", team="team-a", models=("gpt-6-luna",), rpm_limit=1200), + VirtualKey("key-a2", team="team-a", user="alice"), + VirtualKey("key-b1", team="team-b"), + VirtualKey("key-b2", team="team-b", user="dave"), VirtualKey("key-plain"), ), model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"), mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25), - guardrails=(), + guardrails=(Guardrail("perf-pre", mode="pre_call"), Guardrail("perf-post", mode="post_call")), traffic=Traffic( seed=11, rps=5.0, diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py index fe8080959e5..1609617bd73 100644 --- a/tests/e2e/perf/setup/locustfile.py +++ b/tests/e2e/perf/setup/locustfile.py @@ -10,13 +10,14 @@ line), paces to `PERF_RPS` over `PERF_USERS` users, and quits once the last requ from __future__ import annotations import os +import random import time from collections.abc import Iterator from pathlib import Path from typing import Final, cast import gevent -from locust import FastHttpUser, constant_throughput, task +from locust import FastHttpUser, task from locust.contrib.fasthttp import FastResponse, ResponseContextManager from setup.wire import Endpoint, WireRecord, WireRequest, request_body from pydantic import TypeAdapter @@ -28,6 +29,9 @@ _RECORDS_PATH: Final = Path(os.environ["PERF_RECORDS_FILE"]) _RPS: Final = float(os.environ["PERF_RPS"]) _USERS: Final = int(os.environ["PERF_USERS"]) _NEXT: Final[Iterator[WireRequest]] = iter(_PLAN) +_INTERVAL_S: Final = _USERS / _RPS +_JITTER: Final = 0.25 +_RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan") _RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8") _written = 0 BLOCK_REASON: Final = "perf-test-block" @@ -84,6 +88,32 @@ def delta_of_event(endpoint: Endpoint, payload: dict[str, object]) -> str: return "" +_RESPONSES_TOKEN_EVENTS: Final = frozenset( + { + "response.output_text.delta", + "response.refusal.delta", + "response.reasoning_text.delta", + "response.reasoning_summary_text.delta", + "response.function_call_arguments.delta", + } +) + + +def token_in_event(endpoint: Endpoint, payload: dict[str, object]) -> bool: + """Whether one SSE event carries a generated token of any kind: visible text, reasoning or tool-call arguments. + The first such event is the stream's first token. `delta_of_event` stays visible text only because the body + check compares it with the text the mock generated, and a tool-call reply has none.""" + if delta_of_event(endpoint, payload): + return True + if endpoint == "chat_completions": + choices: Final = _as_list(payload.get("choices")) + delta: Final = _as_dict(_as_dict(choices[0]).get("delta")) if choices else {} + return any(delta.get(key) for key in ("reasoning_content", "reasoning", "tool_calls", "function_call")) + if endpoint == "responses": + return payload.get("type") in _RESPONSES_TOKEN_EVENTS + return payload.get("type") == "content_block_delta" + + def _events(buffer: bytes) -> tuple[list[dict[str, object]], bytes]: """Complete SSE events in `buffer` as their JSON data, plus the unfinished tail.""" events: list[dict[str, object]] = [] @@ -169,12 +199,26 @@ def _authorization(user: FastHttpUser, item: WireRequest) -> str: class PlanUser(FastHttpUser): - """One simulated client; all users share the plan and together hold the planned rate.""" + """One simulated client; all users share the plan and together hold the planned rate. Each user starts at + a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the + users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate.""" - wait_time = constant_throughput(_RPS / _USERS) network_timeout = 120.0 connection_timeout = 30.0 + def on_start(self) -> None: + self._last_run = 0.0 + self._last_wait = 0.0 + gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S)) + + def wait_time(self) -> float: + target: Final = _INTERVAL_S * _RNG.uniform(1.0 - _JITTER, 1.0 + _JITTER) + now: Final = time.monotonic() + run_s: Final = now - self._last_run - self._last_wait if self._last_run else 0.0 + self._last_wait = max(0.0, target - run_s) + self._last_run = now + return self._last_wait + @task def send(self) -> None: item: Final = next(_NEXT, None) @@ -221,11 +265,9 @@ class PlanUser(FastHttpUser): buffer += line events, buffer = _events(buffer) for event in events: - delta = delta_of_event(item.endpoint, event) - if delta: - if first_token_s is None: - first_token_s = now - text += delta + if first_token_s is None and token_in_event(item.endpoint, event): + first_token_s = now + text += delta_of_event(item.endpoint, event) else: body_text = context.text or "" first_byte_s = time.perf_counter() - started From 67bfb7a1481362d56f988920f76091a2212358ae Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:17:51 -0700 Subject: [PATCH 08/23] fix(ollama): prevent reused response ids from dropping spend logs (#39227) * fix(spend_tracking): key completion spend logs off the per-call id so duplicate provider response ids cannot drop rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): narrow the provider response id with isinstance instead of cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(gcs_bucket): store request payloads under the spend log request id so the logs viewer can find them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(spend_tracking): validate callback fields for current type gates * test(gcs_pubsub): assert provider response id in spend metadata * fix(ollama): normalize reused response ids automatically --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../convert_dict_to_response.py | 5 +- .../streaming_chunk_builder_utils.py | 7 ++ .../litellm_core_utils/streaming_handler.py | 11 ++- litellm/llms/ollama/common_utils.py | 47 +++++++++- .../llms/openai/chat/gpt_transformation.py | 2 + .../spend/test_spend_logs_metadata_fields.py | 92 +++++++++++++++++++ .../gcs_bucket/test_gcs_bucket_base.py | 29 ++++++ .../test_convert_dict_to_response.py | 48 ++++++++++ .../test_streaming_handler.py | 88 ++++++++++++++++++ 9 files changed, 326 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index f32fba7c6ec..f3c60679c81 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -16,6 +16,7 @@ from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_reasoning_content, ) +from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.databricks import DatabricksTool from litellm.types.llms.openai import ( ChatCompletionReasoningItem, @@ -833,7 +834,9 @@ def convert_to_model_response_object( if k not in _MODEL_RESPONSE_FIELDS: setattr(model_response_object, k, v) - return model_response_object + ollama_ids: Final = OllamaResponseIdNormalizer() + ollama_ids.observe(model_response_object) + return ollama_ids.apply(model_response_object) elif response_type == "embedding" and ( model_response_object is None or isinstance(model_response_object, EmbeddingResponse) ): diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 6c18128a608..d505192184e 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1,4 +1,5 @@ import base64 +import builtins import time from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, replace @@ -118,6 +119,7 @@ class _BaseChunk(TypedDict, total=False): created: ReadOnly[int] model: ReadOnly[str] system_fingerprint: ReadOnly[str | None] + provider_specific_fields: ReadOnly[Mapping[str, builtins.object] | None] service_tier: ReadOnly[str | None] choices: ReadOnly[Required[Sequence[StreamingChoices]]] _hidden_params: ReadOnly[_ChunkHiddenParams] @@ -430,6 +432,11 @@ class ChunkProcessor: "created": created, "model": model, "system_fingerprint": system_fingerprint, + **( + {"provider_specific_fields": chunks[0].get("provider_specific_fields")} + if system_fingerprint == "fp_ollama" + else {} + ), **( MappingProxyType({"service_tier": service_tier}) if service_tier is not None diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 83511456d45..46ad1b70abe 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -14,6 +14,7 @@ from typing import Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx +from openai.types.chat import ChatCompletionChunk from pydantic import BaseModel, ValidationError from typing_extensions import NotRequired, TypedDict @@ -27,6 +28,7 @@ from litellm.litellm_core_utils.model_response_utils import ( ) from litellm.litellm_core_utils.redact_messages import LiteLLMLoggingObject from litellm.litellm_core_utils.thread_pool_executor import executor +from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import OpenAIChatCompletionChunk from litellm.types.router import GenericLiteLLMParams @@ -285,6 +287,7 @@ class CustomStreamWrapper: self._response_headers = _response_headers self.response_id: str | None = None + self._ollama_response_ids = OllamaResponseIdNormalizer() self.received_model_id: str | None = None self.logging_loop = None self.rules = Rules() @@ -824,7 +827,7 @@ class CustomStreamWrapper: pass else: model_response.choices = [StreamingChoices(finish_reason=None)] - return model_response + return self._ollama_response_ids.apply(model_response) def is_delta_empty(self, delta: Delta) -> bool: is_empty = True @@ -1449,8 +1452,14 @@ class CustomStreamWrapper: return _ProviderChunkParsed(response_obj) def chunk_creator(self, chunk: Any): + if self.custom_llm_provider != "cached_response" and isinstance( + chunk, (ChatCompletionChunk, ModelResponseStream) + ): + self._ollama_response_ids.observe(chunk) if hasattr(chunk, "id"): self.response_id = chunk.id + if self._ollama_response_ids.response_id is not None: + self.response_id = self._ollama_response_ids.response_id provider_response_model: Final = _provider_response_model(chunk) if provider_response_model is not None: self._provider_response_model = provider_response_model diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index af8d492316e..b5328b1ba4d 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -1,11 +1,56 @@ import base64 import io -from typing import Final +from typing import Final, overload import httpx +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger +from litellm._uuid import uuid from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.types.utils import ModelResponse, ModelResponseStream + +_PROVIDER_FIELDS: Final = TypeAdapter(dict[str, object]) + + +class OllamaResponseIdNormalizer: + def __init__(self) -> None: + self.response_id: str | None = None + self.original_response_id: str | None = None + + def observe(self, response: object) -> None: + if self.response_id is not None: + return + if not isinstance(response, (ModelResponse, ModelResponseStream, ChatCompletionChunk)): + return + if response.system_fingerprint != "fp_ollama": + return + fields: Final = _PROVIDER_FIELDS.validate_python( + BaseModel.model_dump(response).get("provider_specific_fields") or {} + ) + original_id: Final = fields.get("original_response_id") + self.original_response_id = original_id if isinstance(original_id, str) else response.id + self.response_id = response.id if isinstance(original_id, str) else f"chatcmpl-{uuid.uuid4()}" + + @overload + def apply(self, response: ModelResponse) -> ModelResponse: ... + + @overload + def apply(self, response: ModelResponseStream) -> ModelResponseStream: ... + + def apply(self, response: ModelResponse | ModelResponseStream) -> ModelResponse | ModelResponseStream: + if self.response_id is None: + return response + fields: Final = _PROVIDER_FIELDS.validate_python( + BaseModel.model_dump(response).get("provider_specific_fields") or {} + ) + return response.model_copy( + update={ + "id": self.response_id, + "provider_specific_fields": {**fields, "original_response_id": self.original_response_id}, + } + ) class OllamaError(BaseLLMException): diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 3b70732f3fb..4d7ef833181 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -958,9 +958,11 @@ class OpenAIChatCompletionStreamingHandler(BaseModelResponseIterator): ) choices = chunk.get("choices", []) choices = self._map_reasoning_to_reasoning_content(choices) + fingerprint: Final = _OBJECT_MAPPING.validate_python(chunk).get("system_fingerprint") kwargs: Final[dict[str, Any]] = { "id": chunk.get("id"), + "system_fingerprint": fingerprint, "object": "chat.completion.chunk", "created": chunk.get("created"), "model": chunk.get("model"), diff --git a/tests/integration/spend/test_spend_logs_metadata_fields.py b/tests/integration/spend/test_spend_logs_metadata_fields.py index d7f1dbfd998..9941edf0f26 100644 --- a/tests/integration/spend/test_spend_logs_metadata_fields.py +++ b/tests/integration/spend/test_spend_logs_metadata_fields.py @@ -21,6 +21,98 @@ from pydantic import JsonValue CACHED_PROMPT_TOKENS: Final = 4 +def _ollama_fixed_id_reply(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + assert request.target == "/v1/chat/completions", request.target + assert body["model"] == "gpt-oss:20b", body + common: Final = { + "id": "chatcmpl-0", + "model": "gpt-oss:20b", + "created": 1, + "system_fingerprint": "fp_ollama", + } + usage: Final = {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6} + if body.get("stream") is True: + chunks: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "h"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {"content": "i"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {"choices": [], "usage": usage}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple( + ("data: " + json.dumps({**common, "object": "chat.completion.chunk", **chunk}) + "\n\n").encode() + for chunk in chunks + ) + + (b"data: [DONE]\n\n",), + ) + return Reply( + body=json.dumps( + { + **common, + "object": "chat.completion", + "usage": usage, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + } + ).encode() + ) + + +@pytest.mark.parametrize("stream", [False, True]) +def test_reused_ollama_response_ids_keep_every_spend_row(gateway: Gateway, stream: bool) -> None: + with wire_server(_ollama_fixed_id_reply) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-oss:20b", + api_base=wire.url + "/v1", + api_key="synthetic-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + key: Final = scenario.key(models=[model]) + responses: Final = tuple( + gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"hi {index}"}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + for index in range(3) + ) + assert tuple(response.status_code for response in responses) == (200, 200, 200) + returned: Final = tuple( + tuple( + JSON_OBJECT.validate_json(line[6:]) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + if stream + else (JSON_OBJECT.validate_json(response.content),) + for response in responses + ) + ids: Final = tuple(string_value(chunks[0]["id"]) for chunks in returned) + for chunks, response_id in zip(returned, ids, strict=True): + assert {chunk["id"] for chunk in chunks} == {response_id} + assert object_value(chunks[0]["provider_specific_fields"])["original_response_id"] == "chatcmpl-0" + assert object_value(chunks[-1]["usage"])["total_tokens"] == 6 + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ), + lambda values: len(values) == 3, + seconds=70, + ) + assert len(set(ids)) == 3 + assert {row["request_id"] for row in rows} == set(ids) + assert tuple(float(str(row["spend"])) for row in rows) == pytest.approx((0.007, 0.007, 0.007)) + + def _config( tmp_path: Path, spend_logs_metadata_fields: Mapping[str, JsonValue] | None, diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py index 781b56cae14..07874bdbeb5 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -1,4 +1,5 @@ import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -7,6 +8,34 @@ from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase +def test_ollama_payload_key_matches_returned_response_id() -> None: + from datetime import datetime + + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import ModelResponse + + logger: Final = object.__new__(GCSBucketLogger) + response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-0", + "system_fingerprint": "fp_ollama", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}], + }, + model_response_object=ModelResponse(), + ) + logged: Final = get_logging_payload( + kwargs={"call_type": "acompletion", "response_cost": 0.01}, + response_obj=response, + start_time=datetime(2025, 1, 1), + end_time=datetime(2025, 1, 1), + ) + object_name: Final = logger._get_object_name({}, {"error_str": None}, response) + assert response.id != "chatcmpl-0" + assert object_name.split("/", 1)[1] == response.id + assert logged["request_id"] == response.id + + class TestGCSBucketBase: def test_construct_request_headers_with_project_id(self, monkeypatch): """Test that construct_request_headers correctly uses project_id if passed from env""" diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index f918d4cb79e..0c601a6b2c8 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -14,6 +14,54 @@ from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outco FIXED_TIME: Final[datetime] = datetime(2025, 1, 1, 12, 0, 0) + +def test_reused_ollama_ids_have_distinct_response_and_spend_ids() -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + body: Final = { + "id": "chatcmpl-0", + "system_fingerprint": "fp_ollama", + "model": "gpt-oss:20b", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + "provider_specific_fields": {"extra": "preserved"}, + } + responses: Final = [ + convert_to_model_response_object(response_object=body, model_response_object=ModelResponse()) for _ in range(2) + ] + assert len({response.id for response in responses}) == 2 + for response in responses: + assert response.id != "chatcmpl-0" + assert response.provider_specific_fields == {"extra": "preserved", "original_response_id": "chatcmpl-0"} + logged: Final = get_logging_payload( + kwargs={"call_type": "acompletion", "litellm_call_id": "same-client-call-id", "response_cost": 0.01}, + response_obj=response, + start_time=FIXED_TIME, + end_time=FIXED_TIME, + ) + assert logged["request_id"] == response.id + reparsed: Final = convert_to_model_response_object( + response_object=response.model_dump(), model_response_object=ModelResponse() + ) + assert reparsed.id == response.id + assert body["id"] == "chatcmpl-0" + assert body["provider_specific_fields"] == {"extra": "preserved"} + + +@pytest.mark.parametrize("fingerprint", [None, "fp_openai", "fp_ollama_other"]) +def test_non_ollama_fingerprints_keep_provider_response_ids(fingerprint: str | None) -> None: + response: Final = convert_to_model_response_object( + response_object={ + "id": "chatcmpl-0", + "system_fingerprint": fingerprint, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}}], + }, + model_response_object=ModelResponse(), + ) + assert response.id == "chatcmpl-0" + assert response.model_dump().get("provider_specific_fields") is None + + OPENAI_CUSTOM_TOOL_CALL_RESPONSE = { "id": "chatcmpl-abc", "created": 1784657740, diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 07cdbc15baa..4347a6639bd 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -31,6 +31,94 @@ from litellm import completion from typing import Tuple +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize("fingerprint", ["fp_ollama", "fp_openai", None]) +async def test_ollama_stream_ids_match_usage_and_spend_logs( + async_stream: bool, + transport: str, + fingerprint: str | None, +) -> None: + from datetime import datetime + + from openai.types.chat import ChatCompletionChunk + + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object + from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + parser: Final = OpenAIChatCompletionStreamingHandler(streaming_response=None, sync_stream=True) + bodies: Final = [ + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "h"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {"content": "i"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + {"choices": [], "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}}, + ] + raw_chunks: Final = [ + { + "id": "chatcmpl-0", + "object": "chat.completion.chunk", + "model": "gpt-oss:20b", + "created": 1, + "system_fingerprint": fingerprint, + **body, + } + for body in bodies + ] + chunks: Final = [ + ChatCompletionChunk(**body) if transport == "sdk" else parser.chunk_parser(body) for body in raw_chunks + ] + streams: Final = [ + CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="gpt-oss:20b", + custom_llm_provider="openai", + stream_options={"include_usage": True}, + logging_obj=Logging( + model="gpt-oss:20b", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=1.0, + litellm_call_id="same-client-call-id", + function_id="test", + ), + ) + for _ in range(2) + ] + received_streams: Final = ( + [chunk async for chunk in streams[0]] if async_stream else list(streams[0]), + [chunk async for chunk in streams[1]] if async_stream else list(streams[1]), + ) + for received in received_streams: + assembled: Final = litellm.stream_chunk_builder(chunks=received) + assert assembled.choices[0].message.content == "hi" + assert assembled.usage.total_tokens == 6 + assert len({chunk.id for chunk in received}) == 1 + assert received[-1].usage.total_tokens == 6 + assert assembled.id == received[0].id + reparsed: Final = convert_to_model_response_object( + response_object=assembled.model_dump(), + model_response_object=ModelResponse(), + ) + assert reparsed.id == assembled.id + logged: Final = get_logging_payload( + kwargs={"call_type": "acompletion", "litellm_call_id": "same-client-call-id", "response_cost": 0.01}, + response_obj=assembled, + start_time=datetime(2025, 1, 1), + end_time=datetime(2025, 1, 1), + ) + assert logged["request_id"] == assembled.id + if fingerprint == "fp_ollama": + assert assembled.id != "chatcmpl-0" + assert received[0].provider_specific_fields == {"original_response_id": "chatcmpl-0"} + assert assembled.provider_specific_fields == {"original_response_id": "chatcmpl-0"} + else: + assert assembled.id == "chatcmpl-0" + assert len({chunks[0].id for chunks in received_streams}) == (2 if fingerprint == "fp_ollama" else 1) + assert all(chunk.id == "chatcmpl-0" for chunk in chunks) + + @pytest.fixture def initialized_custom_stream_wrapper() -> CustomStreamWrapper: streaming_handler = CustomStreamWrapper( From d387cdc81628d127657bebda4024d19fd83ee7e6 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:23:57 -0700 Subject: [PATCH 09/23] fix(vertex-ai): end free period for deepseek-v4.1-flash-maas (#46793) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 308dfdeb62a..6f0630dd01b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -75076,14 +75076,14 @@ "supports_vision": true }, "vertex_ai/deepseek-ai/deepseek-v4.1-flash-maas": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_token": 0.0, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-deepseek_models", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 0.0, + "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_regions": [ "global" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 308dfdeb62a..6f0630dd01b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -75076,14 +75076,14 @@ "supports_vision": true }, "vertex_ai/deepseek-ai/deepseek-v4.1-flash-maas": { - "cache_read_input_token_cost": 0.0, - "input_cost_per_token": 0.0, + "cache_read_input_token_cost": 3e-08, + "input_cost_per_token": 3e-07, "litellm_provider": "vertex_ai-deepseek_models", "max_input_tokens": 1048576, "max_output_tokens": 1048576, "max_tokens": 1048576, "mode": "chat", - "output_cost_per_token": 0.0, + "output_cost_per_token": 1.2e-06, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_regions": [ "global" From 6c9c3a909862ad370f37f2cfa353ffdfd130a711 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:25:54 -0700 Subject: [PATCH 10/23] feat(cost-map): add baseten Muse Spark Open and Nemotron 3.5 Super VL (#46794) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 35 +++++++++++++++++++ model_prices_and_context_window.json | 35 +++++++++++++++++++ 2 files changed, 70 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6f0630dd01b..46fd8b21318 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -75128,5 +75128,40 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "baseten/meta-models/muse-spark-open": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://inference.baseten.co/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "baseten/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6f0630dd01b..46fd8b21318 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -75128,5 +75128,40 @@ "supports_vision": true, "supports_xhigh_reasoning_effort": true, "thinking_always_on": true + }, + "baseten/meta-models/muse-spark-open": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.25e-06, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 4.25e-06, + "source": "https://inference.baseten.co/v1/models", + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": false, + "supports_tool_choice": true, + "supports_vision": true + }, + "baseten/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 2e-07, + "litellm_provider": "baseten", + "max_input_tokens": 1048576, + "max_output_tokens": 1048576, + "max_tokens": 1048576, + "mode": "chat", + "output_cost_per_token": 5e-07, + "source": "https://inference.baseten.co/v1/models", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } From 9c62573005c8023b15283153c27bc24f79527d30 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Thu, 15 Oct 2026 17:27:03 -0700 Subject: [PATCH 11/23] fix(proxy): limit gzip to dashboard analytics and offload compression (#46712) * fix(proxy): limit gzip to dashboard analytics and offload compression * fix(proxy): preserve gzip for customer analytics aliases --- litellm/proxy/middleware/gzip_middleware.py | 48 +++++- .../proxy/middleware/test_gzip_middleware.py | 144 +++++++++++++----- 2 files changed, 152 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/middleware/gzip_middleware.py b/litellm/proxy/middleware/gzip_middleware.py index 1b80fbde2fd..4042fd26a18 100644 --- a/litellm/proxy/middleware/gzip_middleware.py +++ b/litellm/proxy/middleware/gzip_middleware.py @@ -4,11 +4,50 @@ from typing import Final import anyio.to_thread from starlette.datastructures import Headers, MutableHeaders +from starlette.routing import get_route_path from starlette.types import ASGIApp, Message, Receive, Scope, Send MINIMUM_SIZE_BYTES: Final = 500 -OFF_LOOP_SIZE_BYTES: Final = 16 * 1024 COMPRESS_LEVEL: Final = 6 +DASHBOARD_GZIP_PATHS: Final = frozenset( + { + "/user/daily/activity/aggregated", + "/user/daily/activity/aggregated/keys", + "/user/daily/activity/aggregated/search", + "/user/daily/activity/aggregated/model_top_keys", + "/user/daily/activity/aggregated/users", + "/user/daily/activity/aggregated/cache_leakage_keys", + "/team/daily/activity/aggregated", + "/team/daily/activity/aggregated/keys", + "/team/daily/activity/aggregated/search", + "/team/daily/activity/aggregated/model_top_keys", + "/organization/daily/activity/aggregated", + "/organization/daily/activity/aggregated/keys", + "/organization/daily/activity/aggregated/search", + "/organization/daily/activity/aggregated/model_top_keys", + "/tag/daily/activity/aggregated", + "/tag/daily/activity/aggregated/keys", + "/tag/daily/activity/aggregated/search", + "/tag/daily/activity/aggregated/model_top_keys", + "/customer/daily/activity/aggregated", + "/customer/daily/activity/aggregated/keys", + "/customer/daily/activity/aggregated/search", + "/customer/daily/activity/aggregated/model_top_keys", + "/end_user/daily/activity/aggregated", + "/end_user/daily/activity/aggregated/keys", + "/end_user/daily/activity/aggregated/search", + "/end_user/daily/activity/aggregated/model_top_keys", + "/agent/daily/activity/aggregated", + "/agent/daily/activity/aggregated/keys", + "/agent/daily/activity/aggregated/search", + "/agent/daily/activity/aggregated/model_top_keys", + "/gateway/daily/activity", + "/gateway/errors/activity", + "/team/spend/by_user", + "/global/spend/provider", + "/auto_router/benchmarks", + } +) def _coding_weight(part: str) -> tuple[str, float]: @@ -26,8 +65,6 @@ def accepts_gzip(accept_encoding: str) -> bool: async def _compress(body: bytes) -> bytes: - if len(body) < OFF_LOOP_SIZE_BYTES: - return gzip.compress(body, compresslevel=COMPRESS_LEVEL) return await anyio.to_thread.run_sync(gzip.compress, body, COMPRESS_LEVEL) @@ -83,11 +120,12 @@ class _BufferedBodyGzipResponder: class GZipBufferedResponseMiddleware: - def __init__(self, app: ASGIApp) -> None: + def __init__(self, app: ASGIApp, paths: frozenset[str] = DASHBOARD_GZIP_PATHS) -> None: self.app = app + self.paths: Final = paths async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: - if scope["type"] != "http": + if scope["type"] != "http" or scope["method"] != "GET" or get_route_path(scope) not in self.paths: await self.app(scope, receive, send) return gzip_accepted: Final = accepts_gzip(Headers(scope=scope).get("accept-encoding", "")) diff --git a/tests/unit/proxy/middleware/test_gzip_middleware.py b/tests/unit/proxy/middleware/test_gzip_middleware.py index a1c294fb80b..6f9864be186 100644 --- a/tests/unit/proxy/middleware/test_gzip_middleware.py +++ b/tests/unit/proxy/middleware/test_gzip_middleware.py @@ -9,17 +9,19 @@ from starlette.applications import Starlette from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route +from starlette.testclient import TestClient from starlette.types import ASGIApp, Message, Receive, Scope, Send from litellm.proxy.middleware import gzip_middleware from litellm.proxy.middleware.gzip_middleware import ( + DASHBOARD_GZIP_PATHS, MINIMUM_SIZE_BYTES, - OFF_LOOP_SIZE_BYTES, GZipBufferedResponseMiddleware, ) LARGE_PAYLOAD = {"rows": [{"date": f"2026-09-{day:02d}", "spend": day * 1.5} for day in range(1, 31)] * 20} BASE64_EMBEDDINGS_RESPONSE_BYTES: Final = 824_000 +HUGE_RESPONSE_BYTES: Final = 2 * 1024 * 1024 STREAM_CHUNKS = tuple(json.dumps({"part": part, "pad": "x" * MINIMUM_SIZE_BYTES}).encode() for part in range(3)) @@ -48,15 +50,15 @@ async def _no_transform(request: Request) -> Response: async def _huge(request: Request) -> Response: - return Response(b"z" * (OFF_LOOP_SIZE_BYTES * 2), media_type="application/json") + return Response(b"z" * HUGE_RESPONSE_BYTES, media_type="application/json") -async def _under_off_loop(request: Request) -> Response: - return Response(b"u" * (OFF_LOOP_SIZE_BYTES - 1), media_type="application/json") +async def _minimum_eligible(request: Request) -> Response: + return Response(b"u" * MINIMUM_SIZE_BYTES, media_type="application/json") -async def _at_off_loop(request: Request) -> Response: - return Response(b"a" * OFF_LOOP_SIZE_BYTES, media_type="application/json") +async def _four_kib(request: Request) -> Response: + return Response(b"a" * 4096, media_type="application/json") async def _base64_embeddings_sized(request: Request) -> Response: @@ -79,19 +81,30 @@ APP = Starlette( Route("/stream", _json_stream), Route("/etag", _with_etag), Route("/huge", _huge), - Route("/under-off-loop", _under_off_loop), - Route("/at-off-loop", _at_off_loop), + Route("/minimum-eligible", _minimum_eligible), + Route("/four-kib", _four_kib), Route("/base64-embeddings-sized", _base64_embeddings_sized), Route("/partial", _partial), Route("/no-transform", _no_transform), ] ) -APP.add_middleware(GZipBufferedResponseMiddleware) +APP.add_middleware( + GZipBufferedResponseMiddleware, paths=frozenset(route.path for route in APP.routes if isinstance(route, Route)) +) -async def _send_messages(path: str, accept_encoding: str | None, app: ASGIApp = APP) -> tuple[Message, ...]: +async def _send_messages( + path: str, accept_encoding: str | None, app: ASGIApp = APP, method: str = "GET", root_path: str = "" +) -> tuple[Message, ...]: headers = [(b"accept-encoding", accept_encoding.encode())] if accept_encoding is not None else [] - scope = {"type": "http", "method": "GET", "path": path, "query_string": b"", "headers": headers} + scope = { + "type": "http", + "method": method, + "path": path, + "root_path": root_path, + "query_string": b"", + "headers": headers, + } sent: list[Message] = [] # mutable-ok: ASGI send callback collects messages in order requests: Final = iter(({"type": "http.request", "body": b"", "more_body": False},)) never_disconnects: Final = asyncio.Event() @@ -133,21 +146,19 @@ async def test_large_buffered_json_is_gzipped_and_round_trips(accept_encoding): @pytest.mark.asyncio -async def test_body_above_off_loop_threshold_round_trips(): +async def test_multi_megabyte_body_round_trips(): messages = await _send_messages("/huge", "gzip") assert _headers(messages)["content-encoding"] == "gzip" - assert gzip.decompress(_body(messages)) == b"z" * (OFF_LOOP_SIZE_BYTES * 2) + assert gzip.decompress(_body(messages)) == b"z" * HUGE_RESPONSE_BYTES @pytest.mark.parametrize( - "path, compressed_on_loop_thread", - [("/under-off-loop", True), ("/at-off-loop", False), ("/base64-embeddings-sized", False), ("/huge", False)], + "path", + ["/minimum-eligible", "/four-kib", "/base64-embeddings-sized", "/huge"], ) @pytest.mark.asyncio -async def test_only_bodies_under_the_off_loop_threshold_compress_on_the_event_loop( - monkeypatch: pytest.MonkeyPatch, path: str, compressed_on_loop_thread: bool -) -> None: +async def test_all_eligible_bodies_compress_on_a_worker_thread(monkeypatch: pytest.MonkeyPatch, path: str) -> None: loop_thread: Final = threading.get_ident() real_compress: Final = gzip.compress @@ -159,7 +170,7 @@ async def test_only_bodies_under_the_off_loop_threshold_compress_on_the_event_lo messages = await _send_messages(path, "gzip") assert _headers(messages)["content-encoding"] == "gzip" - assert gzip.decompress(_body(messages)) == (b"loop" if compressed_on_loop_thread else b"worker") + assert gzip.decompress(_body(messages)) == b"worker" @pytest.mark.parametrize( @@ -223,7 +234,9 @@ async def test_start_message_without_headers_key_is_still_gzipped(): await send({"type": "http.response.start", "status": 200}) await send({"type": "http.response.body", "body": body}) - messages = await _send_messages("/", "gzip", GZipBufferedResponseMiddleware(headerless_app)) + messages = await _send_messages( + "/user/daily/activity/aggregated", "gzip", GZipBufferedResponseMiddleware(headerless_app) + ) assert _headers(messages)["content-encoding"] == "gzip" assert gzip.decompress(_body(messages)) == body @@ -234,28 +247,89 @@ async def test_start_without_a_body_message_is_still_forwarded(): async def start_only_app(scope: Scope, receive: Receive, send: Send) -> None: await send({"type": "http.response.start", "status": 204, "headers": [(b"x-done", b"1")]}) - messages = await _send_messages("/", "gzip", GZipBufferedResponseMiddleware(start_only_app)) + messages = await _send_messages( + "/user/daily/activity/aggregated", "gzip", GZipBufferedResponseMiddleware(start_only_app) + ) assert messages == ({"type": "http.response.start", "status": 204, "headers": [(b"x-done", b"1")]},) -def test_proxy_app_gzips_large_responses_for_clients_that_accept_it(): - from starlette.testclient import TestClient +@pytest.mark.parametrize("path", sorted(DASHBOARD_GZIP_PATHS)) +@pytest.mark.parametrize("root_path", ["", "/gateway"]) +@pytest.mark.asyncio +async def test_dashboard_reads_are_compressed_under_root_paths(path: str, root_path: str) -> None: + async def dashboard(scope: Scope, receive: Receive, send: Send) -> None: + await JSONResponse(LARGE_PAYLOAD)(scope, receive, send) + messages: Final = await _send_messages( + root_path + path, "gzip", GZipBufferedResponseMiddleware(dashboard), root_path=root_path + ) + + assert _headers(messages)["content-encoding"] == "gzip" + assert json.loads(gzip.decompress(_body(messages))) == LARGE_PAYLOAD + + +@pytest.mark.parametrize("suffix", ["", "/keys", "/search", "/model_top_keys"]) +@pytest.mark.asyncio +async def test_customer_analytics_aliases_have_equivalent_compressed_responses(suffix: str) -> None: + async def analytics(scope: Scope, receive: Receive, send: Send) -> None: + await JSONResponse(LARGE_PAYLOAD)(scope, receive, send) + + app: Final = GZipBufferedResponseMiddleware(analytics) + customer: Final = await _send_messages(f"/customer/daily/activity/aggregated{suffix}", "gzip", app) + end_user: Final = await _send_messages(f"/end_user/daily/activity/aggregated{suffix}", "gzip", app) + + assert _headers(customer)["content-encoding"] == _headers(end_user)["content-encoding"] == "gzip" + assert _headers(customer)["vary"] == _headers(end_user)["vary"] == "Accept-Encoding" + assert json.loads(gzip.decompress(_body(customer))) == json.loads(gzip.decompress(_body(end_user))) == LARGE_PAYLOAD + + +@pytest.mark.parametrize( + "path,method", + [ + ("/v1/embeddings", "POST"), + ("/embeddings", "POST"), + ("/v1/chat/completions", "POST"), + ("/v1/responses", "POST"), + ("/v1/messages", "POST"), + ("/openai/deployments/embedding-model/embeddings", "POST"), + ("/custom-passthrough", "GET"), + ("/openapi.json", "GET"), + ("/health/liveliness", "GET"), + ("/end_user/daily/activity/export", "GET"), + ("/user/daily/activity/aggregated/unknown", "GET"), + ("/prefix/user/daily/activity/aggregated", "GET"), + ("/user/daily/activity/aggregated", "POST"), + ("/user/daily/activity/aggregated", "HEAD"), + ], +) +@pytest.mark.asyncio +async def test_non_dashboard_requests_pass_through_without_changing_headers(path: str, method: str) -> None: + body: Final = b"e" * BASE64_EMBEDDINGS_RESPONSE_BYTES + + async def inference(scope: Scope, receive: Receive, send: Send) -> None: + await Response(body, media_type="application/json", headers={"vary": "Origin"})(scope, receive, send) + + messages: Final = await _send_messages(path, "gzip", GZipBufferedResponseMiddleware(inference), method=method) + + assert _headers(messages) == { + "content-type": "application/json", + "content-length": str(len(body)), + "vary": "Origin", + } + assert _body(messages) == body + + +def test_proxy_app_does_not_compress_openapi() -> None: from litellm.proxy.proxy_server import app - client = TestClient(app) - compressed = client.get("/openapi.json", headers={"accept-encoding": "gzip"}) - identity = client.get("/openapi.json", headers={"accept-encoding": "identity"}) + client: Final = TestClient(app) + gzip_accepted: Final = client.get("/openapi.json", headers={"accept-encoding": "gzip"}) + identity: Final = client.get("/openapi.json", headers={"accept-encoding": "identity"}) - assert compressed.status_code == identity.status_code == 200, ( - f"gzip={compressed.status_code} {compressed.text[:MINIMUM_SIZE_BYTES]}, " + assert gzip_accepted.status_code == identity.status_code == 200, ( + f"gzip={gzip_accepted.status_code} {gzip_accepted.text[:MINIMUM_SIZE_BYTES]}, " f"identity={identity.status_code} {identity.text[:MINIMUM_SIZE_BYTES]}" ) - assert "content-encoding" in compressed.headers, ( - f"status={compressed.status_code}, headers={dict(compressed.headers)}, " - f"body_bytes={len(compressed.content)}, body={compressed.text[:MINIMUM_SIZE_BYTES]}" - ) - assert compressed.headers["content-encoding"] == "gzip" - assert int(compressed.headers["content-length"]) < int(identity.headers["content-length"]) - assert compressed.json() == identity.json() + assert "content-encoding" not in gzip_accepted.headers + assert gzip_accepted.content == identity.content From 464bb5a1ca382702feee2355cb2aafdb6865a6ff Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:27:12 -0700 Subject: [PATCH 12/23] chore(pricing): add github_copilot deprecation dates for claude-opus-5 and gpt-5.6-terra (#46792) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 4 ++++ model_prices_and_context_window.json | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 46fd8b21318..d0606578658 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -32498,11 +32498,13 @@ "supports_vision": true }, "github_copilot/claude-opus-5": { + "deprecation_date": "2026-11-16", "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32828,11 +32830,13 @@ "supports_vision": true }, "github_copilot/gpt-5.6-terra": { + "deprecation_date": "2026-11-16", "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", "supported_endpoints": [ "/v1/responses" ], diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 46fd8b21318..d0606578658 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -32498,11 +32498,13 @@ "supports_vision": true }, "github_copilot/claude-opus-5": { + "deprecation_date": "2026-11-16", "litellm_provider": "github_copilot", "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", + "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", "supported_endpoints": [ "/v1/chat/completions", "/v1/messages" @@ -32828,11 +32830,13 @@ "supports_vision": true }, "github_copilot/gpt-5.6-terra": { + "deprecation_date": "2026-11-16", "litellm_provider": "github_copilot", "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "responses", + "source": "https://github.blog/changelog/2026-10-15-upcoming-deprecation-of-selected-github-copilot-models-in-mid-november", "supported_endpoints": [ "/v1/responses" ], From c6a4437e980e1a498066399f69666a14aa618d05 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:30:12 +0000 Subject: [PATCH 13/23] ci(cost map): classify token_count_framing and token_count_fallback_model in the schema generator (#46784) * ci(cost map): classify token_count_framing and token_count_fallback_model in the schema generator * ci(cost map): bound token_count_framing to non-negative and token_count_fallback_model to non-empty * test(cost map): reject a fractional token_count_framing in the generated schema * test(cost map): pin that a token_count_framing of 0 passes the generated schema --- ci_cd/generate_model_prices_schema.py | 9 +++++++ tests/unit/test_model_prices_schema.py | 33 ++++++++++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/ci_cd/generate_model_prices_schema.py b/ci_cd/generate_model_prices_schema.py index 8eec07dadda..e429aae36b9 100644 --- a/ci_cd/generate_model_prices_schema.py +++ b/ci_cd/generate_model_prices_schema.py @@ -189,6 +189,10 @@ INTEGER_KEYS: dict[str, JsonSchema] = { **NONNEG_INTEGER, "description": "Smallest prefix the provider will actually cache; absent means the provider default applies.", }, + "token_count_framing": { + **NONNEG_INTEGER, + "description": "Tokens added once per request to a local count estimate, for the provider's chat framing.", + }, "tpm": {**NONNEG_INTEGER, "description": "Provider default tokens-per-minute limit."}, "rpm": {**NONNEG_INTEGER, "description": "Provider default requests-per-minute limit."}, } @@ -283,6 +287,11 @@ def string_key_schemas(modes: tuple) -> dict[str, JsonSchema]: "enum": ["none", "minimal", "low", "medium", "high", "xhigh"], }, "comment": STRING, + "token_count_fallback_model": { + "type": "string", + "minLength": 1, + "description": "Model that counts tokens for this one when the provider cannot count it.", + }, "audio_transcription_config": STRING, "vertex_ai_audio_api": { "type": "string", diff --git a/tests/unit/test_model_prices_schema.py b/tests/unit/test_model_prices_schema.py index 314a401e88e..b12d3b34c8b 100644 --- a/tests/unit/test_model_prices_schema.py +++ b/tests/unit/test_model_prices_schema.py @@ -299,6 +299,39 @@ def test_generated_off_peak_schema_rejects_malformed_blocks(block: dict): assert not validator.is_valid({"some-model": {**OFF_PEAK_ENTRY, "off_peak_pricing": block}}) +TOKEN_COUNT_ENTRY: Final = MappingProxyType( + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + "token_count_framing": 78, + "token_count_fallback_model": "us.anthropic.claude-opus-5-5", + } +) + + +@pytest.mark.parametrize("framing", [0, 78]) +def test_generator_classifies_token_count_settings(framing: int): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, "token_count_framing": framing}}) + + +@pytest.mark.parametrize( + "setting", + [ + {"token_count_framing": -1}, + {"token_count_framing": "78"}, + {"token_count_framing": 1.5}, + {"token_count_fallback_model": ""}, + {"token_count_fallback_model": 5}, + ], +) +def test_generated_token_count_schema_rejects_malformed_settings(setting: dict): + generator = load_generator() + schema = json.loads(generator.render(generator.build_schema({"some-model": dict(TOKEN_COUNT_ENTRY)}))) + assert not build_validator(schema).is_valid({"some-model": {**TOKEN_COUNT_ENTRY, **setting}}) + + def find_duplicate_keys(path: Path) -> list[str]: duplicates: list[str] = [] From b06f4345e4b3d46019c369f07e62ad894126a1a7 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Thu, 15 Oct 2026 17:40:24 -0700 Subject: [PATCH 14/23] fix(terraform): accept 2xx from the proxy in mcp server, model, key and org requests (#46550) The proxy answers POST /v1/mcp/server with 201 and DELETE /v1/mcp/server/{id} with 202, but handleMCPAPIResponse, handleAPIResponse, Client.sendRequest and resourceLiteLLMMCPServerDelete treated every status other than 200 as a failure. terraform apply created the server on the proxy, then reported an error and never wrote the resource to state, so the next apply created a duplicate, and destroy failed on the 202. Accept any 2xx in the four status checks, treat an empty or null 2xx body in handleAPIResponse as an empty ModelResponse, and cover each path with httptest tables that fail on the 200-only comparison. Re-lands the diff from #40723 on current main. Co-authored-by: matthew-hull-bright Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- terraform/provider/CHANGELOG.md | 1 + terraform/provider/litellm/client.go | 2 +- terraform/provider/litellm/client_test.go | 57 ++++++++ .../litellm/resource_mcp_server_crud.go | 2 +- .../litellm/resource_mcp_server_crud_test.go | 46 +++++++ .../provider/litellm/resource_model_crud.go | 3 + .../provider/litellm/resource_model_test.go | 36 +++++ terraform/provider/litellm/utils.go | 8 +- terraform/provider/litellm/utils_test.go | 128 ++++++++++++++++++ 9 files changed, 279 insertions(+), 4 deletions(-) create mode 100644 terraform/provider/litellm/utils_test.go diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 4c45008623c..4282485714a 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -63,6 +63,7 @@ longer signal it. - **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright - **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead - **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs +- **mcp_server, model, key, organization_member**: Requests now accept any 2xx response instead of requiring exactly HTTP 200; the proxy legitimately returns 201 from `POST /v1/mcp/server` (and other create endpoints), so `litellm_mcp_server` create previously succeeded on the proxy but failed in the provider, leaving the resource out of state and risking a duplicate server on retry - **unified_access_group**: create now accepts any 2xx response instead of requiring exactly HTTP 200; `POST /v1/unified_access_group` legitimately returns 201, so creation previously succeeded on the proxy but failed in the provider, leaving the group out of state and forcing a `terraform import` to recover on the next apply's 409. Same fix and shape as the one already applied to `mcp_server`/`model`/`key`/`organization_member`; `handleResponse` (shared by several other resources) was the one status-check helper that fix didn't reach. The legacy `litellm_access_group` resource calls the unrelated `/access_group/new` endpoint, which already returns 200, so it was never affected - **mcp_server**: A plan right after a clean apply of a minimal litellm_mcp_server no longer shows changes. alias is now computed so the proxy's default (the server name) is kept, spec_version keeps the configured value when the proxy returns none, and an mcp_info block the proxy fills with fields the provider does not model is no longer written to state diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index 33ba4c41271..28a89f1eafd 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -455,7 +455,7 @@ func (c *Client) sendRequest(method, path string, body interface{}) (map[string] log.Printf("Response status: %d", resp.StatusCode) log.Printf("Response body: %s", c.redactSensitiveData(string(bodyBytes))) - if resp.StatusCode != http.StatusOK { + if resp.StatusCode < 200 || resp.StatusCode >= 300 { return nil, &apiError{StatusCode: resp.StatusCode, Body: string(bodyBytes)} } diff --git a/terraform/provider/litellm/client_test.go b/terraform/provider/litellm/client_test.go index 56f76565616..77c7bf817d8 100644 --- a/terraform/provider/litellm/client_test.go +++ b/terraform/provider/litellm/client_test.go @@ -1,10 +1,67 @@ package litellm import ( + "errors" + "net/http" + "net/http/httptest" "strings" "testing" ) +func TestSendRequestAcceptsFullSuccessRange(t *testing.T) { + tests := []struct { + name string + statusCode int + wantErr bool + wantValue string + }{ + {name: "200 OK", statusCode: http.StatusOK, wantErr: false, wantValue: "ok"}, + {name: "201 Created", statusCode: http.StatusCreated, wantErr: false, wantValue: "created"}, + {name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false, wantValue: "accepted"}, + {name: "400 Bad Request", statusCode: http.StatusBadRequest, wantErr: true}, + {name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true}, + {name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(tt.statusCode) + w.Write([]byte(`{"value":"` + tt.wantValue + `"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + result, err := client.sendRequest("POST", "/whatever", map[string]string{"foo": "bar"}) + + if tt.wantErr { + if err == nil { + t.Fatalf("sendRequest returned no error for status %d", tt.statusCode) + } + var apiErr *apiError + if !errors.As(err, &apiErr) { + t.Fatalf("error is not *apiError: %v", err) + } + if apiErr.StatusCode != tt.statusCode { + t.Errorf("apiErr.StatusCode = %d, want %d", apiErr.StatusCode, tt.statusCode) + } + if tt.statusCode == http.StatusNotFound && !isNotFound(err) { + t.Errorf("isNotFound(err) = false for 404, want true") + } + return + } + + if err != nil { + t.Fatalf("sendRequest returned unexpected error: %v", err) + } + if result["value"] != tt.wantValue { + t.Errorf("result[value] = %v, want %q", result["value"], tt.wantValue) + } + }) + } +} + func TestRedactSensitiveDataNestedCredentialValues(t *testing.T) { c := NewClient("http://localhost:4000", "sk-test", false) diff --git a/terraform/provider/litellm/resource_mcp_server_crud.go b/terraform/provider/litellm/resource_mcp_server_crud.go index 9bd91d10218..20605ccb8ef 100644 --- a/terraform/provider/litellm/resource_mcp_server_crud.go +++ b/terraform/provider/litellm/resource_mcp_server_crud.go @@ -298,7 +298,7 @@ func resourceLiteLLMMCPServerDelete(d *schema.ResourceData, m interface{}) error defer resp.Body.Close() // For delete operations, we expect a simple string response - if resp.StatusCode != 200 { + if resp.StatusCode < 200 || resp.StatusCode >= 300 { return fmt.Errorf("failed to delete MCP server: unexpected status code %d", resp.StatusCode) } diff --git a/terraform/provider/litellm/resource_mcp_server_crud_test.go b/terraform/provider/litellm/resource_mcp_server_crud_test.go index 2fb2daab38e..36017e140ce 100644 --- a/terraform/provider/litellm/resource_mcp_server_crud_test.go +++ b/terraform/provider/litellm/resource_mcp_server_crud_test.go @@ -13,6 +13,52 @@ import ( "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" ) +func TestMCPServerDeleteAcceptsFullSuccessRange(t *testing.T) { + tests := []struct { + name string + statusCode int + wantErr bool + }{ + {name: "200 OK", statusCode: http.StatusOK, wantErr: false}, + {name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false}, + {name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true}, + {name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(tt.statusCode) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMMCPServer().Schema, map[string]interface{}{ + "server_name": "gh", + "transport": "stdio", + "command": "npx", + }) + d.SetId("srv-1") + + client := NewClient(srv.URL, "test-key", true) + err := resourceLiteLLMMCPServerDelete(d, client) + + if tt.wantErr { + if err == nil { + t.Fatalf("resourceLiteLLMMCPServerDelete returned no error for status %d", tt.statusCode) + } + return + } + + if err != nil { + t.Fatalf("resourceLiteLLMMCPServerDelete returned unexpected error: %v", err) + } + if d.Id() != "" { + t.Errorf("resource ID not cleared after successful delete, got %q", d.Id()) + } + }) + } +} + func TestMCPServerReadDoesNotPersistServerEnv(t *testing.T) { d := schema.TestResourceDataRaw(t, resourceLiteLLMMCPServer().Schema, map[string]interface{}{ "server_name": "gh", diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index af9fb5eb570..3309dbe26dd 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -337,6 +337,9 @@ func resourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { } return fmt.Errorf("failed to read model: %w", err) } + if modelResp == nil { + return fmt.Errorf("failed to read model: empty response body") + } // Update the state with values from the response or fall back to the data passed in during creation d.Set("model_name", GetStringValue(modelResp.ModelName, d.Get("model_name").(string))) diff --git a/terraform/provider/litellm/resource_model_test.go b/terraform/provider/litellm/resource_model_test.go index bcb11462e09..22c226adfd3 100644 --- a/terraform/provider/litellm/resource_model_test.go +++ b/terraform/provider/litellm/resource_model_test.go @@ -153,6 +153,42 @@ func TestResourceLiteLLMModelReadDisplayName(t *testing.T) { } } +func TestResourceLiteLLMModelReadEmptyBodyKeepsState(t *testing.T) { + cases := []struct { + name string + serverBody string + }{ + {name: "empty body"}, + {name: "null body", serverBody: "null"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/model/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + _, _ = w.Write([]byte(tc.serverBody)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMModel().Schema, map[string]interface{}{ + "model_name": "sonnet-4-5-anthropic", + "custom_llm_provider": "anthropic", + "base_model": "claude-sonnet-4-5", + "display_name": "Saved Name", + }) + d.SetId("model-123") + + if err := resourceLiteLLMModelRead(d, NewClient(srv.URL, "test-key", true)); err == nil { + t.Fatal("expected read error for empty response body") + } + if got := d.Get("display_name").(string); got != "Saved Name" { + t.Errorf("expected display_name to remain %q, got %q", "Saved Name", got) + } + }) + } +} + func updateResourceData(t *testing.T, oldDisplayName, newDisplayName string) *schema.ResourceData { t.Helper() res := resourceLiteLLMModel() diff --git a/terraform/provider/litellm/utils.go b/terraform/provider/litellm/utils.go index 4b3f13d36c4..9cd5f668dac 100644 --- a/terraform/provider/litellm/utils.go +++ b/terraform/provider/litellm/utils.go @@ -46,7 +46,7 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client) return nil, fmt.Errorf("failed to read response body: %v", err) } - if resp.StatusCode != http.StatusOK { + if resp.StatusCode < 200 || resp.StatusCode >= 300 { var errResp ErrorResponse if err := json.Unmarshal(bodyBytes, &errResp); err == nil { if isModelNotFoundError(errResp) { @@ -65,6 +65,10 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client) bodyBytes = envelope.Data[0] } + if len(bodyBytes) == 0 || string(bodyBytes) == "null" { + return nil, nil + } + var modelResp ModelResponse if err := json.Unmarshal(bodyBytes, &modelResp); err != nil { return nil, fmt.Errorf("failed to parse response: %v", err) @@ -143,7 +147,7 @@ func handleMCPAPIResponse(resp *http.Response, result interface{}, client *Clien return fmt.Errorf("failed to read response body: %v", err) } - if resp.StatusCode != http.StatusOK { + if resp.StatusCode < 200 || resp.StatusCode >= 300 { var errResp ErrorResponse if err := json.Unmarshal(bodyBytes, &errResp); err == nil { if isMCPServerNotFoundError(errResp) { diff --git a/terraform/provider/litellm/utils_test.go b/terraform/provider/litellm/utils_test.go new file mode 100644 index 00000000000..c8b4dfd8059 --- /dev/null +++ b/terraform/provider/litellm/utils_test.go @@ -0,0 +1,128 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" +) + +func TestHandleAPIResponseAcceptsFullSuccessRange(t *testing.T) { + tests := []struct { + name string + statusCode int + wantErr bool + }{ + {name: "200 OK", statusCode: http.StatusOK, wantErr: false}, + {name: "201 Created", statusCode: http.StatusCreated, wantErr: false}, + {name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false}, + {name: "400 Bad Request", statusCode: http.StatusBadRequest, wantErr: true}, + {name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true}, + {name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rec := httptest.NewRecorder() + rec.WriteHeader(tt.statusCode) + rec.WriteString(`{"model_name":"gpt-4o"}`) + resp := rec.Result() + + client := NewClient("http://localhost:4000", "test-key", true) + got, err := handleAPIResponse(resp, map[string]interface{}{"model_name": "gpt-4o"}, client) + + if tt.wantErr { + if err == nil { + t.Fatalf("handleAPIResponse returned no error for status %d", tt.statusCode) + } + if !strings.Contains(err.Error(), strconv.Itoa(tt.statusCode)) { + t.Errorf("error %q does not mention status %d", err.Error(), tt.statusCode) + } + return + } + + if err != nil { + t.Fatalf("handleAPIResponse returned unexpected error: %v", err) + } + if got.ModelName != "gpt-4o" { + t.Errorf("got.ModelName = %q, want gpt-4o", got.ModelName) + } + }) + } +} + +func TestHandleAPIResponseAcceptsEmptyBodyOn2xx(t *testing.T) { + rec := httptest.NewRecorder() + rec.WriteHeader(http.StatusNoContent) + httpResp := rec.Result() + + client := NewClient("http://localhost:4000", "test-key", true) + modelResp, err := handleAPIResponse(httpResp, map[string]interface{}{"model_name": "gpt-4o"}, client) + + if err != nil { + t.Fatalf("handleAPIResponse returned unexpected error for empty-body 204: %v", err) + } + if modelResp != nil { + t.Fatalf("handleAPIResponse returned non-nil ModelResponse for empty-body 204: %#v", modelResp) + } +} + +func TestHandleMCPAPIResponseAcceptsFullSuccessRange(t *testing.T) { + tests := []struct { + name string + statusCode int + wantErr bool + }{ + {name: "200 OK", statusCode: http.StatusOK, wantErr: false}, + {name: "201 Created", statusCode: http.StatusCreated, wantErr: false}, + {name: "202 Accepted", statusCode: http.StatusAccepted, wantErr: false}, + {name: "400 Bad Request", statusCode: http.StatusBadRequest, wantErr: true}, + {name: "404 Not Found", statusCode: http.StatusNotFound, wantErr: true}, + {name: "500 Internal Server Error", statusCode: http.StatusInternalServerError, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + rec := httptest.NewRecorder() + rec.WriteHeader(tt.statusCode) + rec.WriteString(`{"server_id":"srv-1","server_name":"gh"}`) + resp := rec.Result() + + client := NewClient("http://localhost:4000", "test-key", true) + var mcpResp MCPServerResponse + err := handleMCPAPIResponse(resp, &mcpResp, client) + + if tt.wantErr { + if err == nil { + t.Fatalf("handleMCPAPIResponse returned no error for status %d", tt.statusCode) + } + if !strings.Contains(err.Error(), strconv.Itoa(tt.statusCode)) { + t.Errorf("error %q does not mention status %d", err.Error(), tt.statusCode) + } + return + } + + if err != nil { + t.Fatalf("handleMCPAPIResponse returned unexpected error: %v", err) + } + if mcpResp.ServerID != "srv-1" { + t.Errorf("mcpResp.ServerID = %q, want srv-1", mcpResp.ServerID) + } + }) + } +} + +func TestHandleMCPAPIResponseRejectsEmptyBodyOn2xx(t *testing.T) { + rec := httptest.NewRecorder() + rec.WriteHeader(http.StatusNoContent) + resp := rec.Result() + + client := NewClient("http://localhost:4000", "test-key", true) + var mcpResp MCPServerResponse + err := handleMCPAPIResponse(resp, &mcpResp, client) + + if err == nil { + t.Fatal("handleMCPAPIResponse returned no error for empty-body 204; every MCP caller (create/read/update) writes the parsed result straight into Terraform state with no fallback, so a silently-accepted empty body would blank out or empty-ID the resource") + } +} From 38f3f8a1a95b4c58aa32c3560628e62b9fb4b36a Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 15 Oct 2026 17:53:52 -0700 Subject: [PATCH 15/23] fix(proxy): cooldown errors carry 429, Sentry starts at boot without request bodies, tracing failures log at ERROR, exporter counters exported (#46344) * fix(router): carry a 429 status on cooldown and no-deployment errors RouterRateLimitError and RouterRateLimitErrorBasic had no status_code, so the proxy answered an all-deployments-cooled-down request as a generic 500 class error and logged the bug report notice for every cooldown. Both now carry status_code 429, which the proxy maps to a rate_limit_error and which the retry policy already treats as retryable * fix(proxy): start Sentry at config load and never attach request bodies Sentry only came up inside set_callbacks on the first LLM call, so boot and background errors were never reported, and general_settings.alerting: [sentry] always raised Missing SENTRY_DSN because the alerting path read a dead litellm.utils pin instead of the live SDK. The proxy now initializes Sentry once while loading the config whenever sentry is a configured callback or alerting target, the alerting and failure handlers read the live instance, and max_request_body_size is pinned to never so prompts and completions stay out of Sentry events * fix(tracing): log Lens failures at ERROR and export row counters to Prometheus A tracing store that fails to start, or a Lens service that rejects a batch, was only logged at WARNING, so operators running at the default ERROR level never saw tracing silently go dark. Both now log at ERROR and name the lens store. The exporter's rows_written and rows_dropped tallies were never exported, so the two counts are now Prometheus counters litellm_lens_export_rows_written_total and litellm_lens_export_rows_dropped_total, registered lazily on the default registry like the other proxy-level counters and charted on the all-metrics Grafana dashboard * test(proxy): drop the litellm import the alerting tests no longer use * fix(proxy): start Sentry once a DSN is set, from any config source Boot-time init used to lock Sentry in even with no SENTRY_DSN, so a DSN loaded later from DB environment_variables was never picked up. initialize_sentry now skips when no DSN is set and re-inits when the DSN changes. alerting_handler calls it too, so alerting: [sentry] set through the DB works. Drops the litellm_settings.callbacks check since that config fails to load with sentry in it * fix(sentry): keep stack frame locals out of events unless PII is opted in Frame locals carried prompts (messages, request bodies, provider options) into every error event. include_local_variables now follows SENTRY_SEND_DEFAULT_PII * test(router): drop the test that only pinned the 429 class attribute * fix(proxy): start Sentry at boot when it is only an input callback * fix(proxy): keep Sentry init and alerting within the strict ruff and basedpyright ceilings --------- Co-authored-by: yuneng-jiang --- .../grafana_dashboard.json | 81 ++++++++++++++++ .../dashboard_all_metrics/readme.md | 4 +- litellm/litellm_core_utils/litellm_logging.py | 42 ++++++--- .../litellm_core_utils/sentry_scrubbing.py | 6 +- litellm/proxy/proxy_server.py | 14 +++ litellm/proxy/tracing_runtime.py | 6 +- litellm/proxy/utils.py | 19 ++-- litellm/tracing/exporter.py | 93 +++++++++++++++---- litellm/types/router.py | 4 + ...test_prometheus_metric_name_consistency.py | 6 ++ .../litellm_core_utils/test_bug_report.py | 3 + .../test_litellm_logging.py | 56 +++++++++++ .../test_sentry_scrubbing.py | 42 ++++++--- .../proxy/proxy_server/test_proxy_config.py | 62 +++++++++++++ .../proxy/test_common_request_processing.py | 16 ++++ .../utils/proxy_logging/test_alerting.py | 50 ++++++++-- tests/unit/tracing/test_exporter.py | 50 ++++++++++ 17 files changed, 481 insertions(+), 73 deletions(-) diff --git a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json index 996e440d77d..cedbbd0f07b 100644 --- a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json +++ b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/grafana_dashboard.json @@ -6563,6 +6563,87 @@ ], "title": "litellm_project_model_rate_limit_used_metric", "type": "timeseries" + }, + { + "collapsed": false, + "gridPos": { + "h": 1, + "w": 24, + "x": 0, + "y": 447 + }, + "id": 115, + "panels": [], + "title": "Agent tracing export (needs LITELLM_LENS_URL)", + "type": "row" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "description": "Request records the Lens exporter delivered to the Lens service against the records it dropped before the service accepted them", + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "drawStyle": "line", + "fillOpacity": 10, + "lineWidth": 1, + "showPoints": "never", + "spanNulls": false + }, + "unit": "short" + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 24, + "x": 0, + "y": 448 + }, + "id": 116, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "multi", + "sort": "desc" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum(rate(litellm_lens_export_rows_written_total[$__rate_interval]))", + "legendFormat": "written", + "range": true, + "refId": "A" + }, + { + "datasource": { + "type": "prometheus", + "uid": "${DS_PROMETHEUS}" + }, + "editorMode": "code", + "expr": "sum(rate(litellm_lens_export_rows_dropped_total[$__rate_interval]))", + "legendFormat": "dropped", + "range": true, + "refId": "B" + } + ], + "title": "litellm_lens_export_rows_written and litellm_lens_export_rows_dropped rate", + "type": "timeseries" } ], "preload": false, diff --git a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md index 5d3d2c159c1..dcac5db6dc0 100644 --- a/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md +++ b/cookbook/litellm_proxy_server/grafana_dashboard/dashboard_all_metrics/readme.md @@ -1,10 +1,10 @@ # LiteLLM All Prometheus Metrics dashboard -Every `litellm_*` metric family the proxy can expose on `/metrics` (141 families across 97 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, and the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes). Panel titles are the metric names so you can grep the JSON for the metric you care about +Every `litellm_*` metric family the proxy can expose on `/metrics` (143 families across 98 panels), grouped into rows: proxy traffic, latency, spend and tokens, cache, LLM API deployments, key and team rate limits, budgets, guardrails, MCP, managed files and batches, users and teams, the Redis circuit breaker, the spend log cleanup job, the `prometheus_system` service callback metrics (per-service latency, request and failure rates, spend update queue sizes), project model rate limits, and the agent tracing export to Lens. Panel titles are the metric names so you can grep the JSON for the metric you care about Import `grafana_dashboard.json` from **Dashboards > New > Import** and pick your Prometheus data source when prompted (the `DS_PROMETHEUS` variable). Counters are plotted as `rate()` over `$__rate_interval`, histograms as p50 / p95 / p99, gauges as the raw value grouped by the most useful label. Every query names the metric exactly as the proxy emits it (counters carry the `_total` suffix the Prometheus client adds), and `tests/unit/integrations/test_prometheus_metric_name_consistency.py` fails if a metric is renamed without updating this dashboard -The first eleven rows need only `callbacks: ["prometheus"]`. The last three rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected +The first eleven rows and the project model rate limits row need only `callbacks: ["prometheus"]`. The other rows and the `litellm_admission_*` panels are emitted by other subsystems and stay empty until those are on: the service callback row needs `service_callback: ["prometheus_system"]` in `litellm_settings`, the circuit breaker row needs a Redis cache, the cleanup row needs spend log retention, the agent tracing export row needs `LITELLM_LENS_URL`, and admission control needs its middleware enabled. Within the base rows, many panels only fill in once the matching feature is in use: budgets need keys, teams, users or orgs with `max_budget` set, cache panels need caching on, guardrail and MCP panels need those features configured, deployment health needs the router with more than one deployment or a failure to record, and `litellm_in_flight_requests` needs traffic at scrape time. An empty panel for a feature you do not use is expected ## Pre-requisites diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0cda24c853b..7038cc531f5 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4507,6 +4507,32 @@ def get_masked_values( _get_masked_values = get_masked_values +@lru_cache(maxsize=1) +def start_sentry(dsn: str) -> None: + try: + import sentry_sdk + except ImportError: + print_verbose("Package 'sentry_sdk' is missing. Installing it...") + subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"]) + import sentry_sdk + from litellm.litellm_core_utils.sentry_scrubbing import build_sentry_init_options + + sentry_sdk.init(**build_sentry_init_options({**os.environ, "SENTRY_DSN": dsn})) + + +def initialize_sentry() -> None: + global sentry_sdk_instance, capture_exception, add_breadcrumb + dsn: Final = os.environ.get("SENTRY_DSN") + if not dsn: + return + start_sentry(dsn) + import sentry_sdk + + sentry_sdk_instance = sentry_sdk # rebind-ok: process-wide Sentry SDK handle read by every logging path + capture_exception = sentry_sdk.capture_exception # rebind-ok: process-wide Sentry SDK handle + add_breadcrumb = sentry_sdk.add_breadcrumb # rebind-ok: process-wide Sentry SDK handle + + def set_callbacks( callback_list: Iterable[str | Callable[..., object] | CustomLogger], function_id: str | None = None, @@ -4515,9 +4541,6 @@ def set_callbacks( Globally sets the callback client """ global \ - sentry_sdk_instance, \ - capture_exception, \ - add_breadcrumb, \ slack_app, \ alerts_channel, \ traceloopLogger, \ @@ -4541,18 +4564,7 @@ def set_callbacks( try: for callback in callback_list: if callback == "sentry": - try: - import sentry_sdk - except ImportError: - print_verbose("Package 'sentry_sdk' is missing. Installing it...") - subprocess.check_call([sys.executable, "-m", "pip", "install", "sentry_sdk"]) - import sentry_sdk - from litellm.litellm_core_utils.sentry_scrubbing import build_sentry_init_options - - sentry_sdk_instance = sentry_sdk - sentry_sdk_instance.init(**build_sentry_init_options(os.environ)) - capture_exception = sentry_sdk_instance.capture_exception - add_breadcrumb = sentry_sdk_instance.add_breadcrumb + initialize_sentry() elif callback == "slack": try: from slack_bolt import App diff --git a/litellm/litellm_core_utils/sentry_scrubbing.py b/litellm/litellm_core_utils/sentry_scrubbing.py index 7147f792411..b2c5b1f5712 100644 --- a/litellm/litellm_core_utils/sentry_scrubbing.py +++ b/litellm/litellm_core_utils/sentry_scrubbing.py @@ -3,7 +3,7 @@ from __future__ import annotations import re from collections.abc import Callable, Mapping, Sequence from functools import reduce -from typing import TYPE_CHECKING, Final, TypeAlias, cast +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast from pydantic import JsonValue from sentry_sdk.scrubber import DEFAULT_DENYLIST, DEFAULT_PII_DENYLIST, EventScrubber @@ -65,6 +65,8 @@ class SentryInitOptions(TypedDict): traces_sample_rate: ReadOnly[float] sample_rate: ReadOnly[float] send_default_pii: ReadOnly[bool] + max_request_body_size: ReadOnly[Literal["never"]] + include_local_variables: ReadOnly[bool] event_scrubber: ReadOnly[EventScrubber] before_send: ReadOnly[EventScrubFn] before_send_transaction: ReadOnly[EventScrubFn] @@ -140,6 +142,8 @@ def build_sentry_init_options(env: Mapping[str, str]) -> SentryInitOptions: traces_sample_rate=float(env.get("SENTRY_API_TRACE_RATE") or "1.0"), sample_rate=float(env.get("SENTRY_API_SAMPLE_RATE") or "1.0"), send_default_pii=send_default_pii, + max_request_body_size="never", + include_local_variables=send_default_pii, event_scrubber=EventScrubber( denylist=list(SECRET_FIELD_NAMES), pii_denylist=list(PII_FIELD_NAMES), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 60842648d16..fb6a3a8a88d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.bug_report import ( ) from litellm.litellm_core_utils.litellm_logging import ( _init_custom_logger_compatible_class, + initialize_sentry, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -4962,6 +4963,16 @@ def _is_remote_module_url(value: object) -> bool: return isinstance(value, str) and (value.startswith("s3://") or value.startswith("gcs://")) +def sentry_is_configured(general_settings: Mapping[str, object]) -> bool: + alerting: Final = general_settings.get("alerting") + return ( + "sentry" in litellm.input_callback + or "sentry" in litellm.success_callback + or "sentry" in litellm.failure_callback + or (isinstance(alerting, list) and "sentry" in alerting) + ) + + def _scrub_guardrail_inner(inner: dict[str, JsonValue]) -> None: """Strip remote-URL entries from a guardrail's ``callbacks`` list and ``guardrail`` (v2 module-path) field. Mutates in place.""" @@ -6911,6 +6922,9 @@ class ProxyConfig: if general_settings is None: general_settings = {} + if sentry_is_configured(general_settings): + initialize_sentry() + typed_general_settings: Final = _GENERAL_SETTINGS_VIEW.validate_python(general_settings) ConfigGeneralSettings.model_validate( MappingProxyType( diff --git a/litellm/proxy/tracing_runtime.py b/litellm/proxy/tracing_runtime.py index 8de151421f2..0ac05216a30 100644 --- a/litellm/proxy/tracing_runtime.py +++ b/litellm/proxy/tracing_runtime.py @@ -46,10 +46,8 @@ async def manage_tracing( return try: connection: Final = LensConnection.from_env() - except ValueError: - verbose_proxy_logger.warning( - "Agent tracing unavailable: configure LITELLM_LENS_URL and LITELLM_LENS_SERVICE_TOKEN" - ) + except ValueError as error: + verbose_proxy_logger.error("Agent tracing store lens failed to start, tracing is disabled: %s", error) yield None return async with client_factory(connection) as client: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 699fde28d4d..9788c4635c9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3250,14 +3250,14 @@ class ProxyLogging: alerting_metadata=alerting_metadata, **extra_kwargs, ) - for client in self.alerting: - if client == "sentry": - if litellm.utils.sentry_sdk_instance is not None: - litellm.utils.sentry_sdk_instance.capture_message(formatted_message) - else: - raise Exception("Missing SENTRY_DSN from environment") + if "sentry" in self.alerting: + litellm.litellm_core_utils.litellm_logging.initialize_sentry() + sentry: Final = litellm.litellm_core_utils.litellm_logging.sentry_sdk_instance + if sentry is None: + raise Exception("Missing SENTRY_DSN from environment") + sentry.capture_message(formatted_message) - async def failure_handler(self, original_exception, duration: float, call_type: str, traceback_str=""): + async def failure_handler(self, original_exception: Exception, duration: float, call_type: str, traceback_str=""): """ Log failed db read/writes @@ -3298,8 +3298,9 @@ class ProxyLogging: call_type=call_type, ) - if litellm.utils.capture_exception: - litellm.utils.capture_exception(error=original_exception) + capture: Final = litellm.litellm_core_utils.litellm_logging.capture_exception + if capture is not None: + capture(error=original_exception) async def post_call_failure_hook( self, diff --git a/litellm/tracing/exporter.py b/litellm/tracing/exporter.py index 5025b5eae2a..68916558443 100644 --- a/litellm/tracing/exporter.py +++ b/litellm/tracing/exporter.py @@ -3,9 +3,11 @@ import json from collections import deque from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from contextlib import suppress +from dataclasses import dataclass from enum import Enum +from functools import cache from io import BytesIO -from typing import Final +from typing import TYPE_CHECKING, Final import httpx from pydantic import TypeAdapter, ValidationError @@ -16,6 +18,9 @@ from litellm.integrations.clickhouse.clickhouse_spend_logger import spend_log_ro from litellm.integrations.custom_logger import CustomLogger from litellm.tracing.types import SpendLogPayload +if TYPE_CHECKING: + from prometheus_client import Counter as PrometheusCounter + MAX_EVENT_BYTES: Final = 1024 * 1024 MAX_BUFFER_BYTES: Final = 32 * 1024 * 1024 MAX_BUFFER_EVENTS: Final = 1000 @@ -89,11 +94,49 @@ def encode_record(value: Mapping[str, object]) -> bytes | ExportFailure: return ExportFailure.INVALID +@dataclass(frozen=True, slots=True) +class ExportCounters: + written: "PrometheusCounter | None" + dropped: "PrometheusCounter | None" + + def add_written(self, count: int) -> None: + if self.written is not None: + self.written.inc(count) + + def add_dropped(self, count: int) -> None: + if self.dropped is not None: + self.dropped.inc(count) + + +@cache +def lens_export_counters() -> ExportCounters: + try: + from prometheus_client import Counter + except ImportError: + return ExportCounters(written=None, dropped=None) + return ExportCounters( + written=Counter( + "litellm_lens_export_rows_written_total", + "Request records the Lens exporter delivered to the Lens service", + ), + dropped=Counter( + "litellm_lens_export_rows_dropped_total", + "Request records the Lens exporter dropped before the Lens service accepted them", + ), + ) + + class LensExporter(CustomLogger): - def __init__(self, client: httpx.AsyncClient, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep) -> None: + def __init__( + self, + client: httpx.AsyncClient, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + counters: ExportCounters | None = None, + ) -> None: super().__init__() self.client: Final = client self.sleep: Final = sleep + self.counters: Final = counters if counters is not None else lens_export_counters() self.queue: Final[deque[bytes]] = deque() # mutable-ok: bounded producer-consumer queue self.wake: Final = asyncio.Event() self.closed = False @@ -108,6 +151,14 @@ class LensExporter(CustomLogger): if self.task is None: self.task = asyncio.create_task(self._run()) + def _wrote(self, count: int) -> None: + self.rows_written += count + self.counters.add_written(count) + + def _drop(self, count: int) -> None: + self.rows_dropped += count + self.counters.add_dropped(count) + def enqueue(self, record: bytes) -> bool: if ( self.closed @@ -115,7 +166,7 @@ class LensExporter(CustomLogger): or self.buffered_events >= MAX_BUFFER_EVENTS or self.buffered_bytes + len(record) > MAX_BUFFER_BYTES ): - self.rows_dropped += 1 + self._drop(1) return False self.queue.append(record) self.buffered_events += 1 @@ -138,13 +189,13 @@ class LensExporter(CustomLogger): if raw is None or self.closed: return if self.buffered_events >= MAX_BUFFER_EVENTS or self.buffered_bytes >= MAX_BUFFER_BYTES: - self.rows_dropped += 1 + self._drop(1) return try: checked: Final = _check_size(raw, MAX_EVENT_BYTES) if isinstance(checked, ExportFailure): - self.rows_dropped += 1 - self._warn(checked.value) + self._drop(1) + self._log_failure(checked.value) return payload: Final = _PAYLOAD.validate_python(raw) if str(payload.get("call_type", "")).startswith(("/v1/traces", "/v1/logs")): @@ -152,24 +203,26 @@ class LensExporter(CustomLogger): row: Final = spend_log_row_from_payload(payload, kwargs) record: Final = encode_record(row) if isinstance(record, ExportFailure): - self.rows_dropped += 1 - self._warn(record.value) + self._drop(1) + self._log_failure(record.value) return self.enqueue(record) except ValidationError as error: - self.rows_dropped += 1 + self._drop(1) fields: Final = tuple( str(issue["loc"][0]) if issue["loc"] else "$" for issue in error.errors(include_input=False, include_context=False, include_url=False)[:5] ) - self._warn("invalid request record fields: " + ", ".join(fields)) + self._log_failure("invalid request record fields: " + ", ".join(fields)) except (ValueError, TypeError, OverflowError, RecursionError) as error: - self.rows_dropped += 1 - self._warn(type(error).__name__) + self._drop(1) + self._log_failure(type(error).__name__) - def _warn(self, reason: str) -> None: + def _log_failure(self, reason: str) -> None: if reason != self.last_error: - verbose_proxy_logger.warning("Lens request export failed (%s); model requests continue", reason) + verbose_proxy_logger.error( + "Agent tracing store lens dropped request records (%s); model requests continue", reason + ) self.last_error = reason def _batch(self) -> tuple[bytes, ...]: @@ -196,13 +249,13 @@ class LensExporter(CustomLogger): self.last_error = "" return True if response.status_code not in (429, 502, 503, 504): - self._warn(f"HTTP {response.status_code}") + self._log_failure(f"HTTP {response.status_code}") return False except httpx.HTTPError: pass if attempt < 2: await self.sleep(float(1 << attempt)) - self._warn("retry limit reached") + self._log_failure("retry limit reached") return False async def _run(self) -> None: @@ -217,11 +270,11 @@ class LensExporter(CustomLogger): batch: Final = self._batch() try: if await self._send(batch): - self.rows_written += len(batch) + self._wrote(len(batch)) else: - self.rows_dropped += len(batch) + self._drop(len(batch)) except asyncio.CancelledError: - self.rows_dropped += len(batch) + self._drop(len(batch)) raise finally: self.buffered_events -= len(batch) @@ -237,7 +290,7 @@ class LensExporter(CustomLogger): self.task.cancel() with suppress(asyncio.CancelledError): await self.task - self.rows_dropped += len(self.queue) + self._drop(len(self.queue)) self.queue.clear() self.buffered_bytes = 0 self.buffered_events = 0 diff --git a/litellm/types/router.py b/litellm/types/router.py index 8b53b7b6e12..0ca39e2400d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1084,6 +1084,8 @@ class RouterRateLimitErrorBasic(ValueError): Raise a basic error inside helper functions. """ + status_code: Final = 429 + def __init__( self, model: str, @@ -1099,6 +1101,8 @@ class RouterErrorTypes(str, enum.Enum): class RouterRateLimitError(ValueError): + status_code: Final = 429 + def __init__( self, model: str, diff --git a/tests/unit/integrations/test_prometheus_metric_name_consistency.py b/tests/unit/integrations/test_prometheus_metric_name_consistency.py index 4db34274e82..2c1def4c613 100644 --- a/tests/unit/integrations/test_prometheus_metric_name_consistency.py +++ b/tests/unit/integrations/test_prometheus_metric_name_consistency.py @@ -26,6 +26,7 @@ from litellm.integrations.prometheus_services import PrometheusServicesLogger from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import SpendLogCleanupMetrics from litellm.proxy.middleware.admission_control_middleware import create_prometheus_admission_metrics from litellm.proxy.middleware.in_flight_requests_middleware import InFlightRequestsMiddleware +from litellm.tracing.exporter import lens_export_counters _GRAFANA_DIR: Final = Path(__file__).parents[3] / "cookbook" / "litellm_proxy_server" / "grafana_dashboard" _ALL_METRICS_DASHBOARD: Final = _GRAFANA_DIR / "dashboard_all_metrics" / "grafana_dashboard.json" @@ -62,6 +63,9 @@ def _lazy_owner_collectors() -> tuple[Collector, ...]: in_flight: Final = InFlightRequestsMiddleware._get_gauge() assert in_flight is not None services: Final = PrometheusServicesLogger() + lens: Final = lens_export_counters() + assert lens.written is not None + assert lens.dropped is not None return ( SpendLogCleanupMetrics.rows_deleted, SpendLogCleanupMetrics.batch_duration, @@ -72,6 +76,8 @@ def _lazy_owner_collectors() -> tuple[Collector, ...]: services.redis_breaker_state, services.redis_breaker_transitions, services.redis_breaker_failures, + lens.written, + lens.dropped, ) diff --git a/tests/unit/litellm_core_utils/test_bug_report.py b/tests/unit/litellm_core_utils/test_bug_report.py index 62d7090b960..03d6c800b60 100644 --- a/tests/unit/litellm_core_utils/test_bug_report.py +++ b/tests/unit/litellm_core_utils/test_bug_report.py @@ -26,6 +26,7 @@ from litellm.litellm_core_utils.bug_report import ( strip_bug_report_notice, ) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic def test_build_bug_report_keeps_only_litellm_frames(): @@ -138,6 +139,8 @@ def test_bug_report_can_be_disabled(monkeypatch: pytest.MonkeyPatch): request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), ), BadRequestError(message="bad input", llm_provider="openai", model="gpt-4"), + RouterRateLimitError(model="gpt-4", cooldown_time=5.0, enable_pre_call_checks=False, cooldown_list=[]), + RouterRateLimitErrorBasic(model="gpt-4"), "not an exception", ], ) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 12cc57ada4c..14334690624 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -32,6 +32,7 @@ from litellm.caching.caching import DualCache from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( Logging, @@ -422,6 +423,13 @@ def test_post_call_serializes_dict_with_datetime(logging_obj): assert "2026-05-11" in serialized +def _reset_sentry_globals(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", None) + litellm_logging.start_sentry.cache_clear() + monkeypatch.setattr(litellm_logging, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "add_breadcrumb", None) + + def test_sentry_environment(monkeypatch): import sentry_sdk @@ -430,15 +438,59 @@ def test_sentry_environment(monkeypatch): monkeypatch.setenv("SENTRY_DSN", "https://test@sentry.io/123456") monkeypatch.delenv("SENTRY_ENVIRONMENT", raising=False) + _reset_sentry_globals(monkeypatch) set_callbacks(["sentry"]) assert mock_init.call_args[1]["environment"] == "production" for environment in ("development", "staging"): monkeypatch.setenv("SENTRY_ENVIRONMENT", environment) mock_init.reset_mock() + _reset_sentry_globals(monkeypatch) set_callbacks(["sentry"]) mock_init.assert_called_once() assert mock_init.call_args[1]["environment"] == environment + + +def test_set_callbacks_initializes_sentry_once_per_dsn(monkeypatch): + import sentry_sdk + + mock_init: Final = MagicMock() + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_DSN", "https://first@sentry.example/1") + _reset_sentry_globals(monkeypatch) + + set_callbacks(["sentry"]) + set_callbacks(["sentry"]) + monkeypatch.setenv("SENTRY_DSN", "https://second@sentry.example/2") + set_callbacks(["sentry"]) + + assert [call.kwargs["dsn"] for call in mock_init.call_args_list] == [ + "https://first@sentry.example/1", + "https://second@sentry.example/2", + ] + assert litellm_logging.sentry_sdk_instance is sentry_sdk + assert litellm_logging.capture_exception is sentry_sdk.capture_exception + + +def test_sentry_started_before_its_dsn_is_loaded_picks_the_dsn_up_later(monkeypatch): + import sentry_sdk + + mock_init: Final = MagicMock() + monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.delenv("SENTRY_DSN", raising=False) + _reset_sentry_globals(monkeypatch) + + litellm_logging.initialize_sentry() + assert litellm_logging.sentry_sdk_instance is None + + monkeypatch.setenv("SENTRY_DSN", "https://from-db@sentry.example/3") + set_callbacks(["sentry"]) + + mock_init.assert_called_once() + assert mock_init.call_args.kwargs["dsn"] == "https://from-db@sentry.example/3" + assert litellm_logging.sentry_sdk_instance is sentry_sdk + + def test_use_custom_pricing_for_model(): from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model @@ -3375,7 +3427,9 @@ def test_sentry_event_scrubber_initialization(monkeypatch): mock_init = MagicMock() monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_DSN", "https://key@sentry.example/1") monkeypatch.delenv("SENTRY_SEND_DEFAULT_PII", raising=False) + _reset_sentry_globals(monkeypatch) set_callbacks(["sentry"]) @@ -3392,7 +3446,9 @@ def test_sentry_send_default_pii_opt_in(monkeypatch): mock_init = MagicMock() monkeypatch.setattr(sentry_sdk, "init", mock_init) + monkeypatch.setenv("SENTRY_DSN", "https://key@sentry.example/1") monkeypatch.setenv("SENTRY_SEND_DEFAULT_PII", "true") + _reset_sentry_globals(monkeypatch) set_callbacks(["sentry"]) diff --git a/tests/unit/litellm_core_utils/test_sentry_scrubbing.py b/tests/unit/litellm_core_utils/test_sentry_scrubbing.py index 9aae3999129..c5bf8334b91 100644 --- a/tests/unit/litellm_core_utils/test_sentry_scrubbing.py +++ b/tests/unit/litellm_core_utils/test_sentry_scrubbing.py @@ -30,6 +30,7 @@ MASTER_KEY: Final = "sk-master-key-under-test" DATABASE_URL: Final = "postgresql://litellm:db-password-under-test@db.internal:5432/litellm" PII_ON: Final = {"SENTRY_DSN": "https://key@sentry.example/1", "SENTRY_SEND_DEFAULT_PII": "true"} PII_OFF: Final = {"SENTRY_DSN": "https://key@sentry.example/1"} +PROMPT: Final = "confidential prompt text under test" class RecordingTransport(Transport): @@ -62,8 +63,12 @@ def raise_with_identity_locals() -> None: def raise_with_source_context_named_locals() -> None: - metadata: Final = {"context_line": f"Bearer {VIRTUAL_KEY}", "pre_context": [EMAIL], "post_context": [KEY_HASH]} - stacktrace: Final = {"frames": [{"context_line": MASTER_KEY, "pre_context": [EMAIL]}]} + metadata: Final = { + "context_line": f"Bearer {VIRTUAL_KEY}", + "pre_context": [MASTER_KEY], + "post_context": [VIRTUAL_KEY], + } + stacktrace: Final = {"frames": [{"context_line": MASTER_KEY, "pre_context": [VIRTUAL_KEY]}]} raise RuntimeError(f"rejected with {len(metadata)} metadata fields and {len(stacktrace)} stack fields") @@ -92,15 +97,21 @@ def test_default_event_carries_no_email_hash_or_secret_anywhere() -> None: assert MASTER_KEY not in serialized assert VIRTUAL_KEY not in serialized assert "db-password-under-test" not in serialized - frame_vars: Final = innermost_frame_vars(serialized) - assert frame_vars["raw_headers"] == {"authorization": FILTERED, "x-api-key": FILTERED, "content-type": "'application/json'"} - assert f"token='{FILTERED}'" in frame_vars["valid_token"] - assert f"user_id='{FILTERED}'" in frame_vars["valid_token"] - assert f"user_email='{FILTERED}'" in frame_vars["user_obj"] - assert frame_vars["general_settings"] == {"master_key": FILTERED, "database_url": FILTERED} - assert frame_vars["data"] == {"metadata": {"user_api_key_hash": FILTERED, "user_api_key_user_email": FILTERED}} - assert "key_name='sk-...test'" in frame_vars["valid_token"] - assert "user_role='internal_user'" in frame_vars["user_obj"] + + +def raise_with_prompt_locals() -> None: + messages: Final = [{"role": "user", "content": PROMPT}] + raise RuntimeError(f"provider rejected {len(messages)} messages") + + +@pytest.mark.parametrize(("env", "locals_sent"), ((PII_OFF, False), (PII_ON, True)), ids=("default", "pii-opt-in")) +def test_frame_locals_carrying_prompts_are_sent_only_with_the_pii_opt_in( + env: Mapping[str, str], locals_sent: bool +) -> None: + serialized: Final = capture_serialized_event(env, raise_with_prompt_locals) + frames: Final = json.loads(serialized)["exception"]["values"][0]["stacktrace"]["frames"] + assert (PROMPT in serialized) is locals_sent + assert any("vars" in frame for frame in frames) is locals_sent def test_source_context_lines_are_left_readable() -> None: @@ -115,11 +126,9 @@ def test_source_context_lines_are_left_readable() -> None: def test_source_context_names_outside_stack_frames_are_scrubbed() -> None: - serialized: Final = capture_serialized_event(PII_OFF, raise_with_source_context_named_locals) + serialized: Final = capture_serialized_event(PII_ON, raise_with_source_context_named_locals) assert VIRTUAL_KEY not in serialized assert MASTER_KEY not in serialized - assert EMAIL not in serialized - assert KEY_HASH not in serialized frame_vars: Final = innermost_frame_vars(serialized) assert frame_vars["metadata"] == { "context_line": f"'Bearer {FILTERED}'", @@ -276,3 +285,8 @@ def test_init_options_defaults() -> None: assert options["traces_sample_rate"] == 1.0 assert options["sample_rate"] == 1.0 assert options["environment"] == "production" + + +@pytest.mark.parametrize("env", ({}, {"SENTRY_SEND_DEFAULT_PII": "true"}), ids=("default", "pii-opt-in")) +def test_init_options_never_attach_request_bodies(env: Mapping[str, str]) -> None: + assert build_sentry_init_options(env)["max_request_body_size"] == "never" diff --git a/tests/unit/proxy/proxy_server/test_proxy_config.py b/tests/unit/proxy/proxy_server/test_proxy_config.py index a1e22780d57..5fd088acd01 100644 --- a/tests/unit/proxy/proxy_server/test_proxy_config.py +++ b/tests/unit/proxy/proxy_server/test_proxy_config.py @@ -186,6 +186,68 @@ async def test_tracing_config_automatically_exports_spend_without_a_storage_depe assert logger.task is not None and logger.task.done() and not logger.task.cancelled() +@pytest.mark.asyncio +async def test_tracing_logs_an_error_naming_the_store_when_lens_is_misconfigured( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + from litellm.proxy.tracing_runtime import manage_tracing + + monkeypatch.delenv("LITELLM_LENS_URL", raising=False) + monkeypatch.delenv("LITELLM_LENS_SERVICE_TOKEN", raising=False) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + async with manage_tracing(enabled=True) as receiver: + assert receiver is None + + failures: Final = [ + record.getMessage() + for record in caplog.records + if record.levelno == logging.ERROR and "lens" in record.getMessage() + ] + assert len(failures) == 1 + assert "LITELLM_LENS_URL" in failures[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "config_text", + ( + "model_list: []\nlitellm_settings:\n failure_callback:\n - sentry\n", + "model_list: []\nlitellm_settings:\n success_callback:\n - sentry\n", + "model_list: []\nlitellm_settings:\n input_callback:\n - sentry\n", + "model_list: []\ngeneral_settings:\n alerting:\n - sentry\n", + ), + ids=("failure_callback", "success_callback", "input_callback", "alerting"), +) +async def test_proxy_config_starts_sentry_at_boot_without_request_bodies( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, config_text: str +) -> None: + import sentry_sdk + + from litellm.litellm_core_utils import litellm_logging + + init: Final = MagicMock() + monkeypatch.setattr(sentry_sdk, "init", init) + monkeypatch.setenv("SENTRY_DSN", "https://key@sentry.example/1") + monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", None) + litellm_logging.start_sentry.cache_clear() + monkeypatch.setattr(litellm_logging, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "add_breadcrumb", None) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + config_file: Final = tmp_path / "sentry.yaml" + config_file.write_text(config_text) + + await ProxyConfig().load_config(router=None, config_file_path=str(config_file)) + + init.assert_called_once() + assert init.call_args.kwargs["max_request_body_size"] == "never" + assert litellm_logging.sentry_sdk_instance is sentry_sdk + # --------------------------------------------------------------------------- # _is_remote_module_url diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index ec8e2e0438a..1d197602543 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.bug_report import ( bug_report_notice, build_bug_report, ) +from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -4674,6 +4675,21 @@ class TestHandleLLMApiExceptionRetryAfter: proxy_exc = await self._invoke(ValueError("some other failure")) assert "retry-after" not in proxy_exc.headers + @pytest.mark.parametrize( + "exc", + ( + RouterRateLimitError(model="gpt-4", cooldown_time=42.3, enable_pre_call_checks=False, cooldown_list=[]), + RouterRateLimitErrorBasic(model="gpt-4"), + ), + ids=("cooldown", "basic"), + ) + async def test_handle_llm_api_exception_answers_a_cooldown_with_429_and_no_bug_notice(self, exc, caplog): + with caplog.at_level("ERROR"): + proxy_exc = await self._invoke(exc) + + assert (proxy_exc.code, proxy_exc.type) == ("429", "rate_limit_error") + assert ISSUE_URL_BASE not in caplog.text + async def test_handle_llm_api_exception_strips_bug_report_notice_from_client_message(self, caplog): report = build_bug_report(RuntimeError("boom"), surface="sdk") notice = bug_report_notice(report) diff --git a/tests/unit/proxy/utils/proxy_logging/test_alerting.py b/tests/unit/proxy/utils/proxy_logging/test_alerting.py index 43ee1094330..91c77fdbe4d 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_alerting.py +++ b/tests/unit/proxy/utils/proxy_logging/test_alerting.py @@ -7,15 +7,15 @@ Covers ``failed_tracking_alert``, ``budget_alerts``, ``alerting_handler``, from __future__ import annotations from datetime import datetime -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException from prisma.errors import PrismaError -import litellm from litellm._service_logger import ServiceTypes +from litellm.litellm_core_utils import litellm_logging from litellm.proxy._types import AlertType, CallInfo @@ -210,11 +210,45 @@ async def test_alerting_handler_sends_to_slack(proxy_logging): @pytest.mark.asyncio async def test_alerting_handler_sentry_without_sdk_error_raises(proxy_logging, monkeypatch): proxy_logging.alerting = ["sentry"] - monkeypatch.setattr(litellm.utils, "sentry_sdk_instance", None) + monkeypatch.delenv("SENTRY_DSN", raising=False) + monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", None) + litellm_logging.start_sentry.cache_clear() with pytest.raises(Exception, match="SENTRY_DSN"): await proxy_logging.alerting_handler(message="x", level="Low", alert_type=AlertType.db_exceptions) +@pytest.mark.asyncio +async def test_alerting_handler_sentry_sends_through_the_initialized_sdk(proxy_logging, monkeypatch): + proxy_logging.alerting = ["sentry"] + sentry = MagicMock() + monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", sentry) + await proxy_logging.alerting_handler(message="db down", level="High", alert_type=AlertType.db_exceptions) + sentry.capture_message.assert_called_once() + assert "db down" in sentry.capture_message.call_args.args[0] + + +@pytest.mark.asyncio +async def test_alerting_handler_starts_sentry_when_the_alert_was_not_configured_at_boot(proxy_logging, monkeypatch): + import sentry_sdk + + monkeypatch.setattr(proxy_logging, "alerting", ["sentry"]) + init: Final = MagicMock() + capture_message: Final = MagicMock() + monkeypatch.setattr(sentry_sdk, "init", init) + monkeypatch.setattr(sentry_sdk, "capture_message", capture_message) + monkeypatch.setenv("SENTRY_DSN", "https://key@sentry.example/1") + monkeypatch.setattr(litellm_logging, "sentry_sdk_instance", None) + litellm_logging.start_sentry.cache_clear() + monkeypatch.setattr(litellm_logging, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "add_breadcrumb", None) + + await proxy_logging.alerting_handler(message="db down", level="High", alert_type=AlertType.db_exceptions) + + assert init.call_args.kwargs["dsn"] == "https://key@sentry.example/1" + capture_message.assert_called_once() + assert "db down" in capture_message.call_args.args[0] + + # --------------------------------------------------------------------------- # failure_handler # --------------------------------------------------------------------------- @@ -235,7 +269,7 @@ async def test_failure_handler_logs_db_error_and_calls_service_logging(proxy_log proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) - monkeypatch.setattr(litellm.utils, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "capture_exception", None) await proxy_logging.failure_handler( original_exception=HTTPException(status_code=500, detail="boom"), duration=1.5, @@ -262,7 +296,7 @@ async def test_failure_handler_alerts_but_leaves_prisma_error_event_to_log_db_me proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) - monkeypatch.setattr(litellm.utils, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "capture_exception", None) await proxy_logging.failure_handler( original_exception=PrismaError("connection reset"), duration=1.0, call_type=call_type ) @@ -278,7 +312,7 @@ async def test_failure_handler_still_emits_db_event_for_wrapped_insert_error(pro proxy_logging.alert_types = [AlertType.db_exceptions] proxy_logging.alerting_handler = AsyncMock() proxy_logging.service_logging_obj = MagicMock(async_service_failure_hook=AsyncMock()) - monkeypatch.setattr(litellm.utils, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "capture_exception", None) await proxy_logging.failure_handler( original_exception=HTTPException(status_code=400, detail={"error": "Foreign Key Constraint failed"}), duration=1.0, @@ -298,7 +332,7 @@ async def test_failure_handler_with_capture_exception_invoked(proxy_logging, mon def fake_capture(error): captured["error"] = error - monkeypatch.setattr(litellm.utils, "capture_exception", fake_capture) + monkeypatch.setattr(litellm_logging, "capture_exception", fake_capture) err = RuntimeError("real") await proxy_logging.failure_handler(original_exception=err, duration=1.0, call_type="db_read") snapshot = { @@ -320,7 +354,7 @@ async def test_failure_handler_propagates_service_logging_error_raises(proxy_log proxy_logging.service_logging_obj = MagicMock( async_service_failure_hook=AsyncMock(side_effect=RuntimeError("svc")) ) - monkeypatch.setattr(litellm.utils, "capture_exception", None) + monkeypatch.setattr(litellm_logging, "capture_exception", None) with pytest.raises(RuntimeError): await proxy_logging.failure_handler( original_exception=Exception("x"), duration=0.0, call_type="db_read" diff --git a/tests/unit/tracing/test_exporter.py b/tests/unit/tracing/test_exporter.py index 5f67fbbb25b..225c96f2c57 100644 --- a/tests/unit/tracing/test_exporter.py +++ b/tests/unit/tracing/test_exporter.py @@ -1,10 +1,12 @@ import asyncio import json +import logging from itertools import chain from typing import Final import httpx import pytest +from prometheus_client import REGISTRY from litellm.tracing.exporter import MAX_BUFFER_EVENTS, MAX_EVENT_BYTES, ExportFailure, LensExporter, encode_record @@ -306,3 +308,51 @@ def test_serialization_rejects_recursive_payloads() -> None: @pytest.mark.parametrize("value", (float("nan"), object(), "\ud800"), ids=("nan", "unsupported", "surrogate")) def test_serialization_returns_a_failure_for_invalid_payloads(value: object) -> None: assert encode_record({"value": value}) is ExportFailure.INVALID + + +def _counter_value(name: str) -> float: + return REGISTRY.get_sample_value(name) or 0.0 + + +async def _drained(exporter: LensExporter) -> None: + while exporter.buffered_events: + await asyncio.sleep(0) + + +@pytest.mark.asyncio +async def test_written_and_dropped_rows_are_exported_as_prometheus_counters() -> None: + statuses: Final = iter((204, 500)) + + def answer(request: httpx.Request) -> httpx.Response: + return httpx.Response(next(statuses)) + + written_before: Final = _counter_value("litellm_lens_export_rows_written_total") + dropped_before: Final = _counter_value("litellm_lens_export_rows_dropped_total") + async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(answer)) as client: + exporter: Final = LensExporter(client) + exporter.start() + for request_id in ("accepted", "rejected"): + await exporter.async_log_success_event( + {"standard_logging_object": {"id": request_id, "response_cost": 0.1}}, None, None, None + ) + await _drained(exporter) + await exporter.aclose() + assert (exporter.rows_written, exporter.rows_dropped) == (1, 1) + assert _counter_value("litellm_lens_export_rows_written_total") - written_before == exporter.rows_written + assert _counter_value("litellm_lens_export_rows_dropped_total") - dropped_before == exporter.rows_dropped + + +@pytest.mark.asyncio +async def test_a_rejecting_store_is_logged_as_an_error_naming_lens(caplog: pytest.LogCaptureFixture) -> None: + def reject(request: httpx.Request) -> httpx.Response: + return httpx.Response(500) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + async with httpx.AsyncClient(base_url="http://lens", transport=httpx.MockTransport(reject)) as client: + exporter: Final = LensExporter(client) + assert exporter.enqueue(b"{}") + exporter.start() + await exporter.aclose() + failures: Final = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + assert len(failures) == 1 + assert "lens" in failures[0] and "HTTP 500" in failures[0] From 123e2071bfe7755bc97747ae79835ac0a88d8258 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:53:56 -0700 Subject: [PATCH 16/23] feat(wandb): add NVIDIA Nemotron 3.5 Super VL to the cost map (#46795) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- .../model_prices_and_context_window_backup.json | 16 ++++++++++++++++ model_prices_and_context_window.json | 16 ++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d0606578658..a204c49967e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -75167,5 +75167,21 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_tokens": 943718, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d0606578658..a204c49967e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -75167,5 +75167,21 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true + }, + "wandb/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B": { + "cache_read_input_token_cost": 2e-08, + "input_cost_per_token": 1.5e-07, + "litellm_provider": "wandb", + "max_input_tokens": 1049000, + "max_tokens": 943718, + "mode": "chat", + "output_cost_per_token": 7.5e-07, + "source": "https://openrouter.ai/api/v1/models/nvidia/nemotron-3.5-super-vl/endpoints", + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true } } From bd60cc8608c6d65f5dc89aea69aec60508d708d0 Mon Sep 17 00:00:00 2001 From: mrinal-berri Date: Thu, 15 Oct 2026 21:10:53 -0400 Subject: [PATCH 17/23] fix(bedrock): forward Nova Sonic realtime tool calls with their arguments (#46214) --- litellm/constants.py | 1 + litellm/llms/bedrock/realtime/handler.py | 67 +- .../llms/bedrock/realtime/transformation.py | 617 +++++++++++++----- litellm/realtime_api/main.py | 3 + litellm/types/llms/openai.py | 11 + .../llm_nonconversational.yaml | 1 + .../realtime/realtime_client.py | 11 +- .../realtime/test_realtime_e2e.py | 84 ++- .../test_realtime_pipecat_audio_e2e.py | 3 + .../realtime/test_realtime_pipecat_e2e.py | 5 +- tests/integration/_support/async_server.py | 8 +- .../_support/bedrock_realtime_peer.py | 429 ++++++++++++ ..._bedrock_nova_sonic_realtime_tool_calls.py | 378 +++++++++++ .../realtime/test_bedrock_realtime_handler.py | 102 +++ .../test_bedrock_realtime_transformation.py | 580 +++++++++++++++- 15 files changed, 2063 insertions(+), 237 deletions(-) create mode 100644 tests/integration/_support/bedrock_realtime_peer.py create mode 100644 tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py diff --git a/litellm/constants.py b/litellm/constants.py index 1f9e81443f7..f7ba55d8ce4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -380,6 +380,7 @@ REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = get_env_int_or_none("REALTIME REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = get_env_float( "REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS", 20.0 ) +BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS: Final = get_env_float("BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS", 5.0) # RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 2ceac5899fd..c0a382b136c 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -17,7 +17,9 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm._logging import redact_string, verbose_proxy_logger +from litellm._uuid import uuid from litellm.constants import ( + BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, BEDROCK_REALTIME_SDK_DISTRIBUTION, @@ -35,7 +37,7 @@ from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig +from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) @@ -232,12 +234,43 @@ async def _ack_session_update( return True +async def _reject_session_update( + client_ws: RealtimeClientWebSocket, + parsed_client_message: Mapping[str, JsonValue], + rejection: BedrockSessionUpdateRejection, +) -> bool: + """Send an error event for a session.update Bedrock cannot apply; False means the client is gone.""" + client_event_id: Final = parsed_client_message.get("event_id") + error_event: Final = { + "type": "error", + "event_id": f"event_{uuid.uuid4()}", + "error": { + "type": "invalid_request_error", + "code": rejection.code, + "message": rejection.message, + "param": rejection.param, + "event_id": client_event_id if isinstance(client_event_id, str) else None, + }, + } + try: + await client_ws.send_text(json.dumps(error_event)) + except Exception as e: # noqa: BLE001 # any send failure means the client is gone + verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) + return False + return True + + class BedrockRealtime(BaseAWSLLM): """Handler for Bedrock Nova Sonic realtime speech-to-speech API.""" - def __init__(self, sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version): + def __init__( + self, + sdk_version_lookup: Callable[[], str | None] = _installed_sdk_version, + close_timeout: float = BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, + ): super().__init__() self._sdk_version_lookup: Final = sdk_version_lookup + self._close_timeout: Final = close_timeout async def async_realtime( self, @@ -259,6 +292,7 @@ class BedrockRealtime(BaseAWSLLM): aws_bedrock_runtime_endpoint: str | None = None, aws_external_id: str | None = None, aws_session_tags: object = None, + drop_params: bool = False, **kwargs: object, ): """ @@ -341,7 +375,9 @@ class BedrockRealtime(BaseAWSLLM): ) try: - await self._run_session(websocket, open_bidirectional_stream, model, logging_obj, pending_session_update) + await self._run_session( + websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params + ) finally: await _close_bedrock_client(bedrock_client) @@ -352,8 +388,9 @@ class BedrockRealtime(BaseAWSLLM): model: str, logging_obj: LiteLLMLogging, pending_session_update: str | None, + drop_params: bool = False, ) -> None: - transformation_config: Final = BedrockRealtimeConfig() + transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) bedrock_stream: Final = await open_bidirectional_stream() @@ -489,6 +526,11 @@ class BedrockRealtime(BaseAWSLLM): parsed_client_message = _parse_client_message(message) is_session_update = _json_str(parsed_client_message.get("type")) == "session.update" if is_session_update: + rejection = transformation_config.session_update_rejection(parsed_client_message) + if rejection is not None: + if not await _reject_session_update(client_ws, parsed_client_message, rejection): + break + continue client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message transformed_messages = transformation_config.transform_realtime_request( @@ -504,11 +546,20 @@ class BedrockRealtime(BaseAWSLLM): ): break finally: - for close_message in transformation_config.session_close_messages(): + + async def close_bedrock_input() -> None: + for close_message in transformation_config.session_close_messages(): + with contextlib.suppress(Exception): + await send_to_bedrock(close_message) with contextlib.suppress(Exception): - await send_to_bedrock(close_message) - with contextlib.suppress(Exception): - await bedrock_stream.input_stream.close() + await bedrock_stream.input_stream.close() + + try: + await asyncio.wait_for(close_bedrock_input(), self._close_timeout) + except (asyncio.TimeoutError, TimeoutError): + verbose_proxy_logger.debug( + "Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", self._close_timeout + ) async def _forward_bedrock_to_client( self, diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index 3038b0831f4..e8bc0482a28 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -5,13 +5,16 @@ Transforms between OpenAI Realtime API format and Bedrock Nova Sonic format. """ import base64 +import dataclasses import json +import math import uuid as uuid_lib -from collections.abc import Mapping -from typing import Final, cast +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final, Literal import httpx -from pydantic import TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -24,6 +27,8 @@ from litellm.types.llms.openai import ( OpenAIRealtimeContentPartDone, OpenAIRealtimeDoneEvent, OpenAIRealtimeEvents, + OpenAIRealtimeFunctionCallArgumentsDelta, + OpenAIRealtimeFunctionCallArgumentsDone, OpenAIRealtimeInputAudioBufferSpeechEvent, OpenAIRealtimeInputAudioTranscriptionCompleted, OpenAIRealtimeInputAudioTranscriptionDelta, @@ -35,7 +40,9 @@ from litellm.types.llms.openai import ( OpenAIRealtimeResponseTextDone, OpenAIRealtimeResponseUsage, OpenAIRealtimeStreamResponseBaseObject, + OpenAIRealtimeStreamResponseOutputItem, OpenAIRealtimeStreamResponseOutputItemAdded, + OpenAIRealtimeStreamResponseOutputItemContent, OpenAIRealtimeStreamSession, OpenAIRealtimeStreamSessionEvents, OpenAIRealtimeUsageTokenDetails, @@ -52,6 +59,26 @@ _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) class BedrockContentEnd(LiteLLMBaseModel): stopReason: str | None = None + type: str | None = None + + +def _tool_function(tool: JsonValue) -> Mapping[str, JsonValue] | None: + if not isinstance(tool, dict) or tool.get("type") != "function": + return None + function: Final = tool.get("function") + return function if isinstance(function, dict) and function else tool + + +class BedrockToolUse(LiteLLMBaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + toolUseId: str = "" + toolName: str = "" + content: JsonValue = "{}" + + @property + def arguments(self) -> str: + return self.content if isinstance(self.content, str) else json.dumps(self.content) class BedrockUsageTokenDetails(LiteLLMBaseModel): @@ -80,12 +107,63 @@ TRIGGER_AUDIO_BYTES_PER_SECOND: Final = TRIGGER_AUDIO_SAMPLE_RATE_HERTZ * 2 TRIGGER_LEADING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND // 2) TRIGGER_TRAILING_SILENCE: Final = bytes(TRIGGER_AUDIO_BYTES_PER_SECOND * 3) TRIGGER_AUDIO_CHUNK_SIZE: Final = 1024 +PROMPT_STARTING_CLIENT_EVENT_TYPES: Final = frozenset( + ("input_audio_buffer.append", "input_audio_buffer.commit", "conversation.item.create", "response.create") +) + + +def _number_setting(value: JsonValue) -> float | None: + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def _integer_setting(value: JsonValue) -> int | None: + number: Final = _number_setting(value) + return None if number is None else int(number) + + +def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]: + temperature: Final = session_config.get("temperature") + max_output_tokens: Final = session_config.get("max_response_output_tokens") + output_rate_hertz: Final = session_config.get("output_sample_rate_hertz") + malformed: Final = ( + ("temperature", temperature is not None and _number_setting(temperature) is None), + ( + "max_response_output_tokens", + max_output_tokens is not None + and max_output_tokens != "inf" + and _integer_setting(max_output_tokens) is None, + ), + ("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None), + ) + return tuple(name for name, is_malformed in malformed if is_malformed) + + +@dataclasses.dataclass(frozen=True, slots=True) +class BedrockSessionUpdateRejection: + code: Literal["invalid_type", "session_update_not_applied"] + message: str + param: str + + +@dataclasses.dataclass(frozen=True, slots=True) +class BedrockLockedSessionSettings: + """The session settings Bedrock Nova Sonic locks in once the prompt starts.""" + + tools_json: str + system_prompt: str + voice_id: str + output_sample_rate_hertz: int + temperature: float + max_tokens: int | None class BedrockRealtimeConfig(BaseRealtimeConfig): """Configuration for Bedrock Nova Sonic realtime transformations.""" - def __init__(self): + def __init__(self, drop_params: bool = False): + self.drop_params = drop_params # Track session state self.prompt_name = str(uuid_lib.uuid4()) self.content_name = str(uuid_lib.uuid4()) @@ -95,7 +173,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): # Default configuration values # Inference configuration - self.max_tokens = 1024 + self.max_tokens: int | None = 1024 self.top_p = 0.9 self.temperature = 0.7 @@ -128,6 +206,19 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self._user_transcript_buffer = "" self._cumulative_usage = BedrockUsageEvent() self._reported_usage = BedrockUsageEvent() + self._response_done_items: Mapping[int, OpenAIRealtimeStreamResponseOutputItem] = MappingProxyType({}) + self._response_item_count = 0 + self._message_output_index = 0 + self._assistant_block_open = False + self._response_done_pending = False + self._started_settings = BedrockLockedSessionSettings( + tools_json="[]", + system_prompt="", + voice_id=self.voice_id, + output_sample_rate_hertz=self.output_sample_rate_hertz, + temperature=self.temperature, + max_tokens=self.max_tokens, + ) def get_error_class( self, @@ -149,7 +240,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): """Bedrock requires session configuration.""" return True - def session_configuration_request(self, model: str, tools: list[dict] | None = None) -> str: + def session_configuration_request(self, model: str, tools: Sequence[JsonValue] | None = None) -> str: """ Create initial session configuration for Bedrock Nova Sonic. @@ -159,19 +250,9 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): Returns JSON string with session start and prompt start events. """ - session_start: Final = { - "event": { - "sessionStart": { - "inferenceConfiguration": { - "maxTokens": self.max_tokens, - "topP": self.top_p, - "temperature": self.temperature, - } - } - } - } + session_start: Final = {"event": {"sessionStart": {"inferenceConfiguration": self._inference_configuration()}}} - prompt_start_config: Final = { + prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, "audioOutputConfiguration": { @@ -183,41 +264,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "encoding": self.output_encoding, "audioType": self.output_audio_type, }, + **( + { + "toolUseOutputConfiguration": {"mediaType": "application/json"}, + "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, + } + if tools + else {} + ), } - # Add tool configuration if tools are provided - if tools: - prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"} - prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)} - prompt_start: Final = {"event": {"promptStart": prompt_start_config}} # Return as a marker that we've sent the configuration return json.dumps({"session_start": session_start, "prompt_start": prompt_start}) - def _transform_tools_to_bedrock_format(self, tools: list[dict]) -> list[dict]: - """ - Transform OpenAI tool format to Bedrock tool format. - - Args: - tools: List of OpenAI format tools - - Returns: - List of Bedrock format tools - """ - bedrock_tools: Final = [] - for tool in tools: - if tool.get("type") == "function": - function = tool.get("function", {}) - bedrock_tool = { - "toolSpec": { - "name": function.get("name", ""), - "description": function.get("description", ""), - "inputSchema": {"json": json.dumps(function.get("parameters", {}))}, - } + def _transform_tools_to_bedrock_format(self, tools: Sequence[JsonValue]) -> tuple[Mapping[str, JsonValue], ...]: + functions: Final = (_tool_function(tool) for tool in tools) + return tuple( + { + "toolSpec": { + "name": function.get("name", ""), + "description": function.get("description", ""), + "inputSchema": {"json": json.dumps(function.get("parameters", {}))}, } - bedrock_tools.append(bedrock_tool) - return bedrock_tools + } + for function in functions + if function is not None + ) def _map_audio_format_to_sample_rate(self, audio_format: str, is_output: bool = True) -> int: """ @@ -252,47 +326,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): messages: Final[list[str]] = [] session_config: Final = json_message.get("session", {}) + locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings) - # Update inference configuration from session if provided - if "max_response_output_tokens" in session_config: - self.max_tokens = session_config["max_response_output_tokens"] - if "temperature" in session_config: - self.temperature = session_config["temperature"] - - # Update audio output configuration from session if provided - if "voice" in session_config: - self.voice_id = session_config["voice"] - if "output_audio_format" in session_config: - output_format: Final = session_config["output_audio_format"] - self.output_sample_rate_hertz = self._map_audio_format_to_sample_rate(output_format, is_output=True) + # Settings Bedrock locks at prompt start may only change before it + if not self.prompt_started: + self.max_tokens = locked_settings.max_tokens + self.temperature = locked_settings.temperature + self.voice_id = locked_settings.voice_id + self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz # Update audio input configuration from session if provided if "input_audio_format" in session_config: input_format: Final = session_config["input_audio_format"] self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) - # Allow direct override of sample rates if provided (custom extension) - if "output_sample_rate_hertz" in session_config: - self.output_sample_rate_hertz = session_config["output_sample_rate_hertz"] + # Allow direct override of the input sample rate if provided (custom extension) if "input_sample_rate_hertz" in session_config: self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"] + if self.prompt_started: + return messages + # Send session start - session_start: Final = { - "event": { - "sessionStart": { - "inferenceConfiguration": { - "maxTokens": self.max_tokens, - "topP": self.top_p, - "temperature": self.temperature, - } - } - } - } + session_start: Final = {"event": {"sessionStart": {"inferenceConfiguration": self._inference_configuration()}}} messages.append(json.dumps(session_start)) # Send prompt start - prompt_start_config: Final = { + tools: Final[Sequence[JsonValue] | None] = session_config.get("tools") + prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, "audioOutputConfiguration": { @@ -304,63 +365,138 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "encoding": self.output_encoding, "audioType": self.output_audio_type, }, + **( + { + "toolUseOutputConfiguration": {"mediaType": "application/json"}, + "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, + } + if tools + else {} + ), } - # Add tool configuration if tools are provided - tools: Final = session_config.get("tools") - if tools: - prompt_start_config["toolUseOutputConfiguration"] = {"mediaType": "application/json"} - prompt_start_config["toolConfiguration"] = {"tools": self._transform_tools_to_bedrock_format(tools)} - prompt_start: Final = {"event": {"promptStart": prompt_start_config}} messages.append(json.dumps(prompt_start)) self.prompt_started = True + self._started_settings = locked_settings - # Send system prompt if provided - instructions: Final = session_config.get("instructions") - if instructions: - text_content_name: Final = str(uuid_lib.uuid4()) + # Bedrock requires the first content block in a prompt to be SYSTEM role + instructions: Final = locked_settings.system_prompt + text_content_name: Final = str(uuid_lib.uuid4()) - # Content start - text_content_start: Final = { - "event": { - "contentStart": { - "promptName": self.prompt_name, - "contentName": text_content_name, - "type": "TEXT", - "interactive": False, - "role": "SYSTEM", - "textInputConfiguration": {"mediaType": self.text_media_type}, - } + # Content start + text_content_start: Final = { + "event": { + "contentStart": { + "promptName": self.prompt_name, + "contentName": text_content_name, + "type": "TEXT", + "interactive": False, + "role": "SYSTEM", + "textInputConfiguration": {"mediaType": self.text_media_type}, } } - messages.append(json.dumps(text_content_start)) + } + messages.append(json.dumps(text_content_start)) - # Text input - text_input: Final = { - "event": { - "textInput": { - "promptName": self.prompt_name, - "contentName": text_content_name, - "content": instructions, - } + # Text input + text_input: Final = { + "event": { + "textInput": { + "promptName": self.prompt_name, + "contentName": text_content_name, + "content": instructions, } } - messages.append(json.dumps(text_input)) + } + messages.append(json.dumps(text_input)) - # Content end - text_content_end: Final = { - "event": { - "contentEnd": { - "promptName": self.prompt_name, - "contentName": text_content_name, - } + # Content end + text_content_end: Final = { + "event": { + "contentEnd": { + "promptName": self.prompt_name, + "contentName": text_content_name, } } - messages.append(json.dumps(text_content_end)) + } + messages.append(json.dumps(text_content_end)) return messages + def _inference_configuration(self) -> Mapping[str, int | float]: + return { + "topP": self.top_p, + "temperature": self.temperature, + **({"maxTokens": self.max_tokens} if self.max_tokens is not None else {}), + } + + def _apply_locked_settings( + self, session_config: Mapping[str, JsonValue], base: BedrockLockedSessionSettings + ) -> BedrockLockedSessionSettings: + output_audio_format: Final = session_config.get("output_audio_format") + temperature: Final = session_config.get("temperature") + max_output_tokens: Final = session_config.get("max_response_output_tokens") + requested_tools: Final = session_config.get("tools") + instructions: Final = session_config.get("instructions") + voice: Final = session_config.get("voice") + tools_json: Final = ( + json.dumps(self._transform_tools_to_bedrock_format(requested_tools)) + if isinstance(requested_tools, list) + else base.tools_json + ) + system_prompt: Final = base.system_prompt if instructions is None else str(instructions) + voice_id: Final = base.voice_id if voice is None else str(voice) + format_sample_rate: Final = ( + base.output_sample_rate_hertz + if output_audio_format is None + else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True) + ) + requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz")) + output_sample_rate_hertz: Final = ( + requested_output_rate if requested_output_rate is not None else format_sample_rate + ) + requested_temperature: Final = _number_setting(temperature) + requested_max_tokens: Final = _integer_setting(max_output_tokens) + locked_temperature: Final = requested_temperature if requested_temperature is not None else base.temperature + explicit_max_tokens: Final = requested_max_tokens if requested_max_tokens is not None else base.max_tokens + max_tokens: Final = None if max_output_tokens == "inf" else explicit_max_tokens + return dataclasses.replace( + base, + tools_json=tools_json, + system_prompt=system_prompt, + voice_id=voice_id, + output_sample_rate_hertz=output_sample_rate_hertz, + temperature=locked_temperature, + max_tokens=max_tokens, + ) + + def session_update_rejection(self, json_message: Mapping[str, JsonValue]) -> BedrockSessionUpdateRejection | None: + """The error for a session.update this session cannot apply, or None when it applies.""" + session_config: Final = json_message.get("session") + if not isinstance(session_config, Mapping): + return None + malformed: Final = () if self.drop_params else _malformed_number_settings(session_config) + if malformed: + return BedrockSessionUpdateRejection( + code="invalid_type", + message=( + f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature " + "and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. " + "Set drop_params to drop it instead" + ), + param=f"session.{malformed[0]}", + ) + if not self.prompt_started: + return None + if self._apply_locked_settings(session_config, self._started_settings) != self._started_settings: + return BedrockSessionUpdateRejection( + code="session_update_not_applied", + message="Bedrock Nova Sonic cannot change tools, instructions, voice, output audio format, temperature or max output tokens after the session has started, so this session.update was not applied", + param="session", + ) + return None + def transform_input_audio_buffer_append_event(self, json_message: dict) -> list[str]: """ Transform input_audio_buffer.append event to Bedrock audio input. @@ -655,18 +791,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return [] message_type: Final = json_message.get("type") - - # Route to appropriate transformation method + prompt_start_messages: Final = ( + self.transform_session_update_event({"session": {}}) + if message_type in PROMPT_STARTING_CLIENT_EVENT_TYPES and not self.prompt_started + else [] + ) if message_type == "session.update": return self.transform_session_update_event(json_message) elif message_type == "input_audio_buffer.append": - return self.transform_input_audio_buffer_append_event(json_message) + return prompt_start_messages + self.transform_input_audio_buffer_append_event(json_message) elif message_type == "input_audio_buffer.commit": - return self.transform_input_audio_buffer_commit_event(json_message) + return prompt_start_messages + self.transform_input_audio_buffer_commit_event(json_message) elif message_type == "conversation.item.create": - return self.transform_conversation_item_create_event(json_message) + return prompt_start_messages + self.transform_conversation_item_create_event(json_message) elif message_type == "response.create": - return self.transform_response_create_event(json_message) + return prompt_start_messages + self.transform_response_create_event(json_message) elif message_type == "response.cancel": return self.transform_response_cancel_event(json_message) else: @@ -718,6 +857,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id: str | None, current_output_item_id: str | None, current_conversation_id: str | None, + current_delta_type: ALL_DELTA_TYPES | None, ) -> tuple[ list[OpenAIRealtimeEvents], str | None, @@ -751,7 +891,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, - None, + current_delta_type, ) generation_stage: Final = self._parse_generation_stage(content_start.get("additionalModelFields")) @@ -776,12 +916,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id = f"resp_{uuid.uuid4()}" if not current_output_item_id: current_output_item_id = f"item_{uuid.uuid4()}" + self._message_output_index = self._response_item_count + self._response_item_count += 1 if not current_conversation_id: current_conversation_id = f"conv_{uuid.uuid4()}" # Determine content type content_type: Final = content_start.get("type", "TEXT") - current_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio" + next_delta_type: Final[ALL_DELTA_TYPES] = "text" if content_type == "TEXT" else "audio" returned_messages: Final[list[OpenAIRealtimeEvents]] = [] @@ -805,7 +947,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded( type="response.output_item.added", response_id=current_response_id, - output_index=0, + output_index=self._message_output_index, item={ "id": current_output_item_id, "object": "realtime.item", @@ -815,18 +957,17 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "content": [], }, ) + self._assistant_block_open = True returned_messages.append(output_item_added) # Send response.content_part.added content_part_added: Final = OpenAIRealtimeResponseContentPartAdded( type="response.content_part.added", content_index=0, - output_index=0, + output_index=self._message_output_index, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - part=( - {"type": "text", "text": ""} if current_delta_type == "text" else {"type": "audio", "transcript": ""} - ), + part=({"type": "text", "text": ""} if next_delta_type == "text" else {"type": "audio", "transcript": ""}), response_id=current_response_id, ) returned_messages.append(content_part_added) @@ -836,7 +977,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, - current_delta_type, + next_delta_type, ) @staticmethod @@ -973,7 +1114,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, delta=text_content, ) @@ -1013,7 +1154,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, delta=audio_content, ) @@ -1049,22 +1190,21 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return [], current_delta_chunks if not current_output_item_id or not current_response_id: return [], current_delta_chunks + if not self._assistant_block_open: + return [], current_delta_chunks + self._assistant_block_open = False returned_messages: Final[list[OpenAIRealtimeEvents]] = [] + accumulated_text: Final = "".join(chunk.get("delta", "") for chunk in current_delta_chunks or []) # Send appropriate done event based on type if current_delta_type == "text": - # Accumulate text - accumulated_text = "" - if current_delta_chunks: - accumulated_text = "".join([chunk.get("delta", "") for chunk in current_delta_chunks]) - text_done: Final = OpenAIRealtimeResponseTextDone( type="response.text.done", content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, text=accumulated_text, ) @@ -1076,7 +1216,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, part={"type": "text", "text": accumulated_text}, response_id=current_response_id, ) @@ -1088,7 +1228,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, ) returned_messages.append(audio_done) @@ -1099,28 +1239,34 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): content_index=0, event_id=f"event_{uuid.uuid4()}", item_id=current_output_item_id, - output_index=0, + output_index=self._message_output_index, part={"type": "audio", "transcript": ""}, response_id=current_response_id, ) returned_messages.append(content_part_done) - # Send output_item.done + completed_parts: Final = self._completed_parts(current_delta_type, accumulated_text) + existing_item: Final = self._response_done_items.get(self._message_output_index) + existing_content: Final = existing_item.get("content", []) if existing_item is not None else [] + completed_message_item: Final[OpenAIRealtimeStreamResponseOutputItem] = { + "id": current_output_item_id, + "object": "realtime.item", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [*existing_content, *completed_parts], + } output_item_done: Final = OpenAIRealtimeOutputItemDone( type="response.output_item.done", event_id=f"event_{uuid.uuid4()}", - output_index=0, + output_index=self._message_output_index, response_id=current_response_id, - item={ - "id": current_output_item_id, - "object": "realtime.item", - "type": "message", - "status": "completed", - "role": "assistant", - "content": [], - }, + item=completed_message_item, ) returned_messages.append(output_item_done) + self._response_done_items = MappingProxyType( + {**self._response_done_items, self._message_output_index: completed_message_item} + ) # Reset delta chunks return returned_messages, None @@ -1165,6 +1311,12 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return [], None, None, None usage: Final = self._take_usage_delta() + output_items: Final = tuple(item for _, item in sorted(self._response_done_items.items())) + self._response_done_items = MappingProxyType({}) + self._response_item_count = 0 + self._message_output_index = 0 + self._assistant_block_open = False + self._response_done_pending = False response_done: Final = OpenAIRealtimeDoneEvent( type="response.done", event_id=f"event_{uuid.uuid4()}", @@ -1172,7 +1324,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): object="realtime.response", id=current_response_id, status="completed", - output=[], + output=list(output_items), conversation_id=current_conversation_id, usage=dict(usage), ), @@ -1184,56 +1336,148 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): def transform_tool_use_event( self, event: dict, - current_output_item_id: str | None, current_response_id: str | None, + current_conversation_id: str | None, ) -> tuple[list[OpenAIRealtimeEvents], str, str]: """ Transform Bedrock toolUse event to OpenAI format. Args: event: Bedrock toolUse event - current_output_item_id: Current output item ID current_response_id: Current response ID + current_conversation_id: Current conversation ID Returns: - Tuple of (events, tool_call_id, tool_name) for tracking + Tuple of (events, response_id, conversation_id) """ verbose_logger.debug("Handling toolUse") - tool_use: Final = event["toolUse"] + tool_use: Final = BedrockToolUse.model_validate(event["toolUse"]) + verbose_logger.debug("Tool use event: %s (ID: %s)", tool_use.toolName, tool_use.toolUseId) - if not current_output_item_id or not current_response_id: - return [], "", "" + response_id: Final = current_response_id or f"resp_{uuid.uuid4()}" + conversation_id: Final = current_conversation_id or f"conv_{uuid.uuid4()}" + item_id: Final = f"item_{uuid.uuid4()}" + output_index: Final = self._response_item_count + self._response_item_count += 1 - # Parse the tool input - tool_input = {} - if "input" in tool_use: - try: - tool_input = json.loads(tool_use["input"]) if isinstance(tool_use["input"], str) else tool_use["input"] - except json.JSONDecodeError: - tool_input = {} + response_created_events: Final[tuple[OpenAIRealtimeEvents, ...]] = ( + ( + OpenAIRealtimeStreamResponseBaseObject( + type="response.created", + event_id=f"event_{uuid.uuid4()}", + response={ + "object": "realtime.response", + "id": response_id, + "status": "in_progress", + "output": [], + "conversation_id": conversation_id, + }, + ), + ) + if current_response_id is None + else () + ) - tool_call_id: Final = tool_use.get("toolUseId", "") - tool_name: Final = tool_use.get("toolName", "") + output_item_added: Final = OpenAIRealtimeStreamResponseOutputItemAdded( + type="response.output_item.added", + response_id=response_id, + output_index=output_index, + item={ + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "in_progress", + "call_id": tool_use.toolUseId, + "name": tool_use.toolName, + "arguments": "", + }, + ) - # Create a function call arguments done event - # This is a custom event format that matches what clients expect - function_call_event: Final[dict[str, object]] = { - "type": "response.function_call_arguments.done", - "event_id": f"event_{uuid.uuid4()}", - "response_id": current_response_id, - "item_id": current_output_item_id, - "output_index": 0, - "call_id": tool_call_id, - "name": tool_name, - "arguments": json.dumps(tool_input), + arguments_delta: Final = OpenAIRealtimeFunctionCallArgumentsDelta( + type="response.function_call_arguments.delta", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + item_id=item_id, + output_index=output_index, + call_id=tool_use.toolUseId, + delta=tool_use.arguments, + ) + + function_call_event: Final = OpenAIRealtimeFunctionCallArgumentsDone( + type="response.function_call_arguments.done", + event_id=f"event_{uuid.uuid4()}", + response_id=response_id, + item_id=item_id, + output_index=output_index, + call_id=tool_use.toolUseId, + name=tool_use.toolName, + arguments=tool_use.arguments, + ) + + completed_item: Final[OpenAIRealtimeStreamResponseOutputItem] = { + "id": item_id, + "object": "realtime.item", + "type": "function_call", + "status": "completed", + "call_id": tool_use.toolUseId, + "name": tool_use.toolName, + "arguments": tool_use.arguments, } + output_item_done: Final = OpenAIRealtimeOutputItemDone( + type="response.output_item.done", + event_id=f"event_{uuid.uuid4()}", + output_index=output_index, + response_id=response_id, + item=completed_item, + ) + + self._response_done_items = MappingProxyType({**self._response_done_items, output_index: completed_item}) return ( - [cast(OpenAIRealtimeEvents, function_call_event)], - tool_call_id, - tool_name, + [ + *response_created_events, + output_item_added, + arguments_delta, + function_call_event, + output_item_done, + ], + response_id, + conversation_id, ) + def _tool_content_end_events( + self, + current_output_item_id: str | None, + current_response_id: str | None, + current_conversation_id: str | None, + current_delta_type: ALL_DELTA_TYPES | None, + ) -> tuple[ + tuple[OpenAIRealtimeEvents, ...], + str | None, + str | None, + ALL_DELTA_TYPES | None, + ]: + if self._assistant_block_open: + self._response_done_pending = True + return (), current_output_item_id, current_response_id, current_delta_type + done_events, output_item_id, response_id, delta_type = self._response_done_events( + current_response_id, current_conversation_id + ) + return tuple(done_events), output_item_id, response_id, delta_type + + @staticmethod + def _completed_parts( + delta_type: str | None, accumulated_text: str + ) -> tuple[OpenAIRealtimeStreamResponseOutputItemContent, ...]: + if delta_type == "text": + return ({"type": "text", "text": accumulated_text},) + elif delta_type == "audio": + return ({"type": "audio", "transcript": ""},) + return () + + def _should_close_response(self, stop_reason: str | None) -> bool: + return stop_reason == "END_TURN" or self._response_done_pending + def transform_conversation_item_create_tool_result_event(self, json_message: dict) -> list[str]: """ Transform conversation.item.create with tool result to Bedrock format. @@ -1363,6 +1607,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_response_id, current_output_item_id, current_conversation_id, + current_delta_type, ) returned_messages.extend(events) @@ -1382,6 +1627,17 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): events = self.transform_audio_output_event(event, current_output_item_id, current_response_id) returned_messages.extend(events) + elif "contentEnd" in event and BedrockContentEnd.model_validate(event["contentEnd"]).type == "TOOL": + ( + done_events, + current_output_item_id, # rebind-ok: same state-threading pattern as the adjacent event branches + current_response_id, # rebind-ok: same state-threading pattern as the adjacent event branches + current_delta_type, # rebind-ok: same state-threading pattern as the adjacent event branches + ) = self._tool_content_end_events( + current_output_item_id, current_response_id, current_conversation_id, current_delta_type + ) + returned_messages.extend(done_events) + elif "contentEnd" in event and self._user_transcript_active: self._user_transcript_active = False self._user_transcript_generation_stage = None @@ -1396,7 +1652,8 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): current_delta_chunks, ) returned_messages.extend(events) - if BedrockContentEnd.model_validate(event["contentEnd"]).stopReason == "END_TURN": + content_end: Final = BedrockContentEnd.model_validate(event["contentEnd"]) + if self._should_close_response(content_end.stopReason): ( done_events, current_output_item_id, @@ -1412,12 +1669,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.transform_usage_event(BedrockUsageEvent.model_validate(event["usageEvent"])) elif "toolUse" in event: - events, tool_call_id, tool_name = self.transform_tool_use_event( - event, current_output_item_id, current_response_id + events, current_response_id, current_conversation_id = self.transform_tool_use_event( + event, current_response_id, current_conversation_id ) returned_messages.extend(events) - # Store tool call info for potential use - verbose_logger.debug("Tool use event: %s (ID: %s)", tool_name, tool_call_id) elif "promptEnd" in event or "completionEnd" in event: ( diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 02b95c80b60..367d5be7bfc 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -33,6 +33,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import CallTypes, LlmProviders from litellm.utils import ProviderConfigManager +from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -584,6 +585,7 @@ async def _arealtime( aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint") aws_external_id: Final = kwargs.get("aws_external_id") aws_session_tags: Final = kwargs.get("aws_session_tags") + drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True await bedrock_realtime.async_realtime( model=model, @@ -604,6 +606,7 @@ async def _arealtime( aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_external_id=aws_external_id, aws_session_tags=aws_session_tags, + drop_params=drop_params, ) elif _custom_llm_provider == "xai": api_base = ( diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 1e4f78b2204..8731bc2905a 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2271,6 +2271,16 @@ class OpenAIRealtimeContentPartDone(TypedDict): type: Literal["response.content_part.done"] +class OpenAIRealtimeFunctionCallArgumentsDelta(TypedDict): + type: ReadOnly[Literal["response.function_call_arguments.delta"]] + event_id: ReadOnly[str] + response_id: ReadOnly[str] + item_id: ReadOnly[str] + output_index: ReadOnly[int] + call_id: ReadOnly[str] + delta: ReadOnly[str] + + class OpenAIRealtimeFunctionCallArgumentsDone(TypedDict): type: Literal["response.function_call_arguments.done"] event_id: str @@ -2438,6 +2448,7 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeResponseAudioDone | OpenAIRealtimeContentPartDone | OpenAIRealtimeOutputItemDone + | OpenAIRealtimeFunctionCallArgumentsDelta | OpenAIRealtimeFunctionCallArgumentsDone | OpenAIRealtimeDoneEvent | OpenAIRealtimeInputAudioBufferSpeechEvent diff --git a/tests/e2e/coverage_registry/llm_nonconversational.yaml b/tests/e2e/coverage_registry/llm_nonconversational.yaml index b05e9c916d2..330ddf88fc5 100644 --- a/tests/e2e/coverage_registry/llm_nonconversational.yaml +++ b/tests/e2e/coverage_registry/llm_nonconversational.yaml @@ -62,6 +62,7 @@ - {id: llm.files.vertex.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Vertex unified file streams the GCS object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.files.bedrock.content.nonstream.works, module: llm, tier: P0, subject_endpoint: files, route: bedrock_converse, capability: basic, streaming: nonstream, assertions: [works], source: "test_batches_e2e.py / LIT-5730", rationale: "GET /v1/files/{id}/content on a Bedrock unified file streams the S3 object back (provider-transformed JSONL, so asserts non-empty JSON lines rather than byte equality)"} - {id: llm.realtime.bedrock_converse.basic.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: basic, streaming: stream, assertions: [works], source: "test_realtime_bedrock_e2e.py", rationale: "Nova Sonic realtime session emits response.done (LIT-2239)"} +- {id: llm.realtime.bedrock_converse.tool_use.stream.works, module: llm, tier: P0, subject_endpoint: realtime, route: bedrock_converse, capability: tool_use, streaming: stream, assertions: [works], source: "test_realtime_e2e.py", rationale: "Nova Sonic realtime forwards each tool call with its arguments and completes the tool-result round trip (LIT-9465)"} - {id: llm.google_native.gemini.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "LIT-4076 / proxy/google_endpoints/endpoints.py", fail_before_fix: proven, rationale: "google-native generateContent must stamp x-litellm-response-cost so SDK traffic reconciles against spend"} - {id: llm.google_native.gemini.basic.stream.works, module: llm, tier: P0, subject_endpoint: google_native, route: gemini, capability: basic, streaming: stream, assertions: [works], source: "PR #28213 / proxy/proxy_server.py async_data_generator", fail_before_fix: proven, rationale: "streamGenerateContent must relay single-prefixed SSE frames with no [DONE] sentinel; doubled data: prefixes and the OpenAI terminator both break the Vertex Java SDK"} - {id: llm.realtime.openai.basic.nonstream.works, module: llm, tier: P1, subject_endpoint: realtime, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "vendor strategy §9.19 / LIT-4778", rationale: "HTTP /v1/realtime/client_secrets returns an ephemeral credential"} diff --git a/tests/e2e/llm_translation/realtime/realtime_client.py b/tests/e2e/llm_translation/realtime/realtime_client.py index c9f06e787f3..4d8336775ef 100644 --- a/tests/e2e/llm_translation/realtime/realtime_client.py +++ b/tests/e2e/llm_translation/realtime/realtime_client.py @@ -85,7 +85,16 @@ PROVIDERS = ( vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", ), ), - # RealtimeProvider("bedrock", "bedrock-realtime", ...) # TODO: Enable when Bedrock is passing + RealtimeProvider( + "bedrock", + "bedrock-realtime", + LiteLLMParamsBody( + model="bedrock/amazon.nova-2-sonic-v1:0", + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + ), + ), # RealtimeProvider( # "xai", # "xai-realtime", diff --git a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py index 622ed9d507f..f5592abb34d 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_e2e.py @@ -29,13 +29,14 @@ from realtime_client import ( JsonSchemaProperty, RealtimeClient, RealtimeProvider, + RealtimeSession, + ReceivedEvent, ResponseCreate, ResponseDone, ServerEnvelope, SessionConfig, SessionUpdate, as_text, - function_call_item, parse_last, realtime_model, transcript, @@ -64,19 +65,26 @@ TEXT_PARAMS: Final = tuple( for p in PROVIDERS ) +BEDROCK_TOOL_COVERS: Final = pytest.mark.covers( + "llm.realtime.bedrock_converse.tool_use.stream.works", exercised_on=["realtime"] +) + TOOL_PARAMS: Final = tuple( pytest.param( p, id=p.id, - marks=meta( - Subject( - domain=Domain.LLM_TRANSLATION, - route=Route.REALTIME, - providers=(Provider(p.id),), - models=(p.litellm_params.model,), - capabilities=(Capability.FUNCTION_CALLING,), - mode=Mode.WEBSOCKET, - ) + marks=( + meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.REALTIME, + providers=(Provider(p.id),), + models=(p.litellm_params.model,), + capabilities=(Capability.FUNCTION_CALLING,), + mode=Mode.WEBSOCKET, + ) + ), + *((BEDROCK_TOOL_COVERS,) if p.id == "bedrock" else ()), ), ) for p in PROVIDERS @@ -164,31 +172,51 @@ def test_tool_call_round_trip( session.send(user_message("What's the weather in Paris right now?")) session.send(ResponseCreate()) - first = session.collect_until("response.done", timeout=60) + final_answer = _tool_call_loop(session, max_rounds=4) - args_event = parse_last( - first, "response.function_call_arguments.done", FunctionCallArgumentsDone + assert "72" in transcript(final_answer), ( + "the answer after the tool results did not use them" ) - assert args_event is not None, "model did not emit a function call" - args = WeatherArgs.model_validate_json(args_event.arguments) - item = function_call_item(first) - assert item is not None, "no completed function_call output item" - assert item.name == "get_weather" - assert item.call_id == args_event.call_id - tool_result = WeatherResult(city=args.city, temperature_f=72) - session.send( - ConversationItemCreate( - item=FunctionCallOutputItem( - call_id=args_event.call_id, output=tool_result.model_dump_json() +def _tool_call_loop(session: RealtimeSession, *, max_rounds: int) -> tuple[ReceivedEvent, ...]: + """Answer every function_call in each response.done until one arrives without any. + + Nova 2 Sonic emits two tool calls per question, each inside its own response.done, + so a one-call round trip cannot wait for "the" follow-up: the second response.done + is another call, not the answer. The loop sends one function_call_output per call + then a single response.create, and returns the events of the first call-free round.""" + for round_index in range(max_rounds): + events = session.collect_until("response.done", timeout=60) + done = parse_last(events, "response.done", ResponseDone) + assert done is not None + calls = tuple( + item for item in done.response.output or () if item.type == "function_call" + ) + if not calls: + assert round_index > 0, "model did not emit a function call" + return events + done_args = tuple( + FunctionCallArgumentsDone.model_validate_json(event.payload) + for event in events + if event.type == "response.function_call_arguments.done" + ) + for item in calls: + assert item.name == "get_weather", item + args_event = next( + event for event in done_args if event.call_id == item.call_id + ) + args = WeatherArgs.model_validate_json(args_event.arguments) + tool_result = WeatherResult(city=args.city, temperature_f=72) + session.send( + ConversationItemCreate( + item=FunctionCallOutputItem( + call_id=args_event.call_id, output=tool_result.model_dump_json() + ) ) ) - ) session.send(ResponseCreate()) - second = session.collect_until("response.done", timeout=60) - - assert "72" in transcript(second), "follow-up did not use the tool result" + raise AssertionError(f"still calling tools after {max_rounds} rounds") _REFUSED_UPSTREAMS = ( diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py index 27d5abd1828..0e4dad7af1a 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_audio_e2e.py @@ -91,6 +91,7 @@ TOOL_PARAMS: Final = tuple( ), ) for p in PROVIDERS + if p.id != "bedrock" ) AUDIO_OUTPUT_PARAMS: Final = tuple( @@ -109,6 +110,7 @@ AUDIO_OUTPUT_PARAMS: Final = tuple( ), ) for p in PROVIDERS + if p.id != "bedrock" ) AUDIO_INPUT_PARAMS: Final = tuple( @@ -127,6 +129,7 @@ AUDIO_INPUT_PARAMS: Final = tuple( ), ) for p in PROVIDERS + if p.id != "bedrock" ) # PCM16 24 kHz mono WAV of "What is the weather in Paris?" (generated via macOS diff --git a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py index 3c22ae1a928..bfb418b32dc 100644 --- a/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py +++ b/tests/e2e/llm_translation/realtime/test_realtime_pipecat_e2e.py @@ -64,6 +64,9 @@ from pipecat_service import LiteLLMRealtimeLLMService # noqa: E402 # Vertex native-audio live is flaky through pipecat tool calling (upstream # pipecat-ai/pipecat#2544); raw-ws tool_call_round_trip[vertex_ai] is the # source of truth for that provider. Keep openai/azure/gemini here. +# Bedrock event shapes are not yet pipecat-compatible (missing GA fields and +# event names, and no tool-result round trip reaches Nova); raw-ws +# tool_call_round_trip[bedrock] is the source of truth for that provider. PROVIDER_PARAMS = [ pytest.param( p, @@ -79,7 +82,7 @@ PROVIDER_PARAMS = [ ) ), ) - for p in PROVIDERS if p.id != "vertex_ai" + for p in PROVIDERS if p.id not in ("vertex_ai", "bedrock") ] WEATHER_TOOL = ToolsSchema( diff --git a/tests/integration/_support/async_server.py b/tests/integration/_support/async_server.py index 324cbaa5da4..53dcbc8d92d 100644 --- a/tests/integration/_support/async_server.py +++ b/tests/integration/_support/async_server.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import ssl from collections.abc import Awaitable, Callable from contextlib import suppress from threading import Thread @@ -16,10 +17,11 @@ async def close_writer(writer: asyncio.StreamWriter) -> None: class LoopbackServer: - def __init__(self, port: int, handler: Handler) -> None: + def __init__(self, port: int, handler: Handler, ssl_context: ssl.SSLContext | None = None) -> None: self.loop: Final = asyncio.new_event_loop() self._port: Final = port self._handler: Final = handler + self._ssl_context: Final = ssl_context self._server: asyncio.Server | None = None self._thread: Final = Thread(target=self._run, name=f"integration-relay-{port}", daemon=True) self._connections: tuple[asyncio.Task[None], ...] = () @@ -42,7 +44,7 @@ class LoopbackServer: finally: self.loop.call_soon_threadsafe(self.loop.stop) self._thread.join(10) - assert not self._thread.is_alive(), f"Database relay thread on port {self._port} did not stop" + assert not self._thread.is_alive(), f"Loopback server thread on port {self._port} did not stop" def _run(self) -> None: asyncio.set_event_loop(self.loop) @@ -53,7 +55,7 @@ class LoopbackServer: self.loop.close() async def _bind(self) -> None: - self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port) + self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port, ssl=self._ssl_context) async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: task: Final = asyncio.current_task() diff --git a/tests/integration/_support/bedrock_realtime_peer.py b/tests/integration/_support/bedrock_realtime_peer.py new file mode 100644 index 00000000000..d30bf1b92f5 --- /dev/null +++ b/tests/integration/_support/bedrock_realtime_peer.py @@ -0,0 +1,429 @@ +"""A scripted Bedrock Nova Sonic realtime peer. + +The handler dials ``api_base`` with ``invoke_model_with_bidirectional_stream`` through +``AWSCRTHTTPClient``, which only speaks h2 over TLS ALPN, so this peer is a loopback TLS +server offering ``h2`` on ``async_server.LoopbackServer``'s thread-owned asyncio loop. +It answers ``POST /model/{id}/invoke-with-bidirectional-stream``, decodes the client's +signed event-stream chunks into the ``{"event": {...}}`` documents Bedrock would read, +and answers each consumed ``Turn`` with event-stream ``chunk`` frames carrying +``{"bytes": }`` payloads, the reverse of what the SDK produces. + +Turn triggers are inbound events. ``user_input`` fires on the ``contentEnd`` closing the +first inbound USER-role TEXT content block, the end of a forwarded conversation.item.create +text message (the injected trigger-audio block is USER/AUDIO, so it never trips it). +``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside +a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration`` +``toolUseId`` names the call being answered. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +import struct +import uuid +from collections.abc import Iterator, Mapping +from contextlib import contextmanager, suppress +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from queue import SimpleQueue +from typing import Final, Literal + +import h2.config +import h2.connection +import h2.events +import h2.exceptions +from integration._support.async_server import Handler, LoopbackServer +from integration._support.process import _free_port +from integration._support.tls import server_context +from integration._support.upstream import aws_event_stream_frame +from pydantic import JsonValue, TypeAdapter + +BedrockEvent = Mapping[str, JsonValue] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _ids() -> tuple[str, str, str, str]: + return (uuid.uuid4().hex, uuid.uuid4().hex, uuid.uuid4().hex, uuid.uuid4().hex) + + +def _block(prompt: str, content: str, session: str, completion: str) -> dict[str, JsonValue]: + return {"promptName": prompt, "contentName": content, "sessionId": session, "completionId": completion} + + +def user_transcript(text: str) -> tuple[BedrockEvent, ...]: + """The USER ASR transcript block Bedrock emits once it has consumed the user turn.""" + prompt, content, session, completion = _ids() + return ( + { + "event": { + "contentStart": { + **_block(prompt, content, session, completion), + "role": "USER", + "type": "TEXT", + "interactive": False, + "textOutputConfiguration": {"mediaType": "text/plain"}, + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {**_block(prompt, content, session, completion), "role": "USER", "content": text}}}, + { + "event": { + "contentEnd": { + **_block(prompt, content, session, completion), + "stopReason": "PARTIAL_TURN", + "type": "TEXT", + } + } + }, + ) + + +def tool_use(tool_use_id: str, name: str, content: str | None) -> tuple[BedrockEvent, ...]: + """The role-TOOL contentStart/toolUse/contentEnd triple Bedrock emits per tool call.""" + prompt, content_name, session, completion = _ids() + tool_use_body: Final[dict[str, JsonValue]] = { + **_block(prompt, content_name, session, completion), + "role": "TOOL", + "toolName": name, + "toolUseId": tool_use_id, + } + if content is not None: + tool_use_body["content"] = content + return ( + { + "event": { + "contentStart": { + **_block(prompt, content_name, session, completion), + "role": "TOOL", + "type": "TOOL", + "toolUseOutputConfiguration": {"mediaType": "application/json"}, + } + } + }, + {"event": {"toolUse": tool_use_body}}, + { + "event": { + "contentEnd": { + **_block(prompt, content_name, session, completion), + "stopReason": "TOOL_USE", + "type": "TOOL", + } + } + }, + ) + + +def assistant_text(text: str, stop_reason: str = "END_TURN") -> tuple[BedrockEvent, ...]: + """An ASSISTANT text block (SPECULATIVE stage) closed with `stop_reason`.""" + prompt, content, session, completion = _ids() + return ( + { + "event": { + "contentStart": { + **_block(prompt, content, session, completion), + "role": "ASSISTANT", + "type": "TEXT", + "interactive": False, + "textOutputConfiguration": {"mediaType": "text/plain"}, + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + { + "event": { + "textOutput": {**_block(prompt, content, session, completion), "role": "ASSISTANT", "content": text} + } + }, + { + "event": { + "contentEnd": { + **_block(prompt, content, session, completion), + "stopReason": stop_reason, + "type": "TEXT", + } + } + }, + ) + + +def usage(input_tokens: int, output_tokens: int) -> BedrockEvent: + total: Final = { + "input": {"speechTokens": 0, "textTokens": input_tokens}, + "output": {"speechTokens": 0, "textTokens": output_tokens}, + } + return { + "event": { + "usageEvent": { + "promptName": uuid.uuid4().hex, + "sessionId": uuid.uuid4().hex, + "completionId": uuid.uuid4().hex, + "details": {"delta": total["output"], "total": total}, + "totalInputTokens": input_tokens, + "totalOutputTokens": output_tokens, + "totalTokens": input_tokens + output_tokens, + } + } + } + + +@dataclass(frozen=True, slots=True) +class Turn: + after: Literal["user_input", "tool_result"] + events: tuple[BedrockEvent, ...] + + +def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]: + body: Final = event["event"] + assert isinstance(body, Mapping) and len(body) == 1, event + name: Final = next(iter(body)) + fields: Final = body[name] + assert isinstance(fields, Mapping), event + return name, fields + + +def _outbound_chunk(event: BedrockEvent) -> bytes: + wrapped: Final = json.dumps({"bytes": base64.b64encode(json.dumps(event).encode()).decode()}).encode() + return aws_event_stream_frame( + {":event-type": "chunk", ":content-type": "application/json", ":message-type": "event"}, wrapped + ) + + +def _parse_frames(buffer: bytes) -> tuple[tuple[bytes, ...], bytes]: + """Slice complete event-stream frames out of `buffer`; return (payloads, remainder).""" + payloads: Final[list[bytes]] = [] + view = memoryview(buffer) + while len(view) >= 12: + total, header_length = struct.unpack_from("!II", view, 0) + if total > len(view): + break + payloads.append(view[12 + header_length : total - 4].tobytes()) + view = view[total:] + return tuple(payloads), view.tobytes() + + +def _frame_payload(raw: bytes) -> bytes: + """The payload slice of one complete framed message (prelude + headers + payload + CRC).""" + if len(raw) < 16: + return b"" + _total, header_length = struct.unpack_from("!II", raw, 0) + return raw[12 + header_length : -4] + + +def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None: + """One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event.""" + inner_payload: Final = _frame_payload(chunk_payload) + try: + wrapped: Final = _JSON_OBJECT.validate_json(inner_payload) + except ValueError: + return None + encoded: Final = wrapped.get("bytes") + if not isinstance(encoded, str): + return None + try: + return _JSON_OBJECT.validate_json(base64.b64decode(encoded)) + except ValueError: + return None + + +class _Inbound: + """Inbound event bookkeeping for turn triggers and accessors.""" + + def __init__(self) -> None: + self.events: Final = SimpleQueue[BedrockEvent]() + self.open_blocks: Final[dict[str, tuple[str, str]]] = {} + self.tool_use_ids: Final[dict[str, str]] = {} + self.tool_results_seen = 0 + self.user_text_closed = False + + def record(self, event: BedrockEvent) -> None: + self.events.put(event) + name, fields = _event_body(event) + content_name = fields.get("contentName") + if name == "contentStart" and isinstance(content_name, str): + self.open_blocks[content_name] = (str(fields.get("role")), str(fields.get("type"))) + configuration = fields.get("toolResultInputConfiguration") + if isinstance(configuration, Mapping) and isinstance(configuration.get("toolUseId"), str): + self.tool_use_ids[content_name] = configuration["toolUseId"] + elif name == "contentEnd" and isinstance(content_name, str): + role_type: Final = self.open_blocks.pop(content_name, ("", "")) + if role_type == ("USER", "TEXT"): + self.user_text_closed = True + elif name == "toolResult": + self.tool_results_seen += 1 + + +class _Script: + """Consumes turns in order as their inbound triggers arrive.""" + + def __init__(self, turns: tuple[Turn, ...]) -> None: + self._turns: Final = turns + self._next = 0 + + @property + def exhausted(self) -> bool: + return self._next >= len(self._turns) + + def ready_events(self, state: _Inbound, event: BedrockEvent) -> tuple[BedrockEvent, ...]: + if self._next >= len(self._turns): + return () + turn: Final = self._turns[self._next] + if turn.after == "user_input": + if not state.user_text_closed: + return () + else: + name, _fields = _event_body(event) + if name != "toolResult": + return () + self._next += 1 + return turn.events + + +class _PeerState: + def __init__(self, turns: tuple[Turn, ...]) -> None: + self.inbound: Final = _Inbound() + self.script: Final = _Script(turns) + + +@dataclass(frozen=True, slots=True) +class BedrockRealtimePeer: + api_base: str + _state: _PeerState + _drained: list[BedrockEvent] + + def received(self) -> tuple[BedrockEvent, ...]: + while not self._state.inbound.events.empty(): + self._drained.append(self._state.inbound.events.get()) # drained log, not a built value + return tuple(self._drained) + + def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]: + return tuple(chain.from_iterable(_tool_specs(event) for event in self.received())) + + def tool_results(self) -> tuple[tuple[str, str], ...]: + """(toolUseId, content) per inbound toolResult in arrival order.""" + return tuple(result for event in self.received() if (result := self._tool_result(event)) is not None) + + def _tool_result(self, event: BedrockEvent) -> tuple[str, str] | None: + name, fields = _event_body(event) + if name != "toolResult": + return None + content_name: Final = fields.get("contentName") + content: Final = fields.get("content") + if not isinstance(content_name, str) or not isinstance(content, str): + return None + return (self._state.inbound.tool_use_ids.get(content_name, ""), content) + + +def _tool_specs(event: BedrockEvent) -> tuple[Mapping[str, JsonValue], ...]: + name, fields = _event_body(event) + if name != "promptStart": + return () + configuration: Final = fields.get("toolConfiguration") + if not isinstance(configuration, Mapping): + return () + tools: Final = configuration.get("tools") + if not isinstance(tools, list): + return () + return tuple( + entry["toolSpec"] + for entry in tools + if isinstance(entry, Mapping) and isinstance(entry.get("toolSpec"), Mapping) + ) + + +class _H2Session: + """One bidirectional-stream connection: decode inbound, answer scripted turns.""" + + def __init__(self, state: _PeerState, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + self._state: Final = state + self._reader: Final = reader + self._writer: Final = writer + self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False)) + self._buffers: Final[dict[int, bytes]] = {} + self._ended: set[int] = set() + self._closed = False + + async def run(self) -> None: + self._h2.initiate_connection() + self._writer.write(self._h2.data_to_send()) + await self._writer.drain() + while True: + incoming: Final = await self._reader.read(65536) + if not incoming: + return + for event in self._h2.receive_data(incoming): + self._apply(event) + pending: Final = self._h2.data_to_send() + if pending: + self._writer.write(pending) + await self._writer.drain() + if self._closed: + self._writer.close() + return + + def _apply(self, event: h2.events.Event) -> None: + if isinstance(event, h2.events.RequestReceived): + self._h2.send_headers( + event.stream_id, + [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], + ) + elif isinstance(event, h2.events.DataReceived): + self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id) + self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data + payloads, remainder = _parse_frames(self._buffers[event.stream_id]) + self._buffers[event.stream_id] = remainder + for payload in payloads: + decoded: Final = _decode_inbound(payload) + if decoded is None: + continue + self._state.inbound.record(decoded) + if event.stream_id not in self._ended: + for outbound in self._state.script.ready_events(self._state.inbound, decoded): + self._h2.send_data(event.stream_id, _outbound_chunk(outbound)) + elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)): + self._close_stream(event.stream_id) + elif isinstance(event, h2.events.ConnectionTerminated): + self._writer.close() + + def _close_stream(self, stream_id: int) -> None: + self._end_stream(stream_id) + self._closed = True + self._h2.close_connection() + + def _end_stream(self, stream_id: int) -> None: + if stream_id in self._ended: + return + self._ended.add(stream_id) + with suppress(h2.exceptions.ProtocolError): + self._h2.end_stream(stream_id) + + +def _session_handler(state: _PeerState) -> Handler: + async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + with suppress(ConnectionError, OSError, asyncio.IncompleteReadError, h2.exceptions.StreamClosedError): + await _H2Session(state, reader, writer).run() + + return serve + + +@contextmanager +def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]: + """Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer. + + `cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE. + SigV4 headers are ignored; the stream closes cleanly when the client ends it. + """ + port: Final = _free_port() + context: Final = server_context(cert[0], cert[1]) + context.set_alpn_protocols(["h2"]) + state: Final = _PeerState(turns) + server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context) + server.start() + try: + yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[]) + finally: + server.stop() diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py new file mode 100644 index 00000000000..ccef839a163 --- /dev/null +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -0,0 +1,378 @@ +"""Bedrock Nova Sonic realtime tool calls through a scripted h2 event-stream peer (LIT-9465). + +Nova Sonic returns a tool call as a role-TOOL contentStart/toolUse/contentEnd triple with the +arguments as a JSON string in ``content`` and no ASSISTANT block before it, and emits one such +triple per call including the two-call turn Nova 2 Sonic produces per question. The cells here +serve those triples from ``_support/bedrock_realtime_peer.py`` (the transport the handler really +dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream) and assert both +sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult +documents the proxy forwarded to the peer. + +The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE`` +pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it. +""" + +from __future__ import annotations + +import asyncio +import json +import time +import uuid +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Final + +import pytest +import websockets +from integration._support.bedrock_realtime_peer import ( + BedrockRealtimePeer, + Turn, + assistant_text, + bedrock_realtime_peer, + tool_use, + usage, + user_transcript, +) +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import owned_proxy +from integration._support.tls import write_self_signed_cert +from pydantic import JsonValue +from websockets.asyncio.client import ClientConnection + +pytestmark: Final = pytest.mark.timeout(180) + +DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" +WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} +WEATHER_FUNCTION: Final = { + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": WEATHER_PARAMETERS, +} +NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} +FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} +TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}' + + +@pytest.fixture(scope="module") +def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: + return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert")) + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + with gateway_from_environment() as base: + with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + yield proxy + + +def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: + name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}" + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": DEPLOYMENT, + "aws_region_name": "us-east-1", + "aws_access_key_id": "integration-dummy-key", + "aws_secret_access_key": "integration-dummy-secret", + "api_base": peer.api_base, + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return name + + +def _session_update(tool: Mapping[str, JsonValue]) -> str: + return json.dumps( + { + "type": "session.update", + "session": { + "instructions": "Always use the get_weather tool for weather questions.", + "tools": [tool], + }, + } + ) + + +def _user_text() -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "What is the weather in New York right now?"}], + }, + } + ) + + +_RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) + + +def _tool_result(call_id: str) -> str: + return json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "function_call_output", "call_id": call_id, "output": TOOL_RESULT_OUTPUT}, + } + ) + + +Event = dict[str, JsonValue] +EventPredicate = Callable[[Event, tuple[Event, ...]], bool] + + +async def _events_until( + socket: ClientConnection, predicate: EventPredicate, *, timeout: float = 60 +) -> tuple[Event, ...]: + collected: Final[list[Event]] = [] + deadline: Final = time.monotonic() + timeout + while True: + try: + raw: Final = await asyncio.wait_for(socket.recv(), deadline - time.monotonic()) + except (TimeoutError, asyncio.TimeoutError): + raise AssertionError(f"timed out waiting; types={[event.get('type') for event in collected]}") + event: Final = JSON_OBJECT.validate_json(raw) + collected.append(event) + if predicate(event, tuple(collected)): + return tuple(collected) + + +def _is_done(events: tuple[Event, ...]) -> bool: + return any(event.get("type") == "response.done" for event in events) + + +def _function_calls(events: tuple[Event, ...]) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == "response.function_call_arguments.done") + + +def _response_dones(events: tuple[Event, ...]) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == "response.done") + + +def _done_ids(events: tuple[Event, ...]) -> tuple[str, ...]: + return tuple(string_value(object_value(event["response"])["id"]) for event in _response_dones(events)) + + +def _tail_types(events: tuple[Event, ...]) -> tuple[str, ...]: + """Event types from the first response.created onward.""" + start: Final = next(index for index, event in enumerate(events) if event.get("type") == "response.created") + return tuple(string_value(event["type"]) for event in events[start:]) + + +@asynccontextmanager +async def _realtime(gateway: Gateway, model: str, key: str) -> AsyncIterator[ClientConnection]: + proxy: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + async with websockets.connect( + f"{proxy}/v1/realtime?model={model}", additional_headers={"Authorization": f"Bearer {key}"} + ) as socket: + yield socket + + +async def _open_session(socket: ClientConnection, tool: Mapping[str, JsonValue]) -> None: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + assert created.get("type") == "session.created", created + await socket.send(_session_update(tool)) + await socket.send(_user_text()) + await socket.send(_RESPONSE_CREATE) + + +async def _one_call_session(gateway: Gateway, model: str, key: str, tool: Mapping[str, JsonValue]) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + await _open_session(socket, tool) + return await _events_until(socket, lambda _event, events: _is_done(events)) + + +def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn( + "user_input", user_transcript("ready") + tool_use("tooluse_A", "get_weather", '{"location":"New York"}') + ) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) + assert _tail_types(events) == ( + "response.created", + "response.output_item.added", + "response.function_call_arguments.done", + "response.output_item.done", + "response.done", + ), events + (call,) = _function_calls(events) + assert string_value(call["call_id"]) == "tooluse_A" + assert string_value(call["name"]) == "get_weather" + assert string_value(call["arguments"]) == '{"location":"New York"}' + (done,) = _response_dones(events) + output: Final = object_value(done["response"])["output"] + assert isinstance(output, list) and len(output) == 1, done + item: Final = object_value(output[0]) + assert string_value(item["type"]) == "function_call" + assert string_value(item["call_id"]) == "tooluse_A" + assert string_value(item["arguments"]) == '{"location":"New York"}' + specs: Final = eventually(lambda: peer.tool_specs(), lambda values: len(values) == 1, seconds=10) + assert string_value(specs[0]["name"]) == "get_weather" + assert json.loads(string_value(object_value(specs[0]["inputSchema"])["json"])) == WEATHER_PARAMETERS + + +def test_flat_tool_shape(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turn: Final = Turn( + "user_input", user_transcript("ready") + tool_use("tooluse_B", "get_weather", '{"location":"New York"}') + ) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, FLAT_TOOL)) + (call,) = _function_calls(events) + assert string_value(call["call_id"]) == "tooluse_B" + assert string_value(call["name"]) == "get_weather" + specs: Final = eventually(lambda: peer.tool_specs(), lambda values: len(values) == 1, seconds=10) + assert string_value(specs[0]["name"]) == "get_weather" + assert string_value(specs[0]["description"]) == "Get current weather for a city" + assert json.loads(string_value(object_value(specs[0]["inputSchema"])["json"])) == WEATHER_PARAMETERS + + +async def _tool_result_round_trip(gateway: Gateway, model: str, key: str) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + await _open_session(socket, NESTED_TOOL) + first: Final = await _events_until(socket, lambda _event, events: _is_done(events)) + (call,) = _function_calls(first) + await socket.send(_tool_result(string_value(call["call_id"]))) + await socket.send(_RESPONSE_CREATE) + second: Final = await _events_until(socket, lambda _event, events: _is_done(events)) + return first + second + + +def test_tool_result_round_trip(candidate: Gateway, cert: tuple[Path, Path]) -> None: + turns: Final = ( + Turn("user_input", user_transcript("ready") + tool_use("tooluse_A", "get_weather", '{"location":"New York"}')), + Turn("tool_result", assistant_text("It is 72 and sunny") + (usage(30, 12),)), + ) + with bedrock_realtime_peer(turns, cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_tool_result_round_trip(candidate, model, key)) + dones: Final = _response_dones(events) + assert len(dones) == 2, events + first_id, second_id = _done_ids(events) + assert first_id != second_id + second_created: Final = tuple(event for event in events if event.get("type") == "response.created") + assert len(second_created) == 2 and object_value(second_created[1]["response"])["id"] == second_id + text: Final = "".join( + string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta" + ) + assert "72" in text + last_output: Final = object_value(dones[-1]["response"])["output"] + assert tuple(object_value(item).get("type") for item in last_output) == ("message",) + (answer_item,) = last_output + assert object_value(answer_item)["content"] == [{"type": "text", "text": "It is 72 and sunny"}] + results: Final = eventually(lambda: peer.tool_results(), lambda values: len(values) == 1, seconds=10) + assert results == (("tooluse_A", TOOL_RESULT_OUTPUT),) + + +async def _two_call_session(gateway: Gateway, model: str, key: str) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + await _open_session(socket, NESTED_TOOL) + events: Final[tuple[Event, ...]] = await _events_until( + socket, lambda _event, collected: len(_response_dones(collected)) == 2 + ) + for call in _function_calls(events): + await socket.send(_tool_result(string_value(call["call_id"]))) + after_answer: Final = await _events_until(socket, lambda _event, collected: _is_done(collected)) + return events + after_answer + + +def test_two_tool_calls_one_turn(candidate: Gateway, cert: tuple[Path, Path]) -> None: + args: Final = '{"location":"New York"}' + turns: Final = ( + Turn( + "user_input", + user_transcript("ready") + + tool_use("tooluse_A", "get_weather", args) + + tool_use("tooluse_B", "get_weather", args), + ), + Turn("tool_result", assistant_text("It is 72 and sunny") + (usage(40, 15),)), + ) + with bedrock_realtime_peer(turns, cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_two_call_session(candidate, model, key)) + calls: Final = _function_calls(events) + assert len(calls) == 2, events + assert tuple(string_value(call["call_id"]) for call in calls) == ("tooluse_A", "tooluse_B") + assert tuple(string_value(call["arguments"]) for call in calls) == (args, args) + item_ids: Final = tuple(string_value(call["item_id"]) for call in calls) + response_ids: Final = tuple(string_value(call["response_id"]) for call in calls) + assert len(set(item_ids)) == 2 and len(set(response_ids)) == 2 + first_dones: Final = _response_dones(events)[:2] + for done, call_id in zip(first_dones, ("tooluse_A", "tooluse_B")): + output: Final = object_value(done["response"])["output"] + assert isinstance(output, list) and len(output) == 1, done + assert string_value(object_value(output[0])["call_id"]) == call_id + results: Final = eventually(lambda: peer.tool_results(), lambda values: len(values) == 2, seconds=10) + assert tuple(result[0] for result in results) == ("tooluse_A", "tooluse_B") + assert all(result[1] == TOOL_RESULT_OUTPUT for result in results) + text: Final = "".join( + string_value(event["delta"]) for event in events if event.get("type") == "response.text.delta" + ) + assert "72" in text + + +def test_tool_call_inside_active_assistant_response(candidate: Gateway, cert: tuple[Path, Path]) -> None: + start, output, _end = assistant_text("It is 72 and sunny") + turn: Final = Turn( + "user_input", + (start, output) + tool_use("tooluse_C", "get_weather", '{"location":"New York"}') + (_end,), + ) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) + created: Final = tuple(event for event in events if event.get("type") == "response.created") + assert len(created) == 1, events + response_id: Final = string_value(object_value(created[0]["response"])["id"]) + (call,) = _function_calls(events) + assert string_value(call["call_id"]) == "tooluse_C" + assert string_value(call["response_id"]) == response_id + assert string_value(call["arguments"]) == '{"location":"New York"}' + assert call["output_index"] == 1, events + assert tuple(string_value(e["text"]) for e in events if e.get("type") == "response.text.done") == ( + "It is 72 and sunny", + ), events + message_done_indexes: Final = tuple( + i + for i, event in enumerate(events) + if event.get("type") == "response.output_item.done" and object_value(event["item"]).get("type") == "message" + ) + response_done_indexes: Final = tuple( + i for i, event in enumerate(events) if event.get("type") == "response.done" + ) + assert len(response_done_indexes) == 1, events + assert len(message_done_indexes) == 1 and response_done_indexes[0] > message_done_indexes[0], events + + +@pytest.mark.parametrize( + "content", + ["", None, "not json", '{"location":"' + "x" * 5000 + '"}'], + ids=["empty", "missing", "not-json", "five-kb"], +) +def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], content: str | None) -> None: + turn: Final = Turn("user_input", tool_use("tooluse_M", "get_weather", content)) + with bedrock_realtime_peer((turn,), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + key: Final = scenario.key() + events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) + (call,) = _function_calls(events) + assert string_value(call["arguments"]) == (content if content is not None else "{}"), events diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index b1836cb6c46..8f95a502137 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -3,6 +3,7 @@ import json import sys import types from types import SimpleNamespace +from typing import Final from unittest.mock import MagicMock import pytest @@ -68,6 +69,20 @@ class FailOnPromptEndStream(FakeInputStream): self.sent.append(event) +class HangingCloseInputStream(FakeInputStream): + """Models a CRT input stream on an already-ended h2 stream: close sends and close() never complete""" + + async def send(self, event): + payload = json.loads(event.value.bytes_.decode("utf-8")) + if "promptEnd" in payload.get("event", {}) or "sessionEnd" in payload.get("event", {}): + await asyncio.Future() + else: + self.sent.append(event) + + async def close(self): + await asyncio.Future() + + class FakeBedrockStream: def __init__(self, input_stream=None): self.input_stream = input_stream if input_stream is not None else FakeInputStream() @@ -417,6 +432,21 @@ class TestBedrockRealtimeHandler: assert stream.input_stream.closed + @pytest.mark.asyncio + async def test_close_flush_is_bounded_when_input_stream_never_completes(self, stub_aws_models): + handler = BedrockRealtime(close_timeout=0.05) + config = BedrockRealtimeConfig() + stream = FakeBedrockStream(input_stream=HangingCloseInputStream()) + client_ws = DisconnectingClientWS( + [json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})] + ) + + await asyncio.wait_for( + handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}), 2 + ) + + assert not stream.input_stream.closed + @pytest.mark.asyncio async def test_close_flush_continues_after_partial_send_failure(self, stub_aws_models): handler = BedrockRealtime() @@ -613,6 +643,78 @@ class TestBedrockRealtimeSessionLifecycle: assert client_ws.sent_to_client == [] + @pytest.mark.asyncio + async def test_conflicting_session_update_gets_error_and_sends_nothing_to_bedrock(self, stub_aws_models): + tool: Final = {"type": "function", "name": "get_weather", "parameters": {"type": "object"}} + other_tool: Final = {"type": "function", "name": "get_time", "parameters": {"type": "object"}} + + def update(tools) -> str: + return json.dumps( + { + "type": "session.update", + "event_id": "evt_update_1", + "session": {"tools": tools, "modalities": ["text"]}, + } + ) + + handler = BedrockRealtime() + config = BedrockRealtimeConfig() + stream = FakeBedrockStream() + client_ws = DisconnectingClientWS([update([tool]), update([other_tool]), update([tool])]) + + await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging()) + + acked: Final = [json.loads(message) for message in client_ws.sent_to_client] + assert [event["type"] for event in acked] == ["session.updated", "error", "session.updated"] + error: Final = acked[1]["error"] + assert error["code"] == "session_update_not_applied" + assert error["event_id"] == "evt_update_1" + prompt_starts: Final = [ + payload + for payload in (json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent) + if "promptStart" in payload.get("event", {}) + ] + assert len(prompt_starts) == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "drop_params,expected_errors", + [(False, [("invalid_type", "session.temperature")]), (True, [])], + ) + async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers( + self, stub_aws_sdk_client, drop_params, expected_errors + ): + stream: Final = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [stream] + user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} + websocket: Final = ConnectedClientWS( + [ + json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}), + json.dumps({"type": "conversation.item.create", "item": user_text}), + ] + ) + + await BedrockRealtime().async_realtime( + model="amazon.nova-sonic-v1:0", + websocket=websocket, + logging_obj=FakeLogging(), + aws_region_name="us-east-1", + aws_access_key_id="k", + aws_secret_access_key="s", + drop_params=drop_params, + ) + + client_events: Final = [json.loads(message) for message in websocket.sent_to_client] + errors: Final = [event["error"] for event in client_events if event["type"] == "error"] + sent_to_bedrock: Final = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent] + session_starts: Final = [ + payload["event"]["sessionStart"] for payload in sent_to_bedrock if "sessionStart" in payload["event"] + ] + assert [(error["code"], error["param"]) for error in errors] == expected_errors + assert [start["inferenceConfiguration"]["temperature"] for start in session_starts] == [0.7] + assert next(iter(sent_to_bedrock[0]["event"])) == "sessionStart" + assert any(payload["event"].get("textInput", {}).get("content") == "hi" for payload in sent_to_bedrock) + class TestBedrockRealtimeProviderFailurePropagation: """Deferred Nova Sonic failures must escape async_realtime so the router can fall back / cool down (LIT-6484)""" diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 9d034c9838c..1cc135a8671 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -15,6 +15,9 @@ from litellm.llms.bedrock.realtime.transformation import ( from litellm.llms.bedrock.realtime.trigger_audio import ready_trigger_pcm from litellm.types.llms.openai import OpenAIRealtimeEventTypes +OBJECT_PARAMETERS: Final = {"type": "object"} +CONFLICT_TOOL: Final = {"type": "function", "name": "get_weather", "parameters": OBJECT_PARAMETERS} + class TestBedrockRealtimeConfig: """Test suite for BedrockRealtimeConfig class""" @@ -110,6 +113,50 @@ class TestBedrockRealtimeConfig: assert schema["type"] == "object" assert "location" in schema["properties"] + def test_transform_tools_flat_realtime_shape(self): + """The flat OpenAI Realtime tool shape {type: function, name, parameters} + produces the same toolSpec as the nested chat-completions shape (LIT-9465)""" + config = BedrockRealtimeConfig() + + nested = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + }, + } + ] + flat = [ + { + "type": "function", + "name": "get_weather", + "description": "Get current weather for a city", + "parameters": { + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + } + ] + + assert config._transform_tools_to_bedrock_format(flat) == config._transform_tools_to_bedrock_format(nested) + + flat_tools = config._transform_tools_to_bedrock_format(flat) + assert flat_tools[0]["toolSpec"]["name"] == "get_weather" + assert flat_tools[0]["toolSpec"]["description"] == "Get current weather for a city" + schema = json.loads(flat_tools[0]["toolSpec"]["inputSchema"]["json"]) + assert schema["required"] == ["location"] + + session_config = config.session_configuration_request("amazon.nova-2-sonic-v1:0", tools=flat) + prompt_start = json.loads(session_config)["prompt_start"]["event"]["promptStart"] + assert prompt_start["toolConfiguration"]["tools"][0]["toolSpec"]["name"] == "get_weather" + def test_audio_format_mapping(self): """Test audio format to sample rate mapping""" config = BedrockRealtimeConfig() @@ -149,6 +196,49 @@ class TestBedrockRealtimeConfig: session_start = json.loads(messages[0]) assert session_start["event"]["sessionStart"]["inferenceConfiguration"]["temperature"] == 0.9 + def test_session_update_without_instructions_sends_empty_system_prompt(self): + config = BedrockRealtimeConfig() + session_update = { + "type": "session.update", + "session": { + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather", + "parameters": {"type": "object", "properties": {}}, + }, + } + ] + }, + } + + messages = config.transform_session_update_event(session_update) + event_names = [next(iter(json.loads(message)["event"])) for message in messages] + assert event_names[:3] == ["sessionStart", "promptStart", "contentStart"] + system_content_start = json.loads(messages[2])["event"]["contentStart"] + assert system_content_start["role"] == "SYSTEM" + text_input = json.loads(messages[3])["event"]["textInput"] + assert text_input["content"] == "" + assert json.loads(messages[4])["event"]["contentEnd"]["contentName"] == system_content_start["contentName"] + + def test_session_update_instructions_become_system_prompt(self): + config = BedrockRealtimeConfig() + session_update = {"type": "session.update", "session": {"instructions": "Answer tersely."}} + + messages = config.transform_session_update_event(session_update) + text_input = json.loads(messages[3])["event"]["textInput"] + assert text_input["content"] == "Answer tersely." + + def test_second_session_update_sends_no_duplicate_prompt_start(self): + config = BedrockRealtimeConfig() + session_update = {"type": "session.update", "session": {"instructions": "You are helpful."}} + + first = config.transform_session_update_event(session_update) + assert any("promptStart" in message for message in first) + assert config.transform_session_update_event(session_update) == [] + def test_transform_session_update_with_tools(self): """Test session.update with tools""" config = BedrockRealtimeConfig() @@ -274,12 +364,14 @@ class TestBedrockRealtimeResponseCreate: "amazon.nova-sonic-v1:0", ) - def test_response_create_before_session_update_is_noop(self): + def test_response_create_before_session_update_starts_prompt_with_defaults(self): config = BedrockRealtimeConfig() messages = config.transform_realtime_request(json.dumps({"type": "response.create"}), "amazon.nova-sonic-v1:0") - assert messages == [] + events = [next(iter(json.loads(message)["event"])) for message in messages] + assert events[:5] == ["sessionStart", "promptStart", "contentStart", "textInput", "contentEnd"] + assert "audioInput" in events[5:] def test_response_create_emits_spoken_trigger_audio(self): config = BedrockRealtimeConfig() @@ -528,7 +620,9 @@ class TestBedrockRealtimeResponseTransformation: assert audio_deltas[0]["delta"] == "base64_audio_content" def test_transform_tool_use_response(self): - """Test toolUse response transformation""" + """toolUse is forwarded with its arguments even when no assistant response is + active yet (Bedrock Nova Sonic sends the call inside a role TOOL block before + any ASSISTANT contentStart, LIT-9465)""" config = BedrockRealtimeConfig() logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -536,9 +630,77 @@ class TestBedrockRealtimeResponseTransformation: tool_use_message = { "event": { "toolUse": { - "toolUseId": "tool_call_123", + "completionId": "comp_1", + "content": '{"location":"New York"}', + "contentId": "content_1", + "promptName": "prompt_1", + "role": "TOOL", + "sessionId": "sess_1", "toolName": "get_weather", - "input": json.dumps({"location": "San Francisco"}), + "toolUseId": "tool_call_123", + } + } + } + + result = config.transform_realtime_response( + json.dumps(tool_use_message), + "amazon.nova-sonic-v1:0", + logging_obj, + realtime_response_transform_input={ + "session_configuration_request": json.dumps({"configured": True}), + "current_output_item_id": "item_message_1", + "current_response_id": None, + "current_conversation_id": None, + "current_delta_chunks": [], + "current_item_chunks": [], + "current_delta_type": None, + }, + ) + + assert [e["type"] for e in result["response"]] == [ + "response.created", + "response.output_item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + ] + response_created = result["response"][0] + assert response_created["response"]["id"] == result["current_response_id"] + + arguments_delta = result["response"][2] + assert arguments_delta["delta"] == '{"location":"New York"}' + assert arguments_delta["call_id"] == "tool_call_123" + assert arguments_delta["item_id"] == result["response"][1]["item"]["id"] + assert arguments_delta["response_id"] == result["current_response_id"] + + function_call = result["response"][3] + assert function_call["call_id"] == "tool_call_123" + assert function_call["name"] == "get_weather" + assert function_call["arguments"] == '{"location":"New York"}' + assert function_call["response_id"] == result["current_response_id"] + assert function_call["item_id"] != "item_message_1" + assert function_call["item_id"] == result["response"][1]["item"]["id"] + assert result["response"][1]["item"]["type"] == "function_call" + + done_item = result["response"][4]["item"] + assert done_item["type"] == "function_call" + assert done_item["status"] == "completed" + assert done_item["call_id"] == "tool_call_123" + assert done_item["arguments"] == '{"location":"New York"}' + + def test_tool_use_reuses_active_response(self): + """A toolUse arriving while an assistant response is in flight must not emit a + second response.created and must reuse the active response id""" + config = BedrockRealtimeConfig() + logging_obj = MagicMock() + logging_obj.litellm_trace_id = "trace_123" + + tool_use_message = { + "event": { + "toolUse": { + "content": '{"location":"New York"}', + "toolName": "get_weather", + "toolUseId": "tool_call_123", } } } @@ -558,20 +720,19 @@ class TestBedrockRealtimeResponseTransformation: }, ) - # Check for function call event - assert len(result["response"]) == 1 - function_call = result["response"][0] - assert function_call["type"] == "response.function_call_arguments.done" - assert function_call["call_id"] == "tool_call_123" - assert function_call["name"] == "get_weather" - - # Verify arguments are properly formatted - args = json.loads(function_call["arguments"]) - assert args["location"] == "San Francisco" + assert [e["type"] for e in result["response"]] == [ + "response.output_item.added", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + "response.output_item.done", + ] + assert result["current_response_id"] == "resp_123" + assert all(e.get("response_id") == "resp_123" for e in result["response"]) def test_transform_content_end_text(self): """Test contentEnd for text response""" config = BedrockRealtimeConfig() + config._assistant_block_open = True logging_obj = MagicMock() logging_obj.litellm_trace_id = "trace_123" @@ -1211,6 +1372,395 @@ class TestBedrockRealtimeUserEventsAndUsage: assert types.count("response.created") == 1 assert types.count("response.done") == 1 + def test_nova2_sonic_tool_call_sequence_forwards_each_call_with_arguments(self): + """The real Nova 2 Sonic event sequence for one question: a USER transcript + block, two role TOOL blocks each carrying a toolUse with a distinct toolUseId, + then the post-tool-result ASSISTANT answer (LIT-9465)""" + args_json: Final = '{"location":"New York"}' + events: Final = self._run( + BedrockRealtimeConfig(), + [ + { + "event": { + "contentStart": { + "role": "USER", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "FINAL"}), + } + } + }, + {"event": {"textOutput": {"content": "ready"}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN", "type": "TEXT"}}}, + ] + + self._tool_block("tooluse_A", args_json) + + self._tool_block("tooluse_B", args_json) + + self._assistant_text_block( + "SPECULATIVE", "The current weather in New York is sunny and 72 degrees.", "END_TURN" + ), + ) + + function_calls: Final = [e for e in events if e["type"] == "response.function_call_arguments.done"] + assert [e["call_id"] for e in function_calls] == ["tooluse_A", "tooluse_B"] + assert all(e["name"] == "get_weather" for e in function_calls) + assert all(e["arguments"] == args_json for e in function_calls) + assert function_calls[0]["item_id"] != function_calls[1]["item_id"] + + response_dones: Final = [e for e in events if e["type"] == "response.done"] + assert len(response_dones) == 3 + tool_call_dones: Final = [ + d for d in response_dones if any(o["type"] == "function_call" for o in d["response"]["output"]) + ] + assert len(tool_call_dones) == 2 + assert tool_call_dones[0]["response"]["output"][0]["call_id"] == "tooluse_A" + assert tool_call_dones[1]["response"]["output"][0]["call_id"] == "tooluse_B" + for done, call in zip(tool_call_dones, function_calls): + assert len(done["response"]["output"]) == 1 + assert done["response"]["output"][0]["type"] == "function_call" + assert done["response"]["output"][0]["status"] == "completed" + assert done["response"]["id"] == call["response_id"] + final_output: Final = response_dones[2]["response"]["output"] + assert [o["type"] for o in final_output] == ["message"] + assert final_output[0]["status"] == "completed" + + response_ids: Final = {d["response"]["id"] for d in response_dones} + assert len(response_ids) == 3 + + @staticmethod + def _tool_block(tool_use_id: str, arguments: str = '{"city":"Paris"}') -> list[dict[str, object]]: + return [ + {"event": {"contentStart": {"role": "TOOL", "type": "TOOL", "contentId": f"c_{tool_use_id}"}}}, + { + "event": { + "toolUse": { + "content": arguments, + "contentId": f"c_{tool_use_id}", + "role": "TOOL", + "toolName": "get_weather", + "toolUseId": tool_use_id, + } + } + }, + {"event": {"contentEnd": {"stopReason": "TOOL_USE", "type": "TOOL", "contentId": f"c_{tool_use_id}"}}}, + ] + + @staticmethod + def _item_id_of(event) -> str | None: + if event["type"].startswith("response.output_item."): + return event["item"]["id"] + return event.get("item_id") + + def test_tool_call_after_assistant_message_uses_output_index_one(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "Let me check.", "PARTIAL_TURN") + self._tool_block("tooluse_A"), + ) + types: Final = [e["type"] for e in events] + assert types.count("response.created") == 1 + + added_items: Final = [e["item"]["id"] for e in events if e["type"] == "response.output_item.added"] + message_id, call_id = added_items + + message_events: Final = [e for e in events if self._item_id_of(e) == message_id] + call_events: Final = [e for e in events if self._item_id_of(e) == call_id] + assert message_events and all(e["output_index"] == 0 for e in message_events) + assert call_events and all(e["output_index"] == 1 for e in call_events) + + response_dones: Final = [e for e in events if e["type"] == "response.done"] + assert len(response_dones) == 1 + assert [o["id"] for o in response_dones[0]["response"]["output"]] == [message_id, call_id] + + def test_tool_call_inside_open_assistant_block_defers_response_done(self): + config: Final = BedrockRealtimeConfig() + events: Final = self._run( + config, + [ + { + "event": { + "contentStart": { + "role": "ASSISTANT", + "type": "TEXT", + "additionalModelFields": json.dumps({"generationStage": "SPECULATIVE"}), + } + } + }, + {"event": {"textOutput": {"content": "Let me check."}}}, + ] + + self._tool_block("tooluse_A") + + [ + {"event": {"textOutput": {"content": " It is 72."}}}, + {"event": {"contentEnd": {"stopReason": "PARTIAL_TURN"}}}, + ], + ) + types: Final = [e["type"] for e in events] + assert types.count("response.done") == 1 + assert types[-1] == "response.done" + assert [e["text"] for e in events if e["type"] == "response.text.done"] == ["Let me check. It is 72."] + + next_events: Final = self._run(config, self._assistant_text_block("SPECULATIVE", "Sunny.", "END_TURN")) + assert [e["text"] for e in next_events if e["type"] == "response.text.done"] == ["Sunny."] + assert all(e["output_index"] == 0 for e in next_events if "output_index" in e) + assert [e["response"]["id"] for e in next_events if e["type"] == "response.created"] != [ + e["response"]["id"] for e in events if e["type"] == "response.created" + ] + + def test_response_done_message_item_carries_the_text_part(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "It is 72.", "END_TURN"), + ) + (done,) = [e for e in events if e["type"] == "response.done"] + (message_item,) = [o for o in done["response"]["output"] if o["type"] == "message"] + assert message_item["content"] == [{"type": "text", "text": "It is 72."}] + + def test_response_done_message_item_merges_text_then_audio_parts_in_order(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "Hi.", "PARTIAL_TURN") + self._assistant_audio_block("END_TURN"), + ) + (done,) = [e for e in events if e["type"] == "response.done"] + assert [o["type"] for o in done["response"]["output"]] == ["message"] + assert done["response"]["output"][0]["content"] == [ + {"type": "text", "text": "Hi."}, + {"type": "audio", "transcript": ""}, + ] + + def test_session_update_rejection_before_prompt_start(self): + assert ( + BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": {"tools": [CONFLICT_TOOL]}} + ) + is None + ) + + @pytest.mark.parametrize( + "session,expects_conflict", + [ + ({"modalities": ["text"]}, False), + ({"tools": [CONFLICT_TOOL]}, False), + ({"tools": [{"parameters": OBJECT_PARAMETERS, "name": "get_weather", "type": "function"}]}, False), + ( + {"tools": [{"type": "function", "function": {"name": "get_weather", "parameters": OBJECT_PARAMETERS}}]}, + False, + ), + ({"tools": None}, False), + ({"voice": None}, False), + ({"output_sample_rate_hertz": 24000.0}, False), + ({"output_sample_rate_hertz": 16000.0}, True), + ({"tools": [{"type": "function", "name": "get_time", "parameters": OBJECT_PARAMETERS}]}, True), + ({"tools": []}, True), + ({"instructions": "Be brief."}, True), + ({"instructions": None}, False), + ({"instructions": ""}, False), + ({"voice": "amy"}, True), + ({"output_audio_format": "g711_ulaw"}, True), + ({"output_sample_rate_hertz": 8000}, True), + ({"temperature": 0.1}, True), + ({"max_response_output_tokens": 512}, True), + ({"voice": "matthew", "temperature": 0.7}, False), + ({"input_audio_format": "g711_alaw"}, False), + ], + ) + def test_session_update_rejection_locked_settings(self, session, expects_conflict): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"tools": [CONFLICT_TOOL]}}) + assert config.prompt_started + + conflict: Final = config.session_update_rejection({"type": "session.update", "session": session}) + assert (conflict.code if conflict else None) == ("session_update_not_applied" if expects_conflict else None) + + def test_session_update_with_inf_max_tokens_omits_max_tokens_from_session_start(self): + config: Final = BedrockRealtimeConfig() + messages: Final = config.transform_session_update_event( + {"type": "session.update", "session": {"max_response_output_tokens": "inf", "modalities": ["text"]}} + ) + session_start: Final = json.loads(messages[0])["event"]["sessionStart"]["inferenceConfiguration"] + assert "maxTokens" not in session_start + assert session_start["topP"] == 0.9 + assert session_start["temperature"] == 0.7 + + @pytest.mark.parametrize( + "first_session,later_session,expects_conflict", + [ + ({"modalities": ["text"]}, {"max_response_output_tokens": "inf"}, True), + ({"max_response_output_tokens": "inf"}, {"max_response_output_tokens": "inf"}, False), + ({"max_response_output_tokens": "inf"}, {"max_response_output_tokens": 2048}, True), + ], + ) + def test_session_update_rejection_inf_max_tokens(self, first_session, later_session, expects_conflict): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": first_session}) + assert config.prompt_started + + conflict: Final = config.session_update_rejection({"type": "session.update", "session": later_session}) + assert (conflict is not None) == expects_conflict + + def test_session_update_after_start_leaves_locked_attributes_unchanged(self): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"modalities": ["text"]}}) + + config.transform_session_update_event( + { + "type": "session.update", + "session": {"voice": "amy", "output_audio_format": "g711_ulaw", "temperature": 0.1}, + } + ) + assert config.voice_id == "matthew" + assert config.output_sample_rate_hertz == 24000 + assert config.temperature == 0.7 + + # Input audio settings still apply mid-session + config.transform_session_update_event( + {"type": "session.update", "session": {"input_audio_format": "g711_alaw"}} + ) + assert config.input_sample_rate_hertz == 8000 + config.transform_session_update_event({"type": "session.update", "session": {"input_sample_rate_hertz": 16000}}) + assert config.input_sample_rate_hertz == 16000 + + def test_stray_non_assistant_content_end_emits_no_duplicate_done_events(self): + events: Final = self._run( + BedrockRealtimeConfig(), + self._assistant_text_block("SPECULATIVE", "Hi.", "PARTIAL_TURN") + + [ + {"event": {"contentStart": {"role": "USER", "type": "AUDIO"}}}, + {"event": {"contentEnd": {}}}, + ], + ) + assert len([e for e in events if e["type"] == "response.text.done"]) == 1 + message_item_dones: Final = [ + e for e in events if e["type"] == "response.output_item.done" and e["item"]["type"] == "message" + ] + assert len(message_item_dones) == 1 + + @staticmethod + def _inference_configuration(config: BedrockRealtimeConfig, session: dict[str, object]) -> dict[str, object]: + messages: Final = config.transform_session_update_event({"type": "session.update", "session": session}) + return json.loads(messages[0])["event"]["sessionStart"]["inferenceConfiguration"] + + @pytest.mark.parametrize("max_output_tokens", [4096.0, 4096.5]) + def test_float_max_tokens_reaches_session_start_as_an_integer(self, max_output_tokens): + inference: Final = self._inference_configuration( + BedrockRealtimeConfig(), {"max_response_output_tokens": max_output_tokens, "temperature": 1} + ) + assert inference["maxTokens"] == 4096 + assert isinstance(inference["maxTokens"], int) + assert inference["temperature"] == 1.0 + + @pytest.mark.parametrize( + "session,param", + [ + ({"temperature": "0.2"}, "session.temperature"), + ({"temperature": True}, "session.temperature"), + ({"max_response_output_tokens": "4096"}, "session.max_response_output_tokens"), + ({"temperature": float("nan")}, "session.temperature"), + ({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"), + ({"max_response_output_tokens": True}, "session.max_response_output_tokens"), + ({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"), + ], + ) + def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param): + rejection: Final = BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": session} + ) + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_type", param) + + @pytest.mark.parametrize( + "session", + [ + {"temperature": "0.2", "max_response_output_tokens": "4096"}, + {"temperature": True, "max_response_output_tokens": float("nan")}, + ], + ) + def test_malformed_number_setting_is_dropped_with_drop_params(self, session): + config: Final = BedrockRealtimeConfig(drop_params=True) + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + inference: Final = self._inference_configuration(config, session) + assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + + def test_null_number_settings_count_as_unset(self): + session: Final = {"temperature": None, "max_response_output_tokens": None} + config: Final = BedrockRealtimeConfig() + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + inference: Final = self._inference_configuration(config, session) + assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + + @staticmethod + def _audio_output_configuration(session: dict[str, object]) -> dict[str, object]: + messages: Final = BedrockRealtimeConfig().transform_session_update_event( + {"type": "session.update", "session": session} + ) + return json.loads(messages[1])["event"]["promptStart"]["audioOutputConfiguration"] + + @pytest.mark.parametrize("output_sample_rate_hertz", [16000, 16000.0]) + def test_output_sample_rate_reaches_prompt_start_as_an_integer(self, output_sample_rate_hertz): + audio_output: Final = self._audio_output_configuration({"output_sample_rate_hertz": output_sample_rate_hertz}) + assert audio_output["sampleRateHertz"] == 16000 + assert isinstance(audio_output["sampleRateHertz"], int) + + def test_null_voice_keeps_the_default_voice(self): + assert self._audio_output_configuration({"voice": None})["voiceId"] == "matthew" + + @pytest.mark.parametrize( + "client_event", + [ + {"type": "conversation.item.create", "item": {"type": "message", "role": "user", "content": []}}, + {"type": "input_audio_buffer.append", "audio": base64.b64encode(bytes(320)).decode()}, + {"type": "response.create"}, + ], + ) + def test_first_input_without_an_applied_session_update_starts_the_prompt_once(self, client_event): + config: Final = BedrockRealtimeConfig() + rejected: Final = {"type": "session.update", "session": {"temperature": "warm"}} + assert config.session_update_rejection(rejected) is not None + + first: Final = config.transform_realtime_request(json.dumps(client_event), "amazon.nova-2-sonic-v1:0") + second: Final = config.transform_realtime_request(json.dumps(client_event), "amazon.nova-2-sonic-v1:0") + + first_events: Final = [json.loads(message)["event"] for message in first] + assert [next(iter(event)) for event in first_events[:2]] == ["sessionStart", "promptStart"] + assert first_events[2]["contentStart"]["role"] == "SYSTEM" + assert first_events[0]["sessionStart"]["inferenceConfiguration"]["temperature"] == 0.7 + assert not any("promptStart" in json.loads(message)["event"] for message in second) + + @pytest.mark.parametrize( + "tool_use,name,arguments", + [ + ({"toolUseId": "t1", "content": '{"city":"Paris"}'}, "", '{"city":"Paris"}'), + ( + {"toolUseId": "t1", "toolName": "get_weather", "content": {"city": "Paris"}}, + "get_weather", + '{"city": "Paris"}', + ), + ], + ) + def test_tool_use_with_missing_name_or_object_content_is_forwarded(self, tool_use, name, arguments): + events: Final = self._run(BedrockRealtimeConfig(), [{"event": {"toolUse": tool_use}}]) + (function_call,) = [e for e in events if e["type"] == "response.function_call_arguments.done"] + assert (function_call["name"], function_call["arguments"]) == (name, arguments) + + def test_tool_use_without_tool_use_id_is_forwarded_with_empty_call_id(self): + events: Final = self._run( + BedrockRealtimeConfig(), + [{"event": {"toolUse": {"content": '{"city":"Paris"}', "toolName": "get_weather"}}}], + ) + (function_call,) = [e for e in events if e["type"] == "response.function_call_arguments.done"] + assert (function_call["call_id"], function_call["arguments"]) == ("", '{"city":"Paris"}') + + def test_tool_call_without_assistant_message_stays_output_index_zero(self): + events: Final = self._run(BedrockRealtimeConfig(), self._tool_block("tooluse_A")) + call_events: Final = [ + e + for e in events + if e["type"] + in ( + "response.output_item.added", + "response.output_item.done", + "response.function_call_arguments.delta", + "response.function_call_arguments.done", + ) + ] + assert call_events and all(e["output_index"] == 0 for e in call_events) + if __name__ == "__main__": pytest.main([__file__, "-v"]) From 238597ed4950fa475ace194b24f45680298be6c4 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:11:57 +0000 Subject: [PATCH 18/23] test(ui): send Connection: close on every integration browser request (#46805) * test(ui): close each federation spec API connection so the proxy cannot drop an idle keep-alive socket * test(ui): send Connection: close on every integration browser spec API call * test(ui): send Connection: close on every integration browser request from the Playwright config --- tests/e2e/ui/integration.config.ts | 2 ++ tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/e2e/ui/integration.config.ts b/tests/e2e/ui/integration.config.ts index e332c865222..55ccb860512 100644 --- a/tests/e2e/ui/integration.config.ts +++ b/tests/e2e/ui/integration.config.ts @@ -23,6 +23,8 @@ export default defineConfig({ use: { ...devices["Desktop Chrome"], baseURL: UI_BASE_URL, + // The proxy closes idle keep-alive sockets after 5s and a reused one fails with "socket hang up" + extraHTTPHeaders: { Connection: "close" }, actionTimeout: 15_000, navigationTimeout: 30_000, trace: "retain-on-failure", diff --git a/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts b/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts index fee09b409c7..11cc1c948a3 100644 --- a/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts +++ b/tests/e2e/ui/tests/integrationCritical/mcpUserEnvVars.spec.ts @@ -12,7 +12,7 @@ import { captureRequestBody } from "../../helpers/roundTrip"; import { masterKey } from "../../helpers/traffic"; const master = masterKey(); -const headers = { Authorization: `Bearer ${master}`, Connection: "close" }; +const headers = { Authorization: `Bearer ${master}` }; const TOKEN = "USER_TOKEN"; type EnvVarStatus = { From 55c8ec4647933438de887e8c3d38405bc6fdc950 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 15 Oct 2026 18:27:06 -0700 Subject: [PATCH 19/23] refactor(rust-bridge): move the route host modules under host/routes (#46786) * refactor(rust-bridge): move the route host modules under host/routes Each native route's Python half lived in its own `/route_host.py`, apart from the shared host modules Rust also imports by name. They now sit together in `litellm/rust_bridge/host/routes/.py`, with their unit tests moved to the mirrored path. No behavior change Co-Authored-By: Claude Opus 5.5 * test(rust-host): reference messages route alias Co-Authored-By: yujonglee.dev@gmail.com --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../src/routes/chat_completions/mod.rs | 2 +- .../python-bridge/src/routes/messages/host.rs | 2 +- .../python-bridge/src/routes/ocr/host.rs | 4 ++-- .../python-bridge/src/routes/responses/mod.rs | 2 +- litellm/rust_bridge/host/AGENTS.md | 2 +- .../rust_bridge/host/routes}/__init__.py | 0 .../routes/chat_completions.py} | 0 .../route_host.py => host/routes/messages.py} | 0 .../{ocr/route_host.py => host/routes/ocr.py} | 0 .../route_host.py => host/routes/responses.py} | 0 tests/unit/rust_bridge/AGENTS.md | 2 +- .../{responses => host/routes}/__init__.py | 0 .../routes/test_chat_completions.py} | 2 +- .../routes/test_messages.py} | 18 +++++++++--------- .../routes/test_ocr.py} | 4 ++-- .../routes/test_responses.py} | 2 +- 16 files changed, 20 insertions(+), 20 deletions(-) rename {tests/unit/rust_bridge/chat_completions => litellm/rust_bridge/host/routes}/__init__.py (100%) rename litellm/rust_bridge/{chat_completions/route_host.py => host/routes/chat_completions.py} (100%) rename litellm/rust_bridge/{messages/route_host.py => host/routes/messages.py} (100%) rename litellm/rust_bridge/{ocr/route_host.py => host/routes/ocr.py} (100%) rename litellm/rust_bridge/{responses/route_host.py => host/routes/responses.py} (100%) rename tests/unit/rust_bridge/{responses => host/routes}/__init__.py (100%) rename tests/unit/rust_bridge/{chat_completions/test_route_host.py => host/routes/test_chat_completions.py} (95%) rename tests/unit/rust_bridge/{messages/test_route_host.py => host/routes/test_messages.py} (85%) rename tests/unit/rust_bridge/{ocr/test_route_host.py => host/routes/test_ocr.py} (95%) rename tests/unit/rust_bridge/{responses/test_route_host.py => host/routes/test_responses.py} (97%) diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs index 00ca5e092a8..d78d141f0f2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -23,7 +23,7 @@ fn run_chat_completions( ) -> PyResult> { let host = InferenceHost::new( call.view()?, - "litellm.rust_bridge.chat_completions.route_host", + "litellm.rust_bridge.host.routes.chat_completions", ); run_inference::( py, diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 008f1bde763..1add8fa485f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -26,7 +26,7 @@ use crate::{ routes::RequestView, }; -const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.messages.route_host"; +const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; const REQUEST_ERROR_MARKER: &str = "messages_request_error"; const BODY_FIELDS: [&str; 22] = [ diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs index 4c8469cc019..fc1b70392e2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/host.rs @@ -77,7 +77,7 @@ impl OcrPythonHost { _ => "", }; let mapped = py - .import("litellm.rust_bridge.ocr.route_host") + .import("litellm.rust_bridge.host.routes.ocr") .and_then(|module| module.getattr("map_failure")) .and_then(|map| map.call1((error.value(py), self.view.request(py)?, provider))) .and_then(|mapped| mapped.extract::>().map_err(PyErr::from)); @@ -115,7 +115,7 @@ impl PythonBinding for OcrPythonHost { py: Python<'_>, response: LiteLLMOcrResponse, ) -> PyResult> { - public_response(py, "litellm.rust_bridge.ocr.route_host", &response) + public_response(py, "litellm.rust_bridge.host.routes.ocr", &response) } fn encode_stream_head( diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 1bf24813103..984f1e74df2 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -21,7 +21,7 @@ use super::{ }; use crate::errors::RustBridgeDeclined; -const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.responses.route_host"; +const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.responses"; fn run_responses(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { let resolved = call.resolved()?; diff --git a/litellm/rust_bridge/host/AGENTS.md b/litellm/rust_bridge/host/AGENTS.md index 44d2909a834..9b751a1ada2 100644 --- a/litellm/rust_bridge/host/AGENTS.md +++ b/litellm/rust_bridge/host/AGENTS.md @@ -8,4 +8,4 @@ Python state Rust needs but does not own, behind one module per question. Rust i - The module path and function signature are a contract. The Rust side pins the path in a `const MODULE` in the crate that asks (`python-bridge/src/routes/messages/host.rs`, `python-bridge/src/secrets/`, `callbacks-legacy-python/src/python.rs`); change both sides together and keep `tests/unit/rust_bridge/host/` green - Inputs and outputs use values and frozen snapshots. Retain caller-owned clients and caches as opaque references when identity is required; Rust never discovers them through module globals - Moving a question into Rust means replacing the function body with a native call, not adding a second resolver. The Python signature stays -- A `*/route_host.py` module is not this: it projects one route's kwargs and maps that route's failures. Shared questions belong here +- `routes/.py` holds the questions one route asks, under the same function names on every route (`response`, `map_failure`, `connection_defaults`), so Rust names only the module. A question two routes ask goes in a top-level module here instead diff --git a/tests/unit/rust_bridge/chat_completions/__init__.py b/litellm/rust_bridge/host/routes/__init__.py similarity index 100% rename from tests/unit/rust_bridge/chat_completions/__init__.py rename to litellm/rust_bridge/host/routes/__init__.py diff --git a/litellm/rust_bridge/chat_completions/route_host.py b/litellm/rust_bridge/host/routes/chat_completions.py similarity index 100% rename from litellm/rust_bridge/chat_completions/route_host.py rename to litellm/rust_bridge/host/routes/chat_completions.py diff --git a/litellm/rust_bridge/messages/route_host.py b/litellm/rust_bridge/host/routes/messages.py similarity index 100% rename from litellm/rust_bridge/messages/route_host.py rename to litellm/rust_bridge/host/routes/messages.py diff --git a/litellm/rust_bridge/ocr/route_host.py b/litellm/rust_bridge/host/routes/ocr.py similarity index 100% rename from litellm/rust_bridge/ocr/route_host.py rename to litellm/rust_bridge/host/routes/ocr.py diff --git a/litellm/rust_bridge/responses/route_host.py b/litellm/rust_bridge/host/routes/responses.py similarity index 100% rename from litellm/rust_bridge/responses/route_host.py rename to litellm/rust_bridge/host/routes/responses.py diff --git a/tests/unit/rust_bridge/AGENTS.md b/tests/unit/rust_bridge/AGENTS.md index ab88d6b49ce..64c9fbbadd3 100644 --- a/tests/unit/rust_bridge/AGENTS.md +++ b/tests/unit/rust_bridge/AGENTS.md @@ -6,4 +6,4 @@ Call each path directly with an explicit decision instead. The Python path is th Rollout policy itself, meaning which rule matches and what `LITELLM_RUST` changes, belongs in `test_catalog.py`, `test_configuration.py` and `test_dispatch.py`, tested against rules the test builds rather than the shipped `catalog.RULES` -Before adding a test here, ask whether it checks something Rust cannot. A `route_host.py` module is the Python half of a native route: it projects Python-only state (the cost map, `litellm.*` settings, request kwargs) into the plain values the Rust side consumes, and maps native failures back onto public exceptions. Those projections are what belongs here, because a wrong key or an ignored provider prefix ships the wrong value to Rust and no Rust test sees it. `messages/test_route_host.py` shows the shape. Behavior that lives in Rust (a request transform given its inputs, header assembly, stream relay) is tested in the crate, and the route end to end is tested against a recording server in `tests/test_litellm_rust/`. A test that only re-checks a Python helper the route host happens to call is a duplicate of that helper's own test and should not be added +Before adding a test here, ask whether it checks something Rust cannot. A `host/routes/.py` module is the Python half of a native route: it projects Python-only state (the cost map, `litellm.*` settings, request kwargs) into the plain values the Rust side consumes, and maps native failures back onto public exceptions. Those projections are what belongs here, because a wrong key or an ignored provider prefix ships the wrong value to Rust and no Rust test sees it. `host/routes/test_messages.py` shows the shape. Behavior that lives in Rust (a request transform given its inputs, header assembly, stream relay) is tested in the crate, and the route end to end is tested against a recording server in `tests/test_litellm_rust/`. A test that only re-checks a Python helper the route host happens to call is a duplicate of that helper's own test and should not be added diff --git a/tests/unit/rust_bridge/responses/__init__.py b/tests/unit/rust_bridge/host/routes/__init__.py similarity index 100% rename from tests/unit/rust_bridge/responses/__init__.py rename to tests/unit/rust_bridge/host/routes/__init__.py diff --git a/tests/unit/rust_bridge/chat_completions/test_route_host.py b/tests/unit/rust_bridge/host/routes/test_chat_completions.py similarity index 95% rename from tests/unit/rust_bridge/chat_completions/test_route_host.py rename to tests/unit/rust_bridge/host/routes/test_chat_completions.py index 402c032f340..3496791f5a0 100644 --- a/tests/unit/rust_bridge/chat_completions/test_route_host.py +++ b/tests/unit/rust_bridge/host/routes/test_chat_completions.py @@ -4,7 +4,7 @@ from typing import Final import pytest import litellm -from litellm.rust_bridge.chat_completions.route_host import connection_defaults, response +from litellm.rust_bridge.host.routes.chat_completions import connection_defaults, response from litellm.types.utils import ModelResponse diff --git a/tests/unit/rust_bridge/messages/test_route_host.py b/tests/unit/rust_bridge/host/routes/test_messages.py similarity index 85% rename from tests/unit/rust_bridge/messages/test_route_host.py rename to tests/unit/rust_bridge/host/routes/test_messages.py index 43a17e5e593..abf0e055e73 100644 --- a/tests/unit/rust_bridge/messages/test_route_host.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -5,8 +5,8 @@ import pytest import litellm from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS -from litellm.rust_bridge.messages import route_host -from litellm.rust_bridge.messages.route_host import response +from litellm.rust_bridge.host.routes import messages as messages_host +from litellm.rust_bridge.host.routes.messages import response def test_response_is_a_detached_public_messages_dict() -> None: @@ -36,7 +36,7 @@ def test_settings_project_caller_configuration_without_resolving_a_model(monkeyp monkeypatch.setattr(litellm, "request_timeout", 1800) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) - projected: Final = route_host.settings({"drop_params": "true", "additional_drop_params": ["metadata.user_id"]}) + projected: Final = messages_host.settings({"drop_params": "true", "additional_drop_params": ["metadata.user_id"]}) assert projected == { "drop_params": True, @@ -60,7 +60,7 @@ def test_request_timeout_carries_only_a_configured_global_timeout( monkeypatch.setattr(litellm, "request_timeout", timeout) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", explicitly_set) - assert route_host.settings({})["request_timeout"] == expected + assert messages_host.settings({})["request_timeout"] == expected @pytest.mark.parametrize( @@ -78,7 +78,7 @@ def test_drop_params_merges_the_global_flag_with_the_request( ) -> None: monkeypatch.setattr(litellm, "drop_params", global_flag) - assert route_host.settings(kwargs)["drop_params"] is expected + assert messages_host.settings(kwargs)["drop_params"] is expected @pytest.mark.parametrize( @@ -90,7 +90,7 @@ def test_drop_params_merges_the_global_flag_with_the_request( ], ) def test_additional_drop_params_keep_only_string_paths(configured: object, expected: tuple[str, ...]) -> None: - settings: Final = route_host.settings({"additional_drop_params": configured}) + settings: Final = messages_host.settings({"additional_drop_params": configured}) assert settings["additional_drop_params"] == expected @@ -117,19 +117,19 @@ def test_native_request_rejections_map_to_the_public_400() -> None: rejected: Final = ValueError("claude-sonnet-5 does not support top_k=5") rejected.messages_request_error = True # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets - mapped: Final = route_host.map_failure(rejected, request.resolved, "anthropic") + mapped: Final = messages_host.map_failure(rejected, request.resolved, "anthropic") assert isinstance(mapped, litellm.BadRequestError) assert mapped.status_code == 400 assert "does not support top_k=5" in mapped.message assert mapped.model == "claude-sonnet-5" assert not isinstance( - route_host.map_failure(ValueError("plain"), request.resolved, "anthropic"), litellm.BadRequestError + messages_host.map_failure(ValueError("plain"), request.resolved, "anthropic"), litellm.BadRequestError ) def test_stream_hidden_params_projects_upstream_headers_the_way_the_python_handler_does() -> None: - hidden: Final = route_host.stream_hidden_params( + hidden: Final = messages_host.stream_hidden_params( (("request-id", "req_upstream_123"), ("x-ratelimit-remaining-requests", "41")) ) diff --git a/tests/unit/rust_bridge/ocr/test_route_host.py b/tests/unit/rust_bridge/host/routes/test_ocr.py similarity index 95% rename from tests/unit/rust_bridge/ocr/test_route_host.py rename to tests/unit/rust_bridge/host/routes/test_ocr.py index 879edb951fc..90040f55ef3 100644 --- a/tests/unit/rust_bridge/ocr/test_route_host.py +++ b/tests/unit/rust_bridge/host/routes/test_ocr.py @@ -3,8 +3,8 @@ from typing import Final import pytest import litellm -from litellm.rust_bridge.ocr.route_host import UpstreamFailure, map_failure -from litellm.rust_bridge.ocr.route_host import response as build_ocr_response +from litellm.rust_bridge.host.routes.ocr import UpstreamFailure, map_failure +from litellm.rust_bridge.host.routes.ocr import response as build_ocr_response from litellm.rust_bridge.public_call import NativeCall REQUEST: Final = NativeCall( diff --git a/tests/unit/rust_bridge/responses/test_route_host.py b/tests/unit/rust_bridge/host/routes/test_responses.py similarity index 97% rename from tests/unit/rust_bridge/responses/test_route_host.py rename to tests/unit/rust_bridge/host/routes/test_responses.py index d93ceb12ec9..99e3726c03a 100644 --- a/tests/unit/rust_bridge/responses/test_route_host.py +++ b/tests/unit/rust_bridge/host/routes/test_responses.py @@ -5,7 +5,7 @@ import pytest from pydantic import ValidationError import litellm -from litellm.rust_bridge.responses.route_host import connection_defaults, map_failure, response +from litellm.rust_bridge.host.routes.responses import connection_defaults, map_failure, response from litellm.types.llms.openai import ResponsesAPIResponse From d66d8ac12742d4669c5331788b177da43e0371d8 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Thu, 15 Oct 2026 18:27:06 -0700 Subject: [PATCH 20/23] refactor(rust-bridge): let each route host pick the body params (#46787) * refactor(rust-bridge): let each route host pick the body params The native routes projected request body params from a name list, hard-coded in Rust for Messages and read from `PARAMETERS` for Chat Completions and Responses. Each route host now answers `body_params(request)` over the keywords the call would send, so the Python side owns which params reach the provider. No behavior change Co-Authored-By: Claude Opus 5.5 * fix(rust-host): use read-only body parameter types Co-Authored-By: yujonglee.dev@gmail.com --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../python-bridge/src/routes/inference.rs | 10 ++---- .../python-bridge/src/routes/messages/host.rs | 27 +--------------- .../crates/python-bridge/src/routes/mod.rs | 18 +++++++++++ .../host/routes/chat_completions.py | 6 +++- litellm/rust_bridge/host/routes/messages.py | 32 ++++++++++++++++++- litellm/rust_bridge/host/routes/responses.py | 6 +++- litellm/rust_bridge/public_call.py | 6 +++- .../host/routes/test_chat_completions.py | 19 ++++++++++- 8 files changed, 85 insertions(+), 39 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index bc442d1cf36..a5be7f51828 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -21,10 +21,7 @@ use serde_json::{Map, Value}; use super::{NativeCall, RequestView}; use crate::{ errors::{RustUpstreamError, route_error_to_pyerr}, - marshal::{ - RouteOptions, optional_timeout, project_optional_fields, public_response, - python_timeout_seconds, - }, + marshal::{RouteOptions, optional_timeout, public_response, python_timeout_seconds}, }; pub(super) struct InferenceHost { @@ -115,10 +112,7 @@ impl InferenceHost { py: Python<'_>, arguments: &Bound<'_, PyDict>, ) -> PyResult> { - let names: Vec = py.import(self.module)?.getattr("PARAMETERS")?.extract()?; - project_optional_fields(names.iter().map(String::as_str), |name| { - self.argument(py, arguments, name) - }) + self.view.body_params(py, arguments, self.module) } pub fn response(&self, py: Python<'_>, response: &impl Serialize) -> PyResult> { diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 1add8fa485f..5e445fcaa7d 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -29,31 +29,6 @@ use crate::{ const ROUTE_HOST_MODULE: &str = "litellm.rust_bridge.host.routes.messages"; const REQUEST_ERROR_MARKER: &str = "messages_request_error"; -const BODY_FIELDS: [&str; 22] = [ - "max_tokens", - "metadata", - "stop_sequences", - "stream", - "system", - "temperature", - "thinking", - "tool_choice", - "tools", - "top_k", - "inference_geo", - "top_p", - "mcp_servers", - "context_management", - "compaction", - "container", - "output_format", - "speed", - "output_config", - "cache_control", - "reasoning_effort", - "safeguards", -]; - fn merge_headers( forwarded: Option>, extra_headers: Option>, @@ -143,7 +118,7 @@ impl MessagesPythonHost { let model = string("model")?.ok_or_else(|| PyValueError::new_err("model is required"))?; let messages = argument("messages")?.ok_or_else(|| PyValueError::new_err("messages is required"))?; - let fields = project_optional_fields(BODY_FIELDS, argument)?; + let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; let body = [ ("model".to_string(), Value::String(model.clone())), ("messages".to_string(), from_py(&messages)?), diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 444ef4f01b7..f51c755c33e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -13,6 +13,7 @@ use litellm_callbacks_legacy_python::{LegacyLogging, LoggingOperation, PublicCal use litellm_host::{call::HostedCompletion, machine::Machine, protocol::Protocol}; use litellm_host_python::{ CallOptions, HookChain, PythonBinding, PythonCallHooks, PythonHostCalls, effective_py_args, + from_py, }; use litellm_host_python::{Error as HostError, present}; use pyo3::{ @@ -20,6 +21,7 @@ use pyo3::{ prelude::*, types::{PyDict, PyMapping, PyTuple}, }; +use serde_json::{Map, Value}; /// The public call as Python bound it: `base` holds the positional arguments by name plus /// the signature defaults, `kwargs` the caller's keyword dict. @@ -92,6 +94,22 @@ impl RequestView { present(arguments, self.request(py)?, name) } + /// The body params `module.body_params` keeps from the keywords the call would send, so + /// each route's Python host decides which params reach the provider. + pub(crate) fn body_params( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + module: &str, + ) -> PyResult> { + let request = effective_py_args(self.request(py)?, arguments)?; + from_py( + &py.import(module)? + .getattr("body_params")? + .call1((request,))?, + ) + } + pub(crate) fn close(&mut self) { self.0 = None; } diff --git a/litellm/rust_bridge/host/routes/chat_completions.py b/litellm/rust_bridge/host/routes/chat_completions.py index d600495c919..5fcb2edf4f7 100644 --- a/litellm/rust_bridge/host/routes/chat_completions.py +++ b/litellm/rust_bridge/host/routes/chat_completions.py @@ -6,7 +6,7 @@ from typing import Final import litellm from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS from litellm.rust_bridge import failures -from litellm.rust_bridge.public_call import optional_str +from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.utils import ModelResponse _TRANSPORT_PARAMETERS: Final = frozenset( @@ -27,6 +27,10 @@ _TRANSPORT_PARAMETERS: Final = frozenset( PARAMETERS: Final = tuple(name for name in OPENAI_CHAT_COMPLETION_PARAMS if name not in _TRANSPORT_PARAMETERS) +def body_params(request: Mapping[str, object]) -> Mapping[str, object]: + return known_params(PARAMETERS, request) + + def connection_defaults(provider: str) -> tuple[str | None, str | None]: if provider == "anthropic": return litellm.anthropic_key or litellm.api_key, litellm.api_base diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index 5dfd3bfd530..1024141d5a2 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -12,10 +12,36 @@ from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout from litellm.llms.anthropic.pass_through.utils import is_reasoning_auto_summary_enabled from litellm.rust_bridge import failures -from litellm.rust_bridge.public_call import optional_str +from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse _DROP_PATHS: Final = TypeAdapter(list[object]) +_BODY_FIELDS: Final = frozenset( + { + "max_tokens", + "metadata", + "stop_sequences", + "stream", + "system", + "temperature", + "thinking", + "tool_choice", + "tools", + "top_k", + "inference_geo", + "top_p", + "mcp_servers", + "context_management", + "compaction", + "container", + "output_format", + "speed", + "output_config", + "cache_control", + "reasoning_effort", + "safeguards", + } +) @dataclass(frozen=True, slots=True) @@ -33,6 +59,10 @@ def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: ) +def body_params(request: Mapping[str, object]) -> Mapping[str, object]: + return known_params(_BODY_FIELDS, request) + + def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]: from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( anthropic_messages_stream_hidden_params, diff --git a/litellm/rust_bridge/host/routes/responses.py b/litellm/rust_bridge/host/routes/responses.py index 218a866cf46..470a89a0e05 100644 --- a/litellm/rust_bridge/host/routes/responses.py +++ b/litellm/rust_bridge/host/routes/responses.py @@ -6,12 +6,16 @@ from typing import Final import litellm from litellm import get_llm_provider from litellm.rust_bridge import failures -from litellm.rust_bridge.public_call import inference_decline_reason, optional_str +from litellm.rust_bridge.public_call import inference_decline_reason, known_params, optional_str from litellm.types.llms.openai import ResponsesAPIOptionalRequestParams, ResponsesAPIResponse PARAMETERS: Final = tuple(ResponsesAPIOptionalRequestParams.__annotations__) +def body_params(request: Mapping[str, object]) -> Mapping[str, object]: + return known_params(PARAMETERS, request) + + def connection_defaults(provider: str) -> tuple[str | None, str | None]: provider_key: Final = litellm.openai_key if provider == "openai" else None return litellm.api_key or provider_key, litellm.api_base diff --git a/litellm/rust_bridge/public_call.py b/litellm/rust_bridge/public_call.py index e5700be89ad..332421c68cc 100644 --- a/litellm/rust_bridge/public_call.py +++ b/litellm/rust_bridge/public_call.py @@ -3,7 +3,7 @@ from __future__ import annotations import inspect -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Collection, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, TypeVar, cast # noqa: TID251 # narrows caller-owned containers without copying them @@ -71,6 +71,10 @@ def optional_sequence(value: object) -> Sequence[object] | None: return cast("Sequence[object]", value) # cast-ok: the same caller-owned object is handed on unchanged +def known_params(names: Collection[str], request: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in request.items() if name in names and value is not None} + + def inference_decline_reason(parameters: tuple[str, ...], kwargs: Mapping[str, object]) -> str | None: if litellm.drop_params or litellm.modify_params: return "native inference does not implement the configured parameter rewrites" diff --git a/tests/unit/rust_bridge/host/routes/test_chat_completions.py b/tests/unit/rust_bridge/host/routes/test_chat_completions.py index 3496791f5a0..ea715a6eb22 100644 --- a/tests/unit/rust_bridge/host/routes/test_chat_completions.py +++ b/tests/unit/rust_bridge/host/routes/test_chat_completions.py @@ -4,7 +4,7 @@ from typing import Final import pytest import litellm -from litellm.rust_bridge.host.routes.chat_completions import connection_defaults, response +from litellm.rust_bridge.host.routes.chat_completions import body_params, connection_defaults, response from litellm.types.utils import ModelResponse @@ -57,3 +57,20 @@ def test_connection_defaults_preserve_provider_precedence( monkeypatch.setattr(litellm, "anthropic_key", provider_key) monkeypatch.setattr(litellm, "api_base", "https://configured.invalid") assert connection_defaults(provider) == (expected_key, expected_base) + + +def test_body_params_keep_request_params_and_leave_out_transport_and_unset_ones() -> None: + body: Final = body_params( + MappingProxyType( + { + "model": "anthropic/claude-sonnet-5-5", + "temperature": 0.2, + "max_tokens": None, + "api_key": "sk-secret", + "api_base": "https://example.invalid", + "timeout": 30, + } + ) + ) + + assert body == {"temperature": 0.2} From b47f7a4712bb28771e7ffa711d9c065870971f25 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 01:28:12 +0000 Subject: [PATCH 21/23] fix(responses): preserve explicit websocket API keys (internal copy of #40853) (#46724) * fix(responses): preserve explicit websocket API keys Refresh the one-line fix from BerriAI/litellm#34253 against the current default branch Co-authored-by: puffer <389053104@qq.com> * test(responses): verify websocket authorization Cover explicit deployment keys, precedence over global keys, and the existing global fallback at the upstream WebSocket transport boundary * test(responses): type the authorization capture Replace the mutable header list with a typed, final Future carrying the upstream authorization value, addressing the review on #40853 * Refactor test_openai_websocket_authorization method * fix(responses): resolve websocket keys with the deployment's provider Pass the deployment custom_llm_provider to get_llm_provider so an api_base host never swaps in another provider's env key, and drop the OpenAI-only global key fallbacks so non-OpenAI websocket deployments resolve their own env keys the way the HTTP route does * fix(responses): let a deployment's own websocket key beat env keys A deployment with an api_base on a host litellm recognizes, and no declared provider, got that host's env key ahead of its own api_key. The explicit key now always wins, matching the HTTP Responses route. Drop the litellm_params api_key fallback, which could never be set because api_key is a named argument * fix(responses): keep openai deployments on other providers' known hosts on the HTTP stream An openai deployment whose api_base is a known OpenAI-compatible provider host (DeepSeek, Groq, ...) opened a native Responses WebSocket there and got a 405, because those hosts serve /responses over HTTP only. Route it through the managed HTTP stream with the deployment's own key instead. * fix(responses): resolve openai deployments on known provider hosts as that provider An openai deployment whose api_base is a known provider host (DeepSeek, Mistral, ...) now resolves to that host's provider on the Responses WebSocket again, as before this branch, so it keeps that provider's chat completions bridge, pricing and env key fallback. Its own api_key still wins over the env key. The known-host test now mocks the upstream HTTP route and checks the route, the bearer and response.completed. * test(responses): prove websocket deployment keys and declared providers reach the provider socket * test(responses): use the reported bare model shape for the websocket deployment key cell --------- Co-authored-by: Ysh <734645956@qq.com> Co-authored-by: puffer <389053104@qq.com> --- litellm/constants.py | 2 +- .../get_llm_provider_logic.py | 4 + litellm/responses/main.py | 26 +-- .../test_openai_responses_websocket_wire.py | 98 +++++++++- .../test_responses_websocket_all_providers.py | 173 ++++++++++++++++++ 5 files changed, 286 insertions(+), 17 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index f7ba55d8ce4..9f29baac562 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -969,7 +969,7 @@ DEFAULT_CHAT_COMPLETION_PARAM_VALUES: Final = { "context_management": None, } -openai_compatible_endpoints: Final[list] = [ +openai_compatible_endpoints: Final[list[str]] = [ "api.perplexity.ai", "api.endpoints.anyscale.com/v1", "api.deepinfra.com/v1/openai", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 74d9344a172..a67c1bb7df8 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -61,6 +61,10 @@ def _endpoint_matches_api_base(endpoint: str, api_base: str) -> bool: return url_path == endpoint_path or url_path.startswith(endpoint_path + "/") +def is_openai_compatible_endpoint(api_base: str) -> bool: + return any(_endpoint_matches_api_base(endpoint, api_base) for endpoint in litellm.openai_compatible_endpoints) + + def is_non_openai_azure_model(model: str) -> bool: try: model_name: Final = model.split("/", 1)[1] diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f9225dc1a79..f566c8f4e02 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -59,8 +59,8 @@ if TYPE_CHECKING: else: ResponseText = str # Fallback for ResponseText import from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.litellm_core_utils.get_llm_provider_logic import is_openai_compatible_endpoint from litellm.llms.openai.data_residency import infer_openai_data_residency -from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.responses.streaming_websocket import ResponsesWebSocketRequestDefaults from litellm.types.router import GenericLiteLLMParams @@ -2461,6 +2461,13 @@ def _build_responses_websocket_request_defaults(kwargs: Mapping[str, object]) -> ) +def _websocket_custom_llm_provider(custom_llm_provider: str | None, api_base: str | None) -> str | None: + is_openai_on_known_provider_host: Final = ( + custom_llm_provider == "openai" and api_base is not None and is_openai_compatible_endpoint(api_base) + ) + return None if is_openai_on_known_provider_host else custom_llm_provider + + @client async def _aresponses_websocket( model: str, @@ -2491,6 +2498,7 @@ async def _aresponses_websocket( dynamic_api_base, ) = litellm.get_llm_provider( model=model, + custom_llm_provider=_websocket_custom_llm_provider(litellm_params.custom_llm_provider, api_base), api_base=api_base, api_key=api_key, litellm_params=litellm_params, @@ -2512,22 +2520,18 @@ async def _aresponses_websocket( ) resolved_api_base: Final = dynamic_api_base or litellm_params.api_base or litellm.api_base or None - responses_api_provider_config: BaseResponsesAPIConfig | None = None - if _custom_llm_provider is not None: - responses_api_provider_config = ProviderConfigManager.get_provider_responses_api_config( + responses_api_provider_config: Final = ( + None + if _custom_llm_provider is None + else ProviderConfigManager.get_provider_responses_api_config( model=resolved_model, provider=litellm.LlmProviders(_custom_llm_provider), api_base=resolved_api_base, ) - - resolved_api_key: Final = ( - dynamic_api_key - or litellm_params.api_key - or litellm.api_key - or litellm.openai_key - or get_secret_str("OPENAI_API_KEY") ) + resolved_api_key: Final = api_key or dynamic_api_key + # Extract params that we're passing explicitly to avoid duplicates in **kwargs _explicit_keys: Final = { "user_api_key_dict", diff --git a/tests/integration/providers/test_openai_responses_websocket_wire.py b/tests/integration/providers/test_openai_responses_websocket_wire.py index 18be6bb9abb..27d328369d3 100644 --- a/tests/integration/providers/test_openai_responses_websocket_wire.py +++ b/tests/integration/providers/test_openai_responses_websocket_wire.py @@ -27,14 +27,23 @@ pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) PROVIDER_MODEL: Final = "ws-peer-model" USAGE: Final = {"input_tokens": 5, "output_tokens": 2, "total_tokens": 7} TERMINAL: Final = frozenset({"response.completed", "response.failed", "error"}) +PROXY_OPENAI_KEY: Final = "sk-proxy-environment-openai-key" +PROXY_AZURE_KEY: Final = "proxy-environment-azure-key" _Script: TypeAlias = Callable[[str], tuple[dict[str, JsonValue], ...]] +@dataclass(frozen=True, slots=True) +class Credentials: + authorization: str | None + api_key: str | None + + @dataclass(frozen=True, slots=True) class Peer: url: str paths: SimpleQueue[str] + credentials: SimpleQueue[Credentials] frames: SimpleQueue[dict[str, JsonValue]] @@ -75,10 +84,17 @@ def _top_level_rate_limit(response_id: str) -> tuple[dict[str, JsonValue], ...]: async def _answer( connection: ServerConnection, paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], frames: SimpleQueue[dict[str, JsonValue]], script: _Script, ) -> None: - paths.put(connection.request.path if connection.request is not None else "") + request: Final = connection.request + paths.put(request.path if request is not None else "") + credentials.put( + Credentials(request.headers.get("Authorization"), request.headers.get("api-key")) + if request is not None + else Credentials(None, None) + ) turns: Final = itertools.count(1) async for raw in connection: frame: Final = json.loads(raw) @@ -92,12 +108,15 @@ async def _answer( async def _serve( tls: ssl.SSLContext, paths: SimpleQueue[str], + credentials: SimpleQueue[Credentials], frames: SimpleQueue[dict[str, JsonValue]], ports: SimpleQueue[int], stop: asyncio.Event, script: _Script, ) -> None: - async with serve(lambda connection: _answer(connection, paths, frames, script), "127.0.0.1", 0, ssl=tls) as server: + async with serve( + lambda connection: _answer(connection, paths, credentials, frames, script), "127.0.0.1", 0, ssl=tls + ) as server: ports.put(next(iter(server.sockets)).getsockname()[1]) await stop.wait() @@ -107,16 +126,17 @@ def responses_peer(cert: tuple[Path, Path], script: _Script = _seven) -> Iterato loop: Final = asyncio.new_event_loop() stop: Final = asyncio.Event() paths: Final = SimpleQueue[str]() + credentials: Final = SimpleQueue[Credentials]() frames: Final = SimpleQueue[dict[str, JsonValue]]() ports: Final = SimpleQueue[int]() thread: Final = threading.Thread( target=loop.run_until_complete, - args=(_serve(server_context(*cert), paths, frames, ports, stop, script),), + args=(_serve(server_context(*cert), paths, credentials, frames, ports, stop, script),), daemon=True, ) thread.start() try: - yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, frames) + yield Peer(f"https://127.0.0.1:{ports.get(timeout=10)}/v1", paths, credentials, frames) finally: loop.call_soon_threadsafe(stop.set) thread.join(timeout=10) @@ -180,7 +200,11 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: @pytest.fixture(scope="module") def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: with gateway_from_environment() as base: - with owned_proxy(base, tmp_path_factory.mktemp("responses-ws"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + with owned_proxy( + base, + tmp_path_factory.mktemp("responses-ws"), + {"SSL_CERT_FILE": str(cert[0]), "OPENAI_API_KEY": PROXY_OPENAI_KEY, "AZURE_API_KEY": PROXY_AZURE_KEY}, + ) as proxy: yield proxy @@ -253,3 +277,67 @@ def test_a_top_level_error_frame_reaches_the_client_and_logs_the_provider_code_a error: Final = object_value(object_value(row["metadata"])["error_information"]) assert str(error["error_code"]) == "429", error assert RATE_LIMIT_MESSAGE in string_value(error["error_message"]), error + + +def _one_turn(candidate: Gateway, key: str, model: str) -> tuple[dict[str, JsonValue], ...]: + (events,) = asyncio.run(_session(str(candidate.client.base_url), key, model, (f"say seven {uuid.uuid4().hex}",))) + return events + + +def test_a_bare_model_declared_as_openai_sends_its_deployment_key_to_the_provider_socket( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + deployment_key: Final = f"sk-deployment-{uuid.uuid4().hex}" + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=PROVIDER_MODEL, custom_llm_provider="openai", api_base=peer.url, api_key=deployment_key + ) + events: Final = _one_turn(candidate, scenario.key(models=[model]), model) + assert _completed(events)["status"] == "completed" + assert peer.credentials.get_nowait().authorization == f"Bearer {deployment_key}" + + +def test_an_openai_deployment_without_its_own_key_uses_the_proxy_openai_key( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model(model=f"openai/{PROVIDER_MODEL}", api_base=peer.url, api_key=None) + events: Final = _one_turn(candidate, scenario.key(models=[model]), model) + assert _completed(events)["status"] == "completed" + assert peer.credentials.get_nowait().authorization == f"Bearer {PROXY_OPENAI_KEY}" + + +def test_an_azure_deployment_without_its_own_key_gets_the_azure_key_and_never_the_openai_key( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model( + model=f"azure/{PROVIDER_MODEL}", api_base=peer.url.removesuffix("/v1"), api_key=None + ) + events: Final = _one_turn(candidate, scenario.key(models=[model]), model) + assert _completed(events)["status"] == "completed" + credentials: Final = peer.credentials.get_nowait() + assert credentials.api_key == PROXY_AZURE_KEY + assert PROXY_OPENAI_KEY not in (credentials.authorization or "") + + +def test_a_declared_openai_provider_wins_over_another_providers_model_prefix( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + gateway_model: Final = f"anthropic/{PROVIDER_MODEL}" + with responses_peer(cert) as peer, candidate.scenario() as scenario: + model: Final = scenario.model(model=gateway_model, custom_llm_provider="openai", api_base=peer.url) + key: Final = scenario.key(models=[model]) + events: Final = _one_turn(candidate, key, model) + assert _completed(events)["status"] == "completed" + (forwarded,) = _drain(peer.frames) + assert forwarded["model"] == gateway_model + (row,) = eventually( + lambda: read_rows( + 'SELECT custom_llm_provider FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (hashlib.sha256(key.encode()).hexdigest(),), + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + assert row["custom_llm_provider"] == "openai", row diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 9f9deb2ca0f..3625bbb7c04 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -8,6 +8,7 @@ Tests that: """ import json +from collections.abc import Mapping from typing import Final from unittest.mock import MagicMock @@ -2805,6 +2806,178 @@ class TestWebSocketChunkTypes: assert "Complete reasoning" in serialized +class TestResponsesWebSocketCredentials: + @staticmethod + async def _upstream_handshake_headers(**call_kwargs: object) -> Mapping[str, str]: + from asyncio import Future, get_running_loop + from types import TracebackType + from unittest.mock import AsyncMock, patch + + from litellm.responses import main as responses_main + + captured_headers: Final[Future[Mapping[str, str]]] = get_running_loop().create_future() + + class FakeConnect: + def __init__( + self, + _url: str, + *, + additional_headers: Mapping[str, str], + **_kwargs: object, + ) -> None: + captured_headers.set_result(dict(additional_headers)) + + async def __aenter__(self) -> None: + raise RuntimeError("stop at upstream handshake") + + async def __aexit__( + self, + _exc_type: type[BaseException] | None, + _exc: BaseException | None, + _tb: TracebackType | None, + ) -> None: + return None + + websocket: Final = MagicMock() + websocket.close = AsyncMock() + with patch("websockets.connect", FakeConnect): + await responses_main._aresponses_websocket.__wrapped__( + websocket=websocket, + litellm_logging_obj=MagicMock(), + **call_kwargs, + ) + + assert captured_headers.done(), "the request never reached the native upstream websocket" + return captured_headers.result() + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "api_base,explicit_key,global_key,expected_key", + [ + ("https://upstream.invalid/v1", "deployment-key", None, "deployment-key"), + ("https://upstream.invalid/v1", "deployment-key", "global-key", "deployment-key"), + ("https://upstream.invalid/v1", None, "global-key", "global-key"), + ], + ) + async def test_openai_websocket_authorization( + self, + monkeypatch: pytest.MonkeyPatch, + api_base: str, + explicit_key: str | None, + global_key: str | None, + expected_key: str, + ) -> None: + import litellm + + monkeypatch.setattr(litellm, "api_key", global_key) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("LITELLM_RUST", "false") + + headers: Final = await self._upstream_handshake_headers( + model="gpt-6-astra", + custom_llm_provider="openai", + api_key=explicit_key, + api_base=api_base, + ) + + assert headers["Authorization"] == f"Bearer {expected_key}" + + @pytest.mark.asyncio + async def test_azure_websocket_without_its_own_key_never_gets_the_openai_key( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import litellm + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setattr(litellm, "azure_key", None) + monkeypatch.setenv("OPENAI_API_KEY", "openai-global-key") + monkeypatch.setenv("AZURE_API_KEY", "azure-env-key") + monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False) + monkeypatch.setenv("LITELLM_RUST", "false") + + headers: Final = await self._upstream_handshake_headers( + model="gpt-6-astra", + custom_llm_provider="azure", + api_base="https://resource.openai.azure.com", + ) + + assert headers["api-key"] == "azure-env-key" + assert "openai-global-key" not in " ".join(headers.values()) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider,api_base,deployment_key,expected_key", + [ + (None, "https://api.deepseek.com/v1", "deployment-key", "deployment-key"), + ("openai", "https://api.deepseek.com/v1", "deployment-key", "deployment-key"), + ("openai", "https://api.deepseek.com/v1", None, "deepseek-env-key"), + ("openai", "https://api.mistral.ai/v1", "deployment-key", "deployment-key"), + ], + ) + async def test_deployment_on_a_known_provider_host_completes_through_that_provider( + self, + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: str | None, + api_base: str, + deployment_key: str | None, + expected_key: str, + ) -> None: + from unittest.mock import AsyncMock, patch + + import httpx + import respx + + import litellm + from litellm.responses import main as responses_main + + chunks: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "OK"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ) + stream: Final = "".join( + f"data: {json.dumps({'id': 'chatcmpl-1', 'object': 'chat.completion.chunk', 'created': 1, 'model': 'astra-chat', **chunk})}\n\n" + for chunk in chunks + ) + + def refuse_native_socket(*_args: object, **_kwargs: object) -> None: + raise ConnectionRefusedError("this host has no Responses WebSocket") + + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setenv("OPENAI_API_KEY", "openai-env-key") + monkeypatch.setenv("DEEPSEEK_API_KEY", "deepseek-env-key") + monkeypatch.setenv("MISTRAL_API_KEY", "mistral-env-key") + monkeypatch.setenv("LITELLM_RUST", "false") + websocket: Final = MagicMock() + websocket.send_text = AsyncMock() + websocket.close = AsyncMock() + websocket.receive_text = AsyncMock(side_effect=RuntimeError("client closed")) + + with respx.mock(assert_all_called=False) as upstream, patch("websockets.connect", refuse_native_socket): + chat_completions: Final = upstream.post(f"{api_base}/chat/completions").mock( + return_value=httpx.Response( + 200, text=stream + "data: [DONE]\n\n", headers={"content-type": "text/event-stream"} + ) + ) + await responses_main._aresponses_websocket.__wrapped__( + model="astra-chat", + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=deployment_key, + first_message=json.dumps({"type": "response.create", "model": "astra-chat", "input": "Reply OK"}), + websocket=websocket, + litellm_logging_obj=MagicMock(), + ) + + sent_types: Final = [json.loads(call.args[0]).get("type") for call in websocket.send_text.await_args_list] + assert chat_completions.called, f"the request never reached {api_base}/chat/completions, client saw {sent_types}" + assert chat_completions.calls.last.request.headers["Authorization"] == f"Bearer {expected_key}" + assert "response.completed" in sent_types + + class TestNativeWebSocketUrlConstruction: """Test that native WebSocket URLs include the model query parameter. From 250e6bd24b16a966e35deb921671e7b5cb7827a5 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 15 Oct 2026 18:28:32 -0700 Subject: [PATCH 22/23] test(integration): build OpenAI and Anthropic SDK models before tests start threads (#46797) K2 in test_logging_only_scope_chaos.py flaked on CircleCI with AttributeError: type object 'BaseModel' has no attribute '__pydantic_core_schema__' raised inside the OpenAI client while its 30-thread control burst parsed responses The SDKs define their response models with deferred schemas, so the first parse of each model rebuilds it, and pydantic's model_rebuild is not safe when several threads hit it at once. Every integration worker now builds those schemas in pytest_configure, before any test can start a thread pool Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- tests/integration/_support/sdk_models.py | 23 ++++ tests/integration/conftest.py | 2 + .../integration_support/test_sdk_models.py | 128 ++++++++++++++++++ 3 files changed, 153 insertions(+) create mode 100644 tests/integration/_support/sdk_models.py create mode 100644 tests/unit/integration_support/test_sdk_models.py diff --git a/tests/integration/_support/sdk_models.py b/tests/integration/_support/sdk_models.py new file mode 100644 index 00000000000..06e94b349d5 --- /dev/null +++ b/tests/integration/_support/sdk_models.py @@ -0,0 +1,23 @@ +from __future__ import annotations + +from collections.abc import Iterator +from typing import Final + +import anthropic +import openai +from pydantic import BaseModel + +_SDK_PACKAGES: Final = frozenset({openai.__name__, anthropic.__name__}) + + +def _subclasses(model: type[BaseModel]) -> Iterator[type[BaseModel]]: + for subclass in model.__subclasses__(): + yield subclass + yield from _subclasses(subclass) + + +def build_sdk_models() -> None: + """Build every pydantic model the OpenAI and Anthropic SDKs define now rather than on its first use.""" + for model in frozenset(_subclasses(BaseModel)): + if model.__module__.partition(".")[0] in _SDK_PACKAGES: + _ = model.model_rebuild(raise_errors=False) diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index a4bd6092a0c..e82430c3add 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -25,6 +25,7 @@ from tests.integration._support.manifest import OWNED_DIRECTORIES from tests.integration._support.ordering import collected_case, fixture_groups, scheduling_groups from tests.integration._support.provider import SharedProvider, shared_provider from tests.integration._support.routing import RoutingPlugin +from tests.integration._support.sdk_models import build_sdk_models from tests.integration._support.workers import gateway_profile, start_worker_stacks, worker_stack from tests.integration.run import GITHUB_FILES @@ -70,6 +71,7 @@ def pytest_configure(config: pytest.Config) -> None: config.stash[REPORTS] = [] config.stash[SHARD_FILES] = _shard_files(SHARD_PATH.validate_python(config.getoption("integration_shard_files"))) config.pluginmanager.register(IntegrationReportPlugin(config)) + build_sdk_models() if os.environ.get("INTEGRATION_ROUTING"): config.pluginmanager.register(RoutingPlugin(config)) diff --git a/tests/unit/integration_support/test_sdk_models.py b/tests/unit/integration_support/test_sdk_models.py new file mode 100644 index 00000000000..285d3d78fff --- /dev/null +++ b/tests/unit/integration_support/test_sdk_models.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path +from typing import Final + +_REPOSITORY: Final = Path(__file__).resolve().parents[3] +_CONCURRENT_FIRST_USE: Final = """ +import sys +import threading + +from anthropic.types import Message, RawContentBlockDeltaEvent +from openai.types.chat import ChatCompletion, ChatCompletionChunk +from openai.types.responses import Response, ResponseTextDeltaEvent + +from tests.integration._support.sdk_models import build_sdk_models + +USAGE = {"input_tokens": 1, "output_tokens": 1} +PAYLOADS = ( + ( + Response, + { + "id": "resp_1", + "object": "response", + "created_at": 1, + "model": "gpt-4.1-mini", + "status": "completed", + "output": [ + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "hi", "annotations": []}], + } + ], + "usage": {**USAGE, "total_tokens": 2}, + }, + ), + ( + ChatCompletion, + { + "id": "chatcmpl_1", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ), + ( + Message, + { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": USAGE, + }, + ), + ( + ResponseTextDeltaEvent, + { + "type": "response.output_text.delta", + "item_id": "msg_1", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + "delta": "hi", + "logprobs": [], + }, + ), + ( + ChatCompletionChunk, + { + "id": "chatcmpl_1", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "finish_reason": None, "delta": {"role": "assistant", "content": "hi"}}], + }, + ), + ( + RawContentBlockDeltaEvent, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + ), +) +THREADS = 30 + +sys.setswitchinterval(1e-6) +build_sdk_models() +gate = threading.Barrier(THREADS) +failures = [] + + +def first_use(): + gate.wait() + for model, payload in PAYLOADS: + try: + model.model_construct(**payload) + except Exception as error: + failures.append(f"{model.__name__}: {type(error).__name__}: {error}") + + +threads = [threading.Thread(target=first_use) for _ in range(THREADS)] +for thread in threads: + thread.start() +for thread in threads: + thread.join() +print(sorted(set(failures))) +""" + + +def test_sdk_models_built_up_front_survive_concurrent_first_use_from_many_threads() -> None: + completed: Final = subprocess.run( + (sys.executable, "-P", "-c", _CONCURRENT_FIRST_USE), + env={**os.environ, "PYTHONPATH": str(_REPOSITORY)}, + capture_output=True, + text=True, + timeout=120, + check=False, + ) + assert completed.returncode == 0, completed.stderr + assert completed.stdout.strip() == "[]", completed.stdout From 627bcab42b2d11dec46b770622b3f7624caa07cf Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 15 Oct 2026 18:29:05 -0700 Subject: [PATCH 23/23] fix(guardrails): honor streaming settings on the Noma v2 guardrail (#46606) * fix(guardrails): honor streaming settings on the Noma v2 guardrail The Noma v2 initializer now passes streaming_buffer_until_moderated, streaming_end_of_stream_only, streaming_sampling_rate and streaming_buffer_release_on_scan from the guardrail config to NomaV2Guardrail, so the unified streaming hook applies them. The constructor validates them, so a custom guardrail class built on NomaV2Guardrail gets the same checks. Unset keys keep the hook defaults, and an invalid value is ignored with a warning naming the guardrail, so the guardrail still loads. Resolves LIT-9633 * test(guardrails): cover Noma v2 streaming settings across clients, failures and workers Adds integration cells for the custom-subclass loader, string and invalid values, unset and false keys, a Noma outage with and without block_failures, an upstream rejection, n=2 choices, pre_call and non-streaming guardrails, the OpenAI and Anthropic SDKs (sync and async), stored-guardrail info, a PATCH while traffic flows, a burst while Noma is down, a slow Noma, and replaced workers. Resolves LIT-9633 * test(guardrails): pin the PATCH probes to one worker at a time Each probe suspends the other worker and opens a fresh connection, so the test checks that every worker applies the patched setting. Resolves LIT-9633 * test(guardrails): check sampled scans by count, not window position /v1/messages hands the streaming hook raw network reads, so where the default sampling windows fall depends on timing; the base leg caught one shifted window. Check the scan count and the final full-text scan instead. Resolves LIT-9633 --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../guardrail_hooks/noma/__init__.py | 5 + .../guardrail_hooks/noma/noma_v2.py | 23 + .../proxy/guardrails/guardrail_hooks/noma.py | 46 +- .../test_noma_v2_streaming_settings.py | 1009 +++++++++++++++++ .../guardrail_hooks/noma/test_noma_v2.py | 121 ++ 5 files changed, 1203 insertions(+), 1 deletion(-) create mode 100644 tests/integration/streaming/test_noma_v2_streaming_settings.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py index 9391cf60cc3..e372039936c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py @@ -38,6 +38,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardrail"): import litellm + extras: Final = litellm_params.model_extra or {} _noma_v2_callback: Final = NomaV2Guardrail( guardrail_name=guardrail.get("guardrail_name", ""), api_key=litellm_params.api_key, @@ -49,6 +50,10 @@ def initialize_guardrail_v2(litellm_params: "LitellmParams", guardrail: "Guardra event_hook=litellm_params.mode, default_on=litellm_params.default_on, timeout=litellm_params.timeout, + streaming_buffer_until_moderated=extras.get("streaming_buffer_until_moderated"), + streaming_end_of_stream_only=extras.get("streaming_end_of_stream_only"), + streaming_sampling_rate=extras.get("streaming_sampling_rate"), + streaming_buffer_release_on_scan=extras.get("streaming_buffer_release_on_scan"), ) litellm.logging_callback_manager.add_litellm_callback(_noma_v2_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 37a33023d54..360752d201a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -26,6 +26,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage from litellm.types.guardrail_base_init import GuardrailBaseInitKwargs from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.noma import NomaV2StreamingParams from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus if TYPE_CHECKING: @@ -63,6 +64,10 @@ class NomaV2Guardrail(CustomGuardrail): monitor_mode: bool | None = None, block_failures: bool | None = None, gateway_name: str | None = None, + streaming_buffer_until_moderated: bool | None = None, + streaming_end_of_stream_only: bool | None = None, + streaming_sampling_rate: int | None = None, + streaming_buffer_release_on_scan: bool | None = None, **kwargs: Any, ) -> None: self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -86,6 +91,24 @@ class NomaV2Guardrail(CustomGuardrail): if self._requires_api_key(api_base=self.api_base) and not self.api_key: raise ValueError("Noma v2 guardrail requires api_key when using Noma SaaS endpoint") + streaming: Final = NomaV2StreamingParams.model_validate( + { + "streaming_buffer_until_moderated": streaming_buffer_until_moderated, + "streaming_end_of_stream_only": streaming_end_of_stream_only, + "streaming_sampling_rate": streaming_sampling_rate, + "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan, + }, + context={"guardrail_name": kwargs.get("guardrail_name")}, + ) + if streaming.streaming_buffer_until_moderated is not None: + self.streaming_buffer_until_moderated: bool = streaming.streaming_buffer_until_moderated + if streaming.streaming_end_of_stream_only is not None: + self.streaming_end_of_stream_only: bool = streaming.streaming_end_of_stream_only + if streaming.streaming_sampling_rate is not None: + self.streaming_sampling_rate: int = streaming.streaming_sampling_rate + if streaming.streaming_buffer_release_on_scan is not None: + self.streaming_buffer_release_on_scan: bool = streaming.streaming_buffer_release_on_scan + kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) base_kwargs: Final[GuardrailBaseInitKwargs] = kwargs diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py index ef22f73810f..992815bf454 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/noma.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/noma.py @@ -1,4 +1,16 @@ -from pydantic import Field +from typing import Final + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + ValidationError, + ValidationInfo, + ValidatorFunctionWrapHandler, + field_validator, +) + +from litellm._logging import verbose_proxy_logger from .base import GuardrailConfigModel @@ -55,3 +67,35 @@ class NomaV2GuardrailConfigModel(GuardrailConfigModel): @staticmethod def ui_friendly_name() -> str: return "Noma Security v2" + + +class NomaV2StreamingParams(BaseModel): + """Streaming settings a Noma v2 guardrail reads from the top level of its litellm_params. + + A key left unset, or set to a value that does not validate, stays None, so the unified + streaming hook keeps its own default for it and the guardrail still loads. + """ + + model_config = ConfigDict(frozen=True) + + streaming_buffer_until_moderated: bool | None = None + streaming_end_of_stream_only: bool | None = None + streaming_sampling_rate: int | None = Field(default=None, ge=1) + streaming_buffer_release_on_scan: bool | None = None + + @field_validator("*", mode="wrap") + @classmethod + def _ignore_invalid_value( + cls, value: object, handler: ValidatorFunctionWrapHandler, info: ValidationInfo + ) -> object: + try: + validated: Final = handler(value) + except ValidationError: + verbose_proxy_logger.warning( + "Noma v2 guardrail %s: ignoring %s=%.100r, which is not a valid value; the default applies", + info.context.get("guardrail_name") if info.context else None, + info.field_name, + value, + ) + return None + return validated diff --git a/tests/integration/streaming/test_noma_v2_streaming_settings.py b/tests/integration/streaming/test_noma_v2_streaming_settings.py new file mode 100644 index 00000000000..6354a2c9135 --- /dev/null +++ b/tests/integration/streaming/test_noma_v2_streaming_settings.py @@ -0,0 +1,1009 @@ +import asyncio +import json +import threading +import time +import uuid +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass, field +from itertools import chain +from pathlib import Path +from typing import Final, Literal, TypeAlias, cast + +import anthropic +import httpx +import openai +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, list_value, object_value, string_value +from integration._support.process import OwnedProxy, WorkerHealthcheck, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +Endpoint: TypeAlias = Literal["chat", "messages", "responses"] +Spelling: TypeAlias = Literal["noma_v2", "noma_use_v2", "stored", "custom"] +Variant: TypeAlias = Literal[ + "default", + "buffered", + "end_only", + "sparse", + "dense", + "release", + "buffered_word", + "sparse_word", + "invalid_rate_buffered", + "invalid_rate_end_only", + "release_invalid", + "unset", + "off", + "off_words", + "invalid", + "release_alone", + "fail_open", + "pre_call", +] +Action: TypeAlias = Literal["pass", "block", "mask", "late", "outage", "slow", "reject"] +Client: TypeAlias = Literal["sync", "async"] +SdkEvent: TypeAlias = tuple[Literal["text", "failed", "status", "error"], str] + +ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +SPELLINGS: Final[tuple[Spelling, ...]] = ("noma_v2", "noma_use_v2", "stored", "custom") +MODELS: Final[Mapping[Endpoint, str]] = { + "chat": "noma-chat", + "messages": "noma-messages", + "responses": "noma-responses", +} +PATHS: Final[Mapping[Endpoint, str]] = { + "chat": "/v1/chat/completions", + "messages": "/v1/messages", + "responses": "/v1/responses", +} +SECRET: Final = "jane@example.com" +MASK: Final = "[EMAIL]" +BLOCKED: Final = "Request blocked by Noma guardrail" +SLOW_SECONDS: Final = 1.5 +SAMPLED_SCANS: Final = range(2, 7) +BUFFERED_BLOCK_STATUS: Final[Mapping[Endpoint, int]] = {"chat": 400, "messages": 400, "responses": 200} +BUFFERED_OUTAGE_STATUS: Final[Mapping[Endpoint, int]] = {"chat": 503, "messages": 503, "responses": 200} +PIECES: Final = ( + " Contact", + " me", + " at", + " work", + " via", + f" {SECRET}", + " for", + " the", + " quarterly", + " report", + " and", + " the", + " budget", + ".", +) +STREAMING_KEYS: Final = ( + "streaming_buffer_until_moderated", + "streaming_end_of_stream_only", + "streaming_sampling_rate", + "streaming_buffer_release_on_scan", +) +INVALID: Final[Mapping[str, JsonValue]] = { + "streaming_buffer_until_moderated": "maybe", + "streaming_end_of_stream_only": "maybe", + "streaming_sampling_rate": 0, + "streaming_buffer_release_on_scan": "maybe", +} +SETTINGS: Final[Mapping[Variant, Mapping[str, JsonValue]]] = { + "default": {}, + "buffered": {"streaming_buffer_until_moderated": True}, + "end_only": {"streaming_end_of_stream_only": True}, + "sparse": {"streaming_sampling_rate": 1000}, + "dense": {"streaming_sampling_rate": 1}, + "release": {"streaming_buffer_until_moderated": True, "streaming_buffer_release_on_scan": True}, + "buffered_word": {"streaming_buffer_until_moderated": "true"}, + "sparse_word": {"streaming_sampling_rate": "1000"}, + "invalid_rate_buffered": {"streaming_buffer_until_moderated": True, "streaming_sampling_rate": -5}, + "invalid_rate_end_only": {"streaming_end_of_stream_only": True, "streaming_sampling_rate": "often"}, + "release_invalid": {"streaming_buffer_until_moderated": True, "streaming_buffer_release_on_scan": "maybe"}, + "unset": dict.fromkeys(STREAMING_KEYS), + "off": { + "streaming_buffer_until_moderated": False, + "streaming_end_of_stream_only": False, + "streaming_buffer_release_on_scan": False, + }, + "off_words": { + "streaming_buffer_until_moderated": "false", + "streaming_end_of_stream_only": "false", + "streaming_buffer_release_on_scan": "false", + }, + "invalid": INVALID, + "release_alone": {"streaming_buffer_release_on_scan": True}, + "fail_open": {"streaming_buffer_until_moderated": True, "block_failures": False}, + "pre_call": {"mode": "pre_call", "streaming_buffer_until_moderated": True}, +} +SUBCLASS_MODULE: Final = "noma_v2_subclass" +SUBCLASS_SOURCE: Final = """from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail + + +class NomaV2Subclass(NomaV2Guardrail): + async def apply_guardrail(self, *args, **kwargs): + return await super().apply_guardrail(*args, **kwargs) +""" +SPELLING_PARAMS: Final[Mapping[Spelling, Mapping[str, JsonValue]]] = { + "noma_v2": {"guardrail": "noma_v2"}, + "noma_use_v2": {"guardrail": "noma", "use_v2": True}, + "stored": {"guardrail": "noma_v2"}, + "custom": {"guardrail": f"{SUBCLASS_MODULE}.NomaV2Subclass"}, +} + + +def _texts(value: JsonValue) -> tuple[str, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(chain.from_iterable(_texts(item) for item in value)) + if isinstance(value, dict): + return tuple(chain.from_iterable(_texts(item) for item in value.values())) + return () + + +def _marker(body: Mapping[str, JsonValue]) -> str: + return next(text for text in _texts(dict(body)) if text.startswith("noma-")) + + +def _full_text(marker: str) -> str: + return marker + "".join(PIECES) + + +def _action(marker: str) -> Action: + return cast(Action, marker.split("-")[1]) + + +def _scan_action(text: str) -> Action: + marker: Final = next((word for word in text.split(" ") if word.startswith("noma-")), None) + return "pass" if marker is None else _action(marker) + + +def _event(name: str | None, payload: Mapping[str, JsonValue]) -> bytes: + prefix: Final = f"event: {name}\n" if name is not None else "" + return f"{prefix}data: {json.dumps(payload, separators=(',', ':'))}\n\n".encode() + + +def _chat_stream(marker: str, choices: int) -> tuple[bytes, ...]: + chunk: Final = {"id": f"chatcmpl-{marker}", "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini"} + deltas: Final = ( + {"role": "assistant", "content": marker}, + *({"content": piece} for piece in PIECES), + ) + return ( + *( + _event(None, {**chunk, "choices": [{"index": index, "delta": delta, "finish_reason": None}]}) + for delta in deltas + for index in range(choices) + ), + _event( + None, + {**chunk, "choices": [{"index": index, "delta": {}, "finish_reason": "stop"} for index in range(choices)]}, + ), + b"data: [DONE]\n\n", + ) + + +def _message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg-{marker}", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5-20250929", + "content": [{"type": "text", "text": _full_text(marker)}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 15}, + } + + +def _messages_stream(marker: str) -> tuple[bytes, ...]: + events: Final[tuple[dict[str, JsonValue], ...]] = ( + { + "type": "message_start", + "message": { + **_message(marker), + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 2, "output_tokens": 0}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + *( + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}} + for text in (marker, *PIECES) + ), + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 15}}, + {"type": "message_stop"}, + ) + return tuple(_event(string_value(event["type"]), event) for event in events) + + +def _response_output(marker: str) -> dict[str, JsonValue]: + return { + "type": "message", + "id": f"msg-{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": _full_text(marker), "annotations": []}], + } + + +def _response(marker: str) -> dict[str, JsonValue]: + return { + "id": f"resp-{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [_response_output(marker)], + "usage": {"input_tokens": 2, "output_tokens": 15, "total_tokens": 17}, + } + + +def _responses_stream(marker: str) -> tuple[bytes, ...]: + item_id: Final = f"msg-{marker}" + text: Final = _full_text(marker) + part: Final[dict[str, JsonValue]] = {"type": "output_text", "text": text, "annotations": []} + output: Final = _response_output(marker) + response: Final = _response(marker) + position: Final[dict[str, JsonValue]] = {"item_id": item_id, "output_index": 0, "content_index": 0} + events: Final[tuple[dict[str, JsonValue], ...]] = ( + {"type": "response.created", "response": {**response, "status": "in_progress", "output": []}}, + {"type": "response.in_progress", "response": {**response, "status": "in_progress", "output": []}}, + { + "type": "response.output_item.added", + "item": {**output, "status": "in_progress", "content": []}, + "output_index": 0, + }, + {"type": "response.content_part.added", **position, "part": {**part, "text": ""}}, + *({"type": "response.output_text.delta", **position, "delta": delta} for delta in (marker, *PIECES)), + {"type": "response.output_text.done", **position, "text": text}, + {"type": "response.content_part.done", **position, "part": part}, + {"type": "response.output_item.done", "item": output, "output_index": 0}, + {"type": "response.completed", "response": response}, + ) + return tuple( + _event(string_value(event["type"]), {**event, "sequence_number": index}) for index, event in enumerate(events) + ) + + +def _completed(target: str, marker: str) -> dict[str, JsonValue]: + if target == PATHS["messages"]: + return _message(marker) + if target == PATHS["responses"]: + return _response(marker) + return { + "id": f"chatcmpl-{marker}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": _full_text(marker)}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 15, "total_tokens": 17}, + } + + +def _rejection(target: str) -> dict[str, JsonValue]: + message: Final = "synthetic upstream rejection" + if target == PATHS["messages"]: + return {"type": "error", "error": {"type": "invalid_request_error", "message": message}} + return {"error": {"message": message, "type": "invalid_request_error", "param": None, "code": None}} + + +def _provider(request: Request) -> Reply: + body: Final = JSON_OBJECT.validate_json(request.body) + marker: Final = _marker(body) + if request.target not in PATHS.values(): + return Reply(status=404, body=json.dumps({"error": f"unexpected target {request.target}"}).encode()) + if _action(marker) == "reject": + return Reply(status=400, body=json.dumps(_rejection(request.target)).encode()) + if body.get("stream") is not True: + return Reply(body=json.dumps(_completed(request.target, marker)).encode()) + choices: Final = body.get("n", 1) + assert isinstance(choices, int), request.body + streams: Final = { + PATHS["chat"]: _chat_stream(marker, choices), + PATHS["messages"]: _messages_stream(marker), + PATHS["responses"]: _responses_stream(marker), + } + return Reply(content_type="text/event-stream", chunks=streams[request.target], pause_between_chunks=0.05) + + +@dataclass(frozen=True, slots=True) +class Scan: + guardrail: str + input_type: str + text: str + + +def _verdict(action: Action, text: str, earlier: int, texts: Sequence[str]) -> dict[str, JsonValue]: + if action == "block" and SECRET in text: + return {"action": "BLOCKED", "reason": "synthetic email policy"} + if action == "late" and earlier > 0: + return {"action": "BLOCKED", "reason": "synthetic late verdict"} + if action == "mask" and SECRET in text: + return {"action": "GUARDRAIL_INTERVENED", "texts": [part.replace(SECRET, MASK) for part in texts]} + return {"action": "NONE"} + + +@dataclass(slots=True) +class NomaDouble: + """Scripted Noma v2 `/litellm/guardrail` endpoint; the action word in the stream's marker picks the verdict.""" + + scans: list[Scan] = field(default_factory=list) + lock: threading.Lock = field(default_factory=threading.Lock) + down: threading.Event = field(default_factory=threading.Event) + + def respond(self, request: Request) -> Reply: + guardrail: Final = request.target.strip("/").split("/")[0] + assert request.target == f"/{guardrail}/litellm/guardrail", request.target + body: Final = JSON_OBJECT.validate_json(request.body) + texts: Final = tuple(string_value(text) for text in list_value(object_value(body["inputs"]).get("texts", []))) + text: Final = "".join(texts) + scan: Final = Scan(guardrail, string_value(body["input_type"]), text) + with self.lock: + earlier: Final = sum(1 for seen in self.scans if seen.text.split(" ")[0] == text.split(" ")[0]) + self.scans.append(scan) + action: Final = _scan_action(text) + if self.down.is_set() or action == "outage": + return Reply(status=503, body=json.dumps({"detail": "synthetic Noma outage"}).encode()) + if action == "slow": + time.sleep(SLOW_SECONDS) + return Reply(body=json.dumps(_verdict(action, text, earlier, texts)).encode()) + + def for_marker(self, marker: str) -> tuple[Scan, ...]: + with self.lock: + return tuple(scan for scan in self.scans if scan.text.startswith(marker)) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: OwnedProxy + noma: NomaDouble + noma_url: str + provider: Wire + names: Mapping[tuple[Spelling, Variant], str] + stored: Mapping[Variant, str] + + @property + def gateway(self) -> Gateway: + return self.proxy.gateway + + +def _guardrail_params(noma_url: str, name: str, spelling: Spelling, variant: Variant) -> dict[str, JsonValue]: + return { + **SPELLING_PARAMS[spelling], + "mode": "post_call", + "default_on": False, + "api_base": f"{noma_url}/{name}", + "api_key": "synthetic-noma-key", + **SETTINGS[variant], + } + + +def _config(provider_url: str, noma_url: str, names: Mapping[tuple[Spelling, Variant], str]) -> dict[str, JsonValue]: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + return { + **base, + "guardrails": [ + {"guardrail_name": name, "litellm_params": _guardrail_params(noma_url, name, spelling, variant)} + for (spelling, variant), name in names.items() + if spelling != "stored" + ], + "model_list": [ + { + "model_name": MODELS["chat"], + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider_url}/v1", + "api_key": "synthetic-provider-key", + }, + }, + { + "model_name": MODELS["messages"], + "litellm_params": { + "model": "anthropic/claude-sonnet-4-5-20250929", + "api_base": provider_url, + "api_key": "synthetic-provider-key", + }, + }, + { + "model_name": MODELS["responses"], + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider_url}/v1", + "api_key": "synthetic-provider-key", + }, + }, + ], + } + + +def _store(gateway: Gateway, noma_url: str, name: str, variant: Variant) -> str: + response: Final = gateway.request( + "POST", + "/guardrails", + {"guardrail": {"guardrail_name": name, "litellm_params": _guardrail_params(noma_url, name, "stored", variant)}}, + ) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["guardrail_id"]) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + noma: Final = NomaDouble() + run: Final = uuid.uuid4().hex[:8] + names: Final = { + (spelling, variant): f"{spelling}-{variant}-{run}".replace("_", "-") + for spelling in SPELLINGS + for variant in SETTINGS + } + with wire_server(_provider) as provider, wire_server(noma.respond) as noma_wire: + stored: Final = { + variant: _store(session_gateway, noma_wire.url, name, variant) + for (spelling, variant), name in names.items() + if spelling == "stored" + } + try: + directory: Final = tmp_path_factory.mktemp("noma-v2-streaming") + (directory / f"{SUBCLASS_MODULE}.py").write_text(SUBCLASS_SOURCE) + config_path: Final = directory / "noma-v2-streaming.yaml" + config_path.write_text(yaml.safe_dump(_config(provider.url, noma_wire.url, names))) + with owned_proxy_process( + session_gateway, + directory, + {}, + config=config_path, + workers=2, + healthcheck=WorkerHealthcheck.LENIENT, + ) as proxy: + yield Rig(proxy, noma, noma_wire.url, provider, names, stored) + finally: + for identity in stored.values(): + deleted = session_gateway.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +@dataclass(frozen=True, slots=True) +class Streamed: + marker: str + response: httpx.Response + content: str + scans: tuple[Scan, ...] + + +def _event_text(event: Mapping[str, JsonValue], endpoint: Endpoint) -> str: + if endpoint == "responses": + return string_value(event.get("delta", "")) if event.get("type") == "response.output_text.delta" else "" + if endpoint == "messages": + return string_value(object_value(event.get("delta", {})).get("text", "")) + return "".join( + string_value(object_value(object_value(choice).get("delta", {})).get("content", "") or "") + for choice in list_value(event.get("choices", [])) + ) + + +def _line_text(line: str, endpoint: Endpoint) -> str: + return ( + _event_text(JSON_OBJECT.validate_json(line.removeprefix("data: ")), endpoint) + if line.startswith("data: {") + else "" + ) + + +def _body(endpoint: Endpoint, marker: str, guardrail: str, *, stream: bool = True) -> dict[str, JsonValue]: + prompt: Final[dict[str, JsonValue]] = ( + {"input": marker} if endpoint == "responses" else {"messages": [{"role": "user", "content": marker}]} + ) + return { + "model": MODELS[endpoint], + **prompt, + **({"max_tokens": 64} if endpoint == "messages" else {}), + "stream": stream, + "guardrails": [guardrail], + } + + +def _new_marker(action: Action) -> str: + return f"noma-{action}-{uuid.uuid4().hex}" + + +def _send( + rig: Rig, + endpoint: Endpoint, + guardrail: str, + action: Action, + *, + input_type: str = "response", + extra: Mapping[str, JsonValue] | None = None, + via: Gateway | None = None, +) -> Streamed: + marker: Final = _new_marker(action) + response: Final = (via or rig.gateway).request( + "POST", PATHS[endpoint], {**_body(endpoint, marker, guardrail), **(extra or {})} + ) + scans: Final = rig.noma.for_marker(marker) + assert all(scan.guardrail == guardrail and scan.input_type == input_type for scan in scans), scans + content: Final = "".join(_line_text(line, endpoint) for line in response.text.splitlines()) + return Streamed(marker, response, content, scans) + + +def _stream( + rig: Rig, + endpoint: Endpoint, + spelling: Spelling, + variant: Variant, + action: Action, + status: int = 200, + *, + input_type: str = "response", + extra: Mapping[str, JsonValue] | None = None, +) -> Streamed: + streamed: Final = _send(rig, endpoint, rig.names[(spelling, variant)], action, input_type=input_type, extra=extra) + assert streamed.response.status_code == status, streamed.response.text + return streamed + + +def _burst(rig: Rig, cases: Sequence[tuple[Endpoint, Action]], guardrail: str) -> tuple[Streamed, ...]: + with ThreadPoolExecutor(max_workers=len(cases)) as pool: + return tuple(pool.map(lambda case: _send(rig, case[0], guardrail, case[1]), cases)) + + +def _attempt(rig: Rig, endpoint: Endpoint, guardrail: str, action: Action) -> str: + try: + return _send(rig, endpoint, guardrail, action).response.text + except httpx.TransportError as error: + return f"transport error: {type(error).__name__}" + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +@pytest.mark.parametrize("variant", ["buffered", "buffered_word", "invalid_rate_buffered", "release_invalid"]) +def test_buffer_until_moderated_keeps_flagged_text_from_the_client( + rig: Rig, endpoint: Endpoint, spelling: Spelling, variant: Variant +) -> None: + unbuffered: Final = _stream(rig, endpoint, spelling, "default", "block") + assert SECRET in unbuffered.content, unbuffered.response.text + assert BLOCKED in unbuffered.response.text, unbuffered.response.text + + buffered: Final = _stream(rig, endpoint, spelling, variant, "block", BUFFERED_BLOCK_STATUS[endpoint]) + assert buffered.content == "", buffered.response.text + assert SECRET not in buffered.response.text, buffered.response.text + assert BLOCKED in buffered.response.text, buffered.response.text + assert [scan.text for scan in buffered.scans] == [_full_text(buffered.marker)], buffered.scans + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +@pytest.mark.parametrize("variant", ["end_only", "sparse", "sparse_word", "invalid_rate_end_only"]) +def test_one_scan_settings_scan_the_finished_stream_once( + rig: Rig, endpoint: Endpoint, spelling: Spelling, variant: Variant +) -> None: + sampled: Final = _stream(rig, endpoint, spelling, "default", "pass") + assert sampled.content == _full_text(sampled.marker), sampled.response.text + assert len(sampled.scans) > 1, sampled.scans + + configured: Final = _stream(rig, endpoint, spelling, variant, "pass") + assert configured.content == _full_text(configured.marker), configured.response.text + assert [scan.text for scan in configured.scans] == [_full_text(configured.marker)], configured.scans + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_sampling_rate_one_scans_every_new_piece(rig: Rig, endpoint: Endpoint, spelling: Spelling) -> None: + sampled: Final = _stream(rig, endpoint, spelling, "default", "pass") + assert len(sampled.scans) < len(PIECES), sampled.scans + + dense: Final = _stream(rig, endpoint, spelling, "dense", "pass") + assert dense.content == _full_text(dense.marker), dense.response.text + scanned: Final = {scan.text for scan in dense.scans} + prefixes: Final = {dense.marker + "".join(PIECES[:count]) for count in range(len(PIECES) + 1)} + assert scanned <= prefixes, sorted(scanned - prefixes) + if endpoint == "messages": + # /v1/messages hands the hook raw network reads, so two deltas read together are one chunk + assert len(scanned) > max(SAMPLED_SCANS), sorted(scanned) + else: + assert prefixes <= scanned, sorted(prefixes - scanned) + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_release_on_scan_releases_only_scanned_windows(rig: Rig, endpoint: Endpoint, spelling: Spelling) -> None: + released: Final = _stream(rig, endpoint, spelling, "release", "late") + assert len(released.scans) == 2, released.scans + assert released.content == released.scans[0].text, (released.response.text, released.scans) + assert released.content and _full_text(released.marker).startswith(released.content), released.response.text + assert BLOCKED in released.response.text, released.response.text + + held: Final = _stream(rig, endpoint, spelling, "buffered", "late") + assert held.content == _full_text(held.marker), held.response.text + assert len(held.scans) == 1, held.scans + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_invalid_release_setting_keeps_the_buffered_stream_held( + rig: Rig, endpoint: Endpoint, spelling: Spelling +) -> None: + held: Final = _stream(rig, endpoint, spelling, "release_invalid", "late") + assert held.content == _full_text(held.marker), held.response.text + assert len(held.scans) == 1, held.scans + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +@pytest.mark.parametrize("variant", ["unset", "off", "off_words", "invalid", "release_alone"]) +def test_unset_off_and_invalid_settings_keep_the_sampled_live_stream( + rig: Rig, endpoint: Endpoint, spelling: Spelling, variant: Variant +) -> None: + configured: Final = _stream(rig, endpoint, spelling, variant, "pass") + assert configured.content == _full_text(configured.marker), configured.response.text + assert len(configured.scans) in SAMPLED_SCANS, configured.scans + assert configured.scans[-1].text == _full_text(configured.marker), configured.scans + + leaked: Final = _stream(rig, endpoint, spelling, variant, "block") + assert SECRET in leaked.content, leaked.response.text + assert BLOCKED in leaked.response.text, leaked.response.text + + +@pytest.mark.parametrize("spelling", SPELLINGS) +def test_invalid_settings_are_logged_with_the_guardrail_name(rig: Rig, spelling: Spelling) -> None: + name: Final = rig.names[(spelling, "invalid")] + log: Final = rig.proxy.log.read_text() + expected: Final = tuple(f"Noma v2 guardrail {name}: ignoring {key}={value!r}" for key, value in INVALID.items()) + assert [line for line in expected if line not in log] == [], name + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_noma_outage_fails_closed_before_any_text_when_buffered( + rig: Rig, endpoint: Endpoint, spelling: Spelling +) -> None: + closed: Final = _stream(rig, endpoint, spelling, "buffered", "outage", BUFFERED_OUTAGE_STATUS[endpoint]) + assert closed.content == "", closed.response.text + assert endpoint != "responses" or "response.failed" in closed.response.text, closed.response.text + assert [scan.text for scan in closed.scans] == [_full_text(closed.marker)], closed.scans + + opened: Final = _stream(rig, endpoint, spelling, "fail_open", "outage") + assert opened.content == _full_text(opened.marker), opened.response.text + assert [scan.text for scan in opened.scans] == [_full_text(opened.marker)], opened.scans + + +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_upstream_rejection_reaches_the_client_without_a_scan(rig: Rig, endpoint: Endpoint) -> None: + rejected: Final = _stream(rig, endpoint, "noma_v2", "buffered", "reject", 400) + assert "synthetic upstream rejection" in rejected.response.text, rejected.response.text + assert rejected.scans == (), rejected.scans + + +@pytest.mark.parametrize("spelling", SPELLINGS) +def test_buffered_block_withholds_every_choice(rig: Rig, spelling: Spelling) -> None: + unbuffered: Final = _stream(rig, "chat", spelling, "default", "block", extra={"n": 2}) + assert SECRET in unbuffered.content, unbuffered.response.text + + buffered: Final = _stream(rig, "chat", spelling, "buffered", "block", 400, extra={"n": 2}) + assert buffered.content == "", buffered.response.text + assert SECRET not in buffered.response.text, buffered.response.text + assert BLOCKED in buffered.response.text, buffered.response.text + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_pre_call_guardrail_scans_the_request_and_streams_live( + rig: Rig, endpoint: Endpoint, spelling: Spelling +) -> None: + streamed: Final = _stream(rig, endpoint, spelling, "pre_call", "pass", input_type="request") + assert streamed.content == _full_text(streamed.marker), streamed.response.text + assert [scan.text for scan in streamed.scans] == [streamed.marker], streamed.scans + + +@pytest.mark.parametrize("endpoint", ENDPOINTS) +@pytest.mark.parametrize("variant", ["default", "buffered", "end_only"]) +def test_non_streaming_post_call_ignores_the_streaming_settings(rig: Rig, endpoint: Endpoint, variant: Variant) -> None: + guardrail: Final = rig.names[("noma_v2", variant)] + passed: Final = _send(rig, endpoint, guardrail, "pass", extra={"stream": False}) + assert passed.response.status_code == 200, passed.response.text + assert _full_text(passed.marker) in passed.response.text, passed.response.text + assert [scan.text for scan in passed.scans] == [_full_text(passed.marker)], passed.scans + + blocked: Final = _send(rig, endpoint, guardrail, "block", extra={"stream": False}) + assert blocked.response.status_code == 400, blocked.response.text + assert SECRET not in blocked.response.text, blocked.response.text + assert BLOCKED in blocked.response.text, blocked.response.text + + +@dataclass(frozen=True, slots=True) +class SdkOutcome: + text: str + status: int | None + failed: bool + + +def _outcome(events: Sequence[SdkEvent]) -> SdkOutcome: + return SdkOutcome( + "".join(text for kind, text in events if kind == "text"), + next((int(text) for kind, text in events if kind == "status"), None), + any(kind in ("failed", "error") for kind, _ in events), + ) + + +def _openai_events(client: openai.OpenAI, endpoint: Endpoint, marker: str, guardrail: str) -> Iterator[SdkEvent]: + extra: Final[dict[str, object]] = {"guardrails": [guardrail]} + if endpoint == "chat": + for chunk in client.chat.completions.create( + model=MODELS["chat"], messages=[{"role": "user", "content": marker}], stream=True, extra_body=extra + ): + yield from (("text", choice.delta.content or "") for choice in chunk.choices) + return + for event in client.responses.create(model=MODELS["responses"], input=marker, stream=True, extra_body=extra): + if event.type == "response.output_text.delta": + yield "text", event.delta + if event.type == "response.failed": + yield "failed", event.type + + +async def _async_openai_events( + client: openai.AsyncOpenAI, endpoint: Endpoint, marker: str, guardrail: str +) -> AsyncIterator[SdkEvent]: + extra: Final[dict[str, object]] = {"guardrails": [guardrail]} + if endpoint == "chat": + async for chunk in await client.chat.completions.create( + model=MODELS["chat"], messages=[{"role": "user", "content": marker}], stream=True, extra_body=extra + ): + for choice in chunk.choices: + yield "text", choice.delta.content or "" + return + async for event in await client.responses.create( + model=MODELS["responses"], input=marker, stream=True, extra_body=extra + ): + if event.type == "response.output_text.delta": + yield "text", event.delta + if event.type == "response.failed": + yield "failed", event.type + + +def _anthropic_events(client: anthropic.Anthropic, marker: str, guardrail: str) -> Iterator[SdkEvent]: + for event in client.messages.create( + model=MODELS["messages"], + max_tokens=64, + messages=[{"role": "user", "content": marker}], + stream=True, + extra_body={"guardrails": [guardrail]}, + ): + if event.type == "content_block_delta" and event.delta.type == "text_delta": + yield "text", event.delta.text + + +async def _async_anthropic_events( + client: anthropic.AsyncAnthropic, marker: str, guardrail: str +) -> AsyncIterator[SdkEvent]: + async for event in await client.messages.create( + model=MODELS["messages"], + max_tokens=64, + messages=[{"role": "user", "content": marker}], + stream=True, + extra_body={"guardrails": [guardrail]}, + ): + if event.type == "content_block_delta" and event.delta.type == "text_delta": + yield "text", event.delta.text + + +def _sync_sdk(gateway: Gateway, endpoint: Endpoint, marker: str, guardrail: str) -> SdkOutcome: + origin: Final = str(gateway.client.base_url).rstrip("/") + + def events() -> Iterator[SdkEvent]: + try: + with httpx.Client(timeout=30, trust_env=False) as transport: + if endpoint == "messages": + with anthropic.Anthropic( + api_key=gateway.key, base_url=origin, max_retries=0, http_client=transport + ) as client: + yield from _anthropic_events(client, marker, guardrail) + return + with openai.OpenAI( + api_key=gateway.key, base_url=f"{origin}/v1", max_retries=0, http_client=transport + ) as client: + yield from _openai_events(client, endpoint, marker, guardrail) + except (openai.APIStatusError, anthropic.APIStatusError) as error: + yield "status", str(error.status_code) + except (openai.APIError, anthropic.APIError) as error: + yield "error", error.message + + return _outcome(tuple(events())) + + +def _async_sdk(gateway: Gateway, endpoint: Endpoint, marker: str, guardrail: str) -> SdkOutcome: + origin: Final = str(gateway.client.base_url).rstrip("/") + + async def events() -> AsyncIterator[SdkEvent]: + try: + async with httpx.AsyncClient(timeout=30, trust_env=False) as transport: + if endpoint == "messages": + async with anthropic.AsyncAnthropic( + api_key=gateway.key, base_url=origin, max_retries=0, http_client=transport + ) as client: + async for event in _async_anthropic_events(client, marker, guardrail): + yield event + return + async with openai.AsyncOpenAI( + api_key=gateway.key, base_url=f"{origin}/v1", max_retries=0, http_client=transport + ) as client: + async for event in _async_openai_events(client, endpoint, marker, guardrail): + yield event + except (openai.APIStatusError, anthropic.APIStatusError) as error: + yield "status", str(error.status_code) + except (openai.APIError, anthropic.APIError) as error: + yield "error", error.message + + async def collect() -> tuple[SdkEvent, ...]: + return tuple([event async for event in events()]) + + return _outcome(asyncio.run(collect())) + + +@pytest.mark.parametrize("client", ["sync", "async"]) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_sdk_clients_receive_no_text_before_a_buffered_block(rig: Rig, endpoint: Endpoint, client: Client) -> None: + marker: Final = _new_marker("block") + outcome: Final = {"sync": _sync_sdk, "async": _async_sdk}[client]( + rig.gateway, endpoint, marker, rig.names[("noma_v2", "buffered")] + ) + assert outcome.text == "", outcome + if endpoint == "responses": + assert outcome.failed and outcome.status is None, outcome + else: + assert outcome.status == 400, outcome + assert [scan.text for scan in rig.noma.for_marker(marker)] == [_full_text(marker)] + + +@pytest.mark.parametrize("variant", ["buffered", "invalid", "unset"]) +def test_stored_guardrail_info_reports_the_settings_as_stored(rig: Rig, variant: Variant) -> None: + info: Final = rig.gateway.get(f"/guardrails/{rig.stored[variant]}/info") + params: Final = object_value(info["litellm_params"]) + assert {key: params.get(key) for key in SETTINGS[variant]} == dict(SETTINGS[variant]), params + + +def _on_one_worker(rig: Rig, pid: int, guardrail: str) -> tuple[Streamed, ...]: + """Two blocked-prompt streams answered by worker `pid` alone, every other worker suspended meanwhile.""" + with ( + ExitStack() as paused, + httpx.Client( + base_url=str(rig.gateway.client.base_url), + timeout=15, + trust_env=False, + limits=httpx.Limits(max_keepalive_connections=0), + ) as client, + ): + for other in sorted(rig.proxy.worker_pids() - {pid}): + paused.enter_context(rig.proxy.suspended(other)) + alone: Final = Gateway(client, rig.gateway.key, rig.gateway.upstream_url) + return tuple(_send(rig, "chat", guardrail, "block", via=alone) for _ in range(2)) + + +def _converged_on(rig: Rig, pid: int, guardrail: str, satisfied: Callable[[Streamed], bool]) -> tuple[Streamed, ...]: + return eventually( + lambda: _on_one_worker(rig, pid, guardrail), + lambda attempts: all(satisfied(attempt) for attempt in attempts), + seconds=120, + ) + + +def test_patching_a_stored_guardrail_applies_the_setting_on_every_worker(rig: Rig) -> None: + name: Final = f"stored-patched-{uuid.uuid4().hex[:8]}" + workers: Final = sorted(eventually(rig.proxy.worker_pids, lambda found: len(found) == 2, seconds=30)) + identity: Final = _store(rig.gateway, rig.noma_url, name, "default") + try: + before: Final = tuple( + chain.from_iterable( + _converged_on(rig, pid, name, lambda attempt: len(attempt.scans) > 0) for pid in workers + ) + ) + assert all(SECRET in attempt.content for attempt in before), [attempt.response.text for attempt in before] + + patched: Final = rig.gateway.request( + "PATCH", f"/guardrails/{identity}", {"litellm_params": {"streaming_buffer_until_moderated": True}} + ) + assert patched.status_code == 200, patched.text + after: Final = tuple( + chain.from_iterable( + _converged_on(rig, pid, name, lambda attempt: attempt.response.status_code == 400) for pid in workers + ) + ) + assert all(attempt.content == "" and SECRET not in attempt.response.text for attempt in after), [ + attempt.response.text for attempt in after + ] + finally: + deleted: Final = rig.gateway.request("DELETE", f"/guardrails/{identity}") + assert deleted.status_code == 200, deleted.text + + +def test_burst_fails_closed_while_noma_is_down_and_recovers_after(rig: Rig) -> None: + guardrail: Final = rig.names[("noma_v2", "buffered")] + cases: Final = tuple( + (endpoint, action) + for endpoint in ENDPOINTS + for action in cast(tuple[Action, ...], ("pass", "block")) + for _ in range(5) + ) + rig.noma.down.set() + try: + down: Final = _burst(rig, cases, guardrail) + finally: + rig.noma.down.clear() + assert [streamed.response.text for streamed in down if streamed.content or SECRET in streamed.response.text] == [] + assert [streamed.scans for streamed in down if len(streamed.scans) != 1] == [] + + up: Final = _burst(rig, cases, guardrail) + delivered: Final = tuple(streamed for streamed in up if _action(streamed.marker) == "pass") + blocked: Final = tuple(streamed for streamed in up if _action(streamed.marker) == "block") + assert [s.response.text for s in delivered if s.content != _full_text(s.marker)] == [] + assert [ + s.response.text for s in blocked if s.content or SECRET in s.response.text or BLOCKED not in s.response.text + ] == [] + assert [streamed.scans for streamed in up if len(streamed.scans) != 1] == [] + + +def _first_text_seconds(rig: Rig, endpoint: Endpoint, guardrail: str) -> tuple[str, float | None, str]: + marker: Final = _new_marker("slow") + started: Final = time.monotonic() + with rig.gateway.client.stream( + "POST", + PATHS[endpoint], + json=_body(endpoint, marker, guardrail), + headers={"Authorization": f"Bearer {rig.gateway.key}"}, + timeout=30, + ) as response: + lines: Final = tuple((time.monotonic() - started, line) for line in response.iter_lines()) + texts: Final = tuple((elapsed, _line_text(line, endpoint)) for elapsed, line in lines) + return marker, next((elapsed for elapsed, text in texts if text), None), "".join(text for _, text in texts) + + +def test_slow_noma_holds_every_stream_until_its_verdict(rig: Rig) -> None: + guardrail: Final = rig.names[("noma_v2", "buffered")] + with ThreadPoolExecutor(max_workers=18) as pool: + streams: Final = tuple(pool.map(lambda endpoint: _first_text_seconds(rig, endpoint, guardrail), ENDPOINTS * 6)) + assert [stream for stream in streams if stream[2] != _full_text(stream[0])] == [] + assert [stream for stream in streams if stream[1] is None or stream[1] < SLOW_SECONDS] == [] + assert [marker for marker, _, _ in streams if len(rig.noma.for_marker(marker)) != 1] == [] + + +def test_replaced_workers_keep_the_streaming_settings(rig: Rig) -> None: + guardrail: Final = rig.names[("noma_v2", "buffered")] + originals: Final = eventually(rig.proxy.worker_pids, lambda found: len(found) == 2, seconds=30) + with ThreadPoolExecutor(max_workers=8) as pool: + inflight: Final = tuple( + pool.submit(_attempt, rig, ENDPOINTS[index % 3], guardrail, "block") for index in range(16) + ) + rig.proxy.kill_worker(min(originals)) + during: Final = tuple(future.result() for future in inflight) + rig.proxy.kill_worker(max(originals)) + assert [text for text in during if SECRET in text] == [] + assert rig.proxy.worker_pids().isdisjoint(originals), rig.proxy.worker_pids() + + after: Final = tuple( + _stream(rig, endpoint, "noma_v2", "buffered", "block", BUFFERED_BLOCK_STATUS[endpoint]) + for endpoint in ENDPOINTS + for _ in range(4) + ) + assert [streamed.response.text for streamed in after if streamed.content or SECRET in streamed.response.text] == [] + assert [streamed.scans for streamed in after if len(streamed.scans) != 1] == [] + + +@pytest.mark.parametrize("spelling", SPELLINGS) +@pytest.mark.parametrize("endpoint", ENDPOINTS) +def test_buffer_until_moderated_delivers_the_masked_text(rig: Rig, endpoint: Endpoint, spelling: Spelling) -> None: + pytest.skip("BUG: LIT-9559 a buffered stream releases the original chunks instead of the guardrail's masked text") + unbuffered: Final = _stream(rig, endpoint, spelling, "default", "mask") + assert SECRET in unbuffered.content, unbuffered.response.text + + buffered: Final = _stream(rig, endpoint, spelling, "buffered", "mask") + assert buffered.content == _full_text(buffered.marker).replace(SECRET, MASK), buffered.response.text + assert SECRET not in buffered.response.text, buffered.response.text diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/noma/test_noma_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/noma/test_noma_v2.py index 6e536f95251..8aca616faed 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/noma/test_noma_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/noma/test_noma_v2.py @@ -1,17 +1,30 @@ import json +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import TypeAlias import httpx import pytest import respx import litellm +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_hooks.noma import ( NomaV2Guardrail, guardrail_initializer_registry, ) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import UnifiedLLMGuardrails +from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler from litellm.types.guardrails import LitellmParams _API_BASE = "https://noma.example.test" +_STREAMING_KEYS = ( + "streaming_buffer_until_moderated", + "streaming_end_of_stream_only", + "streaming_sampling_rate", + "streaming_buffer_release_on_scan", +) @pytest.fixture(autouse=True) @@ -97,3 +110,111 @@ async def test_positional_args_keep_their_meaning_after_gateway_name_was_added(r await guardrail.apply_guardrail( inputs={"texts": ["hello"]}, request_data={"metadata": {}}, input_type="request" ) + + +_Loader: TypeAlias = Callable[[Path, Mapping[str, object]], CustomGuardrail | None] +_SUBCLASS_MODULE = ( + "from litellm.proxy.guardrails.guardrail_hooks.noma.noma_v2 import NomaV2Guardrail\n\n\n" + "class SubclassedNoma(NomaV2Guardrail):\n" + " pass\n" +) + + +def _initializer(guardrail_type: str, fixed: Mapping[str, object]) -> _Loader: + def load(_tmp_path: Path, params: Mapping[str, object]) -> CustomGuardrail | None: + litellm_params = LitellmParams( + guardrail=guardrail_type, mode="post_call", api_base=_API_BASE, **fixed, **params + ) + return guardrail_initializer_registry[guardrail_type](litellm_params, {"guardrail_name": "noma-guard"}) + + return load + + +def _custom_subclass(tmp_path: Path, params: Mapping[str, object]) -> CustomGuardrail | None: + (tmp_path / "noma_subclass.py").write_text(_SUBCLASS_MODULE) + guardrail_type = "noma_subclass.SubclassedNoma" + return InMemoryGuardrailHandler().initialize_custom_guardrail( + guardrail={"guardrail_name": "noma-guard"}, + guardrail_type=guardrail_type, + litellm_params=LitellmParams(guardrail=guardrail_type, mode="post_call", api_base=_API_BASE, **params), + config_file_path=str(tmp_path / "config.yaml"), + ) + + +_LOADERS = pytest.mark.parametrize( + "load", + [_initializer("noma_v2", {}), _initializer("noma", {"use_v2": True}), _custom_subclass], + ids=["noma_v2", "noma-use_v2", "custom-subclass"], +) + + +def _streaming_flags(load: _Loader, tmp_path: Path, params: Mapping[str, object]) -> list[object]: + guardrail = load(tmp_path, params) + return [UnifiedLLMGuardrails().resolve_streaming_flag(guardrail, key, "hook-default") for key in _STREAMING_KEYS] + + +@_LOADERS +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ( + { + "streaming_buffer_until_moderated": True, + "streaming_end_of_stream_only": True, + "streaming_sampling_rate": 3, + "streaming_buffer_release_on_scan": True, + }, + [True, True, 3, True], + ), + ( + { + "streaming_buffer_until_moderated": False, + "streaming_end_of_stream_only": False, + "streaming_sampling_rate": 1, + "streaming_buffer_release_on_scan": False, + }, + [False, False, 1, False], + ), + ({"streaming_sampling_rate": 7}, ["hook-default", "hook-default", 7, "hook-default"]), + ( + {"streaming_buffer_until_moderated": "false", "streaming_sampling_rate": "2"}, + [False, "hook-default", 2, "hook-default"], + ), + ({}, ["hook-default"] * 4), + ], + ids=["all-on", "all-off", "one-key", "string-values", "unset"], +) +def test_streaming_settings_in_guardrail_config_reach_the_streaming_hook( + load: _Loader, tmp_path: Path, configured: dict[str, object], expected: list[object] +) -> None: + assert _streaming_flags(load, tmp_path, configured) == expected + + +@_LOADERS +@pytest.mark.parametrize( + ("configured", "expected"), + [ + ({"streaming_sampling_rate": 0}, ["hook-default"] * 4), + ({"streaming_sampling_rate": -1}, ["hook-default"] * 4), + ({"streaming_sampling_rate": "often"}, ["hook-default"] * 4), + ({"streaming_buffer_until_moderated": "sometimes"}, ["hook-default"] * 4), + ( + {"streaming_buffer_until_moderated": True, "streaming_sampling_rate": 0}, + [True, "hook-default", "hook-default", "hook-default"], + ), + ], + ids=["rate-zero", "rate-negative", "rate-word", "buffer-word", "valid-key-kept"], +) +def test_invalid_streaming_value_is_ignored_and_the_guardrail_still_loads( + load: _Loader, + tmp_path: Path, + configured: dict[str, object], + expected: list[object], + caplog: pytest.LogCaptureFixture, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(verbose_proxy_logger, "propagate", True) + assert _streaming_flags(load, tmp_path, configured) == expected + warnings = [record.getMessage() for record in caplog.records if "not a valid value" in record.getMessage()] + assert warnings, caplog.text + assert all(message.startswith("Noma v2 guardrail noma-guard: ignoring ") for message in warnings), warnings