litellm/tests/unit/proxy/hooks/test_batch_rate_limiter.py
mrinal-berri f516668f96
fix(managed_files): enforce input file ownership on batch and fine-tuning create (#45855)
* test(managed_files): pin ownership checks on batch and fine-tuning create

Tenant B can create a batch or fine-tuning job from tenant A's managed
file id today. These fail on main for acreate_batch input_file_id and
acreate_fine_tuning_job training_file and validation_file

* fix(managed_files): check input file ownership on batch and fine-tuning create

Batch create and fine-tuning create only mapped the managed input file to
provider ids and never checked who owned it, so any key with access to the
model could build a batch or job on another tenant's file and then own the
outputs. Both now run check_file_ids_access like chat and responses do,
which 403s a non-owner and 404s a managed id with no ownership row.
Fine-tuning also checks validation_file, which was never looked at

* test(batches): pin that a pre-call ownership denial stops every batch create branch

LIT-4797 noted an endpoint-level check missed the load-balanced branch
because it dispatches first. The ownership check lives in the pre-call
hook, which runs before branch selection, so a 403 there must keep both
the load-balanced and unified-file branches from reaching the router

* test(integration): pin managed-file ownership on batch and fine-tuning create

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(managed_files): check ownership before the batch rate limiter reads the input file and keep key-owned resources across key regenerate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-12 13:58:25 -04:00

573 lines
24 KiB
Python

"""
Tests for `tpd_limit` (tokens per day) enforcement on batch submissions.
A batch's rows are scheduled by the provider, so a caller cannot keep a large
batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit`
are charged against a 24h token window instead of their minute counters.
"""
import json
import time
from collections.abc import Iterator, Sequence
from datetime import datetime, timezone
from typing import Final, Literal
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
import respx
from fastapi import HTTPException
from typing_extensions import ReadOnly, TypedDict
import litellm
from litellm import DualCache
from litellm.constants import BATCH_TPD_WINDOW_SECONDS
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache, hash_token
from litellm.types.llms.openai import ChatCompletionUserMessage, LiteLLMBatchCreateRequest
class _Clock:
def __init__(self, start: datetime):
self.now = start
def __call__(self) -> datetime:
return self.now
def _make_limiters(clock: _Clock | None = None):
internal_usage_cache = InternalUsageCache(dual_cache=DualCache())
rate_limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache, time_provider=clock)
batch_limiter = rate_limiter._get_batch_rate_limiter()
assert batch_limiter is not None
return internal_usage_cache, rate_limiter, batch_limiter
async def _counter(internal_usage_cache, rate_limiter, descriptor_key, value, rate_limit_type):
cache_key = rate_limiter.create_rate_limit_keys(descriptor_key, value, rate_limit_type)
raw = await internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None, local_only=True)
return int(raw or 0)
@pytest.mark.asyncio
async def test_batch_over_rpm_and_tpm_but_under_tpd_is_accepted():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
api_key = hash_token("tpd-key")
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpm_limit=10, tpd_limit=1000)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=500, request_count=50),
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 500
assert await _counter(internal_usage_cache, rate_limiter, "api_key", api_key, "requests") == 0
assert await _counter(internal_usage_cache, rate_limiter, "api_key", api_key, "tokens") == 0
@pytest.mark.asyncio
async def test_cumulative_batch_tokens_over_tpd_returns_429_with_remaining_daily_window():
window_start = datetime(2026, 9, 13, 8, 0, 0)
clock = _Clock(window_start)
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock)
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("tpd-key-2"), rpm_limit=1, tpd_limit=1000)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
clock.now = datetime(2026, 9, 13, 11, 0, 0)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
assert exc.value.status_code == 429
assert "api_key_tpd" in str(exc.value.detail)
assert "600 tokens but only 400 tokens remaining out of 1000 TPD limit" in str(exc.value.detail)
assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600)
assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC"
@pytest.mark.asyncio
async def test_failed_batch_submission_refunds_tpd_tokens():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
api_key = hash_token("tpd-refund-key")
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpd_limit=1000)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
await rate_limiter.async_post_call_failure_hook(
request_data={},
original_exception=RuntimeError("provider rejected the file"),
user_api_key_dict=user_api_key_dict,
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 0
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=1000, request_count=10),
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 1000
@pytest.mark.asyncio
async def test_tpd_refund_applies_once_and_only_to_daily_counters():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
team_key = UserAPIKeyAuth(
api_key=hash_token("tpd-refund-team-key"),
rpm_limit=100,
tpm_limit=10_000,
team_id="team-r",
team_tpd_limit=5000,
)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=team_key,
data={},
batch_usage=BatchFileUsage(total_tokens=800, request_count=8),
)
await rate_limiter.async_post_call_failure_hook(
request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key
)
await rate_limiter.async_post_call_failure_hook(
request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key
)
assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-r", "tokens") == 0
assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "tokens") == 800
assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "requests") == 8
@pytest.mark.asyncio
async def test_rejected_batch_leaves_nothing_to_refund():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
api_key = hash_token("tpd-rejected-key")
user_api_key_dict = UserAPIKeyAuth(api_key=api_key, tpd_limit=100)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9)
)
with pytest.raises(HTTPException):
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2)
)
await rate_limiter.async_post_call_failure_hook(
request_data={}, original_exception=RuntimeError("429 bubbled up"), user_api_key_dict=user_api_key_dict
)
assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 90
@pytest.mark.asyncio
async def test_batch_without_tpd_still_enforces_minute_rpm():
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters()
user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("rpm-only-key"), rpm_limit=1, tpm_limit=1000)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=50, request_count=5),
)
assert exc.value.status_code == 429
assert "RPM limit" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_team_tpd_replaces_team_minute_limits_but_key_minute_limits_still_apply():
internal_usage_cache, rate_limiter, batch_limiter = _make_limiters()
team_key = UserAPIKeyAuth(
api_key=hash_token("team-key"),
team_id="team-1",
team_rpm_limit=1,
team_tpm_limit=10,
team_tpd_limit=5000,
)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=team_key,
data={},
batch_usage=BatchFileUsage(total_tokens=800, request_count=8),
)
assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-1", "tokens") == 800
assert await _counter(internal_usage_cache, rate_limiter, "team", "team-1", "requests") == 0
key_rpm_in_team_with_tpd = UserAPIKeyAuth(
api_key=hash_token("team-key-2"),
rpm_limit=1,
team_id="team-1",
team_tpd_limit=5000,
)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=key_rpm_in_team_with_tpd,
data={},
batch_usage=BatchFileUsage(total_tokens=10, request_count=2),
)
assert exc.value.status_code == 429
assert "api_key:" in str(exc.value.detail)
assert "RPM limit" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_end_user_tpd_is_enforced_per_end_user():
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters()
first_customer = UserAPIKeyAuth(
api_key=hash_token("shared-key"), end_user_id="customer-a", end_user_rpm_limit=1, end_user_tpd_limit=100
)
second_customer = UserAPIKeyAuth(
api_key=hash_token("shared-key"), end_user_id="customer-b", end_user_rpm_limit=1, end_user_tpd_limit=100
)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=first_customer, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9)
)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=second_customer, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9)
)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=first_customer, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2)
)
assert exc.value.status_code == 429
assert "end_user_tpd: customer-a" in str(exc.value.detail)
def test_tpd_only_key_is_not_skipped_as_having_no_limits():
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters()
descriptors = batch_limiter._create_batch_rate_limit_descriptors(
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("tpd-only"), tpd_limit=100),
data={},
)
assert batch_limiter._has_applicable_batch_rate_limits(descriptors) is True
def test_online_descriptors_ignore_tpd_limit():
_internal_usage_cache, rate_limiter, _batch_limiter = _make_limiters()
api_key = hash_token("online-key")
descriptors = rate_limiter.create_rate_limit_descriptors(
user_api_key_dict=UserAPIKeyAuth(api_key=api_key, rpm_limit=5, tpd_limit=100, team_id="t", team_tpd_limit=9),
data={"model": "gpt-4o"},
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
)
assert [(d["key"], d["rate_limit"]["window_size"]) for d in descriptors] == [("api_key", rate_limiter.window_size)]
@pytest.fixture(params=["Europe/Paris", "Asia/Kolkata", "America/Los_Angeles"])
def process_timezone(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> Iterator[str]:
monkeypatch.setenv("TZ", request.param)
time.tzset()
yield request.param
monkeypatch.undo()
time.tzset()
@pytest.mark.skipif(not hasattr(time, "tzset"), reason="switching the process timezone needs time.tzset()")
@pytest.mark.asyncio
async def test_batch_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_proxy(process_timezone: str) -> None:
window_start: Final = datetime(2026, 9, 13, 8, 0, 0, tzinfo=timezone.utc)
clock: Final = _Clock(window_start)
_internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock)
user_api_key_dict: Final = UserAPIKeyAuth(api_key=hash_token("tpd-key-utc"), rpm_limit=1, tpd_limit=1000)
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
clock.now = datetime(2026, 9, 13, 11, 0, 0, tzinfo=timezone.utc)
with pytest.raises(HTTPException) as exc:
await batch_limiter._check_and_increment_batch_counters(
user_api_key_dict=user_api_key_dict,
data={},
batch_usage=BatchFileUsage(total_tokens=600, request_count=6),
)
assert exc.value.status_code == 429
assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600)
assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC"
assert str(exc.value.detail).endswith("Limit resets at: 2026-09-14 08:00:00 UTC")
_BATCH_MODEL: Final = "gpt-3.5-turbo"
_MANAGED_FILE_ID: Final = (
"bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9vY3RldC1zdHJlYW07dW5pZmllZF9pZCxyZWdyZXNzaW9uLXRlc3QtZmlsZQ=="
)
class _BatchBody(TypedDict):
model: ReadOnly[str]
messages: ReadOnly[Sequence[ChatCompletionUserMessage]]
class _BatchLine(TypedDict):
custom_id: ReadOnly[str]
method: ReadOnly[Literal["POST"]]
url: ReadOnly[Literal["/v1/chat/completions"]]
body: ReadOnly[_BatchBody]
@pytest.fixture
def openai_files(monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> respx.MockRouter:
monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
return respx_mock
def _batch_rows(messages: Sequence[str]) -> tuple[_BatchLine, ...]:
return tuple(
_BatchLine(
custom_id=f"request-{i}",
method="POST",
url="/v1/chat/completions",
body=_BatchBody(model=_BATCH_MODEL, messages=(ChatCompletionUserMessage(role="user", content=message),)),
)
for i, message in enumerate(messages, start=1)
)
def _serve_file(router: respx.MockRouter, file_id: str, rows: Sequence[_BatchLine]) -> respx.Route:
jsonl: Final = "\n".join(json.dumps(row) for row in rows)
return router.get(f"https://api.openai.com/v1/files/{file_id}/content").mock(
return_value=httpx.Response(200, content=jsonl.encode())
)
def _token_counter_total(rows: Sequence[_BatchLine]) -> int:
return sum(litellm.token_counter(model=row["body"]["model"], messages=row["body"]["messages"]) for row in rows)
def _create_batch_data(input_file_id: str) -> LiteLLMBatchCreateRequest:
return LiteLLMBatchCreateRequest(model=_BATCH_MODEL, input_file_id=input_file_id)
@pytest.mark.asyncio
async def test_count_input_file_usage_matches_token_counter(openai_files: respx.MockRouter):
_, _, batch_limiter = _make_limiters()
rows: Final = _batch_rows(("Hello", "Hi there", "Hey"))
content_route: Final = _serve_file(openai_files, "file-abc123", rows)
usage: Final = await batch_limiter.count_input_file_usage(file_id="file-abc123", custom_llm_provider="openai")
assert content_route.call_count == 1
assert usage.request_count == 3
assert usage.total_tokens == _token_counter_total(rows)
@pytest.mark.asyncio
async def test_batch_rate_limit_single_file_under_and_over_tpm(openai_files: respx.MockRouter):
small_rows: Final = _batch_rows(("Hello", "Hi", "Hey"))
big_rows: Final = _batch_rows(
("This is a longer message that will consume more tokens from the rate limit. " * 100,) * 3
)
_serve_file(openai_files, "file-small", small_rows)
_serve_file(openai_files, "file-big", big_rows)
user_api_key_dict: Final = UserAPIKeyAuth(api_key="test-key-123", tpm_limit=200, rpm_limit=10)
_, _, small_limiter = _make_limiters()
data_small: Final = dict(_create_batch_data("file-small"))
result: Final = await small_limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data_small,
call_type="acreate_batch",
)
assert result is data_small
assert data_small["_batch_token_count"] == _token_counter_total(small_rows)
assert data_small["_batch_request_count"] == 3
_, _, big_limiter = _make_limiters()
with pytest.raises(HTTPException) as exc_info:
await big_limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=dict(_create_batch_data("file-big")),
call_type="acreate_batch",
)
assert exc_info.value.status_code == 429
assert "tokens" in str(exc_info.value.detail).lower()
@pytest.mark.asyncio
async def test_batch_rate_limit_cumulative_tpm_rejects_second_request(openai_files: respx.MockRouter):
_, _, batch_limiter = _make_limiters()
user_api_key_dict: Final = UserAPIKeyAuth(api_key="test-key-456", tpm_limit=200, rpm_limit=10)
first_rows: Final = _batch_rows(("This message has some content to reach about 100 tokens total. " * 4,) * 2)
second_rows: Final = _batch_rows(
("This is another message with more content to exceed the remaining limit. " * 11,) * 2
)
_serve_file(openai_files, "file-1", first_rows)
_serve_file(openai_files, "file-2", second_rows)
first_tokens: Final = _token_counter_total(first_rows)
assert first_tokens <= 200 < first_tokens + _token_counter_total(second_rows)
first_data: Final = dict(_create_batch_data("file-1"))
first_result: Final = await batch_limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=first_data,
call_type="acreate_batch",
)
assert first_result is first_data
assert first_data["_batch_token_count"] == first_tokens
with pytest.raises(HTTPException) as exc_info:
await batch_limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=dict(_create_batch_data("file-2")),
call_type="acreate_batch",
)
assert exc_info.value.status_code == 429
assert "tokens" in str(exc_info.value.detail).lower()
@pytest.mark.asyncio
async def test_batch_rate_limiter_reads_a_provider_file_with_user_context(openai_files: respx.MockRouter):
_, _, batch_limiter = _make_limiters()
user_api_key_dict: Final = UserAPIKeyAuth(
api_key="test-key-managed-files", user_id="test-user-abc123", tpm_limit=500, rpm_limit=10
)
rows: Final = _batch_rows(("This is a test message for batch rate limiting with managed files. " * 5,) * 3)
content_route: Final = _serve_file(openai_files, "file-abc123", rows)
data: Final = dict(_create_batch_data("file-abc123"))
result: Final = await batch_limiter.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=DualCache(),
data=data,
call_type="acreate_batch",
)
assert content_route.call_count == 1
assert result is data
assert data["_batch_token_count"] == _token_counter_total(rows)
assert data["_batch_request_count"] == 3
@pytest.mark.asyncio
async def test_batch_rate_limiter_without_user_context(openai_files: respx.MockRouter):
_, _, batch_limiter = _make_limiters()
rows: Final = _batch_rows(("Hello",))
content_route: Final = _serve_file(openai_files, "file-abc123", rows)
usage_without_context: Final = await batch_limiter.count_input_file_usage(
file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=None
)
usage_with_context: Final = await batch_limiter.count_input_file_usage(
file_id="file-abc123",
custom_llm_provider="openai",
user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user-123"),
)
assert content_route.call_count == 2
assert usage_without_context.request_count == usage_with_context.request_count == 1
assert usage_without_context.total_tokens == usage_with_context.total_tokens == _token_counter_total(rows)
async def _install_managed_file_hook(
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rows: str
) -> tuple[respx.Route, respx.Route]:
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
from litellm import Router
from litellm.models.managed_files import LiteLLM_ManagedFileTable
from litellm.proxy import proxy_server
from litellm.proxy.openai_files_endpoints.common_utils import is_base64_encoded_unified_file_id
from litellm.proxy.utils import ProxyLogging
assert is_base64_encoded_unified_file_id(_MANAGED_FILE_ID)
provider_route: Final = _serve_file(openai_files, "file-provider-1", rows)
unrouted_managed_read: Final = _serve_file(openai_files, _MANAGED_FILE_ID, rows)
row: Final = LiteLLM_ManagedFileTable(
unified_file_id=_MANAGED_FILE_ID,
model_mappings={"deployment-1": "file-provider-1"},
flat_model_file_ids=["file-provider-1"],
created_by="test-user-regression",
team_id="team-owner",
)
file_cache: Final = InternalUsageCache(dual_cache=DualCache())
await file_cache.async_set_cache(key=_MANAGED_FILE_ID, value=row.model_dump(), litellm_parent_otel_span=None)
prisma_client: Final = MagicMock()
prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(return_value=row)
proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.proxy_hook_mapping["managed_files"] = _PROXY_LiteLLMManagedFiles(
internal_usage_cache=file_cache, prisma_client=prisma_client
)
router: Final = Router(
model_list=[
{
"model_name": _BATCH_MODEL,
"litellm_params": {"model": f"openai/{_BATCH_MODEL}", "api_key": "sk-test"},
"model_info": {"id": "deployment-1"},
}
]
)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
monkeypatch.setattr(proxy_server, "llm_router", router)
return provider_route, unrouted_managed_read
@pytest.mark.asyncio
async def test_managed_file_is_read_through_the_managed_files_hook_with_user_context(
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
rows: Final = _batch_rows(("Test message for regression",))
provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows)
standard_route: Final = _serve_file(openai_files, "file-abc123", rows)
_, _, batch_limiter = _make_limiters()
user_api_key_dict: Final = UserAPIKeyAuth(
api_key="test-key-regression", user_id="test-user-regression", tpm_limit=1000, rpm_limit=10
)
managed_usage: Final = await batch_limiter.count_input_file_usage(
file_id=_MANAGED_FILE_ID, custom_llm_provider="openai", user_api_key_dict=user_api_key_dict
)
standard_usage: Final = await batch_limiter.count_input_file_usage(
file_id="file-abc123", custom_llm_provider="openai", user_api_key_dict=user_api_key_dict
)
assert provider_route.call_count == 1
assert standard_route.call_count == 1
assert not unrouted_managed_read.called
assert managed_usage.request_count == standard_usage.request_count == 1
assert managed_usage.total_tokens == standard_usage.total_tokens == _token_counter_total(rows)
@pytest.mark.asyncio
async def test_managed_file_of_another_tenant_is_denied_before_it_is_read(
openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
rows: Final = _batch_rows(("Test message for regression",))
provider_route, unrouted_managed_read = await _install_managed_file_hook(openai_files, monkeypatch, rows)
_, _, batch_limiter = _make_limiters()
outsider: Final = UserAPIKeyAuth(api_key="test-key-outsider", user_id="outsider", team_id="team-other", tpm_limit=1)
with pytest.raises(HTTPException) as denied:
await batch_limiter.async_pre_call_hook(
user_api_key_dict=outsider,
cache=DualCache(),
data=dict(_create_batch_data(_MANAGED_FILE_ID)),
call_type="acreate_batch",
)
assert denied.value.status_code == 403
assert not provider_route.called
assert not unrouted_managed_read.called