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,