fix(google): replay Vertex thought signatures #1266
There are no files selected for viewing
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
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
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
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
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
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
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,123 @@ | ||
| import { beforeEach, describe, expect, test } from "bun:test"; | ||
| import { createGoogleAdapter as createGoogleAdapterProduction } from "../src/adapters/google"; | ||
| import { | ||
| __resetAntigravityReplayCache, | ||
| applyAntigravityReplay, | ||
| } from "../src/adapters/google-antigravity-replay"; | ||
| import { antigravitySessionId } from "../src/adapters/google-antigravity-wire"; | ||
| import type { AdapterEvent, OcxParsedRequest, OcxProviderConfig } from "../src/types"; | ||
| import { withTestTranslatorBudget } from "./helpers/translator-budget"; | ||
|
|
||
| const createGoogleAdapter = (...args: Parameters<typeof createGoogleAdapterProduction>) => | ||
| withTestTranslatorBudget(createGoogleAdapterProduction(...args)); | ||
|
|
||
| const SIGNATURE = "CiQAx-vertex-thought-signature-0123456789abcdef"; | ||
| const MODEL = "gemini-3.6-flash"; | ||
|
|
||
| const provider = { | ||
| adapter: "google", | ||
| googleMode: "vertex", | ||
| baseUrl: "https://aiplatform.googleapis.com", | ||
| apiKey: "vertex-test-key", | ||
| } as OcxProviderConfig; | ||
|
|
||
| function request(messages: OcxParsedRequest["context"]["messages"], stream: boolean): OcxParsedRequest { | ||
| return { | ||
| modelId: MODEL, | ||
| stream, | ||
| context: { | ||
| messages, | ||
| systemPrompt: [], | ||
| tools: [{ name: "shell_command", description: "run a command", parameters: { type: "object" } }], | ||
| }, | ||
| options: {}, | ||
| } as unknown as OcxParsedRequest; | ||
| } | ||
|
|
||
| const firstTurn = (stream: boolean) => request([{ role: "user", content: "run pwd" }], stream); | ||
|
|
||
| const continuation = () => request([ | ||
| { role: "user", content: "run pwd" }, | ||
| { | ||
| role: "assistant", | ||
| content: [{ | ||
| type: "toolCall", | ||
| id: "call_shell_1", | ||
| name: "shell_command", | ||
| arguments: { command: "pwd" }, | ||
| }], | ||
| }, | ||
| { | ||
| role: "toolResult", | ||
| toolCallId: "call_shell_1", | ||
| toolName: "shell_command", | ||
| content: "/workspace", | ||
| }, | ||
| ], false); | ||
|
|
||
| function vertexResponseBody(): Record<string, unknown> { | ||
| return { | ||
| candidates: [{ | ||
| content: { | ||
| role: "model", | ||
| parts: [{ | ||
| functionCall: { name: "shell_command", args: { command: "pwd" } }, | ||
| thoughtSignature: SIGNATURE, | ||
| }], | ||
| }, | ||
| finishReason: "STOP", | ||
| }], | ||
| usageMetadata: { promptTokenCount: 5, candidatesTokenCount: 2 }, | ||
| }; | ||
|
Comment on lines
+58
to
+71
Contributor
There was a problem hiding this comment. 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win Test both Vertex thought-signature field spellings.
Parameterize this fixture for both field names. Parse each response form, then assert that the next request replays the exact signature value. 🤖 Prompt for AI Agents |
||
| } | ||
|
|
||
| function replayedFunctionCall(body: string): Record<string, unknown> { | ||
| const parsed = JSON.parse(body) as { contents: Array<{ role?: string; parts?: Record<string, unknown>[] }> }; | ||
| const model = parsed.contents.find(content => content.role === "model"); | ||
| const part = model?.parts?.find(candidate => "functionCall" in candidate); | ||
| if (!part) throw new Error("compiled Vertex request omitted the replayed functionCall"); | ||
| return part; | ||
| } | ||
|
|
||
| describe("Vertex thought-signature continuation (#1254)", () => { | ||
| beforeEach(() => __resetAntigravityReplayCache()); | ||
|
|
||
| test("streaming functionCall signature is replayed on the next tool-result turn", async () => { | ||
| const firstAdapter = createGoogleAdapter(provider); | ||
| await firstAdapter.buildRequest(firstTurn(true)); | ||
| const response = new Response(`data: ${JSON.stringify(vertexResponseBody())}\n\n`, { | ||
| headers: { "content-type": "text/event-stream" }, | ||
| }); | ||
| const events: AdapterEvent[] = []; | ||
| for await (const event of firstAdapter.parseStream(response)) events.push(event); | ||
| expect(events.some(event => event.type === "tool_call_start")).toBe(true); | ||
| expect(events.at(-1)?.type).toBe("done"); | ||
|
|
||
| const followup = await createGoogleAdapter(provider).buildRequest(continuation()); | ||
| expect(replayedFunctionCall(followup.body as string).thoughtSignature).toBe(SIGNATURE); | ||
| }); | ||
|
|
||
| test("non-streaming functionCall signature is replayed unchanged", async () => { | ||
| const firstAdapter = createGoogleAdapter(provider); | ||
| await firstAdapter.buildRequest(firstTurn(false)); | ||
| const events = await firstAdapter.parseResponse!(new Response(JSON.stringify(vertexResponseBody()))); | ||
| expect(events.some(event => event.type === "tool_call_start")).toBe(true); | ||
|
|
||
| const followup = await createGoogleAdapter(provider).buildRequest(continuation()); | ||
| expect(replayedFunctionCall(followup.body as string).thoughtSignature).toBe(SIGNATURE); | ||
| }); | ||
|
|
||
| test("Vertex signatures cannot enter the Antigravity replay namespace", async () => { | ||
| const first = firstTurn(false); | ||
| const adapter = createGoogleAdapter(provider); | ||
| await adapter.buildRequest(first); | ||
| await adapter.parseResponse!(new Response(JSON.stringify(vertexResponseBody()))); | ||
|
|
||
| const contents = [{ | ||
| role: "model", | ||
| parts: [{ functionCall: { name: "shell_command", args: { command: "pwd" } } }], | ||
| }]; | ||
| applyAntigravityReplay(MODEL, antigravitySessionId(first), contents); | ||
| expect((contents[0].parts[0] as { thoughtSignature?: string }).thoughtSignature).toBeUndefined(); | ||
|
Comment on lines
+110
to
+121
Contributor
There was a problem hiding this comment. 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick win Test Vertex project and location isolation. This test proves that a Vertex entry does not enter the Antigravity namespace. It does not prove that two Vertex configurations with the same conversation but different Record a signature with one Vertex 🤖 Prompt for AI Agents |
||
| }); | ||
| }); | ||
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When a Vertex conversation invokes the same compiled function with identical arguments more than once (for example, repeating
pwdor a status check),observeAntigravityReplaystores both occurrences under the same name-plus-arguments key, so the newer signature overwrites the older one. This newly added call then applies that single newest signature to every matching historicalfunctionCall, rather than replaying each model turn's exact opaque value, causing a later tool-result continuation to be rejected by Vertex. Store signatures per occurrence/order or stable call identity, and add a regression test with two identical calls followed by a third turn.AGENTS.md reference: src/AGENTS.md:L19-L19
Useful? React with 👍 / 👎.