test(integration): cover files_batches endpoint contracts (#44703)

* test(integration): cover files_batches endpoint contracts

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

* test(integration): split raw cloud-storage delete test by caller

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

* test(integration): drop mutable polling state in content download test

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

* test(integration): dispatch request spellings without rebinding

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

* test(integration): import assert_never from typing_extensions for python 3.10

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

* test(integration): tighten files_batches wire and access assertions

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

* test(integration): tidy files_batches access tests

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

* test(integration): cover repeated target models, upload expiry, batch aliases

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

* test(integration): prove files streaming path, SDK model denial and admin GCS delete

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

* test(integration): give the streaming head the wire barrier's full budget

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

* test: strengthen files and batches wire coverage

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

* test: drop duplicate parametrize from the model-routed upload helper

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

* test: pass the SDK base url in the skipped repeated-plain upload test

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

* test(integration): allow batch output metadata lookup

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

* test: avoid draining file content requests before wait

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

---------

Co-authored-by: mateo <mateo@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-14 06:57:41 +00:00 • committed by GitHub
parent e226016aee
commit 0ca9f8185c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 1733 additions and 0 deletions

View file

@ -0,0 +1,882 @@
import base64
import json
import re
import threading
import uuid
from collections import deque
from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import ThreadPoolExecutor
from email.message import Message
from email.parser import BytesParser
from email.policy import HTTP
from functools import partial
from queue import Queue
from typing import Final, Literal, cast
from urllib.parse import urlsplit
import httpx
import pytest
from integration._support.client import Gateway, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.wire import Reply, Request, Wire, wire_server
from openai import OpenAI
from pydantic import BaseModel, JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
INPUT_FILENAME: Final = "in.jsonl"
OUTPUT_BYTES: Final = b'{"custom_id": "r1", "result": {"response": {"body": "ok"}}}\n'
STREAM_CHUNK_BYTES: Final = 1024 * 1024
_HEAD_OPEN: Final = b'{"custom_id": "r0", "padding": "'
_HEAD_CLOSE: Final = b'"}\n'
OUTPUT_HEAD: Final = _HEAD_OPEN + b"h" * (STREAM_CHUNK_BYTES - len(_HEAD_OPEN) - len(_HEAD_CLOSE)) + _HEAD_CLOSE
MANAGED_PREFIX: Final = "litellm_proxy:"
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
MANAGED_FILE_ROW: Final = (
'SELECT model_mappings, flat_model_file_ids FROM "LiteLLM_ManagedFileTable" WHERE unified_file_id = %s'
)
_BatchUploadSpelling = Literal["sdk_list", "repeated_bracket", "indexed", "comma_joined", "repeated_plain"]
_ModelRoutedSpelling = Literal[
"sdk_extra_body",
"sdk_proxy_fields",
"multipart_field",
"query_param",
"model_header",
"multipart_proxy_fields",
"sdk_expires_after",
"multipart_expires_after",
]
_ContentRoute = Literal["v1_encoded", "files_encoded", "files_raw_query", "v1_raw_header"]
_BatchCreateSpelling = Literal["encoded_input", "body_model", "query_model", "header_model", "sanitized_metadata"]
def _jsonl(model: str) -> bytes:
line: Final = {
"custom_id": "r1",
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": model, "messages": [{"role": "user", "content": "one"}]},
}
return json.dumps(line, separators=(",", ":")).encode() + b"\n"
def _rewritten_jsonl(model: str) -> bytes:
line: Final = {
"custom_id": "r1",
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": model, "messages": [{"role": "user", "content": "one"}]},
}
return json.dumps(line).encode()
class _FileObject(BaseModel):
id: str
object: str
bytes: int
created_at: int
filename: str
purpose: str
class _Batch(BaseModel):
id: str
object: str
input_file_id: str | None = None
output_file_id: str | None = None
error_file_id: str | None = None
status: str
def _json(response: httpx.Response) -> dict[str, JsonValue]:
return JSON_OBJECT.validate_json(response.content)
def _json_reply(body: Mapping[str, JsonValue], status: int = 200) -> Reply:
return Reply(status=status, body=json.dumps(body).encode())
def _file_object(file_id: str, byte_count: int, filename: str = INPUT_FILENAME) -> dict[str, JsonValue]:
return {
"id": file_id,
"object": "file",
"bytes": byte_count,
"created_at": 1700000000,
"filename": filename,
"purpose": "batch",
"status": "processed",
}
def _expected_file_object(file_id: str, byte_count: int, filename: str = INPUT_FILENAME) -> _FileObject:
return _FileObject(
id=file_id,
object="file",
bytes=byte_count,
created_at=1700000000,
filename=filename,
purpose="batch",
)
def _batch_object(
batch_id: str, input_file_id: str, status: str, output_file_id: str | None = None, error_file_id: str | None = None
) -> dict[str, JsonValue]:
return {
"id": batch_id,
"object": "batch",
"endpoint": "/v1/chat/completions",
"input_file_id": input_file_id,
"output_file_id": output_file_id,
"error_file_id": error_file_id,
"status": status,
"created_at": 1700000000,
"completion_window": "24h",
}
def _multipart_parts(request: Request) -> tuple[Message, ...]:
envelope: Final = f"content-type: {request.headers['content-type']}\r\n\r\n".encode() + request.body
parsed: Final = BytesParser(policy=HTTP).parsebytes(envelope)
assert parsed.is_multipart(), request.headers["content-type"]
return tuple(parsed.iter_parts())
def _text_fields(parts: tuple[Message, ...]) -> dict[str, str]:
return {
part.get_param("name", header="content-disposition"): part.get_payload(decode=True).decode()
for part in parts
if part.get_filename() is None
}
def _file_fields(parts: tuple[Message, ...]) -> dict[str, tuple[str, str, bytes]]:
return {
part.get_param("name", header="content-disposition") or "": (
part.get_filename() or "",
part.get_content_type(),
part.get_payload(decode=True) or b"",
)
for part in parts
if part.get_filename() is not None
}
def _drained_other_than_model_list_probes(wire: Wire) -> tuple[Request, ...]:
return tuple(
request for request in wire.drain() if urlsplit(request.target).path not in {"/v1/models", "/openai/models"}
)
def _bearer(request: Request) -> str:
return request.headers.get("authorization", request.headers.get("api-key", "")).removeprefix("Bearer ")
def _sdk_base_url(gateway: Gateway) -> str:
return str(gateway.client.base_url).rstrip("/") + "/v1"
def _decoded_unified(managed_file_id: str) -> str:
decoded: Final = base64.urlsafe_b64decode(managed_file_id + "=" * (-len(managed_file_id) % 4)).decode()
assert decoded.startswith(MANAGED_PREFIX), decoded
return decoded
def _models_over_a_fresh_connection(gateway: Gateway, _: int) -> frozenset[str]:
with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as client:
listed: Final = client.get("/v1/models", headers={"Authorization": f"Bearer {gateway.key}"})
assert listed.status_code == 200, listed.text
data: Final = _json(listed)["data"]
assert isinstance(data, list), listed.text
return frozenset(string_value(object_value(entry)["id"]) for entry in data)
def _every_worker_serves(gateway: Gateway, *models: str) -> bool:
with ThreadPoolExecutor(max_workers=16) as pool:
rounds: Final = tuple(
tuple(pool.map(partial(_models_over_a_fresh_connection, gateway), range(16))) for _ in range(2)
)
return all(model in seen for model in models for round_ in rounds for seen in round_)
def _wait_until_every_worker_serves(gateway: Gateway, *models: str) -> None:
eventually(lambda: _every_worker_serves(gateway, *models), lambda served: served, seconds=90)
def _deployment_ids(gateway: Gateway, *model_names: str) -> dict[str, str]:
entries: Final = gateway.get("/model/info")["data"]
assert isinstance(entries, list)
by_name: Final = {
string_value(entry["model_name"]): string_value(object_value(entry["model_info"])["id"])
for entry in entries
if isinstance(entry, dict) and entry.get("model_name") in model_names
}
return by_name
def _multiplexed_openai_backends(
created_batch: Mapping[str, JsonValue] | None = None,
content_gate: threading.Event | None = None,
) -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
raw_path: Final = urlsplit(request.target).path
path: Final = "/v1/" + raw_path.removeprefix("/openai/") if raw_path.startswith("/openai/") else raw_path
if request.method == "POST" and path == "/v1/files":
file_part: Final = _file_fields(_multipart_parts(request))["file"]
return _json_reply(_file_object("file-" + _bearer(request), len(file_part[2]), file_part[0]))
if request.method == "POST" and path == "/v1/batches":
if created_batch is not None:
return _json_reply(created_batch)
return _json_reply(_batch_object("batch-" + _bearer(request), "file-pending", "in_progress"))
if request.method == "POST" and re.fullmatch(r"/v1/batches/[^/]+/cancel", path):
return _json_reply(_batch_object("batch-" + _bearer(request), "file-pending", "cancelled"))
if request.method == "GET" and path.startswith("/v1/batches/"):
return _json_reply(
_batch_object(
path.rsplit("/", 1)[1], "file-in", "completed", output_file_id="file-out-" + _bearer(request)
)
)
if request.method == "GET" and path.startswith("/v1/files/") and path.endswith("/content"):
if content_gate is not None:
return Reply(
content_type="application/octet-stream",
chunks=(OUTPUT_HEAD, OUTPUT_BYTES),
gate_after_first=content_gate,
)
return Reply(body=OUTPUT_BYTES, content_type="application/octet-stream")
if request.method == "GET" and path.startswith("/v1/files/"):
return _json_reply(_file_object(path.rsplit("/", 1)[1], len(_jsonl("uploaded-model"))))
if request.method == "DELETE" and path.startswith("/v1/files/"):
return _json_reply({"id": path.rsplit("/", 1)[1], "object": "file", "deleted": True})
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
return respond
def _managed_upload(
gateway: Gateway, spelling: _BatchUploadSpelling, fields: Mapping[str, object], sdk_base_url: str
) -> _FileObject:
key: Final = gateway.key
match spelling:
case "sdk_list":
with OpenAI(base_url=sdk_base_url, api_key=key, max_retries=0) as client:
response: Final = client.files.with_raw_response.create(
file=(INPUT_FILENAME, _jsonl("uploaded-model"), "application/jsonl"),
purpose="batch",
extra_body=fields,
)
assert response.status_code == 200, response.text
return _FileObject.model_validate_json(response.content)
case "repeated_bracket" | "indexed" | "repeated_plain":
posted: Final = gateway.client.post(
"/v1/files",
data={"purpose": "batch", **fields},
files={"file": (INPUT_FILENAME, _jsonl("uploaded-model"), "application/jsonl")},
headers={"Authorization": f"Bearer {key}"},
)
assert posted.status_code == 200, posted.text
return _FileObject.model_validate_json(posted.content)
case "comma_joined":
created: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", **fields},
{"file": (INPUT_FILENAME, _jsonl("uploaded-model"), "application/jsonl")},
)
assert created.status_code == 200, created.text
return _FileObject.model_validate_json(created.content)
case _:
assert_never(spelling)
def _assert_upload_reaches_both_target_models(
gateway: Gateway, spelling: _BatchUploadSpelling, sdk_base_url: str
) -> None:
key_a: Final = "provider-key-a-" + uuid.uuid4().hex[:8]
key_b: Final = "provider-key-b-" + uuid.uuid4().hex[:8]
with wire_server(_multiplexed_openai_backends()) as wire, gateway.scenario() as scenario:
model_a: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_a)
model_b: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_b)
_wait_until_every_worker_serves(gateway, model_a, model_b)
fields: Final[Mapping[str, object]] = {
"sdk_list": {"target_model_names": [model_a, model_b]},
"repeated_bracket": {"target_model_names[]": [model_a, model_b]},
"indexed": {"target_model_names[0]": model_a, "target_model_names[1]": model_b},
"comma_joined": {"target_model_names": f"{model_a},{model_b}"},
"repeated_plain": {"target_model_names": [model_a, model_b]},
}[spelling]
managed_file: Final = _managed_upload(gateway, spelling, fields, sdk_base_url)
assert managed_file.created_at > 0, managed_file.model_dump()
managed: Final = managed_file.id
requests: Final = _drained_other_than_model_list_probes(wire)
assert sorted((r.method, r.target, r.headers["authorization"]) for r in requests) == [
("POST", "/v1/files", f"Bearer {key_a}"),
("POST", "/v1/files", f"Bearer {key_b}"),
], requests
uploads: Final = {request.headers["authorization"]: request for request in requests}
for bearer, underlying in ((key_a, "gpt-4o-mini"), (key_b, "gpt-4.1-mini")):
parts: Final = _multipart_parts(uploads[f"Bearer {bearer}"])
assert len(parts) == 2, [part.get_param("name", header="content-disposition") for part in parts]
assert _text_fields(parts) == {"purpose": "batch"}
file_part: Final = _file_fields(parts)["file"]
assert (file_part[0], file_part[2]) == ("modified_file.jsonl", _rewritten_jsonl(underlying)), file_part
managed_file_part: Final = _file_fields(_multipart_parts(uploads[f"Bearer {key_a}"]))["file"]
expected_file: Final = _expected_file_object(
managed_file.id,
len(managed_file_part[2]),
managed_file_part[0],
)
assert managed_file == expected_file, managed_file.model_dump()
decoded: Final = _decoded_unified(managed)
assert f"target_model_names,{model_a},{model_b}" in decoded, decoded
deployment_ids: Final = _deployment_ids(gateway, model_a, model_b)
rows: Final = eventually(
lambda: read_rows(MANAGED_FILE_ROW, (managed,)),
lambda found: len(found) == 1,
)
assert rows[0]["model_mappings"] == {
deployment_ids[model_a]: f"file-{key_a}",
deployment_ids[model_b]: f"file-{key_b}",
}, rows
assert sorted(rows[0]["flat_model_file_ids"]) == sorted([f"file-{key_a}", f"file-{key_b}"]), rows
readback: Final = gateway.request("GET", f"/v1/files/{managed}")
assert readback.status_code == 200, readback.text
readback_file: Final = _FileObject.model_validate_json(readback.content)
assert readback_file == expected_file, readback.text
assert _drained_other_than_model_list_probes(wire) == ()
@pytest.mark.parametrize(
"spelling",
["sdk_list", "repeated_bracket", "indexed", "comma_joined"],
ids=["sdk-list-extra-body", "multipart-repeated-bracket", "multipart-indexed", "multipart-comma-joined"],
)
def test_batch_upload_to_two_target_models_reaches_each_deployment(
gateway: Gateway, spelling: _BatchUploadSpelling
) -> None:
_assert_upload_reaches_both_target_models(gateway, spelling, _sdk_base_url(gateway))
@pytest.mark.parametrize(
"route",
["/files", "/openai/v1/files"],
ids=["files-alias", "openai-v1-files-alias"],
)
def test_batch_upload_to_two_target_models_through_a_route_alias_reaches_each_deployment(
gateway: Gateway, route: str
) -> None:
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + route.removesuffix("/files")
_assert_upload_reaches_both_target_models(gateway, "sdk_list", sdk_base_url)
def test_batch_upload_with_repeated_plain_target_model_names_reaches_each_deployment(gateway: Gateway) -> None:
pytest.skip("BUG: repeated plain target_model_names keeps only the last value")
_assert_upload_reaches_both_target_models(gateway, "repeated_plain", _sdk_base_url(gateway))
def test_batch_upload_to_two_target_models_keeps_the_jsonl_content_type(gateway: Gateway) -> None:
pytest.skip(
"BUG: managed multi-target upload sends the rewritten batch file as application/octet-stream instead of "
"application/jsonl"
)
key_a: Final = "provider-key-a-" + uuid.uuid4().hex[:8]
key_b: Final = "provider-key-b-" + uuid.uuid4().hex[:8]
with wire_server(_multiplexed_openai_backends()) as wire, gateway.scenario() as scenario:
model_a: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_a)
model_b: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_b)
_wait_until_every_worker_serves(gateway, model_a, model_b)
_managed_upload(
gateway,
"repeated_bracket",
{"target_model_names[]": [model_a, model_b]},
_sdk_base_url(gateway),
)
uploads: Final = {
request.headers["authorization"]: request for request in _drained_other_than_model_list_probes(wire)
}
content_types: Final = {
f"Bearer {key}": _file_fields(_multipart_parts(uploads[f"Bearer {key}"]))["file"][1]
for key in (key_a, key_b)
}
assert content_types == {
f"Bearer {key_a}": "application/jsonl",
f"Bearer {key_b}": "application/jsonl",
}, content_types
def _model_routed_upload(
gateway: Gateway, spelling: _ModelRoutedSpelling, alias: str, uploaded: bytes, sdk_base_url: str
) -> _FileObject:
match spelling:
case "sdk_extra_body":
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
response: Final = client.files.with_raw_response.create(
file=(INPUT_FILENAME, uploaded, "application/jsonl"),
purpose="batch",
extra_body={"model": alias},
)
assert response.status_code == 200, response.text
return _FileObject.model_validate_json(response.content)
case "sdk_proxy_fields":
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
response: Final = client.files.with_raw_response.create(
file=(INPUT_FILENAME, uploaded, "application/jsonl"),
purpose="batch",
extra_body={
"model": alias,
"custom_llm_provider": "openai",
"target_model_names": alias,
},
)
assert response.status_code == 200, response.text
return _FileObject.model_validate_json(response.content)
case "multipart_field":
created: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": alias},
{"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
)
assert created.status_code == 200, created.text
return _FileObject.model_validate_json(created.content)
case "query_param":
queried: Final = gateway.client.post(
"/v1/files",
params={"model": alias},
data={"purpose": "batch"},
files={"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
headers={"Authorization": f"Bearer {gateway.key}"},
)
assert queried.status_code == 200, queried.text
return _FileObject.model_validate_json(queried.content)
case "model_header":
headed: Final = gateway.client.post(
"/v1/files",
data={"purpose": "batch"},
files={"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
headers={"Authorization": f"Bearer {gateway.key}", "x-litellm-model": alias},
)
assert headed.status_code == 200, headed.text
return _FileObject.model_validate_json(headed.content)
case "multipart_proxy_fields":
posted: Final = gateway.client.post(
"/v1/files",
data={
"purpose": "batch",
"model": alias,
"target_model_names": alias,
"custom_llm_provider": "openai",
"litellm_metadata[tag]": "x",
},
files={"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
headers={"Authorization": f"Bearer {gateway.key}"},
)
assert posted.status_code == 200, posted.text
return _FileObject.model_validate_json(posted.content)
case "sdk_expires_after":
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
response: Final = client.files.with_raw_response.create(
file=(INPUT_FILENAME, uploaded, "application/jsonl"),
purpose="batch",
expires_after={"anchor": "created_at", "seconds": 2592000},
extra_body={"model": alias},
)
assert response.status_code == 200, response.text
return _FileObject.model_validate_json(response.content)
case "multipart_expires_after":
created: Final = gateway.request_multipart(
"/v1/files",
{
"purpose": "batch",
"model": alias,
"expires_after[anchor]": "created_at",
"expires_after[seconds]": "2592000",
},
{"file": (INPUT_FILENAME, uploaded, "application/jsonl")},
)
assert created.status_code == 200, created.text
return _FileObject.model_validate_json(created.content)
case _:
assert_never(spelling)
def _assert_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
gateway: Gateway, spelling: _ModelRoutedSpelling, sdk_base_url: str
) -> None:
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
uploaded: Final = _jsonl("uploaded-model")
with wire_server(_multiplexed_openai_backends()) as wire, gateway.scenario() as scenario:
alias: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
_wait_until_every_worker_serves(gateway, alias)
uploaded_file: Final = _model_routed_upload(gateway, spelling, alias, uploaded, sdk_base_url)
expected_encoded: Final = encode_file_id_with_model(file_id=f"file-{key}", model=alias)
assert uploaded_file == _expected_file_object(expected_encoded, len(uploaded)), uploaded_file.model_dump()
encoded: Final = uploaded_file.id
(upload,) = _drained_other_than_model_list_probes(wire)
assert (upload.method, upload.target) == ("POST", "/v1/files"), upload.target
assert upload.headers["authorization"] == f"Bearer {key}", upload.headers
parts: Final = _multipart_parts(upload)
expected_text: Final = {
"sdk_extra_body": {"purpose": "batch"},
"multipart_field": {"purpose": "batch"},
"query_param": {"purpose": "batch"},
"model_header": {"purpose": "batch"},
"multipart_proxy_fields": {"purpose": "batch"},
"sdk_proxy_fields": {"purpose": "batch"},
"sdk_expires_after": {
"purpose": "batch",
"expires_after[anchor]": "created_at",
"expires_after[seconds]": "2592000",
},
"multipart_expires_after": {
"purpose": "batch",
"expires_after[anchor]": "created_at",
"expires_after[seconds]": "2592000",
},
}[spelling]
assert len(parts) == len(expected_text) + 1, [
part.get_param("name", header="content-disposition") for part in parts
]
assert _text_fields(parts) == expected_text
assert _file_fields(parts) == {"file": (INPUT_FILENAME, "application/jsonl", uploaded)}
readback: Final = gateway.request("GET", f"/v1/files/{encoded}")
assert readback.status_code == 200, readback.text
readback_file: Final = _FileObject.model_validate_json(readback.content)
assert readback_file == _expected_file_object(encoded, len(uploaded)), readback.text
(fetched,) = _drained_other_than_model_list_probes(wire)
assert (fetched.method, fetched.target) == ("GET", f"/v1/files/file-{key}"), fetched.target
assert fetched.headers["authorization"] == f"Bearer {key}", fetched.headers
assert fetched.body == b""
@pytest.mark.parametrize(
"spelling",
[
"sdk_extra_body",
"sdk_proxy_fields",
"multipart_field",
"query_param",
"model_header",
"multipart_proxy_fields",
"sdk_expires_after",
"multipart_expires_after",
],
ids=[
"sdk-extra-body-model",
"sdk-extra-body-with-proxy-fields",
"multipart-model-field",
"query-model",
"x-litellm-model-header",
"multipart-model-with-proxy-fields",
"sdk-expires-after",
"multipart-bracketed-expires-after",
],
)
def test_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
gateway: Gateway, spelling: _ModelRoutedSpelling
) -> None:
_assert_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
gateway, spelling, _sdk_base_url(gateway)
)
@pytest.mark.parametrize(
"route",
["/files", "/openai/v1/files"],
ids=["files-alias", "openai-v1-files-alias"],
)
def test_model_routed_upload_through_a_route_alias_reaches_one_deployment(
gateway: Gateway, route: str
) -> None:
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + route.removesuffix("/files")
_assert_model_routed_upload_reaches_one_deployment_and_the_encoded_id_reads_back(
gateway, "sdk_extra_body", sdk_base_url
)
AZURE_API_VERSION: Final = "2024-10-21"
def _encoded_output_file_id(gateway: Gateway, alias: str, bearer_key: str) -> tuple[str, str]:
uploaded_bytes: Final = _jsonl("uploaded-model")
uploaded: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": alias},
{"file": (INPUT_FILENAME, uploaded_bytes, "application/jsonl")},
)
assert uploaded.status_code == 200, uploaded.text
encoded_file: Final = encode_file_id_with_model(file_id=f"file-{bearer_key}", model=alias)
uploaded_file: Final = _FileObject.model_validate_json(uploaded.content)
assert uploaded_file == _expected_file_object(encoded_file, len(uploaded_bytes)), uploaded.text
encoded_input: Final = uploaded_file.id
created: Final = gateway.request(
"POST",
"/v1/batches",
{"input_file_id": encoded_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
)
assert created.status_code == 200, created.text
encoded_batch: Final = string_value(_json(created)["id"])
retrieved: Final = gateway.request("GET", f"/v1/batches/{encoded_batch}")
assert retrieved.status_code == 200, retrieved.text
batch: Final = _Batch.model_validate_json(retrieved.content)
assert batch.output_file_id is not None, retrieved.text
return batch.output_file_id, f"file-out-{bearer_key}"
def _content_request(
gateway: Gateway, route: _ContentRoute, encoded_output: str, raw_output: str, alias: str
) -> httpx.Request:
authorization: Final = {"Authorization": f"Bearer {gateway.key}"}
match route:
case "v1_encoded":
return gateway.client.build_request("GET", f"/v1/files/{encoded_output}/content", headers=authorization)
case "files_encoded":
return gateway.client.build_request("GET", f"/files/{encoded_output}/content", headers=authorization)
case "files_raw_query":
return gateway.client.build_request(
"GET", f"/files/{raw_output}/content", params={"model": alias}, headers=authorization
)
case "v1_raw_header":
return gateway.client.build_request(
"GET", f"/v1/files/{raw_output}/content", headers={**authorization, "x-litellm-model": alias}
)
case _:
assert_never(route)
def _open_and_read_head(
client: httpx.Client, request: httpx.Request
) -> tuple[httpx.Response, Iterator[bytes], bytes]:
response: Final = client.send(request, stream=True)
chunks: Final = response.iter_bytes()
return response, chunks, next(chunks, b"")
def _gated_download(gateway: Gateway, request: httpx.Request, gate: threading.Event) -> tuple[int, bytes, bytes]:
with ThreadPoolExecutor(max_workers=1) as reader:
head_read: Final = reader.submit(_open_and_read_head, gateway.client, request)
try:
head_arrived: Final = eventually(head_read.done, bool, seconds=4.5, return_last_on_timeout=True)
finally:
gate.set()
response, chunks, head = head_read.result()
try:
rest: Final = b"".join(chunks)
finally:
response.close()
assert head_arrived, (
f"client got no content bytes while the upstream held the rest of the file: {response.status_code}"
)
return response.status_code, head, rest
@pytest.mark.parametrize("provider", ["openai", "azure"], ids=["openai-streaming", "azure-non-streaming"])
@pytest.mark.parametrize(
"route",
["v1_encoded", "files_encoded", "files_raw_query", "v1_raw_header"],
ids=["v1-files-encoded", "files-encoded", "files-raw-model-query", "v1-files-raw-model-header"],
)
def test_batch_output_file_content_downloads_through_the_model_encoded_id(
gateway: Gateway, provider: str, route: _ContentRoute
) -> None:
key: Final = f"provider-key-{provider}-" + uuid.uuid4().hex[:8]
gate: Final = threading.Event()
gate.set()
received: Final[Queue[Request]] = Queue()
backend: Final = _multiplexed_openai_backends(content_gate=gate)
def respond(request: Request) -> Reply:
received.put(request)
return backend(request)
with wire_server(respond) as wire, gateway.scenario() as scenario:
deployment: Final = (
{"model": "openai/gpt-4o-mini", "api_base": wire.url + "/v1", "api_key": key}
if provider == "openai"
else {
"model": "azure/gpt-4o-mini",
"api_base": wire.url,
"api_version": AZURE_API_VERSION,
"api_key": key,
}
)
alias: Final = scenario.model(**deployment)
_wait_until_every_worker_serves(gateway, alias)
encoded_output, raw_output = _encoded_output_file_id(gateway, alias, key)
assert encoded_output == encode_file_id_with_model(file_id=raw_output, model=alias), encoded_output
eventually(
lambda: tuple(cast(deque[Request], received.queue)),
lambda pending: any(urlsplit(request.target).path.endswith("/content") for request in pending),
)
settled: Final = _drained_other_than_model_list_probes(wire)
base_prefix: Final = "/v1/" if provider == "openai" else "/openai/"
api_suffix: Final = "" if provider == "openai" else f"?api-version={AZURE_API_VERSION}"
expected_content_target: Final = f"{base_prefix}files/{raw_output}/content{api_suffix}"
expected_requests: Final = (
("POST", f"{base_prefix}files{api_suffix}"),
("POST", f"{base_prefix}batches{api_suffix}"),
("GET", f"{base_prefix}batches/batch-{key}{api_suffix}"),
("GET", expected_content_target),
)
observed_requests: Final = tuple((r.method, r.target) for r in settled)
expected_metadata_request: Final = ("GET", f"{base_prefix}files/{raw_output}{api_suffix}")
assert observed_requests.count(expected_metadata_request) <= 1, observed_requests
assert tuple(request for request in observed_requests if request != expected_metadata_request) == (
expected_requests
), observed_requests
if provider == "openai":
gate.clear()
status, head, rest = _gated_download(
gateway, _content_request(gateway, route, encoded_output, raw_output, alias), gate
)
assert status == 200, (head + rest).decode()
assert head + rest == OUTPUT_HEAD + OUTPUT_BYTES, f"{len(head)} head bytes, {len(rest)} tail bytes"
(content_request,) = _drained_other_than_model_list_probes(wire)
assert (content_request.method, content_request.target) == ("GET", expected_content_target), (
content_request.target
)
assert content_request.body == b"", content_request.body
if provider == "openai":
assert content_request.headers["authorization"] == f"Bearer {key}", content_request.headers
else:
assert content_request.headers["api-key"] == key, content_request.headers
def _encoded_batch_input(
gateway: Gateway, spelling: _BatchCreateSpelling, alias: str, wire: Wire, raw_input: str, bearer_key: str
) -> str:
if spelling == "encoded_input":
uploaded_bytes: Final = _jsonl("uploaded-model")
uploaded: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": alias},
{"file": (INPUT_FILENAME, uploaded_bytes, "application/jsonl")},
)
assert uploaded.status_code == 200, uploaded.text
encoded_file: Final = encode_file_id_with_model(file_id=f"file-{bearer_key}", model=alias)
uploaded_file: Final = _FileObject.model_validate_json(uploaded.content)
assert uploaded_file == _expected_file_object(encoded_file, len(uploaded_bytes)), uploaded.text
wire.drain()
return uploaded_file.id
return encode_file_id_with_model(file_id=raw_input, model=alias)
def _created_batch(
gateway: Gateway,
spelling: _BatchCreateSpelling,
encoded_input: str,
raw_input: str,
alias: str,
sdk_base_url: str,
) -> _Batch:
if spelling in {"query_model", "header_model"}:
created: Final = gateway.request(
"POST",
"/v1/batches",
{
"input_file_id": raw_input,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"metadata": {"job": "x"},
},
params={"model": alias} if spelling == "query_model" else None,
headers={"x-litellm-model": alias} if spelling == "header_model" else None,
)
assert created.status_code == 200, created.text
return _Batch.model_validate_json(created.content)
if spelling == "sanitized_metadata":
created: Final = gateway.request(
"POST",
"/v1/batches",
{
"input_file_id": raw_input,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"metadata": {"job": "x", "attempt": 2},
"model": alias,
},
)
assert created.status_code == 200, created.text
return _Batch.model_validate_json(created.content)
with OpenAI(base_url=sdk_base_url, api_key=gateway.key, max_retries=0) as client:
extra: Final = {} if spelling == "encoded_input" else {"extra_body": {"model": alias}}
response: Final = client.batches.with_raw_response.create(
input_file_id=encoded_input if spelling == "encoded_input" else raw_input,
endpoint="/v1/chat/completions",
completion_window="24h",
metadata={"job": "x"},
**extra,
)
assert response.status_code == 200, response.text
return _Batch.model_validate_json(response.content)
def _assert_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(
gateway: Gateway, spelling: _BatchCreateSpelling, sdk_base_url: str
) -> None:
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
raw_input: Final = f"file-in-{key}"
raw_batch: Final = f"batch-{key}"
created_batch: Final = _batch_object(
raw_batch, raw_input, "validating", output_file_id="file-out", error_file_id="file-err"
)
with wire_server(_multiplexed_openai_backends(created_batch)) as wire, gateway.scenario() as scenario:
alias: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
_wait_until_every_worker_serves(gateway, alias)
encoded_input: Final = _encoded_batch_input(gateway, spelling, alias, wire, raw_input, key)
created_batch_payload: Final = _created_batch(
gateway, spelling, encoded_input, raw_input, alias, sdk_base_url
)
expected_raw_input: Final = f"file-{key}" if spelling == "encoded_input" else raw_input
expected_input: Final = (
encoded_input if spelling == "encoded_input" else encode_file_id_with_model(file_id=raw_input, model=alias)
)
assert created_batch_payload.id == encode_file_id_with_model(file_id=raw_batch, model=alias, id_type="batch"), (
created_batch_payload.id
)
assert created_batch_payload.output_file_id == encode_file_id_with_model(file_id="file-out", model=alias), (
created_batch_payload.output_file_id
)
assert created_batch_payload.error_file_id == encode_file_id_with_model(file_id="file-err", model=alias), (
created_batch_payload.error_file_id
)
assert created_batch_payload.input_file_id == expected_input, created_batch_payload.input_file_id
requests: Final = _drained_other_than_model_list_probes(wire)
assert [(r.method, r.target) for r in requests] == [("POST", "/v1/batches")], requests
(create_request,) = requests
assert create_request.headers["authorization"] == f"Bearer {key}", create_request.headers
expected_body: Final = {
"input_file_id": expected_raw_input,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"metadata": {"job": "x"},
}
assert JSON_OBJECT.validate_json(create_request.body) == expected_body, create_request.body
@pytest.mark.parametrize(
"spelling",
["encoded_input", "body_model", "query_model", "header_model", "sanitized_metadata"],
ids=["encoded-input-id", "body-model", "query-model", "x-litellm-model-header", "non-string-metadata-dropped"],
)
def test_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(
gateway: Gateway, spelling: _BatchCreateSpelling
) -> None:
_assert_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(
gateway, spelling, _sdk_base_url(gateway)
)
@pytest.mark.parametrize("route", ["/batches", "/openai/v1/batches"], ids=["batches-alias", "openai-v1-batches-alias"])
def test_create_batch_through_a_route_alias_reaches_one_deployment(gateway: Gateway, route: str) -> None:
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + route.removesuffix("/batches")
_assert_create_batch_reaches_one_deployment_and_returns_model_encoded_ids(gateway, "encoded_input", sdk_base_url)

View file

@ -0,0 +1,851 @@
import base64
import contextlib
import hashlib
import json
import socket
import socketserver
import threading
import uuid
from collections.abc import Callable, Generator
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from queue import SimpleQueue
from typing import Final, Literal
from urllib.parse import quote, urlsplit
import httpx
import pytest
from botocore.auth import S3SigV4Auth
from botocore.awsrequest import AWSRequest
from botocore.credentials import Credentials
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
from integration._support.database import read_rows
from integration._support.process import owned_proxy
from integration._support.tls import server_context, write_self_signed_cert
from integration._support.wire import Reply, Request, Wire, wire_server
from openai import OpenAI, PermissionDeniedError
from pydantic import BaseModel, JsonValue, TypeAdapter
from typing_extensions import assert_never
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
MANAGED_FILE_ROW: Final = 'SELECT model_mappings, flat_model_file_ids, created_by, team_id FROM "LiteLLM_ManagedFileTable" WHERE unified_file_id = %s'
INPUT_FILENAME: Final = "in.jsonl"
MANAGED_PREFIX: Final = "litellm_proxy:"
BUCKET: Final = "integration-delete-bucket"
VERTEX_PROJECT: Final = "files-batches-delete-project"
VERTEX_LOCATION: Final = "us-central1"
VERTEX_MODEL: Final = "vertex_ai/gemini-2.5-flash"
GCS_HOST: Final = "storage.googleapis.com"
GCS_AUTHORITY: Final = f"{GCS_HOST}:443"
VERTEX_ACCESS_TOKEN: Final = "ya29.scripted-delete-token"
AWS_ACCESS_KEY_ID: Final = "AKIAIOSFODNN7EXAMPLE"
AWS_SECRET_ACCESS_KEY: Final = "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"
AWS_REGION: Final = "us-east-1"
_ForbiddenSurface = Literal[
"upload_sdk_extra_body",
"upload_form",
"upload_query",
"upload_header",
"batch_encoded_input",
"batch_body_model",
"batch_query_model",
"batch_header_model",
"file_get_encoded",
"file_get_query",
"file_get_header",
"file_content_encoded",
"file_content_query",
"file_content_header",
"batch_get_encoded",
"batch_get_query",
"batch_get_header",
]
def _jsonl(model: str) -> bytes:
line: Final = {
"custom_id": "r1",
"method": "POST",
"url": "/v1/chat/completions",
"body": {"model": model, "messages": [{"role": "user", "content": "one"}]},
}
return json.dumps(line, separators=(",", ":")).encode() + b"\n"
class _FileObject(BaseModel):
id: str
object: str
purpose: str
class _Batch(BaseModel):
id: str
object: str
status: str
input_file_id: str | None = None
class _ErrorEnvelope(BaseModel):
error: dict[str, JsonValue]
class _FileDeleted(BaseModel):
id: str
object: str
deleted: bool
def _json(response: httpx.Response) -> dict[str, JsonValue]:
return JSON_OBJECT.validate_json(response.content)
def _json_reply(body: dict[str, JsonValue], status: int = 200) -> Reply:
return Reply(status=status, body=json.dumps(body).encode())
def _drained_other_than_model_list_probes(wire: Wire) -> tuple[Request, ...]:
return tuple(
request for request in wire.drain() if urlsplit(request.target).path not in {"/v1/models", "/openai/models"}
)
def _bearer(request: Request) -> str:
return request.headers.get("authorization", request.headers.get("api-key", "")).removeprefix("Bearer ")
def _file_object(file_id: str) -> dict[str, JsonValue]:
return {
"id": file_id,
"object": "file",
"bytes": 123,
"created_at": 1700000000,
"filename": INPUT_FILENAME,
"purpose": "batch",
"status": "processed",
}
def _batch_object(batch_id: str, status: str) -> dict[str, JsonValue]:
return {
"id": batch_id,
"object": "batch",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-pending",
"status": status,
"created_at": 1700000000,
"completion_window": "24h",
}
def _openai_backend() -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
path: Final = urlsplit(request.target).path
if request.method == "POST" and path == "/v1/files":
return _json_reply(_file_object("file-" + _bearer(request)))
if request.method == "POST" and path == "/v1/batches":
return _json_reply(_batch_object("batch-" + _bearer(request), "in_progress"))
if request.method == "POST" and path.endswith("/cancel"):
return _json_reply(_batch_object(path.rsplit("/", 2)[1], "cancelled"))
if request.method == "GET" and path.startswith("/v1/batches/"):
return _json_reply(_batch_object(path.rsplit("/", 1)[1], "in_progress"))
if request.method == "GET" and path.startswith("/v1/files/") and path.endswith("/content"):
return Reply(body=b"file-content\n", content_type="application/octet-stream")
if request.method == "GET" and path.startswith("/v1/files/"):
return _json_reply(_file_object(path.rsplit("/", 1)[1]))
if request.method == "DELETE" and path.startswith("/v1/files/"):
return _json_reply({"id": path.rsplit("/", 1)[1], "object": "file", "deleted": True})
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
return respond
def _decoded_unified(managed_file_id: str) -> str:
decoded: Final = base64.urlsafe_b64decode(managed_file_id + "=" * (-len(managed_file_id) % 4)).decode()
assert decoded.startswith(MANAGED_PREFIX), decoded
return decoded
def _models_over_a_fresh_connection(gateway: Gateway, _: int) -> frozenset[str]:
with httpx.Client(base_url=gateway.client.base_url, timeout=15, trust_env=False) as client:
listed: Final = client.get("/v1/models", headers={"Authorization": f"Bearer {gateway.key}"})
assert listed.status_code == 200, listed.text
data: Final = _json(listed)["data"]
assert isinstance(data, list), listed.text
return frozenset(string_value(object_value(entry)["id"]) for entry in data)
def _every_worker_serves(gateway: Gateway, *models: str) -> bool:
with ThreadPoolExecutor(max_workers=16) as pool:
rounds: Final = tuple(
tuple(pool.map(partial(_models_over_a_fresh_connection, gateway), range(16))) for _ in range(2)
)
return all(model in seen for model in models for round_ in rounds for seen in round_)
def _wait_until_every_worker_serves(gateway: Gateway, *models: str) -> None:
eventually(lambda: _every_worker_serves(gateway, *models), lambda served: served, seconds=90)
@dataclass(frozen=True, slots=True)
class _Member:
team: str
user: str
key: str
def _member(scenario: Scenario, *models: str) -> _Member:
team: Final = scenario.team(models=list(models))
user: Final = scenario.member(team)
return _Member(team, user, scenario.key(team_id=team, user_id=user))
def _error_message(response: httpx.Response, status: int) -> str:
assert response.status_code == status, response.text
error: Final = _ErrorEnvelope.model_validate_json(response.content)
return str(error.error["message"])
def test_stranger_cannot_read_download_or_delete_another_teams_managed_file(gateway: Gateway) -> None:
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
_wait_until_every_worker_serves(gateway, model)
owner: Final = _member(scenario, model)
stranger: Final = _member(scenario, model)
uploaded: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "target_model_names": model},
{"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")},
key=owner.key,
)
assert uploaded.status_code == 200, uploaded.text
managed: Final = string_value(_json(uploaded)["id"])
wire.drain()
row: Final = eventually(lambda: read_rows(MANAGED_FILE_ROW, (managed,)), lambda rows: len(rows) == 1)[0]
denied: Final = (
gateway.request("GET", f"/v1/files/{managed}", key=stranger.key),
gateway.request("GET", f"/v1/files/{managed}/content", key=stranger.key),
gateway.request("DELETE", f"/v1/files/{managed}", key=stranger.key),
)
for response in denied:
assert (
_error_message(response, 403) == f"User {stranger.user} does not have access to the file {managed}"
), response.text
assert _drained_other_than_model_list_probes(wire) == (), "stranger calls must not reach the provider"
assert read_rows(MANAGED_FILE_ROW, (managed,)) == [row], read_rows(MANAGED_FILE_ROW, (managed,))
readback: Final = gateway.request("GET", f"/v1/files/{managed}", key=owner.key)
assert readback.status_code == 200, readback.text
assert string_value(_json(readback)["id"]) == managed, readback.text
assert _drained_other_than_model_list_probes(wire) == ()
def test_stranger_cannot_retrieve_or_cancel_another_teams_managed_batch(gateway: Gateway) -> None:
key: Final = "provider-key-" + uuid.uuid4().hex[:8]
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key)
_wait_until_every_worker_serves(gateway, model)
owner: Final = _member(scenario, model)
stranger: Final = _member(scenario, model)
uploaded: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "target_model_names": model},
{"file": (INPUT_FILENAME, _jsonl(model), "application/jsonl")},
key=owner.key,
)
assert uploaded.status_code == 200, uploaded.text
managed_input: Final = string_value(_json(uploaded)["id"])
created: Final = gateway.request(
"POST",
"/v1/batches",
{"input_file_id": managed_input, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
key=owner.key,
)
assert created.status_code == 200, created.text
unified_batch: Final = string_value(_json(created)["id"])
wire.drain()
expected_message: Final = f"User {stranger.user} does not have access to the object {unified_batch}"
stranger_get: Final = gateway.request("GET", f"/v1/batches/{unified_batch}", key=stranger.key)
assert _error_message(stranger_get, 403) == expected_message, stranger_get.text
stranger_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=stranger.key)
assert _error_message(stranger_cancel, 403) == expected_message, stranger_cancel.text
assert _drained_other_than_model_list_probes(wire) == (), "stranger batch calls must not reach the provider"
owner_cancel: Final = gateway.request("POST", f"/v1/batches/{unified_batch}/cancel", {}, key=owner.key)
assert owner_cancel.status_code == 200, owner_cancel.text
assert string_value(_json(owner_cancel)["id"]) == unified_batch, owner_cancel.text
(cancel_request,) = _drained_other_than_model_list_probes(wire)
assert (cancel_request.method, cancel_request.target) == ("POST", f"/v1/batches/batch-{key}/cancel"), (
cancel_request.target
)
assert cancel_request.headers["authorization"] == f"Bearer {key}", cancel_request.headers
def _sdk_upload_refusal(gateway: Gateway, forbidden: str, caller_key: str) -> httpx.Response:
sdk_base_url: Final = str(gateway.client.base_url).rstrip("/") + "/v1"
with OpenAI(base_url=sdk_base_url, api_key=caller_key, max_retries=0) as client:
try:
uploaded: Final = client.files.with_raw_response.create(
file=(INPUT_FILENAME, _jsonl(forbidden), "application/jsonl"),
purpose="batch",
extra_body={"model": forbidden},
)
except PermissionDeniedError as refused:
return refused.response
raise AssertionError(f"SDK upload with a forbidden model was accepted: {uploaded.status_code} {uploaded.text}")
def _forbidden_surface_request(
gateway: Gateway,
surface: _ForbiddenSurface,
forbidden: str,
caller_key: str,
raw_file: str,
raw_batch: str,
encoded_file: str,
encoded_batch: str,
) -> httpx.Response:
match surface:
case "upload_sdk_extra_body":
return _sdk_upload_refusal(gateway, forbidden, caller_key)
case "upload_form":
return gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": forbidden},
{"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
key=caller_key,
)
case "upload_query":
return gateway.client.post(
"/v1/files",
params={"model": forbidden},
data={"purpose": "batch"},
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
headers={"Authorization": f"Bearer {caller_key}"},
)
case "upload_header":
return gateway.client.post(
"/v1/files",
data={"purpose": "batch"},
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
headers={"Authorization": f"Bearer {caller_key}", "x-litellm-model": forbidden},
)
case "batch_encoded_input":
return gateway.request(
"POST",
"/v1/batches",
{
"input_file_id": encoded_file,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
key=caller_key,
)
case "batch_body_model":
return gateway.request(
"POST",
"/v1/batches",
{
"input_file_id": raw_file,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"model": forbidden,
},
key=caller_key,
)
case "batch_query_model":
return gateway.request(
"POST",
"/v1/batches",
{"input_file_id": raw_file, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
key=caller_key,
params={"model": forbidden},
)
case "batch_header_model":
return gateway.request(
"POST",
"/v1/batches",
{"input_file_id": raw_file, "endpoint": "/v1/chat/completions", "completion_window": "24h"},
key=caller_key,
headers={"x-litellm-model": forbidden},
)
case "file_get_encoded":
return gateway.request("GET", f"/v1/files/{encoded_file}", key=caller_key)
case "file_get_query":
return gateway.request("GET", f"/v1/files/{raw_file}", key=caller_key, params={"model": forbidden})
case "file_get_header":
return gateway.request(
"GET", f"/v1/files/{raw_file}", key=caller_key, headers={"x-litellm-model": forbidden}
)
case "file_content_encoded":
return gateway.request("GET", f"/v1/files/{encoded_file}/content", key=caller_key)
case "file_content_query":
return gateway.request(
"GET", f"/v1/files/{raw_file}/content", key=caller_key, params={"model": forbidden}
)
case "file_content_header":
return gateway.request(
"GET", f"/v1/files/{raw_file}/content", key=caller_key, headers={"x-litellm-model": forbidden}
)
case "batch_get_encoded":
return gateway.request("GET", f"/v1/batches/{encoded_batch}", key=caller_key)
case "batch_get_query":
return gateway.request("GET", f"/v1/batches/{raw_batch}", key=caller_key, params={"model": forbidden})
case "batch_get_header":
return gateway.request(
"GET", f"/v1/batches/{raw_batch}", key=caller_key, headers={"x-litellm-model": forbidden}
)
case _:
assert_never(surface)
@pytest.mark.parametrize(
"surface",
[
"upload_sdk_extra_body",
"upload_form",
"upload_query",
"upload_header",
"batch_encoded_input",
"batch_body_model",
"batch_query_model",
"batch_header_model",
"file_get_encoded",
"file_get_query",
"file_get_header",
"file_content_encoded",
"file_content_query",
"file_content_header",
"batch_get_encoded",
"batch_get_query",
"batch_get_header",
],
ids=[
"upload-sdk-extra-body-model",
"upload-form-model",
"upload-query-model",
"upload-header-model",
"batch-create-encoded-input",
"batch-create-body-model",
"batch-create-query-model",
"batch-create-model-header",
"file-retrieve-encoded-id",
"file-retrieve-model-query",
"file-retrieve-model-header",
"file-content-encoded-id",
"file-content-model-query",
"file-content-model-header",
"batch-retrieve-encoded-id",
"batch-retrieve-model-query",
"batch-retrieve-model-header",
],
)
def test_a_key_granted_one_model_cannot_spend_another_deployments_files_or_batches(
gateway: Gateway, surface: _ForbiddenSurface
) -> None:
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
_wait_until_every_worker_serves(gateway, allowed, forbidden)
team: Final = scenario.team(models=[allowed, forbidden])
user: Final = scenario.member(team)
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
caller: Final = _Member(team, user, caller_key)
provider_client: Final = httpx.Client(base_url=wire.url, timeout=15, trust_env=False)
with provider_client:
seeded_file: Final = provider_client.post(
"/v1/files",
data={"purpose": "batch"},
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
headers={"Authorization": f"Bearer {key_forbidden}"},
)
assert seeded_file.status_code == 200, seeded_file.text
raw_file_object: Final = _FileObject.model_validate_json(seeded_file.content)
assert raw_file_object.id == f"file-{key_forbidden}", seeded_file.text
seeded_batch: Final = provider_client.post(
"/v1/batches",
json={
"input_file_id": raw_file_object.id,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
},
headers={"Authorization": f"Bearer {key_forbidden}"},
)
assert seeded_batch.status_code == 200, seeded_batch.text
raw_batch_object: Final = _Batch.model_validate_json(seeded_batch.content)
assert raw_batch_object.id == f"batch-{key_forbidden}", seeded_batch.text
wire.drain()
encoded_forbidden_file: Final = encode_file_id_with_model(file_id=raw_file_object.id, model=forbidden)
encoded_forbidden_batch: Final = encode_file_id_with_model(
file_id=raw_batch_object.id, model=forbidden, id_type="batch"
)
response: Final = _forbidden_surface_request(
gateway,
surface,
forbidden,
caller.key,
raw_file_object.id,
raw_batch_object.id,
encoded_forbidden_file,
encoded_forbidden_batch,
)
provider_requests: Final = _drained_other_than_model_list_probes(wire)
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
assert _error_message(response, 403) == (
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
"Check the models available to you and try again."
), response.text
assert provider_requests == (), f"{surface} reached the provider: {provider_requests}"
control: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "model": allowed},
{"file": (INPUT_FILENAME, _jsonl(allowed), "application/jsonl")},
key=caller.key,
)
assert control.status_code == 200, control.text
(control_request,) = _drained_other_than_model_list_probes(wire)
assert control_request.headers["authorization"] == f"Bearer {key_allowed}", control_request.headers
def test_indexed_target_model_upload_cannot_use_a_forbidden_deployment(gateway: Gateway) -> None:
pytest.skip("BUG: target_model_names[0] upload bypasses the key's model restriction and reaches the deployment")
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
_wait_until_every_worker_serves(gateway, allowed, forbidden)
team: Final = scenario.team(models=[allowed, forbidden])
user: Final = scenario.member(team)
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
response: Final = gateway.client.post(
"/v1/files",
data={"purpose": "batch", "target_model_names[0]": forbidden},
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
headers={"Authorization": f"Bearer {caller_key}"},
)
provider_requests: Final = _drained_other_than_model_list_probes(wire)
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
assert _error_message(response, 403) == (
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
"Check the models available to you and try again."
), response.text
assert provider_requests == (), f"indexed target-model upload reached the provider: {provider_requests}"
def test_repeated_target_model_upload_cannot_use_a_forbidden_deployment(gateway: Gateway) -> None:
pytest.skip("BUG: repeated target_model_names[] upload bypasses the key's model restriction and reaches the deployment")
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
_wait_until_every_worker_serves(gateway, allowed, forbidden)
team: Final = scenario.team(models=[allowed, forbidden])
user: Final = scenario.member(team)
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
response: Final = gateway.client.post(
"/v1/files",
data={"purpose": "batch", "target_model_names[]": [allowed, forbidden]},
files={"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
headers={"Authorization": f"Bearer {caller_key}"},
)
provider_requests: Final = _drained_other_than_model_list_probes(wire)
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
assert _error_message(response, 403) == (
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
"Check the models available to you and try again."
), response.text
assert provider_requests == (), f"repeated target-model upload reached the provider: {provider_requests}"
def test_plain_target_model_upload_cannot_use_a_forbidden_deployment(gateway: Gateway) -> None:
key_allowed: Final = "provider-key-allowed-" + uuid.uuid4().hex[:8]
key_forbidden: Final = "provider-key-forbidden-" + uuid.uuid4().hex[:8]
with wire_server(_openai_backend()) as wire, gateway.scenario() as scenario:
allowed: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key=key_allowed)
forbidden: Final = scenario.model(model="openai/gpt-4.1-mini", api_base=wire.url + "/v1", api_key=key_forbidden)
_wait_until_every_worker_serves(gateway, allowed, forbidden)
team: Final = scenario.team(models=[allowed, forbidden])
user: Final = scenario.member(team)
caller_key: Final = scenario.key(team_id=team, user_id=user, models=[allowed])
response: Final = gateway.request_multipart(
"/v1/files",
{"purpose": "batch", "target_model_names": forbidden},
{"file": (INPUT_FILENAME, _jsonl(forbidden), "application/jsonl")},
key=caller_key,
)
provider_requests: Final = _drained_other_than_model_list_probes(wire)
assert response.status_code == 403, f"{response.text}; provider requests={provider_requests}"
assert _error_message(response, 403) == (
f"The requested model '{forbidden}' is not available for this API key, or the model name is invalid. "
"Check the models available to you and try again."
), response.text
assert provider_requests == (), f"plain target-model upload reached the provider: {provider_requests}"
def _s3_backend() -> Callable[[Request], Reply]:
def respond(request: Request) -> Reply:
if request.method in {"DELETE", "GET"}:
return Reply(status=204, body=b"")
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
return respond
def _bedrock_model(scenario: Scenario, wire: Wire) -> str:
return scenario.model(
model="bedrock/anthropic.claude-3-haiku-20240307-v1:0",
api_key=None,
api_base=None,
aws_access_key_id=AWS_ACCESS_KEY_ID,
aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
aws_region_name=AWS_REGION,
s3_bucket_name=BUCKET,
s3_endpoint_url=wire.url,
)
def _vertex_credentials(token_url: str) -> str:
private_key: Final = (
rsa.generate_private_key(public_exponent=65537, key_size=2048)
.private_bytes(
serialization.Encoding.PEM,
serialization.PrivateFormat.PKCS8,
serialization.NoEncryption(),
)
.decode()
)
return json.dumps(
{
"type": "service_account",
"project_id": VERTEX_PROJECT,
"private_key_id": "scripted",
"private_key": private_key,
"client_email": f"scripted@{VERTEX_PROJECT}.iam.gserviceaccount.com",
"client_id": "0",
"auth_uri": f"{token_url}/_oauth/authorize",
"token_uri": f"{token_url}/_oauth/token",
}
)
def _vertex_token_backend(request: Request) -> Reply:
if request.method == "POST" and request.target == "/_oauth/token":
return Reply(status=400, body=b'{"error":"invalid_grant"}', content_type="application/json")
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
def test_non_admin_cannot_delete_a_raw_s3_uri(gateway: Gateway) -> None:
with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario:
model: Final = _bedrock_model(scenario, wire)
_wait_until_every_worker_serves(gateway, model)
caller_user: Final = scenario.user(user_role="internal_user")
caller_key: Final = scenario.key(user_id=caller_user, models=[model])
object_uri: Final = f"s3://{BUCKET}/litellm-bedrock-files-obj.jsonl"
provider_path: Final = "/bedrock/v1/files/" + quote(object_uri, safe="")
denied: Final = gateway.request("DELETE", provider_path, key=caller_key, params={"model": model})
assert denied.status_code == 403, denied.text
assert _error_message(denied, 403) == (
"Raw cloud storage file ids can only be deleted by a proxy admin key. "
"Use the LiteLLM managed file id returned when the file was created."
), denied.text
assert _drained_other_than_model_list_probes(wire) == (), (
"non-admin cloud-storage delete reached the storage backend"
)
def test_non_admin_cannot_delete_a_raw_gcs_uri_before_vertex_auth_or_storage(gateway: Gateway) -> None:
with (
wire_server(_vertex_token_backend) as token_wire,
gateway.scenario() as scenario,
):
model: Final = scenario.model(
model=VERTEX_MODEL,
api_key=None,
vertex_project=VERTEX_PROJECT,
vertex_location=VERTEX_LOCATION,
vertex_credentials=_vertex_credentials(token_wire.url),
gcs_bucket_name=BUCKET,
)
_wait_until_every_worker_serves(gateway, model)
caller_user: Final = scenario.user(user_role="internal_user")
caller_key: Final = scenario.key(user_id=caller_user, models=[model])
object_uri: Final = f"gs://{BUCKET}/litellm-vertex-files/litellm-vertex-files-obj.jsonl"
token_wire.drain()
denied: Final = gateway.request(
"DELETE",
"/vertex_ai/v1/files/" + quote(object_uri, safe=""),
key=caller_key,
params={"model": model},
)
token_requests: Final = token_wire.drain()
assert denied.status_code == 403, f"{denied.text}; token requests={token_requests}"
assert _error_message(denied, 403) == (
"Raw cloud storage file ids can only be deleted by a proxy admin key. "
"Use the LiteLLM managed file id returned when the file was created."
), denied.text
assert token_requests == (), "non-admin raw GCS delete reached the Vertex token endpoint"
def test_proxy_admin_deletes_a_raw_s3_uri_with_one_signed_delete(gateway: Gateway) -> None:
with wire_server(_s3_backend()) as wire, gateway.scenario() as scenario:
model: Final = _bedrock_model(scenario, wire)
_wait_until_every_worker_serves(gateway, model)
object_uri: Final = f"s3://{BUCKET}/litellm-bedrock-files-obj.jsonl"
encoded_path: Final = "/v1/files/" + quote(object_uri, safe="")
deleted: Final = gateway.request("DELETE", encoded_path, params={"model": model})
assert deleted.status_code == 200, deleted.text
parsed: Final = _FileDeleted.model_validate_json(deleted.content)
assert parsed.deleted is True and parsed.id == object_uri, deleted.text
(request,) = _drained_other_than_model_list_probes(wire)
assert request.target == f"/{BUCKET}/litellm-bedrock-files-obj.jsonl", request.target
request_date: Final = request.headers["x-amz-date"]
content_hash: Final = hashlib.sha256(b"").hexdigest()
assert request.headers["x-amz-content-sha256"] == content_hash, request.headers
aws_request: Final = AWSRequest(
method="DELETE",
url=wire.url + request.target,
headers={
"host": request.headers["host"],
"x-amz-date": request_date,
"x-amz-content-sha256": content_hash,
},
)
aws_request.context["timestamp"] = request_date
auth: Final = S3SigV4Auth(Credentials(AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY), "s3", AWS_REGION)
signature: Final = auth.signature(
auth.string_to_sign(aws_request, auth.canonical_request(aws_request)), aws_request
)
expected_authorization: Final = (
f"AWS4-HMAC-SHA256 Credential={AWS_ACCESS_KEY_ID}/{request_date[:8]}/{AWS_REGION}/s3/aws4_request, "
f"SignedHeaders={auth.signed_headers(auth.headers_to_sign(aws_request))}, Signature={signature}"
)
assert request.headers["authorization"] == expected_authorization, request.headers["authorization"]
assert request.method == "DELETE", request.method
def _granting_vertex_token_backend(request: Request) -> Reply:
if request.method == "POST" and request.target == "/_oauth/token":
return _json_reply({"access_token": VERTEX_ACCESS_TOKEN, "expires_in": 3600, "token_type": "Bearer"})
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
def _gcs_backend(request: Request) -> Reply:
if request.method == "DELETE":
return Reply(status=204, body=b"")
return _json_reply({"error": {"message": f"unscripted {request.method} {request.target}"}}, 404)
def _pipe(source: socket.socket, sink: socket.socket) -> None:
with contextlib.suppress(OSError):
for chunk in iter(lambda: source.recv(65536), b""):
sink.sendall(chunk)
with contextlib.suppress(OSError):
sink.shutdown(socket.SHUT_WR)
@dataclass(frozen=True, slots=True)
class _ConnectTunnel:
url: str
authorities: SimpleQueue[str]
@contextmanager
def _gcs_tunnel(destination: Wire) -> Generator[_ConnectTunnel, None, None]:
authorities: Final[SimpleQueue[str]] = SimpleQueue()
destination_port: Final = int(destination.url.rsplit(":", 1)[1])
class Tunnel(socketserver.StreamRequestHandler):
rbufsize = 0
request: socket.socket
def handle(self) -> None:
authority: Final = self.rfile.readline().decode().split()[1]
while self.rfile.readline() not in (b"\r\n", b""):
pass
authorities.put(authority)
if authority != GCS_AUTHORITY:
self.wfile.write(b"HTTP/1.1 403 Forbidden\r\ncontent-length: 0\r\n\r\n")
return
self.wfile.write(b"HTTP/1.1 200 Connection established\r\n\r\n")
self.request.settimeout(10)
with socket.create_connection(("127.0.0.1", destination_port), timeout=10) as upstream:
outbound: Final = threading.Thread(target=_pipe, args=(self.request, upstream))
outbound.start()
_pipe(upstream, self.request)
outbound.join(timeout=12)
with socketserver.ThreadingTCPServer(("127.0.0.1", 0), Tunnel) as server:
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
thread.start()
try:
yield _ConnectTunnel(f"http://127.0.0.1:{server.server_address[1]}", authorities)
finally:
server.shutdown()
thread.join(timeout=6)
def _queued(queue: SimpleQueue[str]) -> tuple[str, ...]:
return tuple(queue.get_nowait() for _ in range(queue.qsize()))
@pytest.mark.timeout(180)
def test_proxy_admin_deletes_a_raw_gcs_uri_with_one_authorized_storage_delete(
gateway: Gateway, tmp_path: Path
) -> None:
cert_file, key_file = write_self_signed_cert(tmp_path, (GCS_HOST,))
object_name: Final = f"litellm-vertex-files/admin-{uuid.uuid4().hex}.jsonl"
object_uri: Final = f"gs://{BUCKET}/{object_name}"
with (
wire_server(_granting_vertex_token_backend) as token_wire,
wire_server(_gcs_backend, tls=server_context(cert_file, key_file)) as storage,
_gcs_tunnel(storage) as tunnel,
owned_proxy(gateway, tmp_path, {"HTTPS_PROXY": tunnel.url, "SSL_VERIFY": "False"}) as candidate,
candidate.scenario() as scenario,
):
model: Final = scenario.model(
model=VERTEX_MODEL,
api_key=None,
vertex_project=VERTEX_PROJECT,
vertex_location=VERTEX_LOCATION,
vertex_credentials=_vertex_credentials(token_wire.url),
gcs_bucket_name=BUCKET,
)
token_wire.drain()
deleted: Final = candidate.request(
"DELETE", "/v1/files/" + quote(object_uri, safe=""), params={"model": model}
)
storage_requests: Final = storage.drain()
assert deleted.status_code == 200, f"{deleted.text}; storage requests={storage_requests}"
parsed: Final = _FileDeleted.model_validate_json(deleted.content)
assert parsed.deleted is True and parsed.id == object_uri, deleted.text
assert [(r.method, r.target) for r in storage_requests] == [
("DELETE", f"/storage/v1/b/{BUCKET}/o/{quote(object_name, safe='')}")
], storage_requests
assert storage_requests[0].headers["authorization"] == f"Bearer {VERTEX_ACCESS_TOKEN}", (
storage_requests[0].headers
)
assert _queued(tunnel.authorities) == (GCS_AUTHORITY,), "raw GCS delete took an unexpected route"