fix: use child task IDs in a2a_mcp workflows

Signed-off-by: Anas Khan <83116240+anxkhn@users.noreply.github.com>

Author: anxkhn

samples/go/agents/deepresearch/go.mod

@@ -39,8 +39,8 @@ require (
 	go.opentelemetry.io/otel/trace v1.43.0 // indirect
 	golang.org/x/crypto v0.51.0 // indirect
 	golang.org/x/mod v0.35.0 // indirect
-	golang.org/x/net v0.54.0 // indirect
-	golang.org/x/sys v0.44.0 // indirect
+	golang.org/x/net v0.55.0 // indirect
+	golang.org/x/sys v0.45.0 // indirect
 	golang.org/x/text v0.37.0 // indirect
 	google.golang.org/api v0.279.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20260511170946-3700d4141b60 // indirect

samples/go/agents/deepresearch/go.sum

@@ -73,12 +73,12 @@ golang.org/x/crypto v0.51.0 h1:IBPXwPfKxY7cWQZ38ZCIRPI50YLeevDLlLnyC5wRGTI=
 golang.org/x/crypto v0.51.0/go.mod h1:8AdwkbraGNABw2kOX6YFPs3WM22XqI4EXEd8g+x7Oc8=
 golang.org/x/mod v0.35.0 h1:Ww1D637e6Pg+Zb2KrWfHQUnH2dQRLBQyAtpr/haaJeM=
 golang.org/x/mod v0.35.0/go.mod h1:+GwiRhIInF8wPm+4AoT6L0FA1QWAad3OMdTRx4tFYlU=
-golang.org/x/net v0.54.0 h1:2zJIZAxAHV/OHCDTCOHAYehQzLfSXuf/5SoL/Dv6w/w=
-golang.org/x/net v0.54.0/go.mod h1:Sj4oj8jK6XmHpBZU/zWHw3BV3abl4Kvi+Ut7cQcY+cQ=
+golang.org/x/net v0.55.0 h1:bcvxaJn3e1U6InsFWt1JUq1aSjnRxLzT2rtD2KfkDF8=
+golang.org/x/net v0.55.0/go.mod h1:L5U2KuzuOe1lY7Z+aWVIKK6qEeJXnXV9yzGA+WCHJww=
 golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4=
 golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
-golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ=
-golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
+golang.org/x/sys v0.45.0 h1:dO4czNzziLiiXplLQgBCEpCvXQ3dnkn0SdaZSYdQ+FY=
+golang.org/x/sys v0.45.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc=
 golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38=

samples/python/agents/a2a_mcp/src/a2a_mcp/agents/orchestrator_agent.py

@@ -1,3 +1,4 @@
+# ruff: noqa
 import json
 import logging
 
@@ -39,9 +40,7 @@ async def generate_summary(self) -> str:
         client = genai.Client()
         response = client.models.generate_content(
             model='gemini-2.0-flash',
-            contents=prompts.SUMMARY_COT_INSTRUCTIONS.replace(
-                '{travel_data}', str(self.results)
-            ),
+            contents=prompts.SUMMARY_COT_INSTRUCTIONS.replace('{travel_data}', str(self.results)),
             config={'temperature': 0.0},
         )
         return response.text
@@ -51,9 +50,7 @@ def answer_user_question(self, question) -> str:
             client = genai.Client()
             response = client.models.generate_content(
                 model='gemini-2.0-flash',
-                contents=prompts.QA_COT_PROMPT.replace(
-                    '{TRIP_CONTEXT}', str(self.travel_context)
-                )
+                contents=prompts.QA_COT_PROMPT.replace('{TRIP_CONTEXT}', str(self.travel_context))
                 .replace('{CONVERSATION_HISTORY}', str(self.query_history))
                 .replace('{TRIP_QUESTION}', question),
                 config={
@@ -66,9 +63,7 @@ def answer_user_question(self, question) -> str:
             logger.info(f'Error answering user question: {e}')
         return '{"can_answer": "no", "answer": "Cannot answer based on provided context"}'
 
-    def set_node_attributes(
-        self, node_id, task_id=None, context_id=None, query=None
-    ):
+    def set_node_attributes(self, node_id, task_id=None, context_id=None, query=None):
         attr_val = {}
         if task_id:
             attr_val['task_id'] = task_id
@@ -81,21 +76,23 @@ def set_node_attributes(
 
     def add_graph_node(
         self,
-        task_id,
         context_id,
         query: str,
         node_id: str = None,
         node_key: str = None,
         node_label: str = None,
     ) -> WorkflowNode:
-        """Add a node to the graph."""
-        node = WorkflowNode(
-            task=query, node_key=node_key, node_label=node_label
-        )
+        """Add a node to the graph.
+
+        Child nodes must not inherit the orchestrator's task_id. A2A
+        treats a present taskId as a resume of a task on the receiving
+        agent, and the child has not created that task yet.
+        """
+        node = WorkflowNode(task=query, node_key=node_key, node_label=node_label)
         self.graph.add_node(node)
         if node_id:
             self.graph.add_edge(node_id, node.id)
-        self.set_node_attributes(node.id, task_id, context_id, query)
+        self.set_node_attributes(node.id, context_id=context_id, query=query)
         return node
 
     def clear_state(self):
@@ -104,9 +101,7 @@ def clear_state(self):
         self.travel_context.clear()
         self.query_history.clear()
 
-    async def stream(
-        self, query, context_id, task_id
-    ) -> AsyncIterable[dict[str, any]]:
+    async def stream(self, query, context_id, task_id) -> AsyncIterable[dict[str, any]]:
         """Execute and stream response."""
         logger.info(
             f'Running {self.agent_name} stream for session {context_id}, task {task_id} - {query}'
@@ -124,7 +119,6 @@ async def stream(
         if not self.graph:
             self.graph = WorkflowGraph()
             planner_node = self.add_graph_node(
-                task_id=task_id,
                 context_id=context_id,
                 query=query,
                 node_key='planner',
@@ -138,59 +132,45 @@ async def stream(
 
         # This loop can be avoided if the workflow graph is dynamic or
         # is built from the results of the planner when the planner
-        # iself is not a part of the graph.
+        # itself is not a part of the graph.
         # TODO: Make the graph dynamically iterable over edges
         while True:
-            # Set attributes on the node so we propagate task and context
+            # Propagate context only. Child task ids are stored on the
+            # node from the child's stream and must not be overwritten
+            # with the orchestrator's task_id.
             self.set_node_attributes(
                 node_id=start_node_id,
-                task_id=task_id,
                 context_id=context_id,
             )
             # Resume workflow, used when the workflow nodes are updated.
             should_resume_workflow = False
-            async for chunk in self.graph.run_workflow(
-                start_node_id=start_node_id
-            ):
+            async for chunk in self.graph.run_workflow(start_node_id=start_node_id):
                 if isinstance(chunk.root, SendStreamingMessageSuccessResponse):
-                    # The graph node retured TaskStatusUpdateEvent
+                    # The graph node returned TaskStatusUpdateEvent
                     # Check if the node is complete and continue to the next node
                     if isinstance(chunk.root.result, TaskStatusUpdateEvent):
                         task_status_event = chunk.root.result
                         context_id = task_status_event.context_id
-                        if (
-                            task_status_event.status.state
-                            == TaskState.completed
-                            and context_id
-                        ):
-                            ## yeild??
+                        if task_status_event.status.state == TaskState.completed and context_id:
+                            ## yield??
                             continue
-                        if (
-                            task_status_event.status.state
-                            == TaskState.input_required
-                        ):
-                            question = task_status_event.status.message.parts[
-                                0
-                            ].root.text
+                        if task_status_event.status.state == TaskState.input_required:
+                            question = task_status_event.status.message.parts[0].root.text
 
                             try:
-                                answer = json.loads(
-                                    self.answer_user_question(question)
-                                )
+                                answer = json.loads(self.answer_user_question(question))
                                 logger.info(f'Agent Answer {answer}')
                                 if answer['can_answer'] == 'yes':
                                     # Orchestrator can answer on behalf of the user set the query
                                     # Resume workflow from paused state.
                                     query = answer['answer']
                                     start_node_id = self.graph.paused_node_id
-                                    self.set_node_attributes(
-                                        node_id=start_node_id, query=query
-                                    )
+                                    self.set_node_attributes(node_id=start_node_id, query=query)
                                     should_resume_workflow = True
                             except Exception:
                                 logger.info('Cannot convert answer data')
 
-                    # The graph node retured TaskArtifactUpdateEvent
+                    # The graph node returned TaskArtifactUpdateEvent
                     # Store the node and continue.
                     if isinstance(chunk.root.result, TaskArtifactUpdateEvent):
                         artifact = chunk.root.result.artifact
@@ -205,11 +185,8 @@ async def stream(
                             )
                             # Define the edges
                             current_node_id = start_node_id
-                            for idx, task_data in enumerate(
-                                artifact_data['tasks']
-                            ):
+                            for idx, task_data in enumerate(artifact_data['tasks']):
                                 node = self.add_graph_node(
-                                    task_id=task_id,
                                     context_id=context_id,
                                     query=task_data['description'],
                                     node_id=current_node_id,

samples/python/agents/a2a_mcp/src/a2a_mcp/common/child_message.py

@@ -0,0 +1,38 @@
+# ruff: noqa
+"""Helpers for child-agent A2A message payloads."""
+
+from typing import Any
+from uuid import uuid4
+
+
+def create_child_message_payload(
+    query: str,
+    context_id: str,
+    task_id: str | None = None,
+) -> dict[str, Any]:
+    """Build a child-agent message payload.
+
+    A2A treats a present taskId as a resume of an existing task on the
+    receiving server. New work must omit taskId so the child creates one.
+    """
+    message: dict[str, Any] = {
+        'role': 'user',
+        'parts': [{'kind': 'text', 'text': query}],
+        'messageId': uuid4().hex,
+        'contextId': context_id,
+    }
+    if task_id:
+        message['taskId'] = task_id
+    return {'message': message}
+
+
+def child_task_id_from_result(result: object) -> str | None:
+    """Return the child-local task id from a streaming result, if present."""
+    event_task_id = getattr(result, 'task_id', None)
+    if isinstance(event_task_id, str) and event_task_id:
+        return event_task_id
+    # Task objects expose id plus a status; request wrappers expose only id.
+    result_id = getattr(result, 'id', None)
+    if isinstance(result_id, str) and result_id and hasattr(result, 'status'):
+        return result_id
+    return None

samples/python/agents/a2a_mcp/src/a2a_mcp/common/workflow.py

@@ -1,3 +1,4 @@
+# ruff: noqa
 import json
 import logging
 import uuid
@@ -19,6 +20,10 @@
     TaskState,
     TaskStatusUpdateEvent,
 )
+from a2a_mcp.common.child_message import (
+    child_task_id_from_result,
+    create_child_message_payload,
+)
 from a2a_mcp.common.utils import get_mcp_server_config
 from a2a_mcp.mcp import client
 
@@ -61,21 +66,15 @@ def __init__(
     async def get_planner_resource(self) -> AgentCard | None:
         logger.info(f'Getting resource for node {self.id}')
         config = get_mcp_server_config()
-        async with client.init_session(
-            config.host, config.port, config.transport
-        ) as session:
-            response = await client.find_resource(
-                session, 'resource://agent_cards/planner_agent'
-            )
+        async with client.init_session(config.host, config.port, config.transport) as session:
+            response = await client.find_resource(session, 'resource://agent_cards/planner_agent')
             data = json.loads(response.contents[0].text)
             return AgentCard(**data['agent_card'][0])
 
     async def find_agent_for_task(self) -> AgentCard | None:
         logger.info(f'Find agent for task - {self.task}')
         config = get_mcp_server_config()
-        async with client.init_session(
-            config.host, config.port, config.transport
-        ) as session:
+        async with client.init_session(config.host, config.port, config.transport) as session:
             result = await client.find_agent(session, self.task)
             agent_card_json = json.loads(result.content[0].text)
             logger.debug(f'Found agent {agent_card_json} for task {self.task}')
@@ -84,7 +83,7 @@ async def find_agent_for_task(self) -> AgentCard | None:
     async def run_node(
         self,
         query: str,
-        task_id: str,
+        task_id: str | None,
         context_id: str,
     ) -> AsyncIterable[dict[str, any]]:
         logger.info(f'Executing node {self.id}')
@@ -96,24 +95,20 @@ async def run_node(
         async with httpx.AsyncClient() as httpx_client:
             client = A2AClient(httpx_client, agent_card)
 
-            payload: dict[str, any] = {
-                'message': {
-                    'role': 'user',
-                    'parts': [{'kind': 'text', 'text': query}],
-                    'messageId': uuid4().hex,
-                    'taskId': task_id,
-                    'contextId': context_id,
-                },
-            }
+            payload = create_child_message_payload(
+                query=query,
+                context_id=context_id,
+                task_id=task_id,
+            )
             request = SendStreamingMessageRequest(
                 id=str(uuid4()), params=MessageSendParams(**payload)
             )
             response_stream = client.send_message_streaming(request)
             async for chunk in response_stream:
                 # Save the artifact as a result of the node
-                if isinstance(
-                    chunk.root, SendStreamingMessageSuccessResponse
-                ) and (isinstance(chunk.root.result, TaskArtifactUpdateEvent)):
+                if isinstance(chunk.root, SendStreamingMessageSuccessResponse) and (
+                    isinstance(chunk.root.result, TaskArtifactUpdateEvent)
+                ):
                     artifact = chunk.root.result.artifact
                     self.results = artifact
                 yield chunk
@@ -142,9 +137,7 @@ def add_edge(self, from_node_id: str, to_node_id: str) -> None:
 
         self.graph.add_edge(from_node_id, to_node_id)
 
-    async def run_workflow(
-        self, start_node_id: str | None = None
-    ) -> AsyncIterable[dict[str, any]]:
+    async def run_workflow(self, start_node_id: str | None = None) -> AsyncIterable[dict[str, any]]:
         logger.info('Executing workflow graph')
         if not start_node_id or start_node_id not in self.nodes:
             start_nodes = [n for n, d in self.graph.in_degree() if d == 0]
@@ -172,21 +165,21 @@ async def run_workflow(
                 # When the workflow node is paused, do not yield any chunks
                 # but, let the loop complete.
                 if node.state != Status.PAUSED:
-                    if isinstance(
-                        chunk.root, SendStreamingMessageSuccessResponse
-                    ) and (
-                        isinstance(chunk.root.result, TaskStatusUpdateEvent)
-                    ):
-                        task_status_event = chunk.root.result
-                        context_id = task_status_event.context_id
-                        if (
-                            task_status_event.status.state
-                            == TaskState.input_required
-                            and context_id
-                        ):
-                            node.state = Status.PAUSED
-                            self.state = Status.PAUSED
-                            self.paused_node_id = node.id
+                    if isinstance(chunk.root, SendStreamingMessageSuccessResponse):
+                        result = chunk.root.result
+                        child_task_id = child_task_id_from_result(result)
+                        if child_task_id:
+                            self.set_node_attribute(node_id, 'task_id', child_task_id)
+                        if isinstance(result, TaskStatusUpdateEvent):
+                            task_status_event = result
+                            context_id = task_status_event.context_id
+                            if (
+                                task_status_event.status.state == TaskState.input_required
+                                and context_id
+                            ):
+                                node.state = Status.PAUSED
+                                self.state = Status.PAUSED
+                                self.paused_node_id = node.id
                     yield chunk
             if self.state == Status.PAUSED:
                 break