Skip to content

Commit 542ffc3

Browse files
asp345asp345
committed
fix(pi): preserve tool call metadata across turns
1 parent 5bb1bcb commit 542ffc3

4 files changed

Lines changed: 175 additions & 17 deletions

File tree

‎packages/pi/src/convert.test.ts‎

Lines changed: 67 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ describe("buildGeminiRequest", () => {
5454
role: "model",
5555
parts: [
5656
{ text: "thinking out loud" },
57-
{ functionCall: { name: "read", args: { path: "a.ts" } } },
57+
{ functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" } },
5858
],
5959
},
6060
])
@@ -93,7 +93,7 @@ describe("buildGeminiRequest", () => {
9393
}),
9494
)
9595
expect(request.contents[0]?.parts[0]).toEqual({
96-
functionCall: { name: "read", args: { path: "a.ts" } },
96+
functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" },
9797
thoughtSignature: "SIG123",
9898
})
9999
})
@@ -125,8 +125,70 @@ describe("buildGeminiRequest", () => {
125125
{
126126
role: "user",
127127
parts: [
128-
{ functionResponse: { name: "read", response: { output: "file A" } } },
129-
{ functionResponse: { name: "grep", response: { output: "match" } } },
128+
{ functionResponse: { name: "read", response: { output: "file A" }, id: "c1" } },
129+
{ functionResponse: { name: "grep", response: { output: "match" }, id: "c2" } },
130+
],
131+
},
132+
])
133+
})
134+
135+
it("preserves matching IDs across parallel tool calls and results", () => {
136+
const request = buildGeminiRequest(
137+
ctx({
138+
messages: [
139+
{
140+
role: "assistant",
141+
content: [
142+
{ type: "toolCall", id: "c1", name: "read", arguments: { path: "a.ts" } },
143+
{ type: "toolCall", id: "c2", name: "grep", arguments: { pattern: "TODO" } },
144+
],
145+
api: "google-generative-ai",
146+
provider: "google-antigravity",
147+
model: "antigravity-claude-opus-4-6-thinking",
148+
usage: {
149+
input: 0,
150+
output: 0,
151+
cacheRead: 0,
152+
cacheWrite: 0,
153+
totalTokens: 0,
154+
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
155+
},
156+
stopReason: "toolUse",
157+
timestamp: 0,
158+
},
159+
{
160+
role: "toolResult",
161+
toolCallId: "c1",
162+
toolName: "read",
163+
content: [{ type: "text", text: "file A" }],
164+
isError: false,
165+
timestamp: 1,
166+
},
167+
{
168+
role: "toolResult",
169+
toolCallId: "c2",
170+
toolName: "grep",
171+
content: [{ type: "text", text: "match" }],
172+
isError: false,
173+
timestamp: 1,
174+
},
175+
],
176+
}),
177+
)
178+
179+
expect(request.contents).toEqual([
180+
{
181+
role: "model",
182+
parts: [
183+
{ functionCall: { name: "read", args: { path: "a.ts" }, id: "c1" } },
184+
{ functionCall: { name: "grep", args: { pattern: "TODO" }, id: "c2" } },
185+
],
186+
},
187+
{
188+
role: "user",
189+
parts: [
190+
{ functionResponse: { name: "read", response: { output: "file A" }, id: "c1" } },
191+
{ functionResponse: { name: "grep", response: { output: "match" }, id: "c2" } },
130192
],
131193
},
132194
])
@@ -148,7 +210,7 @@ describe("buildGeminiRequest", () => {
148210
}),
149211
)
150212
expect(request.contents[0]?.parts[0]).toEqual({
151-
functionResponse: { name: "bash", response: { error: "boom" } },
213+
functionResponse: { name: "bash", response: { error: "boom" }, id: "c1" },
152214
})
153215
})
154216

‎packages/pi/src/convert.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,11 @@ import type {
1414
type GeminiPart =
1515
| { text: string }
1616
| { inlineData: { mimeType: string; data: string } }
17-
| { functionCall: { name: string; args: Record<string, unknown> }; thoughtSignature?: string }
18-
| { functionResponse: { name: string; response: Record<string, unknown> } }
17+
| {
18+
functionCall: { name: string; args: Record<string, unknown>; id: string }
19+
thoughtSignature?: string
20+
}
21+
| { functionResponse: { name: string; response: Record<string, unknown>; id: string } }
1922

2023
interface GeminiContent {
2124
role: "user" | "model"
@@ -64,6 +67,7 @@ function convertAssistantParts(content: Array<TextContent | ThinkingContent | To
6467
functionCall: {
6568
name: block.name,
6669
args: (block.arguments ?? {}) as Record<string, unknown>,
70+
id: block.id,
6771
},
6872
...(block.thoughtSignature ? { thoughtSignature: block.thoughtSignature } : {}),
6973
})
@@ -114,6 +118,7 @@ function convertMessages(messages: Message[]): GeminiContent[] {
114118
functionResponse: {
115119
name: message.toolName,
116120
response: toolResultResponse(message),
121+
id: message.toolCallId,
117122
},
118123
}
119124
// Gemini groups consecutive function responses into one user turn.

‎packages/pi/src/stream.test.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"
22
import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai"
33

44
import {
5+
convertGeminiToolCallPart,
56
finalizePiAntigravityRequest,
67
parseGeminiSse,
78
resolvePiAntigravityModel,
@@ -119,6 +120,73 @@ describe("finalizePiAntigravityRequest", () => {
119120
})
120121
})
121122

123+
describe("convertGeminiToolCallPart", () => {
124+
it("preserves the backend function-call ID", () => {
125+
const state = {}
126+
const toolCall = convertGeminiToolCallPart(
127+
{ functionCall: { name: "read", args: { path: "a.ts" }, id: "toolu_vrtx_123" } },
128+
state,
129+
)
130+
131+
expect(toolCall).toEqual({
132+
type: "toolCall",
133+
id: "toolu_vrtx_123",
134+
name: "read",
135+
arguments: { path: "a.ts" },
136+
})
137+
})
138+
139+
it("generates an ID when the backend omits one", () => {
140+
const toolCall = convertGeminiToolCallPart(
141+
{ functionCall: { name: "read", args: {} } },
142+
{},
143+
)
144+
145+
expect(toolCall?.id).toMatch(/^call_[0-9a-f-]{36}$/)
146+
})
147+
148+
it("carries a preceding thought signature onto the next function call", () => {
149+
const state = {}
150+
151+
expect(convertGeminiToolCallPart(
152+
{ text: "", thought: true, thoughtSignature: "SIG123" },
153+
state,
154+
)).toBeUndefined()
155+
156+
expect(convertGeminiToolCallPart(
157+
{ functionCall: { name: "read", args: {}, id: "c1" } },
158+
state,
159+
)).toEqual({
160+
type: "toolCall",
161+
id: "c1",
162+
name: "read",
163+
arguments: {},
164+
thoughtSignature: "SIG123",
165+
})
166+
expect(convertGeminiToolCallPart(
167+
{ functionCall: { name: "grep", args: {}, id: "c2" } },
168+
state,
169+
)).not.toHaveProperty("thoughtSignature")
170+
})
171+
172+
it("attaches a parallel batch signature only to the first function call", () => {
173+
const state = {}
174+
175+
convertGeminiToolCallPart({ thought: true, thoughtSignature: "SIG1" }, state)
176+
const first = convertGeminiToolCallPart(
177+
{ functionCall: { name: "read", args: {}, id: "c1" } },
178+
state,
179+
)
180+
const second = convertGeminiToolCallPart(
181+
{ functionCall: { name: "grep", args: {}, id: "c2" } },
182+
state,
183+
)
184+
185+
expect(first?.thoughtSignature).toBe("SIG1")
186+
expect(second).not.toHaveProperty("thoughtSignature")
187+
})
188+
})
189+
122190
describe("parseGeminiSse", () => {
123191
it("parses and unwraps the Antigravity response envelope into chunks", async () => {
124192
// Antigravity wraps each chunk under a `response` key (MITM-verified).

‎packages/pi/src/stream.ts‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -94,11 +94,39 @@ function unwrapChunk(raw: unknown): GeminiStreamChunk {
9494
return raw as GeminiStreamChunk
9595
}
9696

97-
interface GeminiResponsePart {
97+
export interface GeminiResponsePart {
9898
text?: string
9999
thought?: boolean
100100
thoughtSignature?: string
101-
functionCall?: { name?: string; args?: Record<string, unknown> }
101+
functionCall?: { name?: string; args?: Record<string, unknown>; id?: string }
102+
}
103+
104+
export interface GeminiToolCallState {
105+
pendingThoughtSignature?: string
106+
}
107+
108+
export function convertGeminiToolCallPart(
109+
part: GeminiResponsePart,
110+
state: GeminiToolCallState,
111+
): ToolCall | undefined {
112+
// Antigravity emits a batch signature on a preceding empty thought part;
113+
// native replay attaches it to the first function call in that batch.
114+
if (part.thought && part.thoughtSignature) {
115+
state.pendingThoughtSignature = part.thoughtSignature
116+
}
117+
118+
if (!part.functionCall) return undefined
119+
120+
const thoughtSignature = part.thoughtSignature ?? state.pendingThoughtSignature
121+
state.pendingThoughtSignature = undefined
122+
123+
return {
124+
type: "toolCall",
125+
id: part.functionCall.id ?? `call_${crypto.randomUUID()}`,
126+
name: part.functionCall.name ?? "",
127+
arguments: (part.functionCall.args ?? {}) as Record<string, unknown>,
128+
...(thoughtSignature ? { thoughtSignature } : {}),
129+
}
102130
}
103131

104132
export function updateUsage(model: Model<Api>, output: AssistantMessage, usage?: GeminiUsageMetadata): void {
@@ -314,6 +342,7 @@ export function streamCortexKitAntigravity(
314342
}
315343

316344
const content = output.content as Array<TextContent | ToolCall>
345+
const toolCallState: GeminiToolCallState = {}
317346
let textIndex = -1
318347
let finished = false
319348

@@ -324,14 +353,8 @@ export function streamCortexKitAntigravity(
324353
const parts = candidate?.content?.parts ?? []
325354

326355
for (const part of parts) {
327-
if (part.functionCall) {
328-
const toolCall: ToolCall = {
329-
type: "toolCall",
330-
id: `call_${crypto.randomUUID()}`,
331-
name: part.functionCall.name ?? "",
332-
arguments: (part.functionCall.args ?? {}) as Record<string, unknown>,
333-
...(part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : {}),
334-
}
356+
const toolCall = convertGeminiToolCallPart(part, toolCallState)
357+
if (toolCall) {
335358
content.push(toolCall)
336359
const idx = content.length - 1
337360
textIndex = -1

0 commit comments

Comments
 (0)