add A2AStreamEvent + handle partial responses in stream_async + modify constructor parameters

awsarron · awsarron · commit 48a006f02243 · 2025-11-18T14:07:54.000-05:00
diff --git a/src/strands/agent/a2a_agent.py b/src/strands/agent/a2a_agent.py
@@ -1,18 +1,22 @@
 """A2A Agent client for Strands Agents.
 
 This module provides the A2AAgent class, which acts as a client wrapper for remote A2A agents,
-allowing them to be used in graphs, swarms, and other multi-agent patterns.
+allowing them to be used standalone or as part of multi-agent patterns.
+
+A2AAgent can be used to get the Agent Card and interact with the agent.
 """
 
 import logging
 from typing import Any, AsyncIterator
 
 import httpx
-from a2a.client import A2ACardResolver, ClientConfig, ClientFactory
-from a2a.types import AgentCard
+from a2a.client import A2ACardResolver, Client, ClientConfig, ClientFactory
+from a2a.types import AgentCard, Message, TaskArtifactUpdateEvent, TaskState, TaskStatusUpdateEvent
 
 from .._async import run_async
 from ..multiagent.a2a.converters import convert_input_to_message, convert_response_to_agent_result
+from ..types._events import AgentResultEvent
+from ..types.a2a import A2AResponse, A2AStreamEvent
 from ..types.agent import AgentInput
 from .agent_result import AgentResult
 
@@ -22,59 +26,46 @@
 
 
 class A2AAgent:
-    """Client wrapper for remote A2A agents.
-
-    Implements the AgentBase protocol to enable remote A2A agents to be used
-    in graphs, swarms, and other multi-agent patterns.
-    """
+    """Client wrapper for remote A2A agents."""
 
     def __init__(
         self,
         endpoint: str,
+        *,
+        name: str | None = None,
+        description: str = "",
         timeout: int = DEFAULT_TIMEOUT,
-        httpx_client_args: dict[str, Any] | None = None,
+        a2a_client_factory: ClientFactory | None = None,
     ):
-        """Initialize A2A agent client.
+        """Initialize A2A agent.
 
         Args:
-            endpoint: The base URL of the remote A2A agent
-            timeout: Timeout for HTTP operations in seconds (defaults to 300)
-            httpx_client_args: Optional dictionary of arguments to pass to httpx.AsyncClient
-                constructor. Allows custom auth, headers, proxies, etc.
-                Example: {"headers": {"Authorization": "Bearer token"}}
+            endpoint: The base URL of the remote A2A agent.
+            name: Agent name. If not provided, will be populated from agent card.
+            description: Agent description. If empty, will be populated from agent card.
+            timeout: Timeout for HTTP operations in seconds (defaults to 300).
+            a2a_client_factory: Optional pre-configured A2A ClientFactory. If provided,
+                it will be used to create the A2A client after discovering the agent card.
         """
         self.endpoint = endpoint
+        self.name = name
+        self.description = description
         self.timeout = timeout
-        self._httpx_client_args: dict[str, Any] = httpx_client_args or {}
-
-        if "timeout" not in self._httpx_client_args:
-            self._httpx_client_args["timeout"] = self.timeout
-
+        self._httpx_client: httpx.AsyncClient | None = None
+        self._owns_client = a2a_client_factory is None
         self._agent_card: AgentCard | None = None
+        self._a2a_client: Client | None = None
+        self._a2a_client_factory: ClientFactory | None = a2a_client_factory
 
     def _get_httpx_client(self) -> httpx.AsyncClient:
-        """Get a fresh httpx client for the current operation.
+        """Get or create the httpx client for this agent.
 
         Returns:
             Configured httpx.AsyncClient instance.
         """
-        return httpx.AsyncClient(**self._httpx_client_args)
-
-    def _get_client_factory(self, streaming: bool = False) -> ClientFactory:
-        """Get a ClientFactory for the current operation.
-
-        Args:
-            streaming: Whether to enable streaming mode.
-
-        Returns:
-            Configured ClientFactory instance.
-        """
-        httpx_client = self._get_httpx_client()
-        config = ClientConfig(
-            httpx_client=httpx_client,
-            streaming=streaming,
-        )
-        return ClientFactory(config)
+        if self._httpx_client is None:
+            self._httpx_client = httpx.AsyncClient(timeout=self.timeout)
+        return self._httpx_client
 
     async def _get_agent_card(self) -> AgentCard:
         """Discover and cache the agent card from the remote endpoint.
@@ -88,15 +79,44 @@ async def _get_agent_card(self) -> AgentCard:
         httpx_client = self._get_httpx_client()
         resolver = A2ACardResolver(httpx_client=httpx_client, base_url=self.endpoint)
         self._agent_card = await resolver.get_agent_card()
-        logger.info("endpoint=<%s> | discovered agent card", self.endpoint)
+
+        # Populate name from card if not set
+        if self.name is None and self._agent_card.name:
+            self.name = self._agent_card.name
+
+        # Populate description from card if not set
+        if not self.description and self._agent_card.description:
+            self.description = self._agent_card.description
+
+        logger.info("agent=<%s>, endpoint=<%s> | discovered agent card", self.name, self.endpoint)
         return self._agent_card
 
-    async def _send_message(self, prompt: AgentInput, streaming: bool) -> AsyncIterator[Any]:
+    async def _get_a2a_client(self) -> Client:
+        """Get or create the A2A client for this agent.
+
+        Returns:
+            Configured A2A client instance.
+        """
+        if self._a2a_client is None:
+            agent_card = await self._get_agent_card()
+
+            if self._a2a_client_factory is not None:
+                # Use provided factory
+                factory = self._a2a_client_factory
+            else:
+                # Create default factory
+                httpx_client = self._get_httpx_client()
+                config = ClientConfig(httpx_client=httpx_client, streaming=False)
+                factory = ClientFactory(config)
+
+            self._a2a_client = factory.create(agent_card)
+        return self._a2a_client
+
+    async def _send_message(self, prompt: AgentInput) -> AsyncIterator[A2AResponse]:
         """Send message to A2A agent.
 
         Args:
             prompt: Input to send to the agent.
-            streaming: Whether to use streaming mode.
 
         Returns:
             Async iterator of A2A events.
@@ -107,13 +127,46 @@ async def _send_message(self, prompt: AgentInput, streaming: bool) -> AsyncItera
         if prompt is None:
             raise ValueError("prompt is required for A2AAgent")
 
-        agent_card = await self._get_agent_card()
-        client = self._get_client_factory(streaming=streaming).create(agent_card)
+        client = await self._get_a2a_client()
         message = convert_input_to_message(prompt)
 
-        logger.info("endpoint=<%s> | %s message", self.endpoint, "streaming" if streaming else "sending")
+        logger.info("agent=<%s>, endpoint=<%s> | sending message", self.name, self.endpoint)
         return client.send_message(message)
 
+    def _is_complete_event(self, event: A2AResponse) -> bool:
+        """Check if an A2A event represents a complete response.
+
+        Args:
+            event: A2A event.
+
+        Returns:
+            True if the event represents a complete response.
+        """
+        # Direct Message is always complete
+        if isinstance(event, Message):
+            return True
+
+        # Handle tuple responses (Task, UpdateEvent | None)
+        if isinstance(event, tuple) and len(event) == 2:
+            task, update_event = event
+
+            # Initial task response (no update event)
+            if update_event is None:
+                return True
+
+            # Artifact update with last_chunk flag
+            if isinstance(update_event, TaskArtifactUpdateEvent):
+                if hasattr(update_event, "last_chunk") and update_event.last_chunk is not None:
+                    return update_event.last_chunk
+                return False
+
+            # Status update with completed state
+            if isinstance(update_event, TaskStatusUpdateEvent):
+                if update_event.status and hasattr(update_event.status, "state"):
+                    return update_event.status.state == TaskState.completed
+
+        return False
+
     async def invoke_async(
         self,
         prompt: AgentInput = None,
@@ -132,7 +185,7 @@ async def invoke_async(
             ValueError: If prompt is None.
             RuntimeError: If no response received from agent.
         """
-        async for event in await self._send_message(prompt, streaming=False):
+        async for event in await self._send_message(prompt):
             return convert_response_to_agent_result(event)
 
         raise RuntimeError("No response received from A2A agent")
@@ -169,10 +222,32 @@ async def stream_async(
             **kwargs: Additional arguments (ignored).
 
         Yields:
-            A2A events wrapped in dictionaries with an 'a2a_event' key.
+            A2A events and a final AgentResult event.
 
         Raises:
             ValueError: If prompt is None.
         """
-        async for event in await self._send_message(prompt, streaming=True):
-            yield {"a2a_event": event}
+        last_event = None
+        last_complete_event = None
+
+        async for event in await self._send_message(prompt):
+            last_event = event
+            if self._is_complete_event(event):
+                last_complete_event = event
+            yield A2AStreamEvent(event)
+
+        # Use the last complete event if available, otherwise fall back to last event
+        final_event = last_complete_event if last_complete_event is not None else last_event
+
+        if final_event is not None:
+            result = convert_response_to_agent_result(final_event)
+            yield AgentResultEvent(result)
+
+    def __del__(self) -> None:
+        """Clean up resources when agent is garbage collected."""
+        if self._owns_client and self._httpx_client is not None:
+            try:
+                client = self._httpx_client
+                run_async(lambda: client.aclose())
+            except Exception:
+                pass  # Best effort cleanup, ignore errors in __del__
diff --git a/src/strands/multiagent/a2a/converters.py b/src/strands/multiagent/a2a/converters.py
@@ -1,18 +1,17 @@
 """Conversion functions between Strands and A2A types."""
 
-from typing import TypeAlias, cast
+from typing import cast
 from uuid import uuid4
 
 from a2a.types import Message as A2AMessage
-from a2a.types import Part, Role, Task, TaskArtifactUpdateEvent, TaskStatusUpdateEvent, TextPart
+from a2a.types import Part, Role, TaskArtifactUpdateEvent, TaskStatusUpdateEvent, TextPart
 
 from ...agent.agent_result import AgentResult
 from ...telemetry.metrics import EventLoopMetrics
+from ...types.a2a import A2AResponse
 from ...types.agent import AgentInput
 from ...types.content import ContentBlock, Message
 
-A2AResponse: TypeAlias = tuple[Task, TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None] | A2AMessage
-
 
 def convert_input_to_message(prompt: AgentInput) -> A2AMessage:
     """Convert AgentInput to A2A Message.
@@ -89,7 +88,21 @@ def convert_response_to_agent_result(response: A2AResponse) -> AgentResult:
 
     if isinstance(response, tuple) and len(response) == 2:
         task, update_event = response
-        if update_event is None and task and hasattr(task, "artifacts") and task.artifacts is not None:
+
+        # Handle artifact updates
+        if isinstance(update_event, TaskArtifactUpdateEvent):
+            if update_event.artifact and hasattr(update_event.artifact, "parts"):
+                for part in update_event.artifact.parts:
+                    if hasattr(part, "root") and hasattr(part.root, "text"):
+                        content.append({"text": part.root.text})
+        # Handle status updates with messages
+        elif isinstance(update_event, TaskStatusUpdateEvent):
+            if update_event.status and hasattr(update_event.status, "message") and update_event.status.message:
+                for part in update_event.status.message.parts:
+                    if hasattr(part, "root") and hasattr(part.root, "text"):
+                        content.append({"text": part.root.text})
+        # Handle initial task or task without update event
+        elif update_event is None and task and hasattr(task, "artifacts") and task.artifacts is not None:
             for artifact in task.artifacts:
                 if hasattr(artifact, "parts"):
                     for part in artifact.parts:
diff --git a/src/strands/types/a2a.py b/src/strands/types/a2a.py
@@ -0,0 +1,26 @@
+"""Additional A2A types."""
+
+from typing import TypeAlias
+
+from a2a.types import Message, Task, TaskArtifactUpdateEvent, TaskStatusUpdateEvent
+
+from ._events import TypedEvent
+
+A2AResponse: TypeAlias = tuple[Task, TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None] | Message
+
+
+class A2AStreamEvent(TypedEvent):
+    """Event that wraps streamed A2A types."""
+
+    def __init__(self, a2a_event: A2AResponse) -> None:
+        """Initialize with A2A event.
+
+        Args:
+            a2a_event: The original A2A event (Task tuple or Message)
+        """
+        super().__init__(
+            {
+                "type": "a2a_stream",
+                "event": a2a_event,  # Nest A2A event to avoid field conflicts
+            }
+        )
diff --git a/tests/strands/agent/test_a2a_agent.py b/tests/strands/agent/test_a2a_agent.py
