From c177ccf16666b09498c30b7dac62e681a9816f4e Mon Sep 17 00:00:00 2001 From: Yucheng Zhu Date: Thu, 20 Aug 2026 11:57:47 -0700 Subject: [PATCH] fix(proxy): give the scan the metadata bag guardrails actually read, and close its spools The narrowed request metadata was installed under `litellm_metadata` only, but a record is scanned as the chat request it describes, and the guardrails that pick a policy from a request header read `metadata` instead. Noma choosing an application and Aim choosing a user both look there, so the header allowlist added for them did not reach either one and a batch record was still evaluated under the fallback policy. The scan metadata now goes into both bags, which are both stripped and restored, so neither survives into the record that ships. The scan spool was closed on the paths that abort, which are exactly the paths where it is empty, and left open on the one path where it holds the rewritten records. Nothing closed the rewrite output either, where before this feature the uploaded handle belonged to Starlette. The upload now owns both and closes them however it exits. --- .../batch_guardrails.py | 13 ++-- .../openai_files_endpoints/files_endpoints.py | 11 +++ .../test_batch_guardrails.py | 20 ++++- .../test_files_endpoint.py | 78 +++++++++++++++++++ 4 files changed, 114 insertions(+), 8 deletions(-) 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()