Merge branch 'main' into feat/add-async-context-manager-base-client-674

Author: cchinchilla-dev

README.md

@@ -34,6 +34,18 @@
 
 ---
 
+## 🧩 Compatibility
+
+This SDK implements the A2A Protocol Specification [`v0.3.0`](https://a2a-protocol.org/v0.3.0/specification).
+
+| Transport | Client | Server |
+| :--- | :---: | :---: |
+| **JSON-RPC** | ✅ | ✅ |
+| **HTTP+JSON/REST** | ✅ | ✅ |
+| **GRPC** | ✅ | ✅ |
+
+---
+
 ## 🚀 Getting Started
 
 ### Prerequisites

pyproject.toml

@@ -105,6 +105,7 @@ dev = [
   "no_implicit_optional",
   "trio",
   "uvicorn>=0.35.0",
+  "pytest-timeout>=2.4.0",
   "a2a-sdk[all]",
 ]
 

src/a2a/client/errors.py

@@ -21,6 +21,14 @@ def __init__(self, status_code: int, message: str):
         self.message = message
         super().__init__(f'HTTP Error {status_code}: {message}')
 
+    def __repr__(self) -> str:
+        """Returns an unambiguous representation showing structured attributes."""
+        return (
+            f'{self.__class__.__name__}('
+            f'status_code={self.status_code!r}, '
+            f'message={self.message!r})'
+        )
+
 
 class A2AClientJSONError(A2AClientError):
     """Client exception for JSON errors during response parsing or validation."""
@@ -34,6 +42,10 @@ def __init__(self, message: str):
         self.message = message
         super().__init__(f'JSON Error: {message}')
 
+    def __repr__(self) -> str:
+        """Returns an unambiguous representation showing structured attributes."""
+        return f'{self.__class__.__name__}(message={self.message!r})'
+
 
 class A2AClientTimeoutError(A2AClientError):
     """Client exception for timeout errors during a request."""
@@ -47,6 +59,10 @@ def __init__(self, message: str):
         self.message = message
         super().__init__(f'Timeout Error: {message}')
 
+    def __repr__(self) -> str:
+        """Returns an unambiguous representation showing structured attributes."""
+        return f'{self.__class__.__name__}(message={self.message!r})'
+
 
 class A2AClientInvalidArgsError(A2AClientError):
     """Client exception for invalid arguments passed to a method."""
@@ -60,6 +76,10 @@ def __init__(self, message: str):
         self.message = message
         super().__init__(f'Invalid arguments error: {message}')
 
+    def __repr__(self) -> str:
+        """Returns an unambiguous representation showing structured attributes."""
+        return f'{self.__class__.__name__}(message={self.message!r})'
+
 
 class A2AClientInvalidStateError(A2AClientError):
     """Client exception for an invalid client state."""
@@ -73,6 +93,10 @@ def __init__(self, message: str):
         self.message = message
         super().__init__(f'Invalid state error: {message}')
 
+    def __repr__(self) -> str:
+        """Returns an unambiguous representation showing structured attributes."""
+        return f'{self.__class__.__name__}(message={self.message!r})'
+
 
 class A2AClientJSONRPCError(A2AClientError):
     """Client exception for JSON-RPC errors returned by the server."""
@@ -85,3 +109,7 @@ def __init__(self, error: JSONRPCErrorResponse):
         """
         self.error = error.error
         super().__init__(f'JSON-RPC Error {error.error}')
+
+    def __repr__(self) -> str:
+        """Returns an unambiguous representation showing the JSON-RPC error object."""
+        return f'{self.__class__.__name__}({self.error!r})'

src/a2a/client/transports/base.py

@@ -1,5 +1,8 @@
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Callable
+from types import TracebackType
+
+from typing_extensions import Self
 
 from a2a.client.middleware import ClientCallContext
 from a2a.types import (
@@ -19,6 +22,19 @@
 class ClientTransport(ABC):
     """Abstract base class for a client transport."""
 
+    async def __aenter__(self) -> Self:
+        """Enters the async context manager, returning the transport itself."""
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exits the async context manager, ensuring close() is called."""
+        await self.close()
+
     @abstractmethod
     async def send_message(
         self,

src/a2a/server/request_handlers/default_request_handler.py

@@ -330,6 +330,7 @@ async def push_notification_callback() -> None:
 
         except Exception:
             logger.exception('Agent execution failed')
+            producer_task.cancel()
             raise
         finally:
             if interrupted_or_non_blocking:
@@ -435,7 +436,12 @@ async def _cleanup_producer(
         task_id: str,
     ) -> None:
         """Cleans up the agent execution task and queue manager entry."""
-        await producer_task
+        try:
+            await producer_task
+        except asyncio.CancelledError:
+            logger.debug(
+                'Producer task %s was cancelled during cleanup', task_id
+            )
         await self._queue_manager.close(task_id)
         async with self._running_agents_lock:
             self._running_agents.pop(task_id, None)

src/a2a/server/request_handlers/grpc_handler.py

@@ -3,17 +3,12 @@
 import logging
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable, Awaitable, Sequence
-from typing import TYPE_CHECKING
+from collections.abc import AsyncIterable, Awaitable
 
 
 try:
     import grpc
     import grpc.aio
-
-    if TYPE_CHECKING:
-        from grpc.aio._typing import MetadataType
-    from grpc.aio import Metadata
 except ImportError as e:
     raise ImportError(
         'GrpcHandler requires grpcio and grpcio-tools to be installed. '
@@ -56,14 +51,16 @@ def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
 def _get_metadata_value(
     context: grpc.aio.ServicerContext, key: str
 ) -> list[str]:
-    md: MetadataType | None = context.invocation_metadata()
-    raw_values: list[str | bytes] = []
+    md = context.invocation_metadata()
+    if md is None:
+        return []
+
     lower_key = key.lower()
-    if isinstance(md, Metadata):
-        raw_values = md.get_all(lower_key)
-    elif isinstance(md, Sequence):
-        raw_values = [e for (k, e) in md if k.lower() == lower_key]
-    return [e if isinstance(e, str) else e.decode('utf-8') for e in raw_values]
+    return [
+        e if isinstance(e, str) else e.decode('utf-8')
+        for k, e in md
+        if k.lower() == lower_key
+    ]
 
 
 class DefaultCallContextBuilder(CallContextBuilder):

tests/client/test_errors.py

@@ -3,6 +3,13 @@
 import pytest
 
 from a2a.client import A2AClientError, A2AClientHTTPError, A2AClientJSONError
+from a2a.client.errors import (
+    A2AClientInvalidArgsError,
+    A2AClientInvalidStateError,
+    A2AClientJSONRPCError,
+    A2AClientTimeoutError,
+)
+from a2a.types import JSONRPCError, JSONRPCErrorResponse
 
 
 class TestA2AClientError:
@@ -35,6 +42,14 @@ def test_message_formatting(self) -> None:
         error = A2AClientHTTPError(500, 'Internal Server Error')
         assert str(error) == 'HTTP Error 500: Internal Server Error'
 
+    def test_repr(self) -> None:
+        """Test that __repr__ shows structured attributes."""
+        error = A2AClientHTTPError(404, 'Not Found')
+        assert (
+            repr(error)
+            == "A2AClientHTTPError(status_code=404, message='Not Found')"
+        )
+
     def test_inheritance(self) -> None:
         """Test that A2AClientHTTPError inherits from A2AClientError."""
         error = A2AClientHTTPError(400, 'Bad Request')
@@ -81,6 +96,13 @@ def test_message_formatting(self) -> None:
         error = A2AClientJSONError('Missing required field')
         assert str(error) == 'JSON Error: Missing required field'
 
+    def test_repr(self) -> None:
+        """Test that __repr__ shows structured attributes."""
+        error = A2AClientJSONError('Invalid JSON format')
+        assert (
+            repr(error) == "A2AClientJSONError(message='Invalid JSON format')"
+        )
+
     def test_inheritance(self) -> None:
         """Test that A2AClientJSONError inherits from A2AClientError."""
         error = A2AClientJSONError('Parsing error')
@@ -108,6 +130,57 @@ def test_with_various_messages(self) -> None:
             assert str(error) == f'JSON Error: {message}'
 
 
+class TestA2AClientTimeoutErrorRepr:
+    """Test __repr__ for A2AClientTimeoutError."""
+
+    def test_repr(self) -> None:
+        """Test that __repr__ shows structured attributes."""
+        error = A2AClientTimeoutError('Request timed out')
+        assert (
+            repr(error) == "A2AClientTimeoutError(message='Request timed out')"
+        )
+
+
+class TestA2AClientInvalidArgsErrorRepr:
+    """Test __repr__ for A2AClientInvalidArgsError."""
+
+    def test_repr(self) -> None:
+        """Test that __repr__ shows structured attributes."""
+        error = A2AClientInvalidArgsError('Missing required param')
+        assert (
+            repr(error)
+            == "A2AClientInvalidArgsError(message='Missing required param')"
+        )
+
+
+class TestA2AClientInvalidStateErrorRepr:
+    """Test __repr__ for A2AClientInvalidStateError."""
+
+    def test_repr(self) -> None:
+        """Test that __repr__ shows structured attributes."""
+        error = A2AClientInvalidStateError('Client not initialized')
+        assert (
+            repr(error)
+            == "A2AClientInvalidStateError(message='Client not initialized')"
+        )
+
+
+class TestA2AClientJSONRPCErrorRepr:
+    """Test __repr__ for A2AClientJSONRPCError."""
+
+    def test_repr(self) -> None:
+        """Test that __repr__ shows the JSON-RPC error object."""
+        response = JSONRPCErrorResponse(
+            id='test-1',
+            error=JSONRPCError(code=-32601, message='Method not found'),
+        )
+        error = A2AClientJSONRPCError(response)
+        assert (
+            repr(error)
+            == "A2AClientJSONRPCError(JSONRPCError(code=-32601, data=None, message='Method not found'))"
+        )
+
+
 class TestExceptionHierarchy:
     """Test the exception hierarchy and relationships."""
 

tests/server/request_handlers/test_default_request_handler.py

@@ -2644,3 +2644,54 @@ async def test_on_message_send_stream_task_id_provided_but_task_not_found():
         f'Task {task_id} was specified but does not exist'
         in exc_info.value.error.message
     )
+
+
+class HelloWorldAgentExecutor(AgentExecutor):
+    """Test Agent Implementation."""
+
+    async def execute(
+        self,
+        context: RequestContext,
+        event_queue: EventQueue,
+    ) -> None:
+        updater = TaskUpdater(
+            event_queue,
+            task_id=context.task_id or str(uuid.uuid4()),
+            context_id=context.context_id or str(uuid.uuid4()),
+        )
+        await updater.update_status(TaskState.working)
+        await updater.complete()
+
+    async def cancel(
+        self, context: RequestContext, event_queue: EventQueue
+    ) -> None:
+        raise NotImplementedError('cancel not supported')
+
+
+# Repro is straight from the https://github.com/a2aproject/a2a-python/issues/609.
+# It uses timeout to test against infinite wait, if it's going to be flaky,
+# we should reconsider the approach.
+@pytest.mark.asyncio
+@pytest.mark.timeout(1)
+async def test_on_message_send_error_does_not_hang():
+    """Test that if the consumer raises an exception during blocking wait, the producer is cancelled and no deadlock occurs."""
+    agent = HelloWorldAgentExecutor()
+    task_store = AsyncMock(spec=TaskStore)
+    task_store.save.side_effect = RuntimeError('This is an Error!')
+
+    request_handler = DefaultRequestHandler(
+        agent_executor=agent, task_store=task_store
+    )
+
+    params = MessageSendParams(
+        message=Message(
+            role=Role.user,
+            message_id='msg_error_blocking',
+            parts=[Part(root=TextPart(text='Test message'))],
+        )
+    )
+
+    with pytest.raises(RuntimeError, match='This is an Error!'):
+        await request_handler.on_message_send(
+            params, create_server_call_context()
+        )

tests/server/request_handlers/test_jsonrpc_handler.py

@@ -322,7 +322,6 @@ async def streaming_coro():
 
             self.assertIsInstance(response.root, JSONRPCErrorResponse)
             assert response.root.error == UnsupportedOperationError()  # type: ignore
-            mock_agent_executor.execute.assert_called_once()
 
     @patch(
         'a2a.server.agent_execution.simple_request_context_builder.SimpleRequestContextBuilder.build'