agent-starter-pack 0.18.2__py3-none-any.whl → 0.21.0__py3-none-any.whl

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

@@ -47,7 +47,7 @@ def start_server() -> subprocess.Popen[str]:
         sys.executable,
         "-m",
         "uvicorn",
-        "{{cookiecutter.agent_directory}}.server:app",
+        "{{cookiecutter.agent_directory}}.fast_api_app:app",
         "--host",
         "0.0.0.0",
         "--port",
@@ -202,8 +202,8 @@ def test_feedback_endpoint(server_fixture: subprocess.Popen[str]) -> None:
     feedback_data = {
         "score": 5,
         "text": "Great response!",
-        "invocation_id": "test-run-123",
-        "user_id": "test-user",
+        "user_id": "test-user-123",
+        "session_id": "test-session-123",
         "log_type": "feedback",
     }
 
@@ -220,12 +220,28 @@ import subprocess
 import sys
 import threading
 import time
+{%- if cookiecutter.is_a2a %}
 import uuid
+{%- endif %}
 from collections.abc import Iterator
 from typing import Any
 
 import pytest
 import requests
+{%- if cookiecutter.is_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 +249,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_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"
@@ -255,7 +274,7 @@ def start_server() -> subprocess.Popen[str]:
         sys.executable,
         "-m",
         "uvicorn",
-        "{{cookiecutter.agent_directory}}.server:app",
+        "{{cookiecutter.agent_directory}}.fast_api_app:app",
         "--host",
         "0.0.0.0",
         "--port",
@@ -264,8 +283,8 @@ def start_server() -> subprocess.Popen[str]:
     env = os.environ.copy()
     env["INTEGRATION_TEST"] = "TRUE"
 {%- if cookiecutter.session_type == "agent_engine" %}
-    # Set test-specific agent engine session name
-    env["AGENT_ENGINE_SESSION_NAME"] = "test-{{cookiecutter.project_name}}"
+    # Use in-memory session for local E2E tests instead of creating Agent Engine
+    env["USE_IN_MEMORY_SESSION"] = "true"
 {%- endif %}
     process = subprocess.Popen(
         command,
@@ -292,7 +311,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_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 +346,82 @@ def server_fixture(request: Any) -> Iterator[subprocess.Popen[str]]:
 
 
 def test_chat_stream(server_fixture: subprocess.Popen[str]) -> None:
+{%- if cookiecutter.is_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 +516,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_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 +620,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:
@@ -443,11 +631,8 @@ def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
     # Create sample feedback data
     feedback_data = {
         "score": 4,
-{%- if cookiecutter.is_adk %}
-        "invocation_id": str(uuid.uuid4()),
-{%- else %}
-        "run_id": str(uuid.uuid4()),
-{%- endif %}
+        "user_id": "test-user-456",
+        "session_id": "test-session-456",
         "text": "Great response!",
     }
 
@@ -455,6 +640,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_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

@@ -12,7 +12,7 @@ Follow these steps to execute load tests on your local machine:
 Launch the FastAPI server in a separate terminal:
 
 ```bash
-uv run uvicorn {{cookiecutter.agent_directory}}.server:app --host 0.0.0.0 --port 8000 --reload
+uv run uvicorn {{cookiecutter.agent_directory}}.fast_api_app:app --host 0.0.0.0 --port 8000 --reload
 ```
 
 **2. (In another tab) Create virtual environment with Locust**
@@ -92,14 +92,14 @@ Follow these steps to execute load tests on your local machine:
 Launch the FastAPI server in a separate terminal:
 
 ```bash
-uv run uvicorn {{cookiecutter.agent_directory}}.server:app --host 0.0.0.0 --port 8000 --reload
+uv run uvicorn {{cookiecutter.agent_directory}}.fast_api_app:app --host 0.0.0.0 --port 8000 --reload
 ```
 
 **2. (In another tab) Create virtual environment with Locust**
 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_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_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

@@ -138,9 +138,23 @@ class RemoteAgentUser(WebSocketUser):
     host = "http://localhost:8000"  # Default for local testing
 {%- else %}
 
+import json
+import logging
 import os
 import time
-{%- if cookiecutter.is_adk %}
+{%- if cookiecutter.is_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 +163,23 @@ from locust import HttpUser, between, task
 
 from locust import HttpUser, between, task
 {%- endif %}
-{% if cookiecutter.is_adk %}
+{%- if cookiecutter.is_a2a %}
+
+ENDPOINT = "/a2a/{{cookiecutter.agent_directory}}"
+{%- elif cookiecutter.is_adk %}
+
 ENDPOINT = "/run_sse"
-{% else %}
+{%- else %}
+
 ENDPOINT = "/stream_messages"
-{% endif %}
+{%- endif %}
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
 
 class ChatStreamUser(HttpUser):
     """Simulates a user interacting with the chat stream API."""
@@ -162,6 +188,34 @@ class ChatStreamUser(HttpUser):
 
     @task
     def chat_stream(self) -> None:
+{%- if cookiecutter.is_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,8 +272,10 @@ class ChatStreamUser(HttpUser):
             stream=True,
             params={"alt": "sse"},
         ) as response:
+{%- endif %}
             if response.status_code == 200:
                 events = []
+                has_error = False
                 for line in response.iter_lines():
                     if line:
                         line_str = line.decode("utf-8")
@@ -234,16 +290,40 @@ class ChatStreamUser(HttpUser):
                                 response=response,
                                 context={},
                             )
+
+                        # Check for error responses in the JSON payload
+                        try:
+                            event_data = json.loads(line_str)
+                            if isinstance(event_data, dict) and "code" in event_data:
+                                # Flag any non-2xx codes as errors
+                                if event_data["code"] >= 400:
+                                    has_error = True
+                                    error_msg = event_data.get(
+                                        "message", "Unknown error"
+                                    )
+                                    response.failure(f"Error in response: {error_msg}")
+                                    logger.error(
+                                        "Received error response: code=%s, message=%s",
+                                        event_data["code"],
+                                        error_msg,
+                                    )
+                        except json.JSONDecodeError:
+                            # If it's not valid JSON, continue processing
+                            pass
+
                 end_time = time.time()
                 total_time = end_time - start_time
-                self.environment.events.request.fire(
-                    request_type="POST",
-                    name=f"{ENDPOINT} end",
-                    response_time=total_time * 1000,  # Convert to milliseconds
-                    response_length=len(events),
-                    response=response,
-                    context={},
-                )
+
+                # Only fire success event if no errors were found
+                if not has_error:
+                    self.environment.events.request.fire(
+                        request_type="POST",
+                        name=f"{ENDPOINT} end",
+                        response_time=total_time * 1000,  # Convert to milliseconds
+                        response_length=len(events),
+                        response=response,
+                        context={},
+                    )
             else:
                 response.failure(f"Unexpected status code: {response.status_code}")
 {%- endif %}