feat: add Agent engine session in Cloud Run (#275)

Author: eliasecchig

.cloudbuild/terraform/build_triggers.tf

@@ -98,6 +98,10 @@ locals {
       name  = "adk_base-cloud_run-alloydb"
       value = "adk_base,cloud_run,--session-type,alloydb"
     },
+    {
+      name  = "adk_base-cloud_run-agent_engine"
+      value = "adk_base,cloud_run,--session-type,agent_engine"
+    },
   ]
 
 agent_testing_included_files = { for combo in local.agent_testing_combinations :

agents/adk_base/.template/templateconfig.yaml

@@ -18,7 +18,7 @@ settings:
   requires_data_ingestion: false
   requires_session: true
   deployment_targets: ["agent_engine", "cloud_run"]
-  extra_dependencies: ["google-adk~=1.6.0"]
+  extra_dependencies: ["google-adk~=1.8.0"]
   tags: ["adk"]
   frontend_type: "None"
 

agents/adk_gemini_fullstack/.template/templateconfig.yaml

@@ -18,7 +18,7 @@ settings:
   requires_data_ingestion: false
   requires_session: true
   deployment_targets: ["agent_engine", "cloud_run"]
-  extra_dependencies: ["google-adk~=1.6.0"]
+  extra_dependencies: ["google-adk~=1.8.0"]
   tags: ["adk"]
   frontend_type: "adk_gemini_fullstack"
   commands:

agents/agentic_rag/.template/templateconfig.yaml

@@ -19,7 +19,7 @@ settings:
   requires_session: true
   deployment_targets: ["agent_engine", "cloud_run"]
   extra_dependencies: [
-    "google-adk~=1.6.0",
+    "google-adk~=1.8.0",
     "langchain-google-vertexai~=2.0.7",
     "langchain~=0.3.24",
     "langchain-core~=0.3.55",

docs/cli/create.md

@@ -22,7 +22,7 @@ The following options will be prompted interactively if not provided via the com
 - `--deployment-target`, `-d`: Deployment target (`agent_engine` or `cloud_run`). Prompts if omitted.
 - `--datastore`, `-ds`: Datastore for RAG agents (`vertex_ai_search` or `vertex_ai_vector_search`). Prompted if `--include-data-ingestion` is specified, or if the selected agent (e.g., `agentic_rag`) requires data ingestion, and this option is omitted.
 - `--region`: GCP region for deployment (default: `us-central1`). Prompts for confirmation if not specified and `--auto-approve` is not used.
-- `--session-type`: Type of session storage to use (`in_memory` or `alloydb`). This is only applicable for agents that require session management and when the deployment target is `cloud_run`.
+- `--session-type`: Type of session storage to use (`in_memory`, `alloydb`, or `agent_engine`). This is only applicable for agents that require session management and when the deployment target is `cloud_run`.
 
 GCP account and project ID are detected automatically (using your active `gcloud config` settings). You will be prompted to confirm or change them unless `--auto-approve` is used.
 

llm.txt

@@ -69,7 +69,7 @@ agent-starter-pack create PROJECT_NAME [OPTIONS]
 *   `-d, --deployment-target`: Target deployment environment (`cloud_run` or `agent_engine`).
 *   `-ds, --datastore`: For RAG agents, the datastore (`vertex_ai_search` or `vertex_ai_vector_search`).
 *   `-i, --include-data-ingestion`: Include data ingestion pipeline scaffolding.
-*   `--session-type`: For agents requiring session management on Cloud Run, specifies the storage type (`in_memory` or `alloydb`).
+*   `--session-type`: For agents requiring session management on Cloud Run, specifies the storage type (`in_memory`, `alloydb`, `agent_engine`).
 *   `--region`: GCP region (e.g., `us-central1`).
 *   `--auto-approve`: **Skips all interactive prompts (crucial for automation).**
 *   `--debug`: Enables debug logging during creation.

src/cli/commands/create.py

@@ -116,7 +116,7 @@ def normalize_project_name(project_name: str) -> str:
 )
 @click.option(
     "--session-type",
-    type=click.Choice(["in_memory", "alloydb"]),
+    type=click.Choice(["in_memory", "alloydb", "agent_engine"]),
     help="Type of session storage to use",
 )
 @click.option("--debug", is_flag=True, help="Enable debug logging")

src/cli/utils/template.py

@@ -230,6 +230,10 @@ def prompt_session_type_selection() -> str:
             "display_name": "AlloyDB",
             "description": "Use AlloyDB for session management. Comes with terraform resources for deployment.",
         },
+        "agent_engine": {
+            "display_name": "Vertex AI Agent Engine",
+            "description": "Managed session service that automatically handles conversation history",
+        },
     }
 
     console.print("\n> Please select a session type:")

src/deployment_targets/cloud_run/app/server.py

@@ -20,6 +20,9 @@
 from google.cloud import logging as google_cloud_logging
 from opentelemetry import trace
 from opentelemetry.sdk.trace import TracerProvider, export
+{%- if cookiecutter.session_type == "agent_engine" %}
+from vertexai import agent_engines
+{%- endif %}
 
 from app.utils.gcs import create_bucket_if_not_exists
 from app.utils.tracing import CloudTraceLoggingSpanExporter
@@ -55,6 +58,22 @@
 session_service_uri = None
 if db_host and db_pass:
     session_service_uri = f"postgresql://{db_user}:{db_pass}@{db_host}:5432/{db_name}"
+{%- elif cookiecutter.session_type == "agent_engine" %}
+# Agent Engine session configuration
+# Use environment variable for agent name, default to project name
+agent_name = os.environ.get("AGENT_ENGINE_SESSION_NAME", "{{cookiecutter.project_name}}")
+
+# Check if an agent with this name already exists
+existing_agents = list(agent_engines.list(filter=f"display_name={agent_name}"))
+
+if existing_agents:
+    # Use the existing agent
+    agent_engine = existing_agents[0]
+else:
+    # Create a new agent if none exists
+    agent_engine = agent_engines.create(display_name=agent_name)
+
+session_service_uri = f"agentengine://{agent_engine.resource_name}"
 {%- else %}
 # In-memory session configuration - no persistent storage
 session_service_uri = None

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

@@ -62,6 +62,10 @@ 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}}"
+{%- endif %}
     process = subprocess.Popen(
         command,
         stdout=subprocess.PIPE,
@@ -250,3 +254,34 @@ 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.session_type == "agent_engine" %}
+
+
+@pytest.fixture(scope="session", autouse=True)
+def cleanup_agent_engine_sessions() -> None:
+    """Cleanup agent engine sessions created during tests."""
+    yield  # Run tests first
+
+    # Cleanup after tests complete
+    from vertexai import agent_engines
+
+    try:
+        # Use same environment variable as server, default to project name
+        agent_name = os.environ.get(
+            "AGENT_ENGINE_SESSION_NAME", "{{cookiecutter.project_name}}"
+        )
+
+        # Find and delete agent engines with this name
+        existing_agents = list(agent_engines.list(filter=f"display_name={agent_name}"))
+
+        for agent_engine in existing_agents:
+            try:
+                agent_engines.delete(resource_name=agent_engine.name)
+                logger.info(f"Cleaned up agent engine: {agent_engine.name}")
+            except Exception as e:
+                logger.warning(
+                    f"Failed to cleanup agent engine {agent_engine.name}: {e}"
+                )
+    except Exception as e:
+        logger.warning(f"Failed to cleanup agent engine sessions: {e}")
+{%- endif %}