fix(ai): recover incomplete streamed tool arguments (#44875)

packages/ai/src/protocols/utils/partial-json.ts

@@ -42,7 +42,61 @@ export function parseJSON(jsonString: string, allowPartial = Allow.ALL): unknown
   try {
     return decodeJson(input)
   } catch {}
-  return _parseJSON(input, allowPartial)
+
+  const repaired = repairJSON(input)
+  if (repaired !== input) {
+    try {
+      return decodeJson(repaired)
+    } catch {}
+  }
+
+  try {
+    return _parseJSON(input, allowPartial)
+  } catch (error) {
+    if (repaired !== input) return _parseJSON(repaired, allowPartial)
+    throw error
+  }
+}
+
+const repairJSON = (input: string) => {
+  let repaired = ""
+  let quoted = false
+
+  for (let index = 0; index < input.length; index++) {
+    const character = input[index]
+    if (!quoted) {
+      repaired += character
+      if (character === '"') quoted = true
+      continue
+    }
+
+    if (character === '"') {
+      repaired += character
+      quoted = false
+      continue
+    }
+
+    if (character === "\\") {
+      const next = input[index + 1]
+      if (next === "u" && /^[0-9a-fA-F]{4}$/.test(input.slice(index + 2, index + 6))) {
+        repaired += input.slice(index, index + 6)
+        index += 5
+        continue
+      }
+      if (next !== undefined && '"\\/bfnrtu'.includes(next)) {
+        repaired += `\\${next}`
+        index++
+        continue
+      }
+      repaired += "\\\\"
+      continue
+    }
+
+    const code = character.charCodeAt(0)
+    repaired += code <= 0x1f ? `\\u${code.toString(16).padStart(4, "0")}` : character
+  }
+
+  return repaired
 }
 
 const _parseJSON = (jsonString: string, allow: number) => {
@@ -148,7 +202,12 @@ const _parseJSON = (jsonString: string, allow: number) => {
         skipBlank()
         index++
         try {
-          object[key] = parseAny()
+          Object.defineProperty(object, key, {
+            value: parseAny(),
+            enumerable: true,
+            configurable: true,
+            writable: true,
+          })
         } catch (error) {
           if (Allow.OBJ & allow) return object
           throw error

packages/ai/src/protocols/utils/tool-stream.ts

@@ -1,5 +1,5 @@
 import { Effect, Option } from "effect"
-import { AIError, LLMEvent, type ProviderMetadata, type ToolCall, type ToolInputError } from "../../schema/index.js"
+import { AIError, LLMEvent, type ProviderMetadata, type ToolCall } from "../../schema/index.js"
 import { eventError, parseToolInput, type ToolAccumulator } from "../shared.js"
 import { parse } from "./partial-json.js"
 
@@ -59,46 +59,44 @@ const inputStart = (tool: PendingTool) =>
     providerMetadata: tool.providerMetadata,
   })
 
-const inputDelta = (tool: PendingTool, text: string) => {
-  const input = parsePartialInput(tool.input)
-  return LLMEvent.toolInputDelta({
+const inputDelta = (tool: PendingTool, text: string) =>
+  LLMEvent.toolInputDelta({
     id: tool.id,
     name: tool.name,
     text,
-    ...(Option.isSome(input) ? { input: input.value } : {}),
+    input: Option.getOrElse(parsePartialInput(tool.input), () => ({})),
   })
-}
 
 const toolCall = (route: string, tool: PendingTool, inputOverride?: string) => {
   const raw = inputOverride ?? tool.input
   return parseToolInput(route, tool.name, raw).pipe(
-    Effect.map((input): ToolCall | ToolInputError =>
-      LLMEvent.toolCall({
-        id: tool.id,
-        name: tool.name,
-        input,
-        providerExecuted: tool.providerExecuted ? true : undefined,
-        providerMetadata: tool.providerMetadata,
-      }),
-    ),
     Effect.catch((error) =>
       tool.providerExecuted
         ? Effect.fail(error)
         : Effect.succeed(
-            LLMEvent.toolInputError({
-              id: tool.id,
-              name: tool.name,
-              raw,
-            }),
+            Option.getOrElse(
+              Option.map(parsePartialInput(raw), (input) => input ?? {}),
+              () => ({}),
+            ),
           ),
     ),
+    Effect.map(
+      (input): ToolCall =>
+        LLMEvent.toolCall({
+          id: tool.id,
+          name: tool.name,
+          input,
+          providerExecuted: tool.providerExecuted ? true : undefined,
+          providerMetadata: tool.providerMetadata,
+        }),
+    ),
   )
 }
 
-const finishEvents = (tool: PendingTool, event: ToolCall | ToolInputError): ReadonlyArray<LLMEvent> =>
-  event.type === "tool-input-error"
-    ? [event]
-    : [LLMEvent.toolInputEnd({ id: tool.id, name: tool.name, providerMetadata: tool.providerMetadata }), event]
+const finishEvents = (tool: PendingTool, event: ToolCall): ReadonlyArray<LLMEvent> => [
+  LLMEvent.toolInputEnd({ id: tool.id, name: tool.name, providerMetadata: tool.providerMetadata }),
+  event,
+]
 
 /** Store the updated tool and produce the optional public delta event. */
 const appendTool = <K extends StreamKey>(
@@ -181,7 +179,7 @@ export const appendExisting = <K extends StreamKey>(
 
 /**
  * Finalize one pending tool call: parse the accumulated raw JSON, remove it
- * from state, and return either a call or a non-executable local input error.
+ * from state, and recover incomplete local arguments when needed.
  * Missing keys are a no-op because some providers emit stop events for
  * non-tool content blocks.
  */

packages/ai/test/partial-json.test.ts

@@ -18,6 +18,21 @@ describe("partial JSON", () => {
     expect(() => parse('"hello', ~Allow.STR)).toThrow(PartialJSON)
   })
 
+  test("repairs invalid escapes and raw control characters", () => {
+    expect(parse('{"path":"A\\H","text":"first\tsecond"}')).toEqual({
+      path: "A\\H",
+      text: "first\tsecond",
+    })
+  })
+
+  test("preserves prototype keys in partial objects", () => {
+    const object = parse('{"__proto__":{"safe":true}') as Record<string, unknown>
+
+    expect(Object.hasOwn(object, "__proto__")).toBe(true)
+    expect(Object.getPrototypeOf(object)).toBe(Object.prototype)
+    expect(object.__proto__).toEqual({ safe: true })
+  })
+
   test("controls partial collection values independently", () => {
     expect(parse('["', Allow.ARR)).toEqual([])
     expect(parse('["', Allow.ARR | Allow.STR)).toEqual([""])

packages/ai/test/provider/bedrock-converse.test.ts

@@ -491,7 +491,7 @@ describe("Bedrock Converse route", () => {
     }),
   )
 
-  it.effect("emits malformed tool input as an unexecuted tool error", () =>
+  it.effect("recovers incomplete tool input at finalization", () =>
     Effect.gen(function* () {
       const body = eventStreamBody(
         ["messageStart", { role: "assistant" }],
@@ -508,10 +508,10 @@ describe("Bedrock Converse route", () => {
       )
       const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
 
-      expect(response.events.find((event) => event.type === "tool-input-error")).toMatchObject({
+      expect(response.events.find((event) => event.type === "tool-call")).toMatchObject({
         id: "tool_1",
         name: "lookup",
-        raw: '{"query":"partial',
+        input: { query: "partial" },
       })
       expect(response.finishReason).toEqual({ normalized: "tool-calls", raw: "end_turn" })
     }),

packages/ai/test/provider/openai-responses.test.ts

@@ -3069,7 +3069,7 @@ describe("OpenAI Responses route", () => {
     }),
   )
 
-  it.effect("emits malformed final function arguments as an unexecuted tool error", () =>
+  it.effect("recovers authoritative incomplete final function arguments", () =>
     Effect.gen(function* () {
       const body = sseEvents(
         {
@@ -3095,18 +3095,17 @@ describe("OpenAI Responses route", () => {
         }),
       ).pipe(Effect.provide(fixedResponse(body)))
 
-      expect(response.events.find(LLMEvent.is.toolInputError)).toEqual({
-        type: "tool-input-error",
+      expect(response.events.find(LLMEvent.is.toolCall)).toMatchObject({
         id: "call_1",
         name: "lookup",
-        raw: '{"query":"partial',
+        input: { query: "partial" },
       })
       expect(response.finishReason.normalized).toBe("tool-calls")
-      expect(response.events.some(LLMEvent.is.toolCall)).toBeFalse()
+      expect(response.events.some(LLMEvent.is.toolInputError)).toBeFalse()
     }),
   )
 
-  it.effect("settles malformed function arguments when output_item.added is absent", () =>
+  it.effect("recovers incomplete function arguments when output_item.added is absent", () =>
     Effect.gen(function* () {
       const body = sseEvents(
         {
@@ -3123,10 +3122,10 @@ describe("OpenAI Responses route", () => {
       )
       const response = yield* LLMClient.generate(request).pipe(Effect.provide(fixedResponse(body)))
 
-      expect(response.events.find(LLMEvent.is.toolInputError)).toMatchObject({
+      expect(response.events.find(LLMEvent.is.toolCall)).toMatchObject({
         id: "call_1",
         name: "lookup",
-        raw: '{"query":"partial',
+        input: { query: "partial" },
       })
       expect(response.finishReason.normalized).toBe("tool-calls")
     }),

packages/ai/test/tool-stream.test.ts

@@ -59,7 +59,7 @@ describe("ToolStream", () => {
     }),
   )
 
-  it.effect("omits partial input when the accumulated value cannot be parsed", () =>
+  it.effect("defaults partial input to an empty object when the accumulated value cannot be parsed", () =>
     Effect.gen(function* () {
       const result = ToolStream.appendOrStart(
         ADAPTER,
@@ -72,7 +72,7 @@ describe("ToolStream", () => {
 
       expect(result.events).toEqual([
         { type: "tool-input-start", id: "call_1", name: "lookup" },
-        { type: "tool-input-delta", id: "call_1", name: "lookup", text: "x" },
+        { type: "tool-input-delta", id: "call_1", name: "lookup", text: "x", input: {} },
       ])
     }),
   )
@@ -132,7 +132,7 @@ describe("ToolStream", () => {
     }),
   )
 
-  it.effect("finalizes malformed local input as a non-executable tool error", () =>
+  it.effect("finalizes incomplete local input using the partial JSON parser", () =>
     Effect.gen(function* () {
       const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
         id: "call_1",
@@ -144,18 +144,46 @@ describe("ToolStream", () => {
       expect(finished).toEqual({
         tools: {},
         events: [
-          {
-            type: "tool-input-error",
-            id: "call_1",
-            name: "lookup",
-            raw: '{"query":"partial',
-          },
+          { type: "tool-input-end", id: "call_1", name: "lookup" },
+          { type: "tool-call", id: "call_1", name: "lookup", input: { query: "partial" } },
         ],
       })
     }),
   )
 
-  it.effect("preserves valid siblings when one parallel input is malformed", () =>
+  it.effect("repairs malformed string escapes in final local input", () =>
+    Effect.gen(function* () {
+      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
+        id: "call_1",
+        name: "lookup",
+        input: '{"path":"A\\H","text":"first\tsecond"}',
+      })
+      const finished = yield* ToolStream.finish(ADAPTER, tools, "item_1")
+
+      expect(finished.events).toEqual([
+        { type: "tool-input-end", id: "call_1", name: "lookup" },
+        { type: "tool-call", id: "call_1", name: "lookup", input: { path: "A\\H", text: "first\tsecond" } },
+      ])
+    }),
+  )
+
+  it.effect("defaults unrecoverable local input to an empty object", () =>
+    Effect.gen(function* () {
+      const tools = ToolStream.start(ToolStream.empty<string>(), "item_1", {
+        id: "call_1",
+        name: "lookup",
+        input: "invalid",
+      })
+      const finished = yield* ToolStream.finish(ADAPTER, tools, "item_1")
+
+      expect(finished.events).toEqual([
+        { type: "tool-input-end", id: "call_1", name: "lookup" },
+        { type: "tool-call", id: "call_1", name: "lookup", input: {} },
+      ])
+    }),
+  )
+
+  it.effect("recovers incomplete input alongside valid parallel tool calls", () =>
     Effect.gen(function* () {
       const valid = ToolStream.start(ToolStream.empty<number>(), 0, {
         id: "call_valid",
@@ -174,12 +202,8 @@ describe("ToolStream", () => {
         events: [
           { type: "tool-input-end", id: "call_valid", name: "lookup" },
           { type: "tool-call", id: "call_valid", name: "lookup", input: { query: "weather" } },
-          {
-            type: "tool-input-error",
-            id: "call_invalid",
-            name: "lookup",
-            raw: '{"query":"partial',
-          },
+          { type: "tool-input-end", id: "call_invalid", name: "lookup" },
+          { type: "tool-call", id: "call_invalid", name: "lookup", input: { query: "partial" } },
         ],
       })
     }),