fix(proxy): count keyless JWT callers per user, require positive file caps, and take the upload slot after request validation

This commit is contained in:
mateo-berri 2026-09-28 17:46:33 -07:00
parent 3b5f1bb122
commit cc3316aa32
7 changed files with 71 additions and 18 deletions

View file

@ -2784,15 +2784,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
)
max_batch_file_records: int | None = Field(
None,
description="max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key or team can carry its own value in metadata, set by a proxy admin, and the lowest applicable value wins. Unset means no limit",
gt=0,
description="max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key's metadata can override it and a team's metadata adds a team cap on top, both set by a proxy admin; the lower of the key's value and the team's value wins. Unset means no limit",
)
max_batch_file_uploads_per_day: int | None = Field(
None,
description="max /v1/files uploads with purpose=batch per key per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
gt=0,
description="max /v1/files uploads with purpose=batch per key (per user for JWT callers) per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
)
max_file_downloads_per_minute: int | None = Field(
None,
description="max GET /v1/files/{file_id}/content calls per key per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
gt=0,
description="max GET /v1/files/{file_id}/content calls per key (per user for JWT callers) per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit",
)
max_file_size_mb: int | None = Field(
None,

View file

@ -25,7 +25,7 @@ FileUsageSetting: TypeAlias = Literal[
"max_file_downloads_per_minute",
]
LimitSource: TypeAlias = Literal["key", "team", "general_settings"]
CounterScope: TypeAlias = Literal["key", "team"]
CounterScope: TypeAlias = Literal["key", "user", "team"]
_COUNTER_PREFIX: Final = "litellm:file_usage"
_DAY_SECONDS: Final = 24 * 60 * 60
@ -65,9 +65,8 @@ def _read_limit(
return FileUsageLimit(setting=setting, value=_LIMIT_ADAPTER.validate_python(raw), source=source)
except ValidationError:
verbose_proxy_logger.warning(
"Ignoring invalid %s value %r in %s; expected a positive integer",
"Ignoring invalid %s in %s; expected a positive integer",
setting,
raw,
source,
)
return None
@ -103,17 +102,24 @@ def batch_file_record_limit(
return min(applicable, key=lambda limit: limit.value, default=None)
def _caller_counter(user_api_key_dict: UserAPIKeyAuth, limit: FileUsageLimit | None) -> ScopedFileUsageLimit | None:
if limit is None:
return None
if user_api_key_dict.api_key:
return ScopedFileUsageLimit(scope="key", scope_id=user_api_key_dict.api_key, limit=limit)
if user_api_key_dict.user_id:
return ScopedFileUsageLimit(scope="user", scope_id=user_api_key_dict.user_id, limit=limit)
return None
def resolve_scoped_limits(
user_api_key_dict: UserAPIKeyAuth,
general_settings: Mapping[str, object],
setting: FileUsageSetting,
) -> tuple[ScopedFileUsageLimit, ...]:
key_limit: Final = _key_limit(user_api_key_dict, general_settings, setting)
team_limit: Final = _team_limit(user_api_key_dict, setting)
candidates: Final = (
ScopedFileUsageLimit(scope="key", scope_id=user_api_key_dict.api_key, limit=key_limit)
if key_limit is not None and user_api_key_dict.api_key
else None,
_caller_counter(user_api_key_dict, _key_limit(user_api_key_dict, general_settings, setting)),
ScopedFileUsageLimit(scope="team", scope_id=user_api_key_dict.team_id, limit=team_limit)
if team_limit is not None and user_api_key_dict.team_id
else None,
@ -172,17 +178,19 @@ def describe_limit_source(source: LimitSource) -> str:
case "general_settings":
return "in general_settings"
case _:
assert_never(source)
return assert_never(source)
def _describe_scope(scoped: ScopedFileUsageLimit) -> str:
match scoped.scope:
case "key":
return "this key"
case "user":
return f"user {scoped.scope_id}"
case "team":
return f"team {scoped.scope_id}"
case _:
assert_never(scoped.scope)
return assert_never(scoped.scope)
def _raise_limit_exceeded(exceeded: FileUsageLimitExceeded, what_ran_out: str, when_it_resets: str) -> NoReturn:

View file

@ -681,9 +681,6 @@ async def create_file(
)
if batch_file_failure is not None:
raise_batch_file_validation_failure(batch_file_failure)
await enforce_batch_file_upload_limit(
proxy_logging_obj.file_usage_cache, user_api_key_dict, general_settings
)
data = {"passthrough": True} if passthrough else {}
@ -757,6 +754,11 @@ async def create_file(
seconds=expires_after_seconds,
)
if purpose == "batch":
await enforce_batch_file_upload_limit(
proxy_logging_obj.file_usage_cache, user_api_key_dict, general_settings
)
# Include original request and headers in the data
data = await add_litellm_data_to_request(
data=data,

View file

@ -4291,12 +4291,19 @@ def test_create_file_batch_uploads_over_daily_limit_get_429_and_only_valid_files
try:
invalid = _upload_batch(b"not json\n")
malformed_expiry = client.post(
"/v1/files",
files={"file": ("batch.jsonl", VALID_BATCH_LINE, "application/jsonl")},
data={"purpose": "batch", "expires_after[anchor]": "created_at"},
headers={"Authorization": "Bearer test-key"},
)
allowed = [_upload_batch(VALID_BATCH_LINE) for _ in range(2)]
rejected = _upload_batch(VALID_BATCH_LINE)
finally:
_teardown_batch_upload_endpoint()
assert invalid.status_code == 400, invalid.text
assert malformed_expiry.status_code == 400, malformed_expiry.text
assert [response.status_code for response in allowed] == [200, 200]
assert rejected.status_code == 429, rejected.text
assert rejected.headers["retry-after"] == str(86400 - 3615)

View file

@ -395,3 +395,15 @@ def test_mcp_metadata_rejects_unavailable_upstream_protocol(revision):
for model in (NewMCPServerRequest, UpdateMCPServerRequest):
with pytest.raises(ValidationError):
model.model_validate(payload)
@pytest.mark.parametrize(
"setting", ["max_batch_file_records", "max_batch_file_uploads_per_day", "max_file_downloads_per_minute"]
)
def test_batch_file_caps_accept_only_positive_limits(setting):
from litellm.proxy._types import ConfigGeneralSettings
for invalid in (0, -5):
with pytest.raises(ValidationError):
ConfigGeneralSettings.model_validate({setting: invalid})
assert getattr(ConfigGeneralSettings.model_validate({setting: 3}), setting) == 3

View file

@ -85,6 +85,12 @@ def test_record_limit_is_the_lowest_applicable_limit(key_metadata, team_metadata
{},
(),
),
(
UserAPIKeyAuth(api_key=None, user_id="u1", team_id="t1"),
{"max_batch_file_uploads_per_day": 5},
(_scoped("user", "u1", 5, "general_settings"),),
),
(UserAPIKeyAuth(api_key=None), {"max_batch_file_uploads_per_day": 5}, ()),
],
)
def test_counter_scopes_pair_each_limit_with_its_own_counter(caller, general_settings, expected):
@ -180,3 +186,18 @@ async def test_no_configured_limit_never_rejects():
for _ in range(50):
await enforce_batch_file_upload_limit(cache, caller, {}, clock=lambda: MIDDAY)
await enforce_file_download_limit(cache, caller, {}, "file-a", clock=lambda: MIDDAY)
async def test_a_caller_without_a_key_is_counted_per_user():
cache: Final = _cache()
general_settings: Final = {"max_file_downloads_per_minute": 1}
jwt_caller: Final = UserAPIKeyAuth(api_key=None, user_id="u1")
await enforce_file_download_limit(cache, jwt_caller, general_settings, "file-a", clock=lambda: MIDDAY)
await enforce_file_download_limit(
cache, UserAPIKeyAuth(api_key=None, user_id="u2"), general_settings, "file-a", clock=lambda: MIDDAY
)
with pytest.raises(ProxyException) as exc:
await enforce_file_download_limit(cache, jwt_caller, general_settings, "file-a", clock=lambda: MIDDAY)
assert "max_file_downloads_per_minute is 1 for user u1 (set in general_settings)" in exc.value.message

View file

@ -28205,7 +28205,7 @@ export interface components {
master_key?: string | null;
/**
* Max Batch File Records
* @description max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key or team can carry its own value in metadata, set by a proxy admin, and the lowest applicable value wins. Unset means no limit
* @description max records (non-blank lines) per batch input file for /v1/files uploads with purpose=batch, applied per key. A key's metadata can override it and a team's metadata adds a team cap on top, both set by a proxy admin; the lower of the key's value and the team's value wins. Unset means no limit
*/
max_batch_file_records?: number | null;
/**
@ -28215,7 +28215,7 @@ export interface components {
max_batch_file_size_mb?: number | null;
/**
* Max Batch File Uploads Per Day
* @description max /v1/files uploads with purpose=batch per key per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
* @description max /v1/files uploads with purpose=batch per key (per user for JWT callers) per UTC day. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
*/
max_batch_file_uploads_per_day?: number | null;
/**
@ -28232,7 +28232,7 @@ export interface components {
} | null;
/**
* Max File Downloads Per Minute
* @description max GET /v1/files/{file_id}/content calls per key per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
* @description max GET /v1/files/{file_id}/content calls per key (per user for JWT callers) per file per minute. A key's metadata can override it and a team's metadata adds a shared team cap, both set by a proxy admin. Unset means no limit
*/
max_file_downloads_per_minute?: number | null;
/**