diff --git a/litellm/proxy/hooks/responses_id_security.py b/litellm/proxy/hooks/responses_id_security.py index 27b5924c540..7c6d34084c2 100644 --- a/litellm/proxy/hooks/responses_id_security.py +++ b/litellm/proxy/hooks/responses_id_security.py @@ -43,7 +43,7 @@ _ADDRESSED_RESPONSE_ID_FIELD_BY_CALL_TYPE: Final[Mapping[str, str]] = MappingPro } ) -_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." @@ -57,7 +57,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 @@ -66,6 +66,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]) @@ -119,7 +133,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 @@ -155,21 +169,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/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index f2605a8ec6f..ad3f0a30f75 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, ValidationError +from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from starlette.websockets import WebSocketState from typing_extensions import ReadOnly, TypedDict @@ -55,6 +55,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 ( @@ -112,6 +113,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 @@ -143,7 +145,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 @@ -2168,9 +2172,103 @@ async def _relay_response_bytes(upstream: httpx.Response) -> AsyncGenerator[byte 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 + + 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) + 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), ) @@ -2184,6 +2282,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) @@ -2202,6 +2301,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, ) @@ -2213,6 +2313,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( @@ -2235,13 +2336,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, ), @@ -2252,7 +2353,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, ), @@ -2317,6 +2418,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), ) @@ -2336,16 +2438,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), ) @@ -2399,21 +2522,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 a5a04fb68d2..21965f073a5 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, @@ -51,9 +52,10 @@ 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.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 @@ -70,6 +72,7 @@ from litellm.types.passthrough_endpoints.managed_id_rewriter import ( PrismaWhere, PrismaWhereValue, ResourceKind, + ResponseIdIssuance, SortOrder, ) @@ -178,8 +181,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 +471,38 @@ 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 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" + 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 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" + return "issued" if issued_count == len(object_ids) else "unissued" + + # --------------------------------------------------------------------------- # OUTPUT path — helpers for minting and storing managed IDs # --------------------------------------------------------------------------- @@ -817,7 +854,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 311dab058a4..9d452f6b6a6 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, @@ -153,6 +154,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() @@ -161,6 +163,7 @@ pass_through_endpoint_logging: Final = PassThroughEndpointLogging() _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) +_REQUEST_BODY_VIEW: Final = TypeAdapter(Mapping[str, object]) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -1374,6 +1377,14 @@ async def pass_through_request( prisma_client as _passthrough_prisma, ) + await refuse_unissued_response_ids( + url.path, + _REQUEST_BODY_VIEW.validate_python(_parsed_body or {}), + _managed_id_provider, + _passthrough_prisma, + user_api_key_dict, + ) + _original_path: Final = url.path _original_query_params: Final = requested_query_params _original_body: Final = _parsed_body @@ -2918,6 +2929,41 @@ 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 ( + 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 + 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: """ Decide from the response headers whether the body must be read into memory. diff --git a/litellm/types/passthrough_endpoints/managed_id_rewriter.py b/litellm/types/passthrough_endpoints/managed_id_rewriter.py index 0603c13d905..14fc00953d0 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]" @@ -80,6 +81,8 @@ class ManagedObjectTable(ManagedTable[ManagedObjectRow], Protocol): async def upsert(self, *, where: PrismaWhere, data: Mapping[str, ManagedRowData]) -> ManagedObjectRow: ... + async def count(self, *, where: Mapping[str, object]) -> int: ... + @runtime_checkable class ManagedFileIdReader(Protocol): 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..1258cb0797d --- /dev/null +++ b/tests/integration/authorization/test_passthrough_response_id_issuance.py @@ -0,0 +1,533 @@ +from __future__ import annotations + +import json +import os +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: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] + +_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"' +) + +_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", + "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 _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": + 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_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_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" + 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) 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 4dbe1f2ef8d..21bd32ea1c8 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 @@ -7503,6 +7503,189 @@ 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), + count=AsyncMock(side_effect=self._count), + 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 _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"])) + + +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 + @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 + + 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 49bc2e9c29e..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 @@ -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( @@ -17,6 +19,7 @@ from litellm.proxy.pass_through_endpoints.managed_id_rewriter import( _MAX_RAW_ID_GUARD_LOOKUPS, _passthrough_provider_marker, _resolve_one, + addressed_response_id_issuance, is_passthrough_list_route, list_passthrough_ids_from_db, rewrite_body_ids, @@ -1709,6 +1712,127 @@ 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 _issued_count(issued: int) -> MagicMock: + pc = _prisma_client_passthrough_managed() + pc.db.litellm_managedobjecttable.count = AsyncMock(return_value=issued) + return pc + + +class TestAddressedResponseIdIssuance: + @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 = _issued_count(0) + + assert await addressed_response_id_issuance(path, body, "openai", pc) == "unissued" + 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 = _issued_count(1) + + assert ( + await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "azure", pc) + == "issued" + ) + 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.count = AsyncMock( + side_effect=lambda where: sum(key.endswith(issued_id) for key in where["model_object_id"]["in"]) + ) + + assert ( + await addressed_response_id_issuance( + f"/v1/responses/{issued_id}", {"previous_response_id": UNISSUED_RESPONSE_ID}, "openai", pc + ) + == "unissued" + ) + + @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 = _issued_count(0) + + assert await addressed_response_id_issuance(path, body, "openai", pc) == "issued" + 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 = _issued_count(0) + + assert ( + await addressed_response_id_issuance( + f"/v1/responses/{managed_id}", {"previous_response_id": managed_id}, "openai", pc + ) + == "issued" + ) + 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.count = AsyncMock(side_effect=lookup_error) + + assert ( + await addressed_response_id_issuance(f"/v1/responses/{UNISSUED_RESPONSE_ID}", {}, "openai", pc) + == "unverified" + ) + + @pytest.mark.asyncio + 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) + + 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): + assert ( + 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 + ) + == "issued" + ) + + @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 5412ea56ac0..b91f7bb0d17 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,179 @@ 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), + count=AsyncMock(return_value=0 if row is None else 1), + 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), + ("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", + {"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 +@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(count=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( + "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): """