fix(openai): bound streamed tool arguments (#117055)

packages/ai/src/providers/openai-completions-tool-calls.ts

@@ -1,9 +1,10 @@
 import { randomUUID } from "node:crypto";
 import { isRecord } from "@openclaw/normalization-core/record-coerce";
 import type { ChatCompletionChunk } from "openai/resources/chat/completions.js";
+import { measureUtf8AppendBytes } from "../transports/openai-transport-shared.js";
 
 type ChatCompletionToolCallDelta = ChatCompletionChunk.Choice.Delta.ToolCall;
-const MAX_BUFFERED_LEGACY_TOOL_CALL_ARGUMENT_BYTES = 256_000;
+const MAX_BUFFERED_TOOL_CALL_ARGUMENT_BYTES = 256_000;
 const MAX_BUFFERED_LEGACY_FOLLOWING_DELTA_BYTES = 256_000;
 const MAX_BUFFERED_LEGACY_FOLLOWING_DELTAS = 1_024;
 
@@ -26,6 +27,8 @@ export function createOpenAICompletionsToolCallDeltaNormalizer(): (
   let pendingLegacyArgumentBytes = 0;
   let pendingFollowingDeltaBytes = 0;
   let pendingLegacyToolCall: ChatCompletionToolCallDelta | undefined;
+  type ModernArgumentState = { bytes: number; pendingHighSurrogate: boolean };
+  const modernArguments = new Map<number | string, ModernArgumentState>();
   const pendingFollowingDeltas: ChatCompletionChunk.Choice.Delta[] = [];
 
   const takePendingFollowingDeltas = (): NormalizedOpenAICompletionsDelta[] => {
@@ -79,6 +82,29 @@ export function createOpenAICompletionsToolCallDeltaNormalizer(): (
   return (delta, finishReason) => {
     const ordinaryDelta = withoutToolCalls(delta);
     if (delta.tool_calls && delta.tool_calls.length > 0) {
+      for (const toolCall of delta.tool_calls) {
+        const index = typeof toolCall.index === "number" ? toolCall.index : undefined;
+        // Both consumers resolve index first, then id; bind every supplied alias
+        // to one byte state so compatible id-only continuations cannot reset the cap.
+        const state = (index === undefined ? undefined : modernArguments.get(index)) ??
+          (toolCall.id ? modernArguments.get(toolCall.id) : undefined) ?? {
+            bytes: 0,
+            pendingHighSurrogate: false,
+          };
+        if (index !== undefined) {
+          modernArguments.set(index, state);
+        }
+        if (toolCall.id) {
+          modernArguments.set(toolCall.id, state);
+        }
+        const argumentDelta = toolCall.function?.arguments ?? "";
+        const append = measureUtf8AppendBytes(state.pendingHighSurrogate, argumentDelta);
+        state.bytes += append.bytes;
+        if (state.bytes > MAX_BUFFERED_TOOL_CALL_ARGUMENT_BYTES) {
+          throw new Error("Exceeded tool-call argument buffer limit");
+        }
+        state.pendingHighSurrogate = append.endsWithHighSurrogate;
+      }
       const precedingDeltas = takePendingFollowingDeltas();
       sawModernToolCall = true;
       pendingLegacyArgumentBytes = 0;
@@ -113,10 +139,7 @@ export function createOpenAICompletionsToolCallDeltaNormalizer(): (
       const nextArgumentBytes =
         Buffer.byteLength(functionCall.arguments ?? "", "utf8") +
         Buffer.byteLength(nextFunctionName ?? "", "utf8");
-      if (
-        pendingLegacyArgumentBytes + nextArgumentBytes >
-        MAX_BUFFERED_LEGACY_TOOL_CALL_ARGUMENT_BYTES
-      ) {
+      if (pendingLegacyArgumentBytes + nextArgumentBytes > MAX_BUFFERED_TOOL_CALL_ARGUMENT_BYTES) {
         throw new Error("Exceeded tool-call argument buffer limit");
       }
       pendingLegacyArgumentBytes += nextArgumentBytes;

packages/ai/src/providers/openai-completions.legacy-function-call.test.ts

@@ -32,6 +32,7 @@ const context = {
     },
   ],
 } satisfies Context;
+const TOOL_ARGUMENT_BYTE_LIMIT = 256_000;
 
 function chunk(
   delta: ChatCompletionChunk.Choice.Delta,
@@ -69,6 +70,46 @@ function toolCallDelta({
   };
 }
 
+function idOnlyToolCallDelta(params: {
+  id: string;
+  name?: string;
+  arguments: string;
+}): ChatCompletionChunk.Choice.Delta.ToolCall {
+  return {
+    id: params.id,
+    type: "function",
+    function: {
+      ...(params.name !== undefined ? { name: params.name } : {}),
+      arguments: params.arguments,
+    },
+  } as ChatCompletionChunk.Choice.Delta.ToolCall;
+}
+
+function argumentsWithByteLength(bytes: number, fill = "a"): string {
+  const prefix = '{"query":"';
+  const suffix = '"}';
+  const availableBytes = bytes - Buffer.byteLength(prefix + suffix, "utf8");
+  const characterBytes = Buffer.byteLength(fill, "utf8");
+  const value =
+    prefix +
+    fill.repeat(Math.floor(availableBytes / characterBytes)) +
+    "a".repeat(availableBytes % characterBytes) +
+    suffix;
+  expect(Buffer.byteLength(value, "utf8")).toBe(bytes);
+  return value;
+}
+
+function splitSurrogateArguments(bytes: number): [string, string] {
+  const prefix = '{"query":"';
+  const suffix = '"}';
+  const emoji = "😀";
+  const padding = bytes - Buffer.byteLength(prefix + suffix + emoji, "utf8");
+  const value = `${prefix}${"a".repeat(padding)}${emoji}${suffix}`;
+  expect(Buffer.byteLength(value, "utf8")).toBe(bytes);
+  const surrogateBoundary = value.indexOf(emoji) + 1;
+  return [value.slice(0, surrogateBoundary), value.slice(surrogateBoundary)];
+}
+
 const modernCallChunk = (
   rawArguments: string,
   { id = "call_modern", index = 0, name = "lookup" } = {},
@@ -476,6 +517,117 @@ describe.each([
     expect(eventTypes).not.toContain("toolcall_start");
   });
 
+  it("bounds an indexed modern call through an id-only continuation", async () => {
+    const value = argumentsWithByteLength(TOOL_ARGUMENT_BYTE_LIMIT + 1);
+    const { eventTypes, result } = await collectFixture([
+      chunk({
+        tool_calls: [
+          toolCallDelta({
+            index: 0,
+            id: "call_alias",
+            name: "lookup",
+            arguments: value.slice(0, 128_000),
+          }),
+        ],
+      }),
+      chunk({
+        tool_calls: [idOnlyToolCallDelta({ id: "call_alias", arguments: value.slice(128_000) })],
+      }),
+      chunk({}, "tool_calls"),
+    ]);
+
+    expect(result.stopReason).toBe("error");
+    expect(result.errorMessage).toContain("Exceeded tool-call argument buffer limit");
+    expect(result.content.filter((block) => block.type === "toolCall")).toHaveLength(0);
+    expect(eventTypes).not.toContain("toolcall_end");
+  });
+
+  it("keeps independent id-only modern calls independently bounded", async () => {
+    const { result } = await collectFixture([
+      chunk({
+        tool_calls: [
+          idOnlyToolCallDelta({
+            id: "call_id_a",
+            name: "lookup",
+            arguments: argumentsWithByteLength(200_000),
+          }),
+          idOnlyToolCallDelta({
+            id: "call_id_b",
+            name: "lookup",
+            arguments: argumentsWithByteLength(200_000),
+          }),
+        ],
+      }),
+      chunk({}, "tool_calls"),
+    ]);
+
+    expect(result.stopReason).toBe("toolUse");
+    expect(
+      result.content
+        .filter((block) => block.type === "toolCall")
+        .map((block) => [block.id, Buffer.byteLength(JSON.stringify(block.arguments), "utf8")]),
+    ).toEqual([
+      ["call_id_a", 200_000],
+      ["call_id_b", 200_000],
+    ]);
+  });
+
+  it.each([
+    {
+      name: "accepts an exact-limit surrogate pair split between deltas",
+      bytes: TOOL_ARGUMENT_BYTE_LIMIT,
+      stopReason: "toolUse",
+    },
+    {
+      name: "rejects an oversized surrogate pair split between deltas",
+      bytes: TOOL_ARGUMENT_BYTE_LIMIT + 1,
+      stopReason: "error",
+    },
+  ] as const)("$name", async ({ bytes, stopReason }) => {
+    const [first, second] = splitSurrogateArguments(bytes);
+    const { eventTypes, result } = await collectFixture([
+      chunk({
+        tool_calls: [
+          toolCallDelta({
+            index: 0,
+            id: "call_surrogate",
+            name: "lookup",
+            arguments: first,
+          }),
+        ],
+      }),
+      chunk({ tool_calls: [toolCallDelta({ index: 0, arguments: second })] }),
+      chunk({}, "tool_calls"),
+    ]);
+
+    expect(result.stopReason).toBe(stopReason);
+    if (stopReason === "error") {
+      expect(result.errorMessage).toContain("Exceeded tool-call argument buffer limit");
+      expect(result.content.filter((block) => block.type === "toolCall")).toHaveLength(0);
+      expect(eventTypes).not.toContain("toolcall_end");
+      return;
+    }
+    expect(result.content[0]).toMatchObject({ id: "call_surrogate", type: "toolCall" });
+    expect(
+      result.content[0]?.type === "toolCall"
+        ? Buffer.byteLength(JSON.stringify(result.content[0].arguments), "utf8")
+        : undefined,
+    ).toBe(bytes);
+  });
+
+  it("measures multibyte modern arguments by UTF-8 bytes", async () => {
+    const { eventTypes, result } = await collectFixture(
+      confirmedModernCallChunks(argumentsWithByteLength(TOOL_ARGUMENT_BYTE_LIMIT + 1, "é"), {
+        id: "call_multibyte",
+      }),
+    );
+
+    expect(result.stopReason).toBe("error");
+    expect(result.errorMessage).toContain("Exceeded tool-call argument buffer limit");
+    expect(result.content.filter((block) => block.type === "toolCall")).toHaveLength(0);
+    expect(eventTypes).not.toContain("toolcall_end");
+  });
+
   it("coalesces tiny provisional legacy fragments into one bounded executable delta", async () => {
     const query = "x".repeat(1_100);
     const { argumentDeltas, result } = await collectFixture([

packages/ai/src/transports/openai-completions-transport.ts

@@ -76,6 +76,7 @@ import {
   createOpenAIResponseHook,
   isOpenAICompletionsThinkingEnabled,
   log,
+  measureUtf8AppendBytes,
   parseOpenAICompletionsUsage,
   readOpenAICompletionsContentDeltas,
   resolvePromptCacheKey,
@@ -406,7 +407,6 @@ async function processOpenAICompletionsStream(
   },
 ) {
   const MAX_POST_TOOL_CALL_BUFFER_BYTES = 256_000;
-  const MAX_TOOL_CALL_ARGUMENT_BUFFER_BYTES = 256_000;
   const emitReasoning = options?.emitReasoning ?? true;
   const compat = getCompat(model as OpenAIModeModel);
   const deepSeekTextFilter = shouldFilterDeepSeekDsmlText(compat)
@@ -435,7 +435,6 @@ async function processOpenAICompletionsStream(
   const toolCallBlocksByIndex = new Map<number, ToolCallBlock>();
   const toolCallBlocksById = new Map<string, ToolCallBlock>();
   const provisionalCommentaryTags: PendingCommentaryTags = new Map();
-  const toolCallBlockBytes = new WeakMap<ToolCallBlock, number>();
   const toolCallBlockIndices = new WeakMap<ToolCallBlock, number>();
   const normalizeToolCallDeltas = createOpenAICompletionsToolCallDeltaNormalizer();
   let sawStopFinishReason = false;
@@ -818,12 +817,6 @@ async function processOpenAICompletionsStream(
             block.thoughtSignature = deltaSig;
           }
           if (toolCall.function?.arguments) {
-            const nextArgumentBytes = measureUtf8Bytes(toolCall.function.arguments);
-            const currentBlockArgBytes = toolCallBlockBytes.get(block) ?? 0;
-            if (currentBlockArgBytes + nextArgumentBytes > MAX_TOOL_CALL_ARGUMENT_BUFFER_BYTES) {
-              throw new Error("Exceeded tool-call argument buffer limit");
-            }
-            toolCallBlockBytes.set(block, currentBlockArgBytes + nextArgumentBytes);
             block.partialArgs += toolCall.function.arguments;
             block.arguments = parseStreamingJson(block.partialArgs);
             pushStreamEvent({
@@ -910,7 +903,7 @@ const DEEPSEEK_DSML_RECOVERY_MAX_BOUNDARY_LEN = Math.max(
   ...DEEPSEEK_DSML_INVOKE_CLOSE_TOKENS.map((token) => token.length),
 );
 
-// Match MAX_TOOL_CALL_ARGUMENT_BUFFER_BYTES / MAX_POST_TOOL_CALL_BUFFER_BYTES.
+// Match the shared Chat tool-argument and post-tool-call buffer limits.
 const MAX_DSML_RECOVERY_BUFFER_BYTES = 256_000;
 const DEEPSEEK_DSML_SCAN_BATCH_CHARS = 64 * 1_024;
 
@@ -1026,7 +1019,7 @@ function createDeepSeekDsmlToolCallRecoverer() {
 
   return {
     push(chunk: string) {
-      const append = utf8ByteLengthForAppend(bufferEndsWithHighSurrogate, chunk);
+      const append = measureUtf8AppendBytes(bufferEndsWithHighSurrogate, chunk);
       bufferBytes += append.bytes;
       bufferEndsWithHighSurrogate = append.endsWithHighSurrogate;
       buffer += chunk;
@@ -1242,23 +1235,6 @@ function scanDeepSeekDsmlToolBlock(
   return { kind: "incomplete" };
 }
 
-function utf8ByteLengthForAppend(bufferEndsWithHighSurrogate: boolean, chunk: string) {
-  let bytes = Buffer.byteLength(chunk, "utf8");
-  if (!chunk) {
-    return { bytes, endsWithHighSurrogate: bufferEndsWithHighSurrogate };
-  }
-  const nextCodeUnit = chunk.charCodeAt(0);
-  if (bufferEndsWithHighSurrogate && nextCodeUnit >= 0xdc00 && nextCodeUnit <= 0xdfff) {
-    // Each isolated surrogate counts as three UTF-8 bytes; the joined scalar is four.
-    bytes -= 2;
-  }
-  const finalCodeUnit = chunk.charCodeAt(chunk.length - 1);
-  return {
-    bytes,
-    endsWithHighSurrogate: finalCodeUnit >= 0xd800 && finalCodeUnit <= 0xdbff,
-  };
-}
-
 function longestDeepSeekDsmlToolOpenPrefixSuffixLength(text: string) {
   const maxLength = Math.min(text.length, DEEPSEEK_DSML_TOOL_MAX_OPEN_TOKEN_LEN - 1);
   for (let length = maxLength; length > 0; length -= 1) {

packages/ai/src/transports/openai-transport-shared.ts

@@ -123,6 +123,24 @@ export function throwIfModelStreamAborted(signal?: AbortSignal): void {
   }
 }
 
+/** Measure one UTF-8 append without double-counting a surrogate pair split across chunks. */
+export function measureUtf8AppendBytes(bufferEndsWithHighSurrogate: boolean, chunk: string) {
+  let bytes = Buffer.byteLength(chunk, "utf8");
+  if (!chunk) {
+    return { bytes, endsWithHighSurrogate: bufferEndsWithHighSurrogate };
+  }
+  const nextCodeUnit = chunk.charCodeAt(0);
+  if (bufferEndsWithHighSurrogate && nextCodeUnit >= 0xdc00 && nextCodeUnit <= 0xdfff) {
+    // Each isolated surrogate counts as three UTF-8 bytes; the joined scalar is four.
+    bytes -= 2;
+  }
+  const finalCodeUnit = chunk.charCodeAt(chunk.length - 1);
+  return {
+    bytes,
+    endsWithHighSurrogate: finalCodeUnit >= 0xd800 && finalCodeUnit <= 0xdbff,
+  };
+}
+
 export function createModelStreamCooperativeScheduler(
   signal?: AbortSignal,
 ): ModelStreamCooperativeScheduler {