Skip to content

Commit 75359f3

Browse files
steipetesteipete
committed
fix: harden smart-quoted argument repair (#86611)
1 parent 20b402d commit 75359f3

2 files changed

Lines changed: 177 additions & 16 deletions

File tree

‎src/agents/embedded-agent-runner/run/attempt.tool-call-argument-repair.test.ts‎

Lines changed: 90 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ async function runToolCallRepairCase(params: {
6161
delta: string;
6262
provider?: string;
6363
modelApi?: string;
64+
includePreamble?: boolean;
65+
preambleToolName?: string;
6466
}): Promise<ToolCallRepairCaseResult> {
6567
const toolName = params.toolName ?? "write";
6668
const partialToolCall = { type: "functionCall", name: toolName, arguments: {} };
@@ -77,12 +79,16 @@ async function runToolCallRepairCase(params: {
7779
baseFn: () =>
7880
createFakeStream({
7981
events: [
80-
{
81-
type: "toolcall_delta",
82-
contentIndex: 0,
83-
delta: `.functions.${toolName}:0 `,
84-
partial: partialMessage,
85-
},
82+
...(params.includePreamble === false
83+
? []
84+
: [
85+
{
86+
type: "toolcall_delta",
87+
contentIndex: 0,
88+
delta: `.functions.${params.preambleToolName ?? toolName}:0 `,
89+
partial: partialMessage,
90+
},
91+
]),
8692
{
8793
type: "toolcall_delta",
8894
contentIndex: 0,
@@ -327,6 +333,84 @@ const re = /\d+/;
327333
expectAllToolCallArgs(result, { path: "safe.txt" });
328334
});
329335

336+
it("repairs smart-quoted non-freeform args before schema-specific option keys", async () => {
337+
const result = await runToolCallRepairCase({
338+
toolName: "read",
339+
delta: "{“path”:“safe.txt”,“offset”:5,“limit”:20}",
340+
});
341+
342+
expectAllToolCallArgs(result, { path: "safe.txt", offset: 5, limit: 20 });
343+
});
344+
345+
it("repairs prefixless smart-quoted read args before schema-specific option keys", async () => {
346+
const result = await runToolCallRepairCase({
347+
toolName: "read",
348+
delta: "{“path”:“safe.txt”,“offset”:5,“limit”:20}",
349+
includePreamble: false,
350+
});
351+
352+
expectAllToolCallArgs(result, { path: "safe.txt", offset: 5, limit: 20 });
353+
});
354+
355+
it("repairs smart-quoted read args with a case-varied structured tool name", async () => {
356+
const result = await runToolCallRepairCase({
357+
toolName: "Read",
358+
delta: "{“path”:“safe.txt”,“offset”:5,“limit”:20}",
359+
includePreamble: false,
360+
});
361+
362+
expectAllToolCallArgs(result, { path: "safe.txt", offset: 5, limit: 20 });
363+
});
364+
365+
it("keeps unknown member-looking prose inside smart-quoted non-freeform args", async () => {
366+
const result = await runToolCallRepairCase({
367+
toolName: "grep",
368+
delta: String.raw` {“pattern”:“Use ”, “foo”: “bar” in prose”,“path”:“safe.txt”}`,
369+
});
370+
371+
expectAllToolCallArgs(result, {
372+
pattern: "Use ”, “foo”: “bar” in prose",
373+
path: "safe.txt",
374+
});
375+
expect(result.finalArgs).not.toHaveProperty("foo");
376+
});
377+
378+
it("keeps known option-looking prose inside unrelated smart-quoted args", async () => {
379+
const result = await runToolCallRepairCase({
380+
toolName: "grep",
381+
delta: String.raw` {“pattern”:“Use ”, “limit”: “bar” in prose”,“path”:“safe.txt”}`,
382+
});
383+
384+
expectAllToolCallArgs(result, {
385+
pattern: "Use ”, “limit”: “bar” in prose",
386+
path: "safe.txt",
387+
});
388+
expect(result.finalArgs).not.toHaveProperty("limit");
389+
});
390+
391+
it("uses the structured tool name over a mismatched smart-quote repair prefix", async () => {
392+
const result = await runToolCallRepairCase({
393+
toolName: "grep",
394+
preambleToolName: "read",
395+
delta: String.raw` {“pattern”:“Use ”, “limit”: “bar” in prose”,“path”:“safe.txt”}`,
396+
});
397+
398+
expectAllToolCallArgs(result, {
399+
pattern: "Use ”, “limit”: “bar” in prose",
400+
path: "safe.txt",
401+
});
402+
expect(result.finalArgs).not.toHaveProperty("limit");
403+
});
404+
405+
it("ignores inherited tool-name successor lookups while repairing smart-quoted args", async () => {
406+
const result = await runToolCallRepairCase({
407+
toolName: "constructor",
408+
delta: "{“length”:“x”,“foo”:1}",
409+
});
410+
411+
expectAllToolCallArgs(result, {});
412+
});
413+
330414
it("decodes JSON escapes inside smart-quoted string args", async () => {
331415
const result = await runToolCallRepairCase({
332416
delta: String.raw` {“path”:“safe.txt”,“content”:“line\nnext \"quoted\" path C:\\tmp mark \u2713 invalid \d”}`,

‎src/agents/embedded-agent-runner/run/attempt.tool-call-argument-repair.ts‎

Lines changed: 87 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,10 @@ const TOOLCALL_REPAIR_FREEFORM_SUCCESSOR_KEYS: Record<string, string> = {
7676
old_string: "new_string",
7777
oldText: "newText",
7878
};
79+
const TOOLCALL_REPAIR_TOOL_VALUE_SUCCESSOR_KEYS = new Map<
80+
string,
81+
ReadonlyMap<string, readonly string[]>
82+
>([["read", new Map([["path", ["offset", "limit"]]])]]);
7983
const TOOLCALL_REPAIR_JSON_STRING_ESCAPES: Record<string, string> = {
8084
'"': '"',
8185
"\\": "\\",
@@ -229,7 +233,41 @@ function readObjectMemberKeyAfterComma(raw: string, commaIndex: number): string
229233
return key.value;
230234
}
231235

232-
function shouldCloseSmartQuotedValueAt(raw: string, quoteIndex: number, valueKey: string): boolean {
236+
function normalizeToolCallRepairToolName(value: string): string | undefined {
237+
const trimmed = value.trim();
238+
if (!/^[a-z0-9_-]{1,128}$/i.test(trimmed)) {
239+
return undefined;
240+
}
241+
return trimmed.toLowerCase();
242+
}
243+
244+
function extractToolNameFromLeadingPrefix(prefix: string): string | undefined {
245+
const match = /(?:^|[.\s])(?:functions?|tools?)[._:/-]?([a-z0-9_-]+)/i.exec(prefix);
246+
return match?.[1] ? normalizeToolCallRepairToolName(match[1]) : undefined;
247+
}
248+
249+
function isToolSpecificValueSuccessor(params: {
250+
toolName?: string;
251+
valueKey: string;
252+
nextKey: string;
253+
}): boolean {
254+
const toolName = params.toolName;
255+
if (!toolName) {
256+
return false;
257+
}
258+
return (
259+
TOOLCALL_REPAIR_TOOL_VALUE_SUCCESSOR_KEYS.get(toolName)
260+
?.get(params.valueKey)
261+
?.includes(params.nextKey) ?? false
262+
);
263+
}
264+
265+
function shouldCloseSmartQuotedValueAt(
266+
raw: string,
267+
quoteIndex: number,
268+
valueKey: string,
269+
toolName?: string,
270+
): boolean {
233271
const nextIndex = skipWhitespace(raw, quoteIndex + 1);
234272
const nextChar = raw[nextIndex];
235273
if (nextIndex >= raw.length || nextChar === "}") {
@@ -244,7 +282,10 @@ function shouldCloseSmartQuotedValueAt(raw: string, quoteIndex: number, valueKey
244282
return false;
245283
}
246284
if (!TOOLCALL_REPAIR_FREEFORM_VALUE_KEYS.has(valueKey)) {
247-
return TOOLCALL_REPAIR_KNOWN_ARG_KEYS.has(nextKey);
285+
return (
286+
TOOLCALL_REPAIR_KNOWN_ARG_KEYS.has(nextKey) ||
287+
isToolSpecificValueSuccessor({ toolName, valueKey, nextKey })
288+
);
248289
}
249290
return TOOLCALL_REPAIR_FREEFORM_SUCCESSOR_KEYS[valueKey] === nextKey;
250291
}
@@ -261,11 +302,12 @@ function readSmartQuotedValue(
261302
raw: string,
262303
startIndex: number,
263304
key: string,
305+
toolName?: string,
264306
): ToolCallRepairJsonValue | undefined {
265307
let value = "";
266308
for (let i = startIndex + 1; i < raw.length; i += 1) {
267309
const char = raw[i];
268-
if (isToolCallRepairSmartQuote(char) && shouldCloseSmartQuotedValueAt(raw, i, key)) {
310+
if (isToolCallRepairSmartQuote(char) && shouldCloseSmartQuotedValueAt(raw, i, key, toolName)) {
269311
return { value: decodeSmartQuotedJsonStringEscapes(value), endIndex: i + 1 };
270312
}
271313
value += char;
@@ -366,13 +408,14 @@ function readObjectValue(
366408
raw: string,
367409
startIndex: number,
368410
key: string,
411+
toolName?: string,
369412
): ToolCallRepairJsonValue | undefined {
370413
const char = raw[startIndex];
371414
if (char === '"') {
372415
return readAsciiQuotedString(raw, startIndex);
373416
}
374417
if (isToolCallRepairSmartQuote(char)) {
375-
return readSmartQuotedValue(raw, startIndex, key);
418+
return readSmartQuotedValue(raw, startIndex, key, toolName);
376419
}
377420
if (key === "edits" && char === "[") {
378421
return readSmartQuotedEditArray(raw, startIndex);
@@ -383,6 +426,7 @@ function readObjectValue(
383426
function parseSmartQuotedToolCallObject(
384427
raw: string,
385428
startIndex: number,
429+
toolName?: string,
386430
): ToolCallRepairParsedObject | undefined {
387431
if (raw[startIndex] !== "{") {
388432
return undefined;
@@ -406,7 +450,7 @@ function parseSmartQuotedToolCallObject(
406450
return undefined;
407451
}
408452

409-
const value = readObjectValue(raw, skipWhitespace(raw, index + 1), key.value);
453+
const value = readObjectValue(raw, skipWhitespace(raw, index + 1), key.value, toolName);
410454
if (!value) {
411455
return undefined;
412456
}
@@ -460,7 +504,10 @@ function tryExtractUsableToolCallArgumentsFromJson(
460504
};
461505
}
462506

463-
function tryExtractSmartQuotedToolCallArguments(raw: string): ToolCallArgumentRepair | undefined {
507+
function tryExtractSmartQuotedToolCallArguments(
508+
raw: string,
509+
toolNameFromContext?: string,
510+
): ToolCallArgumentRepair | undefined {
464511
if (!/[\u201c\u201d\u201e\u201f]/.test(raw)) {
465512
return undefined;
466513
}
@@ -472,7 +519,11 @@ function tryExtractSmartQuotedToolCallArguments(raw: string): ToolCallArgumentRe
472519
if (!isAllowedToolCallRepairLeadingPrefix(leadingPrefix)) {
473520
return undefined;
474521
}
475-
const parsed = parseSmartQuotedToolCallObject(raw, startIndex);
522+
const parsed = parseSmartQuotedToolCallObject(
523+
raw,
524+
startIndex,
525+
toolNameFromContext ?? extractToolNameFromLeadingPrefix(leadingPrefix),
526+
);
476527
if (!parsed) {
477528
return undefined;
478529
}
@@ -491,7 +542,10 @@ function tryExtractSmartQuotedToolCallArguments(raw: string): ToolCallArgumentRe
491542
};
492543
}
493544

494-
function tryExtractUsableToolCallArguments(raw: string): ToolCallArgumentRepair | undefined {
545+
function tryExtractUsableToolCallArguments(
546+
raw: string,
547+
toolNameFromContext?: string,
548+
): ToolCallArgumentRepair | undefined {
495549
if (!raw.trim()) {
496550
return undefined;
497551
}
@@ -506,10 +560,30 @@ function tryExtractUsableToolCallArguments(raw: string): ToolCallArgumentRepair
506560
}
507561

508562
return (
509-
tryExtractUsableToolCallArgumentsFromJson(raw) ?? tryExtractSmartQuotedToolCallArguments(raw)
563+
tryExtractUsableToolCallArgumentsFromJson(raw) ??
564+
tryExtractSmartQuotedToolCallArguments(raw, toolNameFromContext)
510565
);
511566
}
512567

568+
function readToolCallNameInMessage(message: unknown, contentIndex: number): string | undefined {
569+
if (!message || typeof message !== "object") {
570+
return undefined;
571+
}
572+
const content = (message as { content?: unknown }).content;
573+
if (!Array.isArray(content)) {
574+
return undefined;
575+
}
576+
const block = content[contentIndex];
577+
if (!block || typeof block !== "object") {
578+
return undefined;
579+
}
580+
const typedBlock = block as { type?: unknown; name?: unknown };
581+
if (!isToolCallBlockType(typedBlock.type) || typeof typedBlock.name !== "string") {
582+
return undefined;
583+
}
584+
return normalizeToolCallRepairToolName(typedBlock.name);
585+
}
586+
513587
function repairToolCallArgumentsInMessage(
514588
message: unknown,
515589
contentIndex: number,
@@ -635,7 +709,10 @@ function wrapStreamRepairMalformedToolCallArguments(
635709
repairedArgsByIndex.has(event.contentIndex);
636710
if (shouldReevaluateRepair) {
637711
const hadRepairState = repairedArgsByIndex.has(event.contentIndex);
638-
const repair = tryExtractUsableToolCallArguments(nextPartialJson);
712+
const toolName =
713+
readToolCallNameInMessage(event.partial, event.contentIndex) ??
714+
readToolCallNameInMessage(event.message, event.contentIndex);
715+
const repair = tryExtractUsableToolCallArguments(nextPartialJson, toolName);
639716
if (repair) {
640717
if (
641718
!hadRepairState &&

0 commit comments

Comments
 (0)