diff --git a/litellm/llms/a2a/chat/guardrail_translation/handler.py b/litellm/llms/a2a/chat/guardrail_translation/handler.py index 54abcd9b511..71b4f1f3b54 100644 --- a/litellm/llms/a2a/chat/guardrail_translation/handler.py +++ b/litellm/llms/a2a/chat/guardrail_translation/handler.py @@ -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, diff --git a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_handler.py b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_handler.py index 4f958ba15fd..9af0bae00c1 100644 --- a/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/a2a/chat/guardrail_translation/test_handler.py @@ -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."""