Merge branch 'main' into poc/agent-catalog

Author: holtskinner

src/a2a/client/client.py

@@ -1,26 +1,37 @@
 import json
 import logging
+
 from collections.abc import AsyncGenerator
 from typing import Any
 from uuid import uuid4
 
 import httpx
+
 from httpx_sse import SSEError, aconnect_sse
 from pydantic import ValidationError
 
 from a2a.client.errors import A2AClientHTTPError, A2AClientJSONError
-from a2a.types import (AgentCard, CancelTaskRequest, CancelTaskResponse,
-                       GetTaskPushNotificationConfigRequest,
-                       GetTaskPushNotificationConfigResponse, GetTaskRequest,
-                       GetTaskResponse, SendMessageRequest,
-                       SendMessageResponse, SendStreamingMessageRequest,
-                       SendStreamingMessageResponse,
-                       SetTaskPushNotificationConfigRequest,
-                       SetTaskPushNotificationConfigResponse)
+from a2a.types import (
+    AgentCard,
+    CancelTaskRequest,
+    CancelTaskResponse,
+    GetTaskPushNotificationConfigRequest,
+    GetTaskPushNotificationConfigResponse,
+    GetTaskRequest,
+    GetTaskResponse,
+    SendMessageRequest,
+    SendMessageResponse,
+    SendStreamingMessageRequest,
+    SendStreamingMessageResponse,
+    SetTaskPushNotificationConfigRequest,
+    SetTaskPushNotificationConfigResponse,
+)
 from a2a.utils.telemetry import SpanKind, trace_class
 
+
 logger = logging.getLogger(__name__)
 
+
 class A2ACardResolver:
     """Agent Card resolver."""
 
@@ -160,6 +171,7 @@ async def get_client_from_agent_card_url(
             agent_card_path: The path to the agent card endpoint, relative to the base URL.
             http_kwargs: Optional dictionary of keyword arguments to pass to the
                 underlying httpx.get request when fetching the agent card.
+
         Returns:
             An initialized `A2AClient` instance.
 
@@ -169,7 +181,9 @@ async def get_client_from_agent_card_url(
         """
         agent_card: AgentCard = await A2ACardResolver(
             httpx_client, base_url=base_url, agent_card_path=agent_card_path
-        ).get_agent_card(http_kwargs=http_kwargs) # Fetches public card by default
+        ).get_agent_card(
+            http_kwargs=http_kwargs
+        )  # Fetches public card by default
         return A2AClient(httpx_client=httpx_client, agent_card=agent_card)
 
     async def send_message(

src/a2a/server/apps/starlette_app.py

@@ -51,8 +51,20 @@
 
 logger = logging.getLogger(__name__)
 
-# Register Starlette User as an implementation of a2a.auth.user.User
-A2AUser.register(BaseUser)
+
+class StarletteUserProxy(A2AUser):
+    """Adapts the Starlette User class to the A2A user representation."""
+
+    def __init__(self, user: BaseUser):
+        self._user = user
+
+    @property
+    def is_authenticated(self):
+        return self._user.is_authenticated
+
+    @property
+    def user_name(self):
+        return self._user.display_name
 
 
 class CallContextBuilder(ABC):
@@ -70,7 +82,7 @@ def build(self, request: Request) -> ServerCallContext:
         user = UnauthenticatedUser()
         state = {}
         with contextlib.suppress(Exception):
-            user = request.user
+            user = StarletteUserProxy(request.user)
             state['auth'] = request.auth
         return ServerCallContext(user=user, state=state)
 

src/a2a/server/events/event_queue.py

@@ -14,12 +14,7 @@
 logger = logging.getLogger(__name__)
 
 
-Event = (
-    Message
-    | Task
-    | TaskStatusUpdateEvent
-    | TaskArtifactUpdateEvent
-)
+Event = Message | Task | TaskStatusUpdateEvent | TaskArtifactUpdateEvent
 """Type alias for events that can be enqueued."""
 
 

src/a2a/server/tasks/task_updater.py

@@ -1,5 +1,6 @@
 import uuid
 
+from datetime import datetime, timezone
 from typing import Any
 
 from a2a.server.events import EventQueue
@@ -34,15 +35,23 @@ def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
         self.context_id = context_id
 
     def update_status(
-        self, state: TaskState, message: Message | None = None, final=False
+        self,
+        state: TaskState,
+        message: Message | None = None,
+        final=False,
+        timestamp: str | None = None,
     ):
         """Updates the status of the task and publishes a `TaskStatusUpdateEvent`.
 
         Args:
             state: The new state of the task.
             message: An optional message associated with the status update.
             final: If True, indicates this is the final status update for the task.
+            timestamp: Optional ISO 8601 datetime string. Defaults to current time.
         """
+        current_timestamp = (
+            timestamp if timestamp else datetime.now(timezone.utc).isoformat()
+        )
         self.event_queue.enqueue_event(
             TaskStatusUpdateEvent(
                 taskId=self.task_id,
@@ -51,6 +60,7 @@ def update_status(
                 status=TaskStatus(
                     state=state,
                     message=message,
+                    timestamp=current_timestamp,
                 ),
             )
         )
@@ -97,6 +107,10 @@ def failed(self, message: Message | None = None):
         """Marks the task as failed and publishes a final status update."""
         self.update_status(TaskState.failed, message=message, final=True)
 
+    def reject(self, message: Message | None = None):
+        """Marks the task as rejected and publishes a final status update."""
+        self.update_status(TaskState.rejected, message=message, final=True)
+
     def submit(self, message: Message | None = None):
         """Marks the task as submitted and publishes a status update."""
         self.update_status(

src/a2a/types.py

@@ -648,6 +648,10 @@ class PushNotificationConfig(BaseModel):
     """
 
     authentication: PushNotificationAuthenticationInfo | None = None
+    id: str | None = None
+    """
+    Push Notification ID - created by server to support multiple callbacks
+    """
     token: str | None = None
     """
     Token unique to this task/session.

tests/client/test_client.py

@@ -1,24 +1,45 @@
 import json
+
 from collections.abc import AsyncGenerator
 from typing import Any
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import httpx
 import pytest
+
 from httpx_sse import EventSource, ServerSentEvent
-from pydantic import ValidationError as PydanticValidationError
-
-from a2a.client import (A2ACardResolver, A2AClient, A2AClientHTTPError,
-                        A2AClientJSONError, create_text_message_object)
-from a2a.types import (A2ARequest, AgentCapabilities, AgentCard, AgentSkill,
-                       CancelTaskRequest, CancelTaskResponse,
-                       CancelTaskSuccessResponse, GetTaskRequest,
-                       GetTaskResponse, InvalidParamsError,
-                       JSONRPCErrorResponse, MessageSendParams, Role,
-                       SendMessageRequest, SendMessageResponse,
-                       SendMessageSuccessResponse, SendStreamingMessageRequest,
-                       SendStreamingMessageResponse, TaskIdParams,
-                       TaskNotCancelableError, TaskQueryParams)
+
+from a2a.client import (
+    A2ACardResolver,
+    A2AClient,
+    A2AClientHTTPError,
+    A2AClientJSONError,
+    create_text_message_object,
+)
+from a2a.types import (
+    A2ARequest,
+    AgentCapabilities,
+    AgentCard,
+    AgentSkill,
+    CancelTaskRequest,
+    CancelTaskResponse,
+    CancelTaskSuccessResponse,
+    GetTaskRequest,
+    GetTaskResponse,
+    InvalidParamsError,
+    JSONRPCErrorResponse,
+    MessageSendParams,
+    Role,
+    SendMessageRequest,
+    SendMessageResponse,
+    SendMessageSuccessResponse,
+    SendStreamingMessageRequest,
+    SendStreamingMessageResponse,
+    TaskIdParams,
+    TaskNotCancelableError,
+    TaskQueryParams,
+)
+
 
 AGENT_CARD = AgentCard(
     name='Hello World Agent',
@@ -100,7 +121,9 @@ class TestA2ACardResolver:
     BASE_URL = 'http://example.com'
     AGENT_CARD_PATH = '/.well-known/agent.json'
     FULL_AGENT_CARD_URL = f'{BASE_URL}{AGENT_CARD_PATH}'
-    EXTENDED_AGENT_CARD_PATH = '/agent/authenticatedExtendedCard' # Default path
+    EXTENDED_AGENT_CARD_PATH = (
+        '/agent/authenticatedExtendedCard'  # Default path
+    )
 
     @pytest.mark.asyncio
     async def test_init_strips_slashes(self, mock_httpx_client: AsyncMock):
@@ -141,11 +164,12 @@ async def test_get_agent_card_success_public_only(
 
     @pytest.mark.asyncio
     async def test_get_agent_card_success_with_specified_path_for_extended_card(
-        self, mock_httpx_client: AsyncMock):
+        self, mock_httpx_client: AsyncMock
+    ):
         extended_card_response = AsyncMock(spec=httpx.Response)
         extended_card_response.status_code = 200
-        extended_card_response.json.return_value = AGENT_CARD_EXTENDED.model_dump(
-            mode='json'
+        extended_card_response.json.return_value = (
+            AGENT_CARD_EXTENDED.model_dump(mode='json')
         )
 
         # Mock the single call for the extended card
@@ -156,20 +180,26 @@ async def test_get_agent_card_success_with_specified_path_for_extended_card(
             base_url=self.BASE_URL,
             agent_card_path=self.AGENT_CARD_PATH,
         )
-        
+
         # Fetch the extended card by providing its relative path and example auth
-        auth_kwargs = {"headers": {"Authorization": "Bearer test token"}}
+        auth_kwargs = {'headers': {'Authorization': 'Bearer test token'}}
         agent_card_result = await resolver.get_agent_card(
             relative_card_path=self.EXTENDED_AGENT_CARD_PATH,
-            http_kwargs=auth_kwargs
+            http_kwargs=auth_kwargs,
         )
 
-        expected_extended_url = f'{self.BASE_URL}/{self.EXTENDED_AGENT_CARD_PATH.lstrip("/")}'
-        mock_httpx_client.get.assert_called_once_with(expected_extended_url, **auth_kwargs)
+        expected_extended_url = (
+            f'{self.BASE_URL}/{self.EXTENDED_AGENT_CARD_PATH.lstrip("/")}'
+        )
+        mock_httpx_client.get.assert_called_once_with(
+            expected_extended_url, **auth_kwargs
+        )
         extended_card_response.raise_for_status.assert_called_once()
 
         assert isinstance(agent_card_result, AgentCard)
-        assert agent_card_result == AGENT_CARD_EXTENDED # Should return the extended card
+        assert (
+            agent_card_result == AGENT_CARD_EXTENDED
+        )  # Should return the extended card
 
     @pytest.mark.asyncio
     async def test_get_agent_card_validation_error(
@@ -178,19 +208,29 @@ async def test_get_agent_card_validation_error(
         mock_response = AsyncMock(spec=httpx.Response)
         mock_response.status_code = 200
         # Data that will cause a Pydantic ValidationError
-        mock_response.json.return_value = {"invalid_field": "value", "name": "Test Agent"} 
+        mock_response.json.return_value = {
+            'invalid_field': 'value',
+            'name': 'Test Agent',
+        }
         mock_httpx_client.get.return_value = mock_response
 
         resolver = A2ACardResolver(
             httpx_client=mock_httpx_client, base_url=self.BASE_URL
         )
         # The call that is expected to raise an error should be within pytest.raises
         with pytest.raises(A2AClientJSONError) as exc_info:
-            await resolver.get_agent_card() # Fetches from default path
-        
-        assert f'Failed to validate agent card structure from {self.FULL_AGENT_CARD_URL}' in str(exc_info.value)
-        assert 'invalid_field' in str(exc_info.value) # Check if Pydantic error details are present
-        assert mock_httpx_client.get.call_count == 1 # Should only be called once
+            await resolver.get_agent_card()  # Fetches from default path
+
+        assert (
+            f'Failed to validate agent card structure from {self.FULL_AGENT_CARD_URL}'
+            in str(exc_info.value)
+        )
+        assert 'invalid_field' in str(
+            exc_info.value
+        )  # Check if Pydantic error details are present
+        assert (
+            mock_httpx_client.get.call_count == 1
+        )  # Should only be called once
 
     @pytest.mark.asyncio
     async def test_get_agent_card_http_status_error(
@@ -217,7 +257,10 @@ async def test_get_agent_card_http_status_error(
             await resolver.get_agent_card()
 
         assert exc_info.value.status_code == 404
-        assert f'Failed to fetch agent card from {self.FULL_AGENT_CARD_URL}' in str(exc_info.value)
+        assert (
+            f'Failed to fetch agent card from {self.FULL_AGENT_CARD_URL}'
+            in str(exc_info.value)
+        )
         assert 'Not Found' in str(exc_info.value)
         mock_httpx_client.get.assert_called_once_with(self.FULL_AGENT_CARD_URL)
 
@@ -242,7 +285,10 @@ async def test_get_agent_card_json_decode_error(
             await resolver.get_agent_card()
 
         # Assertions using exc_info must be after the with block
-        assert f'Failed to parse JSON for agent card from {self.FULL_AGENT_CARD_URL}' in str(exc_info.value)
+        assert (
+            f'Failed to parse JSON for agent card from {self.FULL_AGENT_CARD_URL}'
+            in str(exc_info.value)
+        )
         assert 'Expecting value' in str(exc_info.value)
         mock_httpx_client.get.assert_called_once_with(self.FULL_AGENT_CARD_URL)
 
@@ -263,7 +309,10 @@ async def test_get_agent_card_request_error(
             await resolver.get_agent_card()
 
         assert exc_info.value.status_code == 503
-        assert f'Network communication error fetching agent card from {self.FULL_AGENT_CARD_URL}' in str(exc_info.value)
+        assert (
+            f'Network communication error fetching agent card from {self.FULL_AGENT_CARD_URL}'
+            in str(exc_info.value)
+        )
         assert 'Network issue' in str(exc_info.value)
         mock_httpx_client.get.assert_called_once_with(self.FULL_AGENT_CARD_URL)
 

tests/server/tasks/test_task_updater.py

@@ -212,9 +212,7 @@ def test_new_agent_message(self, task_updater, sample_parts):
         assert message.parts == sample_parts
         assert message.metadata is None
 
-    def test_new_agent_message_with_metadata(
-        self, task_updater, sample_parts
-    ):
+    def test_new_agent_message_with_metadata(self, task_updater, sample_parts):
         """Test creating a new agent message with metadata and final=True."""
         metadata = {'key': 'value'}