mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(a2a): recurse into nested parts for guardrail scanning
extract_text_from_a2a_message recurses into a part that itself carries a nested 'parts' list, but _extract_texts_from_parts did not, so text/data content inside nested parts could reach the completion text without ever passing through a guardrail. Mirror the recursion (including the same depth guard) so the two extraction paths can't diverge on nested structures either. _apply_text_to_path needed no changes since its path-navigation is already depth-agnostic. Addresses the follow-up finding on this PR: nested data bypasses output guardrails (litellm/llms/a2a/common_utils.py:100).
This commit is contained in:
parent
721babb0e7
commit
a93b5ee3f2
2 changed files with 62 additions and 4 deletions
|
|
@ -36,6 +36,7 @@ class _A2ATextPart(TypedDict, total=False):
|
|||
kind: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
data: ReadOnly[object]
|
||||
parts: ReadOnly[Sequence["_A2ATextPart"]]
|
||||
|
||||
|
||||
class A2AGuardrailHandler(BaseTranslation):
|
||||
|
|
@ -449,8 +450,18 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
path: tuple[str, ...],
|
||||
texts_to_check: list[str],
|
||||
task_mappings: list[tuple[tuple[str, ...], int, str]],
|
||||
depth: int = 0,
|
||||
max_depth: int = 10,
|
||||
) -> None:
|
||||
"""Extract text from message parts, and serialized data from data parts."""
|
||||
"""
|
||||
Extract text from message parts, serialized data from data parts, and
|
||||
recurse into any part that itself carries a nested "parts" list.
|
||||
|
||||
Mirrors `extract_text_from_a2a_message`'s handling (including the
|
||||
recursion depth guard) so the two stay in sync.
|
||||
"""
|
||||
if depth >= max_depth:
|
||||
return
|
||||
for part_idx, part in enumerate(parts):
|
||||
kind = part.get("kind")
|
||||
if kind == "text":
|
||||
|
|
@ -459,10 +470,19 @@ class A2AGuardrailHandler(BaseTranslation):
|
|||
texts_to_check.append(text)
|
||||
task_mappings.append((path, part_idx, "text"))
|
||||
elif kind == "data":
|
||||
data = part.get("data")
|
||||
if data is not None:
|
||||
texts_to_check.append(serialize_a2a_data_part(data))
|
||||
part_data = part.get("data")
|
||||
if part_data is not None:
|
||||
texts_to_check.append(serialize_a2a_data_part(part_data))
|
||||
task_mappings.append((path, part_idx, "data"))
|
||||
elif "parts" in part:
|
||||
self._extract_texts_from_parts(
|
||||
parts=part["parts"],
|
||||
path=path + (str(part_idx), "parts"),
|
||||
texts_to_check=texts_to_check,
|
||||
task_mappings=task_mappings,
|
||||
depth=depth + 1,
|
||||
max_depth=max_depth,
|
||||
)
|
||||
|
||||
def _apply_text_to_path(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -104,6 +104,44 @@ async def test_process_output_response_data_only_still_scanned():
|
|||
assert "PONG" in result["result"]["parts"][0]["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_output_response_scans_nested_parts():
|
||||
"""A part that itself carries a nested "parts" list (grouping sub-parts)
|
||||
must be recursed into, matching extract_text_from_a2a_message's own
|
||||
recursion, instead of being silently skipped as neither text nor data."""
|
||||
handler = A2AGuardrailHandler()
|
||||
guardrail = MockGuardrail()
|
||||
|
||||
response = {
|
||||
"result": {
|
||||
"kind": "message",
|
||||
"parts": [
|
||||
{
|
||||
"kind": "group",
|
||||
"parts": [
|
||||
{"kind": "text", "text": "hello"},
|
||||
{"kind": "data", "data": {"secret": "leak-me"}},
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
result = await handler.process_output_response(
|
||||
response=response,
|
||||
guardrail_to_apply=guardrail,
|
||||
)
|
||||
|
||||
assert guardrail.last_inputs is not None
|
||||
scanned_texts = guardrail.last_inputs["texts"]
|
||||
assert "hello" in scanned_texts
|
||||
assert any("leak-me" in t for t in scanned_texts)
|
||||
|
||||
nested_parts = result["result"]["parts"][0]["parts"]
|
||||
assert nested_parts[0]["text"] == "HELLO"
|
||||
assert "LEAK-ME" in nested_parts[1]["data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_input_messages_scans_data_parts():
|
||||
"""The same bypass existed on the request/input side of the handler."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue