fix(pi): preserve Antigravity state across turns

packages/core/src/agy-request-metadata.test.ts

@@ -43,6 +43,22 @@ describe("agy request metadata", () => {
     expect(other.session.numericSessionId).toBe(first.session.numericSessionId)
   })
 
+  it("records a fresh execution ID only for known sessions", () => {
+    const sessions = new AgyRequestSessionStore("file:///workspace", { now: () => 100 })
+    const session = sessions.beginRequest("session-a").session
+
+    sessions.completeExecution("missing")
+    expect(session.lastExecutionId).toBeUndefined()
+
+    sessions.completeExecution("session-a")
+    const firstExecutionId = session.lastExecutionId
+    expect(firstExecutionId).toMatch(/^[0-9a-f-]{36}$/)
+
+    sessions.completeExecution("session-a")
+    expect(session.lastExecutionId).toMatch(/^[0-9a-f-]{36}$/)
+    expect(session.lastExecutionId).not.toBe(firstExecutionId)
+  })
+
   it("creates stable session IDs with independently generated conversation and trajectory IDs", () => {
     expect(createAgyRequestSessionContext("file:///workspace", {
       conversationId: "conversation-id",
@@ -78,7 +94,7 @@ describe("agy request metadata", () => {
     ])
   })
 
-  it("derives last_step_index from the number of content parts", () => {
+  it("supports part- and content-based step counting", () => {
     const payload = {
       contents: [
         { role: "user", parts: [{ text: "prompt" }] },
@@ -94,7 +110,8 @@ describe("agy request metadata", () => {
     }
 
     expect(countAgyRequestSteps(payload)).toBe(4)
-    expect(countAgyRequestSteps({ contents: [] })).toBe(1)
+    expect(countAgyRequestSteps(payload, "contents")).toBe(3)
+    expect(countAgyRequestSteps({ contents: [] }, "contents")).toBe(1)
     expect(countAgyRequestSteps({})).toBe(1)
   })
 
@@ -132,6 +149,40 @@ describe("agy request metadata", () => {
     })
   })
 
+  it("matches the captured execution-aware step sequence", () => {
+    const session = createAgyRequestSessionContext("", {
+      conversationId: "conversation-id",
+      trajectoryId: "trajectory-id",
+    })
+    const sequence = [
+      { contents: 1, step: 1, executionId: undefined },
+      { contents: 4, step: 5, executionId: "execution-1" },
+      { contents: 7, step: 8, executionId: "execution-2" },
+      { contents: 9, step: 10, executionId: "execution-2" },
+      { contents: 12, step: 13, executionId: "execution-3" },
+      { contents: 15, step: 16, executionId: "execution-4" },
+    ]
+
+    for (const [index, item] of sequence.entries()) {
+      session.lastExecutionId = item.executionId
+      const metadata = buildAgyAgentRequestMetadata(
+        session,
+        { contents: Array.from({ length: item.contents }, () => ({ role: "user", parts: [] })) },
+        "claude-opus-4-6-thinking",
+        index + 1,
+        { stepCountMode: "contents" },
+      )
+
+      expect(metadata.lastStepIndex).toBe(item.step)
+      expect(metadata.requestId.endsWith(`/${item.step + 1}`)).toBe(true)
+      if (item.executionId) {
+        expect(metadata.labels.last_execution_id).toBe(item.executionId)
+      } else {
+        expect(metadata.labels).not.toHaveProperty("last_execution_id")
+      }
+    }
+  })
+
   it("matches every captured agy model enum fixture", () => {
     for (const fixture of MODEL_METADATA_FIXTURES) {
       for (const [model, expected] of Object.entries(fixture.models)) {

packages/core/src/agy-request-metadata.ts

@@ -37,6 +37,7 @@ export interface AgyRequestSessionContext {
   numericSessionId: string
   usedClaude?: boolean
   usedNonGeminiModel?: boolean
+  lastExecutionId?: string
 }
 
 interface StoredAgyRequestSession {
@@ -57,6 +58,7 @@ export interface AgyRequestScope {
 }
 
 export interface AgyRequestLabels {
+  last_execution_id?: string
   last_step_index: string
   model_enum?: string
   trajectory_id: string
@@ -72,6 +74,10 @@ export interface AgyAgentRequestMetadata {
   lastStepIndex: number
 }
 
+export interface AgyAgentRequestMetadataOptions {
+  stepCountMode?: "parts" | "contents"
+}
+
 export function fnv1a64Signed(input: string): string {
   let hash = FNV1A_64_OFFSET_BASIS
   for (const byte of Buffer.from(input, "utf8")) {
@@ -133,6 +139,13 @@ export class AgyRequestSessionStore {
     return { session, timestamp }
   }
 
+  completeExecution(key: string): void {
+    const stored = this.entries.get(key)
+    if (stored) {
+      stored.context.lastExecutionId = randomUUID()
+    }
+  }
+
   has(key: string): boolean {
     return this.entries.has(key)
   }
@@ -194,21 +207,19 @@ export function orderAgyRequestPayloadInPlace(payload: Record<string, unknown>):
   Object.assign(payload, ordered)
 }
 
-export function countAgyRequestSteps(payload: Record<string, unknown>): number {
+export function countAgyRequestSteps(
+  payload: Record<string, unknown>,
+  mode: "parts" | "contents" = "parts",
+): number {
   const contents = payload.contents
-  if (!Array.isArray(contents)) {
-    return 1
-  }
+  if (!Array.isArray(contents)) return 1
+  if (mode === "contents") return Math.max(1, contents.length)
 
   let partCount = 0
   for (const content of contents) {
-    if (!content || typeof content !== "object" || Array.isArray(content)) {
-      continue
-    }
+    if (!content || typeof content !== "object" || Array.isArray(content)) continue
     const parts = (content as Record<string, unknown>).parts
-    if (Array.isArray(parts)) {
-      partCount += parts.length
-    }
+    if (Array.isArray(parts)) partCount += parts.length
   }
   return Math.max(1, partCount)
 }
@@ -218,14 +229,17 @@ export function buildAgyAgentRequestMetadata(
   payload: Record<string, unknown>,
   model: string,
   timestamp = Date.now(),
+  options: AgyAgentRequestMetadataOptions = {},
 ): AgyAgentRequestMetadata {
-  const lastStepIndex = countAgyRequestSteps(payload)
+  const lastStepIndex = countAgyRequestSteps(payload, options.stepCountMode)
+    + (session.lastExecutionId ? 1 : 0)
   const isClaude = model.toLowerCase().startsWith("claude-")
   const isNonGemini = isClaude || model.toLowerCase().startsWith("gpt-")
   session.usedClaude = session.usedClaude === true || isClaude
   session.usedNonGeminiModel = session.usedNonGeminiModel === true || isNonGemini
   const modelEnum = getAgyModelEnum(model)
   const labels: AgyRequestLabels = {
+    ...(session.lastExecutionId ? { last_execution_id: session.lastExecutionId } : {}),
     last_step_index: String(lastStepIndex),
     ...(modelEnum ? { model_enum: modelEnum } : {}),
     trajectory_id: session.trajectoryId,

packages/pi/src/convert.test.ts

@@ -54,7 +54,7 @@ describe("buildGeminiRequest", () => {
         role: "model",
         parts: [
           { text: "thinking out loud" },
-          { functionCall: { name: "read", args: { path: "a.ts" } } },
+          { functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" } },
         ],
       },
     ])
@@ -93,11 +93,118 @@ describe("buildGeminiRequest", () => {
       }),
     )
     expect(request.contents[0]?.parts[0]).toEqual({
-      functionCall: { name: "read", args: { path: "a.ts" } },
+      functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" },
       thoughtSignature: "SIG123",
     })
   })
 
+  it("replays same-model thinking and signed text", () => {
+    const request = buildGeminiRequest(
+      ctx({
+        messages: [
+          {
+            role: "assistant",
+            content: [
+              { type: "thinking", thinking: "reasoning" },
+              { type: "text", text: "answer", textSignature: "SIG123" },
+            ],
+            api: "google-generative-ai",
+            provider: "google-antigravity",
+            model: "antigravity-claude-opus-4-6-thinking",
+            usage: {
+              input: 0,
+              output: 0,
+              cacheRead: 0,
+              cacheWrite: 0,
+              totalTokens: 0,
+              cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+            },
+            stopReason: "stop",
+            timestamp: 0,
+          },
+        ],
+      }),
+      {
+        provider: "google-antigravity",
+        model: "antigravity-claude-opus-4-6-thinking",
+      },
+    )
+
+    expect(request.contents).toEqual([
+      {
+        role: "model",
+        parts: [
+          { text: "reasoning", thought: true },
+          { text: "answer", thoughtSignature: "SIG123" },
+        ],
+      },
+    ])
+  })
+
+  it("strips foreign thinking and signatures and uses model-role tool results", () => {
+    const request = buildGeminiRequest(
+      ctx({
+        messages: [
+          {
+            role: "assistant",
+            content: [
+              { type: "thinking", thinking: "claude reasoning" },
+              { type: "text", text: "before tool", textSignature: "TEXT_SIG" },
+              {
+                type: "toolCall",
+                id: "c1",
+                name: "read",
+                arguments: { path: "a.ts" },
+                thoughtSignature: "TOOL_SIG",
+              },
+            ],
+            api: "google-generative-ai",
+            provider: "google-antigravity",
+            model: "antigravity-claude-opus-4-6-thinking",
+            usage: {
+              input: 0,
+              output: 0,
+              cacheRead: 0,
+              cacheWrite: 0,
+              totalTokens: 0,
+              cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+            },
+            stopReason: "toolUse",
+            timestamp: 0,
+          },
+          {
+            role: "toolResult",
+            toolCallId: "c1",
+            toolName: "read",
+            content: [{ type: "text", text: "file A" }],
+            isError: false,
+            timestamp: 1,
+          },
+        ],
+      }),
+      {
+        provider: "google-antigravity",
+        model: "antigravity-gemini-3.6-flash",
+      },
+    )
+
+    expect(request.contents).toEqual([
+      {
+        role: "model",
+        parts: [
+          { text: "before tool" },
+          { functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" } },
+        ],
+      },
+      {
+        role: "model",
+        parts: [
+          { functionResponse: { name: "read", response: { output: "file A" }, id: "c1" } },
+        ],
+      },
+    ])
+  })
+
   it("groups consecutive tool results into a single user turn", () => {
     const request = buildGeminiRequest(
       ctx({
@@ -125,8 +232,70 @@ describe("buildGeminiRequest", () => {
       {
         role: "user",
         parts: [
-          { functionResponse: { name: "read", response: { output: "file A" } } },
-          { functionResponse: { name: "grep", response: { output: "match" } } },
+          { functionResponse: { name: "read", response: { output: "file A" }, id: "c1" } },
+          { functionResponse: { name: "grep", response: { output: "match" }, id: "c2" } },
+        ],
+      },
+    ])
+  })
+
+  it("preserves matching IDs across parallel tool calls and results", () => {
+    const request = buildGeminiRequest(
+      ctx({
+        messages: [
+          {
+            role: "assistant",
+            content: [
+              { type: "toolCall", id: "c1", name: "read", arguments: { path: "a.ts" } },
+              { type: "toolCall", id: "c2", name: "grep", arguments: { pattern: "TODO" } },
+            ],
+            api: "google-generative-ai",
+            provider: "google-antigravity",
+            model: "antigravity-claude-opus-4-6-thinking",
+            usage: {
+              input: 0,
+              output: 0,
+              cacheRead: 0,
+              cacheWrite: 0,
+              totalTokens: 0,
+              cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+            },
+            stopReason: "toolUse",
+            timestamp: 0,
+          },
+          {
+            role: "toolResult",
+            toolCallId: "c1",
+            toolName: "read",
+            content: [{ type: "text", text: "file A" }],
+            isError: false,
+            timestamp: 1,
+          },
+          {
+            role: "toolResult",
+            toolCallId: "c2",
+            toolName: "grep",
+            content: [{ type: "text", text: "match" }],
+            isError: false,
+            timestamp: 1,
+          },
+        ],
+      }),
+    )
+
+    expect(request.contents).toEqual([
+      {
+        role: "model",
+        parts: [
+          { functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" } },
+          { functionCall: { name: "grep", args: { pattern: "TODO" }, id: "c2" } },
+        ],
+      },
+      {
+        role: "user",
+        parts: [
+          { functionResponse: { name: "read", response: { output: "file A" }, id: "c1" } },
+          { functionResponse: { name: "grep", response: { output: "match" }, id: "c2" } },
         ],
       },
     ])
@@ -148,7 +317,7 @@ describe("buildGeminiRequest", () => {
       }),
     )
     expect(request.contents[0]?.parts[0]).toEqual({
-      functionResponse: { name: "bash", response: { error: "boom" } },
+      functionResponse: { name: "bash", response: { error: "boom" }, id: "c1" },
     })
   })