refactor: Remove redundant JSON-RPC Pydantic types, use jsonrpc library directly

- Update JSONRPCHandler to return dict[str, Any] instead of Pydantic RootModels
- Update response_helpers to build dicts with JSON-RPC 2.0 structure
- Remove unused Pydantic response types from types module
- Fix proto dependency loading in a2a_pb2.py
- Update all tests to check dict responses instead of Pydantic models
- Add TransportProtocol constants to utils module

Author: muscariello

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
   "protobuf>=5.29.5",
   "google-api-core>=1.26.0",
   "json-rpc>=1.15.0",
+  "googleapis-common-protos>=1.70.0",
 ]
 
 classifiers = [

src/a2a/client/auth/interceptor.py

@@ -3,43 +3,10 @@
 
 from a2a.client.auth.credentials import CredentialService
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
-from a2a.types.a2a_pb2 import (
-    AgentCard,
-    APIKeySecurityScheme,
-    HTTPAuthSecurityScheme,
-    MutualTlsSecurityScheme,
-    OAuth2SecurityScheme,
-    OpenIdConnectSecurityScheme,
-    SecurityScheme,
-)
+from a2a.types.a2a_pb2 import AgentCard
 
 logger = logging.getLogger(__name__)
 
-_SecuritySchemeValue = (
-    APIKeySecurityScheme
-    | HTTPAuthSecurityScheme
-    | OAuth2SecurityScheme
-    | OpenIdConnectSecurityScheme
-    | MutualTlsSecurityScheme
-    | None
-)
-
-
-def _get_security_scheme_value(scheme: SecurityScheme) -> _SecuritySchemeValue:
-    """Extract the actual security scheme from the oneof union."""
-    which = scheme.WhichOneof('scheme')
-    if which == 'api_key_security_scheme':
-        return scheme.api_key_security_scheme
-    if which == 'http_auth_security_scheme':
-        return scheme.http_auth_security_scheme
-    if which == 'oauth2_security_scheme':
-        return scheme.oauth2_security_scheme
-    if which == 'open_id_connect_security_scheme':
-        return scheme.open_id_connect_security_scheme
-    if which == 'mtls_security_scheme':
-        return scheme.mtls_security_scheme
-    return None
-
 
 class AuthInterceptor(ClientCallInterceptor):
     """An interceptor that automatically adds authentication details to requests.
@@ -72,54 +39,53 @@ async def intercept(
                     scheme_name, context
                 )
                 if credential and scheme_name in agent_card.security_schemes:
-                    scheme_def_union = agent_card.security_schemes.get(
-                        scheme_name
-                    )
-                    if not scheme_def_union:
-                        continue
-                    scheme_def = _get_security_scheme_value(scheme_def_union)
-                    if not scheme_def:
+                    scheme = agent_card.security_schemes.get(scheme_name)
+                    if not scheme:
                         continue
 
                     headers = http_kwargs.get('headers', {})
 
-                    match scheme_def:
-                        # Case 1a: HTTP Bearer scheme with an if guard
-                        case HTTPAuthSecurityScheme() if (
-                            scheme_def.scheme.lower() == 'bearer'
-                        ):
-                            headers['Authorization'] = f'Bearer {credential}'
-                            logger.debug(
-                                "Added Bearer token for scheme '%s'.",
-                                scheme_name,
-                            )
-                            http_kwargs['headers'] = headers
-                            return request_payload, http_kwargs
-
-                        # Case 1b: OAuth2 and OIDC schemes, which are implicitly Bearer
-                        case (
-                            OAuth2SecurityScheme()
-                            | OpenIdConnectSecurityScheme()
-                        ):
-                            headers['Authorization'] = f'Bearer {credential}'
-                            logger.debug(
-                                "Added Bearer token for scheme '%s'.",
-                                scheme_name,
-                            )
-                            http_kwargs['headers'] = headers
-                            return request_payload, http_kwargs
-
-                        # Case 2: API Key in Header
-                        case APIKeySecurityScheme() if (
-                            scheme_def.location.lower() == 'header'
-                        ):
-                            headers[scheme_def.name] = credential
-                            logger.debug(
-                                "Added API Key Header for scheme '%s'.",
-                                scheme_name,
-                            )
-                            http_kwargs['headers'] = headers
-                            return request_payload, http_kwargs
+                    # HTTP Bearer authentication
+                    if (
+                        scheme.HasField('http_auth_security_scheme')
+                        and scheme.http_auth_security_scheme.scheme.lower()
+                        == 'bearer'
+                    ):
+                        headers['Authorization'] = f'Bearer {credential}'
+                        logger.debug(
+                            "Added Bearer token for scheme '%s'.",
+                            scheme_name,
+                        )
+                        http_kwargs['headers'] = headers
+                        return request_payload, http_kwargs
+
+                    # OAuth2 and OIDC schemes are implicitly Bearer
+                    if scheme.HasField(
+                        'oauth2_security_scheme'
+                    ) or scheme.HasField('open_id_connect_security_scheme'):
+                        headers['Authorization'] = f'Bearer {credential}'
+                        logger.debug(
+                            "Added Bearer token for scheme '%s'.",
+                            scheme_name,
+                        )
+                        http_kwargs['headers'] = headers
+                        return request_payload, http_kwargs
+
+                    # API Key in Header
+                    if (
+                        scheme.HasField('api_key_security_scheme')
+                        and scheme.api_key_security_scheme.location.lower()
+                        == 'header'
+                    ):
+                        headers[scheme.api_key_security_scheme.name] = (
+                            credential
+                        )
+                        logger.debug(
+                            "Added API Key Header for scheme '%s'.",
+                            scheme_name,
+                        )
+                        http_kwargs['headers'] = headers
+                        return request_payload, http_kwargs
 
                 # Note: Other cases like API keys in query/cookie are not handled and will be skipped.
 

src/a2a/client/errors.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from a2a.types.extras import A2AError, JSONRPCErrorResponse
+from a2a.types.extras import A2AError
 
 
 class A2AClientError(Exception):
@@ -81,16 +81,11 @@ class A2AClientJSONRPCError(A2AClientError):
 
     error: dict[str, Any] | A2AError
 
-    def __init__(self, error: JSONRPCErrorResponse | dict[str, Any]):
+    def __init__(self, error: dict[str, Any] | A2AError):
         """Initializes the A2AClientJsonRPCError.
 
         Args:
-            error: The JSON-RPC error object or dict from the jsonrpc library.
+            error: The JSON-RPC error dict from the jsonrpc library, or A2AError object.
         """
-        if isinstance(error, dict):
-            # Raw dict from jsonrpc library: {'code': ..., 'message': ...}
-            self.error = error
-        else:
-            # JSONRPCErrorResponse object
-            self.error = error.error
+        self.error = error
         super().__init__(f'JSON-RPC Error {self.error}')

src/a2a/server/apps/jsonrpc/jsonrpc_app.py

@@ -1,3 +1,5 @@
+"""JSON-RPC application for A2A server."""
+
 import contextlib
 import json
 import logging
@@ -8,7 +10,7 @@
 from typing import TYPE_CHECKING, Any
 
 from google.protobuf.json_format import MessageToDict, ParseDict
-from pydantic import RootModel, ValidationError
+from jsonrpc.jsonrpc2 import JSONRPC20Request, JSONRPC20Response
 
 from a2a.auth.user import UnauthenticatedUser
 from a2a.auth.user import User as A2AUser
@@ -37,10 +39,7 @@
     InvalidParamsError,
     InvalidRequestError,
     JSONParseError,
-    JSONRPCErrorResponse,
-    JSONRPCRequest,
     MethodNotFoundError,
-    SendStreamingMessageResponse,
     TaskResubscriptionRequest,
     UnsupportedOperationError,
 )
@@ -233,10 +232,8 @@ def _generate_error_response(
         Returns:
             A `JSONResponse` object formatted as a JSON-RPC error response.
         """
-        error_resp = JSONRPCErrorResponse(
-            id=request_id,
-            error=error,
-        )
+        error_dict = error.model_dump(exclude_none=True)
+        error_resp = JSONRPC20Response(error=error_dict, _id=request_id)
 
         log_level = (
             logging.ERROR
@@ -247,14 +244,14 @@ def _generate_error_response(
             log_level,
             "Request Error (ID: %s): Code=%s, Message='%s'%s",
             request_id,
-            error_resp.error.code,
-            error_resp.error.message,
-            ', Data=' + str(error_resp.error.data)
-            if error_resp.error.data
+            error_dict.get('code'),
+            error_dict.get('message'),
+            ', Data=' + str(error_dict.get('data'))
+            if error_dict.get('data')
             else '',
         )
         return JSONResponse(
-            error_resp.model_dump(mode='json', exclude_none=True),
+            error_resp.data,
             status_code=200,
         )
 
@@ -274,7 +271,7 @@ def _allowed_content_length(self, request: Request) -> bool:
                     return False
         return True
 
-    async def _handle_requests(self, request: Request) -> Response:  # noqa: PLR0911
+    async def _handle_requests(self, request: Request) -> Response:  # noqa: PLR0911, PLR0912
         """Handles incoming POST requests to the main A2A endpoint.
 
         Parses the request body as JSON, validates it against A2A request types,
@@ -313,17 +310,31 @@ async def _handle_requests(self, request: Request) -> Response:  # noqa: PLR0911
             logger.debug('Request body: %s', body)
             # 1) Validate base JSON-RPC structure only (-32600 on failure)
             try:
-                base_request = JSONRPCRequest.model_validate(body)
-            except ValidationError as e:
+                base_request = JSONRPC20Request.from_data(body)
+                if not isinstance(base_request, JSONRPC20Request):
+                    # Batch requests are not supported
+                    return self._generate_error_response(
+                        request_id,
+                        InvalidRequestError(
+                            message='Batch requests are not supported'
+                        ),
+                    )
+            except Exception as e:
                 logger.exception('Failed to validate base JSON-RPC request')
                 return self._generate_error_response(
                     request_id,
-                    InvalidRequestError(data=json.loads(e.json())),
+                    InvalidRequestError(data=str(e)),
                 )
 
             # 2) Route by method name; unknown -> -32601, known -> validate params (-32602 on failure)
-            method = base_request.method
-            request_id = base_request.id
+            method: str | None = base_request.method
+            request_id = base_request._id  # noqa: SLF001
+
+            if not method:
+                return self._generate_error_response(
+                    request_id,
+                    InvalidRequestError(message='Method is required'),
+                )
 
             model_class = self.METHOD_TO_MODEL.get(method)
             if not model_class:
@@ -483,33 +494,25 @@ async def _process_non_streaming_request(
                 error = UnsupportedOperationError(
                     message=f'Request type {type(request_obj).__name__} is unknown.'
                 )
-                handler_result = JSONRPCErrorResponse(
-                    id=request_id, error=error
-                )
+                return self._generate_error_response(request_id, error)
 
         return self._create_response(context, handler_result)
 
     def _create_response(
         self,
         context: ServerCallContext,
-        handler_result: (
-            AsyncGenerator[SendStreamingMessageResponse]
-            | JSONRPCErrorResponse
-            | RootModel[Any]
-        ),
+        handler_result: AsyncGenerator[dict[str, Any]] | dict[str, Any],
     ) -> Response:
         """Creates a Starlette Response based on the result from the request handler.
 
         Handles:
         - AsyncGenerator for Server-Sent Events (SSE).
-        - JSONRPCErrorResponse for explicit errors returned by handlers.
-        - Pydantic RootModels (like GetTaskResponse) containing success or error
-        payloads.
+        - Dict responses from handlers.
 
         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.
+                async generator for streaming or a dict for non-streaming.
 
         Returns:
             A Starlette JSONResponse or EventSourceResponse.
@@ -518,29 +521,19 @@ def _create_response(
         if exts := context.activated_extensions:
             headers[HTTP_EXTENSION_HEADER] = ', '.join(sorted(exts))
         if isinstance(handler_result, AsyncGenerator):
-            # Result is a stream of SendStreamingMessageResponse objects
+            # Result is a stream of dict objects
             async def event_generator(
-                stream: AsyncGenerator[SendStreamingMessageResponse],
+                stream: AsyncGenerator[dict[str, Any]],
             ) -> AsyncGenerator[dict[str, str]]:
                 async for item in stream:
-                    yield {'data': item.root.model_dump_json(exclude_none=True)}
+                    yield {'data': json.dumps(item)}
 
             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),
-            headers=headers,
-        )
+        # handler_result is a dict (JSON-RPC response)
+        return JSONResponse(handler_result, headers=headers)
 
     async def _handle_get_agent_card(self, request: Request) -> JSONResponse:
         """Handles GET requests for the agent card endpoint.