Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions src/google/adk/agents/remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -514,6 +514,14 @@ async def _handle_a2a_response(
invocation_id=ctx.invocation_id,
branch=ctx.branch,
)
# Filter out thought parts from user-facing response content.
# Intermediate (submitted/working) events have all parts marked as
# thought, so non_thought_parts will be empty and we preserve them.
if event.content is not None and event.content.parts:
non_thought_parts = [p for p in event.content.parts if not p.thought]
if non_thought_parts:
event.content.parts = non_thought_parts

return event
except A2AClientError as e:
logger.error("Failed to handle A2A response: %s", e)
Expand Down
157 changes: 157 additions & 0 deletions tests/unittests/agents/test_remote_a2a_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -1278,6 +1278,163 @@ async def test_handle_a2a_response_with_partial_artifact_update(self):

assert result is None

@pytest.mark.asyncio
async def test_handle_a2a_response_filters_thought_parts_from_completed_task(
self,
):
"""Test that thought parts are filtered from completed task response.

When an A2A server returns a completed task with both thought and
non-thought parts, the client should only include non-thought parts
in the user-facing event. Fixes #4676.
"""
mock_a2a_task = Mock(spec=A2ATask)
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
mock_a2a_task.status.state = TaskState.completed

# Create event with mixed thought/non-thought parts
thought_part = genai_types.Part(text="internal reasoning", thought=True)
answer_part = genai_types.Part(text="final answer")
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
content=genai_types.Content(
role="model", parts=[thought_part, answer_part]
),
)

with patch.object(
remote_a2a_agent,
"convert_a2a_task_to_event",
autospec=True,
) as mock_convert:
mock_convert.return_value = mock_event

result = await self.agent._handle_a2a_response(
(mock_a2a_task, None), self.mock_context
)

# Only non-thought parts should remain
assert len(result.content.parts) == 1
assert result.content.parts[0].text == "final answer"
assert result.content.parts[0].thought is None

@pytest.mark.asyncio
async def test_handle_a2a_response_filters_thought_parts_from_status_update(
self,
):
"""Test that thought parts are filtered from completed status update.

Fixes #4676.
"""
mock_a2a_task = Mock(spec=A2ATask)
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"

mock_update = Mock(spec=TaskStatusUpdateEvent)
mock_update.status = Mock(spec=A2ATaskStatus)
mock_update.status.state = TaskState.completed
mock_update.status.message = Mock(spec=A2AMessage)

# Create event with mixed thought/non-thought parts
thought_part = genai_types.Part(text="thinking...", thought=True)
answer_part = genai_types.Part(text="the answer")
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
content=genai_types.Content(
role="model", parts=[thought_part, answer_part]
),
)

with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

result = await self.agent._handle_a2a_response(
(mock_a2a_task, mock_update), self.mock_context
)

# Only non-thought parts should remain
assert len(result.content.parts) == 1
assert result.content.parts[0].text == "the answer"

@pytest.mark.asyncio
async def test_handle_a2a_response_preserves_all_thought_parts_for_working(
self,
):
"""Test that working state events keep all parts as thoughts.

Intermediate events (working/submitted) should retain all parts
marked as thought for streaming progress display.
"""
mock_a2a_task = Mock(spec=A2ATask)
mock_a2a_task.id = "task-123"
mock_a2a_task.context_id = "context-123"
mock_a2a_task.status = Mock(spec=A2ATaskStatus)
mock_a2a_task.status.state = TaskState.working

part = genai_types.Part(text="still thinking")
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
content=genai_types.Content(role="model", parts=[part]),
)

with patch.object(
remote_a2a_agent,
"convert_a2a_task_to_event",
autospec=True,
) as mock_convert:
mock_convert.return_value = mock_event

result = await self.agent._handle_a2a_response(
(mock_a2a_task, None), self.mock_context
)

# All parts should be marked as thought and preserved
assert len(result.content.parts) == 1
assert result.content.parts[0].thought is True

@pytest.mark.asyncio
async def test_handle_a2a_response_filters_thought_from_a2a_message(self):
"""Test thought filtering for regular A2AMessage responses.

Fixes #4676.
"""
mock_a2a_message = Mock(spec=A2AMessage)
mock_a2a_message.context_id = "context-123"

thought_part = genai_types.Part(text="reasoning", thought=True)
answer_part = genai_types.Part(text="response")
mock_event = Event(
author=self.agent.name,
invocation_id=self.mock_context.invocation_id,
branch=self.mock_context.branch,
content=genai_types.Content(
role="model", parts=[thought_part, answer_part]
),
)

with patch(
"google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
) as mock_convert:
mock_convert.return_value = mock_event

result = await self.agent._handle_a2a_response(
mock_a2a_message, self.mock_context
)

# Only non-thought parts should remain
assert len(result.content.parts) == 1
assert result.content.parts[0].text == "response"


class TestRemoteA2aAgentMessageHandlingFromFactory:
"""Test message handling functionality."""
Expand Down