agent-starter-pack 0.18.1__py3-none-any.whl → 0.19.0__py3-none-any.whl

agent_starter_pack/deployment_targets/agent_engine/tests/integration/test_agent_engine_app.py

@@ -213,18 +213,52 @@ def test_feedback_endpoint(server_fixture: subprocess.Popen[str]) -> None:
     logger.info("Feedback endpoint test passed")
 {% else %}
 
+# mypy: disable-error-code="arg-type"
+{%- if cookiecutter.is_adk_a2a %}
+
+import os
+
+import pytest
+import pytest_asyncio
+from google.adk.artifacts import InMemoryArtifactService
+from google.adk.sessions import InMemorySessionService
+
+from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
+from tests.helpers import (
+    build_get_request,
+    build_post_request,
+    poll_task_completion,
+)
+{%- elif cookiecutter.is_adk %}
+
 import logging
 
 import pytest
-{%- if cookiecutter.is_adk %}
 from google.adk.events.event import Event
 
 from {{cookiecutter.agent_directory}}.agent import root_agent
 from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
 {%- else %}
 
+import logging
+
+import pytest
+
 from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
 {%- endif %}
+{%- if cookiecutter.is_adk_a2a %}
+
+
+@pytest_asyncio.fixture
+async def agent_app() -> AgentEngineApp:
+    """Fixture to create and set up AgentEngineApp instance"""
+    app = await AgentEngineApp.create(
+        artifact_service_builder=lambda: InMemoryArtifactService(),
+        session_service_builder=lambda: InMemorySessionService(),
+    )
+    app.set_up()
+    return app
+{%- else %}
 
 
 @pytest.fixture
@@ -237,8 +271,75 @@ def agent_app() -> AgentEngineApp:
 {%- endif %}
     app.set_up()
     return app
+{% endif %}
+{%- if cookiecutter.is_adk_a2a %}
+
+
+@pytest.mark.asyncio
+async def test_agent_on_message_send(agent_app: AgentEngineApp) -> None:
+    """Test complete A2A message workflow from send to task completion with artifacts."""
+    # Send message
+    message_data = {
+        "message": {
+            "messageId": f"msg-{os.urandom(8).hex()}",
+            "content": [{"text": "What is the capital of France?"}],
+            "role": "ROLE_USER",
+        },
+    }
+    response = await agent_app.on_message_send(
+        request=build_post_request(message_data),
+        context=None,
+    )
+
+    # Verify task creation
+    assert "task" in response and "id" in response["task"], (
+        "Expected task with ID in response"
+    )
+
+    # Poll for completion
+    final_response = await poll_task_completion(agent_app, response["task"]["id"])
+
+    # Verify artifacts
+    assert final_response.get("artifacts"), "Expected artifacts in completed task"
+    artifact = final_response["artifacts"][0]
+    assert artifact.get("parts") and artifact["parts"][0].get("text"), (
+        "Expected artifact with text content"
+    )
+
+
+@pytest.mark.asyncio
+async def test_agent_card(agent_app: AgentEngineApp) -> None:
+    """Test agent card retrieval and validation of required A2A fields."""
+    response = await agent_app.handle_authenticated_agent_card(
+        request=build_get_request(None),
+        context=None,
+    )
+
+    # Verify core agent card fields
+    assert response.get("name") == "root_agent", "Expected agent name 'root_agent'"
+    assert response.get("protocolVersion") == "0.3.0", "Expected protocol version 0.3.0"
+    assert response.get("preferredTransport") == "HTTP+JSON", (
+        "Expected HTTP+JSON transport"
+    )
+
+    # Verify capabilities
+    capabilities = response.get("capabilities", {})
+    assert capabilities.get("streaming") is False, "Expected streaming disabled"
+
+    # Verify skills
+    skills = response.get("skills", [])
+    assert len(skills) > 0, "Expected at least one skill"
+    for skill in skills:
+        assert all(key in skill for key in ["id", "name", "description"]), (
+            "Expected id, name, and description in each skill"
+        )
+
+    # Verify extended card support
+    assert response.get("supportsAuthenticatedExtendedCard") is True, (
+        "Expected supportsAuthenticatedExtendedCard to be True"
+    )
+{% elif cookiecutter.is_adk %}
 
-{% if cookiecutter.is_adk %}
 @pytest.mark.asyncio
 async def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
     """
@@ -293,6 +394,7 @@ def test_agent_feedback(agent_app: AgentEngineApp) -> None:
 
     logging.info("All assertions passed for agent feedback test")
 {% else %}
+
 def test_agent_stream_query(agent_app: AgentEngineApp) -> None:
     """
     Integration test for the agent stream query functionality.

agent_starter_pack/deployment_targets/agent_engine/tests/load_test/load_test.py

@@ -139,6 +139,149 @@ class RemoteAgentUser(WebSocketUser):
     # Set the host via command line: locust -f load_test.py --host=https://your-deployed-service.run.app
     host = "http://localhost:8000"  # Default for local testing
 {%- else %}
+{%- if cookiecutter.is_adk_a2a %}
+
+import json
+import logging
+import os
+import time
+
+from locust import HttpUser, between, task
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+# Initialize Vertex AI and load agent config
+with open("deployment_metadata.json") as f:
+    remote_agent_engine_id = json.load(f)["remote_agent_engine_id"]
+
+parts = remote_agent_engine_id.split("/")
+project_id = parts[1]
+location = parts[3]
+engine_id = parts[5]
+
+# Convert remote agent engine ID to URLs
+base_url = f"https://{location}-aiplatform.googleapis.com"
+a2a_base_path = f"/v1beta1/projects/{project_id}/locations/{location}/reasoningEngines/{engine_id}/a2a/v1"
+
+logger.info("Using remote agent engine ID: %s", remote_agent_engine_id)
+logger.info("Using base URL: %s", base_url)
+logger.info("Using API base path: %s", a2a_base_path)
+
+
+class SendMessageUser(HttpUser):
+    """Simulates a user interacting with the send message API."""
+
+    wait_time = between(1, 3)  # Wait 1-3 seconds between tasks
+    host = base_url  # Set the base host URL for Locust
+
+    @task
+    def send_message_and_poll(self) -> None:
+        """Simulates a chat interaction: sends a message and polls for completion."""
+        headers = {"Content-Type": "application/json"}
+        headers["Authorization"] = f"Bearer {os.environ['_AUTH_TOKEN']}"
+
+        data = {
+            "message": {
+                "messageId": "msg-id",
+                "content": [{"text": "Hello! What's the weather in New York?"}],
+                "role": "ROLE_USER",
+            }
+        }
+
+        e2e_start_time = time.time()
+        with self.client.post(
+            f"{a2a_base_path}/message:send",
+            headers=headers,
+            json=data,
+            catch_response=True,
+            name="/v1/message:send",
+        ) as response:
+            if response.status_code != 200:
+                response.failure(
+                    f"Send failed with status code: {response.status_code}"
+                )
+                return
+
+            response.success()
+            response_data = response.json()
+
+            # Extract task ID
+            try:
+                task_id = response_data["task"]["id"]
+            except (KeyError, TypeError) as e:
+                logger.error(f"Failed to extract task ID: {e}")
+                return
+
+        # Poll for task completion
+        max_polls = 20  # Maximum number of poll attempts
+        poll_interval = 0.5  # Seconds between polls
+        poll_count = 0
+
+        while poll_count < max_polls:
+            poll_count += 1
+            time.sleep(poll_interval)
+
+            with self.client.get(
+                f"{a2a_base_path}/tasks/{task_id}",
+                headers=headers,
+                catch_response=True,
+                name="/v1/tasks/{id}",
+            ) as poll_response:
+                if poll_response.status_code != 200:
+                    poll_response.failure(
+                        f"Poll failed with status code: {poll_response.status_code}"
+                    )
+                    return
+
+                poll_data = poll_response.json()
+
+                try:
+                    task_state = poll_data["status"]["state"]
+                except (KeyError, TypeError) as e:
+                    logger.error(f"Failed to extract task state: {e}")
+                    poll_response.failure(f"Invalid response format: {e}")
+                    return
+
+                # Check if task is complete
+                if task_state in ["TASK_STATE_COMPLETED"]:
+                    poll_response.success()
+
+                    # Measure end-to-end time
+                    e2e_duration = (time.time() - e2e_start_time) * 1000
+
+                    # Fire custom event for end-to-end metrics
+                    self.environment.events.request.fire(
+                        request_type="E2E",
+                        name="message:send_and_complete",
+                        response_time=e2e_duration,
+                        response_length=len(json.dumps(poll_data)),
+                        response=poll_response,
+                        context={"poll_count": poll_count},
+                    )
+                    return
+
+                elif task_state in ["TASK_STATE_WORKING"]:
+                    poll_response.success()
+
+                else:
+                    poll_response.failure(f"Task failed with state: {task_state}")
+                    return
+
+        # Timeout - task didn't complete in time
+        self.environment.events.request.fire(
+            request_type="TIMEOUT",
+            name="message:timeout",
+            response_time=(time.time() - e2e_start_time) * 1000,
+            response_length=0,
+            response=None,
+            context={"poll_count": poll_count},
+            exception=TimeoutError(f"Task did not complete after {max_polls} polls"),
+        )
+{%- else %}
 
 import json
 import logging
@@ -253,3 +396,4 @@ class ChatStreamUser(HttpUser):
             else:
                 response.failure(f"Unexpected status code: {response.status_code}")
 {%- endif %}
+{%- endif %}

agent_starter_pack/deployment_targets/agent_engine/tests/{% if cookiecutter.is_adk_a2a %}helpers.py{% else %}unused_helpers.py{% endif %}

@@ -0,0 +1,138 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+
+# mypy: disable-error-code="arg-type"
+
+"""Helper functions for testing AgentEngineApp with A2A protocol."""
+
+import asyncio
+import json
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Any
+
+from starlette.requests import Request
+
+if TYPE_CHECKING:
+    from {{cookiecutter.agent_directory}}.agent_engine_app import AgentEngineApp
+
+# Test constants
+POLL_MAX_ATTEMPTS = 30
+POLL_INTERVAL_SECONDS = 1.0
+TEST_ARTIFACTS_BUCKET = "test-artifacts-bucket"
+
+
+def receive_wrapper(data: dict[str, Any] | None) -> Callable[[], Awaitable[dict]]:
+    """Creates a mock ASGI receive callable for testing.
+
+    Args:
+        data: Dictionary to encode as JSON request body
+
+    Returns:
+        Async callable that returns mock ASGI receive message
+    """
+
+    async def receive() -> dict:
+        byte_data = json.dumps(data).encode("utf-8")
+        return {"type": "http.request", "body": byte_data, "more_body": False}
+
+    return receive
+
+
+def build_post_request(
+    data: dict[str, Any] | None = None, path_params: dict[str, str] | None = None
+) -> Request:
+    """Builds a mock Starlette Request object for a POST request with JSON data.
+
+    Args:
+        data: JSON data to include in request body
+        path_params: Path parameters to include in request scope
+
+    Returns:
+        Mock Starlette Request object
+    """
+    scope: dict[str, Any] = {
+        "type": "http",
+        "http_version": "1.1",
+        "headers": [(b"content-type", b"application/json")],
+        "app": None,
+    }
+    if path_params:
+        scope["path_params"] = path_params
+    receiver = receive_wrapper(data)
+    return Request(scope, receiver)
+
+
+def build_get_request(path_params: dict[str, str] | None) -> Request:
+    """Builds a mock Starlette Request object for a GET request.
+
+    Args:
+        path_params: Path parameters to include in request scope
+
+    Returns:
+        Mock Starlette Request object
+    """
+    scope: dict[str, Any] = {
+        "type": "http",
+        "http_version": "1.1",
+        "query_string": b"",
+        "app": None,
+    }
+    if path_params:
+        scope["path_params"] = path_params
+
+    async def receive() -> dict:
+        return {"type": "http.disconnect"}
+
+    return Request(scope, receive)
+
+
+async def poll_task_completion(
+    agent_app: "AgentEngineApp",
+    task_id: str,
+    max_attempts: int = POLL_MAX_ATTEMPTS,
+    interval: float = POLL_INTERVAL_SECONDS,
+) -> dict[str, Any]:
+    """Poll for task completion and return final response.
+
+    Args:
+        agent_app: The AgentEngineApp instance to poll
+        task_id: The task ID to poll for
+        max_attempts: Maximum number of polling attempts
+        interval: Seconds to wait between polls
+
+    Returns:
+        Final task response when completed
+
+    Raises:
+        AssertionError: If task fails or times out
+    """
+    for _ in range(max_attempts):
+        poll_request = build_get_request({"id": task_id})
+        response = await agent_app.on_get_task(
+            request=poll_request,
+            context=None,
+        )
+
+        task_state = response.get("status", {}).get("state", "")
+
+        if task_state == "TASK_STATE_COMPLETED":
+            return response
+        elif task_state == "TASK_STATE_FAILED":
+            raise AssertionError(f"Task failed: {response}")
+
+        await asyncio.sleep(interval)
+
+    raise AssertionError(
+        f"Task did not complete within {max_attempts * interval} seconds"
+    )

agent_starter_pack/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/agent_engine_app.py

@@ -14,6 +14,10 @@
 
 # mypy: disable-error-code="attr-defined,arg-type"
 {%- if cookiecutter.is_adk %}
+{%- if cookiecutter.is_adk_a2a %}
+import asyncio
+import copy
+{%- endif %}
 import logging
 import os
 from typing import Any
@@ -21,18 +25,35 @@ from typing import Any
 import click
 import google.auth
 import vertexai
+{%- if cookiecutter.is_adk_a2a %}
+from a2a.types import AgentCapabilities, AgentCard, TransportProtocol
+from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor
+from google.adk.a2a.utils.agent_card_builder import AgentCardBuilder
+from google.adk.apps.app import App
+{%- endif %}
 from google.adk.artifacts import GcsArtifactService
+{%- if cookiecutter.is_adk_a2a %}
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+{%- endif %}
 from google.cloud import logging as google_cloud_logging
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider, export
 from vertexai._genai.types import AgentEngine, AgentEngineConfig{%- if cookiecutter.is_adk_live %}, AgentServerMode{%- endif %}
 {%- if cookiecutter.is_adk_live %}
 from vertexai.preview.reasoning_engines import AdkApp
+{%- elif cookiecutter.is_adk_a2a %}
+from vertexai.preview.reasoning_engines import A2aAgent
 {%- else %}
 from vertexai.agent_engines.templates.adk import AdkApp
 {%- endif %}
+{%- if cookiecutter.is_adk_a2a %}
+
+from {{cookiecutter.agent_directory}}.agent import app
+{%- else %}
 
 from {{cookiecutter.agent_directory}}.agent import root_agent
+{%- endif %}
 from {{cookiecutter.agent_directory}}.utils.deployment import (
     parse_env_vars,
     print_deployment_success,
@@ -41,9 +62,53 @@ from {{cookiecutter.agent_directory}}.utils.deployment import (
 from {{cookiecutter.agent_directory}}.utils.gcs import create_bucket_if_not_exists
 from {{cookiecutter.agent_directory}}.utils.tracing import CloudTraceLoggingSpanExporter
 from {{cookiecutter.agent_directory}}.utils.typing import Feedback
+{%- if cookiecutter.is_adk_a2a %}
+
+
+class AgentEngineApp(A2aAgent):
+    @staticmethod
+    async def create(
+        artifact_service_builder: Any = None,
+        session_service_builder: Any = None,
+    ) -> Any:
+        """Create an AgentEngineApp instance."""
+
+        def create_runner() -> Runner:
+            """Create a Runner for the AgentEngineApp."""
+            return Runner(
+                app=app,
+                session_service=session_service_builder()
+                if session_service_builder
+                else None,
+                artifact_service=artifact_service_builder()
+                if artifact_service_builder
+                else None,
+            )
+
+        return AgentEngineApp(
+            agent_executor_builder=lambda: A2aAgentExecutor(runner=create_runner()),
+            agent_card=await AgentEngineApp.build_agent_card(app=app),
+        )
+
+    @staticmethod
+    async def build_agent_card(app: App) -> AgentCard:
+        """Builds the Agent Card dynamically from the app."""
+        agent_card_builder = AgentCardBuilder(
+            agent=app.root_agent,
+            # Agent Engine does not support streaming yet
+            capabilities=AgentCapabilities(streaming=False),
+            rpc_url="http://localhost:9999/",
+            agent_version=os.getenv("AGENT_VERSION", "0.1.0"),
+        )
+        agent_card = await agent_card_builder.build()
+        agent_card.preferred_transport = TransportProtocol.http_json  # Http Only.
+        agent_card.supports_authenticated_extended_card = True
+        return agent_card
+{% else %}
 
 
 class AgentEngineApp(AdkApp):
+{%- endif %}
     def set_up(self) -> None:
         """Set up logging and tracing for the agent engine app."""
         import logging
@@ -74,6 +139,16 @@ class AgentEngineApp(AdkApp):
         operations = super().register_operations()
         operations[""] = operations.get("", []) + ["register_feedback"]
         return operations
+{%- if cookiecutter.is_adk_a2a %}
+
+    def clone(self) -> "AgentEngineApp":
+        """Returns a clone of the Agent Engine application."""
+        template_attributes = self._tmpl_attrs
+        return self.__class__(
+            agent_card=copy.deepcopy(self.agent_card),
+            agent_executor_builder=self._tmpl_attrs.get("agent_executor_builder"),
+        )
+{%- endif %}
 
 {%- else %}
 import logging
@@ -272,7 +347,7 @@ def deploy_agent_engine_app(
     if not staging_bucket_uri:
         staging_bucket_uri = f"gs://{project}-agent-engine"
     if not artifacts_bucket_name:
-        artifacts_bucket_name = f"gs://{project}-agent-engine"
+        artifacts_bucket_name = f"{project}-agent-engine"
 
 {%- if "adk" in cookiecutter.tags %}
     create_bucket_if_not_exists(
@@ -303,16 +378,25 @@ def deploy_agent_engine_app(
     # Read requirements
     with open(requirements_file) as f:
         requirements = f.read().strip().split("\n")
-{% if cookiecutter.is_adk %}
+{%- if cookiecutter.is_adk_a2a %}
+    agent_engine = asyncio.run(
+        AgentEngineApp.create(
+            artifact_service_builder=lambda: GcsArtifactService(
+                bucket_name=artifacts_bucket_name
+            ),
+            session_service_builder=lambda: InMemorySessionService(),
+        )
+    )
+{%- elif cookiecutter.is_adk %}
     agent_engine = AgentEngineApp(
         agent=root_agent,
         artifact_service_builder=lambda: GcsArtifactService(
             bucket_name=artifacts_bucket_name
         ),
     )
-{% else %}
+{%- else %}
     agent_engine = AgentEngineApp(project_id=project)
-{% endif %}
+{%- endif %}
     # Set worker parallelism to 1
     env_vars["NUM_WORKERS"] = "1"
 

agent_starter_pack/deployment_targets/agent_engine/{{cookiecutter.agent_directory}}/utils/deployment.py

@@ -80,6 +80,10 @@ def print_deployment_success(
 {%- if cookiecutter.is_adk %}
 {%- if cookiecutter.is_adk_live %}
     print("\n✅ Deployment successful! Run your agent with: `make playground-remote`")
+{%- elif cookiecutter.is_adk_a2a %}
+    print(
+        "\n✅ Deployment successful! Test your agent: notebooks/adk_a2a_app_testing.ipynb"
+    )
 {%- else %}
     print(
         "\n✅ Deployment successful! Test your agent: notebooks/adk_app_testing.ipynb"

agent_starter_pack/deployment_targets/cloud_run/Dockerfile

@@ -43,6 +43,9 @@ RUN uv sync --frozen
 ARG COMMIT_SHA=""
 ENV COMMIT_SHA=${COMMIT_SHA}
 
+ARG AGENT_VERSION=0.0.0
+ENV AGENT_VERSION=${AGENT_VERSION}
+
 EXPOSE 8080
 
 CMD ["uv", "run", "uvicorn", "{{cookiecutter.agent_directory}}.server:app", "--host", "0.0.0.0", "--port", "8080"]

agent_starter_pack/deployment_targets/cloud_run/deployment/terraform/dev/service.tf

@@ -150,6 +150,13 @@ resource "google_cloud_run_v2_service" "app" {
     containers {
       image = "us-docker.pkg.dev/cloudrun/container/hello"
 
+{%- if cookiecutter.is_adk_a2a %}
+      env {
+        name  = "APP_URL"
+        value = "https://${var.project_name}-${data.google_project.project.number}.${var.region}.run.app"
+      }
+
+{%- endif %}
       resources {
         limits = {
           cpu    = "4"

agent_starter_pack/deployment_targets/cloud_run/deployment/terraform/service.tf

@@ -148,7 +148,7 @@ resource "google_alloydb_user" "db_user" {
 
 {%- endif %}
 
-resource "google_cloud_run_v2_service" "app_staging" {  
+resource "google_cloud_run_v2_service" "app_staging" {
   name                = var.project_name
   location            = var.region
   project             = var.staging_project_id
@@ -172,6 +172,13 @@ resource "google_cloud_run_v2_service" "app_staging" {
       # Placeholder, will be replaced by the CI/CD pipeline
       image = "us-docker.pkg.dev/cloudrun/container/hello"
 
+{%- if cookiecutter.is_adk_a2a %}
+      env {
+        name  = "APP_URL"
+        value = "https://${var.project_name}-${data.google_project.project["staging"].number}.${var.region}.run.app"
+      }
+
+{%- endif %}
       resources {
         limits = {
           cpu    = "4"
@@ -266,7 +273,7 @@ resource "google_cloud_run_v2_service" "app_staging" {
   depends_on = [google_project_service.deploy_project_services]
 }
 
-resource "google_cloud_run_v2_service" "app_prod" {  
+resource "google_cloud_run_v2_service" "app_prod" {
   name                = var.project_name
   location            = var.region
   project             = var.prod_project_id
@@ -290,6 +297,13 @@ resource "google_cloud_run_v2_service" "app_prod" {
       # Placeholder, will be replaced by the CI/CD pipeline
       image = "us-docker.pkg.dev/cloudrun/container/hello"
 
+{%- if cookiecutter.is_adk_a2a %}
+      env {
+        name  = "APP_URL"
+        value = "https://${var.project_name}-${data.google_project.project["prod"].number}.${var.region}.run.app"
+      }
+
+{%- endif %}
       resources {
         limits = {
           cpu    = "4"