common logic for all bearer based schemes

Author: dmandar

src/a2a/client/auth/interceptor.py

@@ -5,7 +5,7 @@
 
 from a2a.client.auth.credentials import CredentialService
 from a2a.client.middleware import ClientCallContext, ClientCallInterceptor
-from a2a.types import AgentCard, APIKeySecurityScheme, HTTPAuthSecurityScheme, In, OAuth2SecurityScheme
+from a2a.types import AgentCard, APIKeySecurityScheme, HTTPAuthSecurityScheme, In, OAuth2SecurityScheme, OpenIdConnectSecurityScheme
 
 logger = logging.getLogger(__name__)
 
@@ -31,36 +31,35 @@ async def intercept(
             return request_payload, http_kwargs
 
         for requirement in agent_card.security:
-            for scheme_name in requirement: # Iterate through scheme names in the requirement
+            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[scheme_name]
-                    if not scheme_def_union:
-                        continue 
-                    scheme_def = scheme_def_union.root # SecurityScheme is a RootModel
+                    if not scheme_def_union: 
+                        continue
+                    scheme_def = scheme_def_union.root 
 
                     headers = http_kwargs.get('headers', {})
 
-                    if isinstance(scheme_def, HTTPAuthSecurityScheme):
-                        if scheme_def.scheme.lower() == 'bearer':
-                            headers['Authorization'] = f"Bearer {credential}"
-                            logger.debug(f"Added HTTP Bearer Auth for scheme '{scheme_name}'.")
-                            http_kwargs['headers'] = headers
-                            return request_payload, http_kwargs
-                    elif isinstance(scheme_def, OAuth2SecurityScheme): # New condition for OAuth2
-                        # For OAuth2, the credential obtained is the access token, used as a Bearer token.
+                    is_bearer_scheme = False
+                    if isinstance(scheme_def, HTTPAuthSecurityScheme) and scheme_def.scheme.lower() == 'bearer':
+                        is_bearer_scheme = True
+                    elif isinstance(scheme_def, (OAuth2SecurityScheme, OpenIdConnectSecurityScheme)):
+                        is_bearer_scheme = True
+
+                    if is_bearer_scheme:
                         headers['Authorization'] = f"Bearer {credential}"
-                        logger.debug(f"Added OAuth2 Bearer token for scheme '{scheme_name}'.")
+                        logger.debug(f"Added Bearer token for scheme '{scheme_name}' (type: {scheme_def.type}).")
                         http_kwargs['headers'] = headers
                         return request_payload, http_kwargs
                     elif isinstance(scheme_def, APIKeySecurityScheme):
-                        if scheme_def.in_ == In.header: # Use In.header enum member
+                        if scheme_def.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: API keys in query or cookie are not handled by this interceptor modification.
+                        # Note: API keys in query or cookie are not handled here.
 
         return request_payload, http_kwargs

tests/client/test_auth_middleware.py

@@ -25,6 +25,7 @@
     In,
     OAuth2SecurityScheme,
     OAuthFlows,
+    OpenIdConnectSecurityScheme,
     SecurityScheme,
     SendMessageRequest,
 )
@@ -302,4 +303,69 @@ async def test_auth_interceptor_with_oauth2_scheme():
         assert len(respx.calls) == 1
         request_sent = respx.calls.last.request
         assert "Authorization" in request_sent.headers
-        assert request_sent.headers["Authorization"] == f"Bearer {access_token}"
+        assert request_sent.headers["Authorization"] == f"Bearer {access_token}"
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_auth_interceptor_with_oidc_scheme():
+    """
+    Tests the AuthInterceptor with an OpenIdConnectSecurityScheme.
+    Ensures it correctly sets the Authorization: Bearer <token> header.
+    """
+    # Arrange
+    test_url = "http://oidc-agent.com/rpc"
+    context_id = "user-session-oidc"
+    scheme_name = "myOidcScheme"
+    id_token = "secret-oidc-id-token" # Or access_token
+
+    cred_store = InMemoryContextCredentialStore()
+    await cred_store.set_credentials(context_id, scheme_name, id_token)
+
+    auth_interceptor = AuthInterceptor(credential_service=cred_store)
+
+    agent_card = AgentCard(
+        url=test_url,
+        name="OidcBot",
+        description="A bot that uses OpenID Connect",
+        version="1.0",
+        defaultInputModes=[],
+        defaultOutputModes=[],
+        skills=[],
+        capabilities=AgentCapabilities(),
+        security=[{scheme_name: []}], # Security requirement referencing the scheme
+        securitySchemes={
+            scheme_name: SecurityScheme(
+                root=OpenIdConnectSecurityScheme(
+                    type="openIdConnect", 
+                    openIdConnectUrl="http://provider.com/.well-known/openid-configuration"
+                )
+            )
+        },
+    )
+
+    async with httpx.AsyncClient() as http_client:
+        client = A2AClient(
+            httpx_client=http_client,
+            agent_card=agent_card,
+            interceptors=[auth_interceptor]
+        )
+        
+        minimal_success_response = {
+            "jsonrpc": "2.0",
+            "id": "oidc_test_1",
+            "result": {"kind": "message", "messageId": "response-msg-oidc", "role": "agent", "parts": []}
+        }
+        respx.post(test_url).mock(return_value=httpx.Response(200, json=minimal_success_response))
+
+        # Act
+        context = ClientCallContext(state={"contextId": context_id})
+        await client.send_message(
+            request=SendMessageRequest(id="oidc_test_1", params={"message": {"messageId": "msg-oidc", "role": "user", "parts": []}}),
+            context=context
+        )
+
+        # Assert
+        assert len(respx.calls) == 1
+        request_sent = respx.calls.last.request
+        assert "Authorization" in request_sent.headers
+        assert request_sent.headers["Authorization"] == f"Bearer {id_token}"