From 1f2c681c0b48150a569277d3ebd5e95709dc7c39 Mon Sep 17 00:00:00 2001 From: Abhay Singh Date: Wed, 12 Aug 2026 10:17:09 +0530 Subject: [PATCH] fix(proxy/batch): don't crash an OpenAI batch on a valid-JSON non-object line (#2316) ## Description A single malformed line can abort compression for an entire OpenAI batch upload. `_compress_batch_jsonl` parses each JSONL line and immediately reads the request body: ```python request_obj = json.loads(line) body = request_obj.get("body", {}) messages = body.get("messages", []) ... except json.JSONDecodeError as e: ... compressed_lines.append(line) # keep original on error ``` `json.loads` returns a valid JSON *value*, which isn't necessarily an object. A line like `[1, 2, 3]`, `"hello"`, or `null` parses fine, but `request_obj.get(...)` on a list/str/None raises `AttributeError`. Likewise a request object whose `body` is present but not a dict (`{"body": "..."}`) makes `body.get("messages", ...)` raise. The surrounding `except json.JSONDecodeError` doesn't catch `AttributeError`, so the exception propagates out of `_compress_batch_jsonl` and the whole batch-create request fails. This is the OpenAI batch upload path; the file is user-supplied, so a single stray non-object line takes the batch down instead of just being passed through like the other non-compressible cases already are. ## Fix Guard for non-object shapes and pass them through unchanged: ```python request_obj = json.loads(line) if not isinstance(request_obj, dict): compressed_lines.append(line) total_requests += 1 continue body = request_obj.get("body", {}) if not isinstance(body, dict): body = {} ``` Closes # ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `headroom/proxy/handlers/batch.py`: in `_compress_batch_jsonl`, pass through a non-dict parsed line and coalesce a non-dict `body` to `{}`. - `tests/test_proxy_handlers_batch.py`: new test that array / string / null lines and a non-dict `body` pass through without crashing and are preserved. - `CHANGELOG.md`: Bug Fixes entry. ## Testing - [ ] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [x] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [ ] Manual testing performed ### Test Output ```text $ uvx ruff@0.15.17 check headroom/proxy/handlers/batch.py tests/test_proxy_handlers_batch.py All checks passed! $ uvx mypy@1.20.2 --ignore-missing-imports headroom/proxy/handlers/batch.py Success: no issues found in 1 source file ``` ## Real Behavior Proof - Environment: Windows 11, Python 3.12, `uvx ruff@0.15.17` / `uvx mypy@1.20.2`. A full `pytest` OOM-kills this box (ML stack import), so I reproduced the per-line handling with a dependency-free script and left the full pytest to CI. - Exact command / steps: ran lines `[1,2,3]`, `"hello"`, `null`, and `{"body": "not-a-dict"}` through the OLD (bare `.get`) and NEW (isinstance-guarded) logic, plus a normal request and a `not-json` line as controls. - Observed result: OLD raises `AttributeError` on each non-object line and on the non-dict body; NEW passes the non-object lines through unchanged, coalesces the non-dict body to `{}`, still processes a normal request, and still passes `not-json` through as a JSON-decode error. - Not tested: a live OpenAI batch upload end-to-end; full local `pytest` deferred to CI (OOM). ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Additional Notes The "unit tests pass locally" box is unchecked because the full suite imports the ML stack, which I can't run here. The new test reuses the existing `DummyBatchHandler` / `install_batch_support_modules` harness in `tests/test_proxy_handlers_batch.py` (the same one the neighbouring invalid-line test uses), so it runs under the normal CI pytest job; behaviour is additionally verified by the standalone proof above. --------- Co-authored-by: JerrettDavis --- headroom/proxy/handlers/batch.py | 11 ++++++++++ tests/test_proxy_handlers_batch.py | 34 ++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/headroom/proxy/handlers/batch.py b/headroom/proxy/handlers/batch.py index 3681294a4..91c23f640 100644 --- a/headroom/proxy/handlers/batch.py +++ b/headroom/proxy/handlers/batch.py @@ -1104,7 +1104,18 @@ class BatchHandlerMixin: try: request_obj = json.loads(line) + # A JSONL line can be valid JSON that isn't a request object + # (e.g. a stray array, string, or null). `.get` on it would raise + # AttributeError, which the outer `except json.JSONDecodeError` + # below does not catch — crashing the whole batch. Pass such lines + # through unchanged, like the other non-compressible cases. + if not isinstance(request_obj, dict): + compressed_lines.append(line) + total_requests += 1 + continue body = request_obj.get("body", {}) + if not isinstance(body, dict): + body = {} messages = body.get("messages", []) model = body.get("model", "gpt-4") diff --git a/tests/test_proxy_handlers_batch.py b/tests/test_proxy_handlers_batch.py index 5188b718a..a20a18d1b 100644 --- a/tests/test_proxy_handlers_batch.py +++ b/tests/test_proxy_handlers_batch.py @@ -651,6 +651,40 @@ async def test_compress_batch_jsonl_without_optimization_handles_invalid_lines( } +@pytest.mark.asyncio +async def test_compress_batch_jsonl_handles_non_object_lines( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # A JSONL line that is valid JSON but not a request object (array/string/ + # null), or a request whose `body` isn't a dict, must pass through instead + # of crashing the whole batch (`.get` on a non-dict raises AttributeError, + # which the JSONDecodeError guard does not catch). + install_batch_support_modules(monkeypatch, tokenizer_count=12) + handler = DummyBatchHandler() + content = "\n".join( + [ + json.dumps( + {"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}} + ), + json.dumps([1, 2, 3]), + json.dumps("hello"), + "null", + json.dumps({"body": "not-a-dict"}), + ] + ) + + lines, stats = await handler._compress_batch_jsonl(content, "req-1") + + assert len(lines) == 5 + assert json.loads(lines[1]) == [1, 2, 3] + assert json.loads(lines[2]) == "hello" + assert json.loads(lines[3]) is None + assert json.loads(lines[4]) == {"body": "not-a-dict"} + assert stats["total_requests"] == 5 + # None of these are JSON decode errors, so the error counter stays at 0. + assert stats["errors"] == 0 + + @pytest.mark.asyncio async def test_compress_batch_jsonl_uses_pipeline_and_ccr_injection( monkeypatch: pytest.MonkeyPatch,