Merge pull request #673 from cnoe-io/feat/a2a-source-agent-tracking

feat(a2a): add source agent tracking for sub-agent message grouping

Author: sriaradhyula

ai_platform_engineering/multi_agents/platform_engineer/protocol_bindings/a2a/agent.py

@@ -268,6 +268,10 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
           # Format: {tool_name: response_content}
           accumulated_subagent_responses = {}
 
+          # Track current active agent for sub-agent message grouping
+          # This is used by the executor to add sourceAgent metadata to artifacts
+          current_agent: str | None = None
+
           # Check if token-by-token streaming is enabled (default: true)
           # When disabled, uses 'values' mode which waits for complete messages
           enable_streaming = os.getenv("ENABLE_STREAMING", "true").lower() == "true"
@@ -343,6 +347,8 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                               logging.debug("Skipping tool call with empty name (streaming chunk)")
                               continue
 
+                          # Track current agent for sub-agent message grouping
+                          current_agent = tool_name
                           logging.debug(f"Tool call started (from AIMessageChunk): {tool_name}")
 
                           # Stream tool start notification to client with metadata
@@ -351,6 +357,7 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                               "is_task_complete": False,
                               "require_user_input": False,
                               "content": f"🔧 Supervisor: Calling Agent {tool_name_formatted}...\n",
+                              "source_agent": tool_name,
                               "tool_call": {
                                   "name": tool_name,
                                   "status": "started",
@@ -388,12 +395,14 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                       if match:
                           agent_name = match.group(1)
                           purpose = match.group(2)
+                          current_agent = agent_name.lower()  # Update current agent
                           logging.debug(f"Tool update detected: {agent_name} - {purpose}")
                           # Emit as tool_update event
                           yield {
                               "is_task_complete": False,
                               "require_user_input": False,
                               "content": content,
+                              "source_agent": current_agent,
                               "tool_update": {
                                   "name": agent_name.lower(),
                                   "purpose": purpose,
@@ -402,11 +411,12 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                               }
                           }
                       else:
-                          # Regular content - no special handling
+                          # Regular content - include source_agent for grouping
                           yield {
                               "is_task_complete": False,
                               "require_user_input": False,
                               "content": content,
+                              "source_agent": current_agent or "supervisor",
                           }
 
               # Handle AIMessage with tool calls (tool start indicators)
@@ -425,6 +435,8 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                           pending_tool_calls[tool_call_id] = tool_name
                           logging.debug(f"Tracked tool call: {tool_call_id} -> {tool_name}")
 
+                      # Track current agent for sub-agent message grouping
+                      current_agent = tool_name
                       logging.info(f"Tool call started: {tool_name}")
 
                       # Stream tool start notification to client with metadata
@@ -433,6 +445,7 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                           "is_task_complete": False,
                           "require_user_input": False,
                           "content": f"🔧 Supervisor: Calling Agent {tool_name_formatted}...\n",
+                          "source_agent": tool_name,
                           "tool_call": {
                               "name": tool_name,
                               "status": "started",
@@ -494,6 +507,7 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                           yield {
                               "is_task_complete": False,
                               "require_user_input": False,
+                              "source_agent": "supervisor",
                               "artifact": {
                                   "name": "execution_plan_update",
                                   "description": "TODO-based execution plan",
@@ -507,6 +521,7 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                           yield {
                               "is_task_complete": False,
                               "require_user_input": False,
+                              "source_agent": "supervisor",
                               "artifact": {
                                   "name": "execution_plan_status_update",
                                   "description": "TODO progress update",
@@ -557,17 +572,19 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                           logging.warning(f"Failed to parse request_user_input content: {e}")
                           # Fall through to normal handling if parsing fails
                   elif tool_name in rag_tool_names:
-                    # For RAG tools, we don't want to stream the content, as its a LOT of text
+                    # For RAG tools, we don't want to stream the content, as its a LOT of text
                       yield {
                             "is_task_complete": False,
                             "require_user_input": False,
+                            "source_agent": tool_name,
                             "content": f"🔍 {tool_name}...",
                       }
                   # Stream other tool content normally (actual results for user)
                   elif tool_content and tool_content.strip():
                       yield {
                           "is_task_complete": False,
                           "require_user_input": False,
+                          "source_agent": tool_name,
                           "content": tool_content + "\n",
                       }
 
@@ -576,6 +593,7 @@ async def stream(self, query, context_id, trace_id=None) -> AsyncIterable[dict[s
                   yield {
                       "is_task_complete": False,
                       "require_user_input": False,
+                      "source_agent": tool_name,
                       "content": f"✅ Supervisor: Agent task {tool_name_formatted} completed\n",
                       "tool_result": {
                           "name": tool_name,

ai_platform_engineering/multi_agents/platform_engineer/protocol_bindings/a2a/agent_executor.py

@@ -59,6 +59,10 @@ class StreamState:
     task_complete: bool = False
     user_input_required: bool = False
 
+    # Source agent tracking for sub-agent message grouping
+    current_agent: Optional[str] = None
+    agent_streaming_artifact_ids: Dict[str, str] = field(default_factory=dict)
+
 
 class AIPlatformEngineerA2AExecutor(AgentExecutor):
     """AI Platform Engineer A2A Executor."""
@@ -158,7 +162,7 @@ def _extract_final_answer(self, content: str) -> str:
             # Extract everything after the marker
             idx = content.find(marker)
             final_content = content[idx + len(marker):].strip()
-            logger.debug(f"Extracted final answer: {len(final_content)} chars (marker found at pos {idx})")
+            logger.info(f"Extracted final answer: {len(final_content)} chars (marker found at pos {idx})")
             return final_content
         return content
 
@@ -259,19 +263,44 @@ def _normalize_content(self, content) -> str:
     async def _send_artifact(self, event_queue: EventQueue, task: A2ATask,
                              artifact: Artifact, append: bool, last_chunk: bool = False):
         """Send an artifact update event."""
+        # Debug: Log artifact being sent
+        artifact_name = getattr(artifact, 'name', 'unknown')
+        # A2A stores text in parts[0].text, not in a top-level text attribute
+        parts = getattr(artifact, 'parts', [])
+        parts_text = None
+        if parts:
+            # Try different ways to access text
+            first_part = parts[0]
+            if hasattr(first_part, 'text'):
+                parts_text = first_part.text
+            elif hasattr(first_part, 'root') and hasattr(first_part.root, 'text'):
+                parts_text = first_part.root.text
+            elif isinstance(first_part, dict):
+                parts_text = first_part.get('text')
+        text_preview = parts_text[:100] if parts_text else '(no parts.text)'
+        text_len = len(parts_text) if parts_text else 0
+        
+        if artifact_name in ('final_result', 'partial_result'):
+            logger.info(f"📤 FINAL ARTIFACT: parts_count={len(parts)}, text_len={text_len}")
+            logger.info(f"📤 FINAL ARTIFACT preview: {text_preview}...")
+            # Debug: Log the actual artifact structure
+            logger.info(f"📤 FINAL ARTIFACT parts[0] type: {type(parts[0]) if parts else 'NO_PARTS'}")
+            logger.info(f"📤 FINAL ARTIFACT parts[0] attrs: {dir(parts[0]) if parts else 'NO_PARTS'}")
+
         await self._safe_enqueue_event(
             event_queue,
             TaskArtifactUpdateEvent(
                 append=append,
                 context_id=task.context_id,
                 task_id=task.id,
-                last_chunk=last_chunk,
+                lastChunk=last_chunk,
                 artifact=artifact,
             )
         )
 
     async def _send_completion(self, event_queue: EventQueue, task: A2ATask):
         """Send task completion status."""
+        logger.info(f"📤 Sending completion status for task {task.id}")
         await self._safe_enqueue_event(
             event_queue,
             TaskStatusUpdateEvent(
@@ -281,6 +310,7 @@ async def _send_completion(self, event_queue: EventQueue, task: A2ATask):
                 task_id=task.id,
             )
         )
+        logger.info(f"📤 Completion status enqueued for task {task.id}")
 
     async def _send_error(self, event_queue: EventQueue, task: A2ATask, error_msg: str):
         """Send task failure status."""
@@ -312,10 +342,20 @@ async def _handle_sub_agent_artifact(self, event: dict, state: StreamState,
         artifact_name = artifact_data.get('name', 'streaming_result')
         parts = artifact_data.get('parts', [])
 
+        # Extract sourceAgent from artifact metadata, event, or current state
+        existing_metadata = artifact_data.get('metadata', {})
+        source_agent = (
+            existing_metadata.get('sourceAgent') or
+            event.get('source_agent') or
+            state.current_agent or
+            'sub-agent'
+        )
+        logger.debug(f"📦 Sub-agent artifact from: {source_agent}")
+
         # Accumulate final results (complete_result, final_result, partial_result)
         if artifact_name in ('complete_result', 'final_result', 'partial_result'):
             state.sub_agents_completed += 1
-            logger.debug(f"Sub-agent completed with {artifact_name} (total completed: {state.sub_agents_completed})")
+            logger.info(f"Sub-agent completed with {artifact_name} (total completed: {state.sub_agents_completed})")
 
             for part in parts:
                 if isinstance(part, dict):
@@ -335,11 +375,17 @@ async def _handle_sub_agent_artifact(self, event: dict, state: StreamState,
                 elif part.get('data'):
                     artifact_parts.append(Part(root=DataPart(data=part['data'])))
 
+        # Create artifact with sourceAgent metadata for sub-agent message grouping
         artifact = Artifact(
             artifactId=artifact_data.get('artifactId'),
             name=artifact_name,
-            description=artifact_data.get('description', 'From sub-agent'),
-            parts=artifact_parts
+            description=artifact_data.get('description', f'From {source_agent}'),
+            parts=artifact_parts,
+            metadata={
+                'sourceAgent': source_agent,
+                'agentType': 'sub-agent',
+                **existing_metadata  # Preserve any other metadata
+            }
         )
 
         # Track artifact ID for append logic
@@ -474,14 +520,28 @@ async def _handle_streaming_chunk(self, event: dict, state: StreamState,
 
         is_tool_notification = self._is_tool_notification(content, event)
 
+        # Track current agent from tool_call events for sub-agent message grouping
+        if 'tool_call' in event:
+            tool_name = event['tool_call'].get('name', 'unknown')
+            state.current_agent = tool_name
+            logger.info(f"🎯 Current agent set to: {tool_name}")
+        elif 'tool_result' in event:
+            # Tool completed - keep current agent for any remaining content
+            tool_name = event['tool_result'].get('name', state.current_agent)
+            logger.info(f"✅ Tool completed: {tool_name}")
+
+        # Also detect agent from event metadata if provided
+        source_agent = event.get('source_agent') or state.current_agent or 'supervisor'
+
         # Accumulate non-notification content (unless DataPart already received)
         if not is_tool_notification and not state.sub_agent_datapart:
             state.supervisor_content.append(content)
 
-        # Create artifact
+        # Create artifact with sourceAgent metadata
         if is_tool_notification:
             artifact_name, description = self._get_artifact_name_for_notification(content, event)
             artifact = new_text_artifact(name=artifact_name, description=description, text=content)
+            artifact.metadata = {'sourceAgent': source_agent, 'agentType': 'notification'}
             use_append = False
             state.seen_artifact_ids.add(artifact.artifact_id)
         elif state.streaming_artifact_id is None:
@@ -491,6 +551,7 @@ async def _handle_streaming_chunk(self, event: dict, state: StreamState,
                 description='Streaming result',
                 text=content,
             )
+            artifact.metadata = {'sourceAgent': source_agent, 'agentType': 'streaming'}
             state.streaming_artifact_id = artifact.artifact_id
             state.seen_artifact_ids.add(artifact.artifact_id)
             state.first_artifact_sent = True
@@ -503,6 +564,7 @@ async def _handle_streaming_chunk(self, event: dict, state: StreamState,
                 text=content,
             )
             artifact.artifact_id = state.streaming_artifact_id
+            artifact.metadata = {'sourceAgent': source_agent, 'agentType': 'streaming'}
             use_append = True
 
         await self._send_artifact(event_queue, task, artifact, append=use_append)
@@ -514,7 +576,13 @@ async def _handle_stream_end(self, state: StreamState, task: A2ATask,
         # For single-agent scenarios where sub-agent already sent complete_result,
         # we just need to send the completion status (content already forwarded)
 
+        # Debug: Log accumulated content before getting final
+        logger.info(f"📦 Stream end - supervisor_content: {len(state.supervisor_content)} items, {sum(len(c) for c in state.supervisor_content)} chars")
+        logger.info(f"📦 Stream end - sub_agent_content: {len(state.sub_agent_content)} items, {sum(len(c) for c in state.sub_agent_content)} chars")
+        logger.info(f"📦 Stream end - sub_agents_completed: {state.sub_agents_completed}")
+
         final_content, is_datapart = self._get_final_content(state)
+        logger.info(f"📦 Final content for UI: {len(final_content) if isinstance(final_content, str) else 'datapart'} chars, is_datapart={is_datapart}")
 
         # If we have accumulated content (supervisor synthesis or sub-agent content), send it
         if final_content or is_datapart:
@@ -602,6 +670,29 @@ async def execute(self, context: RequestContext, event_queue: EventQueue) -> Non
                             artifact=event.artifact
                         )
                         await self._safe_enqueue_event(event_queue, transformed)
+
+                        # 🔧 CRITICAL FIX: Accumulate content from typed artifacts for final_result
+                        # Without this, _get_final_content returns empty and UI never gets final render
+                        artifact = event.artifact
+                        if artifact and hasattr(artifact, 'parts') and artifact.parts:
+                            artifact_name = getattr(artifact, 'name', 'streaming_result')
+                            is_final_artifact = artifact_name in ('complete_result', 'final_result', 'partial_result')
+
+                            for part in artifact.parts:
+                                part_root = getattr(part, 'root', None)
+                                if part_root and hasattr(part_root, 'text') and part_root.text:
+                                    # Accumulate streaming content
+                                    if artifact_name == 'streaming_result':
+                                        if not self._is_tool_notification(part_root.text, {}):
+                                            state.supervisor_content.append(part_root.text)
+                                    # Accumulate final results from sub-agents
+                                    elif is_final_artifact:
+                                        state.sub_agent_content.append(part_root.text)
+
+                            # Increment sub_agents_completed once per final artifact
+                            if is_final_artifact:
+                                state.sub_agents_completed += 1
+                                logger.info(f"Sub-agent completed via typed event with {artifact_name} (total: {state.sub_agents_completed})")
                     else:
                         corrected = TaskStatusUpdateEvent(
                             context_id=event.context_id,