Skip to content

Commit 86fe9d5

Browse files
vincentkocvincentkoc
committed
fix(agents): harden partial tool call replay
1 parent 8c17be2 commit 86fe9d5

2 files changed

Lines changed: 91 additions & 33 deletions

File tree

‎src/agents/session-transcript-repair.test.ts‎

Lines changed: 61 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -784,24 +784,30 @@ describe("sanitizeToolCallInputs allowed-name filtering", () => {
784784
expect(ids).toEqual(expectedIds);
785785
});
786786

787-
it("keeps finalized OpenAI Responses blocks and drops interrupted partialJson artifacts", () => {
787+
it("keeps finalized OpenAI Responses calls and drops partialJson streaming artifacts", () => {
788788
const input = castAgentMessages([
789789
{
790790
role: "assistant",
791-
stopReason: "stop",
791+
stopReason: "toolUse",
792792
content: [
793793
// complete tool call — kept as-is
794794
{ type: "toolCall", id: "call_ok", name: "read", arguments: { path: "/a" } },
795-
// Finalized OpenAI Responses blocks keep parsed arguments plus
796-
// partialJson in persisted history; repair should strip only
797-
// partialJson and keep the finished tool call.
795+
// Legacy generic Responses transport persisted finalized toolUse
796+
// turns with partialJson; repair strips the scratch field.
798797
{
799798
type: "toolCall",
800799
id: "call_partial|fc_123",
801800
name: "Bash",
802801
arguments: { command: "ls" },
803802
partialJson: '{"command": "ls"}',
804803
},
804+
{
805+
type: "toolCall",
806+
id: "call_empty|fc_789",
807+
name: "session_status",
808+
arguments: {},
809+
partialJson: "",
810+
},
805811
// Anthropic can persist initialized tool calls with arguments: {}
806812
// plus partialJson if the stream aborts before content_block_stop.
807813
// Those incomplete artifacts must be dropped.
@@ -812,6 +818,15 @@ describe("sanitizeToolCallInputs allowed-name filtering", () => {
812818
arguments: {},
813819
partialJson: '{"command":',
814820
},
821+
// An OpenAI-shaped id and parsed partial arguments do not prove that
822+
// response.output_item.done arrived.
823+
{
824+
type: "toolCall",
825+
id: "call_truncated|fc_456",
826+
name: "Bash",
827+
arguments: { command: "ls" },
828+
partialJson: '{"command":"ls"',
829+
},
815830
// Missing required input is also an interrupted artifact and should drop.
816831
{
817832
type: "toolUse",
@@ -827,16 +842,16 @@ describe("sanitizeToolCallInputs allowed-name filtering", () => {
827842
const out = sanitizeToolCallInputs(input);
828843
const toolCalls = getAssistantToolCallBlocks(out);
829844
const ids = toolCalls.map((t) => (t as { id?: unknown }).id);
830-
expect(ids).toEqual(["call_ok", "call_partial|fc_123"]);
831-
const keptPartial = toolCalls.find((t) => (t as { id?: unknown }).id === "call_partial|fc_123");
832-
expect(keptPartial).not.toHaveProperty("partialJson");
845+
expect(ids).toEqual(["call_ok", "call_partial|fc_123", "call_empty|fc_789"]);
846+
expect(toolCalls[1]).not.toHaveProperty("partialJson");
847+
expect(toolCalls[2]).not.toHaveProperty("partialJson");
833848
});
834849

835-
it("strips partialJson and preserves sessions_spawn attachment content", () => {
850+
it("strips finalized partialJson without rewriting sessions_spawn arguments", () => {
836851
const input = castAgentMessages([
837852
{
838853
role: "assistant",
839-
stopReason: "stop",
854+
stopReason: "toolUse",
840855
content: [
841856
{
842857
type: "toolCall",
@@ -852,15 +867,13 @@ describe("sanitizeToolCallInputs allowed-name filtering", () => {
852867
const out = sanitizeToolCallInputs(input);
853868
const toolCalls = getAssistantToolCallBlocks(out);
854869
expect(toolCalls).toHaveLength(1);
855-
const spawn = toolCalls[0] as { id?: unknown; arguments?: unknown };
856-
// partialJson must be stripped
857-
expect(spawn).not.toHaveProperty("partialJson");
858-
// sessions_spawn attachment content must be preserved on current main.
859-
const args = spawn.arguments as { attachments?: Array<{ content?: unknown }> };
860-
expect(args?.attachments?.[0]?.content).toBe("secret data");
870+
expect(toolCalls[0]).not.toHaveProperty("partialJson");
871+
expect((toolCalls[0] as { arguments?: unknown }).arguments).toEqual({
872+
attachments: [{ content: "secret data" }],
873+
});
861874
});
862875

863-
it.each(["aborted", "error"] as const)(
876+
it.each(["stop", "aborted", "error", "length"] as const)(
864877
"drops OpenAI Responses partialJson blocks on %s assistant turns",
865878
(stopReason) => {
866879
const input = castAgentMessages([
@@ -873,7 +886,7 @@ describe("sanitizeToolCallInputs allowed-name filtering", () => {
873886
id: "call_partial|fc_123",
874887
name: "Bash",
875888
arguments: { command: "ls" },
876-
partialJson: '{"command":"ls"',
889+
partialJson: '{"command":"ls"}',
877890
},
878891
],
879892
},
@@ -936,6 +949,36 @@ describe("sanitizeToolCallInputs allowed-name filtering", () => {
936949
expect(out).toStrictEqual([]);
937950
});
938951

952+
it("drops signed-thinking assistant turns with partialJson tool calls", () => {
953+
const input = castAgentMessages([
954+
{
955+
role: "assistant",
956+
stopReason: "toolUse",
957+
content: [
958+
{
959+
type: "thinking",
960+
thinking: "Let me run a command.",
961+
thinkingSignature: "sig_partial",
962+
},
963+
{
964+
type: "toolCall",
965+
id: "call_partial|fc_123",
966+
name: "exec",
967+
arguments: {},
968+
partialJson: '{"command":"ls"}',
969+
},
970+
],
971+
},
972+
]);
973+
974+
const out = sanitizeToolCallInputs(input, {
975+
allowedToolNames: ["exec"],
976+
allowProviderOwnedThinkingReplay: true,
977+
});
978+
979+
expect(out).toStrictEqual([]);
980+
});
981+
939982
it("drops signed-thinking assistant turns when sibling tool calls reuse an id", () => {
940983
const input = castAgentMessages([
941984
{

‎src/agents/session-transcript-repair.ts‎

Lines changed: 30 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -80,23 +80,37 @@ function hasPartialJson(
8080
return typeof block.partialJson === "string";
8181
}
8282

83-
function isFinalizedOpenAIResponsesToolCall(block: RawToolCallBlock): boolean {
84-
if (!hasPartialJson(block) || typeof block.id !== "string" || "input" in block) {
83+
function isCompleteJsonObject(value: string): boolean {
84+
try {
85+
const parsed: unknown = JSON.parse(value);
86+
return parsed !== null && typeof parsed === "object" && !Array.isArray(parsed);
87+
} catch {
8588
return false;
8689
}
90+
}
8791

88-
const separator = block.id.indexOf("|");
89-
if (separator <= 0 || separator === block.id.length - 1) {
92+
function isFinalizedOpenAIResponsesToolCall(
93+
message: AgentMessage,
94+
block: RawToolCallBlock,
95+
): boolean {
96+
if (
97+
message.role !== "assistant" ||
98+
!("stopReason" in message) ||
99+
message.stopReason !== "toolUse" ||
100+
!hasPartialJson(block) ||
101+
typeof block.id !== "string" ||
102+
"input" in block ||
103+
!block.arguments ||
104+
typeof block.arguments !== "object" ||
105+
Array.isArray(block.arguments) ||
106+
(!isCompleteJsonObject(block.partialJson) &&
107+
(block.partialJson.trim() !== "" || Object.keys(block.arguments).length > 0))
108+
) {
90109
return false;
91110
}
92111

93-
return "arguments" in block && block.arguments !== undefined && block.arguments !== null;
94-
}
95-
function isInterruptedAssistantTurn(message: AgentMessage): boolean {
96-
if (message.role !== "assistant" || !("stopReason" in message)) {
97-
return false;
98-
}
99-
return message.stopReason === "aborted" || message.stopReason === "error";
112+
const separator = block.id.indexOf("|");
113+
return separator > 0 && separator < block.id.length - 1;
100114
}
101115

102116
function sanitizeToolCallBlock(block: RawToolCallBlock): RawToolCallBlock {
@@ -143,6 +157,7 @@ function isReplaySafeThinkingAssistantTurn(
143157
const toolCallId = typeof block.id === "string" ? block.id.trim() : "";
144158
if (
145159
!hasToolCallInput(block) ||
160+
hasPartialJson(block) ||
146161
!toolCallId ||
147162
seenToolCallIds.has(toolCallId) ||
148163
!isAllowedToolCallName(block.name, allowedToolNames)
@@ -425,17 +440,17 @@ function repairToolCallInputs(
425440
}
426441
let workBlock = block;
427442
if (isRawToolCallBlock(block) && hasPartialJson(block)) {
428-
if (isInterruptedAssistantTurn(msg) || !isFinalizedOpenAIResponsesToolCall(block)) {
443+
if (!isFinalizedOpenAIResponsesToolCall(msg, block)) {
429444
droppedToolCalls += 1;
430445
droppedInMessage += 1;
431446
changed = true;
432447
messageChanged = true;
433448
continue;
434449
}
435450

436-
// OpenAI Responses persists finalized function calls with both parsed
437-
// arguments and the original partialJson bytes. Strip only the
438-
// redundant partialJson field so replay keeps the finalized call.
451+
// Legacy generic Responses transport persisted successful toolUse turns
452+
// with the scratch buffer intact. Strip it only when terminal state and
453+
// the provider-specific finalized shape both prove completion.
439454
const stripped = { ...block };
440455
delete (stripped as RawToolCallBlock & { partialJson?: unknown }).partialJson;
441456
workBlock = stripped;

0 commit comments

Comments
 (0)