Skip to content

Commit 898ec06

Browse files
canblmz1Cantaltas
authored
fix: clear nativeArgs when tool-call finalize fails (#1221) (#1634)
* fix(task): clear nativeArgs when tool-call finalize fails (#1221) finalizeStreamingToolCall() returns null when a streamed native tool call's arguments are truncated (e.g. the model hits max_tokens mid write_to_file content). Task.ts reuses the same tool-use object the streaming phase had been mutating in place, which still carried nativeArgs built from the incomplete partial-JSON parse, and only set partial = false. presentAssistantMessage.ts already guards against exactly this case (isKnownTool && !block.nativeArgs && !customTool -> structured tool_result instead of execution), but the guard never fired because nativeArgs was never actually cleared. Truncated arguments (e.g. a cut-off content string) could therefore be executed instead of rejected. Clear nativeArgs alongside partial = false at the finalize-null site so the existing guard does what its own comment already said it did. params is left untouched - NativeToolCallParser always initializes it to {} for native tool calls and never puts real data there. Adds truncated-native-tool-args.spec.ts, mirroring the exact Task.ts logic in a small local function per the convention already established in duplicate-tool-use-ids.spec.ts. * test(task): add integration coverage for truncated write_to_file finalize (#1221) Adds an integration-level regression test alongside the existing simulation-based unit tests in truncated-native-tool-args.spec.ts. Drives a truncated write_to_file tool call through the real Task streaming + presentAssistantMessage flow (via recursivelyMakeClineRequests and a mocked attemptApiRequest stream), rather than mirroring the finalize-null logic in an isolated function. Spies on writeToFileTool.handle to confirm it is never invoked with partial: false (the flag that gates real execute()/disk-write behavior in BaseTool.handle) for the truncated call, and spies on pushToolResultToUserContent to confirm the guard's structured error result is emitted instead. Verified this only fails for the intended reason: temporarily reverting the Task.ts fix makes writeToFileTool.handle get called with partial: false (real execution attempted) - confirmed via the test's own failure output, not assumed. An earlier version of this test used a .json target path and was inconclusive, since Architect mode's markdown-only file restriction independently blocked the write before ever reaching the nativeArgs guard; switched to a .md path so the guard under test is what's actually being exercised. Full core/task suite: 27 files, 382 tests, all passing. * fix(task): also clear params on truncated tool-call finalize (#1221) Clearing nativeArgs alone left a second gap: Task.ts records each tool_use into apiConversationHistory via toolUse.nativeArgs || toolUse.params, and params also gets populated during streaming (for handlePartial's UI hooks) with the same incomplete, truncated values - so once nativeArgs was cleared, the fallback just picked up the same truncated data from params instead. Clear params to {} alongside nativeArgs so the history entry for a truncated call doesn't carry that data under a different field. Execution itself was never at risk from this gap - BaseTool.handle only uses nativeArgs for execute(), and throws instead of falling back to params when nativeArgs is missing - but the conversation history leak was real. Adds an assertion to the Task.spec.ts integration test confirming the truncated content doesn't end up in the recorded assistant turn. * test(task): update safeEnsureModelFetched mock after rebase safeEnsureModelFetched's return type changed to Promise<ModelInfo> upstream while this branch was in flight; the new integration test still had mockResolvedValue(undefined) from before that change, which the type checker caught after rebasing onto main. Switched to the same stubModelInfo fixture the other tests in this file already use. --------- Co-authored-by: Can <[email protected]> Co-authored-by: T <[email protected]>
1 parent eb83244 commit 898ec06

3 files changed

Lines changed: 223 additions & 3 deletions

File tree

‎src/core/task/Task.ts‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3826,12 +3826,22 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
38263826
/* v8 ignore next -- streaming presenter; .catch lives in presentAssistantMessageSafe (covered) */
38273827
this.presentAssistantMessageSafe()
38283828
} else if (toolUseIndex !== undefined) {
3829-
// finalizeStreamingToolCall returned null (malformed JSON or missing args)
3830-
// We still need to mark the tool as non-partial so it gets executed
3831-
// The tool's validation will catch any missing required parameters
3829+
// finalizeStreamingToolCall returned null (malformed JSON or missing args).
3830+
// existingToolUse is the same object the streaming phase was mutating in
3831+
// place, so it still carries nativeArgs AND params built from the incomplete
3832+
// partial parse (e.g. a truncated write_to_file `content` string) - both were
3833+
// only ever meant for live progress display, never for execution or for
3834+
// ending up in conversation history. Mark the tool as non-partial so it's
3835+
// presented as complete, and clear both so presentAssistantMessage's
3836+
// `!block.nativeArgs` guard short-circuits with a structured tool_result
3837+
// instead of executing the truncated value, and so the toolUse.nativeArgs ||
3838+
// toolUse.params fallback used when recording history doesn't fall through to
3839+
// the same truncated data under a different name.
38323840
const existingToolUse = this.assistantMessageContent[toolUseIndex]
38333841
if (existingToolUse && existingToolUse.type === "tool_use") {
38343842
existingToolUse.partial = false
3843+
existingToolUse.nativeArgs = undefined
3844+
existingToolUse.params = {}
38353845
// Ensure it has the ID for native protocol
38363846
;(existingToolUse as any).id = event.id
38373847
}

‎src/core/task/__tests__/Task.spec.ts‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import type { ApiMessage } from "../../task-persistence"
3333
import { asyncStreamFrom } from "../../../test-utils/stream"
3434
import { McpHub } from "../../../services/mcp/McpHub"
3535
import { McpServerManager } from "../../../services/mcp/McpServerManager"
36+
import { writeToFileTool } from "../../tools/WriteToFileTool"
3637

3738
type TaskTestAccess = {
3839
getSystemPrompt: (requestState: ProviderState | undefined, requestModelInfo?: ModelInfo) => Promise<string>
@@ -750,6 +751,102 @@ describe("Cline", () => {
750751
},
751752
])
752753
})
754+
755+
it("blocks a truncated write_to_file call instead of executing it (issue #1221)", async () => {
756+
// Regression test for #1221: if the model's stream is cut off mid-way
757+
// through a write_to_file tool call's `content` argument (e.g. it hits
758+
// max_tokens), finalizeStreamingToolCall() can't parse the incomplete
759+
// JSON and returns null. Task.ts must not let the truncated content
760+
// reach writeToFileTool's execution path - it must clear nativeArgs so
761+
// presentAssistantMessage's fail-closed guard emits a structured
762+
// tool_result error instead.
763+
//
764+
// Unlike the simulation-based tests in truncated-native-tool-args.spec.ts,
765+
// this drives the real streaming + presentAssistantMessage flow through
766+
// Task, and spies on the actual tool handler to prove it is never invoked.
767+
const task = new Task({
768+
provider: mockProvider,
769+
apiConfiguration: mockApiConfig,
770+
task: "truncated tool call test",
771+
startTask: false,
772+
})
773+
774+
vi.spyOn(task.diffViewProvider, "reset").mockResolvedValue(undefined)
775+
vi.spyOn(getTaskTestAccess(task), "safeEnsureModelFetched").mockResolvedValue(stubModelInfo)
776+
// presentAssistantMessageSafe is intentionally left un-mocked here (unlike
777+
// the other tests in this block) - the whole point is to exercise the real
778+
// dispatch/guard logic, not just tool_use finalization.
779+
780+
const writeToFileHandleSpy = vi.spyOn(writeToFileTool, "handle")
781+
// Spy directly on the guard's own push, rather than inspecting
782+
// userMessageContent/apiConversationHistory afterwards - the task
783+
// recurses into a follow-up request once the tool_result is ready
784+
// (see the second mocked stream below), which resets those arrays
785+
// for the new turn before this function returns.
786+
const pushToolResultSpy = vi.spyOn(task, "pushToolResultToUserContent")
787+
788+
vi.spyOn(task, "attemptApiRequest")
789+
.mockImplementationOnce(() =>
790+
asyncStreamFrom<ApiStreamChunk>([
791+
{
792+
type: "tool_call_partial",
793+
index: 0,
794+
id: "call_truncated",
795+
name: "write_to_file",
796+
},
797+
{
798+
type: "tool_call_partial",
799+
index: 0,
800+
// Cut off mid-string: no closing quote/brace, and the stream
801+
// ends here with no explicit tool_call_end - exactly what
802+
// happens when the model hits max_tokens mid-argument.
803+
// .md path deliberately used so the only thing that can block
804+
// execution is the nativeArgs guard under test - an arbitrary
805+
// extension could also get caught by unrelated mode-based file
806+
// restrictions (e.g. Architect mode's markdown-only rule),
807+
// which would produce a false pass/fail unrelated to this bug.
808+
arguments: '{"path":"docs/config.md","content":"sk-live-abc123',
809+
},
810+
]),
811+
)
812+
// The task recurses once the error tool_result makes the turn
813+
// "ready" - this bounds that follow-up to a single harmless text
814+
// reply instead of an unmocked second call.
815+
.mockImplementationOnce(() => asyncStreamFrom<ApiStreamChunk>([{ type: "text", text: "" }]))
816+
817+
await task.recursivelyMakeClineRequests([{ type: "text", text: "truncated tool call test" }])
818+
819+
// handle() legitimately gets called with partial: true while the call is
820+
// still streaming (BaseTool.handle short-circuits to a no-op preview hook
821+
// in that case) - that's expected and safe. What must never happen is a
822+
// call with partial: false, which is what actually reaches execute() and
823+
// writes to disk.
824+
const nonPartialCalls = writeToFileHandleSpy.mock.calls.filter(
825+
([, block]) => (block as { partial?: boolean }).partial === false,
826+
)
827+
expect(nonPartialCalls).toHaveLength(0)
828+
829+
// Neither nativeArgs nor params should leak the truncated content into
830+
// the recorded assistant turn - Task.ts builds that entry's `input` via
831+
// `toolUse.nativeArgs || toolUse.params`, so clearing nativeArgs alone
832+
// would just have shifted the leak to params instead of closing it.
833+
const assistantEntry = task.apiConversationHistory.find(
834+
(m) => m.role === "assistant" && Array.isArray(m.content) && m.content[0]?.type === "tool_use",
835+
)
836+
expect(JSON.stringify(assistantEntry)).not.toContain("sk-live-abc123")
837+
838+
// A structured, matching tool_result error must have been pushed for
839+
// the truncated call's ID instead of letting it execute.
840+
const truncatedCallResult = pushToolResultSpy.mock.calls.find(
841+
([result]) => result.tool_use_id === "call_truncated",
842+
)?.[0]
843+
expect(truncatedCallResult).toMatchObject({
844+
type: "tool_result",
845+
tool_use_id: "call_truncated",
846+
is_error: true,
847+
})
848+
expect(JSON.stringify(truncatedCallResult)).toContain("missing nativeArgs")
849+
})
753850
})
754851

755852
describe("constructor", () => {

‎src/core/task/__tests__/truncated-native-tool-args.spec.ts‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/**
2+
* Regression test for issue #1221: truncated tool-call arguments can be silently
3+
* written to disk.
4+
*
5+
* When a streamed native tool call's arguments are cut off mid-value (e.g. the
6+
* model hits max_tokens while still writing write_to_file's `content` string),
7+
* NativeToolCallParser.finalizeStreamingToolCall() returns null. Task.ts
8+
* (~line 3748) reuses the same tool-use object the streaming phase was mutating
9+
* in place and only sets `partial = false` - before the fix it left `nativeArgs`
10+
* (built from the incomplete partial parse) untouched.
11+
*
12+
* presentAssistantMessage.ts (~line 443) is supposed to short-circuit exactly
13+
* this case with a structured tool_result instead of executing the tool - but
14+
* its guard is `isKnownTool && !block.nativeArgs && !customTool`. With
15+
* nativeArgs still populated, the guard never fired and the truncated content
16+
* would be passed straight to write_to_file's execution path.
17+
*
18+
* The fix clears `existingToolUse.nativeArgs` alongside `partial = false` at
19+
* the finalize-null site, so the pre-existing guard actually does what its own
20+
* comment already claimed.
21+
*/
22+
23+
import { isValidToolName } from "../../tools/validateToolUse"
24+
import type { ToolUse, WriteToFileToolUse } from "../../../shared/tools"
25+
26+
describe("Truncated native tool-call args on finalize failure (issue #1221)", () => {
27+
/**
28+
* Simulates the finalize-null branch from Task.ts (~line 3748) as it exists
29+
* after the fix: on finalizeStreamingToolCall() returning null, mark the
30+
* tool non-partial and clear nativeArgs.
31+
*/
32+
function finalizeNullBranch(existingToolUse: ToolUse): ToolUse {
33+
existingToolUse.partial = false
34+
existingToolUse.nativeArgs = undefined
35+
return existingToolUse
36+
}
37+
38+
/**
39+
* Simulates the finalize-null branch as it existed *before* the fix, for a
40+
* companion test proving the old behavior really was the bug (not just an
41+
* assumption).
42+
*/
43+
function finalizeNullBranchBeforeFix(existingToolUse: ToolUse): ToolUse {
44+
existingToolUse.partial = false
45+
return existingToolUse
46+
}
47+
48+
/**
49+
* Simulates the short-circuit guard from presentAssistantMessage.ts (~line
50+
* 443): `isKnownTool && !block.nativeArgs && !customTool`. Returns true when
51+
* the tool call would be blocked (a structured tool_result emitted, no
52+
* execution), false when it would proceed to execution.
53+
*/
54+
function wouldBeBlocked(block: ToolUse, customTool: unknown = undefined): boolean {
55+
const isKnownTool = isValidToolName(String(block.name))
56+
return Boolean(isKnownTool && !block.nativeArgs && !customTool)
57+
}
58+
59+
it("clears nativeArgs so a truncated write_to_file call is blocked instead of executed", () => {
60+
// A write_to_file call whose `content` was cut off mid-stream - exactly
61+
// the scenario in #1221. The streaming phase already populated nativeArgs
62+
// from the incomplete partial-json parse before finalize failed.
63+
const truncated: WriteToFileToolUse = {
64+
type: "tool_use",
65+
name: "write_to_file",
66+
params: {},
67+
partial: true,
68+
nativeArgs: { path: "src/config.json", content: '{"apiKey": "sk-live-abc123' /* cut off mid-string */ },
69+
}
70+
71+
finalizeNullBranch(truncated)
72+
73+
expect(truncated.partial).toBe(false)
74+
expect(truncated.nativeArgs).toBeUndefined()
75+
expect(wouldBeBlocked(truncated)).toBe(true)
76+
})
77+
78+
it("companion: without the fix, the same truncated call would NOT have been blocked", () => {
79+
const truncated: WriteToFileToolUse = {
80+
type: "tool_use",
81+
name: "write_to_file",
82+
params: {},
83+
partial: true,
84+
nativeArgs: { path: "src/config.json", content: '{"apiKey": "sk-live-abc123' },
85+
}
86+
87+
finalizeNullBranchBeforeFix(truncated)
88+
89+
// This is the bug: partial is false (presented as "complete"), but
90+
// nativeArgs still carries the truncated value, so the guard's
91+
// `!block.nativeArgs` never becomes true and the call would proceed to
92+
// execution with the truncated content.
93+
expect(truncated.partial).toBe(false)
94+
expect(truncated.nativeArgs).toEqual({ path: "src/config.json", content: '{"apiKey": "sk-live-abc123' })
95+
expect(wouldBeBlocked(truncated)).toBe(false)
96+
})
97+
98+
it("does not affect a normally-finalized (non-null) tool call", () => {
99+
// When finalizeStreamingToolCall() succeeds, Task.ts replaces the block
100+
// with the freshly-finalized one instead of taking this branch at all -
101+
// this test just confirms a complete, valid nativeArgs is never touched
102+
// by wouldBeBlocked's guard simulation.
103+
const complete: WriteToFileToolUse = {
104+
type: "tool_use",
105+
name: "write_to_file",
106+
params: {},
107+
partial: false,
108+
nativeArgs: { path: "src/config.json", content: '{"apiKey": "sk-live-abc123xyz"}' },
109+
}
110+
111+
expect(wouldBeBlocked(complete)).toBe(false)
112+
})
113+
})

0 commit comments

Comments
 (0)