refactor: replace CallContextBuilder with functional UserBuilder pattern for server request authentication

Author: guglielmo-san

src/a2a/compat/v0_3/grpc_handler.py

@@ -23,8 +23,9 @@
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.grpc_handler import (
     _ERROR_CODE_MAP,
-    CallContextBuilder,
-    DefaultCallContextBuilder,
+    GrpcUserBuilder,
+    build_grpc_server_call_context,
+    default_grpc_user_builder,
 )
 from a2a.server.request_handlers.request_handler import RequestHandler
 from a2a.types.a2a_pb2 import AgentCard
@@ -44,7 +45,7 @@ def __init__(
         self,
         agent_card: AgentCard,
         request_handler: RequestHandler,
-        context_builder: CallContextBuilder | None = None,
+        user_builder: GrpcUserBuilder | None = None,
         card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
         | None = None,
     ):
@@ -54,14 +55,14 @@ def __init__(
             agent_card: The AgentCard describing the agent's capabilities (v1.0).
             request_handler: The underlying `RequestHandler` instance to
                              delegate requests to.
-            context_builder: The CallContextBuilder object. If none the
-                             DefaultCallContextBuilder is used.
+            user_builder: Optional custom user builder to extract user from the
+                          gRPC context.
             card_modifier: An optional callback to dynamically modify the public
               agent card before it is served.
         """
         self.agent_card = agent_card
         self.handler03 = RequestHandler03(request_handler=request_handler)
-        self.context_builder = context_builder or DefaultCallContextBuilder()
+        self.user_builder = user_builder or default_grpc_user_builder
         self.card_modifier = card_modifier
 
     async def _handle_unary(
@@ -72,7 +73,9 @@ async def _handle_unary(
     ) -> TResponse:
         """Centralized error handling and context management for unary calls."""
         try:
-            server_context = self.context_builder.build(context)
+            server_context = build_grpc_server_call_context(
+                context, self.user_builder
+            )
             result = await handler_func(server_context)
             self._set_extension_metadata(context, server_context)
         except A2AError as e:
@@ -88,7 +91,9 @@ async def _handle_stream(
     ) -> AsyncIterable[TResponse]:
         """Centralized error handling and context management for streaming calls."""
         try:
-            server_context = self.context_builder.build(context)
+            server_context = build_grpc_server_call_context(
+                context, self.user_builder
+            )
             async for item in handler_func(server_context):
                 yield item
             self._set_extension_metadata(context, server_context)

src/a2a/compat/v0_3/jsonrpc_adapter.py

@@ -11,7 +11,6 @@
     from starlette.requests import Request
 
     from a2a.server.request_handlers.request_handler import RequestHandler
-    from a2a.server.routes import CallContextBuilder
     from a2a.types.a2a_pb2 import AgentCard
 
     _package_starlette_installed = True
@@ -38,6 +37,11 @@
 from a2a.server.jsonrpc_models import (
     JSONRPCError as CoreJSONRPCError,
 )
+from a2a.server.routes.common import (
+    UserBuilder,
+    build_server_call_context,
+    default_user_builder,
+)
 from a2a.utils import constants
 from a2a.utils.errors import ExtendedAgentCardNotConfiguredError
 from a2a.utils.helpers import maybe_await, validate_version
@@ -67,7 +71,7 @@ def __init__(  # noqa: PLR0913
         agent_card: 'AgentCard',
         http_handler: 'RequestHandler',
         extended_agent_card: 'AgentCard | None' = None,
-        context_builder: 'CallContextBuilder | None' = None,
+        user_builder: 'UserBuilder | None' = None,
         card_modifier: 'Callable[[AgentCard], Awaitable[AgentCard] | AgentCard] | None' = None,
         extended_card_modifier: 'Callable[[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard] | None' = None,
     ):
@@ -78,7 +82,7 @@ def __init__(  # noqa: PLR0913
         self.handler = RequestHandler03(
             request_handler=http_handler,
         )
-        self._context_builder = context_builder
+        self._user_builder = user_builder or default_user_builder
 
     def supports_method(self, method: str) -> bool:
         """Returns True if the v0.3 adapter supports the given method name."""
@@ -126,10 +130,8 @@ async def handle_request(
                     CoreInvalidRequestError(data=str(e)),
                 )
 
-            call_context = (
-                self._context_builder.build(request)
-                if self._context_builder
-                else ServerCallContext()
+            call_context = build_server_call_context(
+                request, self._user_builder
             )
             call_context.tenant = (
                 getattr(specific_request.params, 'tenant', '')

src/a2a/compat/v0_3/rest_adapter.py

@@ -34,7 +34,11 @@
 from a2a.compat.v0_3 import conversions
 from a2a.compat.v0_3.rest_handler import REST03Handler
 from a2a.server.context import ServerCallContext
-from a2a.server.routes import CallContextBuilder, DefaultCallContextBuilder
+from a2a.server.routes.common import (
+    UserBuilder,
+    build_server_call_context,
+    default_user_builder,
+)
 from a2a.utils.error_handlers import (
     rest_error_handler,
     rest_stream_error_handler,
@@ -60,7 +64,7 @@ def __init__(  # noqa: PLR0913
         agent_card: 'AgentCard',
         http_handler: 'RequestHandler',
         extended_agent_card: 'AgentCard | None' = None,
-        context_builder: 'CallContextBuilder | None' = None,
+        user_builder: 'UserBuilder | None' = None,
         card_modifier: 'Callable[[AgentCard], Awaitable[AgentCard] | AgentCard] | None' = None,
         extended_card_modifier: 'Callable[[AgentCard, ServerCallContext], Awaitable[AgentCard] | AgentCard] | None' = None,
     ):
@@ -71,15 +75,15 @@ def __init__(  # noqa: PLR0913
         self.handler = REST03Handler(
             agent_card=agent_card, request_handler=http_handler
         )
-        self._context_builder = context_builder or DefaultCallContextBuilder()
+        self._user_builder = user_builder or default_user_builder
 
     @rest_error_handler
     async def _handle_request(
         self,
         method: 'Callable[[Request, ServerCallContext], Awaitable[Any]]',
         request: Request,
     ) -> Response:
-        call_context = self._context_builder.build(request)
+        call_context = build_server_call_context(request, self._user_builder)
         response = await method(request, call_context)
         return JSONResponse(content=response)
 
@@ -96,7 +100,7 @@ async def _handle_streaming_request(
                 message=f'Failed to pre-consume request body: {e}'
             ) from e
 
-        call_context = self._context_builder.build(request)
+        call_context = build_server_call_context(request, self._user_builder)
 
         async def event_generator(
             stream: AsyncIterable[Any],

src/a2a/server/request_handlers/grpc_handler.py

@@ -1,7 +1,6 @@
 # ruff: noqa: N802
 import logging
 
-from abc import ABC, abstractmethod
 from collections.abc import AsyncIterable, Awaitable, Callable
 from typing import TypeVar
 
@@ -24,7 +23,7 @@
 import a2a.types.a2a_pb2_grpc as a2a_grpc
 
 from a2a import types
-from a2a.auth.user import UnauthenticatedUser
+from a2a.auth.user import UnauthenticatedUser, User
 from a2a.extensions.common import (
     HTTP_EXTENSION_HEADER,
     get_requested_extensions,
@@ -41,15 +40,12 @@
 
 logger = logging.getLogger(__name__)
 
-# For now we use a trivial wrapper on the grpc context object
+GrpcUserBuilder = Callable[[grpc.aio.ServicerContext], User]
 
 
-class CallContextBuilder(ABC):
-    """A class for building ServerCallContexts using the Starlette Request."""
-
-    @abstractmethod
-    def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
-        """Builds a ServerCallContext from a gRPC Request."""
+def default_grpc_user_builder(context: grpc.aio.ServicerContext) -> User:
+    """Default strategy for creating a User from a gRPC context."""
+    return UnauthenticatedUser()
 
 
 def _get_metadata_value(
@@ -67,20 +63,19 @@ def _get_metadata_value(
     ]
 
 
-class DefaultCallContextBuilder(CallContextBuilder):
-    """A default implementation of CallContextBuilder."""
-
-    def build(self, context: grpc.aio.ServicerContext) -> ServerCallContext:
-        """Builds the ServerCallContext."""
-        user = UnauthenticatedUser()
-        state = {'grpc_context': context}
-        return ServerCallContext(
-            user=user,
-            state=state,
-            requested_extensions=get_requested_extensions(
-                _get_metadata_value(context, HTTP_EXTENSION_HEADER)
-            ),
-        )
+def build_grpc_server_call_context(
+    context: grpc.aio.ServicerContext, user_builder: GrpcUserBuilder
+) -> ServerCallContext:
+    """Builds a ServerCallContext from a gRPC ServicerContext."""
+    user = user_builder(context)
+    state = {'grpc_context': context}
+    return ServerCallContext(
+        user=user,
+        state=state,
+        requested_extensions=get_requested_extensions(
+            _get_metadata_value(context, HTTP_EXTENSION_HEADER)
+        ),
+    )
 
 
 _ERROR_CODE_MAP = {
@@ -110,7 +105,7 @@ def __init__(
         self,
         agent_card: AgentCard,
         request_handler: RequestHandler,
-        context_builder: CallContextBuilder | None = None,
+        user_builder: GrpcUserBuilder | None = None,
         card_modifier: Callable[[AgentCard], Awaitable[AgentCard] | AgentCard]
         | None = None,
     ):
@@ -120,14 +115,14 @@ def __init__(
             agent_card: The AgentCard describing the agent's capabilities.
             request_handler: The underlying `RequestHandler` instance to
                              delegate requests to.
-            context_builder: The CallContextBuilder object. If none the
-                             DefaultCallContextBuilder is used.
+            user_builder: Optional custom user builder to extract user from the
+                          gRPC context.
             card_modifier: An optional callback to dynamically modify the public
               agent card before it is served.
         """
         self.agent_card = agent_card
         self.request_handler = request_handler
-        self.context_builder = context_builder or DefaultCallContextBuilder()
+        self.user_builder = user_builder or default_grpc_user_builder
         self.card_modifier = card_modifier
 
     async def _handle_unary(
@@ -451,6 +446,8 @@ def _build_call_context(
         context: grpc.aio.ServicerContext,
         request: message.Message,
     ) -> ServerCallContext:
-        server_context = self.context_builder.build(context)
+        server_context = build_grpc_server_call_context(
+            context, self.user_builder
+        )
         server_context.tenant = getattr(request, 'tenant', '')
         return server_context

src/a2a/server/routes/__init__.py

@@ -1,17 +1,13 @@
 """A2A Routes."""
 
 from a2a.server.routes.agent_card_routes import create_agent_card_routes
-from a2a.server.routes.common import (
-    CallContextBuilder,
-    DefaultCallContextBuilder,
-)
+from a2a.server.routes.common import UserBuilder
 from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes
 from a2a.server.routes.rest_routes import create_rest_routes
 
 
 __all__ = [
-    'CallContextBuilder',
-    'DefaultCallContextBuilder',
+    'UserBuilder',
     'create_agent_card_routes',
     'create_jsonrpc_routes',
     'create_rest_routes',

src/a2a/server/routes/common.py

@@ -1,65 +1,58 @@
-from abc import ABC, abstractmethod
+from collections.abc import Callable
 
-from starlette.authentication import BaseUser
 from starlette.requests import Request
 
-from a2a.auth.user import UnauthenticatedUser
-from a2a.auth.user import User as A2AUser
+from a2a.auth.user import UnauthenticatedUser, User
 from a2a.extensions.common import (
     HTTP_EXTENSION_HEADER,
     get_requested_extensions,
 )
 from a2a.server.context import ServerCallContext
 
 
-class StarletteUserProxy(A2AUser):
-    """Adapts the Starlette User class to the A2A user representation."""
+UserBuilder = Callable[[Request], User]
 
-    def __init__(self, user: BaseUser):
-        self._user = user
 
-    @property
-    def is_authenticated(self) -> bool:
-        """Returns whether the current user is authenticated."""
-        return self._user.is_authenticated
+def default_user_builder(request: Request) -> User:
+    """Default strategy for creating an A2AUser from a Starlette Request."""
+    if 'user' in request.scope:
 
-    @property
-    def user_name(self) -> str:
-        """Returns the user name of the current user."""
-        return self._user.display_name
+        class BaseUser(User):
+            @property
+            def is_authenticated(self) -> bool:
+                return request.user.is_authenticated
 
+            @property
+            def user_name(self) -> str:
+                return request.user.display_name
 
-class CallContextBuilder(ABC):
-    """A class for building ServerCallContexts using the Starlette Request."""
+        return BaseUser()
+    return UnauthenticatedUser()
 
-    @abstractmethod
-    def build(self, request: Request) -> ServerCallContext:
-        """Builds a ServerCallContext from a Starlette Request."""
 
+def build_server_call_context(
+    request: Request, user_builder: UserBuilder
+) -> ServerCallContext:
+    """Builds a ServerCallContext from a Starlette Request.
 
-class DefaultCallContextBuilder(CallContextBuilder):
-    """A default implementation of CallContextBuilder."""
+    Args:
+        request: The incoming Starlette Request object.
+        user_builder: Optional custom user builder.
 
-    def build(self, request: Request) -> ServerCallContext:
-        """Builds a ServerCallContext from a Starlette Request.
+    Returns:
+        A ServerCallContext instance populated with user and state.
+    """
+    user = user_builder(request)
 
-        Args:
-            request: The incoming Starlette Request object.
+    state = {}
+    if 'auth' in request.scope:
+        state['auth'] = request.auth
+    state['headers'] = dict(request.headers)
 
-        Returns:
-            A ServerCallContext instance populated with user and state
-            information from the request.
-        """
-        user: A2AUser = UnauthenticatedUser()
-        state = {}
-        if 'user' in request.scope:
-            user = StarletteUserProxy(request.user)
-            state['auth'] = request.auth
-        state['headers'] = dict(request.headers)
-        return ServerCallContext(
-            user=user,
-            state=state,
-            requested_extensions=get_requested_extensions(
-                request.headers.getlist(HTTP_EXTENSION_HEADER)
-            ),
-        )
+    return ServerCallContext(
+        user=user,
+        state=state,
+        requested_extensions=get_requested_extensions(
+            request.headers.getlist(HTTP_EXTENSION_HEADER)
+        ),
+    )