@@ -1324,6 +1324,163 @@ async def test_handle_a2a_response_with_partial_artifact_update(self):
13241324
13251325 assert result is None
13261326
1327+ @pytest .mark .asyncio
1328+ async def test_handle_a2a_response_filters_thought_parts_from_completed_task (
1329+ self ,
1330+ ):
1331+ """Test that thought parts are filtered from completed task response.
1332+
1333+ When an A2A server returns a completed task with both thought and
1334+ non-thought parts, the client should only include non-thought parts
1335+ in the user-facing event. Fixes #4676.
1336+ """
1337+ mock_a2a_task = Mock (spec = A2ATask )
1338+ mock_a2a_task .id = "task-123"
1339+ mock_a2a_task .context_id = "context-123"
1340+ mock_a2a_task .status = Mock (spec = A2ATaskStatus )
1341+ mock_a2a_task .status .state = TaskState .completed
1342+
1343+ # Create event with mixed thought/non-thought parts
1344+ thought_part = genai_types .Part (text = "internal reasoning" , thought = True )
1345+ answer_part = genai_types .Part (text = "final answer" )
1346+ mock_event = Event (
1347+ author = self .agent .name ,
1348+ invocation_id = self .mock_context .invocation_id ,
1349+ branch = self .mock_context .branch ,
1350+ content = genai_types .Content (
1351+ role = "model" , parts = [thought_part , answer_part ]
1352+ ),
1353+ )
1354+
1355+ with patch .object (
1356+ remote_a2a_agent ,
1357+ "convert_a2a_task_to_event" ,
1358+ autospec = True ,
1359+ ) as mock_convert :
1360+ mock_convert .return_value = mock_event
1361+
1362+ result = await self .agent ._handle_a2a_response (
1363+ (mock_a2a_task , None ), self .mock_context
1364+ )
1365+
1366+ # Only non-thought parts should remain
1367+ assert len (result .content .parts ) == 1
1368+ assert result .content .parts [0 ].text == "final answer"
1369+ assert result .content .parts [0 ].thought is None
1370+
1371+ @pytest .mark .asyncio
1372+ async def test_handle_a2a_response_filters_thought_parts_from_status_update (
1373+ self ,
1374+ ):
1375+ """Test that thought parts are filtered from completed status update.
1376+
1377+ Fixes #4676.
1378+ """
1379+ mock_a2a_task = Mock (spec = A2ATask )
1380+ mock_a2a_task .id = "task-123"
1381+ mock_a2a_task .context_id = "context-123"
1382+
1383+ mock_update = Mock (spec = TaskStatusUpdateEvent )
1384+ mock_update .status = Mock (spec = A2ATaskStatus )
1385+ mock_update .status .state = TaskState .completed
1386+ mock_update .status .message = Mock (spec = A2AMessage )
1387+
1388+ # Create event with mixed thought/non-thought parts
1389+ thought_part = genai_types .Part (text = "thinking..." , thought = True )
1390+ answer_part = genai_types .Part (text = "the answer" )
1391+ mock_event = Event (
1392+ author = self .agent .name ,
1393+ invocation_id = self .mock_context .invocation_id ,
1394+ branch = self .mock_context .branch ,
1395+ content = genai_types .Content (
1396+ role = "model" , parts = [thought_part , answer_part ]
1397+ ),
1398+ )
1399+
1400+ with patch (
1401+ "google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
1402+ ) as mock_convert :
1403+ mock_convert .return_value = mock_event
1404+
1405+ result = await self .agent ._handle_a2a_response (
1406+ (mock_a2a_task , mock_update ), self .mock_context
1407+ )
1408+
1409+ # Only non-thought parts should remain
1410+ assert len (result .content .parts ) == 1
1411+ assert result .content .parts [0 ].text == "the answer"
1412+
1413+ @pytest .mark .asyncio
1414+ async def test_handle_a2a_response_preserves_all_thought_parts_for_working (
1415+ self ,
1416+ ):
1417+ """Test that working state events keep all parts as thoughts.
1418+
1419+ Intermediate events (working/submitted) should retain all parts
1420+ marked as thought for streaming progress display.
1421+ """
1422+ mock_a2a_task = Mock (spec = A2ATask )
1423+ mock_a2a_task .id = "task-123"
1424+ mock_a2a_task .context_id = "context-123"
1425+ mock_a2a_task .status = Mock (spec = A2ATaskStatus )
1426+ mock_a2a_task .status .state = TaskState .working
1427+
1428+ part = genai_types .Part (text = "still thinking" )
1429+ mock_event = Event (
1430+ author = self .agent .name ,
1431+ invocation_id = self .mock_context .invocation_id ,
1432+ branch = self .mock_context .branch ,
1433+ content = genai_types .Content (role = "model" , parts = [part ]),
1434+ )
1435+
1436+ with patch .object (
1437+ remote_a2a_agent ,
1438+ "convert_a2a_task_to_event" ,
1439+ autospec = True ,
1440+ ) as mock_convert :
1441+ mock_convert .return_value = mock_event
1442+
1443+ result = await self .agent ._handle_a2a_response (
1444+ (mock_a2a_task , None ), self .mock_context
1445+ )
1446+
1447+ # All parts should be marked as thought and preserved
1448+ assert len (result .content .parts ) == 1
1449+ assert result .content .parts [0 ].thought is True
1450+
1451+ @pytest .mark .asyncio
1452+ async def test_handle_a2a_response_filters_thought_from_a2a_message (self ):
1453+ """Test thought filtering for regular A2AMessage responses.
1454+
1455+ Fixes #4676.
1456+ """
1457+ mock_a2a_message = Mock (spec = A2AMessage )
1458+ mock_a2a_message .context_id = "context-123"
1459+
1460+ thought_part = genai_types .Part (text = "reasoning" , thought = True )
1461+ answer_part = genai_types .Part (text = "response" )
1462+ mock_event = Event (
1463+ author = self .agent .name ,
1464+ invocation_id = self .mock_context .invocation_id ,
1465+ branch = self .mock_context .branch ,
1466+ content = genai_types .Content (
1467+ role = "model" , parts = [thought_part , answer_part ]
1468+ ),
1469+ )
1470+
1471+ with patch (
1472+ "google.adk.agents.remote_a2a_agent.convert_a2a_message_to_event"
1473+ ) as mock_convert :
1474+ mock_convert .return_value = mock_event
1475+
1476+ result = await self .agent ._handle_a2a_response (
1477+ mock_a2a_message , self .mock_context
1478+ )
1479+
1480+ # Only non-thought parts should remain
1481+ assert len (result .content .parts ) == 1
1482+ assert result .content .parts [0 ].text == "response"
1483+
13271484
13281485class TestRemoteA2aAgentMessageHandlingFromFactory :
13291486 """Test message handling functionality."""
0 commit comments