🐛 fix: make the planner_with_memory → simulator live flow run end-to-end

.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

.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

agents/planner/agent.py

@@ -41,15 +41,18 @@
 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]
 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(

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,
     )
 

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"

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 = {

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"
 

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."""

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)

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],
     )
 

agents/planner_with_memory/memory/embeddings.py

@@ -45,15 +45,15 @@ 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
     ``planner_with_memory`` schemas. Errors propagate; callers decide
     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,

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,