fix: drop obsolete pi-mono workarounds

This commit is contained in:
Peter Steinberger
2026-01-21 19:48:16 +00:00
parent 9f999f6554
commit 41c9c214fc
7 changed files with 57 additions and 354 deletions

View File

@@ -1,71 +0,0 @@
import { describe, expect, it } from "vitest";
import { downgradeGeminiHistory } from "./pi-embedded-helpers.js";
describe("downgradeGeminiHistory", () => {
it("drops unsigned tool calls and matching tool results", () => {
const input = [
{
role: "assistant",
content: [
{ type: "text", text: "hello" },
{ type: "toolCall", id: "call_1", name: "read", arguments: { path: "/tmp" } },
],
},
{
role: "toolResult",
toolCallId: "call_1",
content: [{ type: "text", text: "ok" }],
},
{ role: "user", content: "next" },
];
expect(downgradeGeminiHistory(input)).toEqual([
{
role: "assistant",
content: [{ type: "text", text: "hello" }],
},
{ role: "user", content: "next" },
]);
});
it("keeps signed tool calls and results", () => {
const input = [
{
role: "assistant",
content: [
{
type: "toolCall",
id: "call_2",
name: "read",
arguments: { path: "/tmp" },
thought_signature: "sig_123",
},
],
},
{
role: "toolResult",
toolCallId: "call_2",
content: [{ type: "text", text: "ok" }],
},
];
expect(downgradeGeminiHistory(input)).toEqual(input);
});
it("drops assistant messages that only contain unsigned tool calls", () => {
const input = [
{
role: "assistant",
content: [{ type: "toolCall", id: "call_3", name: "read", arguments: {} }],
},
{
role: "toolResult",
toolCallId: "call_3",
content: [{ type: "text", text: "ok" }],
},
{ role: "user", content: "after" },
];
expect(downgradeGeminiHistory(input)).toEqual([{ role: "user", content: "after" }]);
});
});

View File

@@ -30,12 +30,7 @@ export {
isTimeoutErrorMessage,
parseImageDimensionError,
} from "./pi-embedded-helpers/errors.js";
export {
downgradeGeminiHistory,
downgradeGeminiThinkingBlocks,
isGoogleModelApi,
sanitizeGoogleTurnOrdering,
} from "./pi-embedded-helpers/google.js";
export { isGoogleModelApi, sanitizeGoogleTurnOrdering } from "./pi-embedded-helpers/google.js";
export {
isEmptyAssistantMessageContent,
sanitizeSessionMessagesImages,

View File

@@ -1,5 +1,3 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { sanitizeGoogleTurnOrdering } from "./bootstrap.js";
export function isGoogleModelApi(api?: string | null): boolean {
@@ -14,145 +12,3 @@ export function isAntigravityClaude(api?: string | null, modelId?: string): bool
}
export { sanitizeGoogleTurnOrdering };
/**
* Drops tool calls that are missing `thought_signature` (required by Gemini)
* to prevent 400 INVALID_ARGUMENT errors. Matching tool results are dropped
* so they don't become orphaned in the transcript.
*/
type GeminiToolCallBlock = {
type?: unknown;
thought_signature?: unknown;
thoughtSignature?: unknown;
id?: unknown;
toolCallId?: unknown;
name?: unknown;
toolName?: unknown;
arguments?: unknown;
input?: unknown;
};
type GeminiThinkingBlock = {
type?: unknown;
thinking?: unknown;
thinkingSignature?: unknown;
};
export function downgradeGeminiThinkingBlocks(messages: AgentMessage[]): AgentMessage[] {
const out: AgentMessage[] = [];
for (const msg of messages) {
if (!msg || typeof msg !== "object") {
out.push(msg);
continue;
}
const role = (msg as { role?: unknown }).role;
if (role !== "assistant") {
out.push(msg);
continue;
}
const assistantMsg = msg as Extract<AgentMessage, { role: "assistant" }>;
if (!Array.isArray(assistantMsg.content)) {
out.push(msg);
continue;
}
// Gemini rejects thinking blocks that lack a signature; downgrade to text for safety.
let hasDowngraded = false;
type AssistantContentBlock = (typeof assistantMsg.content)[number];
const nextContent = assistantMsg.content.flatMap((block): AssistantContentBlock[] => {
if (!block || typeof block !== "object") return [block as AssistantContentBlock];
const record = block as GeminiThinkingBlock;
if (record.type !== "thinking") return [block];
const thinkingSig =
typeof record.thinkingSignature === "string" ? record.thinkingSignature.trim() : "";
if (thinkingSig.length > 0) return [block];
const thinking = typeof record.thinking === "string" ? record.thinking : "";
const trimmed = thinking.trim();
hasDowngraded = true;
if (!trimmed) return [];
return [{ type: "text" as const, text: thinking }];
});
if (!hasDowngraded) {
out.push(msg);
continue;
}
if (nextContent.length === 0) {
continue;
}
out.push({ ...assistantMsg, content: nextContent } as AgentMessage);
}
return out;
}
export function downgradeGeminiHistory(messages: AgentMessage[]): AgentMessage[] {
const droppedToolCallIds = new Set<string>();
const out: AgentMessage[] = [];
const resolveToolResultId = (
msg: Extract<AgentMessage, { role: "toolResult" }>,
): string | undefined => {
const toolCallId = (msg as { toolCallId?: unknown }).toolCallId;
if (typeof toolCallId === "string" && toolCallId) return toolCallId;
const toolUseId = (msg as { toolUseId?: unknown }).toolUseId;
if (typeof toolUseId === "string" && toolUseId) return toolUseId;
return undefined;
};
for (const msg of messages) {
if (!msg || typeof msg !== "object") {
out.push(msg);
continue;
}
const role = (msg as { role?: unknown }).role;
if (role === "assistant") {
const assistantMsg = msg as Extract<AgentMessage, { role: "assistant" }>;
if (!Array.isArray(assistantMsg.content)) {
out.push(msg);
continue;
}
let dropped = false;
const nextContent = assistantMsg.content.filter((block) => {
if (!block || typeof block !== "object") return true;
const blockRecord = block as GeminiToolCallBlock;
const type = blockRecord.type;
if (type === "toolCall" || type === "functionCall" || type === "toolUse") {
const signature = blockRecord.thought_signature ?? blockRecord.thoughtSignature;
const hasSignature = Boolean(signature);
if (!hasSignature) {
const id =
typeof blockRecord.id === "string"
? blockRecord.id
: typeof blockRecord.toolCallId === "string"
? blockRecord.toolCallId
: undefined;
if (id) droppedToolCallIds.add(id);
dropped = true;
return false;
}
}
return true;
});
if (dropped && nextContent.length === 0) {
continue;
}
out.push(dropped ? ({ ...assistantMsg, content: nextContent } as AgentMessage) : msg);
continue;
}
if (role === "toolResult") {
const toolMsg = msg as Extract<AgentMessage, { role: "toolResult" }>;
const toolResultId = resolveToolResultId(toolMsg);
if (toolResultId && droppedToolCallIds.has(toolResultId)) {
continue;
}
}
out.push(msg);
}
return out;
}

View File

@@ -4,7 +4,7 @@ import { describe, expect, it } from "vitest";
import { sanitizeSessionHistory } from "./pi-embedded-runner/google.js";
describe("sanitizeSessionHistory (google thinking)", () => {
it("downgrades thinking blocks without signatures for Google models", async () => {
it("keeps thinking blocks without signatures for Google models", async () => {
const sessionManager = SessionManager.inMemory();
const input = [
{
@@ -25,10 +25,10 @@ describe("sanitizeSessionHistory (google thinking)", () => {
});
const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant") as {
content?: Array<{ type?: string; text?: string }>;
content?: Array<{ type?: string; thinking?: string }>;
};
expect(assistant.content?.map((block) => block.type)).toEqual(["text"]);
expect(assistant.content?.[0]?.text).toBe("reasoning");
expect(assistant.content?.map((block) => block.type)).toEqual(["thinking"]);
expect(assistant.content?.[0]?.thinking).toBe("reasoning");
});
it("keeps thinking blocks with signatures for Google models", async () => {
@@ -59,7 +59,7 @@ describe("sanitizeSessionHistory (google thinking)", () => {
expect(assistant.content?.[0]?.thinkingSignature).toBe("sig");
});
it("downgrades thinking blocks with Anthropic-style signatures for Google models", async () => {
it("keeps thinking blocks with Anthropic-style signatures for Google models", async () => {
const sessionManager = SessionManager.inMemory();
const input = [
{
@@ -80,10 +80,10 @@ describe("sanitizeSessionHistory (google thinking)", () => {
});
const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant") as {
content?: Array<{ type?: string; text?: string }>;
content?: Array<{ type?: string; thinking?: string }>;
};
expect(assistant.content?.map((block) => block.type)).toEqual(["text"]);
expect(assistant.content?.[0]?.text).toBe("reasoning");
expect(assistant.content?.map((block) => block.type)).toEqual(["thinking"]);
expect(assistant.content?.[0]?.thinking).toBe("reasoning");
});
it("keeps unsigned thinking blocks for Antigravity Claude", async () => {
@@ -114,7 +114,7 @@ describe("sanitizeSessionHistory (google thinking)", () => {
expect(assistant.content?.[0]?.thinking).toBe("reasoning");
});
it("preserves order when downgrading mixed assistant content", async () => {
it("preserves order for mixed assistant content", async () => {
const sessionManager = SessionManager.inMemory();
const input = [
{
@@ -139,10 +139,10 @@ describe("sanitizeSessionHistory (google thinking)", () => {
});
const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant") as {
content?: Array<{ type?: string; text?: string }>;
content?: Array<{ type?: string; text?: string; thinking?: string }>;
};
expect(assistant.content?.map((block) => block.type)).toEqual(["text", "text", "text"]);
expect(assistant.content?.[1]?.text).toBe("internal note");
expect(assistant.content?.map((block) => block.type)).toEqual(["text", "thinking", "text"]);
expect(assistant.content?.[1]?.thinking).toBe("internal note");
});
it("strips non-base64 thought signatures for OpenRouter Gemini", async () => {
@@ -185,11 +185,22 @@ describe("sanitizeSessionHistory (google thinking)", () => {
});
const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant") as {
content?: Array<{ type?: string; thought_signature?: string; thoughtSignature?: string }>;
content?: Array<{
type?: string;
thought_signature?: string;
thoughtSignature?: string;
thinking?: string;
}>;
};
expect(assistant.content).toEqual([
{ type: "text", text: "hello" },
{ type: "text", text: "ok" },
{ type: "thinking", thinking: "ok", thought_signature: "c2ln" },
{
type: "toolCall",
id: "call_1",
name: "read",
arguments: { path: "/tmp/foo" },
},
{
type: "toolCall",
id: "call_2",
@@ -200,7 +211,7 @@ describe("sanitizeSessionHistory (google thinking)", () => {
]);
});
it("downgrades only unsigned thinking blocks when mixed with signed ones", async () => {
it("keeps mixed signed/unsigned thinking blocks for Google models", async () => {
const sessionManager = SessionManager.inMemory();
const input = [
{
@@ -224,14 +235,14 @@ describe("sanitizeSessionHistory (google thinking)", () => {
});
const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant") as {
content?: Array<{ type?: string; thinking?: string; text?: string }>;
content?: Array<{ type?: string; thinking?: string }>;
};
expect(assistant.content?.map((block) => block.type)).toEqual(["thinking", "text"]);
expect(assistant.content?.map((block) => block.type)).toEqual(["thinking", "thinking"]);
expect(assistant.content?.[0]?.thinking).toBe("signed");
expect(assistant.content?.[1]?.text).toBe("unsigned");
expect(assistant.content?.[1]?.thinking).toBe("unsigned");
});
it("drops empty unsigned thinking blocks for Google models", async () => {
it("keeps empty thinking blocks for Google models", async () => {
const sessionManager = SessionManager.inMemory();
const input = [
{
@@ -251,8 +262,10 @@ describe("sanitizeSessionHistory (google thinking)", () => {
sessionId: "session:google-empty",
});
const assistant = out.find((msg) => (msg as { role?: string }).role === "assistant");
expect(assistant).toBeUndefined();
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"]);
});
it("keeps thinking blocks for non-Google models", async () => {

View File

@@ -13,7 +13,6 @@ vi.mock("./pi-embedded-helpers.js", async () => {
return {
...actual,
isGoogleModelApi: vi.fn(),
downgradeGeminiHistory: vi.fn(),
sanitizeSessionMessagesImages: vi.fn().mockImplementation(async (msgs) => msgs),
};
});
@@ -32,19 +31,14 @@ describe("sanitizeSessionHistory", () => {
beforeEach(async () => {
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" }];
});
vi.resetModules();
({ sanitizeSessionHistory } = await import("./pi-embedded-runner/google.js"));
});
it("should downgrade history for Google models if provider is not google-antigravity", async () => {
it("sanitizes tool call ids for Google model APIs", async () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(true);
const result = await sanitizeSessionHistory({
await sanitizeSessionHistory({
messages: mockMessages,
modelApi: "google-gemini",
provider: "google-vertex",
@@ -53,35 +47,17 @@ describe("sanitizeSessionHistory", () => {
});
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" });
expect(helpers.sanitizeSessionMessagesImages).toHaveBeenCalledWith(
mockMessages,
"session:history",
expect.objectContaining({ sanitizeToolCallIds: true }),
);
});
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 () => {
it("does not sanitize tool call ids for non-Google, non-OpenAI APIs", async () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
const _result = await sanitizeSessionHistory({
await sanitizeSessionHistory({
messages: mockMessages,
modelApi: "anthropic-messages",
provider: "anthropic",
@@ -90,25 +66,14 @@ describe("sanitizeSessionHistory", () => {
});
expect(helpers.isGoogleModelApi).toHaveBeenCalledWith("anthropic-messages");
expect(helpers.downgradeGeminiHistory).not.toHaveBeenCalled();
expect(helpers.sanitizeSessionMessagesImages).toHaveBeenCalledWith(
mockMessages,
"session:history",
expect.objectContaining({ sanitizeToolCallIds: false }),
);
});
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();
});
it("drops reasoning-only assistant messages for openai-responses", async () => {
it("keeps reasoning-only assistant messages for openai-responses", async () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
const messages: AgentMessage[] = [
@@ -135,7 +100,7 @@ describe("sanitizeSessionHistory", () => {
});
expect(helpers.isGoogleModelApi).toHaveBeenCalledWith("openai-responses");
expect(result).toHaveLength(1);
expect(result[0]?.role).toBe("user");
expect(result).toHaveLength(2);
expect(result[1]?.role).toBe("assistant");
});
});

View File

@@ -4,8 +4,6 @@ import type { SessionManager } from "@mariozechner/pi-coding-agent";
import { registerUnhandledRejectionHandler } from "../../infra/unhandled-rejections.js";
import {
downgradeGeminiThinkingBlocks,
downgradeGeminiHistory,
isCompactionFailureError,
isGoogleModelApi,
sanitizeGoogleTurnOrdering,
@@ -52,42 +50,6 @@ function shouldSanitizeToolCallIds(modelApi?: string | null): boolean {
return isGoogleModelApi(modelApi) || OPENAI_TOOL_CALL_ID_APIS.has(modelApi);
}
function filterOpenAIReasoningOnlyMessages(
messages: AgentMessage[],
modelApi?: string | null,
): AgentMessage[] {
if (modelApi !== "openai-responses") return messages;
return messages.filter((msg) => {
if (!msg || typeof msg !== "object") return true;
if ((msg as { role?: unknown }).role !== "assistant") return true;
const assistant = msg as Extract<AgentMessage, { role: "assistant" }>;
const content = assistant.content;
if (!Array.isArray(content) || content.length === 0) return true;
let hasThinking = false;
let hasPairedContent = false;
for (const block of content) {
if (!block || typeof block !== "object") continue;
const type = (block as { type?: unknown }).type;
if (type === "thinking") {
hasThinking = true;
continue;
}
if (type === "toolCall" || type === "toolUse" || type === "functionCall") {
hasPairedContent = true;
break;
}
if (type === "text") {
const text = (block as { text?: unknown }).text;
if (typeof text === "string" && text.trim().length > 0) {
hasPairedContent = true;
break;
}
}
}
return !(hasThinking && !hasPairedContent);
});
}
function findUnsupportedSchemaKeywords(schema: unknown, path: string): string[] {
if (!schema || typeof schema !== "object") return [];
if (Array.isArray(schema)) {
@@ -233,7 +195,6 @@ export async function sanitizeSessionHistory(params: {
const modelId = (params.modelId ?? "").toLowerCase();
const isOpenRouterGemini =
(provider === "openrouter" || provider === "opencode") && modelId.includes("gemini");
const isGeminiLike = isGoogleModelApi(params.modelApi) || isOpenRouterGemini;
const sanitizedImages = await sanitizeSessionMessagesImages(params.messages, "session:history", {
sanitizeToolCallIds: shouldSanitizeToolCallIds(params.modelApi),
enforceToolCallLast: params.modelApi === "anthropic-messages",
@@ -242,26 +203,10 @@ export async function sanitizeSessionHistory(params: {
? { allowBase64Only: true, includeCamelCase: true }
: undefined,
});
// TODO REMOVE when https://github.com/badlogic/pi-mono/pull/838 is merged.
const openaiReasoningFiltered = filterOpenAIReasoningOnlyMessages(
sanitizedImages,
params.modelApi,
);
const repairedTools = sanitizeToolUseResultPairing(openaiReasoningFiltered);
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 = shouldDowngradeThinking
? downgradeGeminiThinkingBlocks(repairedTools)
: repairedTools;
const shouldDowngradeHistory = shouldDowngradeThinking && !isAntigravityProvider;
const downgraded = shouldDowngradeHistory
? downgradeGeminiHistory(downgradedThinking)
: downgradedThinking;
const repairedTools = sanitizeToolUseResultPairing(sanitizedImages);
return applyGoogleTurnOrderingFix({
messages: downgraded,
messages: repairedTools,
modelApi: params.modelApi,
sessionManager: params.sessionManager,
sessionId: params.sessionId,

View File

@@ -36,11 +36,11 @@ describe("injectHistoryImagesIntoMessages", () => {
const didMutate = injectHistoryImagesIntoMessages(messages, new Map([[0, [image]]]));
expect(didMutate).toBe(false);
const content = messages[0]?.content;
if (!Array.isArray(content)) {
const first = messages[0];
if (!first || !Array.isArray(first.content)) {
throw new Error("expected array content");
}
expect(content).toHaveLength(2);
expect(first.content).toHaveLength(2);
});
it("ignores non-user messages and out-of-range indices", () => {