feat: Add server-side support for plumbing requested and activated extensions (#333)

# Description

This adds PR adds support for reading requested extensions from a
client, via the `X-A2A-Extensions` header, and plumbing that through the
ServerCallContext.

Here's my rough thinking on the design choices:

1. Requests to activate extensions are specified in a transport-specific
manner. For HTTP, it's using headers. For gRPC, it's using side-channel
metadata (which are also HTTP headers).
2. However, indicating that an extension was requested is a
transport-independent concept. So, it should be passed to the
transport-independent layers of the SDK.
3. We already have a means of extracting relevant information from the
transport layer (and others) and passing down to the independent layers:
the `ServerCallContext`.
4. So, we can put the list of requested extensions here. The transport
layers are responsible for pulling out the list of requested extensions
in whatever transport specific means they are communicated, and passing
these down via the ServerCallContext.
5. Returning the list of extensions that were activated is slightly more
annoying. I didn't want to have to refactor every layer to plumb both
the request-specific response type (such as SendMessageResponse) AND the
list of extensions activated. That's just too intrusive.
6. So, I need some way to communicate the list of activated extension
from the independent layers back "up" to the transport layers via
something other than the return value. I'm not thrilled with the
decision here, but the answer is to use interior mutability: I put a
mutable set in the ServerCallContext, and allow lower levels to modify
it to add the activated extensions.
7. The transport layers then need to read the list of activated
extensions from the ServerCallContext and pass those back in the
response in the transport-defined method.

The result is that we have full stack plumbing of requested and
activated extensions by way of the ServerCallContext.

---------

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Holt Skinner <[email protected]>
Co-authored-by: Holt Skinner <[email protected]>

Author: 3 people

src/a2a/extensions/common.py

@@ -0,0 +1,27 @@
+from a2a.types import AgentCard, AgentExtension
+
+
+HTTP_EXTENSION_HEADER = 'X-A2A-Extensions'
+
+
+def get_requested_extensions(values: list[str]) -> set[str]:
+    """Get the set of requested extensions from an input list.
+
+    This handles the list containing potentially comma-separated values, as
+    occurs when using a list in an HTTP header.
+    """
+    return {
+        stripped
+        for v in values
+        for ext in v.split(',')
+        if (stripped := ext.strip())
+    }
+
+
+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,10 @@
 
 from a2a.auth.user import UnauthenticatedUser
 from a2a.auth.user import User as A2AUser
+from a2a.extensions.common import (
+    HTTP_EXTENSION_HEADER,
+    get_requested_extensions,
+)
 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 +103,13 @@ 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=get_requested_extensions(
+                request.headers.getlist(HTTP_EXTENSION_HEADER)
+            ),
+        )
 
 
 class JSONRPCApplication(ABC):
@@ -281,7 +291,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 +363,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 +383,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(sorted(exts))
         if isinstance(handler_result, AsyncGenerator):
             # Result is a stream of SendStreamingMessageResponse objects
             async def event_generator(
@@ -386,17 +401,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_factory=set)
+    activated_extensions: set[str] = Field(default_factory=set)

src/a2a/server/request_handlers/grpc_handler.py

@@ -3,12 +3,14 @@
 import logging
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, Sequence
 
 
 try:
     import grpc
     import grpc.aio
+
+    from grpc.aio import Metadata
 except ImportError as e:
     raise ImportError(
         'GrpcHandler requires grpcio and grpcio-tools to be installed. '
@@ -20,6 +22,10 @@
 
 from a2a import types
 from a2a.auth.user import UnauthenticatedUser
+from a2a.extensions.common import (
+    HTTP_EXTENSION_HEADER,
+    get_requested_extensions,
+)
 from a2a.grpc import a2a_pb2
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
@@ -42,6 +48,19 @@ 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
+    raw_values: list[str | bytes] = []
+    if isinstance(md, Metadata):
+        raw_values = md.get_all(key)
+    elif isinstance(md, Sequence):
+        lower_key = key.lower()
+        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]
+
+
 class DefaultCallContextBuilder(CallContextBuilder):
     """A default implementation of CallContextBuilder."""
 
@@ -51,7 +70,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=get_requested_extensions(
+                _get_metadata_value(context, HTTP_EXTENSION_HEADER)
+            ),
+        )
 
 
 class GrpcHandler(a2a_grpc.A2AServiceServicer):
@@ -102,6 +127,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 +166,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 +398,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 sorted(server_context.activated_extensions)
+                ]
+            )

tests/extensions/test_common.py

@@ -0,0 +1,58 @@
+from a2a.extensions.common import (
+    find_extension_by_uri,
+    get_requested_extensions,
+)
+from a2a.types import AgentCapabilities, AgentCard, AgentExtension
+
+
+def test_get_requested_extensions():
+    assert get_requested_extensions([]) == set()
+    assert get_requested_extensions(['foo']) == {'foo'}
+    assert get_requested_extensions(['foo', 'bar']) == {'foo', 'bar'}
+    assert get_requested_extensions(['foo, bar']) == {'foo', 'bar'}
+    assert get_requested_extensions(['foo,bar']) == {'foo', 'bar'}
+    assert get_requested_extensions(['foo', 'bar,baz']) == {'foo', 'bar', 'baz'}
+    assert get_requested_extensions(['foo,, bar', 'baz']) == {
+        'foo',
+        'bar',
+        'baz',
+    }
+    assert get_requested_extensions([' foo , bar ', 'baz']) == {
+        'foo',
+        'bar',
+        'baz',
+    }
+
+
+def test_find_extension_by_uri():
+    ext1 = AgentExtension(uri='foo', description='The Foo extension')
+    ext2 = AgentExtension(uri='bar', description='The Bar extension')
+    card = AgentCard(
+        name='Test Agent',
+        description='Test Agent Description',
+        version='1.0',
+        url='http://test.com',
+        skills=[],
+        default_input_modes=['text/plain'],
+        default_output_modes=['text/plain'],
+        capabilities=AgentCapabilities(extensions=[ext1, ext2]),
+    )
+
+    assert find_extension_by_uri(card, 'foo') == ext1
+    assert find_extension_by_uri(card, 'bar') == ext2
+    assert find_extension_by_uri(card, 'baz') is None
+
+
+def test_find_extension_by_uri_no_extensions():
+    card = AgentCard(
+        name='Test Agent',
+        description='Test Agent Description',
+        version='1.0',
+        url='http://test.com',
+        skills=[],
+        default_input_modes=['text/plain'],
+        default_output_modes=['text/plain'],
+        capabilities=AgentCapabilities(extensions=None),
+    )
+
+    assert find_extension_by_uri(card, 'foo') is None

tests/server/agent_execution/test_context.py

@@ -5,6 +5,7 @@
 import pytest
 
 from a2a.server.agent_execution import RequestContext
+from a2a.server.context import ServerCallContext
 from a2a.types import (
     Message,
     MessageSendParams,
@@ -263,3 +264,16 @@ def test_init_with_context_id_and_existing_context_id_match(
 
         assert context.context_id == mock_task.context_id
         assert context.current_task == mock_task
+
+    def test_extension_handling(self):
+        """Test extension handling in RequestContext."""
+        call_context = ServerCallContext(requested_extensions={'foo', 'bar'})
+        context = RequestContext(call_context=call_context)
+
+        assert context.requested_extensions == {'foo', 'bar'}
+
+        context.add_activated_extension('foo')
+        assert call_context.activated_extensions == {'foo'}
+
+        context.add_activated_extension('baz')
+        assert call_context.activated_extensions == {'foo', 'baz'}