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:
STiFLeR7 2026-07-30 11:11:28 +05:30
parent 721babb0e7
commit a93b5ee3f2
No known key found for this signature in database
2 changed files with 62 additions and 4 deletions

View file

@ -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,

View file

@ -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."""