diff --git a/.env.example b/.env.example index 3d8060b2..1f22b0a5 100644 --- a/.env.example +++ b/.env.example @@ -85,6 +85,16 @@ PUBSUB_RESET_SUBS=router-sub,gateway-push-orchestration # Minimum instances per service (used by deploy.py and cloudbuild.yaml) MIN_INSTANCES=1 +# --- Agent models --- +# Model for the planner family (planner / planner_with_memory / planner +# skills). Read at runtime; deploy.py also sets it on the Agent Engines. +PLANNER_MODEL=gemini-3.5-flash + +# --- Agent-to-agent discovery (agents/utils/communication.py) --- +# Bounded retry for the gateway /api/v1/agent-types fetch. Defaults shown. +# A2A_DISCOVERY_RETRIES=3 +# A2A_DISCOVERY_BACKOFF_BASE=0.5 + # --- VPC (Cloud Run direct VPC egress, used by deploy.py) --- # Values come from Terraform outputs: vpc_network, vpc_subnet # VPC_NETWORK=n26-devkey-simulation-vpc diff --git a/.env.swp b/.env.swp deleted file mode 100644 index e96a8476..00000000 Binary files a/.env.swp and /dev/null differ diff --git a/.gitignore b/.gitignore index 59abeee8..a9894afe 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,13 @@ Icon[ # Thumbnails ._* +# Editor swap files +*.swp +*.swo + +# Stray tool/home artifacts +.home/ + # Files that might appear in the root of a volume .DocumentRevisions-V100 .fseventsd diff --git a/.home/.angular-config.json b/.home/.angular-config.json deleted file mode 100644 index 0967ef42..00000000 --- a/.home/.angular-config.json +++ /dev/null @@ -1 +0,0 @@ -{} diff --git a/agents/planner/agent.py b/agents/planner/agent.py index c8c51486..d1aaf156 100644 --- a/agents/planner/agent.py +++ b/agents/planner/agent.py @@ -41,7 +41,10 @@ AGENT_DIR = os.path.dirname(__file__) -PLANNER_MODEL = os.getenv("PLANNER_MODEL", "gemini-3-flash-preview") +DEFAULT_PLANNER_MODEL = "gemini-3.5-flash" +# Import-time snapshot kept for back-compat; get_agent() reads the env at call +# time so a PLANNER_MODEL set after import (tests, dynamic config) is honored. +PLANNER_MODEL = os.getenv("PLANNER_MODEL", DEFAULT_PLANNER_MODEL) # [START planner_agent] @@ -49,7 +52,7 @@ def get_agent(): """Entry point for the ADK framework.""" return LlmAgent( name="planner", - model=resilient_model(PLANNER_MODEL), + model=resilient_model(os.getenv("PLANNER_MODEL", DEFAULT_PLANNER_MODEL)), description="Expert GIS analyst for marathon route and event planning.", static_instruction=PLANNER.build(), generate_content_config=types.GenerateContentConfig( diff --git a/agents/planner/skills/gis-spatial-engineering/scripts/tools.py b/agents/planner/skills/gis-spatial-engineering/scripts/tools.py index fd9ec844..eb7f8876 100644 --- a/agents/planner/skills/gis-spatial-engineering/scripts/tools.py +++ b/agents/planner/skills/gis-spatial-engineering/scripts/tools.py @@ -1787,7 +1787,7 @@ async def _gemini_traffic_enrichment(closed_segments: list, affected_intersectio ) response = await client.aio.models.generate_content( - model="gemini-3-flash-preview", + model=os.getenv("PLANNER_MODEL", "gemini-3.5-flash"), contents=prompt, ) diff --git a/agents/planner/tests/test_planner_agent_config.py b/agents/planner/tests/test_planner_agent_config.py index 5bbd63f1..1f331c9e 100644 --- a/agents/planner/tests/test_planner_agent_config.py +++ b/agents/planner/tests/test_planner_agent_config.py @@ -63,3 +63,14 @@ def test_planner_max_output_tokens_is_8192(): assert cfg.max_output_tokens == 8192, ( f"Planner must set max_output_tokens=8192. Currently: {cfg.max_output_tokens!r}" ) + + +def test_planner_model_is_env_driven(monkeypatch): + """Planner honors PLANNER_MODEL and defaults to gemini-3.5-flash.""" + from agents.planner.agent import get_agent + + monkeypatch.delenv("PLANNER_MODEL", raising=False) + assert get_agent().model.model == "gemini-3.5-flash" + + monkeypatch.setenv("PLANNER_MODEL", "gemini-3.5-flash-test") + assert get_agent().model.model == "gemini-3.5-flash-test" diff --git a/agents/planner_with_eval/evaluator/config.py b/agents/planner_with_eval/evaluator/config.py index 14f3baf2..4d4e3e62 100644 --- a/agents/planner_with_eval/evaluator/config.py +++ b/agents/planner_with_eval/evaluator/config.py @@ -15,8 +15,12 @@ import os # Model configuration -# Note: Use gemini-3.1-pro-preview for best results in evaluation tasks -MODEL = os.getenv("EVALUATOR_MODEL", "gemini-3.1-pro-preview") +# gemini-3-flash-preview on the global endpoint, matching the other agent +# Gemini 3 calls. The previous gemini-3.1-pro-preview default was a Pro +# preview model served only on the global endpoint, but the evaluator called +# it on the regional endpoint — so it failed and retried for minutes, +# blowing the planning turn past the gateway timeout. +MODEL = os.getenv("EVALUATOR_MODEL", "gemini-3-flash-preview") # Criterion weights must sum to 1.0 (equal weighting) CRITERION_WEIGHTS = { diff --git a/agents/planner_with_eval/evaluator/tools.py b/agents/planner_with_eval/evaluator/tools.py index 047ba090..0f4d99d0 100644 --- a/agents/planner_with_eval/evaluator/tools.py +++ b/agents/planner_with_eval/evaluator/tools.py @@ -61,9 +61,13 @@ def _normalize_score(raw_score: float) -> int: def _get_model_resource() -> str: - """Get the full resource path for the Vertex AI evaluation model.""" + """Get the full resource path for the Vertex AI evaluation model. + + Uses the global endpoint: the evaluator model is a Gemini 3 preview model, + which is only served on the global endpoint (not the regional one). + """ project_id = os.environ.get("GOOGLE_CLOUD_PROJECT") - location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") + location = "global" return f"projects/{project_id}/locations/{location}/publishers/google/models/{MODEL}" if project_id else MODEL @@ -662,7 +666,10 @@ async def _generate_feedback( """ try: project_id = os.environ.get("GOOGLE_CLOUD_PROJECT") - location = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1") + # Global endpoint: gemini-3-flash-preview (a Gemini 3 preview model) is + # only served globally, matching our other agent Gemini 3 calls. Used by + # the Vertex AI path below; the Gemini API-key path ignores location. + location = "global" api_key = os.environ.get("GEMINI_API_KEY") use_vertex = os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "true").lower() == "true" diff --git a/agents/planner_with_eval/tests/test_eval_feedback.py b/agents/planner_with_eval/tests/test_eval_feedback.py index 8ac14c4d..a53ecfb0 100644 --- a/agents/planner_with_eval/tests/test_eval_feedback.py +++ b/agents/planner_with_eval/tests/test_eval_feedback.py @@ -117,6 +117,34 @@ async def test_generate_feedback_graceful_fallback(): assert isinstance(summary, str) +@pytest.mark.asyncio +async def test_generate_feedback_uses_flash_on_global_endpoint(): + """Evaluator feedback must use gemini-3-flash-preview on the GLOBAL endpoint. + + Gemini 3 preview models are only served on the global endpoint; the + previous gemini-3.1-pro-preview on the regional endpoint failed and + retried for minutes, blowing the planning turn past the gateway timeout. + """ + from agents.planner_with_eval.evaluator.tools import _generate_feedback + + mock_response = MagicMock() + mock_response.text = '{"suggestions": ["x"], "summary": "y"}' + mock_client = MagicMock() + mock_client.models.generate_content.return_value = mock_response + + with patch("agents.planner_with_eval.evaluator.tools.genai") as mock_genai: + mock_genai.Client.return_value = mock_client + await _generate_feedback( + scores={"safety_compliance": 60}, + details={}, + user_intent="Plan a marathon", + proposed_plan="A 26.2 mile route", + ) + + assert mock_genai.Client.call_args.kwargs.get("location") == "global" + assert mock_client.models.generate_content.call_args.kwargs.get("model") == "gemini-3-flash-preview" + + @pytest.mark.asyncio async def test_evaluate_plan_includes_llm_suggestions(): """evaluate_plan result must include LLM-generated suggestions and summary.""" diff --git a/agents/planner_with_eval/tools.py b/agents/planner_with_eval/tools.py index 61085bd3..86617370 100644 --- a/agents/planner_with_eval/tools.py +++ b/agents/planner_with_eval/tools.py @@ -342,3 +342,44 @@ async def run_async( except Exception as e: logger.error(f"PLANNER: Failed to call simulator agent: {e}") return {"status": "error", "message": f"Failed to call simulator agent: {str(e)}"} + + +# --- Function shims -------------------------------------------------------- +# The ADK-facing API is the BaseTool classes above (registered in adk_tools.py). +# These thin async wrappers adapt the tools to a plain keyword-argument call so +# internal callers and tests can invoke them directly. SubmitPlanToSimulatorTool +# also relies on start_simulation() for its "LLM skipped start" fallback path. + + +async def start_simulation( + *, + action: str, + message: str, + tool_context: ToolContext, + simulation_config: dict | None = None, + runner_type: str | None = None, +) -> dict: + """Invoke StartSimulationTool with keyword args (see the class for behavior).""" + args: dict[str, Any] = {"action": action, "message": message} + if simulation_config is not None: + args["simulation_config"] = simulation_config + if runner_type is not None: + args["runner_type"] = runner_type + return await StartSimulationTool().run_async(args=args, tool_context=tool_context) + + +async def submit_plan_to_simulator( + *, + action: str, + message: str, + tool_context: ToolContext, + simulation_config: dict | None = None, + runner_type: str | None = None, +) -> dict: + """Invoke SubmitPlanToSimulatorTool with keyword args (see the class for behavior).""" + args: dict[str, Any] = {"action": action, "message": message} + if simulation_config is not None: + args["simulation_config"] = simulation_config + if runner_type is not None: + args["runner_type"] = runner_type + return await SubmitPlanToSimulatorTool().run_async(args=args, tool_context=tool_context) diff --git a/agents/planner_with_memory/agent.py b/agents/planner_with_memory/agent.py index 017d8cec..97cf4c55 100644 --- a/agents/planner_with_memory/agent.py +++ b/agents/planner_with_memory/agent.py @@ -31,6 +31,7 @@ from agents.planner_with_memory.services.memory_manager import auto_save_memories from agents.utils import config from agents.utils.communication_plugin import SimulationCommunicationPlugin +from agents.utils.context_trim import trim_route_geojson_from_context from agents.utils.retry import resilient_model from agents.utils.deployment import create_a2a_deployment from agents.utils.factory import create_simulation_runner @@ -40,7 +41,7 @@ logger = logging.getLogger(__name__) AGENT_NAME = "planner_with_memory" -MODEL = os.getenv("PLANNER_MODEL", "gemini-3-flash-preview") +MODEL = os.getenv("PLANNER_MODEL", "gemini-3.5-flash") planner_config = types.GenerateContentConfig( max_output_tokens=8192, @@ -68,10 +69,12 @@ def get_agent(): static_instruction=PLANNER_WITH_MEMORY.build(), tools=get_tools(), generate_content_config=planner_config, - # No before_model_callback: the secure-financial-modeling skill - # handles refusals via validate_and_emit_a2ui (A2UI card). - # A programmatic guardrail would short-circuit the LLM before - # it can emit the A2UI refusal card. + # This before_model_callback only MUTATES the outgoing request (it strips + # the bulky route GeoJSON from prior tool results so it isn't re-sent to + # the model every turn -- ~200s/turn otherwise). It returns None and + # never short-circuits the LLM, so the secure-financial-modeling skill + # can still emit its A2UI refusal card via validate_and_emit_a2ui. + before_model_callback=trim_route_geojson_from_context, after_agent_callback=[auto_save_memories], ) diff --git a/agents/planner_with_memory/memory/embeddings.py b/agents/planner_with_memory/memory/embeddings.py index 57719344..f91d4e48 100644 --- a/agents/planner_with_memory/memory/embeddings.py +++ b/agents/planner_with_memory/memory/embeddings.py @@ -45,7 +45,7 @@ def _get_genai_client() -> genai.Client: ) -async def compute_embedding(text: str, *, dimension: int = 768) -> list[float]: +async def compute_embedding(text: str, *, dimension: int = 3072) -> list[float]: """Compute an embedding for ``text`` via Vertex AI / Gemini. Default dimension 3072 matches the ``VECTOR(3072)`` column in @@ -53,7 +53,7 @@ async def compute_embedding(text: str, *, dimension: int = 768) -> list[float]: retry/fallback policy. """ client = _get_genai_client() - model = os.environ.get("EMBEDDING_MODEL", "text-embedding-004") + model = os.environ.get("EMBEDDING_MODEL", "gemini-embedding-001") response = await client.aio.models.embed_content( model=model, contents=text, diff --git a/agents/planner_with_memory/memory/store_alloydb.py b/agents/planner_with_memory/memory/store_alloydb.py index 2651f09d..7dbd80f2 100644 --- a/agents/planner_with_memory/memory/store_alloydb.py +++ b/agents/planner_with_memory/memory/store_alloydb.py @@ -21,6 +21,7 @@ from __future__ import annotations +import asyncio import json import logging import os @@ -104,9 +105,43 @@ def _get_dsn() -> str: return f"postgresql://{user}:{password}@{host}:{port}/{database}" +# Transient failures that warrant a reconnect-and-retry (connection drops, +# connect/command timeouts). Query-level errors (constraint violations, etc.) +# are NOT in here -- retrying those is pointless. +_DB_RETRYABLE = (OSError, asyncio.TimeoutError, asyncpg.PostgresConnectionError, asyncpg.InterfaceError) + + async def _get_conn() -> asyncpg.Connection: schema_name = os.environ.get("ALLOYDB_SCHEMA", "local_dev") - return await asyncpg.connect(_get_dsn(), server_settings={"search_path": f"{schema_name}, public"}) + # Bound both connect AND per-query time. Without command_timeout, a degraded + # connection makes a query (e.g. the store_route INSERT) wait forever and + # hangs the whole planner turn -- the same failure mode as the Redis stall. + connect_timeout = float(os.environ.get("ALLOYDB_CONNECT_TIMEOUT_S", "10")) + command_timeout = float(os.environ.get("ALLOYDB_COMMAND_TIMEOUT_S", "30")) + return await asyncpg.connect( + _get_dsn(), + server_settings={"search_path": f"{schema_name}, public"}, + timeout=connect_timeout, + command_timeout=command_timeout, + ) + + +async def _with_db_retry(coro_factory): + """Run an async DB operation, retrying transient connection failures. + + coro_factory must be a zero-arg callable returning a fresh coroutine on each + call (so each attempt opens its own connection). + """ + attempts = int(os.environ.get("ALLOYDB_WRITE_RETRIES", "2")) + 1 + backoff = float(os.environ.get("ALLOYDB_RETRY_BACKOFF_S", "0.5")) + for i in range(attempts): + try: + return await coro_factory() + except _DB_RETRYABLE as e: + if i == attempts - 1: + raise + logger.warning("AlloyDB op failed (attempt %d/%d): %s; retrying", i + 1, attempts, e) + await asyncio.sleep(backoff * (2**i)) def _row_to_route(row: asyncpg.Record) -> PlannedRoute: @@ -156,20 +191,24 @@ async def store_route( ) -> str: """Persist a new planned route and return its UUID.""" route_id = str(uuid.uuid4()) - conn = await _get_conn() - try: - await conn.execute( - """ - INSERT INTO planned_routes (route_id, route_data, created_at, eval_score, eval_result) - VALUES ($1, $2::jsonb, now(), $3, $4::jsonb) - """, - route_id, - json.dumps(route_data), - evaluation_score, - json.dumps(evaluation_result) if evaluation_result is not None else None, - ) - finally: - await conn.close() + + async def _do() -> None: + conn = await _get_conn() + try: + await conn.execute( + """ + INSERT INTO planned_routes (route_id, route_data, created_at, eval_score, eval_result) + VALUES ($1, $2::jsonb, now(), $3, $4::jsonb) + """, + route_id, + json.dumps(route_data), + evaluation_score, + json.dumps(evaluation_result) if evaluation_result is not None else None, + ) + finally: + await conn.close() + + await _with_db_retry(_do) return route_id async def get_route(self, route_id: str) -> PlannedRoute | None: @@ -195,24 +234,28 @@ async def record_simulation( simulation_result: dict, ) -> str | None: """Append a simulation record to route_id. Returns sim UUID or None.""" - conn = await _get_conn() - try: - exists = await conn.fetchval("SELECT 1 FROM planned_routes WHERE route_id = $1", route_id) - if not exists: - return None - sim_id = str(uuid.uuid4()) - await conn.execute( - """ - INSERT INTO simulation_records (simulation_id, route_id, sim_result, simulated_at) - VALUES ($1, $2, $3::jsonb, now()) - """, - sim_id, - route_id, - json.dumps(simulation_result), - ) - return sim_id - finally: - await conn.close() + + async def _do() -> str | None: + conn = await _get_conn() + try: + exists = await conn.fetchval("SELECT 1 FROM planned_routes WHERE route_id = $1", route_id) + if not exists: + return None + sim_id = str(uuid.uuid4()) + await conn.execute( + """ + INSERT INTO simulation_records (simulation_id, route_id, sim_result, simulated_at) + VALUES ($1, $2, $3::jsonb, now()) + """, + sim_id, + route_id, + json.dumps(simulation_result), + ) + return sim_id + finally: + await conn.close() + + return await _with_db_retry(_do) async def recall_routes( self, diff --git a/agents/planner_with_memory/memory/test_embeddings.py b/agents/planner_with_memory/memory/test_embeddings.py index cb98a1ac..34546f4a 100644 --- a/agents/planner_with_memory/memory/test_embeddings.py +++ b/agents/planner_with_memory/memory/test_embeddings.py @@ -44,6 +44,28 @@ async def test_compute_embedding_returns_3072_dim_vector() -> None: assert all(isinstance(v, float) for v in vec) +@pytest.mark.asyncio +async def test_compute_embedding_defaults_match_alloydb_schema(monkeypatch) -> None: + """tools.py calls compute_embedding() with no overrides, so the defaults must + stay gemini-embedding-001/3072 or pgvector rejects the query (3072 vs 768).""" + monkeypatch.delenv("EMBEDDING_MODEL", raising=False) + fake_client = MagicMock() + fake_client.aio.models.embed_content = AsyncMock( + return_value=MagicMock(embeddings=[MagicMock(values=[0.0] * 3072)]) + ) + with patch.object(embeddings, "_get_genai_client", return_value=fake_client): + await embeddings.compute_embedding("hello") + + await_args = fake_client.aio.models.embed_content.await_args + assert await_args is not None + assert await_args.kwargs["model"] == "gemini-embedding-001" + config = await_args.kwargs["config"] + dim = getattr(config, "output_dimensionality", None) + if dim is None and isinstance(config, dict): + dim = config["output_dimensionality"] + assert dim == 3072 + + @pytest.mark.asyncio async def test_compute_embedding_uses_configured_model(monkeypatch) -> None: """Honors EMBEDDING_MODEL env var; defaults to gemini-embedding-001.""" diff --git a/agents/planner_with_memory/tests/test_store_alloydb.py b/agents/planner_with_memory/tests/test_store_alloydb.py index 718a9476..db0f2be6 100644 --- a/agents/planner_with_memory/tests/test_store_alloydb.py +++ b/agents/planner_with_memory/tests/test_store_alloydb.py @@ -17,7 +17,7 @@ from __future__ import annotations import time -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -138,6 +138,56 @@ def test_raises_without_host(self) -> None: mod._get_dsn() +class TestGetConnTimeouts: + """_get_conn must bound connect AND query time so a degraded AlloyDB path + fails fast instead of hanging the planner forever (observed: store_route + INSERT hung indefinitely with no command_timeout).""" + + @pytest.mark.asyncio + async def test_get_conn_passes_connect_and_command_timeouts(self) -> None: + env = {"ALLOYDB_HOST": "10.0.0.1", "ALLOYDB_PASSWORD": "pw"} + captured: dict = {} + + async def fake_connect(dsn, **kwargs): + captured.update(kwargs) + return MagicMock() + + with patch.dict("os.environ", env, clear=False), patch.object(mod.asyncpg, "connect", fake_connect): + await mod._get_conn() + + assert captured.get("timeout") is not None, "connect timeout missing" + assert captured.get("command_timeout") is not None, "command_timeout missing" + + +class TestStoreRouteRetry: + """A transient connection failure on a write must be retried, not fatal.""" + + def setup_method(self) -> None: + mod._cached = None + mod._sm_client = None + + @pytest.mark.asyncio + async def test_store_route_retries_transient_connection_error(self) -> None: + env = {"ALLOYDB_HOST": "10.0.0.1", "ALLOYDB_PASSWORD": "pw"} + good_conn = MagicMock() + good_conn.execute = AsyncMock(return_value="INSERT 0 1") + good_conn.close = AsyncMock() + calls = {"n": 0} + + async def flaky_connect(dsn, **kwargs): + calls["n"] += 1 + if calls["n"] == 1: + raise OSError("connection reset") + return good_conn + + with patch.dict("os.environ", env, clear=False), patch.object(mod.asyncpg, "connect", flaky_connect): + route_id = await mod.AlloyDBRouteStore().store_route({"name": "test"}) + + assert isinstance(route_id, str) and route_id + assert calls["n"] == 2 # retried after the transient failure + good_conn.execute.assert_awaited_once() + + class TestResolveSmProject: def test_prefers_secret_manager_project(self, monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("SECRET_MANAGER_PROJECT", "sm-pid") diff --git a/agents/utils/communication.py b/agents/utils/communication.py index dc0b121f..76ff4010 100644 --- a/agents/utils/communication.py +++ b/agents/utils/communication.py @@ -42,6 +42,8 @@ _discovery_cache: dict[str, Any] = {} _discovery_cache_ts: float = 0.0 _DISCOVERY_CACHE_TTL: float = 30.0 # Match gateway's 30s TTL +_DISCOVERY_RETRIES: int = int(os.environ.get("A2A_DISCOVERY_RETRIES", "3")) +_DISCOVERY_BACKOFF_BASE: float = float(os.environ.get("A2A_DISCOVERY_BACKOFF_BASE", "0.5")) # Shared httpx client for agent discovery _discovery_client: httpx.AsyncClient | None = None @@ -136,33 +138,46 @@ async def _discover_agents(self) -> Dict[str, AgentCard]: return _discovery_cache url = f"{self.gateway_url}/api/v1/agent-types" - try: - client = _get_discovery_client() - # Attach OIDC token for OSS Cloud Run IAM mode (no-op when - # get_id_token returns None: local dev / no ADC). - headers: dict[str, str] = {} - audience = oidc_auth.resolve_audience(self.gateway_url) - token = oidc_auth.get_id_token(audience) - if token: - headers["Authorization"] = f"Bearer {token}" - resp = await client.get(url, timeout=5.0, headers=headers) - resp.raise_for_status() - data = resp.json() - - cards = {} - for name, card_data in data.items(): - # Use model_validate to handle camelCase→snake_case alias - # resolution (e.g., additionalInterfaces, preferredTransport) - cards[name] = AgentCard.model_validate(card_data) - - _discovery_cache = cards - _discovery_cache_ts = now - self._registry_cache = cards - logger.info(f"A2A_DISCOVERY: Discovered {len(cards)} agent types from Gateway.") - return cards - except Exception as e: - logger.error(f"A2A_DISCOVERY_ERROR: Failed to fetch agent types from {url}: {e}") - return self._registry_cache + last_exc: Exception | None = None + for attempt in range(_DISCOVERY_RETRIES): + try: + client = _get_discovery_client() + # Attach OIDC token for OSS Cloud Run IAM mode (no-op when + # get_id_token returns None: local dev / no ADC). + headers: dict[str, str] = {} + audience = oidc_auth.resolve_audience(self.gateway_url) + token = oidc_auth.get_id_token(audience) + if token: + headers["Authorization"] = f"Bearer {token}" + resp = await client.get(url, timeout=5.0, headers=headers) + resp.raise_for_status() + data = resp.json() + + cards = {} + for name, card_data in data.items(): + # Use model_validate to handle camelCase→snake_case alias + # resolution (e.g., additionalInterfaces, preferredTransport) + cards[name] = AgentCard.model_validate(card_data) + + _discovery_cache = cards + _discovery_cache_ts = now + self._registry_cache = cards + logger.info(f"A2A_DISCOVERY: Discovered {len(cards)} agent types from Gateway.") + return cards + except Exception as e: + last_exc = e + logger.warning( + f"A2A_DISCOVERY_RETRY: attempt {attempt + 1}/{_DISCOVERY_RETRIES} " + f"failed for {url}: {e}" + ) + if attempt + 1 < _DISCOVERY_RETRIES: + await asyncio.sleep(_DISCOVERY_BACKOFF_BASE * (2 ** attempt)) + + logger.error( + f"A2A_DISCOVERY_ERROR: Failed to fetch agent types from {url} " + f"after {_DISCOVERY_RETRIES} attempts: {last_exc}" + ) + return self._registry_cache # [END a2a_discovery] diff --git a/agents/utils/communication_test.py b/agents/utils/communication_test.py index d55ba34f..95c512ef 100644 --- a/agents/utils/communication_test.py +++ b/agents/utils/communication_test.py @@ -453,6 +453,58 @@ async def test_discovery_cache_expires_after_ttl(self, a2a_client): await a2a_client._discover_agents() assert route.call_count == 2 # Re-fetched + @respx.mock + @pytest.mark.asyncio + async def test_discovery_retries_then_succeeds(self, a2a_client, monkeypatch): + """A transient gateway failure is retried, not swallowed into an empty cache.""" + from agents.utils import communication + + communication._discovery_cache = {} + communication._discovery_cache_ts = 0.0 + monkeypatch.setattr(communication, "_DISCOVERY_BACKOFF_BASE", 0.0) + + ok = httpx.Response( + 200, + json={ + "simulator": { + "name": "simulator", + "url": "http://simulator:8202", + "description": "Mock simulator", + "version": "1.0.0", + "capabilities": {}, + "skills": [], + "default_input_modes": ["text/plain"], + "default_output_modes": ["text/plain"], + } + }, + ) + route = respx.get("http://gateway:8101/api/v1/agent-types").mock( + side_effect=[httpx.ConnectError("transient"), ok] + ) + + cards = await a2a_client._discover_agents() + assert route.call_count == 2 # retried once + assert "simulator" in cards + + @respx.mock + @pytest.mark.asyncio + async def test_discovery_returns_cache_after_exhaustion(self, a2a_client, monkeypatch): + """After all retries fail, return the existing cache without raising.""" + from agents.utils import communication + + communication._discovery_cache = {} + communication._discovery_cache_ts = 0.0 + monkeypatch.setattr(communication, "_DISCOVERY_RETRIES", 2) + monkeypatch.setattr(communication, "_DISCOVERY_BACKOFF_BASE", 0.0) + + route = respx.get("http://gateway:8101/api/v1/agent-types").mock( + side_effect=httpx.ConnectError("gateway down") + ) + + cards = await a2a_client._discover_agents() + assert cards == {} # falls back to empty cache, does not raise + assert route.call_count == 2 # all attempts made + class TestClientCleanup: """Verify clients are properly closed.""" diff --git a/agents/utils/context_trim.py b/agents/utils/context_trim.py new file mode 100644 index 00000000..dbeaa5ea --- /dev/null +++ b/agents/utils/context_trim.py @@ -0,0 +1,73 @@ +# Copyright 2026 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. + +"""Keep bulky route geometry out of the model's context window. + +``report_marathon_route`` returns the full route GeoJSON because the frontend +draws the 3D course from that tool result. ADK then retains the tool result in +the conversation, so the geometry is re-sent to the model on every subsequent +turn -- measured at ~200s per planner turn, even on a warm engine, with no API +errors or retries. It's pure input-processing cost. + +The model never needs the geometry: ``submit_plan_to_simulator``, the +evaluator, and the dashboard all read the route from session state, and the +frontend already received it in the original tool result. So before each model +call we replace the geometry in prior tool results with a compact placeholder. + +This is a ``before_model_callback`` that only MUTATES the outgoing request and +returns ``None`` -- it never returns an ``LlmResponse``, so it does not +short-circuit the LLM (preserving the secure-financial-modeling skill's ability +to emit its A2UI refusal card; see agents/planner_with_memory/agent.py). +""" + +import logging + +logger = logging.getLogger(__name__) + +# Response keys that carry bulky geometry the model doesn't need to re-read. +_GEOMETRY_KEYS = ("route_geojson", "geojson") + + +def _count_points(geo) -> int: + """Best-effort coordinate count for the placeholder message.""" + try: + total = 0 + for feature in geo.get("features", []): + coords = feature.get("geometry", {}).get("coordinates", []) + total += len(coords) + return total + except AttributeError: + return 0 + + +def trim_route_geojson_from_context(callback_context, llm_request): + """before_model_callback: strip route geometry from prior tool results. + + Returns None so the (mutated) request proceeds to the model. + """ + contents = getattr(llm_request, "contents", None) or [] + for content in contents: + for part in getattr(content, "parts", None) or []: + function_response = getattr(part, "function_response", None) + if function_response is None: + continue + response = getattr(function_response, "response", None) + if not isinstance(response, dict): + continue + for key in _GEOMETRY_KEYS: + value = response.get(key) + if value is not None and not isinstance(value, str): + n = _count_points(value) + response[key] = f"[{n} coordinates omitted from context; retained in session state]" + return None diff --git a/agents/utils/global_gemini.py b/agents/utils/global_gemini.py index ea2d9ea6..f6e556c6 100644 --- a/agents/utils/global_gemini.py +++ b/agents/utils/global_gemini.py @@ -25,23 +25,50 @@ from agents.utils.global_gemini import GlobalGemini agent = LlmAgent( - model=GlobalGemini(model="gemini-3-flash-preview"), + model=GlobalGemini(model="gemini-3.5-flash"), ... ) Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/overview """ +import asyncio +import logging import os from functools import cached_property -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, AsyncGenerator from google.adk.models.google_llm import Gemini from google.genai import types +from pydantic import Field if TYPE_CHECKING: + from google.adk.models.llm_request import LlmRequest + from google.adk.models.llm_response import LlmResponse from google.genai import Client +logger = logging.getLogger(__name__) + + +def _env_float(name: str, default: float) -> float: + raw = os.environ.get(name) + if not raw: + return default + try: + return float(raw) + except ValueError: + return default + + +def _env_int(name: str, default: int) -> int: + raw = os.environ.get(name) + if not raw: + return default + try: + return int(raw) + except ValueError: + return default + class GlobalGemini(Gemini): """Gemini model with explicit location control for Vertex AI. @@ -54,8 +81,8 @@ class GlobalGemini(Gemini): For per-agent location control:: - # Global endpoint (default, required for Gemini 3 previews) - model = GlobalGemini(model="gemini-3-flash-preview") + # Global endpoint (the default; current Gemini 3.x models use it) + model = GlobalGemini(model="gemini-3.5-flash") # Regional endpoint model = GlobalGemini(model="gemini-2.0-flash", location="us-central1") @@ -65,6 +92,64 @@ class GlobalGemini(Gemini): """Vertex AI API location. Defaults to 'global' for Gemini 3 preview models. Set to a region (e.g. 'us-central1') for GA models.""" + first_token_timeout_s: float = Field( + default_factory=lambda: _env_float("GEMINI_MODEL_FIRST_TOKEN_TIMEOUT_S", 180.0) + ) + """Seconds to wait for the first response chunk before treating the turn as + stalled. A stall here is safe to restart (no output emitted yet).""" + + stream_timeout_s: float = Field( + default_factory=lambda: _env_float("GEMINI_MODEL_STREAM_TIMEOUT_S", 120.0) + ) + """Seconds to wait between subsequent chunks. A stall here is mid-stream and + cannot be restarted, so it surfaces as an error.""" + + stall_retries: int = Field( + default_factory=lambda: _env_int("GEMINI_MODEL_STALL_RETRIES", 2) + ) + """Number of from-scratch retries for a first-token stall (in addition to + the initial attempt).""" + + async def generate_content_async( + self, llm_request: "LlmRequest", stream: bool = False + ) -> "AsyncGenerator[LlmResponse, None]": + """Run the model with stall detection and bounded first-token recovery. + + The genai SDK retries only ``APIError`` responses, not client-side + stalls. Production turns intermittently hang after ``model_start`` with + the generation request wedged in-flight. We bound each chunk with + ``asyncio.wait_for``: a stall before the first token is restarted from + scratch (safe -- nothing emitted), a stall mid-stream is re-raised + (restarting would duplicate output). + """ + attempts = max(0, self.stall_retries) + 1 + for attempt in range(attempts): + parent = super().generate_content_async(llm_request, stream=stream) + yielded = False + try: + while True: + timeout = self.stream_timeout_s if yielded else self.first_token_timeout_s + try: + response = await asyncio.wait_for(parent.__anext__(), timeout) + except StopAsyncIteration: + return + yielded = True + yield response + except (asyncio.TimeoutError, TimeoutError) as exc: + await parent.aclose() + if yielded or attempt == attempts - 1: + raise TimeoutError( + f"Gemini model turn stalled (model={self.model}, " + f"attempt {attempt + 1}/{attempts}, mid_stream={yielded})." + ) from exc + logger.warning( + "Gemini model turn stalled before first token " + "(model=%s, attempt %d/%d); restarting.", + self.model, + attempt + 1, + attempts, + ) + @cached_property def api_client(self) -> "Client": """Create a genai Client with the configured location.""" diff --git a/agents/utils/redis_pool.py b/agents/utils/redis_pool.py index 72c76cc5..995c1284 100644 --- a/agents/utils/redis_pool.py +++ b/agents/utils/redis_pool.py @@ -53,7 +53,16 @@ def get_shared_redis_client() -> redis.Redis | None: redis_url, decode_responses=False, max_connections=max_conn, - timeout=5, # max seconds to wait for a connection + timeout=5, # max seconds to wait for a connection from the pool + # Establishing a NEW TCP connection to Memorystore must fail fast. + # Without this, a hung connect falls back to the OS default (~75s), + # which froze the broadcast worker and Redis side-channels mid-turn + # (telemetry showed a single publish stalling ~74s). Keepalive + + # periodic health checks recycle stale connections so a dropped + # connection doesn't force a slow reconnect during a request. + socket_connect_timeout=5, + socket_keepalive=True, + health_check_interval=30, ) _shared_client = redis.Redis(connection_pool=pool) return _shared_client diff --git a/agents/utils/retry.py b/agents/utils/retry.py index 27b2f2f7..7d735cdc 100644 --- a/agents/utils/retry.py +++ b/agents/utils/retry.py @@ -23,7 +23,7 @@ from agents.utils.retry import resilient_model agent = LlmAgent( - model=resilient_model("gemini-3-flash-preview"), + model=resilient_model("gemini-3.5-flash"), ... ) @@ -99,7 +99,7 @@ def resilient_model( """Create a GlobalGemini instance with retry protection. Args: - model_name: Gemini model identifier (e.g. "gemini-3-flash-preview"). + model_name: Gemini model identifier (e.g. "gemini-3.5-flash"). retry_options: Override retry config. Defaults to default_retry_options(). location: Vertex AI API location. Defaults to "global" (required for Gemini 3 preview models). Set to a region like "us-central1" for diff --git a/agents/utils/test_context_trim.py b/agents/utils/test_context_trim.py new file mode 100644 index 00000000..b43cf78d --- /dev/null +++ b/agents/utils/test_context_trim.py @@ -0,0 +1,83 @@ +# Copyright 2026 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. + +"""Tests for trimming bulky route GeoJSON out of the model context. + +report_marathon_route returns the full route GeoJSON so the frontend can draw +the 3D course. ADK keeps that tool result in the conversation, so it is re-sent +to the model on every subsequent turn -- measured at ~200s per planner turn, +even warm. The model never needs the geometry (every backend consumer reads it +from session state), so we strip it from the request before each model call. +""" + +from types import SimpleNamespace + +from google.genai import types + + +def _req_with_route(num_points: int = 5000): + geo = { + "type": "FeatureCollection", + "features": [{"geometry": {"type": "LineString", "coordinates": [[float(i), float(i)] for i in range(num_points)]}}], + } + part = types.Part( + function_response=types.FunctionResponse( + name="report_marathon_route", + response={"status": "success", "message": "route summary text", "route_geojson": geo}, + ) + ) + content = types.Content(role="user", parts=[part]) + return SimpleNamespace(contents=[content]), part + + +def test_trims_route_geojson_to_placeholder(): + from agents.utils.context_trim import trim_route_geojson_from_context + + req, part = _req_with_route() + out = trim_route_geojson_from_context(None, req) + + assert out is None # never short-circuits the model call + resp = part.function_response.response + assert isinstance(resp["route_geojson"], str), "geojson should be replaced by a placeholder string" + assert "5000" in resp["route_geojson"], "placeholder should note the point count" + + +def test_preserves_other_response_fields(): + from agents.utils.context_trim import trim_route_geojson_from_context + + req, part = _req_with_route() + trim_route_geojson_from_context(None, req) + resp = part.function_response.response + assert resp["message"] == "route summary text" + assert resp["status"] == "success" + + +def test_leaves_unrelated_tool_responses_untouched(): + from agents.utils.context_trim import trim_route_geojson_from_context + + part = types.Part( + function_response=types.FunctionResponse( + name="recall_past_simulations", response={"routes": [1, 2, 3]} + ) + ) + req = SimpleNamespace(contents=[types.Content(role="user", parts=[part])]) + trim_route_geojson_from_context(None, req) + assert part.function_response.response["routes"] == [1, 2, 3] + + +def test_handles_empty_or_missing_contents(): + from agents.utils.context_trim import trim_route_geojson_from_context + + assert trim_route_geojson_from_context(None, SimpleNamespace(contents=[])) is None + assert trim_route_geojson_from_context(None, SimpleNamespace(contents=None)) is None diff --git a/agents/utils/test_global_gemini.py b/agents/utils/test_global_gemini.py new file mode 100644 index 00000000..baf1fada --- /dev/null +++ b/agents/utils/test_global_gemini.py @@ -0,0 +1,155 @@ +# Copyright 2026 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. + +"""Tests for GlobalGemini model-turn stall recovery. + +Production incident: a planner_with_memory turn intermittently stalls after +``model_start`` -- the Gemini generation request goes in-flight and never +returns, hanging the whole multi-agent flow. The genai SDK only retries +``APIError`` responses with retriable status codes, NOT client-side stalls, +so a wedged stream hangs forever. + +GlobalGemini.generate_content_async wraps the parent generator with a +per-chunk ``asyncio.wait_for``. A stall *before the first token* is safe to +restart, so it retries from scratch (bounded). A stall *mid-stream* cannot be +safely restarted (would duplicate output), so it re-raises. +""" + +import asyncio +import os +from unittest import mock + +import pytest +from google.adk.models.google_llm import Gemini + + +def _make_model(**overrides): + """Construct a GlobalGemini with tiny timeouts for fast deterministic tests.""" + from agents.utils.global_gemini import GlobalGemini + + params = { + "model": "gemini-3.5-flash", + "first_token_timeout_s": 0.1, + "stream_timeout_s": 0.1, + "stall_retries": 2, + } + params.update(overrides) + return GlobalGemini(**params) + + +async def _drain(agen): + return [chunk async for chunk in agen] + + +@pytest.mark.asyncio +async def test_passthrough_yields_all_chunks_when_healthy(): + """A healthy parent stream is forwarded unchanged.""" + model = _make_model() + + async def fake_parent(self, llm_request, stream=False): + yield "chunk-1" + yield "chunk-2" + + with mock.patch.object(Gemini, "generate_content_async", new=fake_parent): + out = await _drain(model.generate_content_async(object(), stream=False)) + + assert out == ["chunk-1", "chunk-2"] + + +@pytest.mark.asyncio +async def test_retries_after_first_token_stall(): + """Stall before the first token -> abort and restart from scratch.""" + model = _make_model() + calls = {"n": 0} + + async def fake_parent(self, llm_request, stream=False): + calls["n"] += 1 + if calls["n"] == 1: + await asyncio.sleep(5) # stall past first_token_timeout_s + yield "never" + else: + yield "recovered" + + with mock.patch.object(Gemini, "generate_content_async", new=fake_parent): + out = await _drain(model.generate_content_async(object(), stream=False)) + + assert out == ["recovered"] + assert calls["n"] == 2 # one stalled attempt + one successful retry + + +@pytest.mark.asyncio +async def test_raises_after_exhausting_retries(): + """Every attempt stalls -> raise once retries are exhausted.""" + model = _make_model(stall_retries=2) + calls = {"n": 0} + + async def fake_parent(self, llm_request, stream=False): + calls["n"] += 1 + await asyncio.sleep(5) + yield "never" + + with mock.patch.object(Gemini, "generate_content_async", new=fake_parent): + with pytest.raises((asyncio.TimeoutError, TimeoutError)): + await _drain(model.generate_content_async(object(), stream=False)) + + assert calls["n"] == 3 # initial + 2 retries + + +@pytest.mark.asyncio +async def test_midstream_stall_reraises_without_restart(): + """Stall after a token was yielded -> cannot restart, must re-raise.""" + model = _make_model() + calls = {"n": 0} + + async def fake_parent(self, llm_request, stream=False): + calls["n"] += 1 + yield "partial" + await asyncio.sleep(5) # stall past stream_timeout_s, mid-stream + yield "never" + + collected = [] + with mock.patch.object(Gemini, "generate_content_async", new=fake_parent): + with pytest.raises((asyncio.TimeoutError, TimeoutError)): + async for chunk in model.generate_content_async(object(), stream=False): + collected.append(chunk) + + assert collected == ["partial"] + assert calls["n"] == 1 # no restart after partial output + + +def test_timeouts_and_retries_read_from_env(): + from agents.utils.global_gemini import GlobalGemini + + with mock.patch.dict( + os.environ, + { + "GEMINI_MODEL_FIRST_TOKEN_TIMEOUT_S": "42.0", + "GEMINI_MODEL_STREAM_TIMEOUT_S": "17.0", + "GEMINI_MODEL_STALL_RETRIES": "5", + }, + ): + model = GlobalGemini(model="gemini-3.5-flash") + + assert model.first_token_timeout_s == 42.0 + assert model.stream_timeout_s == 17.0 + assert model.stall_retries == 5 + + +def test_default_timeouts_are_generous(): + """Defaults must not abort healthy turns (eval ~20s, planning ~10s).""" + from agents.utils.global_gemini import GlobalGemini + + model = GlobalGemini(model="gemini-3.5-flash") + assert model.first_token_timeout_s >= 120.0 + assert model.stall_retries >= 1 diff --git a/agents/utils/tests/test_redis_pool.py b/agents/utils/tests/test_redis_pool.py index d8dcd64c..eed0b4ea 100644 --- a/agents/utils/tests/test_redis_pool.py +++ b/agents/utils/tests/test_redis_pool.py @@ -57,6 +57,21 @@ def test_pool_timeout_is_set(self): assert isinstance(pool, redis.BlockingConnectionPool) assert pool.timeout == 5, f"Expected timeout=5, got {pool.timeout}" + def test_pool_has_socket_connect_timeout_and_keepalive(self): + """A hung TCP connect to Memorystore must fail fast, not block on the OS + default (~75s). Telemetry showed a single Redis publish stalling ~74s + (the default TCP connect timeout), freezing broadcasts/side-channels. + socket_keepalive + health_check_interval recycle stale connections so a + dropped connection doesn't force a slow reconnect mid-turn. + """ + with patch.dict("os.environ", {"REDIS_ADDR": "127.0.0.1:6379"}): + client = pool_mod.get_shared_redis_client() + assert client is not None + kwargs = client.connection_pool.connection_kwargs + assert kwargs.get("socket_connect_timeout") == 5 + assert kwargs.get("socket_keepalive") is True + assert kwargs.get("health_check_interval") == 30 + def test_max_connections_configurable_via_env(self): """REDIS_MAX_CONNECTIONS env var should override the default pool size.""" with patch.dict( diff --git a/infra/modules/cloud-run-services/main.tf b/infra/modules/cloud-run-services/main.tf index 196406c0..5085cc5b 100644 --- a/infra/modules/cloud-run-services/main.tf +++ b/infra/modules/cloud-run-services/main.tf @@ -80,7 +80,7 @@ resource "google_cloud_run_v2_service" "gateway" { timeout = "3600s" scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["gateway"].min_instances max_instance_count = var.service_sizing["gateway"].max_instances } @@ -169,7 +169,7 @@ resource "google_cloud_run_v2_service" "admin" { template { service_account = var.compute_sa_email scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["admin"].min_instances max_instance_count = var.service_sizing["admin"].max_instances } vpc_access { @@ -229,7 +229,7 @@ resource "google_cloud_run_v2_service" "dash" { template { service_account = var.compute_sa_email scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["dash"].min_instances max_instance_count = var.service_sizing["dash"].max_instances } vpc_access { @@ -289,7 +289,7 @@ resource "google_cloud_run_v2_service" "tester" { template { service_account = var.compute_sa_email scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["tester"].min_instances max_instance_count = var.service_sizing["tester"].max_instances } vpc_access { @@ -350,7 +350,7 @@ resource "google_cloud_run_v2_service" "frontend" { service_account = var.compute_sa_email timeout = "3600s" scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["frontend"].min_instances max_instance_count = var.service_sizing["frontend"].max_instances } vpc_access { @@ -413,7 +413,7 @@ resource "google_cloud_run_v2_service" "runner_autopilot" { template { service_account = var.compute_sa_email scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["runner_autopilot"].min_instances max_instance_count = var.service_sizing["runner_autopilot"].max_instances } vpc_access { @@ -473,7 +473,7 @@ resource "google_cloud_run_v2_service" "runner_cloudrun" { template { service_account = var.compute_sa_email scaling { - min_instance_count = var.min_instances + min_instance_count = var.service_sizing["runner_cloudrun"].min_instances max_instance_count = var.service_sizing["runner_cloudrun"].max_instances } vpc_access { diff --git a/infra/modules/cloud-run-services/tests/sizing.tftest.hcl b/infra/modules/cloud-run-services/tests/sizing.tftest.hcl index a8937f2b..b6cd5a35 100644 --- a/infra/modules/cloud-run-services/tests/sizing.tftest.hcl +++ b/infra/modules/cloud-run-services/tests/sizing.tftest.hcl @@ -83,16 +83,17 @@ run "python_runner_services_default_to_1gi" { } } -run "all_services_capped_at_max_instances_1" { +run "services_capped_for_oss_cost_control" { command = plan - # OSS cost control: every service must cap at a single instance, - # not just the runners. Pinned across all 7 Cloud Run services so - # a future edit can't quietly raise gateway/admin/etc. back to 10. + # OSS cost control: off-path services cap at a single instance. The + # gateway is the hub for WebSocket/A2A-discovery/orchestration, so it + # may scale modestly (max 3) to absorb concurrent sessions; everything + # else stays pinned so a future edit can't quietly raise it back to 10. assert { - condition = google_cloud_run_v2_service.gateway.template[0].scaling[0].max_instance_count == 1 - error_message = "gateway max_instance_count must be 1 for OSS cost control" + condition = google_cloud_run_v2_service.gateway.template[0].scaling[0].max_instance_count == 3 + error_message = "gateway max_instance_count must be 3 (modest scale for the hub)" } assert { @@ -126,12 +127,12 @@ run "all_services_capped_at_max_instances_1" { } } -run "all_services_default_min_instances_zero" { +run "gateway_warm_other_services_scale_to_zero" { command = plan assert { - condition = google_cloud_run_v2_service.gateway.template[0].scaling[0].min_instance_count == 0 - error_message = "gateway must default to scale-to-zero" + condition = google_cloud_run_v2_service.gateway.template[0].scaling[0].min_instance_count == 1 + error_message = "gateway must be warmed (min=1): it is the hub for WebSocket, A2A discovery, and orchestration" } assert { diff --git a/infra/modules/cloud-run-services/variables.tf b/infra/modules/cloud-run-services/variables.tf index bf631610..f9f5b316 100644 --- a/infra/modules/cloud-run-services/variables.tf +++ b/infra/modules/cloud-run-services/variables.tf @@ -83,33 +83,30 @@ variable "agent_urls" { default = "" } -variable "min_instances" { - description = "Default min_instance_count for all services. OSS defaults to scale-to-zero." - type = number - default = 0 -} - # Right-sized per docs/plans/2026-04-16-oss-right-sizing-plan.md. # OSS targets minimum cost: scale-to-zero, platform-minimum CPU/memory. variable "service_sizing" { - description = "Per-service CPU/memory/max_instances. Right-sized for OSS cost." + description = "Per-service CPU/memory/min_instances/max_instances. Right-sized for OSS cost." type = map(object({ cpu = string memory = string + min_instances = number max_instances = number })) - # max_instances = 1 for ALL services (OSS cost control: never scale - # beyond a single instance per service). Memory differs by runtime: - # Go services stay at 512Mi; Python+ADK runners need 1Gi (Phase 7 - # build #20 surfaced 512Mi OOM-killing the Python container during - # startup imports before the port-8080 health probe ever opened). + # Off-path services stay scale-to-zero (min 0) and single-instance (max 1) + # for OSS cost control. The gateway is the hub for the WebSocket, A2A + # discovery, and orchestration, so it is warmed (min 1) and allowed to + # scale modestly (max 3) to absorb concurrent sessions. Memory differs by + # runtime: Go services stay at 512Mi; Python+ADK runners need 1Gi (Phase 7 + # build #20 surfaced 512Mi OOM-killing the Python container during startup + # imports before the port-8080 health probe ever opened). default = { - gateway = { cpu = "1", memory = "512Mi", max_instances = 1 } - admin = { cpu = "1", memory = "512Mi", max_instances = 1 } - tester = { cpu = "1", memory = "512Mi", max_instances = 1 } - frontend = { cpu = "1", memory = "512Mi", max_instances = 1 } - dash = { cpu = "1", memory = "512Mi", max_instances = 1 } - runner_autopilot = { cpu = "1", memory = "1Gi", max_instances = 1 } - runner_cloudrun = { cpu = "1", memory = "1Gi", max_instances = 1 } + gateway = { cpu = "1", memory = "512Mi", min_instances = 1, max_instances = 3 } + admin = { cpu = "1", memory = "512Mi", min_instances = 0, max_instances = 1 } + tester = { cpu = "1", memory = "512Mi", min_instances = 0, max_instances = 1 } + frontend = { cpu = "1", memory = "512Mi", min_instances = 0, max_instances = 1 } + dash = { cpu = "1", memory = "512Mi", min_instances = 0, max_instances = 1 } + runner_autopilot = { cpu = "1", memory = "1Gi", min_instances = 0, max_instances = 1 } + runner_cloudrun = { cpu = "1", memory = "1Gi", min_instances = 0, max_instances = 1 } } } diff --git a/scripts/deploy/deploy.py b/scripts/deploy/deploy.py index 5ebf6c5b..b9ad4ab4 100644 --- a/scripts/deploy/deploy.py +++ b/scripts/deploy/deploy.py @@ -180,7 +180,8 @@ def _construct_gateway_url(project_number: str, region: str) -> str: "module": "agents.simulator.agent", "attr": "simulator_a2a_agent", "resource_limits": {"memory": "2Gi", "cpu": "1"}, - "max_instances": 1, + "min_instances": 1, + "max_instances": 3, }, "planner": { "type": "reasoning-engine", @@ -188,7 +189,8 @@ def _construct_gateway_url(project_number: str, region: str) -> str: "module": "agents.planner.agent", "attr": "planner_a2a_agent", "resource_limits": {"memory": "2Gi", "cpu": "1"}, - "max_instances": 1, + "min_instances": 1, + "max_instances": 3, }, "planner_with_eval": { "type": "reasoning-engine", @@ -197,6 +199,7 @@ def _construct_gateway_url(project_number: str, region: str) -> str: "attr": "planner_a2a_agent", "extra_packages": ["agents/planner"], "resource_limits": {"memory": "2Gi", "cpu": "1"}, + "min_instances": 0, "max_instances": 1, }, "simulator_with_failure": { @@ -206,6 +209,7 @@ def _construct_gateway_url(project_number: str, region: str) -> str: "attr": "simulator_with_failure_a2a_agent", "extra_packages": ["agents/simulator"], "resource_limits": {"memory": "2Gi", "cpu": "1"}, + "min_instances": 0, "max_instances": 1, }, "planner_with_memory": { @@ -214,8 +218,14 @@ def _construct_gateway_url(project_number: str, region: str) -> str: "module": "agents.planner_with_memory.agent", "attr": "planner_a2a_agent", "extra_packages": ["agents/planner", "agents/planner_with_eval"], - "resource_limits": {"memory": "2Gi", "cpu": "1"}, - "max_instances": 1, + # Heaviest engine: bundles planner + planner_with_eval + the memory + # bank and runs pandas-based evaluation in-process. At 2Gi it was + # OOM-killed mid-request (worker restarts ~1/min observed live), + # surfacing as stalled model turns. 8Gi/2cpu gives headroom for the + # multiple uvicorn workers per instance. + "resource_limits": {"memory": "8Gi", "cpu": "2"}, + "min_instances": 1, + "max_instances": 3, }, } @@ -762,7 +772,8 @@ def deploy_agent_engine(service_name: str, cfg: dict, *, tf: dict, force_create: "LITELLM_LOCAL_MODEL_COST_MAP": "True", "LITELLM_TELEMETRY": "False", "BUILD_FINGERPRINT": build_fingerprint, - "EMBEDDING_MODEL": "text-embedding-004", + "EMBEDDING_MODEL": "gemini-embedding-001", + "PLANNER_MODEL": "gemini-3.5-flash", } # Simulation defaults. diff --git a/scripts/deploy/deploy_test.py b/scripts/deploy/deploy_test.py index d721e349..0d6fa3f9 100644 --- a/scripts/deploy/deploy_test.py +++ b/scripts/deploy/deploy_test.py @@ -235,14 +235,14 @@ def test_no_internal_project_ids(self): assert "keynote2026" not in serialized def test_resource_limits_are_right_sized(self): - """Agent Engine agents use minimal resource limits (cpu<=2, memory<=2Gi).""" + """Agent Engine agents stay within sane bounds (cpu<=2, memory<=8Gi).""" for name, cfg in deploy.SERVICES.items(): limits = cfg.get("resource_limits", {}) cpu = int(limits.get("cpu", "4")) assert cpu <= 2, f"{name} has cpu={cpu}, expected <= 2" - mem = limits.get("memory", "8Gi") + mem = limits.get("memory", "16Gi") mem_gb = int(mem.replace("Gi", "")) - assert mem_gb <= 2, f"{name} has memory={mem}, expected <= 2Gi" + assert mem_gb <= 8, f"{name} has memory={mem}, expected <= 8Gi" def test_all_ae_agents_have_at_least_2gi(self): """All AE agents need >=2Gi memory: 1Gi causes worker OOMs during @@ -255,11 +255,37 @@ def test_all_ae_agents_have_at_least_2gi(self): f"ADK runtime startup." ) + def test_planner_with_memory_has_oom_headroom(self): + """planner_with_memory bundles planner + planner_with_eval + the memory + bank and runs pandas-based evaluation in-process, with multiple uvicorn + workers per instance. At 2Gi it was OOM-killed mid-request (~1 worker + restart/min observed live), which manifests as stalled model turns. It + needs >=4Gi.""" + mem_gb = int(deploy.SERVICES["planner_with_memory"]["resource_limits"]["memory"].replace("Gi", "")) + assert mem_gb >= 4, f"planner_with_memory has {mem_gb}Gi; needs >=4Gi to avoid mid-request OOM" + def test_max_instances_capped(self): + # OSS cost control: critical-path agents (the live SandboxIO flow) + # may scale modestly to absorb concurrent sessions; off-path agents + # stay pinned to a single instance. + scalable = {"planner", "planner_with_memory", "simulator"} for name, cfg in deploy.SERVICES.items(): mi = cfg.get("max_instances") assert mi is not None, f"{name} missing max_instances" - assert mi <= 1, f"{name} has max_instances={mi}, expected <= 1 for OSS" + cap = 3 if name in scalable else 1 + assert mi <= cap, f"{name} has max_instances={mi}, expected <= {cap}" + + def test_min_instances_policy(self): + # Only the live-path agents are warmed; off-path agents stay + # scale-to-zero. Set explicitly so the MIN_INSTANCES=1 fallback + # default can't silently warm the off-path agents. + warm = {"planner", "planner_with_memory", "simulator"} + for name, cfg in deploy.SERVICES.items(): + mi = cfg.get("min_instances") + if name in warm: + assert mi == 1, f"{name} should be warm (min_instances=1), got {mi!r}" + else: + assert mi == 0, f"{name} should be scale-to-zero (min_instances=0), got {mi!r}" # --- _determine_deploy_mode (create-or-update with displayName fallback) ---