feat: refactor JSON-RPC dispatcher and common context handling

Author: guglielmo-san

src/a2a/server/routes/__init__.py

@@ -1,7 +1,7 @@
 """A2A Routes."""
 
 from a2a.server.routes.agent_card_routes import create_agent_card_routes
-from a2a.server.routes.jsonrpc_dispatcher import (
+from a2a.server.routes.common import (
     CallContextBuilder,
     DefaultCallContextBuilder,
 )

src/a2a/server/routes/common.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+
+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.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."""
+
+    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
+
+    @property
+    def user_name(self) -> str:
+        """Returns the user name of the current user."""
+        return self._user.display_name
+
+
+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 DefaultCallContextBuilder(CallContextBuilder):
+    """A default implementation of CallContextBuilder."""
+
+    def build(self, request: Request) -> ServerCallContext:
+        """Builds a ServerCallContext from a Starlette Request.
+
+        Args:
+            request: The incoming Starlette Request object.
+
+        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)
+            ),
+        )

src/a2a/server/routes/jsonrpc_dispatcher.py

@@ -4,19 +4,15 @@
 import logging
 import traceback
 
-from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator, Awaitable, Callable
 from typing import TYPE_CHECKING, Any
 
 from google.protobuf.json_format import MessageToDict, ParseDict
 from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response
 
-from a2a.auth.user import UnauthenticatedUser
-from a2a.auth.user import User as A2AUser
 from a2a.compat.v0_3.jsonrpc_adapter import JSONRPC03Adapter
 from a2a.extensions.common import (
     HTTP_EXTENSION_HEADER,
-    get_requested_extensions,
 )
 from a2a.server.context import ServerCallContext
 from a2a.server.jsonrpc_models import (
@@ -31,6 +27,10 @@
 from a2a.server.request_handlers.response_helpers import (
     build_error_response,
 )
+from a2a.server.routes.common import (
+    CallContextBuilder,
+    DefaultCallContextBuilder,
+)
 from a2a.types import A2ARequest
 from a2a.types.a2a_pb2 import (
     AgentCard,
@@ -113,59 +113,6 @@
         HTTP_413_CONTENT_TOO_LARGE = Any
 
 
-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) -> bool:
-        """Returns whether the current user is authenticated."""
-        return self._user.is_authenticated
-
-    @property
-    def user_name(self) -> str:
-        """Returns the user name of the current user."""
-        return self._user.display_name
-
-
-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 DefaultCallContextBuilder(CallContextBuilder):
-    """A default implementation of CallContextBuilder."""
-
-    def build(self, request: Request) -> ServerCallContext:
-        """Builds a ServerCallContext from a Starlette Request.
-
-        Args:
-            request: The incoming Starlette Request object.
-
-        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)
-            ),
-        )
-
-
 @trace_class(kind=SpanKind.SERVER)
 class JsonRpcDispatcher:
     """Base class for A2A JSONRPC applications.

src/a2a/server/routes/jsonrpc_routes.py

@@ -19,10 +19,8 @@
 
 from a2a.server.context import ServerCallContext
 from a2a.server.request_handlers.request_handler import RequestHandler
-from a2a.server.routes.jsonrpc_dispatcher import (
-    CallContextBuilder,
-    JsonRpcDispatcher,
-)
+from a2a.server.routes.common import CallContextBuilder
+from a2a.server.routes.jsonrpc_dispatcher import JsonRpcDispatcher
 from a2a.types.a2a_pb2 import AgentCard
 
 

tests/server/routes/test_jsonrpc_dispatcher.py

@@ -21,12 +21,12 @@
     Role,
 )
 from a2a.server.routes import jsonrpc_dispatcher
-from a2a.server.routes.jsonrpc_dispatcher import (
+from a2a.server.routes.common import (
     CallContextBuilder,
     DefaultCallContextBuilder,
-    JsonRpcDispatcher,
     StarletteUserProxy,
 )
+from a2a.server.routes.jsonrpc_dispatcher import JsonRpcDispatcher
 from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes
 from a2a.server.routes.agent_card_routes import create_agent_card_routes
 from a2a.server.jsonrpc_models import JSONRPCError