From cd409e56677ff377d73fa4152fc1f4f92584438a Mon Sep 17 00:00:00 2001 From: Muhammed Mukhthar CM <56378562+mukhtharcm@users.noreply.github.com> Date: Fri, 16 Jan 2026 13:44:56 +0530 Subject: [PATCH] fix: exclude google-antigravity from history downgrade hack (#894) * Agent: exclude google-antigravity from history downgrade hack * Lint: fix formatting in test * Lint: formatting and unused vars in test * fix: preserve google-antigravity tool calls (#894) (thanks @mukhtharcm) --------- Co-authored-by: Peter Steinberger --- CHANGELOG.md | 1 + ...ed-runner.sanitize-session-history.test.ts | 109 ++++++++++++++++++ src/agents/pi-embedded-runner/google.ts | 9 +- 3 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 src/agents/pi-embedded-runner.sanitize-session-history.test.ts diff --git a/CHANGELOG.md b/CHANGELOG.md index 8f588aaf3..a48109545 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ - Docs: clarify multi-gateway rescue bot guidance. (#969) — thanks @bjesuiter. - Agents: add Current Date & Time system prompt section with configurable time format (auto/12/24). - Agents: avoid false positives when logging unsupported Google tool schema keywords. +- Agents: skip Gemini history downgrades for google-antigravity to preserve tool calls. (#894) — thanks @mukhtharcm. - Status: restore usage summary line for current provider when no OAuth profiles exist. - Tools: normalize Slack/Discord message timestamps with `timestampMs`/`timestampUtc` while keeping raw provider fields. - macOS: add `system.which` for prompt-free remote skill discovery (with gateway fallback to `system.run`). diff --git a/src/agents/pi-embedded-runner.sanitize-session-history.test.ts b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts new file mode 100644 index 000000000..f8fe98eee --- /dev/null +++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts @@ -0,0 +1,109 @@ +import type { AgentMessage } from "@mariozechner/pi-agent-core"; +import type { SessionManager } from "@mariozechner/pi-coding-agent"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import * as helpers from "./pi-embedded-helpers.js"; +import { sanitizeSessionHistory } from "./pi-embedded-runner/google.js"; + +// Mock dependencies +vi.mock("./pi-embedded-helpers.js", async () => { + const actual = await vi.importActual("./pi-embedded-helpers.js"); + return { + ...actual, + isGoogleModelApi: vi.fn(), + downgradeGeminiHistory: vi.fn(), + sanitizeSessionMessagesImages: vi + .fn() + .mockImplementation(async (msgs) => msgs), + }; +}); + +// We don't mock session-transcript-repair.js as it is a pure function and complicates mocking. +// We rely on the real implementation which should pass through our simple messages. + +describe("sanitizeSessionHistory", () => { + const mockSessionManager = { + getEntries: vi.fn().mockReturnValue([]), + appendCustomEntry: vi.fn(), + } as unknown as SessionManager; + + const mockMessages: AgentMessage[] = [{ role: "user", content: "hello" }]; + + beforeEach(() => { + vi.resetAllMocks(); + vi.mocked(helpers.sanitizeSessionMessagesImages).mockImplementation( + async (msgs) => msgs, + ); + // Default mock implementation + vi.mocked(helpers.downgradeGeminiHistory).mockImplementation((msgs) => { + if (!msgs) return []; + return [...msgs, { role: "system", content: "downgraded" }]; + }); + }); + + it("should downgrade history for Google models if provider is not google-antigravity", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(true); + + const result = await sanitizeSessionHistory({ + messages: mockMessages, + modelApi: "google-gemini", + provider: "google-vertex", + sessionManager: mockSessionManager, + sessionId: "test-session", + }); + + expect(helpers.isGoogleModelApi).toHaveBeenCalledWith("google-gemini"); + expect(helpers.downgradeGeminiHistory).toHaveBeenCalled(); + // Check if the result contains the downgraded message + expect(result).toContainEqual({ role: "system", content: "downgraded" }); + }); + + it("should NOT downgrade history for google-antigravity provider", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(true); + + const result = await sanitizeSessionHistory({ + messages: mockMessages, + modelApi: "google-gemini", + provider: "google-antigravity", + sessionManager: mockSessionManager, + sessionId: "test-session", + }); + + expect(helpers.isGoogleModelApi).toHaveBeenCalledWith("google-gemini"); + expect(helpers.downgradeGeminiHistory).not.toHaveBeenCalled(); + // Result should not contain the downgraded message + expect(result).not.toContainEqual({ + role: "system", + content: "downgraded", + }); + }); + + it("should NOT downgrade history for non-Google models", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false); + + const _result = await sanitizeSessionHistory({ + messages: mockMessages, + modelApi: "anthropic-messages", + provider: "anthropic", + sessionManager: mockSessionManager, + sessionId: "test-session", + }); + + expect(helpers.isGoogleModelApi).toHaveBeenCalledWith("anthropic-messages"); + expect(helpers.downgradeGeminiHistory).not.toHaveBeenCalled(); + }); + + it("should downgrade history if provider is undefined but model is Google", async () => { + vi.mocked(helpers.isGoogleModelApi).mockReturnValue(true); + + const _result = await sanitizeSessionHistory({ + messages: mockMessages, + modelApi: "google-gemini", + provider: undefined, + sessionManager: mockSessionManager, + sessionId: "test-session", + }); + + expect(helpers.isGoogleModelApi).toHaveBeenCalledWith("google-gemini"); + expect(helpers.downgradeGeminiHistory).toHaveBeenCalled(); + }); +}); diff --git a/src/agents/pi-embedded-runner/google.ts b/src/agents/pi-embedded-runner/google.ts index 31ee19a59..702a9c3fa 100644 --- a/src/agents/pi-embedded-runner/google.ts +++ b/src/agents/pi-embedded-runner/google.ts @@ -204,12 +204,15 @@ export async function sanitizeSessionHistory(params: { : undefined, }); const repairedTools = sanitizeToolUseResultPairing(sanitizedImages); - const shouldDowngradeGemini = isGeminiLike && !isAntigravityClaudeModel; + const isAntigravityProvider = + provider === "google-antigravity" || params.modelApi === "google-antigravity"; + const shouldDowngradeThinking = isGeminiLike && !isAntigravityClaudeModel; // Gemini rejects unsigned thinking blocks; downgrade them before send to avoid INVALID_ARGUMENT. - const downgradedThinking = shouldDowngradeGemini + const downgradedThinking = shouldDowngradeThinking ? downgradeGeminiThinkingBlocks(repairedTools) : repairedTools; - const downgraded = shouldDowngradeGemini + const shouldDowngradeHistory = shouldDowngradeThinking && !isAntigravityProvider; + const downgraded = shouldDowngradeHistory ? downgradeGeminiHistory(downgradedThinking) : downgradedThinking;