Enhance agent_to_a2a with optional services

src/google/adk/a2a/utils/agent_to_a2a.py

@@ -36,6 +36,10 @@
 from starlette.applications import Starlette
 
 from ...agents.base_agent import BaseAgent
+from ...artifacts.base_artifact_service import BaseArtifactService
+from ...auth.credential_service.base_credential_service import BaseCredentialService
+from ...memory.base_memory_service import BaseMemoryService
+from ...sessions.base_session_service import BaseSessionService
 from ...artifacts.in_memory_artifact_service import InMemoryArtifactService
 from ...auth.credential_service.in_memory_credential_service import InMemoryCredentialService
 from ...cli.utils.logs import setup_adk_logger
@@ -90,6 +94,10 @@ def to_a2a(
     port: int = 8000,
     protocol: str = "http",
     agent_card: Optional[Union[AgentCard, str]] = None,
+    artifact_service: Optional[BaseArtifactService] = None,
+    credential_service: Optional[BaseCredentialService] = None,
+    memory_service: Optional[BaseMemoryService] = None,
+    session_service: Optional[BaseSessionService] = None,
 ) -> Starlette:
   """Convert an ADK agent to a A2A Starlette application.
 
@@ -101,7 +109,14 @@ def to_a2a(
       agent_card: Optional pre-built AgentCard object or path to agent card
                   JSON. If not provided, will be built automatically from the
                   agent.
-
+      artifact_service: Service for artifact management (file storage, logs, etc.).
+                        Defaults to in-memory artifact service.
+      credential_service: Service for authentication/credential management.
+                          Defaults to in-memory credential service.
+      memory_service: Service for conversation or workspace memory.
+                      Defaults to in-memory memory service.
+      session_service: Service for session management. Defaults to in-memory session service.
+
   Returns:
       A Starlette application that can be run with uvicorn
 
@@ -121,11 +136,10 @@ async def create_runner() -> Runner:
     return Runner(
         app_name=agent.name or "adk_agent",
         agent=agent,
-        # Use minimal services - in a real implementation these could be configured
-        artifact_service=InMemoryArtifactService(),
-        session_service=InMemorySessionService(),
-        memory_service=InMemoryMemoryService(),
-        credential_service=InMemoryCredentialService(),
+        artifact_service=artifact_service or InMemoryArtifactService(),
+        session_service=session_service or InMemorySessionService(),
+        memory_service=memory_service or InMemoryMemoryService(),
+        credential_service=credential_service or InMemoryCredentialService(),
     )
 
   # Create A2A components

tests/unittests/a2a/utils/test_agent_to_a2a.py

@@ -34,10 +34,14 @@
   from google.adk.a2a.utils.agent_card_builder import AgentCardBuilder
   from google.adk.a2a.utils.agent_to_a2a import to_a2a
   from google.adk.agents.base_agent import BaseAgent
+  from google.adk.artifacts.base_artifact_service import BaseArtifactService
   from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+  from google.adk.auth.credential_service.base_credential_service import BaseCredentialService
   from google.adk.auth.credential_service.in_memory_credential_service import InMemoryCredentialService
+  from google.adk.memory.base_memory_service import BaseMemoryService
   from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
   from google.adk.runners import Runner
+  from google.adk.sessions.base_session_service import BaseSessionService
   from google.adk.sessions.in_memory_session_service import InMemorySessionService
   from starlette.applications import Starlette
 except ImportError as e:
@@ -59,6 +63,53 @@ def setup_method(self):
     self.mock_agent.name = "test_agent"
     self.mock_agent.description = "Test agent description"
 
+  @pytest.fixture
+  def mock_a2a_components(self):
+    """Fixture that provides all the mocked A2A components for testing."""
+    with (
+        patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor") as mock_agent_executor_class,
+        patch("google.adk.a2a.utils.agent_to_a2a.DefaultRequestHandler") as mock_request_handler_class,
+        patch("google.adk.a2a.utils.agent_to_a2a.InMemoryTaskStore") as mock_task_store_class,
+        patch("google.adk.a2a.utils.agent_to_a2a.AgentCardBuilder") as mock_card_builder_class,
+        patch("google.adk.a2a.utils.agent_to_a2a.Starlette") as mock_starlette_class,
+        patch("google.adk.a2a.utils.agent_to_a2a.Runner") as mock_runner_class,
+    ):
+      # Create mock instances
+      mock_app = Mock(spec=Starlette)
+      mock_starlette_class.return_value = mock_app
+
+      mock_task_store = Mock(spec=InMemoryTaskStore)
+      mock_task_store_class.return_value = mock_task_store
+
+      mock_agent_executor = Mock(spec=A2aAgentExecutor)
+      mock_agent_executor_class.return_value = mock_agent_executor
+
+      mock_request_handler = Mock(spec=DefaultRequestHandler)
+      mock_request_handler_class.return_value = mock_request_handler
+
+      mock_card_builder = Mock(spec=AgentCardBuilder)
+      mock_card_builder_class.return_value = mock_card_builder
+
+      mock_runner = Mock(spec=Runner)
+      mock_runner_class.return_value = mock_runner
+
+      yield {
+          "app": mock_app,
+          "task_store": mock_task_store,
+          "agent_executor": mock_agent_executor,
+          "request_handler": mock_request_handler,
+          "card_builder": mock_card_builder,
+          "runner": mock_runner,
+          "classes": {
+              "starlette": mock_starlette_class,
+              "task_store": mock_task_store_class,
+              "agent_executor": mock_agent_executor_class,
+              "request_handler": mock_request_handler_class,
+              "card_builder": mock_card_builder_class,
+              "runner": mock_runner_class,
+          }
+      }
+
   @patch("google.adk.a2a.utils.agent_to_a2a.A2aAgentExecutor")
   @patch("google.adk.a2a.utils.agent_to_a2a.DefaultRequestHandler")
   @patch("google.adk.a2a.utils.agent_to_a2a.InMemoryTaskStore")
@@ -869,3 +920,217 @@ def test_to_a2a_with_invalid_agent_card_file_path(
     # Act & Assert
     with pytest.raises(ValueError, match="Failed to load agent card from"):
       to_a2a(self.mock_agent, agent_card="/invalid/path.json")
+
+  async def test_to_a2a_with_custom_services(self, mock_a2a_components):
+    """Test to_a2a with custom service implementations."""
+    # Arrange
+    mocks = mock_a2a_components
+    custom_artifact_service = Mock(spec=BaseArtifactService)
+    custom_session_service = Mock(spec=BaseSessionService)
+    custom_memory_service = Mock(spec=BaseMemoryService)
+    custom_credential_service = Mock(spec=BaseCredentialService)
+
+    # Act
+    result = to_a2a(
+        self.mock_agent,
+        artifact_service=custom_artifact_service,
+        session_service=custom_session_service,
+        memory_service=custom_memory_service,
+        credential_service=custom_credential_service,
+    )
+
+    # Assert
+    assert result == mocks["app"]
+    # Get the runner function that was passed to A2aAgentExecutor
+    call_args = mocks["classes"]["agent_executor"].call_args
+    runner_func = call_args[1]["runner"]
+
+    # Call the runner function to verify it creates Runner with custom services
+    await runner_func()
+
+    # Verify Runner was created with custom services
+    mocks["classes"]["runner"].assert_called_once_with(
+        app_name="test_agent",
+        agent=self.mock_agent,
+        artifact_service=custom_artifact_service,
+        session_service=custom_session_service,
+        memory_service=custom_memory_service,
+        credential_service=custom_credential_service,
+    )
+
+  async def test_to_a2a_with_none_services_uses_defaults(self, mock_a2a_components):
+    """Test to_a2a with None services uses default in-memory services."""
+    # Arrange
+    mocks = mock_a2a_components
+
+    # Act
+    result = to_a2a(
+        self.mock_agent,
+        artifact_service=None,
+        session_service=None,
+        memory_service=None,
+        credential_service=None,
+    )
+
+    # Assert
+    assert result == mocks["app"]
+    # Get the runner function that was passed to A2aAgentExecutor
+    call_args = mocks["classes"]["agent_executor"].call_args
+    runner_func = call_args[1]["runner"]
+
+    # Call the runner function to verify it creates Runner with default services
+    await runner_func()
+
+    # Verify Runner was created with default services
+    mocks["classes"]["runner"].assert_called_once()
+    call_args = mocks["classes"]["runner"].call_args[1]
+    assert call_args["app_name"] == "test_agent"
+    assert call_args["agent"] == self.mock_agent
+    # Verify the services are of the correct default types
+    assert isinstance(call_args["artifact_service"], InMemoryArtifactService)
+    assert isinstance(call_args["session_service"], InMemorySessionService)
+    assert isinstance(call_args["memory_service"], InMemoryMemoryService)
+    assert isinstance(call_args["credential_service"], InMemoryCredentialService)
+
+  async def test_to_a2a_with_mixed_services(self, mock_a2a_components):
+    """Test to_a2a with mix of custom and default services."""
+    # Arrange
+    mocks = mock_a2a_components
+    custom_artifact_service = Mock(spec=BaseArtifactService)
+    custom_memory_service = Mock(spec=BaseMemoryService)
+
+    # Act - only provide custom services for artifact and memory
+    result = to_a2a(
+        self.mock_agent,
+        artifact_service=custom_artifact_service,
+        memory_service=custom_memory_service,
+    )
+
+    # Assert
+    assert result == mocks["app"]
+    # Get the runner function that was passed to A2aAgentExecutor
+    call_args = mocks["classes"]["agent_executor"].call_args
+    runner_func = call_args[1]["runner"]
+
+    # Call the runner function to verify it creates Runner with mixed services
+    await runner_func()
+
+    # Verify Runner was created with mixed services
+    mocks["classes"]["runner"].assert_called_once()
+    call_args = mocks["classes"]["runner"].call_args[1]
+    assert call_args["app_name"] == "test_agent"
+    assert call_args["agent"] == self.mock_agent
+    # Verify custom services are used
+    assert call_args["artifact_service"] == custom_artifact_service
+    assert call_args["memory_service"] == custom_memory_service
+    # Verify default services are used for the others
+    assert isinstance(call_args["session_service"], InMemorySessionService)
+    assert isinstance(call_args["credential_service"], InMemoryCredentialService)
+
+  async def test_to_a2a_services_parameter_order_independence(self, mock_a2a_components):
+    """Test that services can be provided in any order and still work correctly."""
+    # Arrange
+    mocks = mock_a2a_components
+    custom_artifact_service = Mock(spec=BaseArtifactService)
+    custom_session_service = Mock(spec=BaseSessionService)
+    custom_memory_service = Mock(spec=BaseMemoryService)
+    custom_credential_service = Mock(spec=BaseCredentialService)
+
+    # Act - provide services in different order
+    result = to_a2a(
+        self.mock_agent,
+        credential_service=custom_credential_service,
+        artifact_service=custom_artifact_service,
+        session_service=custom_session_service,
+        memory_service=custom_memory_service,
+    )
+
+    # Assert
+    assert result == mocks["app"]
+    # Get the runner function that was passed to A2aAgentExecutor
+    call_args = mocks["classes"]["agent_executor"].call_args
+    runner_func = call_args[1]["runner"]
+
+    # Call the runner function to verify it creates Runner with correct services
+    await runner_func()
+
+    # Verify Runner was created with all custom services regardless of order
+    mocks["classes"]["runner"].assert_called_once()
+    call_args = mocks["classes"]["runner"].call_args[1]
+    assert call_args["artifact_service"] == custom_artifact_service
+    assert call_args["session_service"] == custom_session_service
+    assert call_args["memory_service"] == custom_memory_service
+    assert call_args["credential_service"] == custom_credential_service
+
+  async def test_to_a2a_with_agent_without_name_and_custom_services(self, mock_a2a_components):
+    """Test to_a2a with agent that has no name and custom services."""
+    # Arrange
+    mocks = mock_a2a_components
+    self.mock_agent.name = None
+    custom_artifact_service = Mock(spec=BaseArtifactService)
+    custom_session_service = Mock(spec=BaseSessionService)
+
+    # Act
+    result = to_a2a(
+        self.mock_agent,
+        artifact_service=custom_artifact_service,
+        session_service=custom_session_service,
+    )
+
+    # Assert
+    assert result == mocks["app"]
+    # Get the runner function that was passed to A2aAgentExecutor
+    call_args = mocks["classes"]["agent_executor"].call_args
+    runner_func = call_args[1]["runner"]
+
+    # Call the runner function to verify it creates Runner correctly
+    await runner_func()
+
+    # Verify Runner was created with default app_name and custom services
+    mocks["classes"]["runner"].assert_called_once()
+    call_args = mocks["classes"]["runner"].call_args[1]
+    assert call_args["app_name"] == "adk_agent"  # Default name when agent has no name
+    assert call_args["agent"] == self.mock_agent
+    assert call_args["artifact_service"] == custom_artifact_service
+    assert call_args["session_service"] == custom_session_service
+    # Verify default services are used for the others
+    assert isinstance(call_args["memory_service"], InMemoryMemoryService)
+    assert isinstance(call_args["credential_service"], InMemoryCredentialService)
+
+  async def test_to_a2a_service_parameter_validation(self, mock_a2a_components):
+    """Test that to_a2a accepts valid service types and passes them through correctly."""
+    # Arrange
+    mocks = mock_a2a_components
+    custom_artifact_service = Mock(spec=BaseArtifactService)
+    custom_session_service = Mock(spec=BaseSessionService)
+    custom_memory_service = Mock(spec=BaseMemoryService)
+    custom_credential_service = Mock(spec=BaseCredentialService)
+
+    # Act
+    result = to_a2a(
+        self.mock_agent,
+        artifact_service=custom_artifact_service,
+        session_service=custom_session_service,
+        memory_service=custom_memory_service,
+        credential_service=custom_credential_service,
+    )
+
+    # Assert - verify the app is returned
+    assert result == mocks["app"]
+
+    # Get the runner function that was passed to A2aAgentExecutor
+    call_args = mocks["classes"]["agent_executor"].call_args
+    runner_func = call_args[1]["runner"]
+
+    # Call the runner function to verify it creates Runner with correct services
+    await runner_func()
+
+    # Verify Runner was created with the custom services
+    mocks["classes"]["runner"].assert_called_once()
+    call_args = mocks["classes"]["runner"].call_args[1]
+    assert call_args["app_name"] == "test_agent"
+    assert call_args["agent"] == self.mock_agent
+    assert call_args["artifact_service"] == custom_artifact_service
+    assert call_args["session_service"] == custom_session_service
+    assert call_args["memory_service"] == custom_memory_service
+    assert call_args["credential_service"] == custom_credential_service