Skip to content

Commit b04dd6d

Browse files
steipetesteipete
committed
refactor: consolidate session history sanitization
1 parent 2373022 commit b04dd6d

12 files changed

Lines changed: 599 additions & 276 deletions

‎src/agents/pi-embedded-helpers.validate-turns.test.ts‎

Lines changed: 57 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@ function asMessages(messages: unknown[]): AgentMessage[] {
1212

1313
function makeDualToolUseAssistantContent() {
1414
return [
15-
{ type: "toolUse", id: "tool-1", name: "test1", input: {} },
16-
{ type: "toolUse", id: "tool-2", name: "test2", input: {} },
15+
{ type: "toolUse", id: "tool-1", name: "test1", arguments: {} },
16+
{ type: "toolUse", id: "tool-2", name: "test2", arguments: {} },
1717
{ type: "text", text: "Done" },
1818
];
1919
}
@@ -368,7 +368,7 @@ describe("validateAnthropicTurns strips dangling tool_use blocks", () => {
368368
{
369369
role: "assistant",
370370
content: [
371-
{ type: "toolUse", id: "tool-1", name: "test", input: {} },
371+
{ type: "toolUse", id: "tool-1", name: "test", arguments: {} },
372372
{ type: "text", text: "I'll check that" },
373373
],
374374
},
@@ -389,7 +389,7 @@ describe("validateAnthropicTurns strips dangling tool_use blocks", () => {
389389
{
390390
role: "assistant",
391391
content: [
392-
{ type: "toolUse", id: "tool-1", name: "test", input: {} },
392+
{ type: "toolUse", id: "tool-1", name: "test", arguments: {} },
393393
{ type: "text", text: "Here's result" },
394394
],
395395
},
@@ -408,7 +408,7 @@ describe("validateAnthropicTurns strips dangling tool_use blocks", () => {
408408
// tool_use should be preserved because matching tool_result exists
409409
const assistantContent = (result[1] as { content?: unknown[] }).content;
410410
expect(assistantContent).toEqual([
411-
{ type: "toolUse", id: "tool-1", name: "test", input: {} },
411+
{ type: "toolUse", id: "tool-1", name: "test", arguments: {} },
412412
{ type: "text", text: "Here's result" },
413413
]);
414414
});
@@ -418,7 +418,7 @@ describe("validateAnthropicTurns strips dangling tool_use blocks", () => {
418418
{ role: "user", content: [{ type: "text", text: "Use tool" }] },
419419
{
420420
role: "assistant",
421-
content: [{ type: "toolUse", id: "tool-1", name: "test", input: {} }],
421+
content: [{ type: "toolUse", id: "tool-1", name: "test", arguments: {} }],
422422
},
423423
{ role: "user", content: [{ type: "text", text: "Hello" }] },
424424
]);
@@ -431,6 +431,23 @@ describe("validateAnthropicTurns strips dangling tool_use blocks", () => {
431431
expect(assistantContent).toEqual([{ type: "text", text: "[tool calls omitted]" }]);
432432
});
433433

434+
it("leaves aborted tool-only assistant turns empty instead of synthesizing fallback text", () => {
435+
const msgs = asMessages([
436+
{ role: "user", content: [{ type: "text", text: "Use tool" }] },
437+
{
438+
role: "assistant",
439+
stopReason: "aborted",
440+
content: [{ type: "toolCall", id: "tool-1", name: "test", arguments: {} }],
441+
},
442+
{ role: "user", content: [{ type: "text", text: "Hello" }] },
443+
]);
444+
445+
const result = validateAnthropicTurns(msgs);
446+
447+
expect(result).toHaveLength(3);
448+
expect((result[1] as { content?: unknown[] }).content).toEqual([]);
449+
});
450+
434451
it("should handle multiple dangling tool_use blocks", () => {
435452
const msgs = makeDualToolAnthropicTurns([{ type: "text", text: "OK" }]);
436453

@@ -458,28 +475,54 @@ describe("validateAnthropicTurns strips dangling tool_use blocks", () => {
458475
// tool-1 should be preserved (has matching tool_result), tool-2 stripped, text preserved
459476
const assistantContent = (result[1] as { content?: unknown[] }).content;
460477
expect(assistantContent).toEqual([
461-
{ type: "toolUse", id: "tool-1", name: "test1", input: {} },
478+
{ type: "toolUse", id: "tool-1", name: "test1", arguments: {} },
462479
{ type: "text", text: "Done" },
463480
]);
464481
});
465482

466-
it("should not modify messages when next is not user", () => {
483+
it("matches standalone toolResult messages before the next assistant turn", () => {
467484
const msgs = asMessages([
468485
{ role: "user", content: [{ type: "text", text: "Use tool" }] },
469486
{
470487
role: "assistant",
471-
content: [{ type: "toolUse", id: "tool-1", name: "test", input: {} }],
488+
content: [{ type: "toolCall", id: "tool-1", name: "test", arguments: {} }],
472489
},
473-
// Next is assistant, not user - should not strip
474-
{ role: "assistant", content: [{ type: "text", text: "Continue" }] },
490+
{ role: "toolResult", toolCallId: "tool-1", content: [{ type: "text", text: "data" }] },
491+
{ role: "user", content: [{ type: "text", text: "Continue" }] },
475492
]);
476493

477494
const result = validateAnthropicTurns(msgs);
478495

479-
expect(result).toHaveLength(3);
480-
// Original tool_use should be preserved
496+
expect(result).toHaveLength(4);
481497
const assistantContent = (result[1] as { content?: unknown[] }).content;
482-
expect(assistantContent).toEqual([{ type: "toolUse", id: "tool-1", name: "test", input: {} }]);
498+
expect(assistantContent).toEqual([
499+
{ type: "toolCall", id: "tool-1", name: "test", arguments: {} },
500+
]);
501+
});
502+
503+
it("matches tool result blocks across intermediate non-assistant messages", () => {
504+
const msgs = asMessages([
505+
{ role: "user", content: [{ type: "text", text: "Use tool" }] },
506+
{
507+
role: "assistant",
508+
content: [
509+
{ type: "functionCall", id: "tool-1", name: "test", arguments: {} },
510+
{ type: "text", text: "Checking" },
511+
],
512+
},
513+
{ role: "user", content: [{ type: "text", text: "still waiting" }] },
514+
{ role: "tool", toolCallId: "tool-1", content: [{ type: "text", text: "data" }] },
515+
{ role: "user", content: [{ type: "text", text: "Continue" }] },
516+
]);
517+
518+
const result = validateAnthropicTurns(msgs);
519+
520+
expect(result).toHaveLength(5);
521+
const assistantContent = (result[1] as { content?: unknown[] }).content;
522+
expect(assistantContent).toEqual([
523+
{ type: "functionCall", id: "tool-1", name: "test", arguments: {} },
524+
{ type: "text", text: "Checking" },
525+
]);
483526
});
484527

485528
it("is replay-safe across repeated validation passes", () => {

‎src/agents/pi-embedded-helpers/turns.ts‎

Lines changed: 109 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,101 @@
11
import type { AgentMessage } from "@mariozechner/pi-agent-core";
2+
import { extractToolCallsFromAssistant, extractToolResultId } from "../tool-call-id.js";
23

34
type AnthropicContentBlock = {
4-
type: "text" | "toolUse" | "toolResult";
5+
type: "text" | "toolUse" | "toolCall" | "functionCall" | "toolResult" | "tool";
56
text?: string;
67
id?: string;
78
name?: string;
89
toolUseId?: string;
10+
toolCallId?: string;
911
};
1012

13+
function trimNonEmptyString(value: unknown): string | undefined {
14+
if (typeof value !== "string") {
15+
return undefined;
16+
}
17+
const trimmed = value.trim();
18+
return trimmed || undefined;
19+
}
20+
21+
function isToolCallBlock(block: AnthropicContentBlock): boolean {
22+
return block.type === "toolUse" || block.type === "toolCall" || block.type === "functionCall";
23+
}
24+
25+
function isAbortedAssistantTurn(message: AgentMessage): boolean {
26+
const stopReason = (message as { stopReason?: unknown }).stopReason;
27+
return stopReason === "aborted" || stopReason === "error";
28+
}
29+
30+
function extractToolResultIdsFromRecord(record: Record<string, unknown>): string[] {
31+
const ids = [
32+
trimNonEmptyString(record.toolUseId),
33+
trimNonEmptyString(record.toolCallId),
34+
trimNonEmptyString(record.tool_use_id),
35+
trimNonEmptyString(record.tool_call_id),
36+
trimNonEmptyString(record.callId),
37+
trimNonEmptyString(record.call_id),
38+
].filter((value): value is string => typeof value === "string");
39+
return [...new Set(ids)];
40+
}
41+
42+
function collectMatchingToolResultIds(message: AgentMessage): Set<string> {
43+
const ids = new Set<string>();
44+
const role = (message as { role?: unknown }).role;
45+
if (role === "toolResult") {
46+
const toolResultId = extractToolResultId(
47+
message as Extract<AgentMessage, { role: "toolResult" }>,
48+
);
49+
if (toolResultId) {
50+
ids.add(toolResultId);
51+
}
52+
} else if (role === "tool") {
53+
for (const id of extractToolResultIdsFromRecord(message as Record<string, unknown>)) {
54+
ids.add(id);
55+
}
56+
}
57+
58+
const content = (message as { content?: unknown }).content;
59+
if (!Array.isArray(content)) {
60+
return ids;
61+
}
62+
63+
for (const block of content) {
64+
if (!block || typeof block !== "object") {
65+
continue;
66+
}
67+
const record = block as Record<string, unknown>;
68+
if (record.type !== "toolResult" && record.type !== "tool") {
69+
continue;
70+
}
71+
for (const id of extractToolResultIdsFromRecord(record)) {
72+
ids.add(id);
73+
}
74+
}
75+
76+
return ids;
77+
}
78+
79+
function collectFutureToolResultIds(messages: AgentMessage[], startIndex: number): Set<string> {
80+
const ids = new Set<string>();
81+
for (let index = startIndex + 1; index < messages.length; index += 1) {
82+
const candidate = messages[index];
83+
if (!candidate || typeof candidate !== "object") {
84+
continue;
85+
}
86+
if ((candidate as { role?: unknown }).role === "assistant") {
87+
break;
88+
}
89+
for (const id of collectMatchingToolResultIds(candidate)) {
90+
ids.add(id);
91+
}
92+
}
93+
return ids;
94+
}
95+
1196
/**
12-
* Strips dangling tool_use blocks from assistant messages when the immediately
13-
* following user message does not contain a matching tool_result block.
97+
* Strips dangling tool-call blocks from assistant messages when no later
98+
* tool-result span before the next assistant turn resolves them.
1499
* This fixes the "tool_use ids found without tool_result blocks" error from Anthropic.
15100
*/
16101
function stripDanglingAnthropicToolUses(messages: AgentMessage[]): AgentMessage[] {
@@ -32,51 +117,42 @@ function stripDanglingAnthropicToolUses(messages: AgentMessage[]): AgentMessage[
32117
const assistantMsg = msg as {
33118
content?: AnthropicContentBlock[];
34119
};
35-
36-
// Get the next message to check for tool_result blocks
37-
const nextMsg = messages[i + 1];
38-
const nextMsgRole =
39-
nextMsg && typeof nextMsg === "object"
40-
? ((nextMsg as { role?: unknown }).role as string | undefined)
41-
: undefined;
42-
43-
// If next message is not user, keep the assistant message as-is
44-
if (nextMsgRole !== "user") {
120+
const originalContent = Array.isArray(assistantMsg.content) ? assistantMsg.content : [];
121+
if (originalContent.length === 0) {
45122
result.push(msg);
46123
continue;
47124
}
48-
49-
// Collect tool_use_ids from the next user message's tool_result blocks
50-
const nextUserMsg = nextMsg as {
51-
content?: AnthropicContentBlock[];
52-
};
53-
const validToolUseIds = new Set<string>();
54-
if (Array.isArray(nextUserMsg.content)) {
55-
for (const block of nextUserMsg.content) {
56-
if (block && block.type === "toolResult" && block.toolUseId) {
57-
validToolUseIds.add(block.toolUseId);
58-
}
59-
}
125+
if (
126+
extractToolCallsFromAssistant(msg as Extract<AgentMessage, { role: "assistant" }>).length ===
127+
0
128+
) {
129+
result.push(msg);
130+
continue;
60131
}
132+
const validToolUseIds = collectFutureToolResultIds(messages, i);
61133

62-
// Filter out tool_use blocks that don't have matching tool_result
63-
const originalContent = Array.isArray(assistantMsg.content) ? assistantMsg.content : [];
64134
const filteredContent = originalContent.filter((block) => {
65135
if (!block) {
66136
return false;
67137
}
68-
if (block.type !== "toolUse") {
138+
if (!isToolCallBlock(block)) {
69139
return true;
70140
}
71-
// Keep tool_use if its id is in the valid set
72-
return validToolUseIds.has(block.id || "");
141+
const blockId = trimNonEmptyString(block.id);
142+
return blockId ? validToolUseIds.has(blockId) : false;
73143
});
74144

75-
// If all content would be removed, insert a minimal fallback text block
145+
if (filteredContent.length === originalContent.length) {
146+
result.push(msg);
147+
continue;
148+
}
149+
76150
if (originalContent.length > 0 && filteredContent.length === 0) {
77151
result.push({
78152
...assistantMsg,
79-
content: [{ type: "text", text: "[tool calls omitted]" }],
153+
content: isAbortedAssistantTurn(msg)
154+
? []
155+
: ([{ type: "text", text: "[tool calls omitted]" }] as AnthropicContentBlock[]),
80156
} as AgentMessage);
81157
} else {
82158
result.push({
@@ -190,7 +266,7 @@ export function mergeConsecutiveUserTurns(
190266
* Also strips dangling tool_use blocks that lack corresponding tool_result blocks.
191267
*/
192268
export function validateAnthropicTurns(messages: AgentMessage[]): AgentMessage[] {
193-
// First, strip dangling tool_use blocks from assistant messages
269+
// First, strip dangling tool-call blocks from assistant messages.
194270
const stripped = stripDanglingAnthropicToolUses(messages);
195271

196272
return validateTurnsWithConsecutiveMerge({

‎src/agents/tools/chat-history-text.ts‎

Lines changed: 33 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -31,36 +31,52 @@ export function sanitizeTextContent(text: string): string {
3131
);
3232
}
3333

34-
export function extractAssistantText(message: unknown): string | undefined {
34+
export function hasAssistantPhaseMetadata(message: unknown): boolean {
3535
if (!message || typeof message !== "object") {
36-
return undefined;
36+
return false;
3737
}
3838
if ((message as { role?: unknown }).role !== "assistant") {
39-
return undefined;
39+
return false;
4040
}
4141
const content = (message as { content?: unknown }).content;
4242
if (!Array.isArray(content)) {
43-
return undefined;
43+
return false;
4444
}
45-
const hasPhaseMetadata = content.some(
45+
return content.some(
4646
(block) =>
47-
block && typeof block === "object" && typeof (block as { textSignature?: unknown }).textSignature === "string",
47+
block &&
48+
typeof block === "object" &&
49+
typeof (block as { textSignature?: unknown }).textSignature === "string",
4850
);
49-
const joined = hasPhaseMetadata
50-
? (extractAssistantVisibleText(message as Parameters<typeof extractAssistantVisibleText>[0]) ?? "")
51-
: (
52-
extractTextFromChatContent(content, {
53-
sanitizeText: sanitizeTextContent,
54-
joinWith: "",
55-
normalizeText: (text) => text.trim(),
56-
}) ?? ""
51+
}
52+
53+
export function extractAssistantText(message: unknown): string | undefined {
54+
if (!message || typeof message !== "object") {
55+
return undefined;
56+
}
57+
if ((message as { role?: unknown }).role !== "assistant") {
58+
return undefined;
59+
}
60+
if (hasAssistantPhaseMetadata(message)) {
61+
const visibleText = extractAssistantVisibleText(
62+
message as Parameters<typeof extractAssistantVisibleText>[0],
63+
);
64+
return visibleText?.trim() ? visibleText : undefined;
65+
}
66+
const content = (message as { content?: unknown }).content;
67+
const joined = Array.isArray(content)
68+
? (extractTextFromChatContent(content, {
69+
sanitizeText: sanitizeTextContent,
70+
joinWith: "",
71+
normalizeText: (text) => text.trim(),
72+
}) ?? "")
73+
: sanitizeTextContent(
74+
extractAssistantVisibleText(message as Parameters<typeof extractAssistantVisibleText>[0]) ??
75+
"",
5776
);
5877
if (!joined.trim()) {
5978
return undefined;
6079
}
61-
if (hasPhaseMetadata) {
62-
return joined;
63-
}
6480
const stopReason = (message as { stopReason?: unknown }).stopReason;
6581
// Gate on stopReason only — a non-error response with a stale/background errorMessage
6682
// should not have its content rewritten with error templates (#13935).

0 commit comments

Comments
 (0)