Fix load test error detection and add Vertex AI initialization (#493)

* fix: Improve load test error detection and initialize Vertex AI

- Add error detection for HTTP 4xx/5xx responses in SSE streams
- Initialize Vertex AI client with project ID for Agent Engine deployments
- Add logging configuration to track error responses during load tests

## Problem
Load tests were not detecting errors returned within SSE response streams,
causing false positives when the server returned error codes in the JSON
payload. Additionally, Vertex AI client was not properly initialized,
causing potential failures in Agent Engine deployments.

## Solution
- Parse JSON payloads in SSE streams to detect error codes >= 400
- Mark requests as failures and log error details when errors are detected
- Initialize vertexai with project ID and location before creating agent engine
- Add logging configuration to capture error details during load testing

* fix: Correct import order for nest_asyncio in a2a agents

Ruff requires imports to be alphabetically sorted. Move nest_asyncio import
before vertexai import to satisfy linting requirements for adk_a2a agents.

Author: eliasecchig

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

@@ -365,6 +365,7 @@ def chat_stream(self) -> None:
         ) as response:
             if response.status_code == 200:
                 events = []
+                has_error = False
                 for line in response.iter_lines():
                     if line:
                         line_str = line.decode("utf-8")
@@ -379,20 +380,44 @@ def chat_stream(self) -> None:
                                 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",
+
+                # Only fire success event if no errors were found
+                if not has_error:
+                    self.environment.events.request.fire(
+                        request_type="POST",
 {%- if cookiecutter.is_adk %}
-                    name="/streamQuery end",
+                        name="/streamQuery end",
 {%- else %}
-                    name="/stream_messages end",
+                        name="/stream_messages end",
 {%- endif %}
-                    response_time=total_time * 1000,  # Convert to milliseconds
-                    response_length=len(events),
-                    response=response,
-                    context={},
-                )
+                        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 %}

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

@@ -25,6 +25,7 @@
 {%- if cookiecutter.is_adk_a2a %}
 import nest_asyncio
 {%- endif %}
+import vertexai
 {%- if cookiecutter.is_adk_a2a %}
 from a2a.types import AgentCapabilities, AgentCard, TransportProtocol
 from google.adk.a2a.executor.a2a_agent_executor import A2aAgentExecutor
@@ -162,6 +163,7 @@ def clone(self) -> "AgentEngineApp":
 
 
 _, project_id = google.auth.default()
+vertexai.init(project=project_id, location="us-central1")
 artifacts_bucket_name = os.environ.get("ARTIFACTS_BUCKET_NAME")
 {%- if cookiecutter.is_adk_a2a %}
 agent_engine = AgentEngineApp.create(
@@ -193,6 +195,7 @@ def clone(self) -> "AgentEngineApp":
 )
 
 import google.auth
+import vertexai
 from google.cloud import logging as google_cloud_logging
 from langchain_core.runnables import RunnableConfig
 from traceloop.sdk import Instruments, Traceloop
@@ -302,5 +305,6 @@ def register_operations(self) -> dict[str, list[str]]:
 
 
 _, project_id = google.auth.default()
+vertexai.init(project=project_id, location="us-central1")
 agent_engine = AgentEngineApp(project_id=project_id)
 {%- endif %}

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

@@ -138,6 +138,8 @@ class RemoteAgentUser(WebSocketUser):
     host = "http://localhost:8000"  # Default for local testing
 {%- else %}
 
+import json
+import logging
 import os
 import time
 {%- if cookiecutter.is_adk_a2a %}
@@ -172,6 +174,12 @@ class RemoteAgentUser(WebSocketUser):
 ENDPOINT = "/stream_messages"
 {%- 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."""
@@ -267,6 +275,7 @@ def chat_stream(self) -> None:
 {%- endif %}
             if response.status_code == 200:
                 events = []
+                has_error = False
                 for line in response.iter_lines():
                     if line:
                         line_str = line.decode("utf-8")
@@ -281,16 +290,40 @@ def chat_stream(self) -> None:
                                 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 %}