scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(auto_router): route shunt worker calls through the shared request pipeline

Browse source 
The worker call went straight to llm_router.acompletion, so it skipped
everything common_processing_pre_call_logic does for /chat/completions:
registered guardrails, rate limiters, and budget checks. A caller already
over budget or rate limited could keep spending through these endpoints.
Route through ProxyBaseLLMRequestProcessing + route_request instead, and
let _handle_llm_api_exception map failures, so a rejected request also
releases its rate-limit reservation rather than leaking it.

Also bound upload reads. bulk_read and code_write read each UploadFile
without a size limit, so a large multipart body was buffered whole. Add
per-file, aggregate, and file-count budgets, with each read capped by
what the aggregate budget has left.

_read_upload_text rejects a non-positive remaining budget before calling
.read() rather than letting min() produce a non-positive limit:
UploadFile.read treats a negative size as "read the whole file", which
would silently defeat the budget for any caller that passes one.

Rename test_endpoints.py to test_shunt_worker_endpoints.py: pytest
imports test modules by basename with no __init__.py present, so it
collided with credential_endpoints/test_endpoints.py.
This commit is contained in:
moe-berri 2026-09-07 17:38:38 -07:00
parent b4b4f58a76
commit dac683279f
3 changed files with 425 additions and 214 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

153
litellm/proxy/shunt_endpoints/endpoints.py
View file

@ -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,

177
tests/test_litellm/proxy/shunt_endpoints/test_endpoints.py
View file

@ -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

309
tests/test_litellm/proxy/shunt_endpoints/test_shunt_worker_endpoints.py Normal file
View file

@ -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"