feat: Introduce a ServerCallContext (#94)

* Introduce a ServerCallContext parameter

* Finish comment

* Remove Starlette-specific DefaultCallContextBuilder

* Update comment

Author: mikeas1

src/a2a/server/agent_execution/context.py

@@ -1,5 +1,6 @@
 import uuid
 
+from a2a.server.context import ServerCallContext
 from a2a.types import (
     InvalidParamsError,
     Message,
@@ -26,6 +27,7 @@ def __init__(
         context_id: str | None = None,
         task: Task | None = None,
         related_tasks: list[Task] | None = None,
+        call_context: ServerCallContext | None = None,
     ):
         """Initializes the RequestContext.
 
@@ -43,6 +45,7 @@ def __init__(
         self._context_id = context_id
         self._current_task = task
         self._related_tasks = related_tasks
+        self._call_context = call_context
         # If the task id and context id were provided, make sure they
         # match the request. Otherwise, create them
         if self._params:
@@ -125,6 +128,11 @@ def configuration(self) -> MessageSendConfiguration | None:
             return None
         return self._params.configuration
 
+    @property
+    def call_context(self) -> ServerCallContext | None:
+        """The server call context associated with this request."""
+        return self._call_context
+
     def _check_or_generate_task_id(self) -> None:
         """Ensures a task ID is present, generating one if necessary."""
         if not self._params:

src/a2a/server/agent_execution/request_context_builder.py

@@ -1,6 +1,7 @@
 from abc import ABC, abstractmethod
 
 from a2a.server.agent_execution import RequestContext
+from a2a.server.context import ServerCallContext
 from a2a.types import MessageSendParams, Task
 
 
@@ -14,5 +15,6 @@ async def build(
         task_id: str | None = None,
         context_id: str | None = None,
         task: Task | None = None,
+        context: ServerCallContext | None = None,
     ) -> RequestContext:
         pass

src/a2a/server/agent_execution/simple_request_context_builder.py

@@ -1,6 +1,7 @@
 import asyncio
 
 from a2a.server.agent_execution import RequestContext, RequestContextBuilder
+from a2a.server.context import ServerCallContext
 from a2a.server.tasks import TaskStore
 from a2a.types import MessageSendParams, Task
 
@@ -22,6 +23,7 @@ async def build(
         task_id: str | None = None,
         context_id: str | None = None,
         task: Task | None = None,
+        context: ServerCallContext | None = None,
     ) -> RequestContext:
         related_tasks: list[Task] | None = None
 
@@ -45,4 +47,5 @@ async def build(
             context_id=context_id,
             task=task,
             related_tasks=related_tasks,
+            call_context=context,
         )

src/a2a/server/apps/starlette_app.py

@@ -1,13 +1,15 @@
 import logging
 
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator
 from typing import Any
 
 from starlette.applications import Starlette
 from starlette.routing import Route
 
-from a2a.server.request_handlers.request_handler import RequestHandler
+from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.jsonrpc_handler import JSONRPCHandler
-
+from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.types import (
     A2AError,
     A2ARequest,
@@ -34,26 +36,43 @@
 logger = logging.getLogger(__name__)
 
 
-class A2AStarletteApplication(DefaultA2AApplication):
+class CallContextBuilder(ABC):
+    """A class for building ServerCallContexts using the Starlette Request."""
+
+    @abstractmethod
+    def build(self, request: Request) -> ServerCallContext:
+        """Builds a ServerCallContext from a Starlette Request."""
+
+
+class A2AStarletteApplication:
     """A Starlette application implementing the A2A protocol server endpoints.
 
     Handles incoming JSON-RPC requests, routes them to the appropriate
     handler methods, and manages response generation including Server-Sent Events
     (SSE).
     """
 
-    def __init__(self, agent_card: AgentCard, http_handler: RequestHandler):
+    def __init__(
+        self,
+        agent_card: AgentCard,
+        http_handler: RequestHandler,
+        context_builder: CallContextBuilder | None = None,
+    ):
         """Initializes the A2AStarletteApplication.
 
         Args:
             agent_card: The AgentCard describing the agent's capabilities.
             http_handler: The handler instance responsible for processing A2A
               requests via http.
+            context_builder: The CallContextBuilder used to construct the
+              ServerCallContext passed to the http_handler. If None, no
+              ServerCallContext is passed.
         """
         self.agent_card = agent_card
         self.handler = JSONRPCHandler(
             agent_card=agent_card, request_handler=http_handler
         )
+        self._context_builder = context_builder
 
     def _generate_error_response(
         self, request_id: str | int | None, error: JSONRPCError | A2AError
@@ -115,6 +134,11 @@ async def _handle_requests(self, request: Request) -> Response:
         try:
             body = await request.json()
             a2a_request = A2ARequest.model_validate(body)
+            call_context = (
+                self._context_builder.build(request)
+                if self._context_builder
+                else None
+            )
 
             request_id = a2a_request.root.id
             request_obj = a2a_request.root
@@ -124,11 +148,11 @@ async def _handle_requests(self, request: Request) -> Response:
                 TaskResubscriptionRequest | SendStreamingMessageRequest,
             ):
                 return await self._process_streaming_request(
-                    request_id, a2a_request
+                    request_id, a2a_request, call_context
                 )
 
             return await self._process_non_streaming_request(
-                request_id, a2a_request
+                request_id, a2a_request, call_context
             )
         except MethodNotImplementedError:
             traceback.print_exc()
@@ -154,7 +178,10 @@ async def _handle_requests(self, request: Request) -> Response:
             )
 
     async def _process_streaming_request(
-        self, request_id: str | int | None, a2a_request: A2ARequest
+        self,
+        request_id: str | int | None,
+        a2a_request: A2ARequest,
+        context: ServerCallContext,
     ) -> Response:
         """Processes streaming requests (message/stream or tasks/resubscribe).
 
@@ -171,14 +198,21 @@ async def _process_streaming_request(
             request_obj,
             SendStreamingMessageRequest,
         ):
-            handler_result = self.handler.on_message_send_stream(request_obj)
+            handler_result = self.handler.on_message_send_stream(
+                request_obj, context
+            )
         elif isinstance(request_obj, TaskResubscriptionRequest):
-            handler_result = self.handler.on_resubscribe_to_task(request_obj)
+            handler_result = self.handler.on_resubscribe_to_task(
+                request_obj, context
+            )
 
         return self._create_response(handler_result)
 
     async def _process_non_streaming_request(
-        self, request_id: str | int | None, a2a_request: A2ARequest
+        self,
+        request_id: str | int | None,
+        a2a_request: A2ARequest,
+        context: ServerCallContext,
     ) -> Response:
         """Processes non-streaming requests (message/send, tasks/get, tasks/cancel, tasks/pushNotificationConfig/*).
 
@@ -193,18 +227,26 @@ async def _process_non_streaming_request(
         handler_result: Any = None
         match request_obj:
             case SendMessageRequest():
-                handler_result = await self.handler.on_message_send(request_obj)
+                handler_result = await self.handler.on_message_send(
+                    request_obj, context
+                )
             case CancelTaskRequest():
-                handler_result = await self.handler.on_cancel_task(request_obj)
+                handler_result = await self.handler.on_cancel_task(
+                    request_obj, context
+                )
             case GetTaskRequest():
-                handler_result = await self.handler.on_get_task(request_obj)
+                handler_result = await self.handler.on_get_task(
+                    request_obj, context
+                )
             case SetTaskPushNotificationConfigRequest():
                 handler_result = await self.handler.set_push_notification(
-                    request_obj
+                    request_obj,
+                    context,
                 )
             case GetTaskPushNotificationConfigRequest():
                 handler_result = await self.handler.get_push_notification(
-                    request_obj
+                    request_obj,
+                    context,
                 )
             case _:
                 logger.error(

src/a2a/server/context.py

@@ -0,0 +1,24 @@
+"""Defines the ServerCallContext class."""
+
+import collections.abc
+import typing
+
+
+State = collections.abc.MutableMapping[str, typing.Any]
+
+
+class ServerCallContext:
+    """A context passed when calling a server method.
+
+    This class allows storing arbitrary user data in the state attribute.
+    """
+
+    def __init__(self, state: State | None = None):
+        if state is None:
+            state = {}
+        self._state = state
+
+    @property
+    def state(self) -> State:
+        """Get the user-provided state."""
+        return self._state

src/a2a/server/request_handlers/default_request_handler.py

@@ -10,6 +10,7 @@
     RequestContextBuilder,
     SimpleRequestContextBuilder,
 )
+from a2a.server.context import ServerCallContext
 from a2a.server.events import (
     Event,
     EventConsumer,
@@ -70,6 +71,8 @@ def __init__(
             task_store: The `TaskStore` instance to manage task persistence.
             queue_manager: The `QueueManager` instance to manage event queues. Defaults to `InMemoryQueueManager`.
             push_notifier: The `PushNotifier` instance for sending push notifications. Defaults to None.
+            request_context_builder: The `RequestContextBuilder` instance used
+              to build request contexts. Defaults to `SimpleRequestContextBuilder`.
         """
         self.agent_executor = agent_executor
         self.task_store = task_store
@@ -85,14 +88,20 @@ def __init__(
         self._running_agents = {}
         self._running_agents_lock = asyncio.Lock()
 
-    async def on_get_task(self, params: TaskQueryParams) -> Task | None:
+    async def on_get_task(
+        self,
+        params: TaskQueryParams,
+        context: ServerCallContext | None = None,
+    ) -> Task | None:
         """Default handler for 'tasks/get'."""
         task: Task | None = await self.task_store.get(params.id)
         if not task:
             raise ServerError(error=TaskNotFoundError())
         return task
 
-    async def on_cancel_task(self, params: TaskIdParams) -> Task | None:
+    async def on_cancel_task(
+        self, params: TaskIdParams, context: ServerCallContext | None = None
+    ) -> Task | None:
         """Default handler for 'tasks/cancel'.
 
         Attempts to cancel the task managed by the `AgentExecutor`.
@@ -150,7 +159,9 @@ async def _run_event_stream(
         await queue.close()
 
     async def on_message_send(
-        self, params: MessageSendParams
+        self,
+        params: MessageSendParams,
+        context: ServerCallContext | None = None,
     ) -> Message | Task:
         """Default handler for 'message/send' interface (non-streaming).
 
@@ -183,6 +194,7 @@ async def on_message_send(
             task_id=task.id if task else None,
             context_id=params.message.contextId,
             task=task,
+            context=context,
         )
 
         task_id = cast(str, request_context.task_id)
@@ -232,7 +244,9 @@ async def on_message_send(
         return result
 
     async def on_message_send_stream(
-        self, params: MessageSendParams
+        self,
+        params: MessageSendParams,
+        context: ServerCallContext | None = None,
     ) -> AsyncGenerator[Event]:
         """Default handler for 'message/stream' (streaming).
 
@@ -270,6 +284,7 @@ async def on_message_send_stream(
             task_id=task.id if task else None,
             context_id=params.message.contextId,
             task=task,
+            context=context,
         )
 
         task_id = cast(str, request_context.task_id)
@@ -334,7 +349,9 @@ async def _cleanup_producer(
             self._running_agents.pop(task_id, None)
 
     async def on_set_task_push_notification_config(
-        self, params: TaskPushNotificationConfig
+        self,
+        params: TaskPushNotificationConfig,
+        context: ServerCallContext | None = None,
     ) -> TaskPushNotificationConfig:
         """Default handler for 'tasks/pushNotificationConfig/set'.
 
@@ -355,7 +372,9 @@ async def on_set_task_push_notification_config(
         return params
 
     async def on_get_task_push_notification_config(
-        self, params: TaskIdParams
+        self,
+        params: TaskIdParams,
+        context: ServerCallContext | None = None,
     ) -> TaskPushNotificationConfig:
         """Default handler for 'tasks/pushNotificationConfig/get'.
 
@@ -377,7 +396,9 @@ async def on_get_task_push_notification_config(
         )
 
     async def on_resubscribe_to_task(
-        self, params: TaskIdParams
+        self,
+        params: TaskIdParams,
+        context: ServerCallContext | None = None,
     ) -> AsyncGenerator[Event]:
         """Default handler for 'tasks/resubscribe'.