Add VeADK A2A auth switch and utility enhancements

loveyana · loveyana · commit 2e0444692f58 · 2025-11-20T19:21:16.000+08:00
Introduces a new `to_a2a` utility in `veadk.a2a.utils.agent_to_a2a` to wrap Google ADK's A2A conversion with optional VeADK authentication and credential service integration. Adds comprehensive tests for the auth switch, refactors `RemoteVeAgent` to ensure pre-run initialization/auth logic always executes, and renames `credential_service.py` to `ve_credential_service.py` for clarity.
diff --git a/test_to_a2a_auth_switch.py b/test_to_a2a_auth_switch.py
@@ -0,0 +1,109 @@
+"""Test to verify the enable_veadk_auth switch in to_a2a function."""
+
+import pytest
+from veadk import Agent, Runner
+from veadk.a2a.utils.agent_to_a2a import to_a2a
+from veadk.auth.ve_credential_service import VeCredentialService
+from google.adk.auth.credential_service.base_credential_service import BaseCredentialService
+
+
+class MockCredentialService(BaseCredentialService):
+    """Mock credential service for testing."""
+    pass
+
+
+def test_to_a2a_without_veadk_auth():
+    """Test to_a2a with enable_veadk_auth=False (default)."""
+    agent = Agent(name="test_agent")
+    app = to_a2a(agent, enable_veadk_auth=False)
+    
+    # App should be created successfully
+    assert app is not None
+    print("✅ Test 1 passed: to_a2a works without VeADK auth")
+
+
+def test_to_a2a_with_veadk_auth_no_runner():
+    """Test to_a2a with enable_veadk_auth=True and no runner provided."""
+    agent = Agent(name="test_agent")
+    app = to_a2a(agent, enable_veadk_auth=True)
+    
+    # App should be created successfully with VeCredentialService
+    assert app is not None
+    print("✅ Test 2 passed: to_a2a creates runner with VeCredentialService")
+
+
+def test_to_a2a_with_veadk_auth_runner_with_ve_credential_service():
+    """Test to_a2a with enable_veadk_auth=True and runner with VeCredentialService."""
+    agent = Agent(name="test_agent")
+    credential_service = VeCredentialService()
+    runner = Runner(agent=agent, credential_service=credential_service)
+    
+    app = to_a2a(agent, runner=runner, enable_veadk_auth=True)
+    
+    # App should be created successfully
+    assert app is not None
+    # Runner should still have the same credential service
+    assert runner.credential_service is credential_service
+    print("✅ Test 3 passed: to_a2a accepts runner with VeCredentialService")
+
+
+def test_to_a2a_with_veadk_auth_runner_without_credential_service():
+    """Test to_a2a with enable_veadk_auth=True and runner without credential_service."""
+    agent = Agent(name="test_agent")
+    runner = Runner(agent=agent)
+    
+    # Runner initially has no credential_service (or None)
+    initial_credential_service = getattr(runner, 'credential_service', None)
+    
+    app = to_a2a(agent, runner=runner, enable_veadk_auth=True)
+    
+    # App should be created successfully
+    assert app is not None
+    # Runner should now have a VeCredentialService
+    assert hasattr(runner, 'credential_service')
+    assert isinstance(runner.credential_service, VeCredentialService)
+    print("✅ Test 4 passed: to_a2a adds VeCredentialService to runner")
+
+
+def test_to_a2a_with_veadk_auth_runner_with_wrong_credential_service():
+    """Test to_a2a with enable_veadk_auth=True and runner with non-VeCredentialService."""
+    agent = Agent(name="test_agent")
+    mock_credential_service = MockCredentialService()
+    runner = Runner(agent=agent, credential_service=mock_credential_service)
+    
+    # Should raise TypeError
+    with pytest.raises(TypeError) as exc_info:
+        to_a2a(agent, runner=runner, enable_veadk_auth=True)
+    
+    assert "must be a VeCredentialService instance" in str(exc_info.value)
+    assert "MockCredentialService" in str(exc_info.value)
+    print("✅ Test 5 passed: to_a2a raises TypeError for wrong credential service type")
+
+
+def test_to_a2a_without_veadk_auth_accepts_any_credential_service():
+    """Test to_a2a with enable_veadk_auth=False accepts any credential service."""
+    agent = Agent(name="test_agent")
+    mock_credential_service = MockCredentialService()
+    runner = Runner(agent=agent, credential_service=mock_credential_service)
+    
+    # Should work fine when VeADK auth is disabled
+    app = to_a2a(agent, runner=runner, enable_veadk_auth=False)
+    
+    assert app is not None
+    # Runner should still have the mock credential service
+    assert runner.credential_service is mock_credential_service
+    print("✅ Test 6 passed: to_a2a accepts any credential service when auth disabled")
+
+
+if __name__ == "__main__":
+    print("Running to_a2a auth switch tests...\n")
+    
+    test_to_a2a_without_veadk_auth()
+    test_to_a2a_with_veadk_auth_no_runner()
+    test_to_a2a_with_veadk_auth_runner_with_ve_credential_service()
+    test_to_a2a_with_veadk_auth_runner_without_credential_service()
+    test_to_a2a_with_veadk_auth_runner_with_wrong_credential_service()
+    test_to_a2a_without_veadk_auth_accepts_any_credential_service()
+    
+    print("\n🎉 All tests passed!")
+
diff --git a/veadk/a2a/remote_ve_agent.py b/veadk/a2a/remote_ve_agent.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import functools
 from typing import AsyncGenerator, Literal, Optional
 
 from a2a.client.base_client import BaseClient
@@ -221,41 +222,70 @@ def __init__(
         if auth_method:
             self.auth_method = auth_method
 
-    async def _run_async_impl(
-        self, ctx: InvocationContext
-    ) -> AsyncGenerator[Event, None]:
-        """Run the remote agent with credential injection support.
+        # Wrap _run_async_impl with pre-run hook to ensure initialization
+        # and authentication logic always executes, even if users override _run_async_impl
+        self._wrap_run_async_impl()
 
-        This method:
-        1. Ensures the agent is resolved (agent card fetched, client initialized)
-        2. Injects authentication token from credential service if available
-        3. Delegates to parent class for actual execution
+    def _wrap_run_async_impl(self) -> None:
+        """Wrap _run_async_impl with a decorator that ensures pre-run logic executes.
+
+        This method wraps the _run_async_impl method with a decorator that:
+        1. Executes _pre_run before the actual implementation
+        2. Handles errors from _pre_run and yields error events
+        3. Ensures the wrapper works even if users override _run_async_impl
+
+        The wrapper is applied by replacing the bound method on the instance.
+        """
+        # Store the original _run_async_impl method
+        original_run_async_impl = self._run_async_impl
+
+        @functools.wraps(original_run_async_impl)
+        async def wrapped_run_async_impl(
+            ctx: InvocationContext,
+        ) -> AsyncGenerator[Event, None]:
+            """Wrapped version of _run_async_impl with pre-run hook."""
+            # Execute pre-run initialization
+            try:
+                await self._pre_run(ctx)
+            except Exception as e:
+                yield Event(
+                    author=self.name,
+                    error_message=f"Failed to initialize remote A2A agent: {e}",
+                    invocation_id=ctx.invocation_id,
+                    branch=ctx.branch,
+                )
+                return
+
+            # Call the original (or overridden) _run_async_impl
+            async with Aclosing(original_run_async_impl(ctx)) as agen:
+                async for event in agen:
+                    yield event
+
+        # Replace the instance method with the wrapped version
+        self._run_async_impl = wrapped_run_async_impl
+
+    async def _pre_run(self, ctx: InvocationContext) -> None:
+        """Pre-run initialization and authentication setup.
+
+        This method is called before the actual agent execution to:
+        1. Ensure the agent is resolved (agent card fetched, client initialized)
+        2. Inject authentication token from credential service if available
+
+        This method is separated from _run_async_impl to ensure these critical
+        initialization steps are always executed, even if users override _run_async_impl.
 
         Args:
             ctx: Invocation context containing session and user information
 
-        Yields:
-            Events from the remote agent execution
+        Raises:
+            Exception: If agent initialization fails
         """
-        try:
-            await self._ensure_resolved()
-        except Exception as e:
-            yield Event(
-                author=self.name,
-                error_message=f"Failed to initialize remote A2A agent: {e}",
-                invocation_id=ctx.invocation_id,
-                branch=ctx.branch,
-            )
-            return
+        # Ensure agent is resolved
+        await self._ensure_resolved()
 
         # Inject auth token if credential service is available
         await self._inject_auth_token(ctx)
 
-        # Delegate to parent class for execution
-        async with Aclosing(super()._run_async_impl(ctx)) as agen:
-            async for event in agen:
-                yield event
-
     async def _inject_auth_token(self, ctx: InvocationContext) -> None:
         """Inject authentication token from credential service into the HTTP client.
 
diff --git a/veadk/a2a/utils/__init__.py b/veadk/a2a/utils/__init__.py
@@ -0,0 +1,13 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/veadk/a2a/utils/agent_to_a2a.py b/veadk/a2a/utils/agent_to_a2a.py
@@ -0,0 +1,170 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+from typing import Literal, Optional, Union
+
+from a2a.types import AgentCard
+from google.adk.agents import BaseAgent
+from starlette.applications import Starlette
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+from google.adk.a2a.utils.agent_to_a2a import to_a2a as google_adk_to_a2a
+from veadk import Runner
+from veadk.a2a.ve_middlewares import build_a2a_auth_middleware
+from veadk.auth.ve_credential_service import VeCredentialService
+from veadk.consts import DEFAULT_AGENT_NAME
+
+
+def to_a2a(
+    agent: BaseAgent,
+    *,
+    host: str = "localhost",
+    port: int = 8000,
+    protocol: str = "http",
+    agent_card: Optional[Union[AgentCard, str]] = None,
+    runner: Optional[Runner] = None,
+    enable_auth: bool = False,
+    auth_method: Literal["header", "querystring"] = "header",
+) -> Starlette:
+    """Convert an ADK agent to a A2A Starlette application with optional VeADK enhancements.
+
+    This function wraps Google ADK's to_a2a utility and optionally adds:
+    - VeCredentialService for authentication management
+    - A2A authentication middleware for token validation
+
+    Args:
+        agent: The ADK agent to convert to A2A server
+        host: The host for the A2A RPC URL (default: "localhost")
+        port: The port for the A2A RPC URL (default: 8000)
+        protocol: The protocol for the A2A RPC URL (default: "http")
+        agent_card: Optional pre-built AgentCard object or path to agent card
+                    JSON. If not provided, will be built automatically from the
+                    agent.
+        runner: Optional pre-built Runner object. If not provided, a default
+                runner will be created using in-memory services.
+                When enable_auth=True:
+                - If runner is provided and has a credential_service, it must be
+                  a VeCredentialService instance (raises TypeError otherwise)
+                - If runner is provided without credential_service, a new
+                  VeCredentialService will be created and set
+                - If runner is not provided, a new runner with VeCredentialService
+                  will be created
+        auth_method: Authentication method for A2A requests (only used when
+                    enable_auth=True). Options:
+                    - "header": Extract token from Authorization header (default)
+                    - "querystring": Extract token from query parameter
+        enable_auth: Whether to enable VeADK authentication features.
+                          When True, enables credential service and auth middleware.
+                          When False, uses standard Google ADK behavior.
+                          Default: False
+
+    Returns:
+        A Starlette application that can be run with uvicorn
+
+    Raises:
+        TypeError: If enable_auth=True and runner has a credential_service
+                  that is not a VeCredentialService instance
+
+    Example:
+        Basic usage (without VeADK auth):
+        ```python
+        from veadk import Agent
+        from veadk.a2a.utils.agent_to_a2a import to_a2a
+
+        agent = Agent(name="my_agent", tools=[...])
+        app = to_a2a(agent, host="localhost", port=8000)
+        # Run with: uvicorn module:app --host localhost --port 8000
+        ```
+
+        With VeADK authentication enabled:
+        ```python
+        app = to_a2a(agent, enable_auth=True)
+        ```
+
+        With custom runner and VeADK auth:
+        ```python
+        from veadk import Agent, Runner
+        from veadk.memory.short_term_memory import ShortTermMemory
+        from veadk.auth.ve_credential_service import VeCredentialService
+
+        agent = Agent(name="my_agent")
+        runner = Runner(
+            agent=agent,
+            short_term_memory=ShortTermMemory(),
+            app_name="my_app",
+            credential_service=VeCredentialService()  # Optional
+        )
+        app = to_a2a(agent, runner=runner, enable_auth=True)
+        ```
+
+        With querystring authentication:
+        ```python
+        app = to_a2a(agent, enable_auth=True, auth_method="querystring")
+        ```
+    """
+    app_name = agent.name or DEFAULT_AGENT_NAME
+    middleware = None # May need support multiple middlewares in the future
+
+    # Handle VeADK authentication setup
+    if enable_auth:
+        # Create credential service if not provided
+        credential_service = VeCredentialService()
+        if runner is not None:
+            # Check if runner has credential_service
+            if runner.credential_service is not None:
+                # Validate that it's a VeCredentialService
+                if not isinstance(runner.credential_service, VeCredentialService):
+                    raise TypeError(
+                        f"When enable_auth=True, runner.credential_service must be "
+                        f"a VeCredentialService instance, got {type(runner.credential_service).__name__}"
+                    )
+                # Use existing credential service
+                credential_service = runner.credential_service
+            else:
+                # Add credential_service to runner
+                runner.credential_service = credential_service
+        else:
+            # Create runner with credential_service
+            runner = Runner(
+                app_name=app_name,
+                agent=agent,
+                artifact_service=InMemoryArtifactService(),
+                session_service=InMemorySessionService(),
+                memory_service=InMemoryMemoryService(),
+                credential_service=credential_service,
+            )
+
+        middleware = build_a2a_auth_middleware(
+            app_name=app_name,
+            credential_service=credential_service,
+            auth_method=auth_method,
+        )
+
+    # Convert agent to A2A Starlette app using Google ADK utility
+    app: Starlette = google_adk_to_a2a(
+        agent=agent,
+        host=host,
+        port=port,
+        protocol=protocol,
+        agent_card=agent_card,
+        runner=runner,
+    )
+
+    # Add VeADK authentication middleware only if enabled
+    if middleware:
+        app.add_middleware(middleware)
+
+    return app
diff --git a/veadk/auth/ve_credential_service.py b/veadk/auth/ve_credential_service.py
