diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 42a598f1278..d0977b15a3e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -685,19 +685,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): self, scan_result: Mapping[str, object], is_response: bool = False, - also_hide: str | None = None, ) -> Mapping[str, Mapping[str, object]]: - """Build enhanced error detail with scan information. - - ``also_hide`` names one more scan field to withhold, for the caller that knows - its AIRS verdict carries model-generated content under a key that is normally - caller input. - """ + """Build enhanced error detail with scan information.""" action_type: Final = "Response" if is_response else "Prompt" code_suffix: Final = "_response_blocked" if is_response else "_blocked" - hidden_fields: Final = self._CLIENT_HIDDEN_SCAN_FIELDS.union(() if also_hide is None else (also_hide,)) - category: Final = scan_result.get("category", "unknown") default_msg: Final = f"{action_type} blocked by PANW Prisma AI Security policy (Category: {category})" @@ -717,7 +709,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): **{ key: value for key, value in scan_result.items() - if not key.startswith("_") and key not in hidden_fields + if not key.startswith("_") and key not in self._CLIENT_HIDDEN_SCAN_FIELDS }, "message": error_msg, "type": "guardrail_violation", @@ -1498,17 +1490,12 @@ class PanwPrismaAirsHandler(CustomGuardrail): ): self._set_tool_call_arguments(tool_call, masked_args) else: - # tool_event scans are request-side in the AIRS schema, so AIRS returns - # the model's own tool arguments under prompt_masked_data. On a - # response-side block that is generated content, not caller input, and - # the class-level default only withholds response_masked_data — which is - # empty on this path. Withhold it explicitly so the 400 does not become - # the content channel this branch declined to deliver. - error_detail = self._build_error_detail( - scan_result, - is_response=is_response, - also_hide="prompt_masked_data" if is_response else None, - ) + # Tool calls now go out as ordinary prompt/response text, so a + # response-side scan reports the model's arguments under + # response_masked_data, which _CLIENT_HIDDEN_SCAN_FIELDS already + # withholds. prompt_masked_data is the caller's own input again and + # must keep reaching them -- it is one of the fields LIT-5638 asks for. + error_detail = self._build_error_detail(scan_result, is_response=is_response) raise HTTPException(status_code=400, detail=error_detail) @staticmethod diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py index c156c54624c..9bac8bbc732 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_panw_prisma_airs.py @@ -5974,26 +5974,41 @@ class TestPanwAirsErrorDetailWithheldFields: assert detail["error"]["scan_id"] == "scan-2" -class TestPanwAirsToolCallBlockWithholdsGeneratedArgs: - """A response-side tool-call block must not ship the model's tool arguments. +class TestPanwAirsToolCallBlockMaskedDataRouting: + """A tool-call block must withhold model output and keep caller input. - ``_scan_tool_calls_for_guardrail`` calls AIRS with ``is_response=False`` because - tool_event is request-side in the AIRS schema, so AIRS returns the scanned tool - arguments under ``prompt_masked_data``. When the tool calls being scanned are the - model's own output, that key holds generated content, and the class-level - ``_CLIENT_HIDDEN_SCAN_FIELDS`` default (``response_masked_data``, empty on this - path) does not cover it. + Tool calls are scanned as ordinary prompt/response text, so the side of the scan + decides which key holds what: a response-side scan reports the model's generated + arguments under ``response_masked_data`` (withheld by + ``_CLIENT_HIDDEN_SCAN_FIELDS``), while ``prompt_masked_data`` is the caller's own + input and is one of the fields LIT-5638 asks for. + + Regression guard for the interaction with #37036. That PR withheld + ``prompt_masked_data`` on response-side tool blocks, correctly, while tool calls + still went out as a request-side ``tool_event``. Once this PR routes them by side, + that withholding drops a caller-facing audit field instead. The two PRs merge + without a conflict, so nothing but this test catches it. """ - MASKED_ARGS = '{"to_account": "XXXXXXXXXX", "amount": 5000}' + MODEL_ARGS = '{"to_account": "XXXXXXXXXX", "amount": 5000}' + CALLER_INPUT = "my ssn is XXX-XX-XXXX" - SCAN_RESULT = { + RESPONSE_SIDE_SCAN = { "action": "block", "category": "sensitive_data", "scan_id": "scan-tool-1", "prompt_detected": {"dlp": True}, - "prompt_masked_data": {"data": MASKED_ARGS}, - "response_masked_data": {}, + "response_detected": {"dlp": True}, + "prompt_masked_data": {"data": CALLER_INPUT}, + "response_masked_data": {"data": MODEL_ARGS}, + } + + REQUEST_SIDE_SCAN = { + "action": "block", + "category": "sensitive_data", + "scan_id": "scan-tool-2", + "prompt_detected": {"dlp": True}, + "prompt_masked_data": {"data": MODEL_ARGS}, } @staticmethod @@ -6007,9 +6022,9 @@ class TestPanwAirsToolCallBlockWithholdsGeneratedArgs: ), ) - async def _block(self, handler, is_response): + async def _block(self, handler, is_response, scan_result): with patch.object(handler, "_call_panw_api", new_callable=AsyncMock) as mock_api: - mock_api.return_value = dict(self.SCAN_RESULT) + mock_api.return_value = dict(scan_result) with pytest.raises(HTTPException) as exc_info: await handler._scan_tool_calls_for_guardrail( tool_calls=[self._tool_call()], @@ -6027,27 +6042,34 @@ class TestPanwAirsToolCallBlockWithholdsGeneratedArgs: # The block branch is only reached with masking off; guard the premise. assert handler.mask_response_content is False - exc = await self._block(handler, is_response=True) + exc = await self._block(handler, True, self.RESPONSE_SIDE_SCAN) error = exc.detail["error"] assert exc.status_code == 400 - assert "prompt_masked_data" not in error - assert self.MASKED_ARGS not in str(error) + assert "response_masked_data" not in error + assert self.MODEL_ARGS not in str(error) - # The audit fields LIT-5638 asks for are unaffected. + @pytest.mark.asyncio + async def test_response_side_block_still_returns_caller_input(self): + """The caller's own masked input is an audit field, not model output.""" + handler = make_handler(mask_response_content=False) + + exc = await self._block(handler, True, self.RESPONSE_SIDE_SCAN) + error = exc.detail["error"] + + assert error["prompt_masked_data"] == {"data": self.CALLER_INPUT} assert error["scan_id"] == "scan-tool-1" - assert error["prompt_detected"] == {"dlp": True} @pytest.mark.asyncio async def test_request_side_block_still_returns_masked_tool_args(self): """Caller-supplied tool arguments stay in the verdict — that is the ticket's ask.""" handler = make_handler(mask_request_content=False) - exc = await self._block(handler, is_response=False) + exc = await self._block(handler, False, self.REQUEST_SIDE_SCAN) error = exc.detail["error"] - assert error["prompt_masked_data"] == {"data": self.MASKED_ARGS} - assert error["scan_id"] == "scan-tool-1" + assert error["prompt_masked_data"] == {"data": self.MODEL_ARGS} + assert error["scan_id"] == "scan-tool-2" if __name__ == "__main__":