Skip to content

fix: bound streamed tool argument memory #847

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Ingwannu wants to merge 1 commit into
base: dev
Choose a base branch
from
Draft

fix: bound streamed tool argument memory #847

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 50 additions & 3 deletions src/adapters/openai-chat.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,8 @@ import { contentPartsToText } from "./image";
import { neutralizeIdentity } from "./identity";
import { buildNonOpenAIToolCatalogNudgeForTools, shouldInjectNonOpenAIToolCatalogNudge } from "./tool-catalog-nudge";
import { openRouterProviderPayload, resolveOpenRouterRouting } from "../providers/openrouter-routing";
import { MAX_SSE_RECORD_BYTES } from "../lib/sse-decoder";
import { appendToolArguments, utf8ByteLength } from "../lib/tool-argument-bounds";

// Providers may opt into stripping one trailing "[...]" group from the wire model id.
// Z.AI needs this because its OpenAI path rejects glm-5.2[1m] with 400 code 1211;
Expand Down Expand Up @@ -693,8 +695,9 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
// chunks, whole-chunk calls) cannot be represented live without overlapping sequences.
// Keyed by `index` (OpenAI wire standard), falling back to `id`, falling back to the
// last-seen call for providers that omit both on continuation chunks.
interface PendingToolCall { key: string; id: string; name: string; args: string }
interface PendingToolCall { key: string; id: string; name: string; args: string; argsBytes: number }
const pendingToolCalls: PendingToolCall[] = [];
let pendingToolArgumentBytes = 0;
let toolCallSeq = 0;
const flushToolCalls = function* (): Generator<AdapterEvent> {
// Do not treat flushed tool calls as user-facing output for the finish-less EOF
Expand All @@ -706,6 +709,7 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
yield { type: "tool_call_end" };
}
pendingToolCalls.length = 0;
pendingToolArgumentBytes = 0;
};
let pendingUsage: OcxUsage | undefined;
// Track terminal signals so a socket EOF without any terminator can fail closed instead of
Expand All @@ -722,6 +726,17 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
// Yields adapter events and returns "terminate" for a terminal frame ([DONE] / error) that
// must end the stream, or "continue" otherwise. Mutates the closure's terminal-signal state.
const handleDataLine = function* (line: string): Generator<AdapterEvent, "continue" | "terminate"> {
if (utf8ByteLength(line) > MAX_SSE_RECORD_BYTES) {
pendingToolCalls.length = 0;
pendingToolArgumentBytes = 0;
yield {
type: "error",
message: "upstream SSE record exceeded the 4 MiB limit",
status: 502,
errorType: "upstream_error",
};
return "terminate";
}
if (!line.startsWith("data: ")) return "continue";
const payload = line.slice(6).trim();
if (payload === "[DONE]") {
Expand Down Expand Up @@ -803,12 +818,33 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
// splitting into two calls that share one call_id downstream.
if (!call && tc.id) call = pendingToolCalls.find(c => c.id === tc.id);
if (!call) {
call = { key: key ?? `seq:${pendingToolCalls.length}`, id: "", name: "", args: "" };
call = { key: key ?? `seq:${pendingToolCalls.length}`, id: "", name: "", args: "", argsBytes: 0 };
pendingToolCalls.push(call);
}
if (tc.id && !call.id) call.id = tc.id;
if (tc.function?.name && !call.name) call.name = tc.function.name;
if (tc.function?.arguments) call.args += tc.function.arguments;
if (tc.function?.arguments) {
const appended = appendToolArguments(
call.args,
call.argsBytes,
pendingToolArgumentBytes,
tc.function.arguments,
);
if (!appended) {
pendingToolCalls.length = 0;
pendingToolArgumentBytes = 0;
yield {
type: "error",
message: "upstream tool call arguments exceeded the configured byte limit",
status: 502,
errorType: "upstream_error",
};
return "terminate";
}
call.args = appended.value;
call.argsBytes = appended.valueBytes;
pendingToolArgumentBytes = appended.retainedTurnBytes;
}
}
}
}
Expand All @@ -833,6 +869,17 @@ export function createOpenAIChatAdapter(provider: OcxProviderConfig): ProviderAd
for (const line of lines) {
if ((yield* handleDataLine(line)) === "terminate") return;
}
if (utf8ByteLength(buffer) > MAX_SSE_RECORD_BYTES) {
pendingToolCalls.length = 0;
pendingToolArgumentBytes = 0;
yield {
type: "error",
message: "upstream SSE record exceeded the 4 MiB limit",
status: 502,
errorType: "upstream_error",
};
return;
}
}

// Some providers send the terminal `data:` frame (carrying the final delta, finish_reason,
Expand Down
70 changes: 66 additions & 4 deletions src/bridge.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import { encodeCompactionSummary } from "./responses/compaction";
import { encodeReasoningEnvelope, type ReasoningEnvelope } from "./responses/reasoning-envelope";
import { resolveStallTimeoutSec } from "./stall-timeout";
import { usageDisplayTotalTokens } from "./usage/totals";
import { appendToolArguments } from "./lib/tool-argument-bounds";

function uuid(): string {
return crypto.randomUUID().replace(/-/g, "");
Expand Down Expand Up @@ -314,7 +315,8 @@ export function bridgeToResponsesSSE(
// Full assistant text of a compaction turn (across message boundaries) — becomes the
// synthetic compaction item's payload on done.
let compactionText = "";
let currentToolCall: { itemId: string; outputIndex: number; callId: string; name: string; args: string; namespace?: string; freeform?: boolean; toolSearch?: boolean; inputEmitted?: string } | null = null;
let retainedToolArgumentBytes = 0;
let currentToolCall: { itemId: string; outputIndex: number; callId: string; name: string; args: string; argsBytes: number; namespace?: string; freeform?: boolean; toolSearch?: boolean; inputEmitted?: string } | null = null;
// Open native web-search cell (between begin and end). Holds the output index allocated on
// begin so the matching done reuses it; closed as `failed` if the stream terminates early.
let currentWebSearch: { itemId: string; eventId: string; outputIndex: number } | null = null;
Expand Down Expand Up @@ -676,12 +678,42 @@ export function bridgeToResponsesSSE(
? { type: "custom_tool_call", id: itemId, call_id: event.id, name: realName, input: "", status: "in_progress" }
: { type: "function_call", id: itemId, call_id: event.id, name: realName, arguments: "", status: "in_progress", ...(ns ? { namespace: ns } : {}) };
emit("response.output_item.added", { output_index: outputIndex, item });
currentToolCall = { itemId, outputIndex, callId: event.id, name: realName, args: "", namespace: ns, freeform, toolSearch };
currentToolCall = { itemId, outputIndex, callId: event.id, name: realName, args: "", argsBytes: 0, namespace: ns, freeform, toolSearch };
break;
}
case "tool_call_delta": {
if (currentToolCall) {
currentToolCall.args += event.arguments;
const appended = appendToolArguments(
currentToolCall.args,
currentToolCall.argsBytes,
retainedToolArgumentBytes,
event.arguments,
);
if (!appended) {
retainedToolArgumentBytes = Math.max(0, retainedToolArgumentBytes - currentToolCall.argsBytes);
currentToolCall.args = "";
currentToolCall.argsBytes = 0;
currentToolCall.inputEmitted = undefined;
failCurrentToolCall();
const failure = responseError(
502,
"upstream_error",
"upstream tool call arguments exceeded the configured byte limit",
);
emit("response.failed", {
response: {
...responseSnapshot("failed", finishedItems),
error: failure,
last_error: failure,
},
});
reportTerminal("failed");
terminalEvent = true;
break;
}
currentToolCall.args = appended.value;
currentToolCall.argsBytes = appended.valueBytes;
retainedToolArgumentBytes = appended.retainedTurnBytes;
if (!currentToolCall.freeform && !currentToolCall.toolSearch) {
emit("response.function_call_arguments.delta", {
item_id: currentToolCall.itemId, output_index: currentToolCall.outputIndex,
Expand Down Expand Up @@ -1041,6 +1073,9 @@ export function buildResponseJSON(
let currentToolCallId = "";
let currentToolCallName = "";
let currentToolCallArgs = "";
let currentToolCallArgsBytes = 0;
let retainedToolArgumentBytes = 0;
let toolArgumentOverflow = false;
// Web-search citations awaiting the next assistant message (attached as url_citation annotations).
let pendingWebSources: { url: string; title?: string }[] = [];

Expand Down Expand Up @@ -1132,6 +1167,7 @@ export function buildResponseJSON(
currentToolCallId = "";
currentToolCallName = "";
currentToolCallArgs = "";
currentToolCallArgsBytes = 0;
};

for (const e of events) {
Expand Down Expand Up @@ -1186,9 +1222,34 @@ export function buildResponseJSON(
currentToolCallId = e.id;
currentToolCallName = e.name;
currentToolCallArgs = "";
currentToolCallArgsBytes = 0;
break;
case "tool_call_delta":
currentToolCallArgs += e.arguments;
{
const appended = appendToolArguments(
currentToolCallArgs,
currentToolCallArgsBytes,
retainedToolArgumentBytes,
e.arguments,
);
if (!appended) {
retainedToolArgumentBytes = Math.max(0, retainedToolArgumentBytes - currentToolCallArgsBytes);
currentToolCallArgs = "";
currentToolCallArgsBytes = 0;
flushToolCall("incomplete");
errorEvent = {
type: "error",
message: "upstream tool call arguments exceeded the configured byte limit",
status: 502,
errorType: "upstream_error",
};
toolArgumentOverflow = true;
break;
}
currentToolCallArgs = appended.value;
currentToolCallArgsBytes = appended.valueBytes;
retainedToolArgumentBytes = appended.retainedTurnBytes;
}
break;
case "tool_call_end":
if (!toolCallArgumentsUsable(currentToolCallArgs) && currentToolCallId) {
Expand Down Expand Up @@ -1249,6 +1310,7 @@ export function buildResponseJSON(
if (e.stopReason === "max_tokens" || e.stopReason === "content_filter") stopReason = e.stopReason;
break;
}
if (toolArgumentOverflow) break;
}
flushText(cleanDone && !errorEvent && !incompleteEvent ? "final_answer" : undefined);
flushSummaryReasoning();
Expand Down
Loading
Loading