feat: Add middleware to the client SDK (#171)

This PR introduces a middleware framework to the a2a-python client SDK,
which allows for various authentication mechanisms. The new
AuthInterceptor automatically injects authentication credentials into
outgoing requests based on the agent's defined security schemes. This
change simplifies the process of handling different authentication
methods like Bearer tokens and API keys for developers using the SDK.

Author: dmandar

.github/actions/spelling/allow.txt

@@ -45,6 +45,7 @@ opensource
 protoc
 pyi
 pyversions
+respx
 resub
 socio
 sse

.vscode/launch.json

@@ -12,7 +12,12 @@
         "PYTHONPATH": "${workspaceFolder}"
       },
       "cwd": "${workspaceFolder}/examples/helloworld",
-      "args": ["--host", "localhost", "--port", "9999"]
+      "args": [
+        "--host",
+        "localhost",
+        "--port",
+        "9999"
+      ]
     },
     {
       "name": "Debug Currency Agent",
@@ -25,7 +30,24 @@
         "PYTHONPATH": "${workspaceFolder}"
       },
       "cwd": "${workspaceFolder}/examples/langgraph",
-      "args": ["--host", "localhost", "--port", "10000"]
+      "args": [
+        "--host",
+        "localhost",
+        "--port",
+        "10000"
+      ]
+    },
+    {
+      "name": "Pytest All",
+      "type": "debugpy",
+      "request": "launch",
+      "module": "pytest",
+      "args": [
+        "-v",
+        "-s"
+      ],
+      "console": "integratedTerminal",
+      "justMyCode": true
     }
   ]
-}
+}

pyproject.toml

@@ -76,6 +76,7 @@ dev = [
     "pytest-asyncio>=0.26.0",
     "pytest-cov>=6.1.1",
     "pytest-mock>=3.14.0",
+    "respx>=0.20.2",
     "ruff>=0.11.6",
     "uv-dynamic-versioning>=0.8.2",
     "types-protobuf",

src/a2a/client/__init__.py

@@ -1,5 +1,10 @@
 """Client-side components for interacting with an A2A agent."""
 
+from a2a.client.auth import (
+    AuthInterceptor,
+    CredentialService,
+    InMemoryContextCredentialStore,
+)
 from a2a.client.client import A2ACardResolver, A2AClient
 from a2a.client.errors import (
     A2AClientError,
@@ -8,6 +13,7 @@
 )
 from a2a.client.grpc_client import A2AGrpcClient
 from a2a.client.helpers import create_text_message_object
+from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
 
 
 __all__ = [
@@ -17,5 +23,10 @@
     'A2AClientHTTPError',
     'A2AClientJSONError',
     'A2AGrpcClient',
+    'AuthInterceptor',
+    'ClientCallContext',
+    'ClientCallInterceptor',
+    'CredentialService',
+    'InMemoryContextCredentialStore',
     'create_text_message_object',
 ]

src/a2a/client/auth/__init__.py

@@ -0,0 +1,14 @@
+"""Client-side authentication components for the A2A Python SDK."""
+
+from a2a.client.auth.credentials import (
+    CredentialService,
+    InMemoryContextCredentialStore,
+)
+from a2a.client.auth.interceptor import AuthInterceptor
+
+
+__all__ = [
+    'AuthInterceptor',
+    'CredentialService',
+    'InMemoryContextCredentialStore',
+]

src/a2a/client/auth/credentials.py

@@ -0,0 +1,55 @@
+from abc import ABC, abstractmethod
+
+from a2a.client.middleware import ClientCallContext
+
+
+class CredentialService(ABC):
+    """An abstract service for retrieving credentials."""
+
+    @abstractmethod
+    async def get_credentials(
+        self,
+        security_scheme_name: str,
+        context: ClientCallContext | None,
+    ) -> str | None:
+        """
+        Retrieves a credential (e.g., token) for a security scheme.
+        """
+
+
+class InMemoryContextCredentialStore(CredentialService):
+    """A simple in-memory store for session-keyed credentials.
+
+    This class uses the 'sessionId' from the ClientCallContext state to
+    store and retrieve credentials...
+    """
+
+    def __init__(self) -> None:
+        self._store: dict[str, dict[str, str]] = {}
+
+    async def get_credentials(
+        self,
+        security_scheme_name: str,
+        context: ClientCallContext | None,
+    ) -> str | None:
+        """Retrieves credentials from the in-memory store.
+
+        Args:
+            security_scheme_name: The name of the security scheme.
+            context: The client call context.
+
+        Returns:
+            The credential string, or None if not found.
+        """
+        if not context or 'sessionId' not in context.state:
+            return None
+        session_id = context.state['sessionId']
+        return self._store.get(session_id, {}).get(security_scheme_name)
+
+    async def set_credentials(
+        self, session_id: str, security_scheme_name: str, credential: str
+    ) -> None:
+        """Method to populate the store."""
+        if session_id not in self._store:
+            self._store[session_id] = {}
+        self._store[session_id][security_scheme_name] = credential

src/a2a/client/auth/interceptor.py

@@ -0,0 +1,93 @@
+import logging  # noqa: I001
+from typing import Any
+
+from a2a.client.auth.credentials import CredentialService
+from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
+from a2a.types import (
+    AgentCard,
+    APIKeySecurityScheme,
+    HTTPAuthSecurityScheme,
+    In,
+    OAuth2SecurityScheme,
+    OpenIdConnectSecurityScheme,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AuthInterceptor(ClientCallInterceptor):
+    """An interceptor that automatically adds authentication details to requests.
+
+    Based on the agent's security schemes.
+    """
+
+    def __init__(self, credential_service: CredentialService):
+        self._credential_service = credential_service
+
+    async def intercept(
+        self,
+        method_name: str,
+        request_payload: dict[str, Any],
+        http_kwargs: dict[str, Any],
+        agent_card: AgentCard | None,
+        context: ClientCallContext | None,
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Applies authentication headers to the request if credentials are available."""
+        if (
+            agent_card is None
+            or agent_card.security is None
+            or agent_card.securitySchemes is None
+        ):
+            return request_payload, http_kwargs
+
+        for requirement in agent_card.security:
+            for scheme_name in requirement:
+                credential = await self._credential_service.get_credentials(
+                    scheme_name, context
+                )
+                if credential and scheme_name in agent_card.securitySchemes:
+                    scheme_def_union = agent_card.securitySchemes.get(
+                        scheme_name
+                    )
+                    if not scheme_def_union:
+                        continue
+                    scheme_def = scheme_def_union.root
+
+                    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(
+                                f"Added Bearer token for scheme '{scheme_name}' (type: {scheme_def.type})."
+                            )
+                            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(
+                                f"Added Bearer token for scheme '{scheme_name}' (type: {scheme_def.type})."
+                            )
+                            http_kwargs['headers'] = headers
+                            return request_payload, http_kwargs
+
+                        # Case 2: API Key in Header
+                        case APIKeySecurityScheme(in_=In.header):
+                            headers[scheme_def.name] = credential
+                            logger.debug(
+                                f"Added API Key Header for scheme '{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.
+
+        return request_payload, http_kwargs