From 6f93c7c548f870fce53308be4c02c5548d4039f4 Mon Sep 17 00:00:00 2001 From: Harshitmishra001 Date: Thu, 6 Aug 2026 00:27:42 +0530 Subject: [PATCH] fix(core): replace parts[0] in telemetry and tools to support parallel execution This systemic fix replaces hardcoded parts[0] assumptions with correct name-based matching in telemetry tracing, artifact loading, MCP resource loading, and CLI tooling. This resolves issues like #5967 where parallel tools caused infinite loops or dropped telemetry. --- src/google/adk/cli/conformance/cli_record.py | 29 +++---- src/google/adk/cli/conformance/cli_test.py | 29 +++---- src/google/adk/telemetry/tracing.py | 7 +- src/google/adk/tools/load_artifacts_tool.py | 79 ++++++++++--------- .../adk/tools/load_mcp_resource_tool.py | 51 ++++++------ tests/unittests/telemetry/test_spans.py | 4 +- 6 files changed, 97 insertions(+), 102 deletions(-) diff --git a/src/google/adk/cli/conformance/cli_record.py b/src/google/adk/cli/conformance/cli_record.py index eb38c99478..b75c8b0825 100644 --- a/src/google/adk/cli/conformance/cli_record.py +++ b/src/google/adk/cli/conformance/cli_record.py @@ -72,23 +72,18 @@ async def _create_conformance_test_files( # long-running tool. Replace the function call ID with the actual # function call ID. This is needed because the function call ID is not # known when writing the test case. - if ( - user_message.content.parts - and user_message.content.parts[0].function_response - and user_message.content.parts[0].function_response.name - ): - if ( - user_message.content.parts[0].function_response.name - not in function_call_name_to_id_map - ): - raise ValueError( - "Function response for" - f" {user_message.content.parts[0].function_response.name} does" - " not match any pending function call." - ) - content.parts[0].function_response.id = function_call_name_to_id_map[ - user_message.content.parts[0].function_response.name - ] + if user_message.content.parts: + for part in content.parts: + if part.function_response and part.function_response.name: + if part.function_response.name not in function_call_name_to_id_map: + raise ValueError( + "Function response for" + f" {part.function_response.name} does" + " not match any pending function call." + ) + part.function_response.id = function_call_name_to_id_map[ + part.function_response.name + ] elif user_message.text is not None: content = types.UserContent(parts=[types.Part(text=user_message.text)]) else: diff --git a/src/google/adk/cli/conformance/cli_test.py b/src/google/adk/cli/conformance/cli_test.py index bc8337cf0a..d3e4e6bcb1 100644 --- a/src/google/adk/cli/conformance/cli_test.py +++ b/src/google/adk/cli/conformance/cli_test.py @@ -142,23 +142,18 @@ async def _run_user_messages( # long-running tool. Replace the function call ID with the actual # function call ID. This is needed because the function call ID is not # known when writing the test case. - if ( - user_message.content.parts - and user_message.content.parts[0].function_response - and user_message.content.parts[0].function_response.name - ): - if ( - user_message.content.parts[0].function_response.name - not in function_call_name_to_id_map - ): - raise ValueError( - "Function response for" - f" {user_message.content.parts[0].function_response.name} does" - " not match any pending function call." - ) - content.parts[0].function_response.id = function_call_name_to_id_map[ - user_message.content.parts[0].function_response.name - ] + if user_message.content.parts: + for part in content.parts: + if part.function_response and part.function_response.name: + if part.function_response.name not in function_call_name_to_id_map: + raise ValueError( + "Function response for" + f" {part.function_response.name} does" + " not match any pending function call." + ) + part.function_response.id = function_call_name_to_id_map[ + part.function_response.name + ] elif user_message.text is not None: content = types.UserContent(parts=[types.Part(text=user_message.text)]) else: diff --git a/src/google/adk/telemetry/tracing.py b/src/google/adk/telemetry/tracing.py index 800783cec4..0b372216f0 100644 --- a/src/google/adk/telemetry/tracing.py +++ b/src/google/adk/telemetry/tracing.py @@ -260,8 +260,11 @@ def trace_tool_call( and function_response_event.content is not None and function_response_event.content.parts ): - response_parts = function_response_event.content.parts - function_response = response_parts[0].function_response + function_response = None + for part in function_response_event.content.parts: + if part.function_response and part.function_response.name == tool.name: + function_response = part.function_response + break if function_response is not None: if function_response.id is not None: tool_call_id = function_response.id diff --git a/src/google/adk/tools/load_artifacts_tool.py b/src/google/adk/tools/load_artifacts_tool.py index 99de971a4d..2962d1ca24 100644 --- a/src/google/adk/tools/load_artifacts_tool.py +++ b/src/google/adk/tools/load_artifacts_tool.py @@ -286,46 +286,47 @@ async def _append_artifacts_to_llm_request( # Attach the content of the artifacts if the model requests them. # This only adds the content to the model request, instead of the session. if llm_request.contents and llm_request.contents[-1].parts: - function_response = llm_request.contents[-1].parts[0].function_response - if function_response and function_response.name == 'load_artifacts': - response = function_response.response or {} - artifact_names = response.get('artifact_names', []) - for artifact_name in artifact_names: - # Try session-scoped first (default behavior) - artifact = await tool_context.load_artifact(artifact_name) - - # If not found and name doesn't already have user: prefix, - # try cross-session artifacts with user: prefix - if artifact is None and not artifact_name.startswith('user:'): - prefixed_name = f'user:{artifact_name}' - artifact = await tool_context.load_artifact(prefixed_name) - - if artifact is None: - logger.warning('Artifact "%s" not found, skipping', artifact_name) - continue - - artifact_part = _as_safe_part_for_llm(artifact, artifact_name) - if artifact_part is not artifact: - mime_type = ( - artifact.inline_data.mime_type if artifact.inline_data else None - ) - logger.debug( - 'Converted artifact "%s" (mime_type=%s) to text Part', - artifact_name, - mime_type, - ) - - llm_request.contents.append( - types.Content( - role='user', - parts=[ - types.Part.from_text( - text=f'Artifact {artifact_name} is:' - ), - artifact_part, - ], + for part in llm_request.contents[-1].parts: + function_response = part.function_response + if function_response and function_response.name == 'load_artifacts': + response = function_response.response or {} + artifact_names = response.get('artifact_names', []) + for artifact_name in artifact_names: + # Try session-scoped first (default behavior) + artifact = await tool_context.load_artifact(artifact_name) + + # If not found and name doesn't already have user: prefix, + # try cross-session artifacts with user: prefix + if artifact is None and not artifact_name.startswith('user:'): + prefixed_name = f'user:{artifact_name}' + artifact = await tool_context.load_artifact(prefixed_name) + + if artifact is None: + logger.warning('Artifact "%s" not found, skipping', artifact_name) + continue + + artifact_part = _as_safe_part_for_llm(artifact, artifact_name) + if artifact_part is not artifact: + mime_type = ( + artifact.inline_data.mime_type if artifact.inline_data else None + ) + logger.debug( + 'Converted artifact "%s" (mime_type=%s) to text Part', + artifact_name, + mime_type, ) - ) + + llm_request.contents.append( + types.Content( + role='user', + parts=[ + types.Part.from_text( + text=f'Artifact {artifact_name} is:' + ), + artifact_part, + ], + ) + ) load_artifacts_tool = LoadArtifactsTool() diff --git a/src/google/adk/tools/load_mcp_resource_tool.py b/src/google/adk/tools/load_mcp_resource_tool.py index 86eff9182c..f201bdde22 100644 --- a/src/google/adk/tools/load_mcp_resource_tool.py +++ b/src/google/adk/tools/load_mcp_resource_tool.py @@ -124,32 +124,33 @@ async def _append_resources_to_llm_request( # Attach content if llm_request.contents and llm_request.contents[-1].parts: - function_response = llm_request.contents[-1].parts[0].function_response - if function_response and function_response.name == self.name: - response = function_response.response or {} - resource_names = response.get("resource_names", []) - for resource_name in resource_names: - try: - contents = await self._mcp_toolset.read_resource(resource_name) - - for content in contents: - part = self._mcp_content_to_part(content, resource_name) - llm_request.contents.append( - types.Content( - role="user", - parts=[ - types.Part.from_text( - text=f"Resource {resource_name} is:" - ), - part, - ], - ) + for part in llm_request.contents[-1].parts: + function_response = part.function_response + if function_response and function_response.name == self.name: + response = function_response.response or {} + resource_names = response.get("resource_names", []) + for resource_name in resource_names: + try: + contents = await self._mcp_toolset.read_resource(resource_name) + + for content in contents: + part = self._mcp_content_to_part(content, resource_name) + llm_request.contents.append( + types.Content( + role="user", + parts=[ + types.Part.from_text( + text=f"Resource {resource_name} is:" + ), + part, + ], + ) + ) + except Exception as e: + logger.warning( + "Failed to read MCP resource '%s': %s", resource_name, e ) - except Exception as e: - logger.warning( - "Failed to read MCP resource '%s': %s", resource_name, e - ) - continue + continue def _mcp_content_to_part( self, content: Any, resource_name: str diff --git a/tests/unittests/telemetry/test_spans.py b/tests/unittests/telemetry/test_spans.py index 8039e91822..49676f80fc 100644 --- a/tests/unittests/telemetry/test_spans.py +++ b/tests/unittests/telemetry/test_spans.py @@ -542,7 +542,7 @@ def test_trace_tool_call_with_scalar_response( types.Part( function_response=types.FunctionResponse( id=test_tool_call_id, - name='test_function_1', + name=mock_tool_fixture.name, response={'result': scalar_function_response}, ) ), @@ -602,7 +602,7 @@ def test_trace_tool_call_with_dict_response( types.Part( function_response=types.FunctionResponse( id=test_tool_call_id, - name='test_function_1', + name=mock_tool_fixture.name, response=dict_function_response, ) ),