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

agent_starter_pack/deployment_targets/cloud_run/tests/integration/test_server_e2e.py

@@ -226,6 +226,20 @@ from typing import Any
 
 import pytest
 import requests
+{%- if cookiecutter.is_adk_a2a %}
+from a2a.types import (
+    JSONRPCErrorResponse,
+    Message,
+    MessageSendParams,
+    Part,
+    Role,
+    SendMessageRequest,
+    SendMessageResponse,
+    SendStreamingMessageRequest,
+    SendStreamingMessageResponse,
+    TextPart,
+)
+{%- endif %}
 from requests.exceptions import RequestException
 
 # Configure logging
@@ -233,7 +247,10 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 BASE_URL = "http://127.0.0.1:8000/"
-{%- if cookiecutter.is_adk %}
+{%- if cookiecutter.is_adk_a2a %}
+A2A_RPC_URL = BASE_URL + "a2a/{{cookiecutter.agent_directory}}/"
+AGENT_CARD_URL = A2A_RPC_URL + ".well-known/agent-card.json"
+{%- elif cookiecutter.is_adk %}
 STREAM_URL = BASE_URL + "run_sse"
 {%- else %}
 STREAM_URL = BASE_URL + "stream_messages"
@@ -292,7 +309,11 @@ def wait_for_server(timeout: int = 90, interval: int = 1) -> bool:
     start_time = time.time()
     while time.time() - start_time < timeout:
         try:
+{%- if cookiecutter.is_adk_a2a %}
+            response = requests.get(AGENT_CARD_URL, timeout=10)
+{%- else %}
             response = requests.get("http://127.0.0.1:8000/docs", timeout=10)
+{%- endif %}
             if response.status_code == 200:
                 logger.info("Server is ready")
                 return True
@@ -323,6 +344,82 @@ def server_fixture(request: Any) -> Iterator[subprocess.Popen[str]]:
 
 
 def test_chat_stream(server_fixture: subprocess.Popen[str]) -> None:
+{%- if cookiecutter.is_adk_a2a %}
+    """Test the chat stream functionality using A2A JSON-RPC protocol."""
+    logger.info("Starting chat stream test")
+
+    message = Message(
+        message_id=f"msg-user-{uuid.uuid4()}",
+        role=Role.user,
+        parts=[Part(root=TextPart(text="What's the weather in San Francisco?"))],
+    )
+
+    request = SendStreamingMessageRequest(
+        id="test-req-001",
+        params=MessageSendParams(message=message),
+    )
+
+    # Send the request
+    response = requests.post(
+        A2A_RPC_URL,
+        headers=HEADERS,
+        json=request.model_dump(mode="json", exclude_none=True),
+        stream=True,
+        timeout=60,
+    )
+    assert response.status_code == 200
+
+    # Parse streaming JSON-RPC responses
+    responses: list[SendStreamingMessageResponse] = []
+
+    for line in response.iter_lines():
+        if line:
+            line_str = line.decode("utf-8")
+            if line_str.startswith("data: "):
+                event_json = line_str[6:]
+                json_data = json.loads(event_json)
+                streaming_response = SendStreamingMessageResponse.model_validate(
+                    json_data
+                )
+                responses.append(streaming_response)
+
+    assert responses, "No responses received from stream"
+
+    # Check for final status update
+    final_responses = [
+        r.root
+        for r in responses
+        if hasattr(r.root, "result")
+        and hasattr(r.root.result, "final")
+        and r.root.result.final is True
+    ]
+    assert final_responses, "No final response received"
+
+    final_response = final_responses[-1]
+    assert final_response.result.kind == "status-update"
+    assert hasattr(final_response.result, "status")
+    assert final_response.result.status.state == "completed"
+
+    # Check for artifact content
+    artifact_responses = [
+        r.root
+        for r in responses
+        if hasattr(r.root, "result") and r.root.result.kind == "artifact-update"
+    ]
+    assert artifact_responses, "No artifact content received in stream"
+
+    # Verify text content is in the artifact
+    artifact_response = artifact_responses[-1]
+    assert hasattr(artifact_response.result, "artifact")
+    artifact = artifact_response.result.artifact
+    assert artifact.parts, "Artifact has no parts"
+
+    has_text = any(
+        part.root.kind == "text" and hasattr(part.root, "text") and part.root.text
+        for part in artifact.parts
+    )
+    assert has_text, "No text content found in artifact"
+{%- else %}
     """Test the chat stream functionality."""
     logger.info("Starting chat stream test")
 {% if cookiecutter.is_adk %}
@@ -417,6 +514,94 @@ def test_chat_stream(server_fixture: subprocess.Popen[str]) -> None:
             break
     assert has_content, "At least one message should have content"
 {%- endif %}
+{%- endif %}
+
+
+{%- if cookiecutter.is_adk_a2a %}
+
+
+def test_chat_non_streaming(server_fixture: subprocess.Popen[str]) -> None:
+    """Test the non-streaming chat functionality using A2A JSON-RPC protocol."""
+    logger.info("Starting non-streaming chat test")
+
+    message = Message(
+        message_id=f"msg-user-{uuid.uuid4()}",
+        role=Role.user,
+        parts=[Part(root=TextPart(text="What's the weather in San Francisco?"))],
+    )
+
+    request = SendMessageRequest(
+        id="test-req-002",
+        params=MessageSendParams(message=message),
+    )
+
+    response = requests.post(
+        A2A_RPC_URL,
+        headers=HEADERS,
+        json=request.model_dump(mode="json", exclude_none=True),
+        timeout=60,
+    )
+    assert response.status_code == 200
+
+    # Parse the single JSON-RPC response
+    response_data = response.json()
+    message_response = SendMessageResponse.model_validate(response_data)
+    logger.info(f"Received response: {message_response}")
+
+    # For non-streaming, the result is a Task object
+    json_rpc_resp = message_response.root
+    assert hasattr(json_rpc_resp, "result")
+    task = json_rpc_resp.result
+    assert task.kind == "task"
+    assert hasattr(task, "status")
+    assert task.status.state == "completed"
+
+    # Check that we got artifacts (the final agent output)
+    assert hasattr(task, "artifacts")
+    assert task.artifacts, "No artifacts in task"
+
+    # Verify we got text content in the artifact
+    artifact = task.artifacts[0]
+    assert artifact.parts, "Artifact has no parts"
+
+    has_text = any(
+        part.root.kind == "text" and hasattr(part.root, "text") and part.root.text
+        for part in artifact.parts
+    )
+    assert has_text, "No text content found in artifact"
+
+
+def test_chat_stream_error_handling(server_fixture: subprocess.Popen[str]) -> None:
+    """Test the chat stream error handling with invalid A2A request."""
+    logger.info("Starting chat stream error handling test")
+
+    invalid_data = {
+        "jsonrpc": "2.0",
+        "id": "test-error-001",
+        "method": "message/send",
+        "params": {
+            "message": {
+                "role": "user",
+                # Missing required 'parts' field
+                "messageId": f"msg-user-{uuid.uuid4()}",
+            }
+        },
+    }
+
+    response = requests.post(
+        A2A_RPC_URL, headers=HEADERS, json=invalid_data, timeout=10
+    )
+    assert response.status_code == 200
+
+    response_data = response.json()
+    error_response = JSONRPCErrorResponse.model_validate(response_data)
+    assert "error" in response_data, "Expected JSON-RPC error in response"
+
+    # Assert error for invalid parameters
+    assert error_response.error.code == -32602
+
+    logger.info("Error handling test completed successfully")
+{%- else %}
 
 
 def test_chat_stream_error_handling(server_fixture: subprocess.Popen[str]) -> None:
@@ -433,6 +618,7 @@ def test_chat_stream_error_handling(server_fixture: subprocess.Popen[str]) -> No
         f"Expected status code 422, got {response.status_code}"
     )
     logger.info("Error handling test completed successfully")
+{%- endif %}
 
 
 def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
@@ -455,6 +641,37 @@ def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
         FEEDBACK_URL, json=feedback_data, headers=HEADERS, timeout=10
     )
     assert response.status_code == 200
+
+
+{%- if cookiecutter.is_adk_a2a %}
+
+
+def test_a2a_agent_json_generation(server_fixture: subprocess.Popen[str]) -> None:
+    """
+    Test that the agent.json file is automatically generated and served correctly
+    via the well-known URI.
+    """
+    # Verify the A2A endpoint serves the agent card
+    response = requests.get(AGENT_CARD_URL, timeout=10)
+    assert response.status_code == 200, f"A2A endpoint returned {response.status_code}"
+
+    # Validate required fields in served agent card
+    served_agent_card = response.json()
+    required_fields = [
+        "name",
+        "description",
+        "skills",
+        "capabilities",
+        "url",
+        "version",
+    ]
+    for field in required_fields:
+        assert field in served_agent_card, (
+            f"Missing required field in served agent card: {field}"
+        )
+
+
+{%- endif %}
 {%- if cookiecutter.session_type == "agent_engine" %}
 
 

agent_starter_pack/deployment_targets/cloud_run/tests/load_test/README.md

@@ -99,7 +99,7 @@ uv run uvicorn {{cookiecutter.agent_directory}}.server:app --host 0.0.0.0 --port
 Using another terminal tab, This is suggested to avoid conflicts with the existing application python environment.
 
 ```bash
-python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1{%- if cookiecutter.is_adk_a2a %} a2a-sdk~=0.3.9{%- endif %}
 ```
 
 **3. Execute the Load Test:**
@@ -150,7 +150,7 @@ export _ID_TOKEN=$(gcloud auth print-identity-token -q)
 **3. Execute the Load Test:**
 Create virtual environment with Locust:
 ```bash
-python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1
+python3 -m venv .locust_env && source .locust_env/bin/activate && pip install locust==2.31.1{%- if cookiecutter.is_adk_a2a %} a2a-sdk~=0.3.9{%- endif %}
 ```
 
 Execute load tests. The following command executes the same load test parameters as the local test but targets your remote Cloud Run instance.

agent_starter_pack/deployment_targets/cloud_run/tests/load_test/load_test.py

@@ -140,7 +140,19 @@ class RemoteAgentUser(WebSocketUser):
 
 import os
 import time
-{%- if cookiecutter.is_adk %}
+{%- if cookiecutter.is_adk_a2a %}
+import uuid
+
+from a2a.types import (
+    Message,
+    MessageSendParams,
+    Part,
+    Role,
+    SendStreamingMessageRequest,
+    TextPart,
+)
+from locust import HttpUser, between, task
+{%- elif cookiecutter.is_adk %}
 import uuid
 
 import requests
@@ -149,11 +161,17 @@ from locust import HttpUser, between, task
 
 from locust import HttpUser, between, task
 {%- endif %}
-{% if cookiecutter.is_adk %}
+{%- if cookiecutter.is_adk_a2a %}
+
+ENDPOINT = "/a2a/{{cookiecutter.agent_directory}}/"
+{%- elif cookiecutter.is_adk %}
+
 ENDPOINT = "/run_sse"
-{% else %}
+{%- else %}
+
 ENDPOINT = "/stream_messages"
-{% endif %}
+{%- endif %}
+
 
 class ChatStreamUser(HttpUser):
     """Simulates a user interacting with the chat stream API."""
@@ -162,6 +180,34 @@ class ChatStreamUser(HttpUser):
 
     @task
     def chat_stream(self) -> None:
+{%- if cookiecutter.is_adk_a2a %}
+        """Simulates a chat stream interaction using A2A protocol."""
+        headers = {"Content-Type": "application/json"}
+        if os.environ.get("_ID_TOKEN"):
+            headers["Authorization"] = f"Bearer {os.environ['_ID_TOKEN']}"
+
+        message = Message(
+            message_id=f"msg-user-{uuid.uuid4()}",
+            role=Role.user,
+            parts=[Part(root=TextPart(text="Hello! What's the weather in New York?"))],
+        )
+
+        request = SendStreamingMessageRequest(
+            id=f"req-{uuid.uuid4()}",
+            params=MessageSendParams(message=message),
+        )
+
+        start_time = time.time()
+
+        with self.client.post(
+            ENDPOINT,
+            name=f"{ENDPOINT} message",
+            headers=headers,
+            json=request.model_dump(mode="json", exclude_none=True),
+            catch_response=True,
+            stream=True,
+        ) as response:
+{%- else %}
         """Simulates a chat stream interaction."""
         headers = {"Content-Type": "application/json"}
         if os.environ.get("_ID_TOKEN"):
@@ -218,6 +264,7 @@ class ChatStreamUser(HttpUser):
             stream=True,
             params={"alt": "sse"},
         ) as response:
+{%- endif %}
             if response.status_code == 200:
                 events = []
                 for line in response.iter_lines():

agent_starter_pack/deployment_targets/cloud_run/{{cookiecutter.agent_directory}}/server.py

@@ -291,10 +291,28 @@ async def serve_frontend_spa(full_path: str) -> FileResponse:
     )
 {% elif cookiecutter.is_adk %}
 import os
+{%- if cookiecutter.is_adk_a2a %}
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+{%- endif %}
 
 import google.auth
+{%- if cookiecutter.is_adk_a2a %}
+from a2a.server.apps import A2AFastAPIApplication
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.tasks import InMemoryTaskStore
+from a2a.types import AgentCapabilities, AgentCard
+{%- endif %}
 from fastapi import FastAPI
+{%- if cookiecutter.is_adk_a2a %}
+from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor
+from google.adk.a2a.utils.agent_card_builder import AgentCardBuilder
+from google.adk.artifacts.gcs_artifact_service import GcsArtifactService
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+{%- else %}
 from google.adk.cli.fast_api import get_fast_api_app
+{%- endif %}
 from google.cloud import logging as google_cloud_logging
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider, export
@@ -302,6 +320,9 @@ from opentelemetry.sdk.trace import TracerProvider, export
 from vertexai import agent_engines
 {%- endif %}
 
+{% if cookiecutter.is_adk_a2a -%}
+from {{cookiecutter.agent_directory}}.agent import app as adk_app
+{% endif -%}
 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
@@ -309,9 +330,11 @@ from {{cookiecutter.agent_directory}}.utils.typing import Feedback
 _, project_id = google.auth.default()
 logging_client = google_cloud_logging.Client()
 logger = logging_client.logger(__name__)
+{%- if not cookiecutter.is_adk_a2a %}
 allow_origins = (
     os.getenv("ALLOW_ORIGINS", "").split(",") if os.getenv("ALLOW_ORIGINS") else None
 )
+{%- endif %}
 
 bucket_name = f"gs://{project_id}-{{cookiecutter.project_name}}-logs"
 create_bucket_if_not_exists(
@@ -323,6 +346,48 @@ processor = export.BatchSpanProcessor(CloudTraceLoggingSpanExporter())
 provider.add_span_processor(processor)
 trace.set_tracer_provider(provider)
 
+{%- if cookiecutter.is_adk_a2a %}
+
+runner = Runner(
+    app=adk_app,
+    artifact_service=GcsArtifactService(bucket_name=bucket_name),
+    session_service=InMemorySessionService(),
+)
+
+request_handler = DefaultRequestHandler(
+    agent_executor=A2aAgentExecutor(runner=runner), task_store=InMemoryTaskStore()
+)
+
+
+async def build_dynamic_agent_card() -> AgentCard:
+    """Builds the Agent Card dynamically from the root_agent."""
+    agent_card_builder = AgentCardBuilder(
+        agent=adk_app.root_agent,
+        capabilities=AgentCapabilities(streaming=True),
+        rpc_url=f"{os.getenv('APP_URL', 'http://0.0.0.0:8000')}/a2a/{adk_app.name}/",
+        agent_version=os.getenv("AGENT_VERSION", "0.1.0"),
+    )
+    agent_card = await agent_card_builder.build()
+    return agent_card
+
+
+@asynccontextmanager
+async def lifespan(app_instance: FastAPI) -> AsyncIterator[None]:
+    agent_card = await build_dynamic_agent_card()
+    a2a_app = A2AFastAPIApplication(
+        agent_card=agent_card, http_handler=request_handler
+    ).build()
+    app_instance.mount(f"/a2a/{adk_app.name}", a2a_app)
+    yield
+
+
+app = FastAPI(
+    title="{{cookiecutter.project_name}}",
+    description="API for interacting with the Agent {{cookiecutter.project_name}}",
+    lifespan=lifespan,
+)
+{%- else %}
+
 AGENT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
 {%- if cookiecutter.session_type == "alloydb" %}
@@ -366,6 +431,7 @@ app: FastAPI = get_fast_api_app(
 )
 app.title = "{{cookiecutter.project_name}}"
 app.description = "API for interacting with the Agent {{cookiecutter.project_name}}"
+{%- endif %}
 {% else %}
 import logging
 import os