From a50bee438376f36b8efa370ee82823c4abc0ec20 Mon Sep 17 00:00:00 2001
From: Muhammed Mukhthar CM <mukhtharcm@gmail.com>
Date: Wed, 14 Jan 2026 06:48:55 +0000
Subject: [PATCH] Agent: exclude google-antigravity from history downgrade hack

---
 ...ed-runner.sanitize-session-history.test.ts | 104 ++++++++++++++++++
 src/agents/pi-embedded-runner/google.ts       |   9 +-
 2 files changed, 110 insertions(+), 3 deletions(-)
 create mode 100644 src/agents/pi-embedded-runner.sanitize-session-history.test.ts

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..249d483d0
--- /dev/null
+++ b/src/agents/pi-embedded-runner.sanitize-session-history.test.ts
@@ -0,0 +1,104 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { sanitizeSessionHistory } from './pi-embedded-runner/google.js';
+import * as helpers from './pi-embedded-helpers.js';
+import type { AgentMessage } from '@mariozechner/pi-agent-core';
+import type { SessionManager } from '@mariozechner/pi-coding-agent';
+
+// 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;