diff --git a/litellm/proxy/shunt_endpoints/endpoints.py b/litellm/proxy/shunt_endpoints/endpoints.py index 1061357cdff..a89dfe2b2b5 100644 --- a/litellm/proxy/shunt_endpoints/endpoints.py +++ b/litellm/proxy/shunt_endpoints/endpoints.py @@ -16,10 +16,12 @@ from typing import TYPE_CHECKING, Annotated, Final from fastapi import APIRouter, Depends, File, Form, Header, HTTPException, Query, Request, UploadFile from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.env_utils import get_env_int from litellm.proxy._types import LitellmUserRoles, ProxyErrorTypes, ProxyException, UserAPIKeyAuth +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.guardrails.auto_router_shunt import ShuntConfig, shunt_config_for_model from litellm.proxy.guardrails.shunt_capability_token import open_shunt_capability_token -from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.route_llm_request import route_request from litellm.proxy.shunt_endpoints.worker import ( BULK_READ_SYSTEM_PROMPT, CODE_WRITE_SYSTEM_PROMPT, @@ -28,7 +30,8 @@ from litellm.proxy.shunt_endpoints.worker import ( build_code_write_message, strip_code_fences, ) -from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage, ChatCompletionUserMessage +from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage +from litellm.types.utils import Choices, ModelResponse if TYPE_CHECKING: from litellm.router import Router @@ -117,6 +120,7 @@ def _worker_config( async def _worker_text( + request: Request, llm_router: "Router", *, model: str, @@ -127,49 +131,95 @@ async def _worker_text( ) -> str: """The worker model's reply text, or a 502 if it produced none. - One call site for both endpoints, since they differ only in model, system prompt, and - message. Attribution reuses the proxy's own key-metadata builder so the call is billed and - budgeted against the calling key/user/team/org like a normal request. - - `proxy_logging_obj.pre_call_hook` runs first: `llm_router.acompletion` alone skips every - rate-limit and budget callback, since those register as `async_pre_call_hook` and only - `/chat/completions` and friends normally walk that list before routing. Without this call a - caller already over budget or rate-limited could keep spending through this endpoint. + Goes through the same `common_processing_pre_call_logic` + `route_request` pipeline + `/chat/completions` and every other LLM-calling route uses, rather than calling + `llm_router.acompletion` directly: that pipeline is what actually applies model-level + guardrails, budget/rate-limit enforcement, and fallbacks to the model being called, and a + hand-rolled call here would have to re-derive each of those separately and correctly. + Skips only the HTTP response/streaming shaping half of that pipeline, since a worker call + is never itself an HTTP response and is never streamed. """ - from litellm.proxy.proxy_server import proxy_logging_obj + from litellm.proxy.proxy_server import general_settings, proxy_config, proxy_logging_obj system: Final = ChatCompletionSystemMessage(role="system", content=system_prompt) user: Final = ChatCompletionUserMessage(role="user", content=message) - messages: Final[ - list[AllMessageValues] - ] = [ # mutable-ok: shared between pre_call_hook and acompletion, both take a list - system, - user, - ] - key_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key( - user_api_key_dict=user_api_key_dict + processor: Final = ProxyBaseLLMRequestProcessing( + data={"model": model, "messages": [system, user], "temperature": WORKER_TEMPERATURE, "stream": False} ) - metadata: Final = {**key_metadata, "user_api_key": user_api_key_dict.api_key} # mutable-ok: same - request_data: Final = { # mutable-ok: pre_call_hook's own signature takes a plain dict - "model": model, - "messages": messages, - "metadata": metadata, - } - await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=request_data, call_type="acompletion" - ) - - response: Final = await llm_router.acompletion( - model=model, messages=messages, temperature=WORKER_TEMPERATURE, stream=False, metadata=metadata - ) - text: Final = response.choices[0].message.content + try: + data, _logging_obj = await processor.common_processing_pre_call_logic( # pyright: ignore[reportUnknownVariableType] # common_processing_pre_call_logic's own dict/Logging return is unrefined at this call shape + request=request, + general_settings=general_settings, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + proxy_config=proxy_config, + route_type="acompletion", + llm_router=llm_router, + ) + response: Final = await route_request( # pyright: ignore[reportUnknownVariableType] # route_request's own return type is intentionally an untyped union (see its ANN202 suppression) + data=data, + route_type="acompletion", + llm_router=llm_router, + user_model=None, + user_api_key_dict=user_api_key_dict, + ) + processed: Final = await proxy_logging_obj.post_call_success_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, # pyright: ignore[reportArgumentType] # response is the same real ModelResponse a router acompletion call returns; the hook's own signature just can't narrow it here + ) + except Exception as e: # noqa: BLE001 # _handle_llm_api_exception must see every failure mode a real request can hit, same as proxy_server.py's own catch-all here + raise await processor._handle_llm_api_exception( # pyright: ignore[reportPrivateUsage] # same cross-module call proxy_server.py's own /chat/completions and /embeddings routes already make + e=e, user_api_key_dict=user_api_key_dict, proxy_logging_obj=proxy_logging_obj + ) + if not isinstance(processed, ModelResponse): + raise HTTPException(status_code=502, detail=f"shunt {label}: worker model returned no completion") + choice: Final = processed.choices[0] if processed.choices else None + text: Final = choice.message.content if isinstance(choice, Choices) else None if not isinstance(text, str): raise HTTPException(status_code=502, detail=f"shunt {label}: worker model returned no text") return text -async def _read_upload_text(upload: UploadFile) -> str: - content: Final = await upload.read() +# A global request-size limit exists (RequestSizeLimitMiddleware) but is opt-in and +# premium-gated, so these endpoints cannot rely on it: they accept arbitrary caller-supplied +# multipart uploads specifically to hand their contents to a worker model, an authenticated +# caller with a valid capability token could otherwise upload enough data to exhaust a proxy +# worker's memory before the size limit ever runs. +_MAX_UPLOAD_BYTES_PER_FILE: Final = get_env_int("LITELLM_SHUNT_MAX_UPLOAD_BYTES_PER_FILE", 1024 * 1024) +_MAX_UPLOAD_BYTES_TOTAL: Final = get_env_int("LITELLM_SHUNT_MAX_UPLOAD_BYTES_TOTAL", 8 * 1024 * 1024) +_MAX_UPLOAD_FILE_COUNT: Final = get_env_int("LITELLM_SHUNT_MAX_UPLOAD_FILE_COUNT", 20) + + +async def _read_upload_text(upload: UploadFile, *, remaining_total_bytes: int) -> str: + """`upload`'s content as UTF-8 text, reading at most the smaller of the per-file and + remaining-total byte budgets -- never the whole file, so a caller can't force this endpoint + to buffer more than that regardless of how large the real upload is. + + `remaining_total_bytes <= 0` is checked explicitly rather than left to `min()` + + `.read(limit + 1)`: a non-positive `remaining_total_bytes` would make `limit` zero or + negative, and `UploadFile.read` treats a negative size as "read the whole file", which + would silently defeat this budget for any caller of this function that ever passes one. + `_read_upload_texts` below never actually produces a negative value (each read is already + bounded by what was left when it started), so this is the function's own contract holding + regardless of caller, not a path reachable through that call site today. + """ + if remaining_total_bytes <= 0: + raise ProxyException( + message=f"uploads exceed the {_MAX_UPLOAD_BYTES_TOTAL}-byte total limit for this call", + type=ProxyErrorTypes.bad_request_error, + param="paths", + code=400, + ) + limit: Final = min(_MAX_UPLOAD_BYTES_PER_FILE, remaining_total_bytes) + content: Final = await upload.read(limit + 1) + if len(content) > limit: + raise ProxyException( + message=f"'{upload.filename}' exceeds the {limit}-byte upload limit for this call", + type=ProxyErrorTypes.bad_request_error, + param="paths" if upload.filename else "file", + code=400, + ) try: return content.decode("utf-8") except UnicodeDecodeError as e: @@ -181,6 +231,30 @@ async def _read_upload_text(upload: UploadFile) -> str: ) from e +async def _read_upload_texts(uploads: Sequence[UploadFile]) -> tuple[str, ...]: + """Every upload's text, in order, enforcing the aggregate byte and file-count budgets + across the whole request rather than per file. + + A plain accumulator, not a comprehension: each read's byte budget is whatever the + aggregate limit has left after every earlier file in this same request, so the reads are + inherently sequential and the running total has to be rebound as they complete. + """ + if len(uploads) > _MAX_UPLOAD_FILE_COUNT: + raise ProxyException( + message=f"at most {_MAX_UPLOAD_FILE_COUNT} files are allowed per call", + type=ProxyErrorTypes.bad_request_error, + param="paths", + code=400, + ) + texts: tuple[str, ...] = () # rebind-ok: sequential running total, see docstring + remaining_total_bytes: int = _MAX_UPLOAD_BYTES_TOTAL # rebind-ok: same + for upload in uploads: + text = await _read_upload_text(upload, remaining_total_bytes=remaining_total_bytes) + texts = (*texts, text) + remaining_total_bytes -= len(text.encode("utf-8")) + return texts + + @router.post( "/v1/bulk_read", dependencies=_AUTH_DEPENDENCIES, @@ -206,11 +280,15 @@ async def bulk_read( """ llm_router, config = _worker_config(router_name, user_api_key_dict, tags or ()) - files: Final = MappingProxyType({upload.filename or "unnamed": await _read_upload_text(upload) for upload in paths}) + texts: Final = await _read_upload_texts(paths) + files: Final = MappingProxyType( + {upload.filename or "unnamed": text for upload, text in zip(paths, texts, strict=True)} + ) message: Final = build_bulk_read_message(question=question, files=files) verbose_proxy_logger.debug("shunt bulk_read: %s file(s) via %s", len(files), config.bulk_read_model) return await _worker_text( + request, llm_router, model=config.bulk_read_model, system_prompt=BULK_READ_SYSTEM_PROMPT, @@ -246,7 +324,7 @@ async def code_write( """ llm_router, config = _worker_config(router_name, user_api_key_dict, tags or ()) - reference_content: Final = await _read_upload_text(reference) + reference_content: Final = await _read_upload_text(reference, remaining_total_bytes=_MAX_UPLOAD_BYTES_TOTAL) message: Final = build_code_write_message( spec=spec, reference_path=reference.filename or "", reference_content=reference_content ) @@ -254,6 +332,7 @@ async def code_write( verbose_proxy_logger.debug("shunt code_write: reference=%s via %s", reference.filename, config.code_write_model) return strip_code_fences( await _worker_text( + request, llm_router, model=config.code_write_model, system_prompt=CODE_WRITE_SYSTEM_PROMPT, diff --git a/tests/test_litellm/proxy/shunt_endpoints/test_endpoints.py b/tests/test_litellm/proxy/shunt_endpoints/test_endpoints.py deleted file mode 100644 index f1b2bcd06a6..00000000000 --- a/tests/test_litellm/proxy/shunt_endpoints/test_endpoints.py +++ /dev/null @@ -1,177 +0,0 @@ -"""Unit tests for litellm.proxy.shunt_endpoints.endpoints.""" - -import pytest -from fastapi import HTTPException - -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth -from litellm.proxy.guardrails.shunt_capability_token import mint_shunt_capability_token -from litellm.proxy.shunt_endpoints.endpoints import _caller_from_capability_token, _worker_text - - -@pytest.fixture(autouse=True) -def _salt_key(monkeypatch): - monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234-test-salt-key") - - -class TestMissingOrMalformedHeader: - @pytest.mark.asyncio - async def test_no_header_is_rejected(self): - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization=None) - assert exc_info.value.status_code == 401 - - @pytest.mark.asyncio - async def test_non_bearer_header_is_rejected(self): - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization="Basic abc123") - assert exc_info.value.status_code == 401 - - @pytest.mark.asyncio - async def test_malformed_token_is_rejected(self): - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization="Bearer not-a-real-token") - assert exc_info.value.status_code == 401 - - @pytest.mark.asyncio - async def test_expired_token_is_rejected(self, monkeypatch): - token = mint_shunt_capability_token(key_hash="deadbeef", master_key=None, now=1_000_000) - # 121s after mint, one second past the 120s TTL. - import litellm.proxy.guardrails.shunt_capability_token as token_mod - - monkeypatch.setattr(token_mod.time, "time", lambda: 1_000_121) - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization=f"Bearer {token}") - assert exc_info.value.status_code == 401 - - -class TestKeyHashGrant: - @pytest.mark.asyncio - async def test_resolves_the_key_object_for_the_grants_hash(self, monkeypatch): - token = mint_shunt_capability_token(key_hash="deadbeef", master_key=None) - resolved = UserAPIKeyAuth(api_key="deadbeef", team_id="team-1") - - async def _fake_get_key_object(**kwargs): - assert kwargs["hashed_token"] == "deadbeef" - return resolved - - import litellm.proxy.auth.auth_checks as auth_checks - - monkeypatch.setattr(auth_checks, "get_key_object", _fake_get_key_object) - result = await _caller_from_capability_token(authorization=f"Bearer {token}") - assert result is resolved - - @pytest.mark.asyncio - async def test_a_lookup_failure_is_rejected_not_propagated(self, monkeypatch): - token = mint_shunt_capability_token(key_hash="deadbeef", master_key=None) - - async def _raising_get_key_object(**kwargs): - raise Exception("key not found") - - import litellm.proxy.auth.auth_checks as auth_checks - - monkeypatch.setattr(auth_checks, "get_key_object", _raising_get_key_object) - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization=f"Bearer {token}") - assert exc_info.value.status_code == 401 - - -class TestMasterKeyGrant: - @pytest.mark.asyncio - async def test_matching_master_key_resolves_as_proxy_admin(self, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-the-real-master-key") - token = mint_shunt_capability_token(key_hash=None, master_key="sk-the-real-master-key") - result = await _caller_from_capability_token(authorization=f"Bearer {token}") - assert result.user_role == LitellmUserRoles.PROXY_ADMIN - - # Regression: a resolved master-key caller carried the raw master key as its own api_key, - # which _worker_text later places in the outbound request's metadata["user_api_key"] -- - # reachable by any raw-metadata logging callback. Normal master-key auth substitutes a - # stable alias there specifically to keep the real key out of that sink; this must match. - @pytest.mark.asyncio - async def test_resolved_caller_never_carries_the_raw_master_key(self, monkeypatch): - from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS - - monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-the-real-master-key") - token = mint_shunt_capability_token(key_hash=None, master_key="sk-the-real-master-key") - result = await _caller_from_capability_token(authorization=f"Bearer {token}") - assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS - assert result.api_key != "sk-the-real-master-key" - - @pytest.mark.asyncio - async def test_master_key_mismatch_is_rejected(self, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-the-current-master-key") - token = mint_shunt_capability_token(key_hash=None, master_key="sk-a-stale-master-key") - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization=f"Bearer {token}") - assert exc_info.value.status_code == 401 - - @pytest.mark.asyncio - async def test_no_configured_master_key_rejects_a_master_key_grant(self, monkeypatch): - monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) - token = mint_shunt_capability_token(key_hash=None, master_key="sk-anything") - with pytest.raises(HTTPException) as exc_info: - await _caller_from_capability_token(authorization=f"Bearer {token}") - assert exc_info.value.status_code == 401 - - -class _FakeRouter: - def __init__(self, response_text: str): - self._response_text = response_text - - async def acompletion(self, **kwargs): - from litellm.types.utils import Choices, Message, ModelResponse - - return ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content=self._response_text))]) - - -class _FakeProxyLogging: - def __init__(self, *, blocks: bool): - self._blocks = blocks - self.calls = [] - - async def pre_call_hook(self, *, user_api_key_dict, data, call_type): - self.calls.append((user_api_key_dict, data, call_type)) - if self._blocks: - raise HTTPException(status_code=429, detail="rate limited") - return data - - -# Regression: the worker call went straight to llm_router.acompletion, skipping every -# registered rate-limit/budget callback (they run as async_pre_call_hook, which only -# proxy_logging_obj.pre_call_hook walks). A caller already over budget or rate-limited could -# keep spending through this endpoint indefinitely. -class TestWorkerTextEnforcesRateLimitsAndBudget: - @pytest.mark.asyncio - async def test_calls_pre_call_hook_before_the_worker_model(self, monkeypatch): - fake_logging = _FakeProxyLogging(blocks=False) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_logging) - holder = UserAPIKeyAuth(api_key="fakehash1234567890") - text = await _worker_text( - _FakeRouter("the worker's answer"), - model="claude-haiku-4-5", - system_prompt="be precise", - message="what does this do", - user_api_key_dict=holder, - label="bulk_read", - ) - assert text == "the worker's answer" - assert len(fake_logging.calls) == 1 - called_key, _, called_type = fake_logging.calls[0] - assert called_key is holder - assert called_type == "acompletion" - - @pytest.mark.asyncio - async def test_a_blocked_pre_call_hook_prevents_the_worker_call(self, monkeypatch): - fake_logging = _FakeProxyLogging(blocks=True) - monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_logging) - holder = UserAPIKeyAuth(api_key="fakehash1234567890") - with pytest.raises(HTTPException) as exc_info: - await _worker_text( - _FakeRouter("should never be reached"), - model="claude-haiku-4-5", - system_prompt="be precise", - message="what does this do", - user_api_key_dict=holder, - label="bulk_read", - ) - assert exc_info.value.status_code == 429 diff --git a/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py b/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py new file mode 100644 index 00000000000..e529153e346 --- /dev/null +++ b/tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py @@ -0,0 +1,309 @@ +"""Unit tests for litellm.proxy.shunt_endpoints.endpoints.""" + +import io + +import pytest +from fastapi import HTTPException, Request, UploadFile + +import litellm.proxy.shunt_endpoints.endpoints as endpoints_mod +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth +from litellm.proxy.guardrails.shunt_capability_token import mint_shunt_capability_token +from litellm.proxy.shunt_endpoints.endpoints import _caller_from_capability_token, _worker_text + + +@pytest.fixture(autouse=True) +def _salt_key(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234-test-salt-key") + + +class TestMissingOrMalformedHeader: + @pytest.mark.asyncio + async def test_no_header_is_rejected(self): + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization=None) + assert exc_info.value.status_code == 401 + + @pytest.mark.asyncio + async def test_non_bearer_header_is_rejected(self): + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization="Basic abc123") + assert exc_info.value.status_code == 401 + + @pytest.mark.asyncio + async def test_malformed_token_is_rejected(self): + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization="Bearer not-a-real-token") + assert exc_info.value.status_code == 401 + + @pytest.mark.asyncio + async def test_expired_token_is_rejected(self, monkeypatch): + token = mint_shunt_capability_token(key_hash="deadbeef", master_key=None, now=1_000_000) + # 121s after mint, one second past the 120s TTL. + import litellm.proxy.guardrails.shunt_capability_token as token_mod + + monkeypatch.setattr(token_mod.time, "time", lambda: 1_000_121) + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization=f"Bearer {token}") + assert exc_info.value.status_code == 401 + + +class TestKeyHashGrant: + @pytest.mark.asyncio + async def test_resolves_the_key_object_for_the_grants_hash(self, monkeypatch): + token = mint_shunt_capability_token(key_hash="deadbeef", master_key=None) + resolved = UserAPIKeyAuth(api_key="deadbeef", team_id="team-1") + + async def _fake_get_key_object(**kwargs): + assert kwargs["hashed_token"] == "deadbeef" + return resolved + + import litellm.proxy.auth.auth_checks as auth_checks + + monkeypatch.setattr(auth_checks, "get_key_object", _fake_get_key_object) + result = await _caller_from_capability_token(authorization=f"Bearer {token}") + assert result is resolved + + @pytest.mark.asyncio + async def test_a_lookup_failure_is_rejected_not_propagated(self, monkeypatch): + token = mint_shunt_capability_token(key_hash="deadbeef", master_key=None) + + async def _raising_get_key_object(**kwargs): + raise Exception("key not found") + + import litellm.proxy.auth.auth_checks as auth_checks + + monkeypatch.setattr(auth_checks, "get_key_object", _raising_get_key_object) + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization=f"Bearer {token}") + assert exc_info.value.status_code == 401 + + +class TestMasterKeyGrant: + @pytest.mark.asyncio + async def test_matching_master_key_resolves_as_proxy_admin(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-the-real-master-key") + token = mint_shunt_capability_token(key_hash=None, master_key="sk-the-real-master-key") + result = await _caller_from_capability_token(authorization=f"Bearer {token}") + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + + # Regression: a resolved master-key caller carried the raw master key as its own api_key, + # which _worker_text later places in the outbound request's metadata["user_api_key"] -- + # reachable by any raw-metadata logging callback. Normal master-key auth substitutes a + # stable alias there specifically to keep the real key out of that sink; this must match. + @pytest.mark.asyncio + async def test_resolved_caller_never_carries_the_raw_master_key(self, monkeypatch): + from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-the-real-master-key") + token = mint_shunt_capability_token(key_hash=None, master_key="sk-the-real-master-key") + result = await _caller_from_capability_token(authorization=f"Bearer {token}") + assert result.api_key == LITELLM_PROXY_MASTER_KEY_ALIAS + assert result.api_key != "sk-the-real-master-key" + + @pytest.mark.asyncio + async def test_master_key_mismatch_is_rejected(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-the-current-master-key") + token = mint_shunt_capability_token(key_hash=None, master_key="sk-a-stale-master-key") + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization=f"Bearer {token}") + assert exc_info.value.status_code == 401 + + @pytest.mark.asyncio + async def test_no_configured_master_key_rejects_a_master_key_grant(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + token = mint_shunt_capability_token(key_hash=None, master_key="sk-anything") + with pytest.raises(HTTPException) as exc_info: + await _caller_from_capability_token(authorization=f"Bearer {token}") + assert exc_info.value.status_code == 401 + + +# _worker_text no longer calls llm_router.acompletion directly -- route_request is faked at +# module scope instead -- so this only needs to satisfy the type annotation, not do anything. +class _FakeRouter: + pass + + +class _FakeProxyLogging: + """post_call_success_hook is the one real dependency _worker_text calls on this object; + pre_call/rate-limit/budget enforcement now lives inside common_processing_pre_call_logic, + faked separately per test via _patch_pipeline.""" + + def __init__(self): + self.post_call_success_hook_calls = [] + + async def post_call_success_hook(self, *, data, user_api_key_dict, response): + self.post_call_success_hook_calls.append((data, user_api_key_dict, response)) + return response + + async def post_call_failure_hook(self, *, user_api_key_dict, original_exception, request_data): + return None + + +def _fake_request() -> Request: + return Request(scope={"type": "http", "headers": [], "method": "POST", "path": "/"}) + + +# Regression: the worker call went straight to llm_router.acompletion, skipping every +# registered rate-limit/budget callback and guardrail (they only run inside +# common_processing_pre_call_logic + route_request, the same pipeline /chat/completions uses). +# A caller already over budget or rate-limited could keep spending through this endpoint. +class TestWorkerTextGoesThroughTheSharedPipeline: + def _patch_pipeline(self, monkeypatch, *, fake_logging: _FakeProxyLogging, response_text: str): + async def _fake_pre_call_logic(self, **kwargs): + return self.data, object() + + async def _fake_route_request(**kwargs): + from litellm.types.utils import Choices, Message, ModelResponse + + return ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content=response_text))]) + + monkeypatch.setattr( + "litellm.proxy.shunt_endpoints.endpoints.ProxyBaseLLMRequestProcessing.common_processing_pre_call_logic", + _fake_pre_call_logic, + ) + monkeypatch.setattr("litellm.proxy.shunt_endpoints.endpoints.route_request", _fake_route_request) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", fake_logging) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_config", object()) + + @pytest.mark.asyncio + async def test_runs_the_pipeline_and_the_post_call_success_hook(self, monkeypatch): + fake_logging = _FakeProxyLogging() + self._patch_pipeline(monkeypatch, fake_logging=fake_logging, response_text="the worker's answer") + holder = UserAPIKeyAuth(api_key="fakehash1234567890") + text = await _worker_text( + _fake_request(), + _FakeRouter(), + model="claude-haiku-4-5", + system_prompt="be precise", + message="what does this do", + user_api_key_dict=holder, + label="bulk_read", + ) + assert text == "the worker's answer" + assert len(fake_logging.post_call_success_hook_calls) == 1 + + @pytest.mark.asyncio + async def test_a_blocked_pre_call_prevents_the_worker_call(self, monkeypatch): + fake_logging = _FakeProxyLogging() + self._patch_pipeline(monkeypatch, fake_logging=fake_logging, response_text="should never be reached") + + async def _raising_pre_call_logic(self, **kwargs): + raise HTTPException(status_code=429, detail="rate limited") + + monkeypatch.setattr( + "litellm.proxy.shunt_endpoints.endpoints.ProxyBaseLLMRequestProcessing.common_processing_pre_call_logic", + _raising_pre_call_logic, + ) + holder = UserAPIKeyAuth(api_key="fakehash1234567890") + # _worker_text lets a blocked pre-call raise through + # ProxyBaseLLMRequestProcessing._handle_llm_api_exception, the same conversion every + # other LLM route uses, so a raw HTTPException surfaces as the proxy-standard + # ProxyException rather than passing through unmodified. + with pytest.raises(ProxyException) as exc_info: + await _worker_text( + _fake_request(), + _FakeRouter(), + model="claude-haiku-4-5", + system_prompt="be precise", + message="what does this do", + user_api_key_dict=holder, + label="bulk_read", + ) + assert exc_info.value.code == "429" + assert len(fake_logging.post_call_success_hook_calls) == 0 + + +# Regression: _read_upload_text called upload.read() with no size, loading each whole file into +# memory and retaining every decoded file before building the prompt. An authenticated caller +# with a valid capability token could exhaust a proxy worker's memory, since the global +# request-size middleware is opt-in and premium-gated. +class TestUploadLimits: + def _upload(self, name: str, content: bytes) -> UploadFile: + return UploadFile(file=io.BytesIO(content), filename=name) + + @pytest.mark.asyncio + async def test_a_file_over_the_per_file_limit_is_rejected(self, monkeypatch): + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_PER_FILE", 10) + with pytest.raises(ProxyException) as exc_info: + await endpoints_mod._read_upload_texts([self._upload("big.py", b"x" * 50)]) + assert exc_info.value.code == "400" + + @pytest.mark.asyncio + async def test_files_under_the_limits_are_read_in_full(self, monkeypatch): + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_PER_FILE", 100) + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_TOTAL", 100) + texts = await endpoints_mod._read_upload_texts( + [self._upload("a.py", b"hello"), self._upload("b.py", b"world")] + ) + assert texts == ("hello", "world") + + @pytest.mark.asyncio + async def test_too_many_files_is_rejected_before_any_read(self, monkeypatch): + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_FILE_COUNT", 2) + with pytest.raises(ProxyException) as exc_info: + await endpoints_mod._read_upload_texts( + [self._upload(f"{i}.py", b"x") for i in range(3)] + ) + assert exc_info.value.code == "400" + + @pytest.mark.asyncio + async def test_files_individually_under_but_together_over_the_total_are_rejected(self, monkeypatch): + """The aggregate budget is what a many-small-files flood would otherwise slip past.""" + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_PER_FILE", 100) + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_TOTAL", 12) + with pytest.raises(ProxyException) as exc_info: + await endpoints_mod._read_upload_texts( + [self._upload("a.py", b"x" * 10), self._upload("b.py", b"y" * 10)] + ) + assert exc_info.value.code == "400" + + @pytest.mark.asyncio + async def test_an_exhausted_total_budget_rejects_the_next_file(self, monkeypatch): + """The running total in _read_upload_texts can only ever fall to exactly zero, never + below it (each read is already bounded by whatever was left when it started), so this + exercises that real, reachable boundary: two files exactly filling the total budget, + then a third that must be rejected outright rather than read at all.""" + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_PER_FILE", 100) + monkeypatch.setattr(endpoints_mod, "_MAX_UPLOAD_BYTES_TOTAL", 10) + third = self._upload("c.py", b"this file must never actually be read") + real_read = third.read + read_calls: list[int] = [] + + async def _tracking_read(size: int = -1): + read_calls.append(size) + return await real_read(size) + + third.read = _tracking_read # rebind-ok: test spy on this one instance + with pytest.raises(ProxyException) as exc_info: + await endpoints_mod._read_upload_texts( + [self._upload("a.py", b"aaaaaa"), self._upload("b.py", b"bbbb"), third] + ) + assert exc_info.value.code == "400" + assert read_calls == [], f"third file's read() was called with sizes {read_calls}, expected no call at all" + + @pytest.mark.asyncio + async def test_read_upload_text_itself_rejects_a_negative_budget_without_reading(self): + """_read_upload_texts's own loop can never pass a negative remaining_total_bytes (see + the test above), but _read_upload_text is called with a caller-supplied budget and + must reject one directly rather than pass it to UploadFile.read(), which treats a + negative size as "read the whole file" and would silently defeat this limit.""" + upload = self._upload("c.py", b"must never actually be read") + real_read = upload.read + read_calls: list[int] = [] + + async def _tracking_read(size: int = -1): + read_calls.append(size) + return await real_read(size) + + upload.read = _tracking_read # rebind-ok: test spy on this one instance + with pytest.raises(ProxyException) as exc_info: + await endpoints_mod._read_upload_text(upload, remaining_total_bytes=-5) + assert exc_info.value.code == "400" + assert read_calls == [], f"read() was called with sizes {read_calls}, expected no call at all" + + @pytest.mark.asyncio + async def test_non_utf8_content_is_rejected(self): + with pytest.raises(ProxyException) as exc_info: + await endpoints_mod._read_upload_texts([self._upload("bin.dat", b"\xff\xfe\x00binary")]) + assert exc_info.value.code == "400"