diff --git a/src/agents/pi-embedded-helpers/google.ts b/src/agents/pi-embedded-helpers/google.ts index fe3290412..e3c6200bd 100644 --- a/src/agents/pi-embedded-helpers/google.ts +++ b/src/agents/pi-embedded-helpers/google.ts @@ -8,6 +8,12 @@ export function isGoogleModelApi(api?: string | null): boolean { ); } +export function isAntigravityClaude(api?: string | null, modelId?: string): boolean { + if (api !== "google-antigravity") return false; + return modelId?.toLowerCase().includes("claude") ?? false; +} + + export { sanitizeGoogleTurnOrdering }; /** diff --git a/src/agents/pi-embedded-helpers/images.ts b/src/agents/pi-embedded-helpers/images.ts index 56130abec..bd390b842 100644 --- a/src/agents/pi-embedded-helpers/images.ts +++ b/src/agents/pi-embedded-helpers/images.ts @@ -30,7 +30,7 @@ function isEmptyAssistantErrorMessage( export async function sanitizeSessionMessagesImages( messages: AgentMessage[], label: string, - options?: { sanitizeToolCallIds?: boolean; enforceToolCallLast?: boolean }, + options?: { sanitizeToolCallIds?: boolean; enforceToolCallLast?: boolean; preserveSignatures?: boolean }, ): Promise { // We sanitize historical session messages because Anthropic can reject a request // if the transcript contains oversized base64 images (see MAX_IMAGE_DIMENSION_PX). @@ -76,7 +76,10 @@ export async function sanitizeSessionMessagesImages( } const content = assistantMsg.content; if (Array.isArray(content)) { - const strippedContent = stripThoughtSignatures(content); + const strippedContent = options?.preserveSignatures + ? content // Keep signatures for Antigravity Claude + : stripThoughtSignatures(content); // Strip for Gemini + const filteredContent = strippedContent.filter((block) => { if (!block || typeof block !== "object") return true; const rec = block as { type?: unknown; text?: unknown }; diff --git a/src/agents/pi-embedded-runner.google-sanitize-thinking.test.ts b/src/agents/pi-embedded-runner.google-sanitize-thinking.test.ts index 0f4aa90ff..900679b9c 100644 --- a/src/agents/pi-embedded-runner.google-sanitize-thinking.test.ts +++ b/src/agents/pi-embedded-runner.google-sanitize-thinking.test.ts @@ -59,6 +59,34 @@ describe("sanitizeSessionHistory (google thinking)", () => { expect(assistant.content?.[0]?.thinkingSignature).toBe("sig"); }); + it("keeps unsigned thinking blocks for Antigravity Claude", async () => { + const sessionManager = SessionManager.inMemory(); + const input = [ + { + role: "user", + content: "hi", + }, + { + role: "assistant", + content: [{ type: "thinking", thinking: "reasoning" }], + }, + ] satisfies AgentMessage[]; + + const out = await sanitizeSessionHistory({ + messages: input, + modelApi: "google-antigravity", + modelId: "anthropic/claude-3.5-sonnet", + sessionManager, + sessionId: "session:antigravity-claude", + }); + + const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant") as { + content?: Array<{ type?: string; thinking?: string }>; + }; + expect(assistant.content?.map((block) => block.type)).toEqual(["thinking"]); + expect(assistant.content?.[0]?.thinking).toBe("reasoning"); + }); + it("preserves order when downgrading mixed assistant content", async () => { const sessionManager = SessionManager.inMemory(); const input = [ diff --git a/src/agents/pi-embedded-runner/compact.ts b/src/agents/pi-embedded-runner/compact.ts index 8c88c1945..de871ebee 100644 --- a/src/agents/pi-embedded-runner/compact.ts +++ b/src/agents/pi-embedded-runner/compact.ts @@ -302,6 +302,7 @@ export async function compactEmbeddedPiSession(params: { const prior = await sanitizeSessionHistory({ messages: session.messages, modelApi: model.api, + modelId, sessionManager, sessionId: params.sessionId, }); diff --git a/src/agents/pi-embedded-runner/google.ts b/src/agents/pi-embedded-runner/google.ts index 4a3a7c3f6..be6e66a90 100644 --- a/src/agents/pi-embedded-runner/google.ts +++ b/src/agents/pi-embedded-runner/google.ts @@ -13,6 +13,7 @@ import { import { sanitizeToolUseResultPairing } from "../session-transcript-repair.js"; import { log } from "./logger.js"; import { describeUnknownError } from "./utils.js"; +import { isAntigravityClaude } from "../pi-embedded-helpers/google.js"; const GOOGLE_TURN_ORDERING_CUSTOM_TYPE = "google-turn-ordering-bootstrap"; const GOOGLE_SCHEMA_UNSUPPORTED_KEYWORDS = new Set([ @@ -152,19 +153,23 @@ export function applyGoogleTurnOrderingFix(params: { export async function sanitizeSessionHistory(params: { messages: AgentMessage[]; modelApi?: string | null; + modelId?: string; sessionManager: SessionManager; sessionId: string; }): Promise { const sanitizedImages = await sanitizeSessionMessagesImages(params.messages, "session:history", { sanitizeToolCallIds: shouldSanitizeToolCallIds(params.modelApi), enforceToolCallLast: params.modelApi === "anthropic-messages", + preserveSignatures: params.modelApi === "google-antigravity" && isAntigravityClaude(params.modelId), }); const repairedTools = sanitizeToolUseResultPairing(sanitizedImages); + const shouldDowngradeGemini = + isGoogleModelApi(params.modelApi) && !isAntigravityClaude(params.modelId); // Gemini rejects unsigned thinking blocks; downgrade them before send to avoid INVALID_ARGUMENT. - const downgradedThinking = isGoogleModelApi(params.modelApi) + const downgradedThinking = shouldDowngradeGemini ? downgradeGeminiThinkingBlocks(repairedTools) : repairedTools; - const downgraded = isGoogleModelApi(params.modelApi) + const downgraded = shouldDowngradeGemini ? downgradeGeminiHistory(downgradedThinking) : downgradedThinking; diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 93e881c59..3d5313bdc 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -303,6 +303,7 @@ export async function runEmbeddedAttempt( const prior = await sanitizeSessionHistory({ messages: activeSession.messages, modelApi: params.model.api, + modelId: params.modelId, sessionManager, sessionId: params.sessionId, });