refactor(agents): extract transcript repair module

This commit is contained in:
Peter Steinberger
2026-01-10 22:03:42 +00:00
parent 708f04b02f
commit 08cc8f2281
5 changed files with 269 additions and 267 deletions

View File

@@ -14,7 +14,6 @@ import {
sanitizeGoogleTurnOrdering,
sanitizeSessionMessagesImages,
sanitizeToolCallId,
sanitizeToolUseResultPairing,
validateGeminiTurns,
} from "./pi-embedded-helpers.js";
import {
@@ -585,99 +584,6 @@ describe("sanitizeSessionMessagesImages", () => {
});
});
describe("sanitizeToolUseResultPairing", () => {
it("moves tool results directly after tool calls and inserts missing results", () => {
const input = [
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: {} },
{ type: "toolCall", id: "call_2", name: "bash", arguments: {} },
],
},
{ role: "user", content: "user message that should come after tool use" },
{
role: "toolResult",
toolCallId: "call_2",
toolName: "bash",
content: [{ type: "text", text: "ok" }],
isError: false,
},
] satisfies AgentMessage[];
const out = sanitizeToolUseResultPairing(input);
expect(out[0]?.role).toBe("assistant");
expect(out[1]?.role).toBe("toolResult");
expect((out[1] as { toolCallId?: string }).toolCallId).toBe("call_1");
expect(out[2]?.role).toBe("toolResult");
expect((out[2] as { toolCallId?: string }).toolCallId).toBe("call_2");
expect(out[3]?.role).toBe("user");
});
it("drops duplicate tool results for the same id within a span", () => {
const input = [
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: {} },
],
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "first" }],
isError: false,
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "second" }],
isError: false,
},
{ role: "user", content: "ok" },
] satisfies AgentMessage[];
const out = sanitizeToolUseResultPairing(input);
expect(out.filter((m) => m.role === "toolResult")).toHaveLength(1);
});
it("drops duplicate tool results for the same id across the transcript", () => {
const input = [
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: {} },
],
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "first" }],
isError: false,
},
{ role: "assistant", content: [{ type: "text", text: "ok" }] },
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "second (duplicate)" }],
isError: false,
},
] satisfies AgentMessage[];
const out = sanitizeToolUseResultPairing(input);
const results = out.filter((m) => m.role === "toolResult") as Array<{
toolCallId?: string;
content?: unknown;
}>;
expect(results).toHaveLength(1);
expect(results[0]?.toolCallId).toBe("call_1");
});
});
describe("normalizeTextForComparison", () => {
it("lowercases text", () => {
expect(normalizeTextForComparison("Hello World")).toBe("hello world");

View File

@@ -218,178 +218,6 @@ export async function sanitizeSessionMessagesImages(
return out;
}
type ToolCallLike = {
id: string;
name?: string;
};
function extractToolCallsFromAssistant(
msg: Extract<AgentMessage, { role: "assistant" }>,
): ToolCallLike[] {
const content = msg.content;
if (!Array.isArray(content)) return [];
const toolCalls: ToolCallLike[] = [];
for (const block of content) {
if (!block || typeof block !== "object") continue;
const rec = block as { type?: unknown; id?: unknown; name?: unknown };
if (typeof rec.id !== "string" || !rec.id) continue;
if (
rec.type === "toolCall" ||
rec.type === "toolUse" ||
rec.type === "functionCall"
) {
toolCalls.push({
id: rec.id,
name: typeof rec.name === "string" ? rec.name : undefined,
});
}
}
return toolCalls;
}
function extractToolResultId(
msg: Extract<AgentMessage, { role: "toolResult" }>,
): string | null {
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 null;
}
function makeMissingToolResult(params: {
toolCallId: string;
toolName?: string;
}): Extract<AgentMessage, { role: "toolResult" }> {
return {
role: "toolResult",
toolCallId: params.toolCallId,
toolName: params.toolName ?? "unknown",
content: [
{
type: "text",
text: "[clawdbot] missing tool result in session history; inserted synthetic error result for transcript repair.",
},
],
isError: true,
timestamp: Date.now(),
} as Extract<AgentMessage, { role: "toolResult" }>;
}
export function sanitizeToolUseResultPairing(
messages: AgentMessage[],
): AgentMessage[] {
// Anthropic (and Cloud Code Assist) reject transcripts where assistant tool calls are not
// immediately followed by matching tool results. Session files can end up with results
// displaced (e.g. after user turns) or duplicated. Repair by:
// - moving matching toolResult messages directly after their assistant toolCall turn
// - inserting synthetic error toolResults for missing ids
// - dropping duplicate toolResults for the same id (anywhere in the transcript)
const out: AgentMessage[] = [];
const seenToolResultIds = new Set<string>();
const pushToolResult = (
msg: Extract<AgentMessage, { role: "toolResult" }>,
) => {
const id = extractToolResultId(msg);
if (id && seenToolResultIds.has(id)) return;
if (id) seenToolResultIds.add(id);
out.push(msg);
};
for (let i = 0; i < messages.length; i += 1) {
const msg = messages[i] as AgentMessage;
if (!msg || typeof msg !== "object") {
out.push(msg);
continue;
}
const role = (msg as { role?: unknown }).role;
if (role !== "assistant") {
if (role === "toolResult") {
pushToolResult(msg as Extract<AgentMessage, { role: "toolResult" }>);
} else {
out.push(msg);
}
continue;
}
const assistant = msg as Extract<AgentMessage, { role: "assistant" }>;
const toolCalls = extractToolCallsFromAssistant(assistant);
if (toolCalls.length === 0) {
out.push(msg);
continue;
}
const toolCallIds = new Set(toolCalls.map((t) => t.id));
const spanResultsById = new Map<
string,
Extract<AgentMessage, { role: "toolResult" }>
>();
const remainder: AgentMessage[] = [];
let j = i + 1;
for (; j < messages.length; j += 1) {
const next = messages[j] as AgentMessage;
if (!next || typeof next !== "object") {
remainder.push(next);
continue;
}
const nextRole = (next as { role?: unknown }).role;
if (nextRole === "assistant") break;
if (nextRole === "toolResult") {
const toolResult = next as Extract<
AgentMessage,
{ role: "toolResult" }
>;
const id = extractToolResultId(toolResult);
if (id && toolCallIds.has(id)) {
if (seenToolResultIds.has(id)) {
continue;
}
if (!spanResultsById.has(id)) {
spanResultsById.set(id, toolResult);
}
continue;
}
}
remainder.push(next);
}
out.push(msg);
for (const call of toolCalls) {
const existing = spanResultsById.get(call.id);
pushToolResult(
existing ??
makeMissingToolResult({ toolCallId: call.id, toolName: call.name }),
);
}
for (const rem of remainder) {
if (!rem || typeof rem !== "object") {
out.push(rem);
continue;
}
const remRole = (rem as { role?: unknown }).role;
if (remRole === "toolResult") {
pushToolResult(rem as Extract<AgentMessage, { role: "toolResult" }>);
continue;
}
out.push(rem);
}
i = j - 1;
}
return out;
}
const GOOGLE_TURN_ORDER_BOOTSTRAP_TEXT = "(session bootstrap)";
export function isGoogleModelApi(api?: string | null): boolean {

View File

@@ -88,7 +88,6 @@ import {
pickFallbackThinkingLevel,
sanitizeGoogleTurnOrdering,
sanitizeSessionMessagesImages,
sanitizeToolUseResultPairing,
validateGeminiTurns,
} from "./pi-embedded-helpers.js";
import {
@@ -106,6 +105,7 @@ import { makeToolPrunablePredicate } from "./pi-extensions/context-pruning/tools
import { toToolDefinitions } from "./pi-tool-definition-adapter.js";
import { createClawdbotCodingTools } from "./pi-tools.js";
import { resolveSandboxContext } from "./sandbox.js";
import { sanitizeToolUseResultPairing } from "./session-transcript-repair.js";
import {
applySkillEnvOverrides,
applySkillEnvOverridesFromSnapshot,

View File

@@ -0,0 +1,95 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import { describe, expect, it } from "vitest";
import { sanitizeToolUseResultPairing } from "./session-transcript-repair.js";
describe("sanitizeToolUseResultPairing", () => {
it("moves tool results directly after tool calls and inserts missing results", () => {
const input = [
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: {} },
{ type: "toolCall", id: "call_2", name: "bash", arguments: {} },
],
},
{ role: "user", content: "user message that should come after tool use" },
{
role: "toolResult",
toolCallId: "call_2",
toolName: "bash",
content: [{ type: "text", text: "ok" }],
isError: false,
},
] satisfies AgentMessage[];
const out = sanitizeToolUseResultPairing(input);
expect(out[0]?.role).toBe("assistant");
expect(out[1]?.role).toBe("toolResult");
expect((out[1] as { toolCallId?: string }).toolCallId).toBe("call_1");
expect(out[2]?.role).toBe("toolResult");
expect((out[2] as { toolCallId?: string }).toolCallId).toBe("call_2");
expect(out[3]?.role).toBe("user");
});
it("drops duplicate tool results for the same id within a span", () => {
const input = [
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: {} },
],
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "first" }],
isError: false,
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "second" }],
isError: false,
},
{ role: "user", content: "ok" },
] satisfies AgentMessage[];
const out = sanitizeToolUseResultPairing(input);
expect(out.filter((m) => m.role === "toolResult")).toHaveLength(1);
});
it("drops duplicate tool results for the same id across the transcript", () => {
const input = [
{
role: "assistant",
content: [
{ type: "toolCall", id: "call_1", name: "read", arguments: {} },
],
},
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "first" }],
isError: false,
},
{ role: "assistant", content: [{ type: "text", text: "ok" }] },
{
role: "toolResult",
toolCallId: "call_1",
toolName: "read",
content: [{ type: "text", text: "second (duplicate)" }],
isError: false,
},
] satisfies AgentMessage[];
const out = sanitizeToolUseResultPairing(input);
const results = out.filter((m) => m.role === "toolResult") as Array<{
toolCallId?: string;
}>;
expect(results).toHaveLength(1);
expect(results[0]?.toolCallId).toBe("call_1");
});
});

View File

@@ -0,0 +1,173 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
type ToolCallLike = {
id: string;
name?: string;
};
function extractToolCallsFromAssistant(
msg: Extract<AgentMessage, { role: "assistant" }>,
): ToolCallLike[] {
const content = msg.content;
if (!Array.isArray(content)) return [];
const toolCalls: ToolCallLike[] = [];
for (const block of content) {
if (!block || typeof block !== "object") continue;
const rec = block as { type?: unknown; id?: unknown; name?: unknown };
if (typeof rec.id !== "string" || !rec.id) continue;
if (
rec.type === "toolCall" ||
rec.type === "toolUse" ||
rec.type === "functionCall"
) {
toolCalls.push({
id: rec.id,
name: typeof rec.name === "string" ? rec.name : undefined,
});
}
}
return toolCalls;
}
function extractToolResultId(
msg: Extract<AgentMessage, { role: "toolResult" }>,
): string | null {
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 null;
}
function makeMissingToolResult(params: {
toolCallId: string;
toolName?: string;
}): Extract<AgentMessage, { role: "toolResult" }> {
return {
role: "toolResult",
toolCallId: params.toolCallId,
toolName: params.toolName ?? "unknown",
content: [
{
type: "text",
text: "[clawdbot] missing tool result in session history; inserted synthetic error result for transcript repair.",
},
],
isError: true,
timestamp: Date.now(),
} as Extract<AgentMessage, { role: "toolResult" }>;
}
export function sanitizeToolUseResultPairing(
messages: AgentMessage[],
): AgentMessage[] {
// Anthropic (and Cloud Code Assist) reject transcripts where assistant tool calls are not
// immediately followed by matching tool results. Session files can end up with results
// displaced (e.g. after user turns) or duplicated. Repair by:
// - moving matching toolResult messages directly after their assistant toolCall turn
// - inserting synthetic error toolResults for missing ids
// - dropping duplicate toolResults for the same id (anywhere in the transcript)
const out: AgentMessage[] = [];
const seenToolResultIds = new Set<string>();
const pushToolResult = (
msg: Extract<AgentMessage, { role: "toolResult" }>,
) => {
const id = extractToolResultId(msg);
if (id && seenToolResultIds.has(id)) return;
if (id) seenToolResultIds.add(id);
out.push(msg);
};
for (let i = 0; i < messages.length; i += 1) {
const msg = messages[i] as AgentMessage;
if (!msg || typeof msg !== "object") {
out.push(msg);
continue;
}
const role = (msg as { role?: unknown }).role;
if (role !== "assistant") {
if (role === "toolResult") {
pushToolResult(msg as Extract<AgentMessage, { role: "toolResult" }>);
} else {
out.push(msg);
}
continue;
}
const assistant = msg as Extract<AgentMessage, { role: "assistant" }>;
const toolCalls = extractToolCallsFromAssistant(assistant);
if (toolCalls.length === 0) {
out.push(msg);
continue;
}
const toolCallIds = new Set(toolCalls.map((t) => t.id));
const spanResultsById = new Map<
string,
Extract<AgentMessage, { role: "toolResult" }>
>();
const remainder: AgentMessage[] = [];
let j = i + 1;
for (; j < messages.length; j += 1) {
const next = messages[j] as AgentMessage;
if (!next || typeof next !== "object") {
remainder.push(next);
continue;
}
const nextRole = (next as { role?: unknown }).role;
if (nextRole === "assistant") break;
if (nextRole === "toolResult") {
const toolResult = next as Extract<
AgentMessage,
{ role: "toolResult" }
>;
const id = extractToolResultId(toolResult);
if (id && toolCallIds.has(id)) {
if (seenToolResultIds.has(id)) {
continue;
}
if (!spanResultsById.has(id)) {
spanResultsById.set(id, toolResult);
}
continue;
}
}
remainder.push(next);
}
out.push(msg);
for (const call of toolCalls) {
const existing = spanResultsById.get(call.id);
pushToolResult(
existing ??
makeMissingToolResult({ toolCallId: call.id, toolName: call.name }),
);
}
for (const rem of remainder) {
if (!rem || typeof rem !== "object") {
out.push(rem);
continue;
}
const remRole = (rem as { role?: unknown }).role;
if (remRole === "toolResult") {
pushToolResult(rem as Extract<AgentMessage, { role: "toolResult" }>);
continue;
}
out.push(rem);
}
i = j - 1;
}
return out;
}