From 0ce93c9f1abd0b9a884677afa1424929692aa040 Mon Sep 17 00:00:00 2001
From: Neerav Makwana <261249544+neeravmakwana@users.noreply.github.com>
Date: Fri, 24 Apr 2026 22:38:32 -0400
Subject: [PATCH] fix: enable incomplete-turn recovery for Gemini

---
 .../run.incomplete-turn.test.ts               | 64 +++++++++++++++++++
 .../pi-embedded-runner/run/incomplete-turn.ts | 35 +++++++++-
 2 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/src/agents/pi-embedded-runner/run.incomplete-turn.test.ts b/src/agents/pi-embedded-runner/run.incomplete-turn.test.ts
index 0ddd2c0d0c1..7a94f134d2b 100644
--- a/src/agents/pi-embedded-runner/run.incomplete-turn.test.ts
+++ b/src/agents/pi-embedded-runner/run.incomplete-turn.test.ts
@@ -735,6 +735,33 @@ describe("runEmbeddedPiAgent incomplete-turn safety", () => {
     expect(retryInstruction).toBe(REASONING_ONLY_RETRY_INSTRUCTION);
   });
 
+  it("detects reasoning-only Gemini turns from signed thinking blocks", () => {
+    const retryInstruction = resolveReasoningOnlyRetryInstruction({
+      provider: "google",
+      modelId: "gemini-2.5-pro",
+      aborted: false,
+      timedOut: false,
+      attempt: makeAttemptResult({
+        assistantTexts: [],
+        lastAssistant: {
+          role: "assistant",
+          stopReason: "end_turn",
+          provider: "google",
+          model: "gemini-2.5-pro",
+          content: [
+            {
+              type: "thinking",
+              thinking: "internal reasoning",
+              thinkingSignature: JSON.stringify({ id: "gemini_rs_helper", type: "reasoning" }),
+            },
+          ],
+        } as unknown as EmbeddedRunAttemptResult["lastAssistant"],
+      }),
+    });
+
+    expect(retryInstruction).toBe(REASONING_ONLY_RETRY_INSTRUCTION);
+  });
+
   it("treats exact NO_REPLY as a deliberate silent assistant reply", () => {
     const incompleteTurnText = resolveIncompleteTurnPayloadText({
       payloadCount: 0,
@@ -916,6 +943,28 @@ describe("runEmbeddedPiAgent incomplete-turn safety", () => {
     expect(DEFAULT_EMPTY_RESPONSE_RETRY_LIMIT).toBe(1);
   });
 
+  it("detects generic empty Gemini turns without visible text", () => {
+    const retryInstruction = resolveEmptyResponseRetryInstruction({
+      provider: "google-vertex",
+      modelId: "google/gemini-3.1-flash",
+      payloadCount: 0,
+      aborted: false,
+      timedOut: false,
+      attempt: makeAttemptResult({
+        assistantTexts: [],
+        lastAssistant: {
+          role: "assistant",
+          stopReason: "end_turn",
+          provider: "google-vertex",
+          model: "gemini-3.1-flash",
+          content: [{ type: "text", text: "" }],
+        } as unknown as EmbeddedRunAttemptResult["lastAssistant"],
+      }),
+    });
+
+    expect(retryInstruction).toBe(EMPTY_RESPONSE_RETRY_INSTRUCTION);
+  });
+
   it("does not retry generic empty GPT turns after side effects", () => {
     const retryInstruction = resolveEmptyResponseRetryInstruction({
       provider: "openai",
@@ -985,6 +1034,21 @@ describe("runEmbeddedPiAgent incomplete-turn safety", () => {
     expect(result.meta.livenessState).toBe("working");
   });
 
+  it("detects replay-safe planning-only Gemini turns", () => {
+    const retryInstruction = resolvePlanningOnlyRetryInstruction({
+      provider: "google-gemini-cli",
+      modelId: "gemini-3.1-pro",
+      prompt: "Please inspect the code, make the change, and run the checks.",
+      aborted: false,
+      timedOut: false,
+      attempt: makeAttemptResult({
+        assistantTexts: ["I'll inspect the code, make the change, and run the checks."],
+      }),
+    });
+
+    expect(retryInstruction).toContain("Do not restate the plan");
+  });
+
   it("does not misclassify a direct answer that says 'i'm not going to' as planning-only", () => {
     const retryInstruction = resolvePlanningOnlyRetryInstruction({
       provider: "openai-codex",
diff --git a/src/agents/pi-embedded-runner/run/incomplete-turn.ts b/src/agents/pi-embedded-runner/run/incomplete-turn.ts
index 961b5b8fddd..1b9429eb23a 100644
--- a/src/agents/pi-embedded-runner/run/incomplete-turn.ts
+++ b/src/agents/pi-embedded-runner/run/incomplete-turn.ts
@@ -77,6 +77,13 @@ const SINGLE_ACTION_RETRY_SAFE_TOOL_NAMES = new Set([
   "glob",
   "ls",
 ]);
+const GEMINI_INCOMPLETE_TURN_PROVIDER_IDS = new Set([
+  "google",
+  "google-vertex",
+  "google-antigravity",
+  "google-gemini-cli",
+]);
+const GEMINI_INCOMPLETE_TURN_MODEL_ID_PATTERN = /^gemini(?:[.-]|$)/i;
 const DEFAULT_PLANNING_ONLY_RETRY_LIMIT = 1;
 const STRICT_AGENTIC_PLANNING_ONLY_RETRY_LIMIT = 2;
 // Allow one immediate continuation plus one follow-up continuation before
@@ -394,12 +401,38 @@ function shouldApplyPlanningOnlyRetryGuard(params: {
   if (params.executionContract === "strict-agentic") {
     return true;
   }
-  return isStrictAgenticSupportedProviderModel({
+  return isIncompleteTurnRecoverySupportedProviderModel({
     provider: params.provider,
     modelId: params.modelId,
   });
 }
 
+function stripProviderPrefix(modelId: string): string {
+  const normalizedModelId = modelId.trim();
+  const match = /^([^/:]+)[/:](.+)$/.exec(normalizedModelId);
+  return (match?.[2] ?? normalizedModelId).toLowerCase();
+}
+
+function isIncompleteTurnRecoverySupportedProviderModel(params: {
+  provider?: string;
+  modelId?: string;
+}): boolean {
+  if (
+    isStrictAgenticSupportedProviderModel({
+      provider: params.provider,
+      modelId: params.modelId,
+    })
+  ) {
+    return true;
+  }
+  const provider = normalizeLowercaseStringOrEmpty(params.provider ?? "");
+  if (!GEMINI_INCOMPLETE_TURN_PROVIDER_IDS.has(provider)) {
+    return false;
+  }
+  const modelId = typeof params.modelId === "string" ? params.modelId : "";
+  return GEMINI_INCOMPLETE_TURN_MODEL_ID_PATTERN.test(stripProviderPrefix(modelId));
+}
+
 function normalizeAckPrompt(text: string): string {
   const normalized = text
     .normalize("NFKC")