From 01a32a3d0788b3f71d9a120d3ad4fd7d5abc895f Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 21 Aug 2026 11:15:59 -0700 Subject: [PATCH] fix(proxy): read batch records the same way the upload validation does (#37776) * fix(proxy): read batch records the same way the upload validation does The upload validation parses each JSONL line as bytes, where the json module sniffs the encoding itself and accepts a leading byte order mark or a lone surrogate. The guardrail scan that runs immediately after decoded each line to text first, which is stricter, so a file the validation had just accepted could fail the scan. A `.jsonl` written by any of the editors that emit a BOM, which includes PowerShell's Out-File and classic Notepad, uploaded fine until a pre_call guardrail was configured and then returned 500 with a decode error and no indication of which line or why. The scan now parses the same bytes the validation did, and an untouched record is copied through as the bytes it arrived as rather than re-encoded. A numeric custom_id was reported as null. The spec asks for a string, but callers do send numbers, and null leaves the one field a caller reconciles on empty for exactly the records that need it. * fix(proxy): read the load-balancing record the same way, so a byte order mark keeps its routing The first record is parsed to pick a deployment when batch load balancing is on, and it was decoded to text before parsing, which rejects a leading byte order mark. The lookup returns None on any parse failure, so such a file silently lost its routing and went to the default provider rather than the configured one. That was already reachable for an upload no guardrail changed, since the original bytes are passed straight through, and preserving the mark through a rewrite widens it. Parsed as bytes now, like the validation and the scan. * fix(proxy): find the routing record past a blank first line The upload validation and the guardrail scan both skip blank lines, but deployment selection read only the first physical line, so a file starting with a blank line lost its routing model and went to the default provider rather than the configured one. It now skips blanks the way the other two readers do, reading lazily so a large file is not read past its first record. * fix(proxy): do not crash deployment selection on a record whose body is not an object The upload validation checks that a record has a `body`, not that it is an object, so a record can carry a string or a list there. Deployment selection called `.get` on it unconditionally and raised, returning 500. That was already reachable for a plain file, and reading past a byte order mark or a blank first line widened it to files that previously fell through to the default provider instead. A record whose body names no readable model now resolves to no model, which is the same answer the default-provider branch already handled. * fix(proxy): keep a custom_id that cannot be encoded from failing the whole upload A record identifier is echoed back in the create response. JSON parses a lone surrogate happily but it cannot be encoded again, so a file the upload validation accepts returned 500 from the response renderer rather than a report. Unencodable characters are replaced, which leaves every ordinary identifier untouched and keeps a pathological one reconcilable. This predates the reader change; reading past a byte order mark only altered which error the same file produced first. * fix(proxy): treat a url the parser rejects as one we do not recognize Resolving a record's call type from its url runs the url through urlsplit, which raises on a few malformed authorities such as an unclosed bracket. That happens before the try that wraps the guardrail call, so it escaped the scan and returned 500 on a file the upload validation had just accepted. An unreadable url is simply one we cannot recognize, which the body-shape fallback already handles, so the record is still scanned rather than lost. Reachable on staging today for a proxy running any guardrail. Enabling the scan for a proxy that runs only a content-enforcing CustomLogger widens it to that configuration too, which is why it is fixed here rather than left. --- .../batch_guardrails.py | 50 +++++--- .../openai_files_endpoints/files_endpoints.py | 28 +++-- .../test_batch_guardrails.py | 115 ++++++++++++++++++ 3 files changed, 171 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index 5c886ca0e9b..53d51db2b7f 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -260,18 +260,24 @@ def _describe(custom_id: str | None) -> str: return f" (custom_id {safe})" -def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, str]]: - """Yield every non-blank line with its 1-based number, so both passes number records alike.""" +def _iter_lines(source: BinaryIO) -> Iterator[tuple[int, bytes]]: + """ + Yield every non-blank line with its 1-based number, so both passes number records alike. + + Bytes, not text. The upload validation immediately before this parses each line as bytes, + where the json module sniffs the encoding itself and accepts a leading byte order mark or a + lone surrogate. Decoding to `str` first is stricter than that, so a file written by any of + the editors that emit a BOM would pass validation and then fail the scan. + """ for line_number, raw_line in enumerate(source, start=1): - text = raw_line.decode("utf-8") - if text.strip(): - yield line_number, text + if raw_line.strip(): + yield line_number, raw_line def _iter_records(source: BinaryIO) -> Iterator[_ParsedRecord]: """Yield one record per line, relying on the upload validation that already ran.""" - for line_number, text in _iter_lines(source): - yield _ParsedRecord(line_number=line_number, payload=json.loads(text)) + for line_number, raw_line in _iter_lines(source): + yield _ParsedRecord(line_number=line_number, payload=json.loads(raw_line)) def _call_type_from_url(url: str) -> CallTypesLiteral | None: @@ -282,7 +288,13 @@ def _call_type_from_url(url: str) -> CallTypesLiteral | None: ``/v1/responses`` in full would fall through to its body, where ``input`` reads as an embedding and the record gets scanned as the wrong call type rather than the right one. """ - path: Final = urlsplit(url).path.split("?")[0].rstrip("/") + try: + path: Final = urlsplit(url).path.split("?")[0].rstrip("/") + except ValueError: + # urlsplit rejects a few malformed authorities outright, and the validation that ran + # before this only checks the key is present. An unreadable url is one we do not + # recognize, which is what falling back to the body shape already handles. + return None call_types: Final = get_call_types_for_route(path) if call_types is None: return None @@ -308,8 +320,18 @@ def _scannable_call_type(url: object, body: Mapping[str, object]) -> CallTypesLi def _custom_id_of(payload: Mapping[str, object]) -> str | None: + """ + The record's identifier, rendered as text. + + The batch spec asks for a string, but callers do send numbers, and reporting those as null + would leave the one field a caller reconciles on empty for exactly the records it needs. + """ custom_id: Final = payload.get("custom_id") - return custom_id if isinstance(custom_id, str) else None + if isinstance(custom_id, str): + # A lone surrogate parses out of the file but cannot be encoded back out, and this value + # is echoed in the response, so rendering it would fail the whole upload with a 500. + return custom_id.encode("utf-8", "replace").decode("utf-8") + return str(custom_id) if isinstance(custom_id, (int, float)) and not isinstance(custom_id, bool) else None def _fingerprint(body: Mapping[str, object], keys: frozenset[str]) -> str: @@ -507,9 +529,9 @@ async def scan_batch_input_file( ) -def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> str: +def _read_spooled(redactions: BinaryIO, change: RecordRedacted) -> bytes: redactions.seek(change.offset) - return redactions.read(change.length).decode("utf-8") + return redactions.read(change.length) def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> BinaryIO: @@ -532,12 +554,12 @@ def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> ) wrote_any = False # rebind-ok: tracks whether a separator is needed try: - for line_number, text in _iter_lines(file_source): + for line_number, raw_line in _iter_lines(file_source): if line_number in dropped: continue change = redacted.get(line_number) - line = text.rstrip("\n") if change is None else _read_spooled(result.redactions, change) - output.write((("\n" if wrote_any else "") + line).encode("utf-8")) + line = raw_line.rstrip(b"\n") if change is None else _read_spooled(result.redactions, change) + output.write(b"\n" + line if wrote_any else line) wrote_any = True except BaseException: output.close() diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 37cfd9d073d..813ce9630a5 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -145,25 +145,37 @@ async def _scan_batch_upload( def get_first_json_object(file_source: bytes | BinaryIO) -> dict | None: + """ + The first record, used to pick a deployment when batch load balancing is on. + + Read the way the upload validation reads it, since a file it accepted must not lose its + routing here: blank lines are not records and are skipped, and the line is parsed as bytes so + the json module sniffs the encoding rather than rejecting a leading byte order mark. Either + difference makes this return None, which silently sends the batch to the default provider. + """ try: if isinstance(file_source, (bytes, bytearray)): - newline: Final = file_source.find(b"\n") - raw: Final = file_source if newline == -1 else file_source[:newline] - first_line = raw.decode("utf-8") + first_record: bytes | None = next((line for line in file_source.splitlines() if line.strip()), None) else: + # lazily, so a batch file that can be gigabytes is not read past its first record file_source.seek(0) - first_line = file_source.readline().decode("utf-8") + first_record = next((line for line in file_source if line.strip()), None) file_source.seek(0) - return json.loads(first_line.strip()) + return None if first_record is None else json.loads(first_record.strip()) except (json.JSONDecodeError, UnicodeDecodeError, OSError, ValueError): return None def get_model_from_json_obj(json_object: dict) -> str | None: - body: Final = json_object.get("body", {}) or {} - model: Final = body.get("model") + """ + The model a record names, or None when it does not name one readably. - return model + The upload validation only checks that `body` is present, not that it is an object, so a + record can carry a string there and reach this. Returning None sends the upload down the + default-provider branch, which is what a record with no resolvable model already did. + """ + body: Final = json_object.get("body") + return body.get("model") if isinstance(body, dict) else None async def _deprecated_loadbalanced_create_file( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py index 6ce7af1e2ee..a06b1110aa5 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_batch_guardrails.py @@ -363,6 +363,121 @@ async def test_an_absolute_url_resolves_by_path_not_by_body_shape(url, expected_ assert logging_obj.seen[0][0] == expected_call_type +@pytest.mark.asyncio +@pytest.mark.parametrize( + "prefix, label", + [(b"\xef\xbb\xbf", "utf-8 BOM"), (b"", "plain")], + ids=["utf8_bom", "plain"], +) +async def test_a_file_the_upload_validation_accepts_is_a_file_the_scan_can_read(prefix, label): + """The validator parses each line as bytes, which tolerates a BOM; the scan must match it.""" + from litellm.proxy.openai_files_endpoints.batch_file_validation import check_batch_file_upload + + payload = prefix + (json.dumps(_record("a")) + "\n").encode() + assert check_batch_file_upload("in.jsonl", io.BytesIO(payload), None) is None, f"{label} rejected upfront" + + logging_obj = FakeProxyLogging() + assert await _scan(io.BytesIO(payload), logging_obj) is None + assert logging_obj.seen, f"{label} was never scanned" + + +@pytest.mark.asyncio +async def test_a_bom_file_is_rewritten_without_losing_the_untouched_records(): + source = io.BytesIO(b"\xef\xbb\xbf" + ("\n".join( + json.dumps(r) for r in (_record("keep"), _record("dirty", content="my secret is here")) + ) + "\n").encode()) + + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + rewritten = rewrite_batch_input_file(source, result).read().decode("utf-8-sig") + + rows = [json.loads(line) for line in rewritten.splitlines()] + assert [row["custom_id"] for row in rows] == ["keep", "dirty"] + assert rows[1]["body"]["messages"][0]["content"] == "my *** is here" + + +@pytest.mark.parametrize( + "prefix", + [b"", b"\xef\xbb\xbf", b"\n", b"\n\xef\xbb\xbf", b" \n"], + ids=["plain", "utf8_bom", "leading_blank", "blank_then_bom", "whitespace_line"], +) +def test_load_balancing_finds_the_routing_record_in_any_file_the_upload_accepts(prefix): + """A file whose routing model cannot be read is silently sent to the default provider.""" + from litellm.proxy.openai_files_endpoints.batch_file_validation import check_batch_file_upload + from litellm.proxy.openai_files_endpoints.files_endpoints import get_first_json_object + + payload = prefix + (json.dumps(_record("a")) + "\n").encode() + assert check_batch_file_upload("in.jsonl", io.BytesIO(payload), None) is None, "rejected upfront" + + assert get_first_json_object(io.BytesIO(payload))["body"]["model"] == "gpt-4o-mini" + assert get_first_json_object(payload)["body"]["model"] == "gpt-4o-mini" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("url", ["http://[", "http://[::1", "https://["], ids=["open_bracket", "unclosed_v6", "https_bracket"]) +async def test_a_malformed_url_does_not_escape_the_scan(url): + """Validation only checks the url key is present, and urlsplit rejects some authorities.""" + record = {**_record("m"), "url": url} + + result = await _scan_full(_jsonl(record), FakeProxyLogging()) + + assert result.changes == () + assert result.scanned_records == 1, "the record should still be scanned by its body shape" + + +@pytest.mark.parametrize( + "custom_id, expected", + [("req-1", "req-1"), ("caf\u00e9-42", "caf\u00e9-42"), ("a\ud800b", "a?b")], + ids=["ascii", "unicode", "lone_surrogate"], +) +def test_a_reported_custom_id_can_always_be_rendered(custom_id, expected): + """The id is echoed in the response; one that cannot be encoded back out would 500 the upload.""" + from litellm.proxy.openai_files_endpoints.batch_guardrails import _custom_id_of + + rendered = _custom_id_of({"custom_id": custom_id}) + + assert rendered == expected + assert json.dumps({"custom_id": rendered}, ensure_ascii=False).encode("utf-8") + + +@pytest.mark.parametrize( + "body", + ["summarize this", ["a"], None, 12345], + ids=["string", "list", "null", "number"], +) +def test_a_record_whose_body_is_not_an_object_does_not_crash_deployment_selection(body): + """Validation only checks that `body` is present, so a record can carry anything there.""" + from litellm.proxy.openai_files_endpoints.batch_file_validation import check_batch_file_upload + from litellm.proxy.openai_files_endpoints.files_endpoints import ( + get_first_json_object, + get_model_from_json_obj, + ) + + record = {"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": body} + payload = b"\xef\xbb\xbf" + (json.dumps(record) + "\n").encode() + assert check_batch_file_upload("in.jsonl", io.BytesIO(payload), None) is None, "rejected upfront" + + found = get_first_json_object(io.BytesIO(payload)) + assert get_model_from_json_obj(json_object=found) is None + + +@pytest.mark.parametrize("payload", [b"", b"\n\n\n"], ids=["empty", "blanks_only"]) +def test_load_balancing_returns_none_when_there_is_no_record(payload): + from litellm.proxy.openai_files_endpoints.files_endpoints import get_first_json_object + + assert get_first_json_object(io.BytesIO(payload)) is None + assert get_first_json_object(payload) is None + + +@pytest.mark.asyncio +async def test_a_numeric_custom_id_is_still_reported(): + """The spec asks for a string, but callers send numbers, and null would break reconciliation.""" + record = {**_record("x", content="tripwire"), "custom_id": 12345} + + result = await _scan_full(_jsonl(record), FakeProxyLogging(_blocking("tripwire"))) + + assert result.changes == (RecordDropped(line_number=1, custom_id="12345", guardrail="block-guard"),) + + @pytest.mark.asyncio async def test_query_string_on_a_known_url_does_not_change_the_call_type(): """The body carries `messages`, so only stripping the query string can yield aembedding."""