fix(google): replay Vertex thought signatures - #1266
Conversation
📝 WalkthroughWalkthroughThe Google adapter now preserves and replays Gemini ChangesVertex thought-signature replay
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant GoogleAdapter
participant ReplayCache
participant VertexAPI
participant ToolContinuation
GoogleAdapter->>ReplayCache: Apply cached thoughtSignature values
GoogleAdapter->>VertexAPI: Send compiled request
VertexAPI-->>GoogleAdapter: Return functionCall and thoughtSignature
GoogleAdapter->>ReplayCache: Record thoughtSignature
ToolContinuation->>GoogleAdapter: Submit tool result
GoogleAdapter->>VertexAPI: Replay thoughtSignature in continuation
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
✅ Deterministic PR hygiene checks passed. |
ae28b69 to
c0ffaef
Compare
✅ READY
UI screenshot waived by a maintainer comment. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/adapters/google.ts`:
- Around line 540-545: Centralize replay cleanup in a helper that selects the
Antigravity or Vertex model/session from provider.googleMode and matches
signature or invalid-argument errors, then invoke it from formatErrorBody, the
SSE chunk.error branch, and the non-streaming raw.error branch. Remove the
duplicated inline selection/match logic while preserving existing behavior, and
add regression coverage verifying a non-streaming or HTTP 400 rejection clears
replay state before a clean continuation.
In `@tests/google-vertex-thought-signature.test.ts`:
- Around line 58-71: Update the vertexResponseBody fixture to accept a parameter
selecting thoughtSignature or thought_signature, then parameterize the
continuation tests over both response spellings. Parse each variant and assert
that the subsequent request replays the exact SIGNATURE value.
- Around line 110-121: Extend the test around applyAntigravityReplay to record a
Vertex signature using one project or location, then build the same continuation
with a different value and assert its functionCall has no thoughtSignature.
Ensure the assertions specifically verify isolation across Vertex configuration
identifiers, not only between Vertex and Antigravity namespaces.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 42235640-70b4-42c8-b807-b85f13d961bb
📒 Files selected for processing (9)
docs-site/src/content/docs/ja/reference/adapters.mddocs-site/src/content/docs/ko/reference/adapters.mddocs-site/src/content/docs/reference/adapters.mddocs-site/src/content/docs/ru/reference/adapters.mddocs-site/src/content/docs/zh-cn/reference/adapters.mdsrc/adapters/google-antigravity-replay.tssrc/adapters/google.tsstructure/04_transports-and-sidecars.mdtests/google-vertex-thought-signature.test.ts
| const replayModel = provider.googleMode === "cloud-code-assist" ? antigravityModel : vertexReplayModel; | ||
| const replaySession = provider.googleMode === "cloud-code-assist" ? antigravitySession : vertexReplaySession; | ||
| if ((provider.googleMode === "cloud-code-assist" || provider.googleMode === "vertex") | ||
| && replayModel && replaySession | ||
| && /signature|invalid_argument|invalid argument/i.test(err?.message ?? "")) { | ||
| clearAntigravityReplay(antigravityModel, antigravitySession); | ||
| clearAntigravityReplay(replayModel, replaySession); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Clear replay state for all signature rejection paths.
Lines 540-545 clear state only for an in-band SSE chunk.error. A non-streaming raw.error returns at line 778 without clearing state. An HTTP 400 is formatted through formatErrorBody at lines 327-329 without clearing state.
The stale signature is replayed on the next continuation until the TTL expires. Centralize the model/session selection and signature-error match in a helper. Call it from formatErrorBody, the SSE error branch, and the non-streaming raw.error branch. Add regression coverage for a non-streaming or HTTP 400 rejection followed by a clean continuation.
Proposed fix
+ const clearReplayOnSignatureError = (message: string): void => {
+ const replayModel = provider.googleMode === "cloud-code-assist"
+ ? antigravityModel
+ : vertexReplayModel;
+ const replaySession = provider.googleMode === "cloud-code-assist"
+ ? antigravitySession
+ : vertexReplaySession;
+ if ((provider.googleMode === "cloud-code-assist" || provider.googleMode === "vertex")
+ && replayModel && replaySession
+ && /signature|invalid_argument|invalid argument/i.test(message)) {
+ clearAntigravityReplay(replayModel, replaySession);
+ }
+ };
+
...(provider.googleMode === "vertex" || provider.googleMode === "cloud-code-assist"
? {
fetchResponse: ...,
- formatErrorBody: (status, _headers, payloadText) =>
- ...,
+ formatErrorBody: (status, _headers, payloadText) => {
+ clearReplayOnSignatureError(payloadText);
+ return ...;
+ },
}
: {}),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/adapters/google.ts` around lines 540 - 545, Centralize replay cleanup in
a helper that selects the Antigravity or Vertex model/session from
provider.googleMode and matches signature or invalid-argument errors, then
invoke it from formatErrorBody, the SSE chunk.error branch, and the
non-streaming raw.error branch. Remove the duplicated inline selection/match
logic while preserving existing behavior, and add regression coverage verifying
a non-streaming or HTTP 400 rejection clears replay state before a clean
continuation.
| 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 }, | ||
| }; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
Test both Vertex thought-signature field spellings.
vertexResponseBody() only emits thoughtSignature. The stated replay contract also accepts thought_signature. A regression in the snake_case response parser would pass both continuation tests and fail against affected Vertex responses.
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
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/google-vertex-thought-signature.test.ts` around lines 58 - 71, Update
the vertexResponseBody fixture to accept a parameter selecting thoughtSignature
or thought_signature, then parameterize the continuation tests over both
response spellings. Parse each variant and assert that the subsequent request
replays the exact SIGNATURE value.
| 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(); |
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 project or location stay isolated. If either identifier is removed from the replay key, this suite still passes and can replay provider-private state across GCP tenants or regions.
Record a signature with one Vertex project or location. Build the same continuation with a different value. Assert that its functionCall has no thoughtSignature.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/google-vertex-thought-signature.test.ts` around lines 110 - 121, Extend
the test around applyAntigravityReplay to record a Vertex signature using one
project or location, then build the same continuation with a different value and
assert its functionCall has no thoughtSignature. Ensure the assertions
specifically verify isolation across Vertex configuration identifiers, not only
between Vertex and Antigravity namespaces.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c0ffaef643
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| applyAntigravityReplay( | ||
| vertexReplayModel, | ||
| vertexReplaySession, | ||
| (compiled.body as { contents: unknown[] }).contents, | ||
| ); |
There was a problem hiding this comment.
When a Vertex conversation invokes the same compiled function with identical arguments more than once (for example, repeating pwd or a status check), observeAntigravityReplay stores 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 historical functionCall, 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 👍 / 👎.
|
Not gui |
3 participants
Closes #1254
Summary
thoughtSignaturevalues across Codex tool-result continuations by reusing the existing bounded TTL/LRU Google replay storeRoot cause
The Google adapter wired signature observation and replay only for
googleMode: "cloud-code-assist". Vertex used the same Gemini function-call wire but bypassed the replay store, so the first call succeeded and the next tool-result request omitted the required opaque signature and received HTTP 400.Verification
bun test tests/google-vertex-thought-signature.test.ts tests/google-antigravity-replay.test.ts tests/google-antigravity-wire.test.ts tests/google-adapter.test.ts— 94 pass, 0 failbun run typecheckbun run privacy:scanbun run --cwd docs-site build— 221 pagesgit diff --checkbun run testwas also exercised under a two-core affinity. The changed Google suites passed; the host run exposed an unrelated existingcodex-shimservice-token path isolation failure. The initially missing GUI dependency was installed and its affected provider-rail suite then passed. Clean GitHub CI remains the authoritative full-suite gate.Checklist
Summary by CodeRabbit
New Features
Documentation