diff --git a/src/agents/pi-embedded-runner/run/attempt.test.ts b/src/agents/pi-embedded-runner/run/attempt.test.ts index 33a4f9654..020372122 100644 --- a/src/agents/pi-embedded-runner/run/attempt.test.ts +++ b/src/agents/pi-embedded-runner/run/attempt.test.ts @@ -560,6 +560,51 @@ describe("wrapStreamFnRepairMalformedToolCallArguments", () => { expect(partialToolCall.arguments).toEqual({}); expect(streamedToolCall.arguments).toEqual({}); }); + + it("clears a cached repair when later deltas make the trailing suffix invalid", async () => { + const partialToolCall = { type: "toolCall", name: "read", arguments: {} }; + const streamedToolCall = { type: "toolCall", name: "read", arguments: {} }; + const partialMessage = { role: "assistant", content: [partialToolCall] }; + const baseFn = vi.fn(() => + createFakeStream({ + events: [ + { + type: "toolcall_delta", + contentIndex: 0, + delta: '{"path":"/tmp/report.txt"}', + partial: partialMessage, + }, + { + type: "toolcall_delta", + contentIndex: 0, + delta: "x", + partial: partialMessage, + }, + { + type: "toolcall_delta", + contentIndex: 0, + delta: "yzq", + partial: partialMessage, + }, + { + type: "toolcall_end", + contentIndex: 0, + toolCall: streamedToolCall, + partial: partialMessage, + }, + ], + resultMessage: { role: "assistant", content: [partialToolCall] }, + }), + ); + + const stream = await invokeWrappedStream(baseFn); + for await (const _item of stream) { + // drain + } + + expect(partialToolCall.arguments).toEqual({}); + expect(streamedToolCall.arguments).toEqual({}); + }); }); describe("isOllamaCompatProvider", () => { diff --git a/src/agents/pi-embedded-runner/run/attempt.ts b/src/agents/pi-embedded-runner/run/attempt.ts index 790323b82..2f77b46af 100644 --- a/src/agents/pi-embedded-runner/run/attempt.ts +++ b/src/agents/pi-embedded-runner/run/attempt.ts @@ -557,6 +557,25 @@ function repairToolCallArgumentsInMessage( typedBlock.arguments = repairedArgs; } +function clearToolCallArgumentsInMessage(message: unknown, contentIndex: number): void { + if (!message || typeof message !== "object") { + return; + } + const content = (message as { content?: unknown }).content; + if (!Array.isArray(content)) { + return; + } + const block = content[contentIndex]; + if (!block || typeof block !== "object") { + return; + } + const typedBlock = block as { type?: unknown; arguments?: unknown }; + if (!isToolCallBlockType(typedBlock.type)) { + return; + } + typedBlock.arguments = {}; +} + function repairMalformedToolCallArgumentsInMessage( message: unknown, repairedArgsByIndex: Map>, @@ -637,6 +656,10 @@ function wrapStreamRepairMalformedToolCallArguments( `repairing kimi-coding tool call arguments after ${repair.trailingSuffix.length} trailing chars`, ); } + } else { + repairedArgsByIndex.delete(event.contentIndex); + clearToolCallArgumentsInMessage(event.partial, event.contentIndex); + clearToolCallArgumentsInMessage(event.message, event.contentIndex); } } }