From 175991c3752066cee6ddcd7934c73f59e5ee0ed8 Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Thu, 20 Aug 2026 12:14:37 -0700 Subject: [PATCH] fix(proxy): register the scan spool before the rewrite can fail The scan spool was added to the request's cleanup list only after the rewrite returned, so a rewrite that raised, which for a spilled file can be as ordinary as the disk filling up, jumped to the handler with the list still empty and left the scan's own handle open. The rewrite also left its half-written output behind on that path, since nothing owns that handle until it is returned. Both now close. --- .../batch_guardrails.py | 3 ++ .../openai_files_endpoints/files_endpoints.py | 4 +-- .../test_batch_guardrails.py | 32 +++++++++++++++++++ 3 files changed, 37 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index 78aa04c2e65..5c886ca0e9b 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -539,6 +539,9 @@ def rewrite_batch_input_file(file_source: BinaryIO, result: BatchScanResult) -> 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")) wrote_any = True + except BaseException: + output.close() + raise finally: file_source.seek(0) output.seek(0) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 92f176f2f4f..37cfd9d073d 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -535,13 +535,13 @@ async def create_file( ) # Prepare the file data according to FileTypes + if scan_result is not None: + spools.append(scan_result.redactions) upload_source: Final = ( await asyncio.to_thread(rewrite_batch_input_file, file_source, scan_result) if scan_result is not None and scan_result.changes else file_source ) - if scan_result is not None: - spools.append(scan_result.redactions) if upload_source is not file_source: spools.append(upload_source) file_data: Final = (file.filename, upload_source, file.content_type) 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 0e07249f7a7..a05b8ae530c 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 @@ -758,6 +758,38 @@ async def test_the_scan_spool_is_closed_when_the_upload_is_refused(): assert spools and all(handle.closed for handle in spools) +@pytest.mark.asyncio +async def test_the_rewrite_closes_its_own_output_when_it_cannot_finish(): + """A half-written rewrite spool has no owner yet, so it has to clean up after itself.""" + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + + source = _jsonl(_record("a"), _record("b", content="my secret is here")) + result = await _scan_full(source, FakeProxyLogging(_redact_containing("secret"))) + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + def _boom(*args, **kwargs): + raise OSError("no space left on device") + + bg.tempfile.SpooledTemporaryFile = _tracking + original_read = bg._read_spooled + bg._read_spooled = _boom + try: + with pytest.raises(OSError): + rewrite_batch_input_file(source, result) + finally: + bg.tempfile.SpooledTemporaryFile = real + bg._read_spooled = original_read + + assert spools and all(handle.closed for handle in spools) + + @pytest.mark.asyncio async def test_the_scan_spool_is_closed_when_a_record_escapes_the_iterator(): """A raise from inside the read loop bypasses the per-record outcome path entirely."""