OAuth test

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
+from a2a.types import AgentCard, APIKeySecurityScheme, HTTPAuthSecurityScheme, In, OAuth2SecurityScheme
 
 logger = logging.getLogger(__name__)
 
@@ -27,35 +27,40 @@ async def intercept(
         agent_card: AgentCard | None,
         context: ClientCallContext | None,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
-        """
-        Adds authentication headers to the request if credentials can be found.
-        """
         if not agent_card or not agent_card.security or not agent_card.securitySchemes:
             return request_payload, http_kwargs
 
         for requirement in agent_card.security:
-            for scheme_name in requirement:
+            for scheme_name in requirement: # Iterate through scheme names in the requirement
                 credential = await self._credential_service.get_credentials(
                     scheme_name, context
                 )
                 if credential and scheme_name in agent_card.securitySchemes:
-                    scheme_def = agent_card.securitySchemes[scheme_name].root
+                    scheme_def_union = agent_card.securitySchemes[scheme_name]
+                    if not scheme_def_union:
+                        continue 
+                    scheme_def = scheme_def_union.root # SecurityScheme is a RootModel
+
                     headers = http_kwargs.get('headers', {})
 
                     if isinstance(scheme_def, HTTPAuthSecurityScheme):
-                        headers['Authorization'] = f"{scheme_def.scheme} {credential}"
+                        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.
+                        headers['Authorization'] = f"Bearer {credential}"
+                        logger.debug(f"Added OAuth2 Bearer token for scheme '{scheme_name}'.")
                         http_kwargs['headers'] = headers
-                        logger.debug(f"Added HTTP Auth for scheme '{scheme_name}'.")
                         return request_payload, http_kwargs
                     elif isinstance(scheme_def, APIKeySecurityScheme):
-                        if scheme_def.in_ == 'header':
+                        if scheme_def.in_ == In.header: # Use In.header enum member
                             headers[scheme_def.name] = credential
-                            http_kwargs['headers'] = headers
                             logger.debug(f"Added API Key Header for scheme '{scheme_name}'.")
+                            http_kwargs['headers'] = headers
                             return request_payload, http_kwargs
-                        else:
-                            logger.warning(
-                                f"API Key in '{scheme_def.in_}' not supported by this interceptor."
-                            )
-
+                        # Note: API keys in query or cookie are not handled by this interceptor modification.
+        
         return request_payload, http_kwargs

tests/client/test_auth_middleware.py

@@ -20,8 +20,11 @@
     AgentCard,
     AgentCapabilities,
     APIKeySecurityScheme,
+    AuthorizationCodeOAuthFlow,
     HTTPAuthSecurityScheme,
     In,
+    OAuth2SecurityScheme,
+    OAuthFlows,
     SecurityScheme,
     SendMessageRequest,
 )
@@ -230,4 +233,73 @@ async def test_auth_interceptor_with_api_key():
         assert len(respx.calls) == 1
         request = respx.calls.last.request
         assert "x-api-key" in request.headers
-        assert request.headers["x-api-key"] == api_key
+        assert request.headers["x-api-key"] == api_key
+
+
+@pytest.mark.asyncio
+@respx.mock
+async def test_auth_interceptor_with_oauth2_scheme():
+    """
+    Tests the AuthInterceptor with an OAuth2 security scheme defined in AgentCard.
+    Ensures it correctly sets the Authorization: Bearer <token> header.
+    """
+    test_url = "http://oauth-agent.com/rpc"
+    context_id = "user-session-oauth"
+    scheme_name = "myOAuthScheme" 
+    access_token = "secret-oauth-access-token"
+
+    cred_store = InMemoryContextCredentialStore()
+    await cred_store.set_credentials(context_id, scheme_name, access_token)
+
+    auth_interceptor = AuthInterceptor(credential_service=cred_store)
+
+    # Define a minimal OAuth2 flow
+    oauth_flows = OAuthFlows(
+        authorizationCode=AuthorizationCodeOAuthFlow(
+            authorizationUrl="http://provider.com/auth",
+            tokenUrl="http://provider.com/token",
+            scopes={"read": "Read scope"}
+        )
+    )
+
+    agent_card = AgentCard(
+        url=test_url,
+        name="OAuthBot",
+        description="A bot that uses OAuth2",
+        version="1.0",
+        defaultInputModes=[],
+        defaultOutputModes=[],
+        skills=[],
+        capabilities=AgentCapabilities(),
+        security=[{scheme_name: ["read"]}], 
+        securitySchemes={
+            scheme_name: SecurityScheme(root=OAuth2SecurityScheme(type="oauth2", flows=oauth_flows))
+        },
+    )
+
+    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": "oauth_test_1",
+            "result": {"kind": "message", "messageId": "response-msg-oauth", "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="oauth_test_1", params={"message": {"messageId": "msg-oauth", "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 {access_token}"