Merge branch 'main' into add-fastapi-app

Author: martimfasantos

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -45,8 +45,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):
@@ -64,7 +76,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)
 
@@ -139,7 +151,7 @@ def _generate_error_response(
             log_level,
             f'Request Error (ID: {request_id}): '
             f"Code={error_resp.error.code}, Message='{error_resp.error.message}'"
-            f'{", Data=" + str(error_resp.error.data) if hasattr(error, "data") and error_resp.error.data else ""}',
+            f'{", Data=" + str(error_resp.error.data) if error_resp.error.data else ""}',
         )
         return JSONResponse(
             error_resp.model_dump(mode='json', exclude_none=True),

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,
                 ),
             )
         )

src/a2a/types.py

@@ -583,6 +583,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.

src/a2a/utils/helpers.py

@@ -69,7 +69,7 @@ def append_artifact_to_task(task: Task, event: TaskArtifactUpdateEvent) -> None:
 
     # Find existing artifact by its id
     for i, art in enumerate(task.artifacts):
-        if hasattr(art, 'artifactId') and art.artifactId == artifact_id:
+        if art.artifactId == artifact_id:
             existing_artifact = art
             existing_artifact_list_index = i
             break

tests/server/test_integration.py

@@ -5,6 +5,15 @@
 
 import pytest
 
+from starlette.authentication import (
+    AuthCredentials,
+    AuthenticationBackend,
+    BaseUser,
+    SimpleUser,
+)
+from starlette.middleware import Middleware
+from starlette.middleware.authentication import AuthenticationMiddleware
+from starlette.requests import HTTPConnection
 from starlette.responses import JSONResponse
 from starlette.routing import Route
 from starlette.testclient import TestClient
@@ -21,8 +30,12 @@
     InternalError,
     InvalidRequestError,
     JSONParseError,
+    Message,
     Part,
     PushNotificationConfig,
+    Role,
+    SendMessageResponse,
+    SendMessageSuccessResponse,
     Task,
     TaskArtifactUpdateEvent,
     TaskPushNotificationConfig,
@@ -122,21 +135,20 @@ def handler():
 def starlette_app(agent_card: AgentCard, handler: mock.AsyncMock):
     return A2AStarletteApplication(agent_card, handler)
 
-
 @pytest.fixture
-def starlette_client(app: A2AStarletteApplication):
+def starlette_client(app: A2AStarletteApplication, **kwargs):
     """Create a test client with the Starlette app."""
-    return TestClient(app.build())
+    return TestClient(app.build(**kwargs))
+
 
 @pytest.fixture
 def fastapi_app(agent_card: AgentCard, handler: mock.AsyncMock):
     return A2AFastAPIApplication(agent_card, handler)
 
-
 @pytest.fixture
-def fastapi_client(app: A2AFastAPIApplication):
+def fastapi_client(app: A2AFastAPIApplication, **kwargs):
     """Create a test client with the FastAPI app."""
-    return TestClient(app.build())
+    return TestClient(app.build(**kwargs))
 
 
 # === BASIC FUNCTIONALITY TESTS ===
@@ -345,7 +357,6 @@ def test_send_message(client: TestClient, handler: mock.AsyncMock):
     mock_task = Task(
         id='task1',
         contextId='session-xyz',
-        state='completed',
         status=task_status,
     )
     handler.on_message_send.return_value = mock_task
@@ -514,6 +525,67 @@ def test_get_push_notification_config(
     handler.on_get_task_push_notification_config.assert_awaited_once()
 
 
+def test_server_auth(app: A2AStarletteApplication, handler: mock.AsyncMock):
+    class TestAuthMiddleware(AuthenticationBackend):
+        async def authenticate(
+            self, conn: HTTPConnection
+        ) -> tuple[AuthCredentials, BaseUser] | None:
+            # For the purposes of this test, all requests are authenticated!
+            return (AuthCredentials(['authenticated']), SimpleUser('test_user'))
+
+    client = TestClient(
+        app.build(
+            middleware=[
+                Middleware(
+                    AuthenticationMiddleware, backend=TestAuthMiddleware()
+                )
+            ]
+        )
+    )
+
+    # Set the output message to be the authenticated user name
+    handler.on_message_send.side_effect = lambda params, context: Message(
+        contextId='session-xyz',
+        messageId='112',
+        role=Role.agent,
+        parts=[
+            Part(TextPart(text=context.user.user_name)),
+        ],
+    )
+
+    # Send request
+    response = client.post(
+        '/',
+        json={
+            'jsonrpc': '2.0',
+            'id': '123',
+            'method': 'message/send',
+            'params': {
+                'message': {
+                    'role': 'agent',
+                    'parts': [{'kind': 'text', 'text': 'Hello'}],
+                    'messageId': '111',
+                    'kind': 'message',
+                    'taskId': 'task1',
+                    'contextId': 'session-xyz',
+                }
+            },
+        },
+    )
+
+    # Verify response
+    assert response.status_code == 200
+    result = SendMessageResponse.model_validate(response.json())
+    assert isinstance(result.root, SendMessageSuccessResponse)
+    assert isinstance(result.root.result, Message)
+    message = result.root.result
+    assert isinstance(message.parts[0].root, TextPart)
+    assert message.parts[0].root.text == 'test_user'
+
+    # Verify handler was called
+    handler.on_message_send.assert_awaited_once()
+
+
 # === STREAMING TESTS ===