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."""