diff --git a/litellm/proxy/openai_files_endpoints/batch_guardrails.py b/litellm/proxy/openai_files_endpoints/batch_guardrails.py index f8325ff5d36..78aa04c2e65 100644 --- a/litellm/proxy/openai_files_endpoints/batch_guardrails.py +++ b/litellm/proxy/openai_files_endpoints/batch_guardrails.py @@ -46,6 +46,7 @@ _CUSTOM_ID_LOG_LIMIT: Final = 128 _SUMMARY_LIMIT: Final = 50 _SCAN_METADATA_KEY: Final = "litellm_metadata" +_SCAN_METADATA_BAGS: Final = (_SCAN_METADATA_KEY, "metadata") # Set by pre_call_hook when a guardrail rerouted the request to a different model. _ROUTE_APPLIED_KEY: Final = "sensitive_data_routing_applied" @@ -361,11 +362,13 @@ async def _scan_record( own_injected: Final = MappingProxyType({key: body[key] for key in _INJECTED_KEYS if key in body}) for injected in _INJECTED_KEYS: scan_input.pop(injected, None) - # Deep, and per record: `headers` and `tags` are nested containers shared with the upload - # request and with every other record in the window, and a guardrail that writes into one in - # place would otherwise leak across records and back into the request. The narrowing above - # already removed the values that cannot be copied. - scan_input[_SCAN_METADATA_KEY] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here + # Both bags, because guardrails read whichever one their own route populates and a record + # scanned as chat reaches ones that only ever look at `metadata`; both are injected keys, so + # neither survives into the record that ships. Deep, and per bag per record, because `headers` + # and `tags` are nested containers otherwise shared with the upload request and with every + # other record in the window. The narrowing above already removed what cannot be copied. + for injected in _SCAN_METADATA_BAGS: + scan_input[injected] = copy.deepcopy(dict(scan_metadata)) # mutable-ok: guardrails write here try: # The chain hands back the body it produced, which may be a replacement for the dict it was diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 9dbd62fb7b5..92f176f2f4f 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -371,6 +371,10 @@ async def create_file( ) data: dict = {} + # Spools this request owns. Starlette owns the upload handle; anything the guardrail scan + # opens is ours, and a batch upload that fails after the scan would otherwise hold the + # descriptor and its disk blocks until the collector runs. + spools: Final[list[BinaryIO]] = [] # mutable-ok: filled as the scan opens handles try: # Batch uploads can be gigabytes. Starlette has already spooled the upload # to disk, so stream from that handle instead of reading it into memory. @@ -536,6 +540,10 @@ async def create_file( 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) ## check if model is a loadbalanced model @@ -652,6 +660,9 @@ async def create_file( param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), ) + finally: + for spool in spools: + spool.close() @router.get( 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 758f213b081..0e07249f7a7 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 @@ -140,14 +140,28 @@ async def test_records_own_metadata_is_left_out_of_the_scan_and_the_diff(): seen = [] def _write_bookkeeping(data): - seen.append("metadata" in data) + seen.append(dict(data.get("metadata") or {})) data.setdefault("metadata", {})["applied_guardrails"] = ["g"] - assert await _scan(_jsonl(record), FakeProxyLogging(_write_bookkeeping)) is None - assert seen == [False], "the record's own metadata must not be handed to guardrail dispatch" + assert await _scan(_jsonl(record), FakeProxyLogging(_write_bookkeeping), metadata={"tags": ["t"]}) is None + assert seen == [{"tags": ["t"]}], "dispatch sees the proxy's metadata, never the record's own" assert record["body"]["metadata"] == {"team": "finance"} +@pytest.mark.asyncio +async def test_the_scan_metadata_reaches_guardrails_that_only_read_the_metadata_bag(): + """noma and aim read `metadata["headers"]`; a record scanned as chat must reach them too.""" + seen = [] + + await _scan( + _jsonl(_record("a")), + FakeProxyLogging(lambda d: seen.append((d.get("metadata") or {}).get("headers"))), + metadata={"guardrails": ["g"], "headers": {"x-noma-application-id": "app-1"}}, + ) + + assert seen == [{"x-noma-application-id": "app-1"}] + + @pytest.mark.asyncio async def test_request_metadata_is_narrowed_to_what_guardrails_read(): """An OTel-enabled proxy puts a lock-bearing span here; a per-record copy of it is a crash.""" diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index ca073f9bb10..bf9323cdc6a 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -3680,3 +3680,81 @@ def test_batch_upload_redacts_per_record(monkeypatch, llm_router: Router): finally: app.dependency_overrides.pop(ps.user_api_key_auth, None) ProxyLogging._callback_capabilities_cache.clear() + + +def test_batch_upload_closes_the_spools_it_opened(monkeypatch, llm_router: Router): + """The scan and the rewrite each open a spool; the request owns both and must not leak them.""" + import json as _json + + import litellm + import litellm.proxy.openai_files_endpoints.batch_guardrails as bg + import litellm.proxy.openai_files_endpoints.files_endpoints as fe + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import LitellmUserRoles + from litellm.proxy.utils import ProxyLogging + + class _Redactor(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + for message in data.get("messages") or []: + if "leak" in (message.get("content") or ""): + message["content"] = message["content"].replace("leak", "***") + return data + + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr(litellm, "callbacks", [_Redactor(guardrail_name="g", default_on=True)]) + ProxyLogging._callback_capabilities_cache.clear() + + spools = [] + real = bg.tempfile.SpooledTemporaryFile + + def _tracking(*args, **kwargs): + handle = real(*args, **kwargs) + spools.append(handle) + return handle + + monkeypatch.setattr(bg.tempfile, "SpooledTemporaryFile", _tracking) + + async def fake_route_create_file(**kwargs): + return OpenAIFileObject( + id="dummy-id", + object="file", + bytes=0, + created_at=1234567890, + filename="batch.jsonl", + purpose="batch", + status="uploaded", + ) + + monkeypatch.setattr(fe, "route_create_file", fake_route_create_file) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="test-user" + ) + + def _row(custom_id, content): + return _json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": content}]}, + } + ) + + content = ("\n".join([_row("keep", "fine"), _row("dirty", "please leak this")])).encode() + try: + resp = client.post( + "/v1/files", + files={"file": ("batch.jsonl", content, "application/jsonl")}, + data={"purpose": "batch"}, + headers={"Authorization": "Bearer test-key"}, + ) + assert resp.status_code == 200, resp.text + assert len(spools) == 2, f"expected a scan spool and a rewrite spool, saw {len(spools)}" + assert all(handle.closed for handle in spools), "the request must close every spool it opened" + finally: + app.dependency_overrides.pop(ps.user_api_key_auth, None) + ProxyLogging._callback_capabilities_cache.clear()