From 19c236aaac999c0c321362793ca10b40d1bdff96 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 25 Jun 2026 17:35:15 -0700 Subject: [PATCH 1/4] fix(proxy/client): redact api key from key/info client error messages (#31342) * fix(proxy/client): redact api key from key/info client error messages The keys management client builds GET /key/info?key= and lets the requests HTTPError propagate. str(HTTPError) renders the failing request URL verbatim ("... for url: .../key/info?key=sk-..."), so any caller that logs the exception leaks the full key; the 401 branch leaked the same way through UnauthorizedError(str(orig_exception)) Redact both branches with the existing redact_secrets helper so the secret-bearing query param is scrubbed to ?REDACTED while the status code, reason, and response object are preserved. Server-side responses already mask the key, so this closes the remaining client-side surface * fix: preserve key info unauthorized response --------- Co-authored-by: Cursor Agent (cherry picked from commit 71ee1a852a02787097c025d71f24202fdfc6b8c0) --- litellm/proxy/client/exceptions.py | 20 +++- litellm/proxy/client/keys.py | 9 +- tests/test_litellm/proxy/client/test_keys.py | 96 +++++++++++++++++++- 3 files changed, 118 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/client/exceptions.py b/litellm/proxy/client/exceptions.py index fffd1b78b7e..c4089381e30 100644 --- a/litellm/proxy/client/exceptions.py +++ b/litellm/proxy/client/exceptions.py @@ -2,18 +2,30 @@ from typing import Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + + +def _redact_orig_exception( + orig_exception: Union[requests.exceptions.HTTPError, str], +) -> Union[requests.exceptions.HTTPError, str]: + if isinstance(orig_exception, requests.exceptions.HTTPError): + return requests.exceptions.HTTPError( + redact_string(str(orig_exception)), response=orig_exception.response + ) + return redact_string(str(orig_exception)) + class UnauthorizedError(Exception): """Exception raised when the API returns a 401 Unauthorized response.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) class NotFoundError(Exception): """Exception raised when the API returns a 404 Not Found response or indicates a resource was not found.""" def __init__(self, orig_exception: Union[requests.exceptions.HTTPError, str]): - self.orig_exception = orig_exception - super().__init__(str(orig_exception)) + self.orig_exception = _redact_orig_exception(orig_exception) + super().__init__(str(self.orig_exception)) diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index d8687cbad16..845b49d1581 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Union import requests +from litellm.litellm_core_utils.secret_redaction import redact_string + from .exceptions import UnauthorizedError @@ -314,6 +316,9 @@ class KeysManagementClient: response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: + redacted_message = redact_string(str(e)) if e.response.status_code == 401: - raise UnauthorizedError(e) - raise + raise UnauthorizedError(e) from None + raise requests.exceptions.HTTPError( + redacted_message, response=e.response + ) from None diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 136408e01ac..620daefb39e 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,5 +1,6 @@ import os import sys +import traceback import pytest import requests @@ -11,7 +12,7 @@ sys.path.insert( import responses -from litellm.proxy.client.exceptions import UnauthorizedError +from litellm.proxy.client.exceptions import NotFoundError, UnauthorizedError from litellm.proxy.client.keys import KeysManagementClient @@ -420,3 +421,96 @@ def test_info_server_error(client): ) with pytest.raises(requests.exceptions.HTTPError): client.info(key="test-key") + + +LEAKY_KEY = "sk-1234567890abcdefghijklmnop" + + +def _render_full_traceback(exc: BaseException) -> str: + return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)) + + +@responses.activate +def test_info_not_found_redacts_key_everywhere(client): + """A 404 must not echo the raw key embedded in the request URL. + + Covers str(exc) and the rendered traceback, since the chain through + __cause__ / __context__ is what logging.exception() and the default + excepthook print. + """ + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=404, + json={"error": {"message": "Key not found", "code": "404"}}, + ) + with pytest.raises(requests.exceptions.HTTPError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert exc.response is not None + assert exc.response.status_code == 404 + assert exc.request is not None + # Known residual: the live request URL still carries the key, since the + # response is preserved so callers keep status_code / text. str(exc) and the + # traceback are scrubbed; the URL-borne key is the root issue tracked in + # LIT-4013 (move the lookup key out of the query string server-side). + assert LEAKY_KEY in exc.response.request.url + + +@responses.activate +def test_info_unauthorized_redacts_key_everywhere(client): + """A 401 surfaced as UnauthorizedError must not echo the raw key in the + message, the retained original, or the rendered traceback chain.""" + responses.add( + responses.GET, + f"{client._base_url}/key/info?key={LEAKY_KEY}", + status=401, + json={"error": "Unauthorized"}, + ) + with pytest.raises(UnauthorizedError) as excinfo: + client.info(key=LEAKY_KEY) + + exc = excinfo.value + assert LEAKY_KEY not in str(exc) + assert "REDACTED" in str(exc) + assert LEAKY_KEY not in str(exc.orig_exception) + assert LEAKY_KEY not in _render_full_traceback(exc) + assert exc.__cause__ is None and exc.__suppress_context__ + assert isinstance(exc.orig_exception, requests.exceptions.HTTPError) + assert exc.orig_exception.response is not None + assert exc.orig_exception.response.status_code == 401 + + +def _http_error_with_key(prefix: str, status: int) -> requests.exceptions.HTTPError: + resp = requests.Response() + resp.status_code = status + return requests.exceptions.HTTPError( + f"{prefix} for url: http://x/key/info?key={LEAKY_KEY}", response=resp + ) + + +def test_unauthorized_error_redacts_wrapped_key(): + """UnauthorizedError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = UnauthorizedError( + _http_error_with_key("401 Client Error: Unauthorized", 401) + ) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 401 + + +def test_not_found_error_redacts_wrapped_key(): + """NotFoundError scrubs the key in str(exc) and in the retained + orig_exception, while preserving the response for structured access.""" + wrapped = NotFoundError(_http_error_with_key("404 Client Error: Not Found", 404)) + assert LEAKY_KEY not in str(wrapped) + assert "REDACTED" in str(wrapped) + assert LEAKY_KEY not in str(wrapped.orig_exception) + assert wrapped.orig_exception.response.status_code == 404 From 7b026008b9dd60cacee0d0242f4e3dbe911ca1b5 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 29 Jun 2026 17:31:32 -0700 Subject: [PATCH 2/4] fix(vertex_ai/files): single media upload for batch files to fix 499s on large uploads (#31653) * fix(vertex_ai/files): upload batch files in a single media request to fix 499s on large uploads PR #31036 switched the vertex batch file upload from a single GCS media upload to a chunked resumable session. The resumable path sends the body as many sequential PUTs, each waiting a full round-trip to GCS before the next, so a multi-GB upload accumulates hundreds of round-trips and overruns the client/load-balancer request timeout, surfacing as 499s (client closed connection) on files as small as 500MB. This was a regression from the last-known-good commit, where the upload completed as one continuous request. Revert the batch upload to a single uploadType=media request, but stage the transformed payload to a temp file first so peak memory stays bounded (the goal of the resumable rewrite) without the per-chunk round-trips. The temp file is closed deterministically (TemporaryFile unlinks on close), not left to the GC. The now-unused resumable chunked-upload plumbing is removed. Also swap the per-row transform's stdlib json for orjson (parse + serialize), which is ~4x faster on this hot path; the streaming body now emits compact orjson bytes. The request stays synchronous, so the returned file object is real and POST /v1/batches keeps working immediately against the uploaded object. Tests: single media request carries the whole payload with a real Content-Length (no chunked transfer-encoding); failed upload raises; the staged temp file is closed deterministically; byte-for-byte transform parity. * test(vertex_ai/files): mock single media upload POST instead of removed resumable method test_avertex_batch_prediction patched BaseLLMHTTPHandler._aresumable_chunked_upload, which was removed when the batch jsonl upload moved from a chunked resumable GCS session to a single uploadType=media request. Patch the raw httpx.AsyncClient.post that _astage_and_upload_media issues so the real staging, upload and response transform run while the GCS object response is mocked, and assert the media URL and Content-Type. * fix(vertex_ai/files): forward request timeout to media upload, drop orjson, sort imports Forward the per-request timeout through _stage_and_upload_media / _astage_and_upload_media to the GCS POST. Every other upload branch forwards it; the new media path was dropping it, so a caller-provided timeout was silently ignored (the files path passes 600s by default, but a custom request_timeout would not have reached this upload). Regression test asserts the resolved timeout reaches the request (mutation-verified). Revert the orjson swap in the batch transform: importing orjson at module load in this core-path file broke `import litellm` on environments without orjson (the Windows import test). Back to stdlib json; the upload leg dominates large uploads anyway, so the transform-side win was marginal. Fix import ordering in llm_http_handler.py (I001) introduced by the new imports. * fix(vertex_ai/files): stream batch upload to GCS instead of staging to a temp file Addresses a disk-exhaustion concern: staging the full transformed batch body to a local temp file before the GCS request meant an authenticated user could fill the proxy's temp volume with large concurrent uploads (on top of Starlette's input spool). GCS's simple/media upload accepts chunked transfer-encoding, so stream the transform straight to the single media request instead. Each block is produced on a worker thread (the transform never runs on the event loop) and sent chunked, so the body is neither buffered in memory nor written to disk, and the upload is still one continuous request (no per-chunk round-trips, no 499). Drops the temp-file staging, the tempfile/IO imports, and Content-Length computation. Regression test asserts the upload streams (chunked transfer-encoding, no Content-Length) and creates no temp file; mutation-verified that reintroducing staging fails it. (cherry picked from commit 85840aef513ce93e11d7c203996426c14d4e4582) --- litellm/llms/custom_httpx/llm_http_handler.py | 285 +++++------------ .../llms/vertex_ai/files/transformation.py | 25 +- litellm/types/files.py | 17 +- .../test_openai_batches_and_files.py | 38 ++- .../test_vertex_ai_binary_file_upload.py | 43 +-- .../files/test_vertex_ai_files_streaming.py | 294 ++++++------------ 6 files changed, 217 insertions(+), 485 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 14c3b7736a3..1bcf1b571bd 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -41,7 +41,10 @@ from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.containers.transformation import BaseContainerConfig from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.evals.transformation import BaseEvalsAPIConfig -from litellm.llms.base_llm.files.transformation import BaseFilesConfig +from litellm.llms.base_llm.files.transformation import ( + BaseFilesConfig, + BaseFileUploadStream, +) from litellm.llms.base_llm.google_genai.transformation import ( BaseGoogleGenAIGenerateContentConfig, ) @@ -80,7 +83,7 @@ from litellm.types.containers.main import ( ContainerObject, DeleteContainerResult, ) -from litellm.types.files import TwoStepFileUploadConfig +from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig from litellm.types.integrations.custom_logger import ( AgenticLoopPlan, AgenticLoopRequestPatch, @@ -3232,18 +3235,15 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) - elif ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ): + elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request: + media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"]) try: - upload_response = self._resumable_chunked_upload( + upload_response = self._upload_media( client=sync_httpx_client, - initiate_url=api_base, + url=api_base, base_headers=headers, - config=cast(Dict[str, Any], transformed_request)[ - "resumable_chunked_upload" - ], + body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]), + content_type=media_cfg.get("content_type") or "application/octet-stream", timeout=timeout, ) except Exception as e: @@ -3330,13 +3330,12 @@ class BaseLLMHTTPHandler: input="", api_key="", additional_args={ - # A resumable upload config holds a reference to the (potentially + # A streaming upload config holds a reference to the (potentially # huge) upload payload; logging deep-copies additional_args, so log # a placeholder instead of re-materializing the payload. "complete_input_dict": ( - "" - if isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request + "" + if isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request else transformed_request ), "api_base": api_base, @@ -3415,18 +3414,15 @@ class BaseLLMHTTPHandler: data=presigned_request["data"], timeout=timeout, ) - elif ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ): + elif isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request: + media_cfg = cast(StreamingMediaUploadConfig, transformed_request["streaming_media_upload"]) try: - upload_response = await self._aresumable_chunked_upload( + upload_response = await self._aupload_media( client=async_httpx_client, - initiate_url=api_base, + url=api_base, base_headers=headers, - config=cast(Dict[str, Any], transformed_request)[ - "resumable_chunked_upload" - ], + body_stream=cast(BaseFileUploadStream, media_cfg["body_stream"]), + content_type=media_cfg.get("content_type") or "application/octet-stream", timeout=timeout, ) except Exception as e: @@ -3475,222 +3471,81 @@ class BaseLLMHTTPHandler: litellm_params=litellm_params, ) - # 8 MiB; a 256 KiB multiple, which GCS requires for every non-final chunk. - _RESUMABLE_CHUNK_SIZE = 8 * 1024 * 1024 + # The fine-grained transform stream (one piece per JSONL row) is regrouped + # into blocks of this size before upload, so the request yields a manageable + # number of chunks; never more than one block is buffered. + _MEDIA_UPLOAD_BLOCK_SIZE = 4 * 1024 * 1024 @staticmethod - def _iter_resumable_chunks( - byte_iter: Iterator[bytes], chunk_size: int - ) -> Iterator[bytes]: - """Regroup a byte stream into ``chunk_size`` pieces, yielding a final - partial piece only when it is non-empty. Every full piece is exactly - ``chunk_size`` bytes (kept a 256 KiB multiple for GCS) and never more than - one chunk is buffered. An exactly chunk-aligned stream yields only full - chunks, so the upload finalizes on its last data chunk instead of making - an extra empty request; a 0-byte stream yields nothing and the caller - finalizes with a single empty request. - """ + def _iter_in_blocks(byte_iter: Iterator[bytes], block_size: int) -> Iterator[bytes]: buf = bytearray() for piece in byte_iter: buf.extend(piece) - while len(buf) >= chunk_size: - yield bytes(buf[:chunk_size]) - del buf[:chunk_size] + while len(buf) >= block_size: + yield bytes(buf[:block_size]) + del buf[:block_size] if buf: yield bytes(buf) - @staticmethod - def _resumable_content_range(offset: int, data_len: int, is_final: bool) -> str: - if not is_final: - return f"bytes {offset}-{offset + data_len - 1}/*" - total = offset + data_len - if data_len == 0: - return f"bytes */{total}" - return f"bytes {offset}-{total - 1}/{total}" + def _check_media_upload_response(self, resp: httpx.Response) -> None: + if resp.status_code not in (200, 201): + resp.raise_for_status() + raise ValueError(f"media upload: unexpected status {resp.status_code}") - @staticmethod - def _resumable_request_kwargs( - headers: dict, - content: bytes, - timeout: Optional[Union[float, httpx.Timeout]], - ) -> dict: - kwargs: Dict[str, Any] = {"headers": headers, "content": content} - if timeout is not None: - kwargs["timeout"] = timeout - return kwargs - - def _resumable_chunked_upload( + def _upload_media( self, *, client: HTTPHandler, - initiate_url: str, - base_headers: dict, - config: dict, - timeout: Optional[Union[float, httpx.Timeout]], - ) -> httpx.Response: - """Open a GCS resumable session, then PUT the body in bounded chunks so a - large upload is never held in memory in full.""" - stream = config["body_stream"] - chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) - session_url_header = config.get("session_url_header", "location") - httpx_client = client.client - - init_headers = {**base_headers, **config.get("initiate_headers", {})} - init_req = httpx_client.build_request( - "POST", - initiate_url, - **self._resumable_request_kwargs(init_headers, b"", timeout), - ) - init_resp = httpx_client.send(init_req, follow_redirects=False) - init_resp.read() - if init_resp.status_code not in (200, 201): - init_resp.raise_for_status() - session_url = init_resp.headers.get(session_url_header) - if not session_url: - raise ValueError( - f"resumable upload: no session URL in '{session_url_header}' header" - ) - - offset = 0 - pending: Optional[bytes] = None - for chunk in self._iter_resumable_chunks(stream.iter_bytes(), chunk_size): - if pending is not None: - self._send_resumable_chunk( - httpx_client, - session_url, - base_headers, - pending, - offset, - is_final=False, - timeout=timeout, - ) - offset += len(pending) - pending = chunk - return self._send_resumable_chunk( - httpx_client, - session_url, - base_headers, - pending or b"", - offset, - is_final=True, - timeout=timeout, - ) - - def _send_resumable_chunk( - self, - httpx_client: httpx.Client, url: str, - base_headers: dict, - data: bytes, - offset: int, - *, - is_final: bool, + base_headers: Dict[str, str], + body_stream: BaseFileUploadStream, + content_type: str, timeout: Optional[Union[float, httpx.Timeout]], ) -> httpx.Response: - headers = { - **base_headers, - "Content-Range": self._resumable_content_range(offset, len(data), is_final), + headers = {**base_headers, "Content-Type": content_type} + kwargs: Dict[str, Any] = { + "headers": headers, + "content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE), } - req = httpx_client.build_request( - "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) - ) - resp = httpx_client.send(req, follow_redirects=False) - resp.read() - if resp.status_code not in ((200, 201) if is_final else (308,)): - # 4xx/5xx raise here; the ValueError catches an unexpected success - # status (e.g. a 200 where the protocol expects a 308 between chunks). - resp.raise_for_status() - raise ValueError(f"resumable upload: unexpected status {resp.status_code}") + if timeout is not None: + kwargs["timeout"] = timeout + resp = client.client.post(url, **kwargs) + self._check_media_upload_response(resp) return resp - async def _aresumable_chunked_upload( + async def _aupload_media( self, *, client: AsyncHTTPHandler, - initiate_url: str, - base_headers: dict, - config: dict, - timeout: Optional[Union[float, httpx.Timeout]], - ) -> httpx.Response: - stream = config["body_stream"] - chunk_size = config.get("chunk_size", self._RESUMABLE_CHUNK_SIZE) - session_url_header = config.get("session_url_header", "location") - httpx_client = client.client - - init_headers = {**base_headers, **config.get("initiate_headers", {})} - init_req = httpx_client.build_request( - "POST", - initiate_url, - **self._resumable_request_kwargs(init_headers, b"", timeout), - ) - init_resp = await httpx_client.send(init_req, follow_redirects=False) - await init_resp.aread() - if init_resp.status_code not in (200, 201): - init_resp.raise_for_status() - session_url = init_resp.headers.get(session_url_header) - if not session_url: - raise ValueError( - f"resumable upload: no session URL in '{session_url_header}' header" - ) - - offset = 0 - pending: Optional[bytes] = None - # Producing each chunk runs the synchronous per-row transform for that - # chunk's worth of rows. Pull it off the event loop thread so a large - # upload does not block other concurrent requests between PUTs. - chunk_iter = self._iter_resumable_chunks(stream.iter_bytes(), chunk_size) - done = object() - while True: - chunk = await asyncio.to_thread(next, chunk_iter, done) - if chunk is done: - break - if pending is not None: - await self._asend_resumable_chunk( - httpx_client, - session_url, - base_headers, - pending, - offset, - is_final=False, - timeout=timeout, - ) - offset += len(pending) - pending = chunk - return await self._asend_resumable_chunk( - httpx_client, - session_url, - base_headers, - pending or b"", - offset, - is_final=True, - timeout=timeout, - ) - - async def _asend_resumable_chunk( - self, - httpx_client: httpx.AsyncClient, url: str, - base_headers: dict, - data: bytes, - offset: int, - *, - is_final: bool, + base_headers: Dict[str, str], + body_stream: BaseFileUploadStream, + content_type: str, timeout: Optional[Union[float, httpx.Timeout]], ) -> httpx.Response: - headers = { - **base_headers, - "Content-Range": self._resumable_content_range(offset, len(data), is_final), - } - req = httpx_client.build_request( - "PUT", url, **self._resumable_request_kwargs(headers, data, timeout) - ) - resp = await httpx_client.send(req, follow_redirects=False) + """Stream the transformed body straight to a single media upload. Each + block is produced on a worker thread (the transform never runs on the + event loop) and sent with chunked transfer-encoding, so the body is + neither buffered in memory nor staged to disk, and the upload is one + continuous request rather than the many sequential round-trips of the + resumable path that overran client/LB timeouts.""" + headers = {**base_headers, "Content-Type": content_type} + block_iter = iter(self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE)) + done = object() + + async def _abody() -> AsyncIterator[bytes]: + while True: + block = await asyncio.to_thread(next, block_iter, done) + if block is done: + break + yield cast(bytes, block) + + kwargs: Dict[str, Any] = {"headers": headers, "content": _abody()} + if timeout is not None: + kwargs["timeout"] = timeout + resp = await client.client.post(url, **kwargs) await resp.aread() - if resp.status_code not in ((200, 201) if is_final else (308,)): - # 4xx/5xx raise here; the ValueError catches an unexpected success - # status (e.g. a 200 where the protocol expects a 308 between chunks). - resp.raise_for_status() - raise ValueError(f"resumable upload: unexpected status {resp.status_code}") + self._check_media_upload_response(resp) return resp def create_batch( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index d5164d8c1c2..5abaedae5a8 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -60,7 +60,7 @@ from litellm.types.llms.openai import ( OpenAIFileObject, PathLike, ) -from litellm.types.files import ResumableChunkedUploadConfig +from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.vertex_ai import GcsBucketResponse from litellm.types.utils import LlmProviders, ModelResponse @@ -389,21 +389,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - _, content_type = extract_file_metadata(file_data) object_name = self.get_object_name(file_data, purpose) if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name = encode_gcs_object_name_for_url(object_name) - # Batch jsonl is streamed via a resumable session (bounded memory on - # large uploads); everything else is a single simple-media upload. - upload_type = ( - "resumable" - if FilesAPIUtils.is_batch_jsonl_request( - create_file_data=data, content_type=content_type - ) - else "media" - ) - endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType={upload_type}&name={encoded_object_name}" + endpoint = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") @@ -455,8 +445,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ 2 Cases: 1. Handle basic file upload - 2. Handle batch file upload (.jsonl), streamed to a GCS resumable - session so large uploads stay memory-bounded. + 2. Handle batch file upload (.jsonl), staged to a temp file and uploaded + in a single media request so large uploads stay memory-bounded without + the per-chunk round-trips of a resumable session. """ file_data = create_file_data.get("file") if file_data is None: @@ -468,14 +459,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): content_type=content_type, ): return { - "resumable_chunked_upload": ResumableChunkedUploadConfig( + "streaming_media_upload": StreamingMediaUploadConfig( body_stream=_OpenAIToVertexBatchUploadStream( file_data, self._map_openai_to_vertex_params, ), - initiate_headers={ - "X-Upload-Content-Type": "application/json", - }, + content_type="application/json", ) } diff --git a/litellm/types/files.py b/litellm/types/files.py index 1b2d7e30f1f..97c4dbb3abc 100644 --- a/litellm/types/files.py +++ b/litellm/types/files.py @@ -323,19 +323,18 @@ class TwoStepFileUploadConfig(TypedDict, total=False): upload_url_key: str -class ResumableChunkedUploadConfig(TypedDict, total=False): - """Drives a memory-bounded resumable upload (GCS JSON API). +class StreamingMediaUploadConfig(TypedDict, total=False): + """Drives a memory-bounded single-request upload (GCS simple/media upload). - The handler POSTs to the upload URL to open a session, reads the session URI - from ``session_url_header``, then PUTs ``body_stream`` to that URI in - ``chunk_size``-byte chunks (a 256 KiB multiple) using Content-Range, so the - payload is never buffered in full and the transfer is resumable. + The handler stages ``body_stream`` to a temp file off the event loop (so peak + memory stays bounded), then PUTs/POSTs it in one request with a known + Content-Length. Unlike a resumable chunked upload this incurs no per-chunk + round-trips, so a multi-GB upload finishes in one continuous transfer instead + of hundreds of sequential PUTs that overrun client/LB timeouts. ``body_stream`` is a ``BaseFileUploadStream``; it is typed ``Any`` here to avoid importing the llms layer into types. """ body_stream: Required[Any] - chunk_size: int - session_url_header: str - initiate_headers: Dict[str, str] + content_type: str diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 8a2d5f33805..0a49b3d77d1 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -513,25 +513,26 @@ async def test_avertex_batch_prediction(monkeypatch): mock_response.status_code = 200 return mock_response - # Batch jsonl file creation now streams to a GCS resumable session via - # _aresumable_chunked_upload (httpx send), not AsyncHTTPHandler.post, so mock - # that entry point to return the GCS object response. The resumable protocol - # itself is covered in test_vertex_ai_files_streaming.py. - mock_upload_response = httpx.Response( - 200, - json=mock_file_response, - request=httpx.Request("PUT", "https://storage.googleapis.com/upload"), - ) + # Batch jsonl creation now stages the body to a temp file and issues a single + # uploadType=media POST against the raw httpx.AsyncClient (client.client) inside + # _astage_and_upload_media, not AsyncHTTPHandler.post. Patch that raw POST so the + # real staging/upload + response transform run while the GCS object response is + # mocked; AsyncHTTPHandler.post still handles the batch-prediction call. with ( patch( "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", side_effect=mock_side_effect, - ) as mock_global_post, - patch( - "litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler._aresumable_chunked_upload", - new_callable=AsyncMock, - return_value=mock_upload_response, ), + patch.object( + httpx.AsyncClient, + "post", + new_callable=AsyncMock, + return_value=httpx.Response( + 200, + json=mock_file_response, + request=httpx.Request("POST", "https://storage.googleapis.com/upload"), + ), + ) as mock_gcs_upload, ): litellm.set_verbose = True litellm._turn_on_debug() @@ -552,6 +553,15 @@ async def test_avertex_batch_prediction(monkeypatch): == "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/5f7b99ad-9203-4430-98bf-3b45451af4cb" ) + mock_gcs_upload.assert_awaited_once() + upload_url = str(mock_gcs_upload.call_args.args[0]) + assert "uploadType=media" in upload_url + assert "/b/litellm-local/o" in upload_url + assert ( + mock_gcs_upload.call_args.kwargs["headers"]["Content-Type"] + == "application/json" + ) + # Create batch create_batch_response = await litellm.acreate_batch( completion_window="24h", diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py index 122518d4acb..d2ee9d7d659 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_binary_file_upload.py @@ -37,9 +37,7 @@ class TestVertexAIBinaryFileUpload: # Create mock PDF binary data (with non-UTF-8 bytes) # PDF files start with %PDF- and contain binary data mock_pdf_content = b"%PDF-1.4\n%\xc4\xe5\xf2\xe5\xeb\xa7\xf3\xa0\xd0\xc4\xc6\n" - mock_pdf_content += ( - b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 - ) # Add more binary data + mock_pdf_content += b"\x00\x01\x02\x03\xff\xfe\xfd" * 100 # Add more binary data # Create file object file_obj = io.BytesIO(mock_pdf_content) @@ -60,14 +58,12 @@ class TestVertexAIBinaryFileUpload: ) # Verify the transformation returns bytes (not string) - assert isinstance( - transformed_request, bytes - ), f"Expected bytes for binary file, got {type(transformed_request)}" + assert isinstance(transformed_request, bytes), ( + f"Expected bytes for binary file, got {type(transformed_request)}" + ) # Verify the bytes match the original content - assert ( - transformed_request == mock_pdf_content - ), "Transformed request should preserve binary content exactly" + assert transformed_request == mock_pdf_content, "Transformed request should preserve binary content exactly" # Verify that the bytes contain non-UTF-8 characters # This should raise UnicodeDecodeError if we try to decode @@ -132,16 +128,14 @@ class TestVertexAIBinaryFileUpload: pytest.fail(f"httpx should accept bytes in data parameter: {e}") # Document the expected behavior - assert isinstance( - mock_binary_data, bytes - ), "Binary file data should remain as bytes" + assert isinstance(mock_binary_data, bytes), "Binary file data should remain as bytes" @pytest.mark.asyncio - async def test_jsonl_file_upload_returns_resumable_stream(self): + async def test_jsonl_file_upload_returns_streaming_body(self): """ - Test that JSONL batch files are transformed into a resumable-upload config + Test that JSONL batch files are transformed into a streaming-media config carrying a streaming body (not a buffered bytes payload), so the handler - can stream the upload to GCS in bounded chunks. + can stage the upload to a temp file and send it in one media request. """ # Create mock JSONL content mock_jsonl_content = ( @@ -164,16 +158,13 @@ class TestVertexAIBinaryFileUpload: litellm_params={}, ) - assert ( - isinstance(transformed_request, dict) - and "resumable_chunked_upload" in transformed_request - ), f"Expected a resumable upload config for JSONL, got {type(transformed_request)}" + assert isinstance(transformed_request, dict) and "streaming_media_upload" in transformed_request, ( + f"Expected a streaming media upload config for JSONL, got {type(transformed_request)}" + ) - stream = transformed_request["resumable_chunked_upload"]["body_stream"] + stream = transformed_request["streaming_media_upload"]["body_stream"] decoded = json.loads(b"".join(stream.iter_bytes()).decode("utf-8")) - assert ( - "request" in decoded - ), "JSONL transform must wrap each row in {'request': ...}" + assert "request" in decoded, "JSONL transform must wrap each row in {'request': ...}" @pytest.mark.asyncio async def test_mixed_file_types_in_sequence(self): @@ -214,7 +205,7 @@ class TestVertexAIBinaryFileUpload: optional_params={}, litellm_params={}, ) - assert isinstance(result2, dict) and "resumable_chunked_upload" in result2 + assert isinstance(result2, dict) and "streaming_media_upload" in result2 # Test 3: Upload another binary file binary_content2 = b"\xc4\xe5\xf2\xe5\xeb" @@ -264,7 +255,5 @@ class TestVertexAIBinaryFileUpload: }, } - assert ( - expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" - ) + assert expected_behavior["binary_files"]["encoding"] == "none - preserve raw bytes" assert expected_behavior["text_files"]["encoding"] == "UTF-8" diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index cd556c48b6b..2e3280c0ed1 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -20,6 +20,7 @@ replaced by a list-based pipeline: import gc import io import json +import tempfile import time import tracemalloc @@ -41,15 +42,15 @@ from litellm.llms.vertex_ai.files.transformation import ( from litellm.types.llms.openai import CreateFileRequest -def _resumable_stream(transformed) -> BaseFileUploadStream: - """Pull the streaming body out of a resumable-upload transform result.""" - return transformed["resumable_chunked_upload"]["body_stream"] +def _upload_stream(transformed) -> BaseFileUploadStream: + """Pull the streaming body out of the upload transform result.""" + return transformed["streaming_media_upload"]["body_stream"] def _join_upload_body(transformed) -> bytes: """Materialize a transform result's upload body for byte-level assertions.""" - if isinstance(transformed, dict) and "resumable_chunked_upload" in transformed: - return b"".join(_resumable_stream(transformed).iter_bytes()) + if isinstance(transformed, dict) and "streaming_media_upload" in transformed: + return b"".join(_upload_stream(transformed).iter_bytes()) if isinstance(transformed, BaseFileUploadStream): return b"".join(transformed.iter_bytes()) if isinstance(transformed, str): @@ -83,17 +84,13 @@ def _reference_vertex_jsonl_string(cfg: VertexAIFilesConfig, content: str) -> st transform, so the streaming path can be checked against it for parity.""" entries = [json.loads(line) for line in content.splitlines() if line.strip()] return "\n".join( - json.dumps( - _openai_batch_jsonl_entry_to_vertex_wrapped_request( - entry, cfg._map_openai_to_vertex_params - ) - ) + json.dumps(_openai_batch_jsonl_entry_to_vertex_wrapped_request(entry, cfg._map_openai_to_vertex_params)) for entry in entries ) class TestStreamingOutputParity: - def test_transform_create_file_request_returns_resumable_stream_parity(self): + def test_transform_create_file_request_returns_streaming_body_parity(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(300) request: CreateFileRequest = { @@ -105,14 +102,12 @@ class TestStreamingOutputParity: model="", create_file_data=request, optional_params={}, litellm_params={} ) - # A batch upload must be a resumable-upload config carrying a streaming - # body, so the handler can chunk it; a buffered bytes/str return would - # defeat the OOM fix. - assert isinstance(out, dict) and "resumable_chunked_upload" in out - assert isinstance(_resumable_stream(out), BaseFileUploadStream) - assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string( - cfg, raw.decode("utf-8") - ) + # A batch upload must be a streaming-media config carrying a streaming + # body, so the handler can stream it to GCS; a buffered bytes/str return + # would defeat the OOM fix. + assert isinstance(out, dict) and "streaming_media_upload" in out + assert isinstance(_upload_stream(out), BaseFileUploadStream) + assert _join_upload_body(out).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) class TestFileLikeInputNotPartiallyConsumed: @@ -215,9 +210,7 @@ class TestStreamingLineIterator: def seek(self, *args): raise io.UnsupportedOperation("not seekable") - handle = _NonSeekable( - b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n' - ) + handle = _NonSeekable(b'{"custom_id": "request-0"}\n{"custom_id": "request-1"}\n') with pytest.raises(ValueError, match="seekable"): list(_iter_openai_jsonl_lines(handle)) @@ -235,13 +228,8 @@ class TestGetObjectNameLazyParse: cfg = VertexAIFilesConfig() # Tail rows are deliberately not valid JSON. Parsing the whole payload # would raise here; a first-row-only parse must not. - raw = ( - b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\n' - b"garbage line that is not json\n" - ) - object_name = cfg.get_object_name( - ("batch.jsonl", raw, "application/jsonl"), purpose="batch" - ) + raw = b'{"custom_id": "r-0", "body": {"model": "gemini-2.5-flash"}}\ngarbage line that is not json\n' + object_name = cfg.get_object_name(("batch.jsonl", raw, "application/jsonl"), purpose="batch") assert "gemini-2.5-flash" in object_name @@ -278,15 +266,11 @@ class TestStreamingPeakMemory: def drain_stream(): # Consume the upload body one row at a time, as the chunked uploader # does, without accumulating it. - for _ in _OpenAIToVertexBatchUploadStream( - raw, cfg._map_openai_to_vertex_params - ).iter_bytes(): + for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes(): pass streaming_peak = self._measure(drain_stream) - list_peak = self._measure( - lambda: _reference_vertex_jsonl_string(cfg, content_str) - ) + list_peak = self._measure(lambda: _reference_vertex_jsonl_string(cfg, content_str)) # Core guard: the lazily consumed streaming body peaks well under a list # pipeline that materializes every transformed row. Building full @@ -305,9 +289,7 @@ class TestStreamingPeakMemory: # first-row parse should allocate only a small fraction of the payload; # parsing every row would blow past this bound. peak = self._measure(lambda: cfg.get_object_name(file_data, purpose="batch")) - assert ( - peak / len(raw) < 2.0 - ), "get_object_name should not copy the whole payload" + assert peak / len(raw) < 2.0, "get_object_name should not copy the whole payload" class TestPathSourcedStreaming: @@ -341,12 +323,10 @@ class TestPathSourcedStreaming: litellm_params={"gcs_bucket_name": "test-bucket"}, data=data, ) - assert "uploadType=resumable" in url + assert "uploadType=media" in url - out = cfg.transform_create_file_request( - model="", create_file_data=data, optional_params={}, litellm_params={} - ) - assert isinstance(out, dict) and "resumable_chunked_upload" in out + out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) + assert isinstance(out, dict) and "streaming_media_upload" in out body = _join_upload_body(out).decode("utf-8") assert body == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) lines = body.splitlines() @@ -371,7 +351,7 @@ class TestPathSourcedStreaming: out = cfg.transform_create_file_request( model="", create_file_data=data, optional_params={}, litellm_params={} ) - for _ in _resumable_stream(out).iter_bytes(): + for _ in _upload_stream(out).iter_bytes(): pass # drain without accumulating gc.collect() @@ -384,20 +364,15 @@ class TestPathSourcedStreaming: # Streaming from disk must not materialize the payload. Reading the whole # file into bytes (the pre-fix path) would push peak past the file size. - assert peak < len(raw) * 0.3, ( - f"peak {peak} not bounded vs payload {len(raw)} " - f"(ratio {peak / len(raw):.2f})" - ) + assert peak < len(raw) * 0.3, f"peak {peak} not bounded vs payload {len(raw)} (ratio {peak / len(raw):.2f})" def test_path_source_stream_is_reiterable(self, tmp_path): cfg = VertexAIFilesConfig() path, _ = self._write_jsonl(tmp_path, 50) data = self._batch_request(path) - out = cfg.transform_create_file_request( - model="", create_file_data=data, optional_params={}, litellm_params={} - ) - stream = _resumable_stream(out) + out = cfg.transform_create_file_request(model="", create_file_data=data, optional_params={}, litellm_params={}) + stream = _upload_stream(out) first = b"".join(stream.iter_bytes()) second = b"".join(stream.iter_bytes()) assert first == second and len(first) > 0 @@ -436,25 +411,20 @@ def _logging_obj() -> Logging: ) -def _gcs_resumable_mock(session_url: str, final_status: int = 200): - """A fake GCS resumable endpoint: POST opens a session (URI in Location), - each PUT appends and returns 308 until the final chunk returns 200/201.""" - state = {"received": bytearray(), "ranges": [], "methods": [], "urls": []} +def _gcs_media_mock(status: int = 200): + """A fake GCS simple-media endpoint: one request carries the whole object; + capture the body and headers and return the object resource.""" + state = {"received": bytearray(), "methods": [], "urls": [], "headers": [], "timeouts": []} async def handler(request: httpx.Request) -> httpx.Response: state["methods"].append(request.method) state["urls"].append(str(request.url)) - if request.method == "POST": - return httpx.Response(200, headers={"location": session_url}) - body = await request.aread() - content_range = request.headers["content-range"] - state["ranges"].append(content_range) - state["received"].extend(body) - if content_range.rsplit("/", 1)[-1] == "*": - return httpx.Response( - 308, headers={"range": f"bytes=0-{len(state['received']) - 1}"} - ) - return httpx.Response(final_status, json=_GCS_OBJECT_JSON) + state["headers"].append(dict(request.headers)) + # httpx records the resolved per-request timeout here, so the test can + # assert the caller's timeout was forwarded rather than the client default. + state["timeouts"].append(request.extensions.get("timeout")) + state["received"].extend(await request.aread()) + return httpx.Response(status, json=_GCS_OBJECT_JSON) return handler, state @@ -465,8 +435,8 @@ def _async_handler_with(mock) -> AsyncHTTPHandler: return handler -class TestResumableUploadUrl: - def test_batch_jsonl_uses_resumable_upload_type(self): +class TestUploadUrl: + def test_batch_jsonl_uses_media_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "application/jsonl"), @@ -480,29 +450,12 @@ class TestResumableUploadUrl: litellm_params={"gcs_bucket_name": "test-bucket"}, data=request, ) - assert "uploadType=resumable" in url - assert "uploadType=media" not in url + # A single media upload is one continuous transfer (no per-chunk + # round-trips), which is what keeps large uploads under client/LB timeouts. + assert "uploadType=media" in url + assert "uploadType=resumable" not in url - def test_batch_text_plain_uses_resumable_upload_type(self): - # Clients often label a .jsonl batch upload as text/plain; it must still - # take the streaming/resumable path, not the buffered media path. - cfg = VertexAIFilesConfig() - request: CreateFileRequest = { - "file": ("batch.jsonl", _make_openai_jsonl_bytes(3), "text/plain"), - "purpose": "batch", - } - url = cfg.get_complete_file_url( - api_base=None, - api_key=None, - model="", - optional_params={}, - litellm_params={"gcs_bucket_name": "test-bucket"}, - data=request, - ) - assert "uploadType=resumable" in url - assert "uploadType=media" not in url - - def test_binary_upload_stays_simple_media(self): + def test_binary_upload_uses_media_upload_type(self): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("doc.pdf", b"%PDF-1.4 binary", "application/pdf"), @@ -520,14 +473,12 @@ class TestResumableUploadUrl: assert "uploadType=resumable" not in url -class TestResumableStreamBody: +class TestUploadStreamBody: def test_stream_matches_legacy_pipeline(self): cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(120) stream = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params) - assert b"".join(stream.iter_bytes()).decode( - "utf-8" - ) == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) + assert b"".join(stream.iter_bytes()).decode("utf-8") == _reference_vertex_jsonl_string(cfg, raw.decode("utf-8")) def test_stream_is_reiterable_for_retries(self): # A one-shot generator would make a transport retry upload an empty body; @@ -545,59 +496,19 @@ class TestResumableStreamBody: # an empty body silently. cfg = VertexAIFilesConfig() raw = _make_openai_jsonl_bytes(40) - stream = _OpenAIToVertexBatchUploadStream( - io.BytesIO(raw), cfg._map_openai_to_vertex_params - ) + stream = _OpenAIToVertexBatchUploadStream(io.BytesIO(raw), cfg._map_openai_to_vertex_params) first = b"".join(stream.iter_bytes()) second = b"".join(stream.iter_bytes()) assert first == second and len(first) > 0 -class TestResumableChunking: - def test_intermediate_chunks_are_exactly_chunk_size(self): - pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 10]), 4)) - assert pieces == [b"xxxx", b"xxxx", b"xx"] - - def test_exact_multiple_yields_no_trailing_empty(self): - # An exactly chunk-aligned stream yields only full chunks; the upload - # finalizes on the last data chunk instead of an extra empty request. - pieces = list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([b"x" * 8]), 4)) - assert pieces == [b"xxxx", b"xxxx"] - - def test_empty_stream_yields_nothing(self): - # A 0-byte stream yields no chunks; the caller finalizes with one empty - # request (bytes */0). - assert list(BaseLLMHTTPHandler._iter_resumable_chunks(iter([]), 4)) == [] - - def test_default_chunk_size_is_256kib_multiple(self): - assert BaseLLMHTTPHandler._RESUMABLE_CHUNK_SIZE % (256 * 1024) == 0 - - def test_content_range_intermediate_uses_star_total(self): - assert ( - BaseLLMHTTPHandler._resumable_content_range(0, 4096, is_final=False) - == "bytes 0-4095/*" - ) - - def test_content_range_final_uses_real_total(self): - assert ( - BaseLLMHTTPHandler._resumable_content_range(8192, 100, is_final=True) - == "bytes 8192-8291/8292" - ) - - def test_content_range_empty_finalize(self): - assert ( - BaseLLMHTTPHandler._resumable_content_range(8192, 0, is_final=True) - == "bytes */8192" - ) - - @pytest.mark.asyncio -class TestResumableUploadProtocol: - """End-to-end against a faked GCS resumable endpoint. These are the tests - that fail if the handler buffers the whole body, drops bytes, mislabels a - Content-Range, follows the 308 instead of continuing, or skips finalize.""" +class TestStreamingMediaUpload: + """End-to-end against a faked GCS media endpoint. These fail if the handler + buffers the payload in memory, drops bytes, omits Content-Length (which would + flip httpx to chunked transfer-encoding), or makes more than one request.""" - async def _run(self, raw: bytes, chunk_size: int, final_status: int = 200): + async def _run(self, raw: bytes, status: int = 200, timeout=None): cfg = VertexAIFilesConfig() request: CreateFileRequest = { "file": ("batch.jsonl", raw, "application/jsonl"), @@ -614,11 +525,8 @@ class TestResumableUploadProtocol: transformed = cfg.transform_create_file_request( model="", create_file_data=request, optional_params={}, litellm_params={} ) - transformed["resumable_chunked_upload"]["chunk_size"] = chunk_size expected = _join_upload_body(transformed) - - session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" - mock, state = _gcs_resumable_mock(session_url, final_status=final_status) + mock, state = _gcs_media_mock(status=status) response = await BaseLLMHTTPHandler().async_create_file( transformed_request=transformed, litellm_params={}, @@ -627,70 +535,52 @@ class TestResumableUploadProtocol: api_base=api_base, logging_obj=_logging_obj(), client=_async_handler_with(mock), - timeout=None, + timeout=timeout, ) - return expected, state, response, session_url, api_base + return expected, state, response - async def test_streams_in_chunks_and_reassembles(self): + async def test_single_request_carries_whole_payload(self): raw = _make_openai_jsonl_bytes(300) - chunk_size = 4096 - expected, state, response, session_url, api_base = await self._run( - raw, chunk_size - ) + expected, state, response = await self._run(raw) - # One session-open POST, then a sequence of chunk PUTs. - assert state["methods"][0] == "POST" - assert set(state["methods"][1:]) == {"PUT"} - assert state["methods"].count("PUT") >= 2, "payload must span multiple chunks" + # Exactly one request (the single media upload), and it lands on the + # media endpoint, not a resumable session. + assert state["methods"] == ["POST"] + assert "uploadType=media" in state["urls"][0] - # POST opens a resumable session; every chunk goes to the session URI. - assert "uploadType=resumable" in state["urls"][0] - assert all(u == session_url for u in state["urls"][1:]) - - # Every non-final chunk is exactly chunk_size with an unknown-total range; - # the final chunk carries the real total. - intermediate = state["ranges"][:-1] - for index, content_range in enumerate(intermediate): - assert ( - content_range - == f"bytes {index * chunk_size}-{(index + 1) * chunk_size - 1}/*" - ) - total = len(expected) - last_offset = len(intermediate) * chunk_size - if last_offset == total: # payload landed on a chunk boundary - assert state["ranges"][-1] == f"bytes */{total}" - else: - assert state["ranges"][-1] == f"bytes {last_offset}-{total - 1}/{total}" - - # The bytes GCS received are exactly the transformed batch payload. + # The body is streamed with chunked transfer-encoding and no + # Content-Length, which is what proves it is neither buffered in memory + # nor staged to a temp file (the disk-exhaustion guard) before sending. + headers = state["headers"][0] + assert headers.get("transfer-encoding") == "chunked" + assert "content-length" not in headers + # httpx reassembles the chunked body; GCS receives exactly the transform. assert bytes(state["received"]) == expected assert response.object == "file" - async def test_exact_multiple_finalizes_on_last_data_chunk(self): - # A body that is an exact multiple of the chunk size finalizes on its - # last data chunk (bytes (TOTAL-chunk)-(TOTAL-1)/TOTAL), with no extra - # empty finalize request. - chunk_size = 256 - total = chunk_size * 3 - stream = _FixedBytesStream(b"a" * total) - config = {"body_stream": stream, "chunk_size": chunk_size} - session_url = "https://storage.googleapis.com/upload/sess?upload_id=SID" - mock, state = _gcs_resumable_mock(session_url) - - response = await BaseLLMHTTPHandler()._aresumable_chunked_upload( - client=_async_handler_with(mock), - initiate_url="https://storage.googleapis.com/upload?uploadType=resumable", - base_headers={"Authorization": "Bearer x"}, - config=config, - timeout=None, - ) - - assert state["ranges"][-1] == f"bytes {total - chunk_size}-{total - 1}/{total}" - assert "*" not in state["ranges"][-1] - assert bytes(state["received"]) == b"a" * total - assert response.status_code == 200 - - async def test_failed_chunk_raises(self): + async def test_failed_upload_raises(self): raw = _make_openai_jsonl_bytes(80) with pytest.raises(Exception): - await self._run(raw, chunk_size=4096, final_status=403) + await self._run(raw, status=403) + + async def test_request_timeout_is_forwarded(self): + # The caller's per-request timeout must reach the GCS upload; every other + # upload branch forwards it. httpx records the resolved timeout in + # request.extensions["timeout"]; a dropped timeout would show the client + # default instead of the value passed here. + raw = _make_openai_jsonl_bytes(20) + _, state, _ = await self._run(raw, timeout=httpx.Timeout(137.0)) + forwarded = state["timeouts"][0] + assert forwarded is not None + assert forwarded.get("read") == 137.0 and forwarded.get("write") == 137.0 + + async def test_upload_does_not_stage_to_disk(self, monkeypatch): + # Disk-exhaustion guard: the transformed body must stream to GCS, never be + # written to a temp file first. If any tempfile is created during the + # upload, an attacker could fill the proxy's temp volume with large + # concurrent uploads. + created = [] + real_tempfile = tempfile.TemporaryFile + monkeypatch.setattr(tempfile, "TemporaryFile", lambda *a, **k: (created.append(1), real_tempfile(*a, **k))[1]) + await self._run(_make_openai_jsonl_bytes(50)) + assert created == [] From c0b71895bb0787773cb7f58d78599766177c4b52 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 29 Jun 2026 17:53:10 -0700 Subject: [PATCH 3/4] =?UTF-8?q?bump:=20version=201.89.4=20=E2=86=92=201.89?= =?UTF-8?q?.5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 9eb07d8f108..1a9845bd636 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.89.4" +version = "1.89.5" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.14" @@ -261,7 +261,7 @@ source-exclude = [ profile = "black" [tool.commitizen] -version = "1.89.4" +version = "1.89.5" version_files = [ "pyproject.toml:^version", ] From c6268eecce9352880f67310bb808a98b4e3384cb Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 29 Jun 2026 17:53:21 -0700 Subject: [PATCH 4/4] chore: refresh uv.lock for 1.89.5 --- uv.lock | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/uv.lock b/uv.lock index a8d563931e2..faaa5e5e741 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-06-22T01:17:31.93796Z" +exclude-newer = "2026-06-27T00:53:21.064066Z" exclude-newer-span = "P3D" [manifest] @@ -3294,7 +3294,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.89.4" +version = "1.89.5" source = { editable = "." } dependencies = [ { name = "aiohttp" },