From 542ffc394e57bc4c2dc8ca8a503e669d9c0575c4 Mon Sep 17 00:00:00 2001 From: Chanyeong Lim Date: Fri, 24 Jul 2026 11:46:59 +0900 Subject: [PATCH 1/3] fix(pi): preserve tool call metadata across turns --- packages/pi/src/convert.test.ts | 72 ++++++++++++++++++++++++++++++--- packages/pi/src/convert.ts | 9 ++++- packages/pi/src/stream.test.ts | 68 +++++++++++++++++++++++++++++++ packages/pi/src/stream.ts | 43 +++++++++++++++----- 4 files changed, 175 insertions(+), 17 deletions(-) diff --git a/packages/pi/src/convert.test.ts b/packages/pi/src/convert.test.ts index 0f1a50a8..6a0764fd 100644 --- a/packages/pi/src/convert.test.ts +++ b/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,7 +93,7 @@ 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", }) }) @@ -125,8 +125,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 +210,7 @@ describe("buildGeminiRequest", () => { }), ) expect(request.contents[0]?.parts[0]).toEqual({ - functionResponse: { name: "bash", response: { error: "boom" } }, + functionResponse: { name: "bash", response: { error: "boom" }, id: "c1" }, }) }) diff --git a/packages/pi/src/convert.ts b/packages/pi/src/convert.ts index 5e330658..c9beb727 100644 --- a/packages/pi/src/convert.ts +++ b/packages/pi/src/convert.ts @@ -14,8 +14,11 @@ import type { type GeminiPart = | { text: string } | { inlineData: { mimeType: string; data: string } } - | { functionCall: { name: string; args: Record }; thoughtSignature?: string } - | { functionResponse: { name: string; response: Record } } + | { + functionCall: { name: string; args: Record; id: string } + thoughtSignature?: string + } + | { functionResponse: { name: string; response: Record; id: string } } interface GeminiContent { role: "user" | "model" @@ -64,6 +67,7 @@ function convertAssistantParts(content: Array, + id: block.id, }, ...(block.thoughtSignature ? { thoughtSignature: block.thoughtSignature } : {}), }) @@ -114,6 +118,7 @@ function convertMessages(messages: Message[]): GeminiContent[] { functionResponse: { name: message.toolName, response: toolResultResponse(message), + id: message.toolCallId, }, } // Gemini groups consecutive function responses into one user turn. diff --git a/packages/pi/src/stream.test.ts b/packages/pi/src/stream.test.ts index 12be08d9..daad4b65 100644 --- a/packages/pi/src/stream.test.ts +++ b/packages/pi/src/stream.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest" import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai" import { + convertGeminiToolCallPart, finalizePiAntigravityRequest, parseGeminiSse, resolvePiAntigravityModel, @@ -119,6 +120,73 @@ describe("finalizePiAntigravityRequest", () => { }) }) +describe("convertGeminiToolCallPart", () => { + it("preserves the backend function-call ID", () => { + const state = {} + const toolCall = convertGeminiToolCallPart( + { functionCall: { name: "read", args: { path: "a.ts" }, id: "toolu_vrtx_123" } }, + state, + ) + + expect(toolCall).toEqual({ + type: "toolCall", + id: "toolu_vrtx_123", + name: "read", + arguments: { path: "a.ts" }, + }) + }) + + it("generates an ID when the backend omits one", () => { + const toolCall = convertGeminiToolCallPart( + { functionCall: { name: "read", args: {} } }, + {}, + ) + + expect(toolCall?.id).toMatch(/^call_[0-9a-f-]{36}$/) + }) + + it("carries a preceding thought signature onto the next function call", () => { + const state = {} + + expect(convertGeminiToolCallPart( + { text: "", thought: true, thoughtSignature: "SIG123" }, + state, + )).toBeUndefined() + + expect(convertGeminiToolCallPart( + { functionCall: { name: "read", args: {}, id: "c1" } }, + state, + )).toEqual({ + type: "toolCall", + id: "c1", + name: "read", + arguments: {}, + thoughtSignature: "SIG123", + }) + expect(convertGeminiToolCallPart( + { functionCall: { name: "grep", args: {}, id: "c2" } }, + state, + )).not.toHaveProperty("thoughtSignature") + }) + + it("attaches a parallel batch signature only to the first function call", () => { + const state = {} + + convertGeminiToolCallPart({ thought: true, thoughtSignature: "SIG1" }, state) + const first = convertGeminiToolCallPart( + { functionCall: { name: "read", args: {}, id: "c1" } }, + state, + ) + const second = convertGeminiToolCallPart( + { functionCall: { name: "grep", args: {}, id: "c2" } }, + state, + ) + + expect(first?.thoughtSignature).toBe("SIG1") + expect(second).not.toHaveProperty("thoughtSignature") + }) +}) + describe("parseGeminiSse", () => { it("parses and unwraps the Antigravity response envelope into chunks", async () => { // Antigravity wraps each chunk under a `response` key (MITM-verified). diff --git a/packages/pi/src/stream.ts b/packages/pi/src/stream.ts index 9fae68e7..73f217bb 100644 --- a/packages/pi/src/stream.ts +++ b/packages/pi/src/stream.ts @@ -94,11 +94,39 @@ function unwrapChunk(raw: unknown): GeminiStreamChunk { return raw as GeminiStreamChunk } -interface GeminiResponsePart { +export interface GeminiResponsePart { text?: string thought?: boolean thoughtSignature?: string - functionCall?: { name?: string; args?: Record } + functionCall?: { name?: string; args?: Record; id?: string } +} + +export interface GeminiToolCallState { + pendingThoughtSignature?: string +} + +export function convertGeminiToolCallPart( + part: GeminiResponsePart, + state: GeminiToolCallState, +): ToolCall | undefined { + // Antigravity emits a batch signature on a preceding empty thought part; + // native replay attaches it to the first function call in that batch. + if (part.thought && part.thoughtSignature) { + state.pendingThoughtSignature = part.thoughtSignature + } + + if (!part.functionCall) return undefined + + const thoughtSignature = part.thoughtSignature ?? state.pendingThoughtSignature + state.pendingThoughtSignature = undefined + + return { + type: "toolCall", + id: part.functionCall.id ?? `call_${crypto.randomUUID()}`, + name: part.functionCall.name ?? "", + arguments: (part.functionCall.args ?? {}) as Record, + ...(thoughtSignature ? { thoughtSignature } : {}), + } } export function updateUsage(model: Model, output: AssistantMessage, usage?: GeminiUsageMetadata): void { @@ -314,6 +342,7 @@ export function streamCortexKitAntigravity( } const content = output.content as Array + const toolCallState: GeminiToolCallState = {} let textIndex = -1 let finished = false @@ -324,14 +353,8 @@ export function streamCortexKitAntigravity( const parts = candidate?.content?.parts ?? [] for (const part of parts) { - if (part.functionCall) { - const toolCall: ToolCall = { - type: "toolCall", - id: `call_${crypto.randomUUID()}`, - name: part.functionCall.name ?? "", - arguments: (part.functionCall.args ?? {}) as Record, - ...(part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : {}), - } + const toolCall = convertGeminiToolCallPart(part, toolCallState) + if (toolCall) { content.push(toolCall) const idx = content.length - 1 textIndex = -1 From b27a24e7362a2a1c50bc06ec0c95fbb1756477fc Mon Sep 17 00:00:00 2001 From: Chanyeong Lim Date: Fri, 24 Jul 2026 13:25:39 +0900 Subject: [PATCH 2/3] fix(core): track Antigravity execution metadata --- .../core/src/agy-request-metadata.test.ts | 55 ++++++++++++++++++- packages/core/src/agy-request-metadata.ts | 36 ++++++++---- 2 files changed, 78 insertions(+), 13 deletions(-) diff --git a/packages/core/src/agy-request-metadata.test.ts b/packages/core/src/agy-request-metadata.test.ts index eb9cfcc6..0f75b9ea 100644 --- a/packages/core/src/agy-request-metadata.test.ts +++ b/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)) { diff --git a/packages/core/src/agy-request-metadata.ts b/packages/core/src/agy-request-metadata.ts index 19798af7..7c6b51a4 100644 --- a/packages/core/src/agy-request-metadata.ts +++ b/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): Object.assign(payload, ordered) } -export function countAgyRequestSteps(payload: Record): number { +export function countAgyRequestSteps( + payload: Record, + 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).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, 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, From d3c3bd66690abd56d32124faed8f44b65ecc2fe3 Mon Sep 17 00:00:00 2001 From: Chanyeong Lim Date: Fri, 24 Jul 2026 13:25:40 +0900 Subject: [PATCH 3/3] fix(pi): preserve signed history across model turns --- packages/pi/src/convert.test.ts | 107 +++++++++++++++++ packages/pi/src/convert.ts | 79 ++++++++++--- packages/pi/src/stream.test.ts | 201 +++++++++++++++++++++++++++++++- packages/pi/src/stream.ts | 176 +++++++++++++++++++++++----- 4 files changed, 516 insertions(+), 47 deletions(-) diff --git a/packages/pi/src/convert.test.ts b/packages/pi/src/convert.test.ts index 6a0764fd..b7de42d8 100644 --- a/packages/pi/src/convert.test.ts +++ b/packages/pi/src/convert.test.ts @@ -98,6 +98,113 @@ describe("buildGeminiRequest", () => { }) }) + 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({ diff --git a/packages/pi/src/convert.ts b/packages/pi/src/convert.ts index c9beb727..1e9d3619 100644 --- a/packages/pi/src/convert.ts +++ b/packages/pi/src/convert.ts @@ -1,5 +1,6 @@ import { toGeminiSchema } from "@cortexkit/antigravity-auth-core" import type { + AssistantMessage, Context, ImageContent, Message, @@ -12,7 +13,7 @@ import type { /** Gemini `contents` part shapes. */ type GeminiPart = - | { text: string } + | { text: string; thought?: boolean; thoughtSignature?: string } | { inlineData: { mimeType: string; data: string } } | { functionCall: { name: string; args: Record; id: string } @@ -55,26 +56,39 @@ function convertUserParts(content: Array): GeminiPar return parts } -function convertAssistantParts(content: Array): GeminiPart[] { +function convertAssistantParts( + message: AssistantMessage, + preserveSignedHistory: boolean, +): GeminiPart[] { const parts: GeminiPart[] = [] - for (const block of content) { - if (block.type === "text" && block.text.trim()) { - parts.push({ text: sanitize(block.text) }) + for (const block of message.content) { + if (block.type === "thinking") { + if (preserveSignedHistory && block.thinking) { + parts.push({ + text: sanitize(block.thinking), + thought: true, + ...(block.thinkingSignature ? { thoughtSignature: block.thinkingSignature } : {}), + }) + } + } else if (block.type === "text" && block.text.trim()) { + parts.push({ + text: sanitize(block.text), + ...(preserveSignedHistory && block.textSignature + ? { thoughtSignature: block.textSignature } + : {}), + }) } else if (block.type === "toolCall") { - // Antigravity requires the prior functionCall to echo its - // thoughtSignature on replay (400 INVALID_ARGUMENT otherwise). parts.push({ functionCall: { name: block.name, args: (block.arguments ?? {}) as Record, id: block.id, }, - ...(block.thoughtSignature ? { thoughtSignature: block.thoughtSignature } : {}), + ...(preserveSignedHistory && block.thoughtSignature + ? { thoughtSignature: block.thoughtSignature } + : {}), }) } - // Thinking blocks are intentionally not replayed: OpenCode/pi history does - // not carry replayable signed Antigravity thinking, and unsigned thinking - // is rejected. The model regenerates thinking each turn. } return parts } @@ -90,8 +104,35 @@ function toolResultResponse(message: ToolResultMessage): Record return { output: text } } -function convertMessages(messages: Message[]): GeminiContent[] { +export interface BuildGeminiRequestOptions { + provider?: string + model?: string +} + +function isSameTargetModel( + message: AssistantMessage, + options: BuildGeminiRequestOptions | undefined, +): boolean { + if (!options?.provider || !options.model) return true + return message.provider === options.provider && message.model === options.model +} + +function convertMessages( + messages: Message[], + options?: BuildGeminiRequestOptions, +): GeminiContent[] { const contents: GeminiContent[] = [] + const callMatchesTarget = new Map() + + for (const message of messages) { + if (message?.role !== "assistant") continue + const matchesTarget = isSameTargetModel(message, options) + for (const block of message.content) { + if (block.type === "toolCall") { + callMatchesTarget.set(block.id, matchesTarget) + } + } + } for (const message of messages) { if (!message) continue @@ -108,12 +149,13 @@ function convertMessages(messages: Message[]): GeminiContent[] { } if (message.role === "assistant") { - const parts = convertAssistantParts(message.content) + const parts = convertAssistantParts(message, isSameTargetModel(message, options)) if (parts.length) contents.push({ role: "model", parts }) continue } if (message.role === "toolResult") { + const role = callMatchesTarget.get(message.toolCallId) === false ? "model" : "user" const part: GeminiPart = { functionResponse: { name: message.toolName, @@ -123,10 +165,10 @@ function convertMessages(messages: Message[]): GeminiContent[] { } // Gemini groups consecutive function responses into one user turn. const last = contents[contents.length - 1] - if (last && last.role === "user" && last.parts.every((p) => "functionResponse" in p)) { + if (last && last.role === role && last.parts.every((p) => "functionResponse" in p)) { last.parts.push(part) } else { - contents.push({ role: "user", parts: [part] }) + contents.push({ role, parts: [part] }) } } } @@ -154,9 +196,12 @@ function convertTools(tools: Tool[] | undefined): GeminiTool[] | undefined { * Convert a pi `Context` into a Gemini `generateContent` request body * (the inner `request` object of the Antigravity envelope). */ -export function buildGeminiRequest(context: Context): GeminiRequest { +export function buildGeminiRequest( + context: Context, + options?: BuildGeminiRequestOptions, +): GeminiRequest { const request: GeminiRequest = { - contents: convertMessages(context.messages), + contents: convertMessages(context.messages, options), } const tools = convertTools(context.tools) diff --git a/packages/pi/src/stream.test.ts b/packages/pi/src/stream.test.ts index daad4b65..e9a7972e 100644 --- a/packages/pi/src/stream.test.ts +++ b/packages/pi/src/stream.test.ts @@ -1,17 +1,31 @@ -import { describe, expect, it } from "vitest" -import type { Api, AssistantMessage, Model } from "@earendil-works/pi-ai" +import { afterEach, describe, expect, it, vi } from "vitest" +import { + ensureProjectContext, + fetchWithAgyCliTransport, +} from "@cortexkit/antigravity-auth-core" +import type { Api, AssistantMessage, Context, Model } from "@earendil-works/pi-ai" + +vi.mock("@cortexkit/antigravity-auth-core", async (importOriginal) => { + const actual = await importOriginal() + return { + ...actual, + ensureProjectContext: vi.fn(async () => ({ effectiveProjectId: "test-project" })), + fetchWithAgyCliTransport: vi.fn(), + } +}) import { convertGeminiToolCallPart, finalizePiAntigravityRequest, parseGeminiSse, resolvePiAntigravityModel, + streamCortexKitAntigravity, updateUsage, } from "./stream.ts" -function fakeModel(): Model { +function fakeModel(id = "antigravity-gemini-3.5-flash"): Model { return { - id: "antigravity-gemini-3.5-flash", + id, api: "google-generative-ai", provider: "google-antigravity", cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, @@ -51,6 +65,72 @@ function sseResponse(frames: string[]): Response { return new Response(body, { status: 200 }) } +function openSseResponse(frame: string): { + response: Response + wasCancelled: () => boolean +} { + let cancelled = false + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(frame)) + }, + cancel() { + cancelled = true + }, + }) + return { + response: new Response(body, { status: 200 }), + wasCancelled: () => cancelled, + } +} + +function stalledSseResponse(frame: string): { + response: Response + abort: () => void +} { + let streamController: ReadableStreamDefaultController | undefined + const body = new ReadableStream({ + start(controller) { + streamController = controller + controller.enqueue(new TextEncoder().encode(frame)) + }, + }) + return { + response: new Response(body, { status: 200 }), + abort: () => streamController?.error(new Error("aborted")), + } +} + +function userContext(): Context { + return { messages: [{ role: "user", content: "test", timestamp: 1 }] } +} + +async function runStream( + model: Model, + response: Response, + sessionId: string, + onAbort?: () => void, +) { + vi.mocked(fetchWithAgyCliTransport).mockImplementationOnce(async (_url, _init, transportOptions) => { + if (onAbort) transportOptions?.signal?.addEventListener("abort", onAbort, { once: true }) + return response + }) + const eventStream = streamCortexKitAntigravity(model, userContext(), { + apiKey: "test-token", + sessionId, + }) + const events = [] + for await (const event of eventStream) { + events.push(event) + } + return { events, result: await eventStream.result() } +} + +afterEach(() => { + vi.mocked(fetchWithAgyCliTransport).mockReset() + vi.mocked(ensureProjectContext).mockClear() +}) + describe("resolvePiAntigravityModel", () => { const gemini36 = { ...fakeModel(), @@ -187,6 +267,119 @@ describe("convertGeminiToolCallPart", () => { }) }) +describe("streamCortexKitAntigravity", () => { + it("streams thinking and transfers a pending signature to visible text", async () => { + const { events, result } = await runStream( + fakeModel("antigravity-claude-opus-4-6-thinking"), + sseResponse([ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"reasoning","thought":true}]}}]}}\n\n', + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"","thought":true,"thoughtSignature":"SIG123"}]}}]}}\n\n', + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]}}]}}\n\n', + 'data: {"response":{"candidates":[{"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,"thoughtsTokenCount":3}}}\n\n', + ]), + "thinking-text-signature", + ) + + expect(events.map((event) => event.type)).toEqual([ + "start", + "thinking_start", + "thinking_delta", + "thinking_end", + "text_start", + "text_delta", + "text_end", + "done", + ]) + expect(result.content).toEqual([ + { type: "thinking", thinking: "reasoning" }, + { type: "text", text: "answer", textSignature: "SIG123" }, + ]) + }) + + it("adds execution metadata after a completed turn", async () => { + const terminal = () => sseResponse([ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}}\n\n', + ]) + + await runStream(fakeModel(), terminal(), "execution-metadata") + await runStream(fakeModel(), terminal(), "execution-metadata") + + const calls = vi.mocked(fetchWithAgyCliTransport).mock.calls + const firstBody = JSON.parse(String(calls.at(-2)?.[1]?.body)) + const secondBody = JSON.parse(String(calls.at(-1)?.[1]?.body)) + expect(firstBody.request.labels).not.toHaveProperty("last_execution_id") + expect(firstBody.request.labels.last_step_index).toBe("1") + expect(secondBody.request.labels.last_execution_id).toMatch(/^[0-9a-f-]{36}$/) + expect(secondBody.request.labels.last_step_index).toBe("2") + }) + + it("does not rotate execution metadata on a tool-call turn", async () => { + await runStream( + fakeModel("antigravity-claude-opus-4-6-thinking"), + sseResponse([ + 'data: {"response":{"candidates":[{"content":{"parts":[{"functionCall":{"name":"read","args":{},"id":"c1"}}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}}\n\n', + ]), + "tool-execution-metadata", + ) + await runStream( + fakeModel("antigravity-claude-opus-4-6-thinking"), + sseResponse([ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}}\n\n', + ]), + "tool-execution-metadata", + ) + + const secondBody = JSON.parse(String(vi.mocked(fetchWithAgyCliTransport).mock.calls.at(-1)?.[1]?.body)) + expect(secondBody.request.labels).not.toHaveProperty("last_execution_id") + expect(secondBody.request.labels.last_step_index).toBe("1") + }) + + it("releases and cancels an open response body after STOP", async () => { + const open = openSseResponse( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2}}}\n\n', + ) + const { result } = await runStream( + fakeModel("antigravity-claude-opus-4-6-thinking"), + open.response, + "open-response-cleanup", + ) + + expect(result.stopReason).toBe("stop") + expect(open.wasCancelled()).toBe(true) + }) + + it("consumes GPT usage metadata sent after STOP", async () => { + const { result } = await runStream( + fakeModel("antigravity-gpt-oss-120b-medium"), + sseResponse([ + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}]}}\n\n', + 'data: {"response":{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":3,"totalTokenCount":13}}}\n\n', + ]), + "gpt-trailing-usage", + ) + + expect(result.stopReason).toBe("stop") + expect(result.usage.input).toBe(10) + expect(result.usage.output).toBe(3) + expect(result.usage.totalTokens).toBe(13) + }) + + it("finishes a GPT turn when trailing usage never arrives", async () => { + const stalled = stalledSseResponse( + 'data: {"response":{"candidates":[{"content":{"parts":[{"text":"answer"}]},"finishReason":"STOP"}]}}\n\n', + ) + const { result } = await runStream( + fakeModel("antigravity-gpt-oss-120b-medium"), + stalled.response, + "gpt-missing-trailing-usage", + stalled.abort, + ) + + expect(result.stopReason).toBe("stop") + expect(result.content).toEqual([{ type: "text", text: "answer" }]) + }) +}) + describe("parseGeminiSse", () => { it("parses and unwraps the Antigravity response envelope into chunks", async () => { // Antigravity wraps each chunk under a `response` key (MITM-verified). diff --git a/packages/pi/src/stream.ts b/packages/pi/src/stream.ts index 73f217bb..f6568be8 100644 --- a/packages/pi/src/stream.ts +++ b/packages/pi/src/stream.ts @@ -20,6 +20,7 @@ import { type SimpleStreamOptions, type StopReason, type TextContent, + type ThinkingContent, type ThinkingLevel, type ToolCall, } from "@earendil-works/pi-ai" @@ -29,8 +30,30 @@ import { getPackedRefresh } from "./credential-cache.ts" const STREAM_ACTION = "streamGenerateContent" const FALLBACK_SESSION_KEY = "__default__" +const TRAILING_USAGE_TIMEOUT_MS = 1_000 const requestSessions = new AgyRequestSessionStore("") +async function nextWithTimeout( + iterator: AsyncIterator, + timeoutMs: number, + onTimeout: () => void, +): Promise | undefined> { + let timer: ReturnType | undefined + try { + return await Promise.race([ + iterator.next(), + new Promise((resolve) => { + timer = setTimeout(() => { + onTimeout() + resolve(undefined) + }, timeoutMs) + }), + ]) + } finally { + if (timer) clearTimeout(timer) + } +} + function mapFinishReason(reason: string | null | undefined): StopReason { switch (reason) { case "STOP": @@ -111,7 +134,7 @@ export function convertGeminiToolCallPart( ): ToolCall | undefined { // Antigravity emits a batch signature on a preceding empty thought part; // native replay attaches it to the first function call in that batch. - if (part.thought && part.thoughtSignature) { + if (part.thought && !part.text && part.thoughtSignature) { state.pendingThoughtSignature = part.thoughtSignature } @@ -211,6 +234,7 @@ export function finalizePiAntigravityRequest( request, wireModel, scope.timestamp, + { stepCountMode: "contents" }, ) request.labels = metadata.labels request.sessionId = metadata.sessionId @@ -243,6 +267,8 @@ async function sendAntigravityRequest(options: { context: Context streamOptions?: SimpleStreamOptions accessToken: string + sessionKey: string + signal?: AbortSignal }): Promise { const resolved = resolvePiAntigravityModel(options.model, options.streamOptions?.reasoning) const wireModel = resolved.actualModel @@ -259,7 +285,10 @@ async function sendAntigravityRequest(options: { expires: Date.now() + 60_000, }) - const request = buildGeminiRequest(options.context) as unknown as Record + const request = buildGeminiRequest(options.context, { + provider: options.model.provider, + model: options.model.id, + }) as unknown as Record const generationConfig: Record = {} if (resolved.thinkingLevel) { @@ -283,9 +312,7 @@ async function sendAntigravityRequest(options: { request.generationConfig = generationConfig } - const requestScope = requestSessions.beginRequest( - getRequestSessionKey(options.context, options.streamOptions), - ) + const requestScope = requestSessions.beginRequest(options.sessionKey) const requestId = finalizePiAntigravityRequest(request, wireModel, requestScope) const envelope = { @@ -311,7 +338,7 @@ async function sendAntigravityRequest(options: { }, body: JSON.stringify(envelope), }, - { signal: options.streamOptions?.signal ?? null }, + { signal: options.signal ?? options.streamOptions?.signal ?? null }, ) } @@ -330,46 +357,142 @@ export function streamCortexKitAntigravity( const accessToken = options?.apiKey ?? "" if (!accessToken) throw new Error("Missing Antigravity OAuth access token") + const sessionKey = getRequestSessionKey(context, options) + const trailingUsageAbort = new AbortController() + const requestSignal = options?.signal + ? AbortSignal.any([options.signal, trailingUsageAbort.signal]) + : trailingUsageAbort.signal const response = await sendAntigravityRequest({ model, context, streamOptions: options, accessToken, + sessionKey, + signal: requestSignal, }) if (!response.ok) { throw new Error(`Antigravity request failed: HTTP ${response.status} ${await response.text()}`) } - const content = output.content as Array + const content = output.content as Array const toolCallState: GeminiToolCallState = {} let textIndex = -1 - let finished = false + let thinkingIndex = -1 + let terminalSeen = false + + const closeText = () => { + if (textIndex === -1) return + const block = content[textIndex] + if (block?.type === "text") { + stream.push({ type: "text_end", contentIndex: textIndex, content: block.text, partial: output }) + } + textIndex = -1 + } - for await (const chunk of parseGeminiSse(response)) { + const closeThinking = () => { + if (thinkingIndex === -1) return + const block = content[thinkingIndex] + if (block?.type === "thinking") { + stream.push({ + type: "thinking_end", + contentIndex: thinkingIndex, + content: block.thinking, + partial: output, + }) + } + thinkingIndex = -1 + } + + const chunkIterator = parseGeminiSse(response)[Symbol.asyncIterator]() + while (true) { + const next = terminalSeen + ? await nextWithTimeout( + chunkIterator, + TRAILING_USAGE_TIMEOUT_MS, + () => trailingUsageAbort.abort(), + ) + : await chunkIterator.next() + if (!next) break + if (next.done) break + + const chunk = next.value updateUsage(model, output, chunk.usageMetadata) + if (terminalSeen) { + if (chunk.usageMetadata) break + continue + } + const candidate = chunk.candidates?.[0] const parts = candidate?.content?.parts ?? [] for (const part of parts) { + if (!part.thought && !part.functionCall && !part.text && part.thoughtSignature) { + const block = textIndex === -1 ? undefined : content[textIndex] + if (block?.type === "text") { + block.textSignature = part.thoughtSignature + } else { + toolCallState.pendingThoughtSignature = part.thoughtSignature + } + continue + } + const toolCall = convertGeminiToolCallPart(part, toolCallState) if (toolCall) { + closeText() + closeThinking() content.push(toolCall) const idx = content.length - 1 - textIndex = -1 stream.push({ type: "toolcall_start", contentIndex: idx, partial: output }) stream.push({ type: "toolcall_end", contentIndex: idx, toolCall, partial: output }) output.stopReason = "toolUse" - } else if (typeof part.text === "string" && part.text.length > 0 && !part.thought) { + continue + } + + if (part.thought) { + if (typeof part.text !== "string" || part.text.length === 0) continue + closeText() + if (thinkingIndex === -1) { + content.push({ + type: "thinking", + thinking: "", + ...(part.thoughtSignature ? { thinkingSignature: part.thoughtSignature } : {}), + }) + thinkingIndex = content.length - 1 + stream.push({ type: "thinking_start", contentIndex: thinkingIndex, partial: output }) + } + const block = content[thinkingIndex] + if (block?.type === "thinking") { + block.thinking += part.text + if (part.thoughtSignature) block.thinkingSignature = part.thoughtSignature + stream.push({ + type: "thinking_delta", + contentIndex: thinkingIndex, + delta: part.text, + partial: output, + }) + } + continue + } + + if (typeof part.text === "string" && part.text.length > 0) { + closeThinking() if (textIndex === -1) { - content.push({ type: "text", text: "" }) + const textSignature = part.thoughtSignature ?? toolCallState.pendingThoughtSignature + toolCallState.pendingThoughtSignature = undefined + content.push({ + type: "text", + text: "", + ...(textSignature ? { textSignature } : {}), + }) textIndex = content.length - 1 stream.push({ type: "text_start", contentIndex: textIndex, partial: output }) } const block = content[textIndex] - if (block && block.type === "text") { + if (block?.type === "text") { block.text += part.text + if (part.thoughtSignature) block.textSignature = part.thoughtSignature stream.push({ type: "text_delta", contentIndex: textIndex, @@ -381,25 +504,23 @@ export function streamCortexKitAntigravity( } if (candidate?.finishReason) { - if (textIndex !== -1) { - const block = content[textIndex] - if (block && block.type === "text") { - stream.push({ type: "text_end", contentIndex: textIndex, content: block.text, partial: output }) - } - textIndex = -1 - } + closeText() + closeThinking() if (output.stopReason !== "toolUse") { output.stopReason = mapFinishReason(candidate.finishReason) } - // The AGY raw-socket transport may keep the response body open after - // the terminal chunk, which would hang the SSE reader. Stop consuming - // once a finishReason arrives so the turn terminates promptly. - finished = true - break + terminalSeen = true + const needsTrailingUsage = model.id.toLowerCase().includes("gpt-oss") && !chunk.usageMetadata + if (!needsTrailingUsage) break } } - if (finished) { + if (terminalSeen) { + try { + await chunkIterator.return?.(undefined) + } catch (error) { + if (!trailingUsageAbort.signal.aborted) throw error + } await response.body?.cancel().catch(() => {}) } @@ -410,6 +531,9 @@ export function streamCortexKitAntigravity( reason: output.stopReason as "stop" | "length" | "toolUse", message: output, }) + if (output.stopReason === "stop" || output.stopReason === "length") { + requestSessions.completeExecution(sessionKey) + } stream.end() } catch (error) { output.stopReason = options?.signal?.aborted ? "aborted" : "error"