fix(vertex_ai): drop stale transfer headers when GCS serves an encoded file body

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-17 00:06:43 +00:00
parent bddb64ddc5
commit 5ddc96e560
3 changed files with 40 additions and 3 deletions

View file

@ -317,6 +317,25 @@ async def _aiter_bytes_then_close(response: httpx.Response, *, chunk_size: int)
await response.aclose()
_DECODED_BODY_STALE_HEADERS: Final[frozenset[str]] = frozenset({"content-encoding", "content-length"})
def _decoded_body_headers(response: httpx.Response) -> httpx.Headers:
"""
`aiter_bytes` yields the decoded body, so the upstream transfer headers only
describe the bytes on the wire when no content-encoding was applied.
"""
if response.headers.get("content-encoding", "identity").lower() == "identity":
return response.headers
return httpx.Headers(
[
(name, value)
for name, value in response.headers.multi_items()
if name.lower() not in _DECODED_BODY_STALE_HEADERS
]
)
def _collect_ws_project_quota_callbacks() -> tuple[ProjectQuotaCallback, ...]:
"""Duck-type discover proxy hooks exposing per-frame project ITPM/OTPM
enforcement, so the Responses WebSocket loop can charge every
@ -5312,7 +5331,7 @@ class BaseLLMHTTPHandler:
return await provider_config.transform_file_content_stream(
stream_iterator=_aiter_bytes_then_close(response, chunk_size=chunk_size),
headers=response.headers,
headers=_decoded_body_headers(response),
request_url=str(response.request.url),
logging_obj=logging_obj,
litellm_params=litellm_params,

View file

@ -304,8 +304,7 @@ async def _peek_first_jsonl_line(
buffered: bytes = b"" # rebind-ok: accumulates the prefix read while looking for the first newline
async for chunk in chunks:
buffered = buffered + chunk
*complete_lines, _partial = buffered.split(_JSONL_NEWLINE)
first_line = _first_non_empty_jsonl_line(complete_lines)
first_line = _first_non_empty_jsonl_line(buffered.split(_JSONL_NEWLINE)[:-1])
if first_line is not None:
return first_line, buffered
if len(buffered) > peek_limit_bytes:

View file

@ -23,6 +23,7 @@ replaced by a list-based pipeline:
import asyncio
import gc
import gzip
import io
import json
import tempfile
@ -725,6 +726,24 @@ class TestFileContentStreaming:
assert state["served"] < len(raw_chunks)
assert state["closed"] is False
async def test_gzip_encoded_object_is_decoded_without_stale_transfer_headers(self):
raw = b'{"line": 1}\n{"line": 2}\n' * 200
encoded = gzip.compress(raw)
upstream = {
"content-type": "application/octet-stream",
"content-encoding": "gzip",
"content-length": str(len(encoded)),
}
result, state = await self._open([encoded[i : i + 64] for i in range(0, len(encoded), 64)], upstream)
streamed = b"".join([chunk async for chunk in result.stream_iterator])
assert streamed == raw
assert result.headers["content-type"] == "application/octet-stream"
assert "content-encoding" not in result.headers
assert "content-length" not in result.headers
assert state["closed"] is True
async def test_vertex_batch_output_is_transformed_row_by_row(self):
rows = [_vertex_batch_output_row(f"request-{i}", f"answer {i}") for i in range(30)]
raw = b"\n".join(rows) + b"\n"