From 8cc77672ed9c22f138b37bdb73ace1c41bc8b638 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 05:43:23 -0700 Subject: [PATCH 1/9] fix(proxy): refuse response ids the managed OpenAI pass-through never issued With passthrough_managed_object_ids on, the OpenAI and Azure pass-through Responses routes forwarded a raw resp_ id that the pass-through never handed out (retrieve, input_items, cancel, delete, and previous_response_id on create and compact). They now answer the same 403 the native routes give since the fail-closed rule, with the same allow_unmanaged_response_ids, disable_responses_id_security, and proxy admin exemptions. Ids the pass-through issued keep working for their owner. --- litellm/proxy/hooks/responses_id_security.py | 35 +++-- .../managed_id_rewriter.py | 41 ++++- .../pass_through_endpoints.py | 17 ++ .../test_managed_id_rewriter.py | 100 ++++++++++++ .../test_pass_through_endpoints.py | 147 +++++++++++++++++- 5 files changed, 321 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index e554512ec91..9f95ec5eab7 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -32,7 +32,7 @@ if TYPE_CHECKING: _RESPONSES_API_PROVIDER_PREFIX: Final = "/openai" _RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"}) -_UNMANAGED_RESPONSE_ID_DETAIL: Final = ( +UNMANAGED_RESPONSE_ID_DETAIL: Final = ( "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. " "To let keys address responses this proxy did not issue, set " "general_settings::allow_unmanaged_response_ids to True in the config.yaml file." @@ -46,7 +46,7 @@ def _proxy_general_settings() -> Mapping[str, object]: return general_settings -def _proxy_signing_key() -> str | None: +def proxy_signing_key() -> str | None: import os from litellm.proxy.proxy_server import master_key @@ -55,6 +55,20 @@ def _proxy_signing_key() -> str | None: return master_key if salt_key is None else salt_key +def unmanaged_response_ids_allowed( + general_settings: Mapping[str, object], + signing_key: str | None, + user_api_key_dict: "UserAPIKeyAuth", +) -> bool: + if general_settings.get("disable_responses_id_security", False): + return True + if general_settings.get("allow_unmanaged_response_ids", False): + return True + if signing_key is None: + return True + return user_api_key_dict.user_role in _PROXY_ADMIN_ROLES + + _RESPONSE_PAYLOAD_ADAPTER: Final = TypeAdapter(Mapping[str, object]) @@ -108,7 +122,7 @@ class ResponsesIDSecurity(CustomLogger): def __init__( self, general_settings_reader: Callable[[], Mapping[str, object]] = _proxy_general_settings, - signing_key_reader: Callable[[], str | None] = _proxy_signing_key, + signing_key_reader: Callable[[], str | None] = proxy_signing_key, ) -> None: self._general_settings_reader: Final = general_settings_reader self._signing_key_reader: Final = signing_key_reader @@ -152,21 +166,10 @@ class ResponsesIDSecurity(CustomLogger): self.check_user_access_to_response_id(user_id, team_id, user_api_key_dict) return original_response_id - if self._unmanaged_response_ids_allowed(user_api_key_dict): + if unmanaged_response_ids_allowed(self._general_settings_reader(), self._get_signing_key(), user_api_key_dict): return response_id - raise HTTPException(status_code=403, detail=_UNMANAGED_RESPONSE_ID_DETAIL) - - def _unmanaged_response_ids_allowed(self, user_api_key_dict: "UserAPIKeyAuth") -> bool: - general_settings: Final = self._general_settings_reader() - - if general_settings.get("disable_responses_id_security", False): - return True - if general_settings.get("allow_unmanaged_response_ids", False): - return True - if self._get_signing_key() is None: - return True - return user_api_key_dict.user_role in _PROXY_ADMIN_ROLES + raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL) def check_user_access_to_response_id( self, diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index a5a04fb68d2..cb8fcc6badc 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -33,6 +33,7 @@ from __future__ import annotations import json import re from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from itertools import pairwise from typing import ( TYPE_CHECKING, Final, @@ -178,8 +179,10 @@ BUILTIN_OUTPUT_ID_FIELD_MAP: Final[dict[_MapKey, list[_FieldSpec]]] = { # Prefixes that live in the *file* table rather than the object table. _FILE_PREFIXES: Final[frozenset[str]] = frozenset({"file-"}) +_RESPONSE_ID_PREFIX: Final = "resp_" + # Raw provider-ID prefixes that live in the object table (batches, responses). -_OBJECT_PREFIXES: Final[frozenset[str]] = frozenset({"batch_", "resp_"}) +_OBJECT_PREFIXES: Final[frozenset[str]] = frozenset({"batch_", _RESPONSE_ID_PREFIX}) # Guards request-body rewriting against stack exhaustion from adversarially # deep payloads. Real OpenAI files/batches bodies nest only a few levels. @@ -466,6 +469,41 @@ async def _guard_raw_provider_id( raise HTTPException(status_code=404, detail="Managed resource not found.") +def _addressed_raw_response_ids(path: str, body: Mapping[str, object]) -> frozenset[str]: + segments: Final = tuple(unquote(seg) for seg in path.split("/")) + path_ids: Final = frozenset( + seg for parent, seg in pairwise(segments) if parent == "responses" and seg.startswith(_RESPONSE_ID_PREFIX) + ) + previous_response_id: Final = body.get("previous_response_id") + if isinstance(previous_response_id, str) and previous_response_id.startswith(_RESPONSE_ID_PREFIX): + return path_ids | {previous_response_id} + return path_ids + + +async def _passthrough_issued_response_id(raw_id: str, provider: str, prisma_client: PrismaClient) -> bool: + row: Final = await _object_table(prisma_client).find_first( + where={"model_object_id": f"passthrough:{provider}:{raw_id}"} + ) + return row is not None + + +async def addresses_unissued_response_id( + path: str, + body: Mapping[str, object], + provider: str, + prisma_client: PrismaClient | None, +) -> bool: + addressed: Final = _addressed_raw_response_ids(path, body) + if not addressed: + return False + if prisma_client is None: + return True + for raw_id in addressed: + if not await _passthrough_issued_response_id(raw_id, provider, prisma_client): + return True + return False + + # --------------------------------------------------------------------------- # OUTPUT path — helpers for minting and storing managed IDs # --------------------------------------------------------------------------- @@ -817,7 +855,6 @@ async def rewrite_response_ids( return mutated if changed else body -_RESPONSE_ID_PREFIX: Final = "resp_" _STREAMED_RESPONSE_ID_SPEC: Final[_FieldSpec] = ("id", _RESPONSE_ID_PREFIX) _SSE_DATA_PREFIX: Final = "data:" _SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue]) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1663907b556..6449b67b903 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -160,6 +160,7 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging() _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) +_REQUEST_BODY_VIEW: Final = TypeAdapter(Mapping[str, object]) # Global registry to track registered pass-through routes and prevent memory leaks _registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | Mapping[str, object]]]] = {} @@ -1339,7 +1340,13 @@ async def pass_through_request( ) _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook("managed_files") if _passthrough_managed_hook is not None: + from litellm.proxy.hooks.responses_id_security import ( + UNMANAGED_RESPONSE_ID_DETAIL, + proxy_signing_key, + unmanaged_response_ids_allowed, + ) from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( + addresses_unissued_response_id, rewrite_body_ids, rewrite_path_ids, rewrite_query_ids, @@ -1348,6 +1355,16 @@ async def pass_through_request( prisma_client as _passthrough_prisma, ) + if not unmanaged_response_ids_allowed( + general_settings_view(), proxy_signing_key(), user_api_key_dict + ) and await addresses_unissued_response_id( + url.path, + _REQUEST_BODY_VIEW.validate_python(_parsed_body or {}), + _managed_id_provider, + _passthrough_prisma, + ): + raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL) + _original_path: Final = url.path _original_query_params: Final = requested_query_params _original_body: Final = _parsed_body diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index 49bc2e9c29e..ba3399107cf 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -17,6 +17,7 @@ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import( _MAX_RAW_ID_GUARD_LOOKUPS, _passthrough_provider_marker, _resolve_one, + addresses_unissued_response_id, is_passthrough_list_route, list_passthrough_ids_from_db, rewrite_body_ids, @@ -1709,6 +1710,105 @@ class TestRawProviderIdGuardBudget: assert exc_info.value.status_code == 400 assert pc.db.litellm_managedfiletable.find_many.call_count == _MAX_RAW_ID_GUARD_LOOKUPS + +UNISSUED_RESPONSE_ID = "resp_0000000000000000unissued" + + +def _object_lookup(row: object) -> MagicMock: + pc = _prisma_client_passthrough_managed() + pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=row) + return pc + + +class TestAddressesUnissuedResponseId: + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("path", "body"), + [ + (f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}), + (f"/v1/responses/{UNISSUED_RESPONSE_ID}/input_items", {}), + (f"/v1/responses/{UNISSUED_RESPONSE_ID}/cancel", {}), + (f"/openai/responses/{UNISSUED_RESPONSE_ID}", {}), + ("/v1/responses/compact", {"model": "gpt-5.6", "previous_response_id": UNISSUED_RESPONSE_ID}), + ("/v1/responses", {"model": "gpt-5.6", "previous_response_id": UNISSUED_RESPONSE_ID}), + ], + ) + async def test_response_id_without_a_passthrough_row_is_unissued(self, path, body): + pc = _object_lookup(None) + + assert await addresses_unissued_response_id(path, body, "openai", pc) is True + pc.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( + where={"model_object_id": f"passthrough:openai:{UNISSUED_RESPONSE_ID}"} + ) + + @pytest.mark.asyncio + async def test_response_id_with_a_passthrough_row_is_issued(self): + pc = _object_lookup(_batch_row(encode("azure", "user-1", UNISSUED_RESPONSE_ID))) + + assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "azure", pc) is False + pc.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( + where={"model_object_id": f"passthrough:azure:{UNISSUED_RESPONSE_ID}"} + ) + + @pytest.mark.asyncio + async def test_any_unissued_id_among_several_is_unissued(self): + issued_id = "resp_1111111111111111issued" + pc = _prisma_client_passthrough_managed() + pc.db.litellm_managedobjecttable.find_first = AsyncMock( + side_effect=lambda where: _batch_row("x") if where["model_object_id"].endswith(issued_id) else None + ) + + assert ( + await addresses_unissued_response_id( + f"/v1/responses/{issued_id}", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", pc + ) + is True + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("path", "body"), + [ + ("/v1/responses", {"model": "gpt-5.6", "input": f"what does {UNISSUED_RESPONSE_ID} mean?"}), + ("/v1/responses", {"input": [{"role": "user", "content": UNISSUED_RESPONSE_ID}]}), + (f"/v1/files/{UNISSUED_RESPONSE_ID}", {}), + ("/v1/responses", {"previous_response_id": "msg_not_a_response"}), + ("/v1/responses", {}), + ], + ) + async def test_requests_that_address_no_response_id_skip_the_lookup(self, path, body): + pc = _object_lookup(None) + + assert await addresses_unissued_response_id(path, body, "openai", pc) is False + pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + @pytest.mark.asyncio + async def test_managed_ids_are_left_to_the_managed_id_resolver(self): + managed_id = new_managed_id("openai", UNISSUED_RESPONSE_ID) + pc = _object_lookup(None) + + assert ( + await addresses_unissued_response_id( + f"/v1/responses/{managed_id}", {"previous_response_id": managed_id}, "openai", pc + ) + is False + ) + pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + @pytest.mark.asyncio + async def test_failed_lookup_is_never_read_as_issued(self): + pc = _prisma_client_passthrough_managed() + pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=RuntimeError("db down")) + + with pytest.raises(RuntimeError, match="db down"): + await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) + + @pytest.mark.asyncio + async def test_no_database_counts_as_unissued_only_when_an_id_is_addressed(self): + assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", None) is True + assert await addresses_unissued_response_id("/v1/responses", {"model": "gpt-5.6"}, "openai", None) is False + + @pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate") class TestFlagOff: """ diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index eb31797d460..e4857d35a24 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -35,7 +35,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._lazy_features import LazyFeature, attach_lazy_features -from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import router as llm_passthrough_router from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( @@ -3678,6 +3678,151 @@ async def test_pass_through_request_merge_query_params_rewrites_managed_ids_on_t } +_UNISSUED_RESPONSE_ID: Final = "resp_0000000000000000unissued" + + +def _passthrough_object_rows(owner_user_id: str | None) -> SimpleNamespace: + row: Final = None if owner_user_id is None else SimpleNamespace(created_by=owner_user_id, team_id=None) + return SimpleNamespace( + db=SimpleNamespace( + litellm_managedobjecttable=SimpleNamespace( + find_first=AsyncMock(return_value=row), upsert=AsyncMock(), update=AsyncMock() + ) + ) + ) + + +async def _run_managed_responses_pass_through( + method: str, + target: str, + body: dict[str, object] | None, + general_settings: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + prisma_client: SimpleNamespace, +) -> tuple[str, int]: + from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.types.llms.custom_http import httpxSpecialProvider + + recorded_requests: Final[list[httpx.Request]] = [] + + def transport_handler(upstream_request: httpx.Request) -> httpx.Response: + recorded_requests.append(upstream_request) + return httpx.Response(200, json={"object": "list", "data": []}) + + real_handler: Final = get_async_httpx_client( + llm_provider=httpxSpecialProvider.PassThroughEndpoint, + params={"timeout": resolve_pass_through_request_timeout(None)}, + ) + cache_dict: Final = litellm.in_memory_llm_clients_cache.cache_dict + cache_key: Final = next(key for key, cached in cache_dict.items() if cached is real_handler) + cache_dict[cache_key] = SimpleNamespace(client=httpx.AsyncClient(transport=httpx.MockTransport(transport_handler))) + + mock_request: Final = MagicMock(spec=Request) + mock_request.method = method + mock_request.headers = Headers({"content-type": "application/json"}) + mock_request.query_params = QueryParams("") + mock_request.body = AsyncMock(return_value=b"") + + mock_proxy_logging: Final = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock( + side_effect=lambda user_api_key_dict, data, call_type, endpoint_type=None: data + ) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + mock_proxy_logging.get_proxy_hook = MagicMock( + return_value=_FakeManagedFilesHook(SimpleNamespace(created_by="user-1", team_id=None)) + ) + + try: + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging), + patch( + "litellm.proxy.proxy_server.general_settings", + {"passthrough_managed_object_ids": True, **general_settings}, + ), + patch("litellm.proxy.proxy_server.prisma_client", prisma_client), + patch("litellm.proxy.proxy_server.master_key", "sk-test-master-key"), + ): + response: Final = await pass_through_request( + request=mock_request, + target=target, + custom_headers={}, + user_api_key_dict=user_api_key_dict, + custom_body=body, + custom_llm_provider="openai", + ) + return str(response.status_code), len(recorded_requests) + except ProxyException as refusal: + return f"{refusal.code} {refusal.message}", len(recorded_requests) + finally: + cache_dict[cache_key] = real_handler + + +_RESPONSE_ID_ROUTES: Final = ( + ("GET", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}", None), + ("GET", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}/input_items", None), + ( + "POST", + "https://api.openai.com/v1/responses/compact", + {"model": "gpt-5.6", "previous_response_id": _UNISSUED_RESPONSE_ID}, + ), + ( + "POST", + "https://api.openai.com/v1/responses", + {"model": "gpt-5.6", "previous_response_id": _UNISSUED_RESPONSE_ID}, + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("method", "target", "body"), _RESPONSE_ID_ROUTES) +async def test_managed_pass_through_refuses_a_response_id_it_never_issued(method, target, body): + outcome, upstream_calls = await _run_managed_responses_pass_through( + method, target, body, {}, UserAPIKeyAuth(user_id="user-2"), _passthrough_object_rows(None) + ) + + assert outcome.startswith("403 ") + assert "allow_unmanaged_response_ids" in outcome + assert upstream_calls == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("method", "target", "body"), _RESPONSE_ID_ROUTES) +@pytest.mark.parametrize( + ("general_settings", "user_api_key_dict", "owner_user_id"), + [ + ({"allow_unmanaged_response_ids": True}, UserAPIKeyAuth(user_id="user-2"), None), + ({"disable_responses_id_security": True}, UserAPIKeyAuth(user_id="user-2"), None), + ({}, UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), None), + ({}, UserAPIKeyAuth(user_id="user-1"), "user-1"), + ], + ids=["allow_unmanaged_response_ids", "disable_responses_id_security", "proxy_admin", "issued_to_the_caller"], +) +async def test_managed_pass_through_forwards_ids_it_may_address( + method, target, body, general_settings, user_api_key_dict, owner_user_id +): + outcome, upstream_calls = await _run_managed_responses_pass_through( + method, target, body, general_settings, user_api_key_dict, _passthrough_object_rows(owner_user_id) + ) + + assert (outcome, upstream_calls) == ("200", 1) + + +@pytest.mark.asyncio +async def test_managed_pass_through_still_hides_a_response_issued_to_someone_else(): + outcome, upstream_calls = await _run_managed_responses_pass_through( + "GET", + f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}", + None, + {"allow_unmanaged_response_ids": True}, + UserAPIKeyAuth(user_id="user-2"), + _passthrough_object_rows("user-1"), + ) + + assert outcome.startswith("404 ") + assert upstream_calls == 0 + + @pytest.mark.asyncio async def test_pass_through_request_follows_redirect_to_final_response(httpx_transport): """ From 5b2b9c88922163ed56c1ce4f5f8550aa6402be9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:17:34 -0700 Subject: [PATCH 2/9] fix(proxy): skip the pass-through response id check without a database and cover delete, cancel, and input_tokens --- .../pass_through_endpoints/managed_id_rewriter.py | 7 ++----- .../test_managed_id_rewriter.py | 11 ++++++++--- .../test_pass_through_endpoints.py | 7 +++++++ 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index cb8fcc6badc..8b82000786e 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -493,12 +493,9 @@ async def addresses_unissued_response_id( provider: str, prisma_client: PrismaClient | None, ) -> bool: - addressed: Final = _addressed_raw_response_ids(path, body) - if not addressed: - return False if prisma_client is None: - return True - for raw_id in addressed: + return False + for raw_id in _addressed_raw_response_ids(path, body): if not await _passthrough_issued_response_id(raw_id, provider, prisma_client): return True return False diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index ba3399107cf..c9171f0290b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -1804,9 +1804,14 @@ class TestAddressesUnissuedResponseId: await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) @pytest.mark.asyncio - async def test_no_database_counts_as_unissued_only_when_an_id_is_addressed(self): - assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", None) is True - assert await addresses_unissued_response_id("/v1/responses", {"model": "gpt-5.6"}, "openai", None) is False + async def test_without_a_database_the_pass_through_issues_only_raw_ids_so_none_is_refused(self): + assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", None) is False + assert ( + await addresses_unissued_response_id( + "/v1/responses", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", None + ) + is False + ) @pytest.mark.usefixtures("_drain_logging_worker", "_vcr_outcome_gate") diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index e4857d35a24..24a8132c49f 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -3761,6 +3761,13 @@ async def _run_managed_responses_pass_through( _RESPONSE_ID_ROUTES: Final = ( ("GET", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}", None), ("GET", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}/input_items", None), + ("DELETE", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}", None), + ("POST", f"https://api.openai.com/v1/responses/{_UNISSUED_RESPONSE_ID}/cancel", None), + ( + "POST", + "https://api.openai.com/v1/responses/input_tokens", + {"model": "gpt-5.6", "previous_response_id": _UNISSUED_RESPONSE_ID}, + ), ( "POST", "https://api.openai.com/v1/responses/compact", From 9c023cacb69d482efbb52dd55e44b673ff6b569e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 07:20:26 -0700 Subject: [PATCH 3/9] fix(proxy): own and check response ids on the Azure model_list pass-through relay --- .../llm_passthrough_endpoints.py | 154 ++++++++++++++--- .../managed_id_rewriter.py | 16 +- .../pass_through_endpoints.py | 38 +++-- .../test_llm_pass_through_endpoints.py | 156 ++++++++++++++++++ .../test_managed_id_rewriter.py | 10 +- 5 files changed, 332 insertions(+), 42 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 1eefb58d7fa..6f95737cd61 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -25,7 +25,7 @@ import httpx import openai from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket from fastapi.responses import StreamingResponse -from pydantic import ConfigDict, TypeAdapter +from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from starlette.websockets import WebSocketState from typing_extensions import ReadOnly, TypedDict @@ -51,6 +51,7 @@ from litellm.llms.azure.passthrough.transformation import ( foreign_azure_deployment, is_azure_body_model_inference_endpoint, ) +from litellm.llms.base_llm.passthrough.transformation import strip_leading_model_segment from litellm.llms.bedrock.passthrough.transformation import is_bedrock_streaming_endpoint from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.deepgram.common_utils import ( @@ -100,6 +101,7 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, create_pass_through_route, create_websocket_passthrough_route, + refuse_unissued_response_ids, websocket_passthrough_request, ) from litellm.proxy.utils import ProxyLogging as ProxyLoggingType @@ -131,7 +133,9 @@ from litellm.utils import ProviderConfigManager from .passthrough_endpoint_router import PassthroughEndpointRouter if TYPE_CHECKING: + from litellm.integrations.custom_logger import CustomLogger from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig + from litellm.proxy.utils import PrismaClient from litellm.router import Router ProxyConfig = _ProxyConfig # rebind-ok: conditional type alias @@ -2058,9 +2062,101 @@ async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> Async await upstream.aclose() -async def _relay_upstream_response(upstream: httpx.Response) -> Response: +_RELAYED_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _RouterRelayManagedIds: + provider: str + prisma_client: PrismaClient + managed_files_hook: CustomLogger + user_api_key_dict: UserAPIKeyAuth + method: str + route: str + + async def resolve_request( + self, endpoint: str, request_body: Mapping[str, object] + ) -> tuple[str, Mapping[str, object]]: + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import rewrite_body_ids, rewrite_path_ids + + await refuse_unissued_response_ids( + endpoint, request_body, self.provider, self.prisma_client, self.user_api_key_dict + ) + resolved_endpoint: Final = await rewrite_path_ids( + endpoint, self.provider, self.user_api_key_dict, self.prisma_client, self.managed_files_hook + ) + resolved_body: Final = await rewrite_body_ids( + dict(request_body), self.provider, self.user_api_key_dict, self.prisma_client, self.managed_files_hook + ) + return resolved_endpoint, request_body if resolved_body is None else resolved_body + + async def own_response_body(self, content: bytes, status_code: int) -> bytes: + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import rewrite_response_ids + + if status_code >= 300: + return content + try: + body: Final = _RELAYED_JSON_OBJECT.validate_json(content) + except ValidationError: + return content + owned: Final = await rewrite_response_ids( + provider=self.provider, + method=self.method, + route=self.route, + body=body, + user_api_key_dict=self.user_api_key_dict, + prisma_client=self.prisma_client, + managed_files_hook=self.managed_files_hook, + ) + return content if owned is body else json.dumps(owned).encode("utf-8") + + def own_stream(self, stream: AsyncGenerator[bytes, None]) -> AsyncGenerator[bytes, None]: + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import rewrite_streamed_response_ids + + return rewrite_streamed_response_ids( + stream=stream, + provider=self.provider, + method=self.method, + route=self.route, + user_api_key_dict=self.user_api_key_dict, + prisma_client=self.prisma_client, + ) + + +def _router_relay_managed_ids( + provider: str, method: str, route: str, user_api_key_dict: UserAPIKeyAuth +) -> _RouterRelayManagedIds | None: + from litellm.proxy.proxy_server import general_settings_view, prisma_client, proxy_logging_obj + + managed_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") + if ( + not general_settings_view().get("passthrough_managed_object_ids", False) + or prisma_client is None + or managed_files_hook is None + ): + return None + return _RouterRelayManagedIds( + provider=provider, + prisma_client=prisma_client, + managed_files_hook=managed_files_hook, + user_api_key_dict=user_api_key_dict, + method=method, + route=route, + ) + + +def _owned_relay_stream( + stream: AsyncGenerator[bytes, None], managed_ids: _RouterRelayManagedIds | None +) -> AsyncGenerator[bytes, None]: + return stream if managed_ids is None else managed_ids.own_stream(stream) + + +async def _relay_upstream_response( + upstream: httpx.Response, managed_ids: _RouterRelayManagedIds | None = None +) -> Response: + content: Final = await upstream.aread() return Response( - content=await upstream.aread(), + content=content if managed_ids is None else await managed_ids.own_response_body(content, upstream.status_code), status_code=upstream.status_code, headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None), ) @@ -2074,6 +2170,7 @@ async def _relay_azure_router_model( request_body: Mapping[str, object], is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, + managed_ids: _RouterRelayManagedIds | None, ) -> Response: foreign_deployment: Final = foreign_azure_deployment( endpoint, model, lambda: _models_served_by_group(llm_router, model) @@ -2092,6 +2189,7 @@ async def _relay_azure_router_model( request_body=request_body, is_streaming_request=is_streaming_request, user_api_key_dict=user_api_key_dict, + managed_ids=managed_ids, ) @@ -2103,6 +2201,7 @@ async def _relay_router_model( request_body: Mapping[str, object], is_streaming_request: bool, user_api_key_dict: UserAPIKeyAuth, + managed_ids: _RouterRelayManagedIds | None, ) -> Response: try: result: Final = await llm_router.allm_passthrough_route( @@ -2125,13 +2224,13 @@ async def _relay_router_model( return await _relay_upstream_response(upstream_error.response) if not is_streaming_request: - return await _relay_upstream_response(cast(httpx.Response, result)) + return await _relay_upstream_response(cast(httpx.Response, result), managed_ids) if inspect.isasyncgen(result): sse_headers: Final = {"content-type": "text/event-stream"} return StreamingResponse( content=wrap_passthrough_sse_bytes_with_keepalive_pings( - stream=_relay_upstream_bytes(result), + stream=_owned_relay_stream(_relay_upstream_bytes(result), managed_ids), ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, upstream_headers=sse_headers, ), @@ -2142,7 +2241,7 @@ async def _relay_router_model( upstream_stream: Final = cast(AsyncPassthroughStreamingResponse, result) return StreamingResponse( content=wrap_passthrough_sse_bytes_with_keepalive_pings( - stream=_relay_upstream_bytes(upstream_stream), + stream=_owned_relay_stream(_relay_upstream_bytes(upstream_stream), managed_ids), ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds, upstream_headers=upstream_stream.headers, ), @@ -2207,6 +2306,7 @@ async def relay_nvidia_nim_request( request_body=request_body, is_streaming_request=is_streaming_request, user_api_key_dict=user_api_key_dict, + managed_ids=None, ), ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None), ) @@ -2226,16 +2326,37 @@ async def _relay_azure_body_model_group( model: Final = _optional_str(request_body.get("model")) if model is None or not is_passthrough_request_using_router_model(request_body, llm_router): return None + return await _open_azure_router_relay(llm_router, model, endpoint, request, request_body, user_api_key_dict) + + +async def _open_azure_router_relay( + llm_router: litellm.Router, + model: str, + endpoint: str, + request: Request, + request_body: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, +) -> Response: is_streaming_request: Final = is_passthrough_request_streaming(request_body) + managed_ids: Final = _router_relay_managed_ids( + LlmProviders.AZURE.value, + request.method, + f"/azure/{strip_leading_model_segment(endpoint, (model,))}", + user_api_key_dict, + ) + relayed_endpoint, relayed_body = ( + (endpoint, request_body) if managed_ids is None else await managed_ids.resolve_request(endpoint, request_body) + ) return await open_sse_before_first_byte( _relay_azure_router_model( llm_router=llm_router, model=model, - endpoint=endpoint, + endpoint=relayed_endpoint, request=request, - request_body=request_body, + request_body=relayed_body, is_streaming_request=is_streaming_request, user_api_key_dict=user_api_key_dict, + managed_ids=managed_ids, ), ping_interval_seconds=(litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None), ) @@ -2289,21 +2410,8 @@ async def azure_proxy_route( ) if is_router_model: - request_body = await get_request_body(request) - is_streaming_request = is_passthrough_request_streaming(request_body) - return await open_sse_before_first_byte( - _relay_azure_router_model( - llm_router=llm_router, - model=part, - endpoint=endpoint, - request=request, - request_body=request_body, - is_streaming_request=is_streaming_request, - user_api_key_dict=user_api_key_dict, - ), - ping_interval_seconds=( - litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None - ), + return await _open_azure_router_relay( + llm_router, part, endpoint, request, await get_request_body(request), user_api_key_dict ) elif is_vector_store_index: # get the api key from the provider config diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 8b82000786e..9453e4993e2 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -52,7 +52,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( can_access_resource, resolve_resource_owner_id, ) -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import DB_CONNECTION_ERROR_TYPES, UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames from litellm.repositories.managed_file_repository import ManagedFileRepository @@ -481,9 +481,17 @@ def _addressed_raw_response_ids(path: str, body: Mapping[str, object]) -> frozen async def _passthrough_issued_response_id(raw_id: str, provider: str, prisma_client: PrismaClient) -> bool: - row: Final = await _object_table(prisma_client).find_first( - where={"model_object_id": f"passthrough:{provider}:{raw_id}"} - ) + from prisma.errors import PrismaError + + try: + row: Final = await _object_table(prisma_client).find_first( + where={"model_object_id": f"passthrough:{provider}:{raw_id}"} + ) + except (PrismaError, *DB_CONNECTION_ERROR_TYPES): + verbose_proxy_logger.warning( + "managed_id_rewriter: response id ownership lookup failed, refusing the id", exc_info=True + ) + return False return row is not None diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 6449b67b903..7882b7f72ed 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -153,6 +153,7 @@ from .upstream_usage_headers import ( if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import PrismaClient router: Final = APIRouter() @@ -1340,13 +1341,7 @@ async def pass_through_request( ) _passthrough_managed_hook = proxy_logging_obj.get_proxy_hook("managed_files") if _passthrough_managed_hook is not None: - from litellm.proxy.hooks.responses_id_security import ( - UNMANAGED_RESPONSE_ID_DETAIL, - proxy_signing_key, - unmanaged_response_ids_allowed, - ) from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( - addresses_unissued_response_id, rewrite_body_ids, rewrite_path_ids, rewrite_query_ids, @@ -1355,15 +1350,13 @@ async def pass_through_request( prisma_client as _passthrough_prisma, ) - if not unmanaged_response_ids_allowed( - general_settings_view(), proxy_signing_key(), user_api_key_dict - ) and await addresses_unissued_response_id( + await refuse_unissued_response_ids( url.path, _REQUEST_BODY_VIEW.validate_python(_parsed_body or {}), _managed_id_provider, _passthrough_prisma, - ): - raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL) + user_api_key_dict, + ) _original_path: Final = url.path _original_query_params: Final = requested_query_params @@ -2905,6 +2898,29 @@ def _own_streamed_managed_ids( ) +async def refuse_unissued_response_ids( + path: str, + body: Mapping[str, object], + managed_id_provider: str, + prisma_client: "PrismaClient | None", + user_api_key_dict: UserAPIKeyAuth, +) -> None: + from litellm.proxy.hooks.responses_id_security import ( + UNMANAGED_RESPONSE_ID_DETAIL, + proxy_signing_key, + unmanaged_response_ids_allowed, + ) + from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( + addresses_unissued_response_id, + ) + from litellm.proxy.proxy_server import general_settings_view + + if unmanaged_response_ids_allowed(general_settings_view(), proxy_signing_key(), user_api_key_dict): + return + if await addresses_unissued_response_id(path, body, managed_id_provider, prisma_client): + raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL) + + def _should_buffer_passthrough_response(response: httpx.Response) -> bool: """ Decide from the response headers whether the body must be read into memory. diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index d25b4470668..2b0867ff866 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7067,6 +7067,162 @@ class TestAzureBodyModelGroupRelay: assert route["target"] == "https://operator.openai.azure.com/openai/v1/fine_tuning/jobs" +_VENDOR_RESPONSE_ID: Final = "resp_0000000000000000vendor" + + +class _ManagedObjectStore: + def __init__(self) -> None: + self.rows: list[SimpleNamespace] = [] + self.litellm_managedobjecttable = SimpleNamespace( + find_first=AsyncMock(side_effect=self._find_first), + upsert=AsyncMock(side_effect=self._upsert), + update=AsyncMock(), + ) + + @property + def db(self) -> "_ManagedObjectStore": + return self + + async def _find_first(self, where: Mapping[str, str]) -> SimpleNamespace | None: + ((column, value),) = where.items() + return next((row for row in self.rows if getattr(row, column) == value), None) + + async def _upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str | None]]) -> None: + self.rows.append(SimpleNamespace(**data["create"])) + + +class TestAzureRouterRelayManagedIds: + def _install( + self, + monkeypatch: pytest.MonkeyPatch, + bodies: Iterator[dict], + upstream: _AzureRelayUpstream = _azure_relay_json_upstream, + general_settings: Mapping[str, object] = MappingProxyType({}), + ) -> tuple[list[dict], _ManagedObjectStore]: + import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep + from litellm.proxy import proxy_server + + captured: list[dict] = [] + store: Final = _ManagedObjectStore() + + async def next_request_body(_request: Request) -> dict: + return next(bodies) + + monkeypatch.setattr(proxy_server, "llm_router", _AzureBodyModelGroupRouter(captured, upstream)) + monkeypatch.setattr(ep, "get_request_body", next_request_body) + monkeypatch.setattr( + proxy_server, "general_settings", {"passthrough_managed_object_ids": True, **general_settings} + ) + monkeypatch.setattr(proxy_server, "prisma_client", store) + monkeypatch.setattr(proxy_server, "master_key", "sk-test-master-key") + monkeypatch.setattr(proxy_server, "proxy_logging_obj", SimpleNamespace(get_proxy_hook=lambda name: object())) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + return captured, store + + def _request(self, method: str) -> Request: + request = MagicMock(spec=Request) + request.method = method + request.headers = {"content-type": "application/json"} + request.query_params = {"api-version": "2025-03-01-preview"} + return request + + async def _call(self, method: str, endpoint: str, user_id: str) -> Response: + return await azure_proxy_route( + endpoint=endpoint, + request=self._request(method), + fastapi_response=MagicMock(spec=Response), + user_api_key_dict=UserAPIKeyAuth(api_key=f"hashed-{user_id}", user_id=user_id), + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("method", "endpoint", "body"), + [ + ("POST", "openai/v1/responses", {"model": "gpt-5.4", "previous_response_id": _VENDOR_RESPONSE_ID}), + ( + "POST", + "gpt-5.4/openai/v1/responses", + {"model": "gpt-5.4", "previous_response_id": _VENDOR_RESPONSE_ID}, + ), + ("GET", f"gpt-5.4/openai/v1/responses/{_VENDOR_RESPONSE_ID}", {}), + ("POST", "gpt-5.4/openai/v1/responses/compact", {"previous_response_id": _VENDOR_RESPONSE_ID}), + ], + ids=["body_model_create", "path_model_create", "path_model_retrieve", "path_model_compact"], + ) + async def test_router_relay_refuses_a_response_id_the_pass_through_never_issued( + self, monkeypatch, method, endpoint, body + ): + captured, _ = self._install(monkeypatch, iter([body])) + + with pytest.raises(HTTPException) as refusal: + await self._call(method, endpoint, "user-2") + + assert refusal.value.status_code == 403 + assert "allow_unmanaged_response_ids" in refusal.value.detail + assert captured == [] + + @pytest.mark.asyncio + async def test_router_relay_lets_a_vendor_id_through_with_allow_unmanaged_response_ids(self, monkeypatch): + body = {"model": "gpt-5.4", "previous_response_id": _VENDOR_RESPONSE_ID} + captured, _ = self._install(monkeypatch, iter([body]), general_settings={"allow_unmanaged_response_ids": True}) + + result = await self._call("POST", "openai/v1/responses", "user-2") + + assert result.status_code == 200 + assert captured[0]["json"]["previous_response_id"] == _VENDOR_RESPONSE_ID + + @pytest.mark.asyncio + async def test_an_id_the_router_relay_issued_keeps_working_for_its_owner_only(self, monkeypatch): + from litellm.proxy.pass_through_endpoints.managed_id_codec import decode + + captured, store = self._install(monkeypatch, iter([{"model": "gpt-5.4", "input": "ping"}, {}, {}, {}])) + + created = await self._call("POST", "openai/v1/responses", "user-1") + issued_id = json.loads(created.body)["id"] + payload = decode(issued_id) + assert payload is not None + assert (payload.provider, payload.raw_provider_id) == ("azure", "resp_1") + assert [(row.model_object_id, row.created_by) for row in store.rows] == [("passthrough:azure:resp_1", "user-1")] + + await self._call("GET", f"gpt-5.4/openai/v1/responses/{issued_id}", "user-1") + assert captured[-1]["endpoint"] == "gpt-5.4/openai/v1/responses/resp_1" + + await self._call("GET", "gpt-5.4/openai/v1/responses/resp_1", "user-1") + assert captured[-1]["endpoint"] == "gpt-5.4/openai/v1/responses/resp_1" + + relayed_before_the_stranger: Final = len(captured) + with pytest.raises(HTTPException) as refusal: + await self._call("GET", f"gpt-5.4/openai/v1/responses/{issued_id}", "user-2") + assert refusal.value.status_code == 403 + assert len(captured) == relayed_before_the_stranger + + @pytest.mark.asyncio + async def test_a_streamed_create_through_the_router_relay_hands_out_an_owned_id(self, monkeypatch): + from litellm.proxy.pass_through_endpoints.managed_id_codec import decode + + async def upstream_events() -> AsyncIterator[bytes]: + yield b'event: response.created\ndata: {"type":"response.created","response":{"id":"resp_1"}}\n\n' + yield b'event: response.completed\ndata: {"type":"response.completed","response":{"id":"resp_1"}}\n\n' + + async def streaming_upstream() -> AsyncIterator[bytes]: + return upstream_events() + + _, store = self._install( + monkeypatch, iter([{"model": "gpt-5.4", "input": "ping", "stream": True}]), streaming_upstream + ) + + result = await self._call("POST", "openai/v1/responses", "user-1") + + assert isinstance(result, StreamingResponse) + streamed = b"".join([chunk async for chunk in result.body_iterator]).decode() + assert "resp_1" not in streamed + ((row),) = store.rows + assert row.created_by == "user-1" + assert streamed.count(row.unified_object_id) == 2 + payload = decode(row.unified_object_id) + assert payload is not None and payload.raw_provider_id == "resp_1" + + AZURE_SPEECH_SHORT_AUDIO_ENDPOINT: Final = "/speech/recognition/conversation/cognitiveservices/v1" AZURE_SPEECH_BATCH_ENDPOINT: Final = "/speechtotext/v3.2/transcriptions" AZURE_SPEECH_FAST_ENDPOINT: Final = "/speechtotext/transcriptions:transcribe" diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index c9171f0290b..cab36251b70 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -3,7 +3,9 @@ import json from collections.abc import AsyncIterator, Iterable from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +from prisma.errors import PrismaError from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.managed_id_codec import( @@ -1796,12 +1798,12 @@ class TestAddressesUnissuedResponseId: pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() @pytest.mark.asyncio - async def test_failed_lookup_is_never_read_as_issued(self): + @pytest.mark.parametrize("lookup_error", [httpx.ConnectError("db down"), PrismaError("query failed")]) + async def test_failed_lookup_is_never_read_as_issued(self, lookup_error): pc = _prisma_client_passthrough_managed() - pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=RuntimeError("db down")) + pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=lookup_error) - with pytest.raises(RuntimeError, match="db down"): - await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) + assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) is True @pytest.mark.asyncio async def test_without_a_database_the_pass_through_issues_only_raw_ids_so_none_is_refused(self): From 2417cf348e5fe68871c4b8b332d668a576ac8e5b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:57:22 -0700 Subject: [PATCH 4/9] fix(proxy): keep the model_list relay's managed ids on Responses routes and answer 503 when a response id lookup fails --- .../llm_passthrough_endpoints.py | 2 + .../managed_id_rewriter.py | 30 +++++----- .../pass_through_endpoints.py | 19 ++++++- .../managed_id_rewriter.py | 1 + .../test_llm_pass_through_endpoints.py | 22 +++++++ .../test_managed_id_rewriter.py | 57 ++++++++++++++----- .../test_pass_through_endpoints.py | 18 ++++++ 7 files changed, 119 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6f95737cd61..197caa89fe5 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -2128,6 +2128,8 @@ def _router_relay_managed_ids( ) -> _RouterRelayManagedIds | None: from litellm.proxy.proxy_server import general_settings_view, prisma_client, proxy_logging_obj + if "responses" not in route.split("/"): + return None managed_files_hook: Final = proxy_logging_obj.get_proxy_hook("managed_files") if ( not general_settings_view().get("passthrough_managed_object_ids", False) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 9453e4993e2..501e3f108be 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -71,6 +71,7 @@ from litellm.types.passthrough_endpoints.managed_id_rewriter import ( PrismaWhere, PrismaWhereValue, ResourceKind, + ResponseIdIssuance, SortOrder, ) @@ -480,7 +481,7 @@ def _addressed_raw_response_ids(path: str, body: Mapping[str, object]) -> frozen return path_ids -async def _passthrough_issued_response_id(raw_id: str, provider: str, prisma_client: PrismaClient) -> bool: +async def _response_id_issuance(raw_id: str, provider: str, prisma_client: PrismaClient) -> ResponseIdIssuance: from prisma.errors import PrismaError try: @@ -488,25 +489,28 @@ async def _passthrough_issued_response_id(raw_id: str, provider: str, prisma_cli where={"model_object_id": f"passthrough:{provider}:{raw_id}"} ) except (PrismaError, *DB_CONNECTION_ERROR_TYPES): - verbose_proxy_logger.warning( - "managed_id_rewriter: response id ownership lookup failed, refusing the id", exc_info=True - ) - return False - return row is not None + verbose_proxy_logger.warning("managed_id_rewriter: response id ownership lookup failed", exc_info=True) + return "unverified" + return "unissued" if row is None else "issued" -async def addresses_unissued_response_id( +async def addressed_response_id_issuance( path: str, body: Mapping[str, object], provider: str, prisma_client: PrismaClient | None, -) -> bool: +) -> ResponseIdIssuance: if prisma_client is None: - return False - for raw_id in _addressed_raw_response_ids(path, body): - if not await _passthrough_issued_response_id(raw_id, provider, prisma_client): - return True - return False + return "issued" + outcomes: Final = frozenset( + [ + await _response_id_issuance(raw_id, provider, prisma_client) + for raw_id in _addressed_raw_response_ids(path, body) + ] + ) + if "unissued" in outcomes: + return "unissued" + return "unverified" if "unverified" in outcomes else "issued" # --------------------------------------------------------------------------- diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7882b7f72ed..5d4c023f702 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -30,6 +30,7 @@ from pydantic import TypeAdapter from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.routing import BaseRoute, Route from starlette.websockets import WebSocketState +from typing_extensions import assert_never from websockets.asyncio.client import connect from websockets.exceptions import ( ConnectionClosedError, @@ -2911,14 +2912,26 @@ async def refuse_unissued_response_ids( unmanaged_response_ids_allowed, ) from litellm.proxy.pass_through_endpoints.managed_id_rewriter import ( - addresses_unissued_response_id, + addressed_response_id_issuance, ) from litellm.proxy.proxy_server import general_settings_view if unmanaged_response_ids_allowed(general_settings_view(), proxy_signing_key(), user_api_key_dict): return - if await addresses_unissued_response_id(path, body, managed_id_provider, prisma_client): - raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL) + issuance: Final = await addressed_response_id_issuance(path, body, managed_id_provider, prisma_client) + match issuance: + case "issued": + return + case "unissued": + raise HTTPException(status_code=403, detail=UNMANAGED_RESPONSE_ID_DETAIL) + case "unverified": + raise HTTPException( + status_code=503, + detail="The proxy could not check who owns this response id because its database lookup failed. " + "Retry the request.", + ) + case _: + assert_never(issuance) def _should_buffer_passthrough_response(response: httpx.Response) -> bool: diff --git a/litellm/types/passthrough_endpoints/managed_id_rewriter.py b/litellm/types/passthrough_endpoints/managed_id_rewriter.py index 0603c13d905..488b9305141 100644 --- a/litellm/types/passthrough_endpoints/managed_id_rewriter.py +++ b/litellm/types/passthrough_endpoints/managed_id_rewriter.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: SortOrder: TypeAlias = Literal["asc", "desc"] ResourceKind: TypeAlias = Literal["files", "batches"] +ResponseIdIssuance: TypeAlias = Literal["issued", "unissued", "unverified"] PrismaWhereValue: TypeAlias = ( "str | int | bool | datetime | None | Mapping[str, PrismaWhereValue] | Sequence[PrismaWhereValue]" diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 2b0867ff866..2e4514752a8 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7196,6 +7196,28 @@ class TestAzureRouterRelayManagedIds: assert refusal.value.status_code == 403 assert len(captured) == relayed_before_the_stranger + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("method", "endpoint", "upstream_body"), + [ + ("GET", "gpt-5.4/openai/v1/batches/batch_1", {"id": "batch_1", "object": "batch"}), + ("POST", "gpt-5.4/openai/v1/batches/batch_1/cancel", {"id": "batch_1", "object": "batch"}), + ], + ) + async def test_router_relay_leaves_ids_off_the_responses_routes_as_they_were( + self, monkeypatch, method, endpoint, upstream_body + ): + async def upstream() -> httpx.Response: + return httpx.Response(200, json=upstream_body) + + captured, store = self._install(monkeypatch, iter([{}]), upstream) + + result = await self._call(method, endpoint, "user-1") + + assert json.loads(result.body) == upstream_body + assert captured[-1]["endpoint"] == endpoint + assert store.rows == [] + @pytest.mark.asyncio async def test_a_streamed_create_through_the_router_relay_hands_out_an_owned_id(self, monkeypatch): from litellm.proxy.pass_through_endpoints.managed_id_codec import decode diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index cab36251b70..fabdea0ec1d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -19,7 +19,7 @@ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import( _MAX_RAW_ID_GUARD_LOOKUPS, _passthrough_provider_marker, _resolve_one, - addresses_unissued_response_id, + addressed_response_id_issuance, is_passthrough_list_route, list_passthrough_ids_from_db, rewrite_body_ids, @@ -1722,7 +1722,7 @@ def _object_lookup(row: object) -> MagicMock: return pc -class TestAddressesUnissuedResponseId: +class TestAddressedResponseIdIssuance: @pytest.mark.asyncio @pytest.mark.parametrize( ("path", "body"), @@ -1738,7 +1738,7 @@ class TestAddressesUnissuedResponseId: async def test_response_id_without_a_passthrough_row_is_unissued(self, path, body): pc = _object_lookup(None) - assert await addresses_unissued_response_id(path, body, "openai", pc) is True + assert await addressed_response_id_issuance(path, body, "openai", pc) == "unissued" pc.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( where={"model_object_id": f"passthrough:openai:{UNISSUED_RESPONSE_ID}"} ) @@ -1747,7 +1747,10 @@ class TestAddressesUnissuedResponseId: async def test_response_id_with_a_passthrough_row_is_issued(self): pc = _object_lookup(_batch_row(encode("azure", "user-1", UNISSUED_RESPONSE_ID))) - assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "azure", pc) is False + assert ( + await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "azure", pc) + == "issued" + ) pc.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( where={"model_object_id": f"passthrough:azure:{UNISSUED_RESPONSE_ID}"} ) @@ -1761,10 +1764,10 @@ class TestAddressesUnissuedResponseId: ) assert ( - await addresses_unissued_response_id( + await addressed_response_id_issuance( f"/v1/responses/{issued_id}", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", pc ) - is True + == "unissued" ) @pytest.mark.asyncio @@ -1781,7 +1784,7 @@ class TestAddressesUnissuedResponseId: async def test_requests_that_address_no_response_id_skip_the_lookup(self, path, body): pc = _object_lookup(None) - assert await addresses_unissued_response_id(path, body, "openai", pc) is False + assert await addressed_response_id_issuance(path, body, "openai", pc) == "issued" pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() @pytest.mark.asyncio @@ -1790,29 +1793,55 @@ class TestAddressesUnissuedResponseId: pc = _object_lookup(None) assert ( - await addresses_unissued_response_id( + await addressed_response_id_issuance( f"/v1/responses/{managed_id}", {"previous_response_id": managed_id}, "openai", pc ) - is False + == "issued" ) pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() @pytest.mark.asyncio @pytest.mark.parametrize("lookup_error", [httpx.ConnectError("db down"), PrismaError("query failed")]) - async def test_failed_lookup_is_never_read_as_issued(self, lookup_error): + async def test_failed_lookup_is_unverified_rather_than_issued_or_unissued(self, lookup_error): pc = _prisma_client_passthrough_managed() pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=lookup_error) - assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) is True + assert ( + await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) + == "unverified" + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize(("issued_row", "expected"), [(None, "unissued"), (_batch_row("x"), "unverified")]) + async def test_a_failed_lookup_beside_another_id_never_hides_an_unissued_one(self, issued_row, expected): + failing_id = "resp_2222222222222222failing" + + async def find_first(where): + if where["model_object_id"].endswith(failing_id): + raise PrismaError("query failed") + return issued_row + + pc = _prisma_client_passthrough_managed() + pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=find_first) + + assert ( + await addressed_response_id_issuance( + f"/v1/responses/{failing_id}", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", pc + ) + == expected + ) @pytest.mark.asyncio async def test_without_a_database_the_pass_through_issues_only_raw_ids_so_none_is_refused(self): - assert await addresses_unissued_response_id(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", None) is False assert ( - await addresses_unissued_response_id( + await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", None) + == "issued" + ) + assert ( + await addressed_response_id_issuance( "/v1/responses", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", None ) - is False + == "issued" ) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 24a8132c49f..f8d78b79851 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -3815,6 +3815,24 @@ async def test_managed_pass_through_forwards_ids_it_may_address( assert (outcome, upstream_calls) == ("200", 1) +@pytest.mark.asyncio +@pytest.mark.parametrize(("method", "target", "body"), _RESPONSE_ID_ROUTES) +async def test_managed_pass_through_answers_a_retryable_503_when_it_cannot_check_a_response_id(method, target, body): + unreachable_database: Final = SimpleNamespace( + db=SimpleNamespace( + litellm_managedobjecttable=SimpleNamespace(find_first=AsyncMock(side_effect=httpx.ConnectError("db down"))) + ) + ) + + outcome, upstream_calls = await _run_managed_responses_pass_through( + method, target, body, {}, UserAPIKeyAuth(user_id="user-1"), unreachable_database + ) + + assert outcome.startswith("503 ") + assert "Retry the request" in outcome + assert upstream_calls == 0 + + @pytest.mark.asyncio async def test_managed_pass_through_still_hides_a_response_issued_to_someone_else(): outcome, upstream_calls = await _run_managed_responses_pass_through( From 20d9c99b921b978715805210a967ff374683fb18 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:21:39 -0700 Subject: [PATCH 5/9] fix(passthrough): check every addressed response id in one database query --- .../managed_id_rewriter.py | 33 ++++-------- .../managed_id_rewriter.py | 2 + .../test_llm_pass_through_endpoints.py | 5 ++ .../test_managed_id_rewriter.py | 52 +++++++------------ .../test_pass_through_endpoints.py | 7 ++- 5 files changed, 43 insertions(+), 56 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 501e3f108be..8013d7cf184 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -481,36 +481,25 @@ def _addressed_raw_response_ids(path: str, body: Mapping[str, object]) -> frozen return path_ids -async def _response_id_issuance(raw_id: str, provider: str, prisma_client: PrismaClient) -> ResponseIdIssuance: - from prisma.errors import PrismaError - - try: - row: Final = await _object_table(prisma_client).find_first( - where={"model_object_id": f"passthrough:{provider}:{raw_id}"} - ) - except (PrismaError, *DB_CONNECTION_ERROR_TYPES): - verbose_proxy_logger.warning("managed_id_rewriter: response id ownership lookup failed", exc_info=True) - return "unverified" - return "unissued" if row is None else "issued" - - async def addressed_response_id_issuance( path: str, body: Mapping[str, object], provider: str, prisma_client: PrismaClient | None, ) -> ResponseIdIssuance: + from prisma.errors import PrismaError + if prisma_client is None: return "issued" - outcomes: Final = frozenset( - [ - await _response_id_issuance(raw_id, provider, prisma_client) - for raw_id in _addressed_raw_response_ids(path, body) - ] - ) - if "unissued" in outcomes: - return "unissued" - return "unverified" if "unverified" in outcomes else "issued" + object_ids: Final = sorted(f"passthrough:{provider}:{raw_id}" for raw_id in _addressed_raw_response_ids(path, body)) + if not object_ids: + return "issued" + try: + issued_count: Final = await _object_table(prisma_client).count(where={"model_object_id": {"in": object_ids}}) + except (PrismaError, *DB_CONNECTION_ERROR_TYPES): + verbose_proxy_logger.warning("managed_id_rewriter: response id ownership lookup failed", exc_info=True) + return "unverified" + return "issued" if issued_count == len(object_ids) else "unissued" # --------------------------------------------------------------------------- diff --git a/litellm/types/passthrough_endpoints/managed_id_rewriter.py b/litellm/types/passthrough_endpoints/managed_id_rewriter.py index 488b9305141..f98042afc85 100644 --- a/litellm/types/passthrough_endpoints/managed_id_rewriter.py +++ b/litellm/types/passthrough_endpoints/managed_id_rewriter.py @@ -81,6 +81,8 @@ class ManagedObjectTable(ManagedTable[ManagedObjectRow], Protocol): async def upsert(self, *, where: PrismaWhere, data: Mapping[str, ManagedRowData]) -> ManagedObjectRow: ... + async def count(self, *, where: PrismaWhere) -> int: ... + @runtime_checkable class ManagedFileIdReader(Protocol): diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 2e4514752a8..e1e94d48b00 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -7075,6 +7075,7 @@ class _ManagedObjectStore: self.rows: list[SimpleNamespace] = [] self.litellm_managedobjecttable = SimpleNamespace( find_first=AsyncMock(side_effect=self._find_first), + count=AsyncMock(side_effect=self._count), upsert=AsyncMock(side_effect=self._upsert), update=AsyncMock(), ) @@ -7087,6 +7088,10 @@ class _ManagedObjectStore: ((column, value),) = where.items() return next((row for row in self.rows if getattr(row, column) == value), None) + async def _count(self, where: Mapping[str, Mapping[str, list[str]]]) -> int: + ((column, condition),) = where.items() + return sum(1 for row in self.rows if getattr(row, column) in condition["in"]) + async def _upsert(self, where: Mapping[str, str], data: Mapping[str, Mapping[str, str | None]]) -> None: self.rows.append(SimpleNamespace(**data["create"])) diff --git a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py index fabdea0ec1d..6502e139929 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py +++ b/tests/unit/proxy/pass_through_endpoints/test_managed_id_rewriter.py @@ -1716,9 +1716,9 @@ class TestRawProviderIdGuardBudget: UNISSUED_RESPONSE_ID = "resp_0000000000000000unissued" -def _object_lookup(row: object) -> MagicMock: +def _issued_count(issued: int) -> MagicMock: pc = _prisma_client_passthrough_managed() - pc.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=row) + pc.db.litellm_managedobjecttable.count = AsyncMock(return_value=issued) return pc @@ -1736,31 +1736,31 @@ class TestAddressedResponseIdIssuance: ], ) async def test_response_id_without_a_passthrough_row_is_unissued(self, path, body): - pc = _object_lookup(None) + pc = _issued_count(0) assert await addressed_response_id_issuance(path, body, "openai", pc) == "unissued" - pc.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( - where={"model_object_id": f"passthrough:openai:{UNISSUED_RESPONSE_ID}"} + pc.db.litellm_managedobjecttable.count.assert_awaited_once_with( + where={"model_object_id": {"in": [f"passthrough:openai:{UNISSUED_RESPONSE_ID}"]}} ) @pytest.mark.asyncio async def test_response_id_with_a_passthrough_row_is_issued(self): - pc = _object_lookup(_batch_row(encode("azure", "user-1", UNISSUED_RESPONSE_ID))) + pc = _issued_count(1) assert ( await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "azure", pc) == "issued" ) - pc.db.litellm_managedobjecttable.find_first.assert_awaited_once_with( - where={"model_object_id": f"passthrough:azure:{UNISSUED_RESPONSE_ID}"} + pc.db.litellm_managedobjecttable.count.assert_awaited_once_with( + where={"model_object_id": {"in": [f"passthrough:azure:{UNISSUED_RESPONSE_ID}"]}} ) @pytest.mark.asyncio async def test_any_unissued_id_among_several_is_unissued(self): issued_id = "resp_1111111111111111issued" pc = _prisma_client_passthrough_managed() - pc.db.litellm_managedobjecttable.find_first = AsyncMock( - side_effect=lambda where: _batch_row("x") if where["model_object_id"].endswith(issued_id) else None + pc.db.litellm_managedobjecttable.count = AsyncMock( + side_effect=lambda where: sum(key.endswith(issued_id) for key in where["model_object_id"]["in"]) ) assert ( @@ -1782,15 +1782,15 @@ class TestAddressedResponseIdIssuance: ], ) async def test_requests_that_address_no_response_id_skip_the_lookup(self, path, body): - pc = _object_lookup(None) + pc = _issued_count(0) assert await addressed_response_id_issuance(path, body, "openai", pc) == "issued" - pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() + pc.db.litellm_managedobjecttable.count.assert_not_awaited() @pytest.mark.asyncio async def test_managed_ids_are_left_to_the_managed_id_resolver(self): managed_id = new_managed_id("openai", UNISSUED_RESPONSE_ID) - pc = _object_lookup(None) + pc = _issued_count(0) assert ( await addressed_response_id_issuance( @@ -1798,13 +1798,13 @@ class TestAddressedResponseIdIssuance: ) == "issued" ) - pc.db.litellm_managedobjecttable.find_first.assert_not_awaited() + pc.db.litellm_managedobjecttable.count.assert_not_awaited() @pytest.mark.asyncio @pytest.mark.parametrize("lookup_error", [httpx.ConnectError("db down"), PrismaError("query failed")]) async def test_failed_lookup_is_unverified_rather_than_issued_or_unissued(self, lookup_error): pc = _prisma_client_passthrough_managed() - pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=lookup_error) + pc.db.litellm_managedobjecttable.count = AsyncMock(side_effect=lookup_error) assert ( await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) @@ -1812,24 +1812,12 @@ class TestAddressedResponseIdIssuance: ) @pytest.mark.asyncio - @pytest.mark.parametrize(("issued_row", "expected"), [(None, "unissued"), (_batch_row("x"), "unverified")]) - async def test_a_failed_lookup_beside_another_id_never_hides_an_unissued_one(self, issued_row, expected): - failing_id = "resp_2222222222222222failing" + async def test_a_path_naming_many_response_ids_costs_one_lookup(self): + path = "/v1" + "".join(f"/responses/resp_{n:040d}" for n in range(500)) + pc = _issued_count(0) - async def find_first(where): - if where["model_object_id"].endswith(failing_id): - raise PrismaError("query failed") - return issued_row - - pc = _prisma_client_passthrough_managed() - pc.db.litellm_managedobjecttable.find_first = AsyncMock(side_effect=find_first) - - assert ( - await addressed_response_id_issuance( - f"/v1/responses/{failing_id}", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", pc - ) - == expected - ) + assert await addressed_response_id_issuance(path, {}, "openai", pc) == "unissued" + pc.db.litellm_managedobjecttable.count.assert_awaited_once() @pytest.mark.asyncio async def test_without_a_database_the_pass_through_issues_only_raw_ids_so_none_is_refused(self): diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f8d78b79851..e3afcf10046 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -3686,7 +3686,10 @@ def _passthrough_object_rows(owner_user_id: str | None) -> SimpleNamespace: return SimpleNamespace( db=SimpleNamespace( litellm_managedobjecttable=SimpleNamespace( - find_first=AsyncMock(return_value=row), upsert=AsyncMock(), update=AsyncMock() + find_first=AsyncMock(return_value=row), + count=AsyncMock(return_value=0 if row is None else 1), + upsert=AsyncMock(), + update=AsyncMock(), ) ) ) @@ -3820,7 +3823,7 @@ async def test_managed_pass_through_forwards_ids_it_may_address( async def test_managed_pass_through_answers_a_retryable_503_when_it_cannot_check_a_response_id(method, target, body): unreachable_database: Final = SimpleNamespace( db=SimpleNamespace( - litellm_managedobjecttable=SimpleNamespace(find_first=AsyncMock(side_effect=httpx.ConnectError("db down"))) + litellm_managedobjecttable=SimpleNamespace(count=AsyncMock(side_effect=httpx.ConnectError("db down"))) ) ) From dd41af8b71412b7ca4cb616d04766086bf581a6d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:39:33 -0700 Subject: [PATCH 6/9] fix(passthrough): chunk the response id ownership count through count_in --- litellm/proxy/pass_through_endpoints/managed_id_rewriter.py | 5 +++-- litellm/types/passthrough_endpoints/managed_id_rewriter.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 8013d7cf184..21965f073a5 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -55,6 +55,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( from litellm.proxy._types import DB_CONNECTION_ERROR_TYPES, UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames +from litellm.repositories.chunked_in import count_in from litellm.repositories.managed_file_repository import ManagedFileRepository from litellm.repositories.table_repositories import ManagedObjectRepository from litellm.types.llms.openai import BATCH_GUARDRAIL_RESPONSE_FIELD, OpenAIFileObject @@ -491,11 +492,11 @@ async def addressed_response_id_issuance( if prisma_client is None: return "issued" - object_ids: Final = sorted(f"passthrough:{provider}:{raw_id}" for raw_id in _addressed_raw_response_ids(path, body)) + object_ids: Final = tuple(f"passthrough:{provider}:{raw_id}" for raw_id in _addressed_raw_response_ids(path, body)) if not object_ids: return "issued" try: - issued_count: Final = await _object_table(prisma_client).count(where={"model_object_id": {"in": object_ids}}) + issued_count: Final = await count_in(_object_table(prisma_client), "model_object_id", object_ids) except (PrismaError, *DB_CONNECTION_ERROR_TYPES): verbose_proxy_logger.warning("managed_id_rewriter: response id ownership lookup failed", exc_info=True) return "unverified" diff --git a/litellm/types/passthrough_endpoints/managed_id_rewriter.py b/litellm/types/passthrough_endpoints/managed_id_rewriter.py index f98042afc85..14fc00953d0 100644 --- a/litellm/types/passthrough_endpoints/managed_id_rewriter.py +++ b/litellm/types/passthrough_endpoints/managed_id_rewriter.py @@ -81,7 +81,7 @@ class ManagedObjectTable(ManagedTable[ManagedObjectRow], Protocol): async def upsert(self, *, where: PrismaWhere, data: Mapping[str, ManagedRowData]) -> ManagedObjectRow: ... - async def count(self, *, where: PrismaWhere) -> int: ... + async def count(self, *, where: Mapping[str, object]) -> int: ... @runtime_checkable From 79f79ac8bfc324fdc940b02ddd25120e39594c67 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:44:35 -0700 Subject: [PATCH 7/9] test(integration): cover pass-through and relay response id issuance checks against a live proxy --- .../test_passthrough_response_id_issuance.py | 423 ++++++++++++++++++ 1 file changed, 423 insertions(+) create mode 100644 tests/integration/authorization/test_passthrough_response_id_issuance.py diff --git a/tests/integration/authorization/test_passthrough_response_id_issuance.py b/tests/integration/authorization/test_passthrough_response_id_issuance.py new file mode 100644 index 00000000000..1f18f56c9e6 --- /dev/null +++ b/tests/integration/authorization/test_passthrough_response_id_issuance.py @@ -0,0 +1,423 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal, assert_never +from urllib.parse import urlsplit + +import httpx +import psycopg +import pytest +from integration._support.client import Gateway, Scenario +from integration._support.database import scratch_database +from integration._support.process import graceful_stop_seconds, owned_proxy, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from pydantic import JsonValue, TypeAdapter + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_UPSTREAM_KEY: Final = "synthetic-openai-key" +_MODEL: Final = "gpt-5" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MANAGED: Final = "general_settings:\n passthrough_managed_object_ids: true\n" +_ALLOW_UNMANAGED: Final = ( + "general_settings:\n passthrough_managed_object_ids: true\n allow_unmanaged_response_ids: true\n" +) +_UNMANAGED_DETAIL: Final = ( + "Forbidden. This response id was not issued by this proxy, so the proxy cannot tell who owns it. " + "To let keys address responses this proxy did not issue, set " + "general_settings::allow_unmanaged_response_ids to True in the config.yaml file." +) +_LOOKUP_FAILED_DETAIL: Final = ( + "The proxy could not check who owns this response id because its database lookup failed. Retry the request." +) +_ACCESS_DENIED: Final = "Access denied to managed resource." +_NOT_FOUND: Final = "Managed resource not found." +_TAKE_OWNER_TABLE_OFFLINE: Final = ( + 'ALTER TABLE "LiteLLM_ManagedObjectTable" RENAME TO "LiteLLM_ManagedObjectTable_offline"' +) +_BRING_OWNER_TABLE_BACK: Final = ( + 'ALTER TABLE "LiteLLM_ManagedObjectTable_offline" RENAME TO "LiteLLM_ManagedObjectTable"' +) + +Route = Literal[ + "retrieve", + "input_items", + "delete", + "cancel", + "create", + "compact", + "input_tokens", + "openai_passthrough", + "percent_encoded", +] + + +def _vendor_id() -> str: + return f"resp_{uuid.uuid4().hex}" + + +def _response(identity: str) -> bytes: + return json.dumps( + { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": _MODEL, + "output": [], + } + ).encode() + + +def _requested_vendor_id(request: Request) -> str | None: + if not request.body.startswith(b"{"): + return None + requested: Final = _JSON_OBJECT.validate_json(request.body).get("input") + return requested if isinstance(requested, str) and requested.startswith("resp_") else None + + +def _respond(request: Request) -> Reply: + addressed: Final = next( + (segment for segment in urlsplit(request.target).path.split("/") if segment.startswith("resp_")), None + ) + return Reply(body=_response(addressed or _requested_vendor_id(request) or _vendor_id())) + + +def _forwarded(wire: Wire) -> tuple[tuple[str, str], ...]: + return tuple((request.method, urlsplit(request.target).path) for request in wire.drain()) + + +def _send(candidate: Gateway, key: str, method: str, path: str, body: bytes | None = None) -> httpx.Response: + json_header: Final = {"Content-Type": "application/json"} if body is not None else {} + return candidate.client.request( + method, path, content=body, headers={"Authorization": f"Bearer {key}", **json_header} + ) + + +def _continue(previous_response_id: str) -> bytes: + return json.dumps({"model": _MODEL, "input": "continue", "previous_response_id": previous_response_id}).encode() + + +def _user_key(scenario: Scenario) -> str: + return scenario.key(user_id=scenario.user()) + + +def _create(candidate: Gateway, key: str, vendor_id: str, prefix: str = "/openai") -> str: + created: Final = _send( + candidate, key, "POST", f"{prefix}/responses", json.dumps({"model": _MODEL, "input": vendor_id}).encode() + ) + assert created.status_code == 200, created.text + issued: Final = created.json()["id"] + assert isinstance(issued, str) and issued != vendor_id, created.text + return issued + + +def _assert_unmanaged_refusal(response: httpx.Response) -> None: + assert response.status_code == 403, response.text + assert response.json() == { + "error": {"message": _UNMANAGED_DETAIL, "type": "permission_error", "param": None, "code": "403"} + }, response.text + + +@contextmanager +def _proxy( + gateway: Gateway, + directory: Path, + upstream: Wire, + settings: str = _MANAGED, + overrides: Mapping[str, str] = MappingProxyType({}), +) -> Iterator[Gateway]: + config: Final = directory / f"passthrough-{uuid.uuid4().hex}.yaml" + config.write_text(settings) + with owned_proxy( + gateway, + directory, + {"OPENAI_API_BASE": upstream.url, "OPENAI_API_KEY": _UPSTREAM_KEY, **overrides}, + config=config, + ) as candidate: + yield candidate + + +@contextmanager +def _owner_table_offline(database_url: str) -> Iterator[None]: + with psycopg.connect(database_url, autocommit=True) as admin: + admin.execute(_TAKE_OWNER_TABLE_OFFLINE) + try: + yield + finally: + admin.execute(_BRING_OWNER_TABLE_BACK) + + +def _call(route: Route, vendor_id: str) -> tuple[str, str, bytes | None]: + match route: + case "retrieve": + return ("GET", f"/openai/responses/{vendor_id}", None) + case "input_items": + return ("GET", f"/openai/responses/{vendor_id}/input_items", None) + case "delete": + return ("DELETE", f"/openai/responses/{vendor_id}", None) + case "cancel": + return ("POST", f"/openai/responses/{vendor_id}/cancel", b"{}") + case "create": + return ("POST", "/openai/responses", _continue(vendor_id)) + case "compact": + return ("POST", "/openai/responses/compact", _continue(vendor_id)) + case "input_tokens": + return ("POST", "/openai/responses/input_tokens", _continue(vendor_id)) + case "openai_passthrough": + return ("GET", f"/openai_passthrough/v1/responses/{vendor_id}", None) + case "percent_encoded": + return ("GET", f"/openai/responses/%72{vendor_id[1:]}", None) + case _: + assert_never(route) + + +def test_the_owner_keeps_using_the_gateway_id_and_another_key_cannot(gateway: Gateway, tmp_path: Path) -> None: + vendor_id: Final = _vendor_id() + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + owner: Final = _user_key(scenario) + other: Final = _user_key(scenario) + with _proxy(gateway, tmp_path, upstream) as candidate: + issued: Final = _create(candidate, owner, vendor_id) + assert _forwarded(upstream) == (("POST", "/responses"),) + + retrieved: Final = _send(candidate, owner, "GET", f"/openai/responses/{issued}") + assert retrieved.status_code == 200, retrieved.text + assert retrieved.json()["id"] == issued, retrieved.text + assert _forwarded(upstream) == (("GET", f"/responses/{vendor_id}"),) + + continued: Final = _send(candidate, owner, "POST", "/openai/responses", _continue(issued)) + assert continued.status_code == 200, continued.text + sent: Final = upstream.drain() + assert [(request.method, request.target) for request in sent] == [("POST", "/responses")], sent + assert _JSON_OBJECT.validate_json(sent[0].body)["previous_response_id"] == vendor_id, sent[0].body + + refused: Final = _send(candidate, other, "GET", f"/openai/responses/{issued}") + assert refused.status_code == 403, refused.text + assert refused.json()["error"]["message"] == _ACCESS_DENIED, refused.text + assert _forwarded(upstream) == () + + +@pytest.mark.parametrize( + "route", + ( + "retrieve", + "input_items", + "delete", + "cancel", + "create", + "compact", + "input_tokens", + "openai_passthrough", + "percent_encoded", + ), +) +def test_a_vendor_id_the_gateway_never_issued_is_refused_before_it_reaches_openai( + gateway: Gateway, tmp_path: Path, route: Route +) -> None: + method, path, body = _call(route, _vendor_id()) + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + key: Final = _user_key(scenario) + with _proxy(gateway, tmp_path, upstream) as candidate: + _assert_unmanaged_refusal(_send(candidate, key, method, path, body)) + assert _forwarded(upstream) == () + + +def test_another_keys_vendor_id_is_not_found_and_never_forwarded(gateway: Gateway, tmp_path: Path) -> None: + vendor_id: Final = _vendor_id() + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + owner: Final = _user_key(scenario) + other: Final = _user_key(scenario) + with _proxy(gateway, tmp_path, upstream) as candidate: + _create(candidate, owner, vendor_id) + assert _forwarded(upstream) == (("POST", "/responses"),) + refusals: Final = tuple( + _send(candidate, other, *_call(route, vendor_id)) for route in ("retrieve", "create") + ) + assert [response.status_code for response in refusals] == [404, 404], [r.text for r in refusals] + assert _forwarded(upstream) == () + + +@pytest.mark.parametrize("exemption", ("admin", "allow_unmanaged", "flag_off")) +def test_exempt_callers_reach_openai_with_the_vendor_id( + gateway: Gateway, tmp_path: Path, exemption: Literal["admin", "allow_unmanaged", "flag_off"] +) -> None: + vendor_id: Final = _vendor_id() + settings: Final = {"admin": _MANAGED, "allow_unmanaged": _ALLOW_UNMANAGED, "flag_off": "general_settings: {}\n"}[ + exemption + ] + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + key: Final = gateway.key if exemption == "admin" else _user_key(scenario) + with _proxy(gateway, tmp_path, upstream, settings) as candidate: + retrieved: Final = _send(candidate, key, "GET", f"/openai/responses/{vendor_id}") + assert retrieved.status_code == 200, retrieved.text + continued: Final = _send(candidate, key, "POST", "/openai/responses", _continue(vendor_id)) + assert continued.status_code == 200, continued.text + sent: Final = upstream.drain() + assert [(request.method, request.target) for request in sent] == [ + ("GET", f"/responses/{vendor_id}"), + ("POST", "/responses"), + ], sent + assert _JSON_OBJECT.validate_json(sent[1].body)["previous_response_id"] == vendor_id, sent[1].body + + +@pytest.mark.parametrize( + "value", + (123, ["resp_listed"], "", None), + ids=("int", "list", "empty", "null"), +) +def test_a_previous_response_id_that_is_not_an_id_string_is_forwarded_unchanged( + gateway: Gateway, tmp_path: Path, value: JsonValue +) -> None: + body: Final = {"model": _MODEL, "input": "hostile", "previous_response_id": value} + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + key: Final = _user_key(scenario) + with _proxy(gateway, tmp_path, upstream) as candidate: + response: Final = _send(candidate, key, "POST", "/openai/responses", json.dumps(body).encode()) + assert response.status_code == 200, response.text + sent: Final = upstream.drain() + assert [(request.method, request.target) for request in sent] == [("POST", "/responses")], sent + assert _JSON_OBJECT.validate_json(sent[0].body) == body, sent[0].body + + +def test_an_oversized_vendor_id_is_refused_before_it_reaches_openai(gateway: Gateway, tmp_path: Path) -> None: + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + key: Final = _user_key(scenario) + with _proxy(gateway, tmp_path, upstream) as candidate: + _assert_unmanaged_refusal( + _send(candidate, key, "POST", "/openai/responses", _continue("resp_" + "a" * 5000)) + ) + assert _forwarded(upstream) == () + + +def test_a_duplicated_previous_response_id_is_judged_and_forwarded_as_one_value( + gateway: Gateway, tmp_path: Path +) -> None: + owned_vendor_id: Final = _vendor_id() + foreign_vendor_id: Final = _vendor_id() + with wire_server(_respond) as upstream, gateway.scenario() as scenario: + key: Final = _user_key(scenario) + with _proxy(gateway, tmp_path, upstream) as candidate: + issued: Final = _create(candidate, key, owned_vendor_id) + upstream.drain() + vendor_last: Final = ( + '{"model":"%s","input":"twice","previous_response_id":"%s","previous_response_id":"%s"}' + % (_MODEL, issued, foreign_vendor_id) + ).encode() + _assert_unmanaged_refusal(_send(candidate, key, "POST", "/openai/responses", vendor_last)) + assert _forwarded(upstream) == () + + issued_last: Final = ( + '{"model":"%s","input":"twice","previous_response_id":"%s","previous_response_id":"%s"}' + % (_MODEL, foreign_vendor_id, issued) + ).encode() + continued: Final = _send(candidate, key, "POST", "/openai/responses", issued_last) + assert continued.status_code == 200, continued.text + sent: Final = upstream.drain() + assert [(request.method, request.target) for request in sent] == [("POST", "/responses")], sent + assert sent[0].body.count(b"previous_response_id") == 1, sent[0].body + assert _JSON_OBJECT.validate_json(sent[0].body)["previous_response_id"] == owned_vendor_id, sent[0].body + + +def test_a_failed_owner_lookup_answers_a_retryable_503_and_forwards_nothing(gateway: Gateway, tmp_path: Path) -> None: + vendor_id: Final = _vendor_id() + with ( + scratch_database("passthrough_lookup") as database_url, + wire_server(_respond) as upstream, + _proxy(gateway, tmp_path, upstream, overrides={"DATABASE_URL": database_url}) as candidate, + candidate.scenario() as scenario, + ): + key: Final = _user_key(scenario) + with _owner_table_offline(database_url): + failed: Final = _send(candidate, key, "GET", f"/openai/responses/{vendor_id}") + assert failed.status_code == 503, failed.text + assert failed.json()["error"]["message"] == _LOOKUP_FAILED_DETAIL, failed.text + assert _forwarded(upstream) == () + _assert_unmanaged_refusal(_send(candidate, key, "GET", f"/openai/responses/{vendor_id}")) + assert _forwarded(upstream) == () + + +def test_the_model_list_relay_hands_out_and_checks_gateway_ids(gateway: Gateway, tmp_path: Path) -> None: + vendor_id: Final = _vendor_id() + with ( + wire_server(_respond) as azure, + wire_server(_respond) as openai_upstream, + _proxy(gateway, tmp_path, openai_upstream) as candidate, + candidate.scenario() as scenario, + ): + group: Final = scenario.model(model=f"azure/{_MODEL}", api_base=azure.url, api_key=_UPSTREAM_KEY) + wait_for_worker_readiness(candidate, group, 1, 1) + owner: Final = _user_key(scenario) + other: Final = _user_key(scenario) + relay: Final = f"/azure/{group}/openai/v1" + issued: Final = _create(candidate, owner, vendor_id, relay) + assert _forwarded(azure) == (("POST", "/openai/v1/responses"),) + + retrieved: Final = _send(candidate, owner, "GET", f"{relay}/responses/{issued}") + assert retrieved.status_code == 200, retrieved.text + assert retrieved.json()["id"] == issued, retrieved.text + assert _forwarded(azure) == (("GET", f"/openai/v1/responses/{vendor_id}"),) + + refusals: Final = tuple( + _send(candidate, other, "GET", path) + for path in ( + f"/azure/{group}/openai/v1/responses/{_vendor_id()}", + f"/azure_ai/{group}/openai/v1/responses/{_vendor_id()}", + f"{relay}/responses/{issued}", + f"{relay}/responses/{vendor_id}", + ) + ) + assert [(response.status_code, response.json()) for response in refusals] == [ + (403, {"detail": _UNMANAGED_DETAIL}), + (403, {"detail": _UNMANAGED_DETAIL}), + (403, {"detail": _ACCESS_DENIED}), + (404, {"detail": _NOT_FOUND}), + ], [response.text for response in refusals] + assert _forwarded(azure) == () + assert _forwarded(openai_upstream) == () + + +def test_every_verdict_holds_through_a_burst_after_a_worker_is_replaced(gateway: Gateway, tmp_path: Path) -> None: + vendor_ids: Final = tuple(_vendor_id() for _ in range(10)) + config: Final = tmp_path / "passthrough-burst.yaml" + config.write_text(_MANAGED) + with ( + wire_server(_respond) as upstream, + gateway.scenario() as scenario, + owned_proxy_process( + gateway, + tmp_path, + {"OPENAI_API_BASE": upstream.url, "OPENAI_API_KEY": _UPSTREAM_KEY}, + config=config, + workers=2, + ) as owned, + ): + owner: Final = _user_key(scenario) + other: Final = _user_key(scenario) + issued: Final = tuple(_create(owned.gateway, owner, vendor_id) for vendor_id in vendor_ids) + upstream.drain() + workers: Final = owned.worker_pids() + assert len(workers) >= 2, workers + owned.kill_worker(min(workers)) + + calls: Final = ( + *((owner, f"/openai/responses/{identity}", 200) for identity in issued), + *((other, f"/openai/responses/{_vendor_id()}", 403) for _ in range(10)), + *((other, f"/openai/responses/{vendor_id}", 404) for vendor_id in vendor_ids), + ) + + def call(entry: tuple[str, str, int]) -> tuple[int, int]: + key, path, expected = entry + return expected, _send(owned.gateway, key, "GET", path).status_code + + with ThreadPoolExecutor(max_workers=10) as pool: + outcomes: Final = tuple(pool.map(call, calls)) + assert all(expected == observed for expected, observed in outcomes), outcomes + assert sorted(_forwarded(upstream)) == sorted(("GET", f"/responses/{vendor_id}") for vendor_id in vendor_ids) From 1a99d6971c749b3eb77b409e266ffb25c970cca3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 13:44:52 -0700 Subject: [PATCH 8/9] test(integration): cover the Azure AI Search index branch of the managed pass-through --- .../test_passthrough_response_id_issuance.py | 75 +++++++++++++++++++ 1 file changed, 75 insertions(+) diff --git a/tests/integration/authorization/test_passthrough_response_id_issuance.py b/tests/integration/authorization/test_passthrough_response_id_issuance.py index 1f18f56c9e6..fb03aaf6ae4 100644 --- a/tests/integration/authorization/test_passthrough_response_id_issuance.py +++ b/tests/integration/authorization/test_passthrough_response_id_issuance.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import uuid from collections.abc import Iterator, Mapping from concurrent.futures import ThreadPoolExecutor @@ -46,6 +47,10 @@ _BRING_OWNER_TABLE_BACK: Final = ( 'ALTER TABLE "LiteLLM_ManagedObjectTable_offline" RENAME TO "LiteLLM_ManagedObjectTable"' ) +_SEARCH_STORE: Final = "integration-search-store" +_SEARCH_KEY: Final = "synthetic-search-key" +_SEARCH_RESULTS: Final = Reply(body=json.dumps({"value": []}).encode()) + Route = Literal[ "retrieve", "input_items", @@ -155,6 +160,51 @@ def _owner_table_offline(database_url: str) -> Iterator[None]: admin.execute(_BRING_OWNER_TABLE_BACK) +def _search_settings(search: Wire) -> str: + return ( + _MANAGED + + "vector_store_registry:\n" + + f" - vector_store_name: {_SEARCH_STORE}\n" + + " litellm_params:\n" + + f" vector_store_id: {_SEARCH_STORE}\n" + + " custom_llm_provider: azure_ai\n" + + f" api_key: {_SEARCH_KEY}\n" + + f" api_base: {search.url}\n" + ) + + +@contextmanager +def _search_index(gateway: Gateway) -> Iterator[str]: + index: Final = f"integration-index-{uuid.uuid4().hex}" + gateway.post( + "/v1/indexes", + {"index_name": index, "litellm_params": {"vector_store_index": index, "vector_store_name": _SEARCH_STORE}}, + ) + try: + yield index + finally: + with psycopg.connect(os.environ["DATABASE_URL"], autocommit=True) as admin: + admin.execute('DELETE FROM "LiteLLM_ManagedVectorStoreIndexTable" WHERE index_name = %s', (index,)) + + +@contextmanager +def _search_rig(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, Wire, str, str]]: + with ( + wire_server(lambda _request: _SEARCH_RESULTS) as search, + wire_server(_respond) as openai_upstream, + _search_index(gateway) as index, + _proxy(gateway, directory, openai_upstream, settings=_search_settings(search)) as candidate, + candidate.scenario() as scenario, + ): + scenario.model(model=f"azure/{_MODEL}", api_base=openai_upstream.url, api_key=_UPSTREAM_KEY) + reader: Final = scenario.key( + user_id=scenario.user(), + metadata={"allowed_vector_store_indexes": [{"index_name": index, "index_permissions": ["read"]}]}, + ) + yield candidate, search, index, reader + assert _forwarded(openai_upstream) == () + + def _call(route: Route, vendor_id: str) -> tuple[str, str, bytes | None]: match route: case "retrieve": @@ -384,6 +434,31 @@ def test_the_model_list_relay_hands_out_and_checks_gateway_ids(gateway: Gateway, assert _forwarded(openai_upstream) == () +def test_an_azure_ai_search_query_reaches_the_index_unchanged(gateway: Gateway, tmp_path: Path) -> None: + query: Final = {"search": "mango", "top": 3} + with _search_rig(gateway, tmp_path) as (candidate, search, index, reader): + answered: Final = _send( + candidate, reader, "POST", f"/azure_ai/indexes/{index}/docs/search", json.dumps(query).encode() + ) + assert answered.status_code == 200, answered.text + assert answered.json() == {"value": []}, answered.text + received: Final = search.drain() + assert [(request.method, urlsplit(request.target).path) for request in received] == [ + ("POST", f"/indexes/{index}/docs/search") + ] + assert [_JSON_OBJECT.validate_json(request.body) for request in received] == [query] + assert [request.headers.get("api-key") for request in received] == [_SEARCH_KEY] + + +def test_an_azure_ai_search_query_carrying_a_vendor_response_id_is_refused(gateway: Gateway, tmp_path: Path) -> None: + query: Final = {"search": "mango", "previous_response_id": _vendor_id()} + with _search_rig(gateway, tmp_path) as (candidate, search, index, reader): + _assert_unmanaged_refusal( + _send(candidate, reader, "POST", f"/azure_ai/indexes/{index}/docs/search", json.dumps(query).encode()) + ) + assert _forwarded(search) == () + + def test_every_verdict_holds_through_a_burst_after_a_worker_is_replaced(gateway: Gateway, tmp_path: Path) -> None: vendor_ids: Final = tuple(_vendor_id() for _ in range(10)) config: Final = tmp_path / "passthrough-burst.yaml" From 8ef7210311cb4dbe59832857022869641f09331b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:02:15 -0700 Subject: [PATCH 9/9] test(integration): cover error and non-object bodies on the model_list relay --- .../test_passthrough_response_id_issuance.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/integration/authorization/test_passthrough_response_id_issuance.py b/tests/integration/authorization/test_passthrough_response_id_issuance.py index fb03aaf6ae4..20a9386e945 100644 --- a/tests/integration/authorization/test_passthrough_response_id_issuance.py +++ b/tests/integration/authorization/test_passthrough_response_id_issuance.py @@ -459,6 +459,41 @@ def test_an_azure_ai_search_query_carrying_a_vendor_response_id_is_refused(gatew assert _forwarded(search) == () +def test_the_model_list_relay_passes_error_and_non_object_bodies_through_unchanged( + gateway: Gateway, tmp_path: Path +) -> None: + vendor_id: Final = _vendor_id() + refused: Final = json.dumps( + {"error": {"message": "invalid input", "type": "invalid_request_error"}, "id": vendor_id} + ) + listed: Final = json.dumps([{"id": vendor_id, "object": "response"}]) + replies: Final = MappingProxyType( + {"refused": Reply(status=400, body=refused.encode()), "listed": Reply(body=listed.encode())} + ) + with ( + wire_server(lambda request: replies[_JSON_OBJECT.validate_json(request.body)["input"]]) as azure, + wire_server(_respond) as openai_upstream, + _proxy(gateway, tmp_path, openai_upstream) as candidate, + candidate.scenario() as scenario, + ): + group: Final = scenario.model(model=f"azure/{_MODEL}", api_base=azure.url, api_key=_UPSTREAM_KEY) + wait_for_worker_readiness(candidate, group, 1, 1) + key: Final = _user_key(scenario) + answers: Final = tuple( + _send( + candidate, + key, + "POST", + f"/azure/{group}/openai/v1/responses", + json.dumps({"model": _MODEL, "input": kind}).encode(), + ) + for kind in ("refused", "listed") + ) + assert [(answer.status_code, answer.text) for answer in answers] == [(400, refused), (200, listed)] + assert _forwarded(azure) == (("POST", "/openai/v1/responses"), ("POST", "/openai/v1/responses")) + assert _forwarded(openai_upstream) == () + + def test_every_verdict_holds_through_a_burst_after_a_worker_is_replaced(gateway: Gateway, tmp_path: Path) -> None: vendor_ids: Final = tuple(_vendor_id() for _ in range(10)) config: Final = tmp_path / "passthrough-burst.yaml"