refactor(agents): extract transcript repair module
This commit is contained in:
@@ -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");
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
95
src/agents/session-transcript-repair.test.ts
Normal file
95
src/agents/session-transcript-repair.test.ts
Normal 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");
|
||||
});
|
||||
});
|
||||
173
src/agents/session-transcript-repair.ts
Normal file
173
src/agents/session-transcript-repair.ts
Normal 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;
|
||||
}
|
||||
Reference in New Issue
Block a user