mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(guardrails): return the full PANW AIRS scan response on blocked requests (#37036)
* fix(guardrails): return the full PANW AIRS scan response on blocked requests The blocked-request error detail was assembled from a hardcoded allowlist, so audit fields like prompt_detection_details, prompt_masked_data, source, transaction_id and session_id never reached the client even though AIRS returned them. Resolves LIT-5638 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(guardrails): drop redundant comment in AIRS error detail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(panw_prisma_airs): withhold response_masked_data from the blocked-response error The full AIRS passthrough also reached the response-side block path, where response_masked_data carries the model's own generation. That branch is only reached when mask_response_content is False, so the operator had explicitly declined to deliver that text, and the error body handed it back anyway. Withhold response_masked_data from the client-visible detail. prompt_masked_data stays: it is the caller's own input and one of the fields the ticket asks for. Every other AIRS field, including prompt_detection_details, source, transaction_id and session_id, is unchanged. * fix(panw_prisma_airs): withhold generated tool args from response-side blocks _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 _CLIENT_HIDDEN_SCAN_FIELDS default (response_masked_data, empty on this path) does not cover it. With the default mask_response_content=False the block branch then shipped the model's masked tool arguments in the 400 -- the same content channel this PR closed for response_masked_data. _build_error_detail takes an extra_hidden_fields argument so the withholding stays in one place, and the tool-call block branch passes prompt_masked_data when is_response is True. Request-side blocks are unchanged and still carry prompt_masked_data, which is what LIT-5638 asks for. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * style(panw_prisma_airs): apply ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Yucheng Zhu <yucheng@berri.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
540caa6574
commit
d7d10be063
2 changed files with 266 additions and 23 deletions
|
|
@ -71,6 +71,14 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
|
||||
_PROVIDER_NAME = "panw_prisma_airs"
|
||||
|
||||
#: AIRS fields withheld from the client-visible error detail.
|
||||
#: ``response_masked_data`` is the model's own generation. The block branch that builds
|
||||
#: this detail is only reached when ``mask_response_content`` is False, so echoing it
|
||||
#: back would hand the caller exactly the text the operator declined to deliver.
|
||||
#: ``prompt_masked_data`` is deliberately NOT withheld: it is the caller's own input,
|
||||
#: and it is one of the fields the ticket asks for.
|
||||
_CLIENT_HIDDEN_SCAN_FIELDS: Final = frozenset({"response_masked_data"})
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: str,
|
||||
|
|
@ -632,12 +640,21 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
choice.message.function_call.arguments = masked_text
|
||||
|
||||
def _build_error_detail(
|
||||
self, scan_result: Mapping[str, object], is_response: bool = False
|
||||
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."""
|
||||
"""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.
|
||||
"""
|
||||
action_type: Final = "Response" if is_response else "Prompt"
|
||||
code_suffix: Final = "_response_blocked" if is_response else "_blocked"
|
||||
detection_key: Final = "response_detected" if is_response else "prompt_detected"
|
||||
|
||||
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})"
|
||||
|
|
@ -653,8 +670,13 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
},
|
||||
)
|
||||
|
||||
error_detail: Final[dict[str, dict[str, object]]] = {
|
||||
return {
|
||||
"error": {
|
||||
**{
|
||||
key: value
|
||||
for key, value in scan_result.items()
|
||||
if not key.startswith("_") and key not in hidden_fields
|
||||
},
|
||||
"message": error_msg,
|
||||
"type": "guardrail_violation",
|
||||
"code": f"panw_prisma_airs{code_suffix}",
|
||||
|
|
@ -663,24 +685,6 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
}
|
||||
}
|
||||
|
||||
# Add optional fields if present
|
||||
optional_fields: Final = [
|
||||
"scan_id",
|
||||
"report_id",
|
||||
"profile_name",
|
||||
"profile_id",
|
||||
"tr_id",
|
||||
]
|
||||
for field in optional_fields:
|
||||
if scan_result.get(field):
|
||||
error_detail["error"][field] = scan_result[field]
|
||||
|
||||
# Add detection details
|
||||
if scan_result.get(detection_key):
|
||||
error_detail["error"][detection_key] = scan_result[detection_key]
|
||||
|
||||
return error_detail
|
||||
|
||||
def _record_scan_id(self, request_data: dict[str, Any], scan_result: Mapping[str, object]) -> None:
|
||||
"""Surface the AIRS scan id on the response, so allowed calls are auditable too."""
|
||||
scan_id: Final = scan_result.get("scan_id")
|
||||
|
|
@ -1481,7 +1485,17 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
|||
):
|
||||
self._set_tool_call_arguments(tool_call, masked_text)
|
||||
else:
|
||||
error_detail = self._build_error_detail(scan_result, is_response=is_response)
|
||||
# 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,
|
||||
)
|
||||
raise HTTPException(status_code=400, detail=error_detail)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
|
|
@ -5647,6 +5647,235 @@ class TestPanwAirsScanIdExposure:
|
|||
|
||||
assert "guardrail_scan_ids" in _UNTRUSTED_METADATA_CONTROL_FIELDS
|
||||
assert "guardrail_scan_ids" in _UNTRUSTED_ROOT_CONTROL_FIELDS
|
||||
class TestPanwAirsBlockedErrorDetailPassthrough:
|
||||
"""Regression tests for the full AIRS scan response on blocks.
|
||||
|
||||
Before the fix, the error detail was built from a hardcoded allowlist
|
||||
(scan_id, report_id, profile_name, profile_id, tr_id, prompt/response_detected),
|
||||
so audit-relevant fields such as prompt_detection_details, prompt_masked_data,
|
||||
source, transaction_id and session_id never reached the client.
|
||||
"""
|
||||
|
||||
_FULL_BLOCK_RESPONSE = {
|
||||
"action": "block",
|
||||
"category": "malicious",
|
||||
"scan_id": "b2f0a4be-1f6f-4f9a-9f3d-4b6a9d8b1c0e",
|
||||
"report_id": "R0000000000000000000",
|
||||
"tr_id": "test-call-id",
|
||||
"profile_id": "6f5c9f6e-2d0b-4d3f-8a1e-9b7c5d4e3f2a",
|
||||
"profile_name": "test_profile",
|
||||
"source": "prisma_airs",
|
||||
"transaction_id": "4b8c1e2f-5a6d-4c3b-9e8f-1a2b3c4d5e6f",
|
||||
"session_id": "3a2b1c0d-9e8f-4a7b-8c6d-5e4f3a2b1c0d",
|
||||
"timeout": False,
|
||||
"errors": [],
|
||||
"prompt_detected": {"dlp": True, "injection": False, "url_cats": False},
|
||||
"prompt_detection_details": {
|
||||
"dlp_report": {
|
||||
"dlp_report_id": "1234567890",
|
||||
"dlp_profile_name": "Sensitive Content",
|
||||
"data_pattern_rule1_verdict": "MATCHED",
|
||||
}
|
||||
},
|
||||
"prompt_masked_data": {"data": "my ssn is XXX-XX-XXXX"},
|
||||
"response_detected": {"dlp": False, "url_cats": False},
|
||||
"response_detection_details": {},
|
||||
"response_masked_data": {},
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_response", [False, True])
|
||||
async def test_block_returns_every_airs_field(
|
||||
self, base_handler, user_api_key_dict, safe_prompt_data, is_response
|
||||
):
|
||||
response = ModelResponse(
|
||||
id="test_id",
|
||||
choices=[
|
||||
Choices(index=0, message=Message(role="assistant", content="Test response")),
|
||||
],
|
||||
model="gpt-3.5-turbo",
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
base_handler, "_call_panw_api", return_value=copy.deepcopy(self._FULL_BLOCK_RESPONSE)
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
if is_response:
|
||||
await base_handler.async_post_call_success_hook(
|
||||
data=safe_prompt_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
)
|
||||
else:
|
||||
await base_handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=None,
|
||||
data=safe_prompt_data,
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
error = exc_info.value.detail["error"]
|
||||
for field, value in self._FULL_BLOCK_RESPONSE.items():
|
||||
if field == "category":
|
||||
continue
|
||||
if field in PanwPrismaAirsHandler._CLIENT_HIDDEN_SCAN_FIELDS:
|
||||
# Withheld on purpose, covered by TestPanwAirsErrorDetailWithheldFields
|
||||
continue
|
||||
assert error[field] == value, f"{field} missing or altered in blocked-request error"
|
||||
|
||||
assert error["category"] == "malicious"
|
||||
assert error["type"] == "guardrail_violation"
|
||||
assert error["guardrail"] == "test_panw_airs"
|
||||
assert error["code"] == ("panw_prisma_airs_response_blocked" if is_response else "panw_prisma_airs_blocked")
|
||||
assert "PANW Prisma AI Security policy" in error["message"]
|
||||
|
||||
def test_internal_control_flags_are_not_leaked(self, base_handler):
|
||||
detail = base_handler._build_error_detail(
|
||||
{
|
||||
"action": "block",
|
||||
"category": "malicious",
|
||||
"scan_id": "scan-1",
|
||||
"_always_block": True,
|
||||
"_is_transient": True,
|
||||
}
|
||||
)
|
||||
|
||||
assert "_always_block" not in detail["error"]
|
||||
assert "_is_transient" not in detail["error"]
|
||||
assert detail["error"]["scan_id"] == "scan-1"
|
||||
|
||||
|
||||
class TestPanwAirsErrorDetailWithheldFields:
|
||||
"""The blocked-request passthrough must not become a content channel.
|
||||
|
||||
``response_masked_data`` is the model's own generation. The block branch is only
|
||||
reached when ``mask_response_content`` is False, so echoing it back would hand the
|
||||
caller exactly the text the operator declined to deliver. ``error`` is AIRS's own
|
||||
message about the operator's Strata Cloud Manager profile configuration.
|
||||
|
||||
``prompt_masked_data`` is deliberately NOT withheld by default: it is the caller's
|
||||
own input, and it is one of the fields LIT-5638 asks for. The one exception is the
|
||||
response-side tool-call path, covered by
|
||||
``TestPanwAirsToolCallBlockWithholdsGeneratedArgs`` below — tool_event scans are
|
||||
request-side in the AIRS schema, so there the key holds model output instead.
|
||||
"""
|
||||
|
||||
@pytest.mark.parametrize("is_response", [False, True])
|
||||
def test_response_masked_data_never_reaches_client(self, base_handler, is_response):
|
||||
detail = base_handler._build_error_detail(
|
||||
{
|
||||
"action": "block",
|
||||
"category": "sensitive_data",
|
||||
"scan_id": "scan-1",
|
||||
"response_detected": {"dlp": True},
|
||||
"response_masked_data": {"data": "routing number XXXXXXXXXX"},
|
||||
"prompt_masked_data": {"data": "my ssn is XXX-XX-XXXX"},
|
||||
"prompt_detection_details": {"dlp_report": {"dlp_report_id": "1"}},
|
||||
},
|
||||
is_response=is_response,
|
||||
)
|
||||
error = detail["error"]
|
||||
|
||||
assert "response_masked_data" not in error
|
||||
assert "routing number" not in str(error)
|
||||
|
||||
# The audit fields LIT-5638 asks for still come through untouched.
|
||||
assert error["scan_id"] == "scan-1"
|
||||
assert error["response_detected"] == {"dlp": True}
|
||||
assert error["prompt_masked_data"] == {"data": "my ssn is XXX-XX-XXXX"}
|
||||
assert error["prompt_detection_details"] == {"dlp_report": {"dlp_report_id": "1"}}
|
||||
|
||||
def test_upstream_airs_error_field_still_passes_through(self, base_handler):
|
||||
"""A 2xx AIRS body can carry its own ``error`` (see _call_panw_api's
|
||||
profile-misconfiguration branch, which only logs and then blocks). It is
|
||||
diagnostic rather than content, so it stays in the passthrough."""
|
||||
detail = base_handler._build_error_detail(
|
||||
{
|
||||
"action": "block",
|
||||
"category": "malicious",
|
||||
"scan_id": "scan-2",
|
||||
"error": "profile not found",
|
||||
}
|
||||
)
|
||||
|
||||
assert detail["error"]["error"] == "profile not found"
|
||||
assert detail["error"]["scan_id"] == "scan-2"
|
||||
|
||||
|
||||
class TestPanwAirsToolCallBlockWithholdsGeneratedArgs:
|
||||
"""A response-side tool-call block must not ship the model's tool arguments.
|
||||
|
||||
``_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.
|
||||
"""
|
||||
|
||||
MASKED_ARGS = '{"to_account": "XXXXXXXXXX", "amount": 5000}'
|
||||
|
||||
SCAN_RESULT = {
|
||||
"action": "block",
|
||||
"category": "sensitive_data",
|
||||
"scan_id": "scan-tool-1",
|
||||
"prompt_detected": {"dlp": True},
|
||||
"prompt_masked_data": {"data": MASKED_ARGS},
|
||||
"response_masked_data": {},
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _tool_call():
|
||||
return ChatCompletionMessageToolCall(
|
||||
id="call_1",
|
||||
type="function",
|
||||
function=Function(
|
||||
name="transfer_funds",
|
||||
arguments='{"to_account": "ACME-VENDOR-001", "amount": 5000}',
|
||||
),
|
||||
)
|
||||
|
||||
async def _block(self, handler, is_response):
|
||||
with patch.object(handler, "_call_panw_api", new_callable=AsyncMock) as mock_api:
|
||||
mock_api.return_value = dict(self.SCAN_RESULT)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await handler._scan_tool_calls_for_guardrail(
|
||||
tool_calls=[self._tool_call()],
|
||||
is_response=is_response,
|
||||
metadata={},
|
||||
call_id="test-call-id",
|
||||
request_data={"metadata": {}},
|
||||
start_time=datetime.now(),
|
||||
)
|
||||
return exc_info.value
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_side_block_withholds_generated_tool_args(self):
|
||||
handler = make_handler(mask_response_content=False)
|
||||
# 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)
|
||||
error = exc.detail["error"]
|
||||
|
||||
assert exc.status_code == 400
|
||||
assert "prompt_masked_data" not in error
|
||||
assert self.MASKED_ARGS not in str(error)
|
||||
|
||||
# The audit fields LIT-5638 asks for are unaffected.
|
||||
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)
|
||||
error = exc.detail["error"]
|
||||
|
||||
assert error["prompt_masked_data"] == {"data": self.MASKED_ARGS}
|
||||
assert error["scan_id"] == "scan-tool-1"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue