Add server support for propagating extensions, both input and output.

This commit adds support for extracting extensions requested by a
client, and marking an extension as activated, which causes it to be
returned to clients in a header.

This is purely plumbing, no particular support for actually using
extensions.

Author: mikeas1

src/a2a/extensions/common.py

@@ -0,0 +1,13 @@
+from a2a.types import AgentCard, AgentExtension
+
+
+HTTP_EXTENSION_HEADER = 'X-A2A-Extensions'
+
+
+def find_extension_by_uri(card: AgentCard, uri: str) -> AgentExtension | None:
+    """Find an AgentExtension in an AgentCard given a uri."""
+    for ext in card.capabilities.extensions or []:
+        if ext.uri == uri:
+            return ext
+
+    return None

src/a2a/server/agent_execution/context.py

@@ -143,6 +143,24 @@ def metadata(self) -> dict[str, Any]:
             return {}
         return self._params.metadata or {}
 
+    def add_activated_extension(self, uri: str) -> None:
+        """Add an extension to the set of activated extensions for this request.
+
+        This causes the extension to be indicated back to the client in the
+        response.
+        """
+        if self._call_context:
+            self._call_context.activated_extensions.add(uri)
+
+    @property
+    def requested_extensions(self) -> set[str]:
+        """Extensions that the client requested to activate."""
+        return (
+            self._call_context.requested_extensions
+            if self._call_context
+            else set()
+        )
+
     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/apps/jsonrpc/jsonrpc_app.py

@@ -19,6 +19,7 @@
 
 from a2a.auth.user import UnauthenticatedUser
 from a2a.auth.user import User as A2AUser
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.jsonrpc_handler import JSONRPCHandler
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -99,7 +100,15 @@ def build(self, request: Request) -> ServerCallContext:
             user = StarletteUserProxy(request.user)
             state['auth'] = request.auth
         state['headers'] = dict(request.headers)
-        return ServerCallContext(user=user, state=state)
+        return ServerCallContext(
+            user=user,
+            state=state,
+            requested_extensions={
+                ext
+                for h in request.headers.getlist(HTTP_EXTENSION_HEADER)
+                for ext in h.split(', ')
+            },
+        )
 
 
 class JSONRPCApplication(ABC):
@@ -281,7 +290,7 @@ async def _process_streaming_request(
                 request_obj, context
             )
 
-        return self._create_response(handler_result)
+        return self._create_response(context, handler_result)
 
     async def _process_non_streaming_request(
         self,
@@ -353,10 +362,11 @@ async def _process_non_streaming_request(
                     id=request_id, error=error
                 )
 
-        return self._create_response(handler_result)
+        return self._create_response(context, handler_result)
 
     def _create_response(
         self,
+        context: ServerCallContext,
         handler_result: (
             AsyncGenerator[SendStreamingMessageResponse]
             | JSONRPCErrorResponse
@@ -372,12 +382,16 @@ def _create_response(
         payloads.
 
         Args:
+            context: The ServerCallContext provided to the request handler.
             handler_result: The result from a request handler method. Can be an
                 async generator for streaming or a Pydantic model for non-streaming.
 
         Returns:
             A Starlette JSONResponse or EventSourceResponse.
         """
+        headers = {}
+        if exts := context.activated_extensions:
+            headers[HTTP_EXTENSION_HEADER] = ', '.join(exts)
         if isinstance(handler_result, AsyncGenerator):
             # Result is a stream of SendStreamingMessageResponse objects
             async def event_generator(
@@ -386,17 +400,21 @@ async def event_generator(
                 async for item in stream:
                     yield {'data': item.root.model_dump_json(exclude_none=True)}
 
-            return EventSourceResponse(event_generator(handler_result))
+            return EventSourceResponse(
+                event_generator(handler_result), headers=headers
+            )
         if isinstance(handler_result, JSONRPCErrorResponse):
             return JSONResponse(
                 handler_result.model_dump(
                     mode='json',
                     exclude_none=True,
-                )
+                ),
+                headers=headers,
             )
 
         return JSONResponse(
-            handler_result.root.model_dump(mode='json', exclude_none=True)
+            handler_result.root.model_dump(mode='json', exclude_none=True),
+            headers=headers,
         )
 
     async def _handle_get_agent_card(self, request: Request) -> JSONResponse:

src/a2a/server/context.py

@@ -21,3 +21,5 @@ class ServerCallContext(BaseModel):
 
     state: State = Field(default={})
     user: User = Field(default=UnauthenticatedUser())
+    requested_extensions: set[str] = Field(default=set())
+    activated_extensions: set[str] = Field(default=set())

src/a2a/server/request_handlers/grpc_handler.py

@@ -3,7 +3,7 @@
 import logging
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, Sequence
 
 
 try:
@@ -16,10 +16,13 @@
         "'pip install a2a-sdk[grpc]'"
     ) from e
 
+from grpc.aio import Metadata
+
 import a2a.grpc.a2a_pb2_grpc as a2a_grpc
 
 from a2a import types
 from a2a.auth.user import UnauthenticatedUser
+from a2a.extensions.common import HTTP_EXTENSION_HEADER
 from a2a.grpc import a2a_pb2
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -42,6 +45,25 @@ def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
         """Builds a ServerCallContext from a gRPC Request."""
 
 
+def _get_metadata_value(
+    context: grpc.aio.ServicerContext, key: str
+) -> list[str]:
+    md = context.invocation_metadata
+    vs = []
+    if isinstance(md, Metadata):
+        vs = [
+            e if isinstance(e, str) else e.decode('utf-8')
+            for e in md.get_all(key)
+        ]
+    elif isinstance(md, Sequence):
+        vs = [
+            e if isinstance(e, str) else e.decode('utf-8')
+            for (k, e) in md
+            if k == key.lower()
+        ]
+    return vs
+
+
 class DefaultCallContextBuilder(CallContextBuilder):
     """A default implementation of CallContextBuilder."""
 
@@ -51,7 +73,13 @@ def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
         state = {}
         with contextlib.suppress(Exception):
             state['grpc_context'] = context
-        return ServerCallContext(user=user, state=state)
+        return ServerCallContext(
+            user=user,
+            state=state,
+            requested_extensions=set(
+                _get_metadata_value(context, HTTP_EXTENSION_HEADER)
+            ),
+        )
 
 
 class GrpcHandler(a2a_grpc.A2AServiceServicer):
@@ -102,6 +130,7 @@ async def SendMessage(
             task_or_message = await self.request_handler.on_message_send(
                 a2a_request, server_context
             )
+            self._set_extension_metadata(context, server_context)
             return proto_utils.ToProto.task_or_message(task_or_message)
         except ServerError as e:
             await self.abort_context(e, context)
@@ -140,6 +169,7 @@ async def SendStreamingMessage(
                 a2a_request, server_context
             ):
                 yield proto_utils.ToProto.stream_response(event)
+            self._set_extension_metadata(context, server_context)
         except ServerError as e:
             await self.abort_context(e, context)
         return
@@ -371,3 +401,16 @@ async def abort_context(
                     grpc.StatusCode.UNKNOWN,
                     f'Unknown error type: {error.error}',
                 )
+
+    def _set_extension_metadata(
+        self,
+        context: grpc.aio.ServicerContext,
+        server_context: ServerCallContext,
+    ) -> None:
+        if server_context.activated_extensions:
+            context.set_trailing_metadata(
+                [
+                    (HTTP_EXTENSION_HEADER, e)
+                    for e in server_context.activated_extensions
+                ]
+            )