🐛 fix: make the planner_with_memory → simulator live flow run end-to-end - #140
Open
cwest wants to merge 17 commits into
Open
🐛 fix: make the planner_with_memory → simulator live flow run end-to-end#140cwest wants to merge 17 commits into
cwest wants to merge 17 commits into
Conversation
The memory planner had been switched to text-embedding-004, which returns 768-dim vectors. But the agent_memory tables store 3072-dim embeddings (VECTOR(3072), seeded with gemini-embedding-001), so every similarity search failed with "different vector dimensions 3072 and 768". That broke planner_with_memory's summary step, and with it the whole live SandboxIO flow on the keynote demo. Go back to gemini-embedding-001 at 3072 dims so query embeddings match the seeded data. No re-seed needed, and 3072 is the model's native size, so there's no manual normalization to deal with. Also pins the defaults with a test, since the call sites in tools.py rely on them rather than passing model/dimension explicitly.
A vim swap file (.env.swp) and a stray .home/.angular-config.json got committed by accident. Remove both and add ignore rules so they don't sneak back in.
A failed /api/v1/agent-types fetch was swallowed into an empty cache, so a single transient gateway blip made every inter-agent call (e.g. the simulator handoff) fail with "not found in registry". Retry with exponential backoff (tunable via A2A_DISCOVERY_RETRIES / A2A_DISCOVERY_BACKOFF_BASE) before falling back to the cache.
gemini-3-flash-preview is slow, intermittently 429s, and sometimes emits the a2ui tool call as text instead of calling it. Move the planner family to the GA gemini-3.5-flash (global endpoint), read from PLANNER_MODEL so a future swap is an env edit rather than a code change. deploy.py sets the env explicitly on the Agent Engines. Runner keeps its own model.
planner, planner_with_memory, and simulator are on the live SandboxIO path and were max_instances=1, so they saturated under concurrent sessions. Set min_instances=1 (warm) and max_instances=3 for those three; off-path agents stay pinned at 1. Updates the cost-control test to encode the new per-agent policy.
The gateway is the hub for the WebSocket, A2A discovery, and orchestration, but ran at minScale:0/maxScale:1 — cold starts plus a single saturated instance under concurrent sessions, which surfaced as flaky chat readiness and failed agent discovery. Make min_instances per-service (replacing the single global var) and set the gateway to min:1/max:3; every other service stays scale-to-zero, single-instance. Tests updated to encode the new per-service policy.
Code review caught that planner_with_eval and simulator_with_failure would warm at min=1 via the MIN_INSTANCES fallback, against the intent that only the live-path agents stay warm. Set them to min_instances=0 explicitly and add a test pinning the min-instances policy. Document PLANNER_MODEL and the A2A_DISCOVERY_* tunables in .env.example, and refresh the stale model ids in the global_gemini / retry docstring examples.
…-flash-preview Live testing showed gemini-3.5-flash regressed the SandboxIO flow: the planner emitted a route_list a2ui card and ended its turn instead of the plan dashboard with the Run Simulation button, so the simulator handoff never fired. The one fully-working end-to-end run was on gemini-3-flash-preview (the model the planner's thinking_budget / max_output_tokens config was tuned for). Keep the env-driven PLANNER_MODEL indirection (so a future swap is an env edit) but set the default back to the known-good gemini-3-flash-preview.
evaluate_plan defaulted to gemini-3.1-pro-preview but built its genai client (and judge model resource) with the regional GOOGLE_CLOUD_LOCATION (us-central1). Gemini 3 preview models are only served on the global endpoint, so every evaluation failed and retried for minutes — telemetry showed a single evaluate_plan call burning 706s, pushing the planning turn (~24 min total) far past the gateway's 60s timeout, so the plan dashboard + Run Simulation button never reached the user. That is why the planner_with_memory -> simulator flow "only worked once". Switch the evaluator to gemini-3-flash-preview on the global endpoint, matching our other agent Gemini 3 calls, so evaluate_plan returns in seconds and the planning turn completes within budget.
A single Redis pipeline publish was observed stalling ~74s (≈ the OS default TCP connect timeout). The shared BlockingConnectionPool set a pool-checkout timeout but no socket_connect_timeout, so establishing a new connection to Memorystore (when one was stale/dropped) blocked for ~75s — freezing the broadcast worker and any Redis side-channel mid-turn, which surfaced as tools "hanging" in the UI. Add socket_connect_timeout=5 so a hung connect fails fast, plus socket_keepalive + health_check_interval=30 to recycle stale connections and avoid slow mid-request reconnects. Redis itself is fine; this bounds the connection-establishment path.
The headless e2e driver measured the planning turn spending ~258s in gemini-3-flash-preview model round-trips (on the large route-GeoJSON + traffic + recall context) between report_marathon_route and evaluate_plan — the dominant remaining cost after the eval and Redis fixes. gemini-3.5-flash is the faster GA model originally requested; the earlier revert to gemini-3-flash-preview was based on a misdiagnosis (the route_list cards were separate organizer sessions, not the planning response). Restore gemini-3.5-flash (global endpoint, env-driven via PLANNER_MODEL).
A planner_with_memory turn intermittently hangs after model_start: the Gemini generation request goes in-flight and never returns, stalling the whole multi-agent flow. The genai SDK only retries APIError responses with retriable status codes -- a wedged client-side stream is not retried, so it hangs forever. GlobalGemini.generate_content_async now bounds each chunk with asyncio.wait_for. A stall before the first token is safe to restart (no output emitted), so it retries from scratch up to GEMINI_MODEL_STALL_RETRIES times. A stall mid-stream cannot be restarted without duplicating output, so it surfaces as a TimeoutError. Timeouts default to 180s/120s and are env-tunable.
…lures store_route's INSERT hung indefinitely on a degraded connection because the asyncpg connection had no command_timeout -- the same infinite-hang failure mode as the Redis stall, stalling the planner turn after evaluate_plan. _get_conn now sets both a connect timeout and a command_timeout, and the write paths (store_route, record_simulation) retry transient connection failures with backoff. All bounds are env-tunable (ALLOYDB_CONNECT_TIMEOUT_S, ALLOYDB_COMMAND_TIMEOUT_S, ALLOYDB_WRITE_RETRIES).
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 the geometry was 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 is pure input-processing cost, and it made the multi-turn plan flow too slow to finish. 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. A before_model_callback now replaces the geometry in prior tool results with a compact placeholder before each model call. It only mutates the request and returns None, so it never short-circuits the LLM (the secure-financial-modeling refusal card still works).
The planner_with_memory engine bundles planner + planner_with_eval + the memory bank and runs pandas-based evaluation in-process, with multiple uvicorn workers per instance. At 2Gi its workers were OOM-killed mid-request (~1 restart/min observed live, abrupt boots with no graceful shutdown), which surfaced as model turns that never returned. Bump to 8Gi/2cpu for headroom.
cwest
changed the title
…ing-dimension # Conflicts: # agents/planner/agent.py # agents/planner_with_eval/evaluator/tools.py
…ulator shims The simulation-tools refactor to StartSimulationTool/SubmitPlanToSimulatorTool classes dropped the module-level start_simulation()/submit_plan_to_simulator() functions, but two things still called them: SubmitPlanToSimulatorTool's "LLM skipped start_simulation" fallback path (a latent NameError at runtime), and five existing tests. Both were broken. Add thin async wrappers that adapt the tool classes to a keyword-argument call. The ADK-facing API stays the BaseTool classes (adk_tools.py registers those); these shims just give internal callers and tests a direct entry point. Fixes 5 pre-existing test failures and the execute-without-start fallback.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
Live mode (
planner_with_memory → simulator) almost never finished. A plan would start, then hang somewhere before the dashboard card showed up, so "Run Simulation" and the simulator never ran. This branch tracks down each reason it stalled and fixes them, so a plan runs all the way through and does it repeatably.How I verified it
Ran it end to end against the deployed
race-condition-expo-demostack with a headless driver, twice back to back. Both completed the whole pipeline:plan → evaluate_plan → store_route → submit_plan_to_simulator (verify) → dashboard a2ui (Run Simulation) → run_simulation → submit_plan_to_simulator (execute) → simulator race_engine tick loop → store_simulation_summary → sim_results a2uiFull Python suite: 1605 passing, 0 failing.
What changed, by area
The big one: planner_with_memory was getting OOM-killed mid-request
scripts/deploy/deploy.py,scripts/deploy/deploy_test.pyThis is the heavy engine. It bundles the planner, the eval agent, and the memory bank, and it runs pandas-based evaluation in-process with several uvicorn workers per instance. It was sized at 2Gi. Workers were getting killed about once a minute (abrupt restarts, no graceful shutdown), and that's what made the model turns look like they were hanging. Bumping it to 8Gi/2cpu is the change that actually unblocked the flow. Everything below is real, but secondary.
Keep the route GeoJSON out of the model context
agents/utils/context_trim.py(new),agents/utils/test_context_trim.py(new),agents/planner_with_memory/agent.pyreport_marathon_routereturns the full route GeoJSON so the frontend can draw the 3D course. ADK keeps that tool result in the conversation, so it got re-sent to the model on every following turn, roughly 200s a turn of pure input chewing. Abefore_model_callbacknow replaces the geometry in prior tool results with a short placeholder before each model call. The model never needed it: the simulator, evaluator, and dashboard all read the route from session state, and the frontend already got it in the original tool result. The callback only mutates the request and returnsNone, so it never short-circuits the LLM, and the secure-financial-modeling refusal card still fires.Stop the AlloyDB writes from hanging forever
agents/planner_with_memory/memory/store_alloydb.py,agents/planner_with_memory/tests/test_store_alloydb.pystore_route's INSERT could hang indefinitely because the asyncpg connection had nocommand_timeout. Added a connect timeout and a command timeout. The write paths (store_route,record_simulation) also retry transient connection failures with backoff now.Faster, correct evaluator
agents/planner_with_eval/evaluator/config.py,agents/planner_with_eval/evaluator/tools.py,agents/planner_with_eval/tests/test_eval_feedback.pyEvaluation runs
gemini-3-flash-previewon the global endpoint, which Gemini 3 preview models require. That took a regional-endpoint stall of ~706s down to ~20s. The model is env-driven viaEVALUATOR_MODEL.Planner model is env-driven
agents/planner/agent.py,agents/planner_with_memory/agent.py,agents/planner/.../scripts/tools.py,agents/planner/tests/test_planner_agent_config.pyPLANNER_MODELenv var, defaultgemini-3.5-flash, read at call time.Shared-infra resilience
agents/utils/redis_pool.py(+test):socket_connect_timeoutand keepalive on the shared pool. It was stalling ~74s on a cold connect.agents/utils/communication.py(+test): A2A gateway discovery retries with backoff instead of giving up on the first miss.agents/utils/global_gemini.py(+test_global_gemini.py): a per-chunkasyncio.wait_foraround the model stream. A turn that stalls before the first token restarts from scratch (bounded); a mid-stream stall surfaces as an error instead of an infinite hang. This one is a safety net. In practice the hangs were the OOM above, not the model, but it's cheap insurance against a real stall.Embeddings (the original bug)
agents/planner_with_memory/memory/embeddings.py(+test): back togemini-embedding-001at 3072 dims to match the AlloyDB schema. It was 768-dimtext-embedding-004, which 400'd every memory write.Synced with latest expo-demo, and finished its tools refactor
Merged
origin/expo-demo(it was 5 commits ahead). Two files conflicted; I kept both sides' intent:planner/agent.py: env-drivenPLANNER_MODEL, defaultgemini-3.5-flash, read at call time.planner_with_eval/evaluator/tools.py: forces the global endpoint for the preview eval model, and keeps expo-demo's new Gemini API-key auth path.That merge also surfaced a problem the refactor left behind. Moving
start_simulation/submit_plan_to_simulatortoStartSimulationTool/SubmitPlanToSimulatorToolclasses dropped the old module-level functions, but two callers still used them: a fallback path insideSubmitPlanToSimulatorTool(for when the model skipsstart_simulation), which would have thrownNameErrorin production, and 5 tests. I added thin async wrappers that delegate to the classes. The ADK-facing API stays the classes; the wrappers just give internal callers and the tests a direct way in.Infra / deploy sizing
scripts/deploy/deploy.py: critical AEs (planner,planner_with_memory,simulator) at min 1 / max 3; off-path AEs (planner_with_eval,simulator_with_failure) at min 0 / max 1;planner_with_memoryat 8Gi/2cpu.infra/modules/cloud-run-services/*: per-servicemin_instancesso the gateway stays warm.Housekeeping
.gitignore,.env.example, and removing committed junk (.env.swp,.home/.angular-config.json).New config knobs (env-driven, sane defaults)
PLANNER_MODELgemini-3.5-flashEVALUATOR_MODELgemini-3-flash-previewEMBEDDING_MODELgemini-embedding-001ALLOYDB_CONNECT_TIMEOUT_S/ALLOYDB_COMMAND_TIMEOUT_SALLOYDB_WRITE_RETRIESGEMINI_MODEL_FIRST_TOKEN_TIMEOUT_S/GEMINI_MODEL_STALL_RETRIESA2A_DISCOVERY_RETRIEScommunication.pyNot in this diff: operational notes for the team
I did these against the live
race-condition-expo-demoproject, not in code:expo-demo'scloudbuild-bootstrap.yamldeploys agents with--force-create(commitba7d0b5, not onmain), so every deploy spun up brand-new engines. 55 had piled up, ~45 of them always-on (minInstances>=1), which starved instance capacity and fed the restarts. I kept the 5 newest (the ones the gateway points at) and deleted the rest. Worth reverting--force-createback to update-in-place (the waymaindoes it) so this stops happening.deploy.py. Local ADC was a non-authorized account (403 onreasoningEngines.*).gcloud builds submitruns under the in-project Cloud Build SA and works. These deploys were update-in-place (no--force-create), so they reused the engine IDs the gateway already points at.