Add VeADK A2A auth switch and utility enhancements

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.

Author: loveyana

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!")
+

tests/auth/test_credential_service.py

@@ -26,7 +26,7 @@
 from google.adk.auth.auth_tool import AuthConfig
 from google.adk.agents.callback_context import CallbackContext
 
-from veadk.auth.credential_service import VeCredentialService
+from veadk.auth.ve_credential_service  import VeCredentialService
 
 
 @pytest.fixture

tests/test_ve_a2a_middlewares.py

@@ -21,7 +21,7 @@
 from starlette.responses import Response
 
 from veadk.a2a.ve_middlewares import A2AAuthMiddleware, build_a2a_auth_middleware
-from veadk.auth.credential_service import VeCredentialService
+from veadk.auth.ve_credential_service  import VeCredentialService
 from veadk.utils.auth import VE_TIP_TOKEN_HEADER
 
 

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.
 

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.