refactor: Migrate authentication logic to the new interceptor pattern and add streaming support to base client interceptors.

Author: guglielmo-san

src/a2a/client/auth/interceptor.py

@@ -1,14 +1,17 @@
 import logging  # noqa: I001
-from typing import Any
 
 from a2a.client.auth.credentials import CredentialService
 from a2a.client.client import ClientCallContext
-from a2a.types.a2a_pb2 import AgentCard
+from a2a.client.interceptors import (
+    ClientCallInterceptor,
+    UnionAfterArgs,
+    UnionBeforeArgs,
+)
 
 logger = logging.getLogger(__name__)
 
 
-class AuthInterceptor:
+class AuthInterceptor(ClientCallInterceptor):
     """An interceptor that automatically adds authentication details to requests.
 
     Based on the agent's security schemes.
@@ -17,79 +20,79 @@ class AuthInterceptor:
     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]]:
+    async def before(self, args: UnionBeforeArgs) -> None:
         """Applies authentication headers to the request if credentials are available."""
+        agent_card = args.agent_card
+
         # Proto3 repeated fields (security) and maps (security_schemes) do not track presence.
         # HasField() raises ValueError for them.
         # We check for truthiness to see if they are non-empty.
         if (
-            agent_card is None
-            or not agent_card.security_requirements
+            not agent_card.security_requirements
             or not agent_card.security_schemes
         ):
-            return request_payload, http_kwargs
+            return
 
         for requirement in agent_card.security_requirements:
             for scheme_name in requirement.schemes:
                 credential = await self._credential_service.get_credentials(
-                    scheme_name, context
+                    scheme_name, args.context
                 )
                 if credential and scheme_name in agent_card.security_schemes:
                     scheme = agent_card.security_schemes.get(scheme_name)
                     if not scheme:
                         continue
 
-                    headers = http_kwargs.get('headers', {})
+                    if args.context is None:
+                        args.context = ClientCallContext()
+
+                    if args.context.service_parameters is None:
+                        args.context.service_parameters = {}
 
                     # HTTP Bearer authentication
                     if (
                         scheme.HasField('http_auth_security_scheme')
                         and scheme.http_auth_security_scheme.scheme.lower()
                         == 'bearer'
                     ):
-                        headers['Authorization'] = f'Bearer {credential}'
+                        args.context.service_parameters['Authorization'] = (
+                            f'Bearer {credential}'
+                        )
                         logger.debug(
                             "Added Bearer token for scheme '%s'.",
                             scheme_name,
                         )
-                        http_kwargs['headers'] = headers
-                        return request_payload, http_kwargs
+                        return
 
                     # 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}'
+                        args.context.service_parameters['Authorization'] = (
+                            f'Bearer {credential}'
+                        )
                         logger.debug(
                             "Added Bearer token for scheme '%s'.",
                             scheme_name,
                         )
-                        http_kwargs['headers'] = headers
-                        return request_payload, http_kwargs
+                        return
 
                     # 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
-                        )
+                        args.context.service_parameters[
+                            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
+                        return
 
                 # Note: Other cases like API keys in query/cookie are not handled and will be skipped.
 
-        return request_payload, http_kwargs
+    async def after(self, args: UnionAfterArgs) -> None:
+        """Invoked after the method is executed."""

src/a2a/client/base_client.py

@@ -515,6 +515,7 @@ async def _intercept_after(
     async def _format_stream_event(
         self, stream_response: StreamResponse, tracker: ClientTaskManager
     ) -> ClientEvent:
+        client_event: ClientEvent
         if stream_response.HasField('message'):
             client_event = (stream_response, None)
             await self.consume(client_event, self._card)

tests/client/test_auth_middleware.py tests/client/test_auth_interceptor.pytests/client/test_auth_middleware.py renamed to tests/client/test_auth_interceptor.py

@@ -17,7 +17,7 @@
     ClientFactory,
     InMemoryContextCredentialStore,
 )
-from a2a.utils.constants import TransportProtocol
+from a2a.client.interceptors import BeforeArgs, ClientCallInput
 from a2a.types.a2a_pb2 import (
     APIKeySecurityScheme,
     AgentCapabilities,
@@ -36,6 +36,7 @@
     SendMessageResponse,
     StringList,
 )
+from a2a.utils.constants import TransportProtocol
 
 
 class HeaderInterceptor(ClientCallInterceptor):
@@ -64,7 +65,6 @@ async def intercept(
 
 def build_success_response(request: httpx.Request) -> httpx.Response:
     """Creates a valid JSON-RPC success response based on the request."""
-    from a2a.types.a2a_pb2 import SendMessageResponse
 
     request_payload = json.loads(request.content)
     message = Message(
@@ -120,19 +120,17 @@ async def test_auth_interceptor_skips_when_no_agent_card(
     store: InMemoryContextCredentialStore,
 ) -> None:
     """Tests that the AuthInterceptor does not modify the request when no AgentCard is provided."""
-    request_payload = {'foo': 'bar'}
-    http_kwargs = {'fizz': 'buzz'}
     auth_interceptor = AuthInterceptor(credential_service=store)
-
-    new_payload, new_kwargs = await auth_interceptor.intercept(
-        method_name='SendMessage',
-        request_payload=request_payload,
-        http_kwargs=http_kwargs,
-        agent_card=None,
-        context=ClientCallContext(state={}),
+    request = SendMessageRequest(message=Message())
+    context = ClientCallContext(state={})
+    args = BeforeArgs(
+        input=ClientCallInput(method='send_message', value=request),
+        agent_card=AgentCard(),
+        context=context,
     )
-    assert new_payload == request_payload
-    assert new_kwargs == http_kwargs
+
+    await auth_interceptor.before(args)
+    assert context.service_parameters is None
 
 
 @pytest.mark.asyncio
@@ -210,14 +208,13 @@ def wrap_security_scheme(scheme: Any) -> SecurityScheme:
     """Wraps a security scheme in the correct SecurityScheme proto field."""
     if isinstance(scheme, APIKeySecurityScheme):
         return SecurityScheme(api_key_security_scheme=scheme)
-    elif isinstance(scheme, HTTPAuthSecurityScheme):
+    if isinstance(scheme, HTTPAuthSecurityScheme):
         return SecurityScheme(http_auth_security_scheme=scheme)
-    elif isinstance(scheme, OAuth2SecurityScheme):
+    if isinstance(scheme, OAuth2SecurityScheme):
         return SecurityScheme(oauth2_security_scheme=scheme)
-    elif isinstance(scheme, OpenIdConnectSecurityScheme):
+    if isinstance(scheme, OpenIdConnectSecurityScheme):
         return SecurityScheme(open_id_connect_security_scheme=scheme)
-    else:
-        raise ValueError(f'Unknown security scheme type: {type(scheme)}')
+    raise ValueError(f'Unknown security scheme type: {type(scheme)}')
 
 
 @dataclass
@@ -363,8 +360,6 @@ async def test_auth_interceptor_skips_when_scheme_not_in_security_schemes(
     scheme_name = 'missing'
     session_id = 'session-id'
     credential = 'test-token'
-    request_payload = {'foo': 'bar'}
-    http_kwargs = {'fizz': 'buzz'}
     await store.set_credentials(session_id, scheme_name, credential)
     auth_interceptor = AuthInterceptor(credential_service=store)
     agent_card = AgentCard(
@@ -386,13 +381,13 @@ async def test_auth_interceptor_skips_when_scheme_not_in_security_schemes(
         ],
         security_schemes={},
     )
-
-    new_payload, new_kwargs = await auth_interceptor.intercept(
-        method_name='SendMessage',
-        request_payload=request_payload,
-        http_kwargs=http_kwargs,
+    request = SendMessageRequest(message=Message())
+    context = ClientCallContext(state={'sessionId': session_id})
+    args = BeforeArgs(
+        input=ClientCallInput(method='send_message', value=request),
         agent_card=agent_card,
-        context=ClientCallContext(state={'sessionId': session_id}),
+        context=context,
     )
-    assert new_payload == request_payload
-    assert new_kwargs == http_kwargs
+
+    await auth_interceptor.before(args)
+    assert context.service_parameters is None

tests/client/test_base_client_interceptors.py

@@ -17,6 +17,8 @@
     AgentCapabilities,
     AgentCard,
     AgentInterface,
+    Message,
+    StreamResponse,
 )
 
 
@@ -142,3 +144,101 @@ async def mock_before_with_early_return(args: BeforeArgs):
         assert isinstance(after_args, AfterArgs)
         assert after_args.result.value == 'early_result'
         assert after_args.context == context
+
+    @pytest.mark.asyncio
+    async def test_execute_stream_with_interceptors_normal_flow(
+        self,
+        base_client: BaseClient,
+        mock_interceptor: AsyncMock,
+    ):
+        input_data = ClientCallInput(
+            method='send_message_streaming', value=MagicMock()
+        )
+        context = MagicMock()
+
+        async def mock_transport_call(*args, **kwargs):
+            yield StreamResponse(message=Message(message_id='1'))
+
+        # Set up mock interceptor to just pass through
+        mock_interceptor.before.return_value = None
+
+        events = [
+            e
+            async for e in base_client._execute_stream_with_interceptors(
+                input_data=input_data,
+                context=context,
+                transport_call=mock_transport_call,
+            )
+        ]
+
+        assert len(events) == 1
+
+        # Verify before was called
+        mock_interceptor.before.assert_called_once()
+        before_args = mock_interceptor.before.call_args[0][0]
+        assert isinstance(before_args, BeforeArgs)
+        assert before_args.input == input_data
+        assert before_args.context == context
+
+        # Verify after was called
+        mock_interceptor.after.assert_called_once()
+        after_args = mock_interceptor.after.call_args[0][0]
+        assert isinstance(after_args, AfterArgs)
+        assert after_args.result.method == 'send_message_streaming'
+
+    @pytest.mark.asyncio
+    async def test_execute_stream_with_interceptors_early_return(
+        self,
+        base_client: BaseClient,
+        mock_interceptor: AsyncMock,
+    ):
+        input_data = ClientCallInput(
+            method='send_message_streaming', value=MagicMock()
+        )
+        context = MagicMock()
+        mock_transport_call = AsyncMock()
+
+        # Set up early return in before
+        early_return_result = ClientCallResult(
+            method='send_message_streaming',
+            value=StreamResponse(message=Message(message_id='2')),
+        )
+
+        async def mock_before_with_early_return(args: BeforeArgs):
+            args.early_return = early_return_result
+            return {
+                'early_return': early_return_result,
+                'executed': [mock_interceptor],
+            }
+
+        mock_interceptor.before.side_effect = mock_before_with_early_return
+
+        # Override BaseClient's _intercept_before to respect our early return setup
+        # as the test's mock interceptor replaces the actual list items
+        base_client._intercept_before = AsyncMock(
+            return_value={
+                'early_return': early_return_result,
+                'executed': [mock_interceptor],
+            }
+        )
+
+        events = [
+            e
+            async for e in base_client._execute_stream_with_interceptors(
+                input_data=input_data,
+                context=context,
+                transport_call=mock_transport_call,
+            )
+        ]
+
+        assert len(events) == 1
+
+        # Verify transport call was NOT made
+        mock_transport_call.assert_not_called()
+
+        # Verify after was called with early return value
+        mock_interceptor.after.assert_called_once()
+        after_args = mock_interceptor.after.call_args[0][0]
+        assert isinstance(after_args, AfterArgs)
+        assert after_args.result.method == 'send_message_streaming'
+        assert after_args.context == context