test(agents): centralize AgentMessage fixtures and remove unsafe casts

This commit is contained in:
Peter Steinberger
2026-03-03 02:13:43 +00:00
parent 15a0455d04
commit 70db52de71
14 changed files with 216 additions and 139 deletions

View File

@@ -5,6 +5,7 @@ import {
sanitizeGoogleTurnOrdering,
sanitizeSessionMessagesImages,
} from "./pi-embedded-helpers.js";
import { castAgentMessages } from "./test-helpers/agent-message-fixtures.js";
let testTimestamp = 1;
const nextTimestamp = () => testTimestamp++;
@@ -93,7 +94,7 @@ describe("sanitizeSessionMessagesImages", () => {
});
it("does not synthesize tool call input when missing", async () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read" }],
@@ -111,7 +112,7 @@ describe("sanitizeSessionMessagesImages", () => {
stopReason: "toolUse",
timestamp: nextTimestamp(),
},
] as unknown as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test");
const assistant = out[0] as { content?: Array<Record<string, unknown>> };
@@ -122,7 +123,7 @@ describe("sanitizeSessionMessagesImages", () => {
});
it("removes empty assistant text blocks but preserves tool calls", async () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -143,7 +144,7 @@ describe("sanitizeSessionMessagesImages", () => {
stopReason: "toolUse",
timestamp: nextTimestamp(),
},
] as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test");
@@ -153,7 +154,7 @@ describe("sanitizeSessionMessagesImages", () => {
});
it("sanitizes tool ids in strict mode (alphanumeric only)", async () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -171,7 +172,7 @@ describe("sanitizeSessionMessagesImages", () => {
toolUseId: "call_abc|item:123",
content: [{ type: "text", text: "ok" }],
},
] as unknown as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test", {
sanitizeToolCallIds: true,
@@ -188,7 +189,7 @@ describe("sanitizeSessionMessagesImages", () => {
});
it("sanitizes tool IDs in images-only mode when explicitly enabled", async () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_123|fc_456", name: "read", arguments: {} }],
@@ -214,7 +215,7 @@ describe("sanitizeSessionMessagesImages", () => {
isError: false,
timestamp: nextTimestamp(),
},
] as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test", {
sanitizeMode: "images-only",
@@ -236,7 +237,7 @@ describe("sanitizeSessionMessagesImages", () => {
}
});
it("filters whitespace-only assistant text blocks", async () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -257,7 +258,7 @@ describe("sanitizeSessionMessagesImages", () => {
stopReason: "stop",
timestamp: nextTimestamp(),
},
] as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test");
@@ -266,7 +267,7 @@ describe("sanitizeSessionMessagesImages", () => {
});
});
it("drops assistant messages that only contain empty text", async () => {
const input = [
const input = castAgentMessages([
{ role: "user", content: "hello", timestamp: nextTimestamp() } satisfies UserMessage,
{
role: "assistant",
@@ -285,7 +286,7 @@ describe("sanitizeSessionMessagesImages", () => {
stopReason: "stop",
timestamp: nextTimestamp(),
} satisfies AssistantMessage,
];
]);
const out = await sanitizeSessionMessagesImages(input, "test");
@@ -293,7 +294,7 @@ describe("sanitizeSessionMessagesImages", () => {
expect(out[0]?.role).toBe("user");
});
it("keeps empty assistant error messages", async () => {
const input = [
const input = castAgentMessages([
{ role: "user", content: "hello", timestamp: nextTimestamp() } satisfies UserMessage,
{
role: "assistant",
@@ -329,7 +330,7 @@ describe("sanitizeSessionMessagesImages", () => {
},
timestamp: nextTimestamp(),
} satisfies AssistantMessage,
] as unknown as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test");
@@ -360,7 +361,7 @@ describe("sanitizeSessionMessagesImages", () => {
describe("thought_signature stripping", () => {
it("strips msg_-prefixed thought_signature from assistant message content blocks", async () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -372,7 +373,7 @@ describe("sanitizeSessionMessagesImages", () => {
},
],
},
] as unknown as AgentMessage[];
]);
const out = await sanitizeSessionMessagesImages(input, "test");
@@ -387,19 +388,19 @@ describe("sanitizeSessionMessagesImages", () => {
describe("sanitizeGoogleTurnOrdering", () => {
it("prepends a synthetic user turn when history starts with assistant", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "exec", arguments: {} }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeGoogleTurnOrdering(input);
expect(out[0]?.role).toBe("user");
expect(out[1]?.role).toBe("assistant");
});
it("is a no-op when history starts with user", () => {
const input = [{ role: "user", content: "hi" }] as unknown as AgentMessage[];
const input = castAgentMessages([{ role: "user", content: "hi" }]);
const out = sanitizeGoogleTurnOrdering(input);
expect(out).toBe(input);
});

View File

@@ -2,13 +2,14 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { SessionManager } from "@mariozechner/pi-coding-agent";
import { describe, expect, it, vi } from "vitest";
import { applyGoogleTurnOrderingFix } from "./pi-embedded-runner.js";
import { castAgentMessage } from "./test-helpers/agent-message-fixtures.js";
describe("applyGoogleTurnOrderingFix", () => {
const makeAssistantFirst = (): AgentMessage[] => [
{
castAgentMessage({
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "exec", arguments: {} }],
} as unknown as AgentMessage,
}),
];
it("prepends a bootstrap once and records a marker for Google models", () => {

View File

@@ -5,6 +5,7 @@ import {
makeModelSnapshotEntry,
} from "./pi-embedded-runner.sanitize-session-history.test-harness.js";
import { sanitizeSessionHistory } from "./pi-embedded-runner/google.js";
import { castAgentMessage } from "./test-helpers/agent-message-fixtures.js";
describe("sanitizeSessionHistory openai tool id preservation", () => {
const makeSessionManager = () =>
@@ -17,7 +18,7 @@ describe("sanitizeSessionHistory openai tool id preservation", () => {
]);
const makeMessages = (withReasoning: boolean): AgentMessage[] => [
{
castAgentMessage({
role: "assistant",
content: [
...(withReasoning
@@ -31,14 +32,14 @@ describe("sanitizeSessionHistory openai tool id preservation", () => {
: []),
{ type: "toolCall", id: "call_123|fc_123", name: "noop", arguments: {} },
],
} as unknown as AgentMessage,
{
}),
castAgentMessage({
role: "toolResult",
toolCallId: "call_123|fc_123",
toolName: "noop",
content: [{ type: "text", text: "ok" }],
isError: false,
} as unknown as AgentMessage,
}),
];
it.each([

View File

@@ -15,6 +15,7 @@ import {
sanitizeWithOpenAIResponses,
TEST_SESSION_ID,
} from "./pi-embedded-runner.sanitize-session-history.test-harness.js";
import { castAgentMessage, castAgentMessages } from "./test-helpers/agent-message-fixtures.js";
import { makeZeroUsageSnapshot } from "./usage.js";
vi.mock("./pi-embedded-helpers.js", async () => ({
@@ -136,12 +137,12 @@ describe("sanitizeSessionHistory", () => {
});
const makeCompactionSummaryMessage = (tokensBefore: number, timestamp: string) =>
({
castAgentMessage({
role: "compactionSummary",
summary: "compressed",
tokensBefore,
timestamp,
}) as unknown as AgentMessage;
});
const sanitizeOpenAIHistory = async (
messages: AgentMessage[],
@@ -258,7 +259,7 @@ describe("sanitizeSessionHistory", () => {
setNonGoogleModelApi();
const messages: AgentMessage[] = [
{
castAgentMessage({
role: "user",
content: "forwarded instruction",
provenance: {
@@ -266,7 +267,7 @@ describe("sanitizeSessionHistory", () => {
sourceSessionKey: "agent:main:req",
sourceTool: "sessions_send",
},
} as unknown as AgentMessage,
}),
];
const result = await sanitizeSessionHistory({
@@ -287,14 +288,14 @@ describe("sanitizeSessionHistory", () => {
it("drops stale assistant usage snapshots kept before latest compaction summary", async () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
const messages = [
const messages = castAgentMessages([
{ role: "user", content: "old context" },
makeAssistantUsageMessage({
text: "old answer",
usage: makeUsage(191_919, 2_000, 193_919),
}),
makeCompactionSummaryMessage(191_919, new Date().toISOString()),
] as unknown as AgentMessage[];
]);
const result = await sanitizeOpenAIHistory(messages);
@@ -308,7 +309,7 @@ describe("sanitizeSessionHistory", () => {
it("preserves fresh assistant usage snapshots created after latest compaction summary", async () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
const messages = [
const messages = castAgentMessages([
makeAssistantUsageMessage({
text: "pre-compaction answer",
usage: makeUsage(120_000, 3_000, 123_000),
@@ -319,7 +320,7 @@ describe("sanitizeSessionHistory", () => {
text: "fresh answer",
usage: makeUsage(1_000, 250, 1_250),
}),
] as unknown as AgentMessage[];
]);
const result = await sanitizeOpenAIHistory(messages);
@@ -333,14 +334,14 @@ describe("sanitizeSessionHistory", () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
const compactionTs = Date.parse("2026-02-26T12:00:00.000Z");
const messages = [
const messages = castAgentMessages([
makeCompactionSummaryMessage(191_919, new Date(compactionTs).toISOString()),
makeAssistantUsageMessage({
text: "kept pre-compaction answer",
timestamp: compactionTs - 1_000,
usage: makeUsage(191_919, 2_000, 193_919),
}),
] as unknown as AgentMessage[];
]);
const result = await sanitizeOpenAIHistory(messages);
@@ -354,7 +355,7 @@ describe("sanitizeSessionHistory", () => {
vi.mocked(helpers.isGoogleModelApi).mockReturnValue(false);
const compactionTs = Date.parse("2026-02-26T12:00:00.000Z");
const messages = [
const messages = castAgentMessages([
makeCompactionSummaryMessage(123_000, new Date(compactionTs).toISOString()),
makeAssistantUsageMessage({
text: "kept pre-compaction answer",
@@ -367,7 +368,7 @@ describe("sanitizeSessionHistory", () => {
timestamp: compactionTs + 2_000,
usage: makeUsage(1_000, 250, 1_250),
}),
] as unknown as AgentMessage[];
]);
const result = await sanitizeOpenAIHistory(messages);
@@ -431,13 +432,13 @@ describe("sanitizeSessionHistory", () => {
{
name: "missing input or arguments",
makeMessages: () =>
[
{
castAgentMessages([
castAgentMessage({
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read" }],
} as unknown as AgentMessage,
}),
makeUserMessage("hello"),
] as AgentMessage[],
]),
overrides: { sessionId: "test-session" } as Partial<
Parameters<typeof sanitizeOpenAIHistory>[1]
>,
@@ -445,7 +446,7 @@ describe("sanitizeSessionHistory", () => {
{
name: "invalid or overlong names",
makeMessages: () =>
[
castAgentMessages([
makeAssistantMessage(
[
{
@@ -464,7 +465,7 @@ describe("sanitizeSessionHistory", () => {
{ stopReason: "toolUse" },
),
makeUserMessage("hello"),
] as AgentMessage[],
]),
overrides: {} as Partial<Parameters<typeof sanitizeOpenAIHistory>[1]>,
},
])("drops malformed tool calls: $name", async ({ makeMessages, overrides }) => {

View File

@@ -1,5 +1,5 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { castAgentMessage } from "../../test-helpers/agent-message-fixtures.js";
import {
selectCompactionTimeoutSnapshot,
shouldFlagCompactionTimeout,
@@ -32,8 +32,8 @@ describe("compaction-timeout helpers", () => {
});
it("uses pre-compaction snapshot when compaction timeout occurs", () => {
const pre = [{ role: "assistant", content: "pre" } as unknown as AgentMessage] as const;
const current = [{ role: "assistant", content: "current" } as unknown as AgentMessage] as const;
const pre = [castAgentMessage({ role: "assistant", content: "pre" })] as const;
const current = [castAgentMessage({ role: "assistant", content: "current" })] as const;
const selected = selectCompactionTimeoutSnapshot({
timedOutDuringCompaction: true,
preCompactionSnapshot: [...pre],
@@ -47,7 +47,7 @@ describe("compaction-timeout helpers", () => {
});
it("falls back to current snapshot when pre-compaction snapshot is unavailable", () => {
const current = [{ role: "assistant", content: "current" } as unknown as AgentMessage] as const;
const current = [castAgentMessage({ role: "assistant", content: "current" })] as const;
const selected = selectCompactionTimeoutSnapshot({
timedOutDuringCompaction: true,
preCompactionSnapshot: null,

View File

@@ -1,6 +1,7 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { ImageContent } from "@mariozechner/pi-ai";
import { describe, expect, it } from "vitest";
import { castAgentMessage } from "../../test-helpers/agent-message-fixtures.js";
import { PRUNED_HISTORY_IMAGE_MARKER, pruneProcessedHistoryImages } from "./history-image-prune.js";
describe("pruneProcessedHistoryImages", () => {
@@ -8,14 +9,14 @@ describe("pruneProcessedHistoryImages", () => {
it("prunes image blocks from user messages that already have assistant replies", () => {
const messages: AgentMessage[] = [
{
castAgentMessage({
role: "user",
content: [{ type: "text", text: "See /tmp/photo.png" }, { ...image }],
} as AgentMessage,
{
}),
castAgentMessage({
role: "assistant",
content: "got it",
} as unknown as AgentMessage,
}),
];
const didMutate = pruneProcessedHistoryImages(messages);
@@ -31,10 +32,10 @@ describe("pruneProcessedHistoryImages", () => {
it("does not prune latest user message when no assistant response exists yet", () => {
const messages: AgentMessage[] = [
{
castAgentMessage({
role: "user",
content: [{ type: "text", text: "See /tmp/photo.png" }, { ...image }],
} as AgentMessage,
}),
];
const didMutate = pruneProcessedHistoryImages(messages);
@@ -50,10 +51,10 @@ describe("pruneProcessedHistoryImages", () => {
it("does not change messages when no assistant turn exists", () => {
const messages: AgentMessage[] = [
{
castAgentMessage({
role: "user",
content: "noop",
} as AgentMessage,
}),
];
const didMutate = pruneProcessedHistoryImages(messages);

View File

@@ -1,15 +1,16 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { castAgentMessage } from "../test-helpers/agent-message-fixtures.js";
import { dropThinkingBlocks, isAssistantMessageWithContent } from "./thinking.js";
describe("isAssistantMessageWithContent", () => {
it("accepts assistant messages with array content and rejects others", () => {
const assistant = {
const assistant = castAgentMessage({
role: "assistant",
content: [{ type: "text", text: "ok" }],
} as AgentMessage;
const user = { role: "user", content: "hi" } as AgentMessage;
const malformed = { role: "assistant", content: "not-array" } as unknown as AgentMessage;
});
const user = castAgentMessage({ role: "user", content: "hi" });
const malformed = castAgentMessage({ role: "assistant", content: "not-array" });
expect(isAssistantMessageWithContent(assistant)).toBe(true);
expect(isAssistantMessageWithContent(user)).toBe(false);
@@ -20,8 +21,8 @@ describe("isAssistantMessageWithContent", () => {
describe("dropThinkingBlocks", () => {
it("returns the original reference when no thinking blocks are present", () => {
const messages: AgentMessage[] = [
{ role: "user", content: "hello" } as AgentMessage,
{ role: "assistant", content: [{ type: "text", text: "world" }] } as AgentMessage,
castAgentMessage({ role: "user", content: "hello" }),
castAgentMessage({ role: "assistant", content: [{ type: "text", text: "world" }] }),
];
const result = dropThinkingBlocks(messages);
@@ -30,13 +31,13 @@ describe("dropThinkingBlocks", () => {
it("drops thinking blocks while preserving non-thinking assistant content", () => {
const messages: AgentMessage[] = [
{
castAgentMessage({
role: "assistant",
content: [
{ type: "thinking", thinking: "internal" },
{ type: "text", text: "final" },
],
} as unknown as AgentMessage,
}),
];
const result = dropThinkingBlocks(messages);
@@ -47,10 +48,10 @@ describe("dropThinkingBlocks", () => {
it("keeps assistant turn structure when all content blocks were thinking", () => {
const messages: AgentMessage[] = [
{
castAgentMessage({
role: "assistant",
content: [{ type: "thinking", thinking: "internal-only" }],
} as unknown as AgentMessage,
}),
];
const result = dropThinkingBlocks(messages);

View File

@@ -1,5 +1,6 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { castAgentMessage } from "../test-helpers/agent-message-fixtures.js";
import {
CONTEXT_LIMIT_TRUNCATION_NOTICE,
PREEMPTIVE_TOOL_RESULT_COMPACTION_PLACEHOLDER,
@@ -7,35 +8,35 @@ import {
} from "./tool-result-context-guard.js";
function makeUser(text: string): AgentMessage {
return {
return castAgentMessage({
role: "user",
content: text,
timestamp: Date.now(),
} as unknown as AgentMessage;
});
}
function makeToolResult(id: string, text: string): AgentMessage {
return {
return castAgentMessage({
role: "toolResult",
toolCallId: id,
toolName: "read",
content: [{ type: "text", text }],
isError: false,
timestamp: Date.now(),
} as unknown as AgentMessage;
});
}
function makeLegacyToolResult(id: string, text: string): AgentMessage {
return {
return castAgentMessage({
role: "tool",
tool_call_id: id,
tool_name: "read",
content: text,
} as unknown as AgentMessage;
});
}
function makeToolResultWithDetails(id: string, text: string, detailText: string): AgentMessage {
return {
return castAgentMessage({
role: "toolResult",
toolCallId: id,
toolName: "read",
@@ -49,7 +50,7 @@ function makeToolResultWithDetails(id: string, text: string, detailText: string)
},
isError: false,
timestamp: Date.now(),
} as unknown as AgentMessage;
});
}
function getToolResultText(msg: AgentMessage): string {
@@ -199,11 +200,10 @@ describe("installToolResultContextGuard", () => {
it("wraps an existing transformContext and guards the transformed output", async () => {
const agent = makeGuardableAgent((messages) => {
return messages.map(
(msg) =>
({
...(msg as unknown as Record<string, unknown>),
}) as unknown as AgentMessage,
return messages.map((msg) =>
castAgentMessage({
...(msg as unknown as Record<string, unknown>),
}),
);
});
const contextForNextCall = makeTwoToolResultOverflowContext();
@@ -254,10 +254,10 @@ describe("installToolResultContextGuard", () => {
await agent.transformContext?.(contextForNextCall, new AbortController().signal);
const oldResult = contextForNextCall[1] as unknown as {
const oldResult = contextForNextCall[1] as {
details?: unknown;
};
const newResult = contextForNextCall[2] as unknown as {
const newResult = contextForNextCall[2] as {
details?: unknown;
};
const oldResultText = getToolResultText(contextForNextCall[1]);

View File

@@ -5,6 +5,7 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { Api, Model } from "@mariozechner/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { describe, expect, it, vi } from "vitest";
import { castAgentMessage } from "../test-helpers/agent-message-fixtures.js";
import {
getCompactionSafeguardRuntime,
setCompactionSafeguardRuntime,
@@ -218,11 +219,11 @@ describe("computeAdaptiveChunkRatio", () => {
// Small messages: 1000 tokens each, well under 10% of context
const messages: AgentMessage[] = [
{ role: "user", content: "x".repeat(1000), timestamp: Date.now() },
{
castAgentMessage({
role: "assistant",
content: [{ type: "text", text: "y".repeat(1000) }],
timestamp: Date.now(),
} as unknown as AgentMessage,
}),
];
const ratio = computeAdaptiveChunkRatio(messages, CONTEXT_WINDOW);
@@ -233,11 +234,11 @@ describe("computeAdaptiveChunkRatio", () => {
// Large messages: ~50K tokens each (25% of context)
const messages: AgentMessage[] = [
{ role: "user", content: "x".repeat(50_000 * 4), timestamp: Date.now() },
{
castAgentMessage({
role: "assistant",
content: [{ type: "text", text: "y".repeat(50_000 * 4) }],
timestamp: Date.now(),
} as unknown as AgentMessage,
}),
];
const ratio = computeAdaptiveChunkRatio(messages, CONTEXT_WINDOW);

View File

@@ -2,6 +2,7 @@ import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { SessionManager } from "@mariozechner/pi-coding-agent";
import { describe, expect, it } from "vitest";
import { installSessionToolResultGuard } from "./session-tool-result-guard.js";
import { castAgentMessage } from "./test-helpers/agent-message-fixtures.js";
type AppendMessage = Parameters<SessionManager["appendMessage"]>[0];
@@ -388,10 +389,10 @@ describe("installSessionToolResultGuard", () => {
return undefined;
}
return {
message: {
message: castAgentMessage({
...(message as unknown as Record<string, unknown>),
content: [{ type: "text", text: "rewritten by hook" }],
} as unknown as AgentMessage,
}),
};
},
});
@@ -425,10 +426,10 @@ describe("installSessionToolResultGuard", () => {
installSessionToolResultGuard(sm, {
transformMessageForPersistence: (message) =>
(message as { role?: string }).role === "user"
? ({
? castAgentMessage({
...(message as unknown as Record<string, unknown>),
provenance: { kind: "inter_session", sourceTool: "sessions_send" },
} as unknown as AgentMessage)
})
: message,
});

View File

@@ -1,9 +1,10 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, it, expect } from "vitest";
import { sanitizeToolCallInputs } from "./session-transcript-repair.js";
import { castAgentMessage, castAgentMessages } from "./test-helpers/agent-message-fixtures.js";
function mkSessionsSpawnToolCall(content: string): AgentMessage {
return {
return castAgentMessage({
role: "assistant",
content: [
{
@@ -23,7 +24,7 @@ function mkSessionsSpawnToolCall(content: string): AgentMessage {
},
],
timestamp: Date.now(),
} as unknown as AgentMessage;
});
}
describe("sanitizeToolCallInputs redacts sessions_spawn attachments", () => {
@@ -44,7 +45,7 @@ describe("sanitizeToolCallInputs redacts sessions_spawn attachments", () => {
it("redacts attachments content from tool input payloads too", () => {
const secret = "INPUT_SECRET_SHOULD_NOT_PERSIST";
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -59,7 +60,7 @@ describe("sanitizeToolCallInputs redacts sessions_spawn attachments", () => {
},
],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallInputs(input);
const msg = out[0] as { content?: unknown[] };

View File

@@ -6,6 +6,7 @@ import {
repairToolUseResultPairing,
stripToolResultDetails,
} from "./session-transcript-repair.js";
import { castAgentMessage, castAgentMessages } from "./test-helpers/agent-message-fixtures.js";
const TOOL_CALL_BLOCK_TYPES = new Set(["toolCall", "toolUse", "functionCall"]);
@@ -25,7 +26,7 @@ describe("sanitizeToolUseResultPairing", () => {
middleMessage?: unknown;
secondText?: string;
}): AgentMessage[] =>
[
castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }],
@@ -37,7 +38,7 @@ describe("sanitizeToolUseResultPairing", () => {
content: [{ type: "text", text: "first" }],
isError: false,
},
...(opts?.middleMessage ? [opts.middleMessage as AgentMessage] : []),
...(opts?.middleMessage ? [castAgentMessage(opts.middleMessage)] : []),
{
role: "toolResult",
toolCallId: "call_1",
@@ -45,10 +46,10 @@ describe("sanitizeToolUseResultPairing", () => {
content: [{ type: "text", text: opts?.secondText ?? "second" }],
isError: false,
},
] as unknown as AgentMessage[];
]);
it("moves tool results directly after tool calls and inserts missing results", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -64,7 +65,7 @@ describe("sanitizeToolUseResultPairing", () => {
content: [{ type: "text", text: "ok" }],
isError: false,
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolUseResultPairing(input);
expect(out[0]?.role).toBe("assistant");
@@ -76,7 +77,7 @@ describe("sanitizeToolUseResultPairing", () => {
});
it("repairs blank tool result names from matching tool calls", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read", arguments: {} }],
@@ -88,7 +89,7 @@ describe("sanitizeToolUseResultPairing", () => {
content: [{ type: "text", text: "ok" }],
isError: false,
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolUseResultPairing(input);
const toolResult = out.find((message) => message.role === "toolResult") as {
@@ -99,10 +100,10 @@ describe("sanitizeToolUseResultPairing", () => {
});
it("drops duplicate tool results for the same id within a span", () => {
const input = [
const input = castAgentMessages([
...buildDuplicateToolResultInput(),
{ role: "user", content: "ok" },
] as AgentMessage[];
]);
const out = sanitizeToolUseResultPairing(input);
expect(out.filter((m) => m.role === "toolResult")).toHaveLength(1);
@@ -123,7 +124,7 @@ describe("sanitizeToolUseResultPairing", () => {
});
it("drops orphan tool results that do not match any tool call", () => {
const input = [
const input = castAgentMessages([
{ role: "user", content: "hello" },
{
role: "toolResult",
@@ -136,7 +137,7 @@ describe("sanitizeToolUseResultPairing", () => {
role: "assistant",
content: [{ type: "text", text: "ok" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolUseResultPairing(input);
expect(out.some((m) => m.role === "toolResult")).toBe(false);
@@ -147,14 +148,14 @@ describe("sanitizeToolUseResultPairing", () => {
// When an assistant message has stopReason: "error", its tool_use blocks may be
// incomplete/malformed. We should NOT create synthetic tool_results for them,
// as this causes API 400 errors: "unexpected tool_use_id found in tool_result blocks"
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_error", name: "exec", arguments: {} }],
stopReason: "error",
},
{ role: "user", content: "something went wrong" },
] as unknown as AgentMessage[];
]);
const result = repairToolUseResultPairing(input);
@@ -169,14 +170,14 @@ describe("sanitizeToolUseResultPairing", () => {
it("skips tool call extraction for assistant messages with stopReason 'aborted'", () => {
// When a request is aborted mid-stream, the assistant message may have incomplete
// tool_use blocks (with partialJson). We should NOT create synthetic tool_results.
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_aborted", name: "Bash", arguments: {} }],
stopReason: "aborted",
},
{ role: "user", content: "retrying after abort" },
] as unknown as AgentMessage[];
]);
const result = repairToolUseResultPairing(input);
@@ -190,14 +191,14 @@ describe("sanitizeToolUseResultPairing", () => {
it("still repairs tool results for normal assistant messages with stopReason 'toolUse'", () => {
// Normal tool calls (stopReason: "toolUse" or "stop") should still be repaired
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_normal", name: "read", arguments: {} }],
stopReason: "toolUse",
},
{ role: "user", content: "user message" },
] as unknown as AgentMessage[];
]);
const result = repairToolUseResultPairing(input);
@@ -210,7 +211,7 @@ describe("sanitizeToolUseResultPairing", () => {
// When an assistant message is aborted, any tool results that follow should be
// dropped as orphans (since we skip extracting tool calls from aborted messages).
// This addresses the edge case where a partial tool result was persisted before abort.
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_aborted", name: "exec", arguments: {} }],
@@ -224,7 +225,7 @@ describe("sanitizeToolUseResultPairing", () => {
isError: false,
},
{ role: "user", content: "retrying" },
] as unknown as AgentMessage[];
]);
const result = repairToolUseResultPairing(input);
@@ -244,12 +245,12 @@ describe("sanitizeToolCallInputs", () => {
options?: Parameters<typeof sanitizeToolCallInputs>[1],
) {
return sanitizeToolCallInputs(
[
castAgentMessages([
{
role: "assistant",
content,
},
] as unknown as AgentMessage[],
]),
options,
);
}
@@ -262,13 +263,13 @@ describe("sanitizeToolCallInputs", () => {
}
it("drops tool calls missing input or arguments", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call_1", name: "read" }],
},
{ role: "user", content: "hello" },
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallInputs(input);
expect(out.map((m) => m.role)).toEqual(["user"]);
@@ -325,7 +326,7 @@ describe("sanitizeToolCallInputs", () => {
});
it("keeps valid tool calls and preserves text blocks", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -334,7 +335,7 @@ describe("sanitizeToolCallInputs", () => {
{ type: "toolCall", id: "call_drop", name: "read" },
],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallInputs(input);
const assistant = out[0] as Extract<AgentMessage, { role: "assistant" }>;
@@ -384,7 +385,7 @@ describe("sanitizeToolCallInputs", () => {
});
it("preserves toolUse input shape for sessions_spawn when no attachments are present", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -396,7 +397,7 @@ describe("sanitizeToolCallInputs", () => {
},
],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallInputs(input);
const toolCalls = getAssistantToolCallBlocks(out) as Array<Record<string, unknown>>;
@@ -408,7 +409,7 @@ describe("sanitizeToolCallInputs", () => {
});
it("redacts sessions_spawn attachments for mixed-case and padded tool names", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -423,7 +424,7 @@ describe("sanitizeToolCallInputs", () => {
},
],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallInputs(input);
const toolCalls = getAssistantToolCallBlocks(out) as Array<Record<string, unknown>>;
@@ -448,7 +449,7 @@ describe("sanitizeToolCallInputs", () => {
describe("stripToolResultDetails", () => {
it("removes details only from toolResult messages", () => {
const input = [
const input = castAgentMessages([
{
role: "toolResult",
toolCallId: "call_1",
@@ -458,7 +459,7 @@ describe("stripToolResultDetails", () => {
},
{ role: "assistant", content: [{ type: "text", text: "keep me" }], details: { no: "touch" } },
{ role: "user", content: "hello" },
] as unknown as AgentMessage[];
]);
const out = stripToolResultDetails(input) as unknown as Array<Record<string, unknown>>;
@@ -472,7 +473,7 @@ describe("stripToolResultDetails", () => {
});
it("returns the same array reference when there are no toolResult details", () => {
const input = [
const input = castAgentMessages([
{ role: "assistant", content: [{ type: "text", text: "a" }] },
{
role: "toolResult",
@@ -481,7 +482,7 @@ describe("stripToolResultDetails", () => {
content: [{ type: "text", text: "ok" }],
},
{ role: "user", content: "b" },
] as unknown as AgentMessage[];
]);
const out = stripToolResultDetails(input);
expect(out).toBe(input);

View File

@@ -0,0 +1,66 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ToolResultMessage, Usage, UserMessage } from "@mariozechner/pi-ai";
const ZERO_USAGE: Usage = {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
totalTokens: 0,
cost: {
input: 0,
output: 0,
cacheRead: 0,
cacheWrite: 0,
total: 0,
},
};
export function castAgentMessage(message: unknown): AgentMessage {
return message as AgentMessage;
}
export function castAgentMessages(messages: unknown[]): AgentMessage[] {
return messages as AgentMessage[];
}
export function makeAgentUserMessage(
overrides: Partial<UserMessage> & Pick<UserMessage, "content">,
): UserMessage {
return {
role: "user",
timestamp: 0,
...overrides,
};
}
export function makeAgentAssistantMessage(
overrides: Partial<AssistantMessage> & Pick<AssistantMessage, "content">,
): AssistantMessage {
return {
role: "assistant",
api: "openai-responses",
provider: "openai",
model: "test-model",
usage: ZERO_USAGE,
stopReason: "stop",
timestamp: 0,
...overrides,
};
}
export function makeAgentToolResultMessage(
overrides: Partial<ToolResultMessage> &
Pick<ToolResultMessage, "toolCallId" | "toolName" | "content">,
): ToolResultMessage {
const { toolCallId, toolName, content, ...rest } = overrides;
return {
role: "toolResult",
toolCallId,
toolName,
content,
isError: false,
timestamp: 0,
...rest,
};
}

View File

@@ -1,12 +1,13 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { castAgentMessages } from "./test-helpers/agent-message-fixtures.js";
import {
isValidCloudCodeAssistToolId,
sanitizeToolCallIdsForCloudCodeAssist,
} from "./tool-call-id.js";
const buildDuplicateIdCollisionInput = () =>
[
castAgentMessages([
{
role: "assistant",
content: [
@@ -26,7 +27,7 @@ const buildDuplicateIdCollisionInput = () =>
toolName: "read",
content: [{ type: "text", text: "two" }],
},
] as unknown as AgentMessage[];
]);
function expectCollisionIdsRemainDistinct(
out: AgentMessage[],
@@ -65,7 +66,7 @@ function expectSingleToolCallRewrite(
describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
describe("strict mode (default)", () => {
it("is a no-op for already-valid non-colliding IDs", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call1", name: "read", arguments: {} }],
@@ -76,14 +77,14 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
toolName: "read",
content: [{ type: "text", text: "ok" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallIdsForCloudCodeAssist(input);
expect(out).toBe(input);
});
it("strips non-alphanumeric characters from tool call IDs", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "call|item:123", name: "read", arguments: {} }],
@@ -94,7 +95,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
toolName: "read",
content: [{ type: "text", text: "ok" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallIdsForCloudCodeAssist(input);
expect(out).not.toBe(input);
@@ -113,7 +114,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
it("caps tool call IDs at 40 chars while preserving uniqueness", () => {
const longA = `call_${"a".repeat(60)}`;
const longB = `call_${"a".repeat(59)}b`;
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -133,7 +134,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
toolName: "read",
content: [{ type: "text", text: "two" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallIdsForCloudCodeAssist(input);
const { aId, bId } = expectCollisionIdsRemainDistinct(out, "strict");
@@ -144,7 +145,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
describe("strict mode (alphanumeric only)", () => {
it("strips underscores and hyphens from tool call IDs", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -162,7 +163,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
toolName: "login",
content: [{ type: "text", text: "ok" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict");
expect(out).not.toBe(input);
@@ -184,7 +185,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
describe("strict9 mode (Mistral tool call IDs)", () => {
it("is a no-op for already-valid 9-char alphanumeric IDs", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [{ type: "toolCall", id: "abc123XYZ", name: "read", arguments: {} }],
@@ -195,14 +196,14 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
toolName: "read",
content: [{ type: "text", text: "ok" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
expect(out).toBe(input);
});
it("enforces alphanumeric IDs with length 9", () => {
const input = [
const input = castAgentMessages([
{
role: "assistant",
content: [
@@ -222,7 +223,7 @@ describe("sanitizeToolCallIdsForCloudCodeAssist", () => {
toolName: "read",
content: [{ type: "text", text: "two" }],
},
] as unknown as AgentMessage[];
]);
const out = sanitizeToolCallIdsForCloudCodeAssist(input, "strict9");
expect(out).not.toBe(input);