mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(panw_prisma_airs): route tool-block masked data by scan side, not by key name
Merging #37036 (already on staging) with this PR produces no conflict and a silent bug. #37036 withholds prompt_masked_data on response-side tool blocks, which was right while tool calls went out as a request-side tool_event: AIRS reported the model's arguments under that key. This PR scans tool calls as ordinary prompt/response text, so the side of the scan now decides which key holds what. The model's arguments arrive under response_masked_data, already covered by _CLIENT_HIDDEN_SCAN_FIELDS, and prompt_masked_data goes back to being the caller's own input -- one of the audit fields LIT-5638 asks for. Left as merged, a response-side tool block drops that field with nothing to flag it. - Tool-path block branch calls _build_error_detail without also_hide - also_hide parameter removed; after this change it has no callers - Regression test asserts both directions: model output withheld, caller input preserved. It fails against the auto-merged combination. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
a8dde0ce2f
commit
fd9f6396e5
2 changed files with 52 additions and 43 deletions
|
|
@ -685,19 +685,11 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
||||||
self,
|
self,
|
||||||
scan_result: Mapping[str, object],
|
scan_result: Mapping[str, object],
|
||||||
is_response: bool = False,
|
is_response: bool = False,
|
||||||
also_hide: str | None = None,
|
|
||||||
) -> Mapping[str, Mapping[str, object]]:
|
) -> 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"
|
action_type: Final = "Response" if is_response else "Prompt"
|
||||||
code_suffix: Final = "_response_blocked" if is_response else "_blocked"
|
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")
|
category: Final = scan_result.get("category", "unknown")
|
||||||
default_msg: Final = f"{action_type} blocked by PANW Prisma AI Security policy (Category: {category})"
|
default_msg: Final = f"{action_type} blocked by PANW Prisma AI Security policy (Category: {category})"
|
||||||
|
|
||||||
|
|
@ -717,7 +709,7 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
||||||
**{
|
**{
|
||||||
key: value
|
key: value
|
||||||
for key, value in scan_result.items()
|
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,
|
"message": error_msg,
|
||||||
"type": "guardrail_violation",
|
"type": "guardrail_violation",
|
||||||
|
|
@ -1498,17 +1490,12 @@ class PanwPrismaAirsHandler(CustomGuardrail):
|
||||||
):
|
):
|
||||||
self._set_tool_call_arguments(tool_call, masked_args)
|
self._set_tool_call_arguments(tool_call, masked_args)
|
||||||
else:
|
else:
|
||||||
# tool_event scans are request-side in the AIRS schema, so AIRS returns
|
# Tool calls now go out as ordinary prompt/response text, so a
|
||||||
# the model's own tool arguments under prompt_masked_data. On a
|
# response-side scan reports the model's arguments under
|
||||||
# response-side block that is generated content, not caller input, and
|
# response_masked_data, which _CLIENT_HIDDEN_SCAN_FIELDS already
|
||||||
# the class-level default only withholds response_masked_data — which is
|
# withholds. prompt_masked_data is the caller's own input again and
|
||||||
# empty on this path. Withhold it explicitly so the 400 does not become
|
# must keep reaching them -- it is one of the fields LIT-5638 asks for.
|
||||||
# the content channel this branch declined to deliver.
|
error_detail = self._build_error_detail(scan_result, is_response=is_response)
|
||||||
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)
|
raise HTTPException(status_code=400, detail=error_detail)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
|
|
||||||
|
|
@ -5974,26 +5974,41 @@ class TestPanwAirsErrorDetailWithheldFields:
|
||||||
assert detail["error"]["scan_id"] == "scan-2"
|
assert detail["error"]["scan_id"] == "scan-2"
|
||||||
|
|
||||||
|
|
||||||
class TestPanwAirsToolCallBlockWithholdsGeneratedArgs:
|
class TestPanwAirsToolCallBlockMaskedDataRouting:
|
||||||
"""A response-side tool-call block must not ship the model's tool arguments.
|
"""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 calls are scanned as ordinary prompt/response text, so the side of the scan
|
||||||
tool_event is request-side in the AIRS schema, so AIRS returns the scanned tool
|
decides which key holds what: a response-side scan reports the model's generated
|
||||||
arguments under ``prompt_masked_data``. When the tool calls being scanned are the
|
arguments under ``response_masked_data`` (withheld by
|
||||||
model's own output, that key holds generated content, and the class-level
|
``_CLIENT_HIDDEN_SCAN_FIELDS``), while ``prompt_masked_data`` is the caller's own
|
||||||
``_CLIENT_HIDDEN_SCAN_FIELDS`` default (``response_masked_data``, empty on this
|
input and is one of the fields LIT-5638 asks for.
|
||||||
path) does not cover it.
|
|
||||||
|
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",
|
"action": "block",
|
||||||
"category": "sensitive_data",
|
"category": "sensitive_data",
|
||||||
"scan_id": "scan-tool-1",
|
"scan_id": "scan-tool-1",
|
||||||
"prompt_detected": {"dlp": True},
|
"prompt_detected": {"dlp": True},
|
||||||
"prompt_masked_data": {"data": MASKED_ARGS},
|
"response_detected": {"dlp": True},
|
||||||
"response_masked_data": {},
|
"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
|
@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:
|
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:
|
with pytest.raises(HTTPException) as exc_info:
|
||||||
await handler._scan_tool_calls_for_guardrail(
|
await handler._scan_tool_calls_for_guardrail(
|
||||||
tool_calls=[self._tool_call()],
|
tool_calls=[self._tool_call()],
|
||||||
|
|
@ -6027,27 +6042,34 @@ class TestPanwAirsToolCallBlockWithholdsGeneratedArgs:
|
||||||
# The block branch is only reached with masking off; guard the premise.
|
# The block branch is only reached with masking off; guard the premise.
|
||||||
assert handler.mask_response_content is False
|
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"]
|
error = exc.detail["error"]
|
||||||
|
|
||||||
assert exc.status_code == 400
|
assert exc.status_code == 400
|
||||||
assert "prompt_masked_data" not in error
|
assert "response_masked_data" not in error
|
||||||
assert self.MASKED_ARGS not in str(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["scan_id"] == "scan-tool-1"
|
||||||
assert error["prompt_detected"] == {"dlp": True}
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_request_side_block_still_returns_masked_tool_args(self):
|
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."""
|
"""Caller-supplied tool arguments stay in the verdict — that is the ticket's ask."""
|
||||||
handler = make_handler(mask_request_content=False)
|
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"]
|
error = exc.detail["error"]
|
||||||
|
|
||||||
assert error["prompt_masked_data"] == {"data": self.MASKED_ARGS}
|
assert error["prompt_masked_data"] == {"data": self.MODEL_ARGS}
|
||||||
assert error["scan_id"] == "scan-tool-1"
|
assert error["scan_id"] == "scan-tool-2"
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue