From 63a883a4f770c5b8c7f8a9d37ac99718ef3b03c9 Mon Sep 17 00:00:00 2001 From: Dax Date: Sun, 23 Aug 2026 13:14:18 -0400 Subject: [PATCH] refactor(console): stream zen bodies without modifiers (#44472) --- .../app/src/routes/zen/util/handler.ts | 194 ++++++------------ .../app/src/routes/zen/util/requestBody.ts | 191 +++++++++++++++++ packages/console/app/test/requestBody.test.ts | 122 +++++++++++ 3 files changed, 378 insertions(+), 129 deletions(-) create mode 100644 packages/console/app/src/routes/zen/util/requestBody.ts create mode 100644 packages/console/app/test/requestBody.test.ts diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 92c5515a646..cd7e795f56b 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -30,7 +30,6 @@ import { } from "./error" import { buildCostChunk, - createBodyConverter, createStreamPartConverter, createResponseConverter, UsageInfo, @@ -53,12 +52,10 @@ import { createProviderBudgetTracker } from "./providerBudgetTracker" import { accumulateUsage, HOT_WORKSPACES } from "./usageBatcher" import { Workspace } from "@opencode-ai/console-core/workspace.js" import { countryFromRequest, isModelCountryRestricted } from "~/lib/request-country" +import { prepareRequestBody } from "./requestBody" type ZenData = Awaited> -type RetryOptions = { - excludeProviders: string[] - retryCount: number -} +type PreparedBody = Awaited> type BillingSource = "anonymous" | "free" | "byok" | "subscription" | "lite" | "balance" function resolve(text: string, params?: Record) { @@ -86,8 +83,6 @@ export async function handler( type ProviderInfo = Awaited> type CostInfo = ReturnType - const MAX_FAILOVER_RETRIES = 3 - const MAX_RETRYABLE_STATUS_RETRIES = 3 const dict = i18n(localeFromRequest(input.request)) const t = (key: Key, params?: Record) => resolve(dict[key], params) const ADMIN_WORKSPACES = [ @@ -96,12 +91,15 @@ export async function handler( "wrk_01KKZDKDWCS1VTJF8QTX62DD50", // contributors ] + let requestBody: PreparedBody | undefined try { const url = input.request.url - const body = await input.request.text() + const body = input.request.body + if (!body) throw new Error("Missing request body") + requestBody = opts.format === "google" ? undefined : await prepareRequestBody(body) const model = - opts.format === "google" ? opts.parseModel(url, undefined) : (body.match(/"model"\s*:\s*"([^"]+)"/)?.[1] ?? "") - const isStream = opts.format === "google" ? opts.parseIsStream(url, undefined) : /"stream"\s*:\s*true/.test(body) + opts.format === "google" ? opts.parseModel(url, undefined) : (requestBody?.model ?? "") + const googleStream = opts.format === "google" ? opts.parseIsStream(url, undefined) : undefined const rawIp = input.request.headers.get("x-real-ip") ?? "" const ip = rawIp.includes(":") ? rawIp.split(":").slice(0, 4).join(":") : rawIp const rawZenApiKey = opts.parseApiKey(input.request.headers) @@ -112,7 +110,6 @@ export async function handler( const projectId = input.request.headers.get("x-opencode-project") ?? "" const userAgent = input.request.headers.get("user-agent") ?? "" logger.metric({ - is_stream: isStream, session: sessionId, request: requestId, client: ocClient, @@ -174,7 +171,7 @@ export async function handler( ) const providerBudget = await providerBudgetTracker?.check() - const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => { + const providerRequest = async () => { const providerInfo = selectProvider( model, zenData, @@ -182,7 +179,6 @@ export async function handler( modelInfo, stickyId, trialProviders, - retry, stickyProvider, modelTpmLimits, modelTpsLimits, @@ -198,95 +194,65 @@ export async function handler( }) const startTimestamp = Date.now() - const reqUrl = providerInfo.modifyUrl(providerInfo.api, isStream) - const directBody = (() => { - const specialAnthropic = - providerInfo.format === "anthropic" && - (providerInfo.model.startsWith("arn:aws:bedrock:") || - providerInfo.model.startsWith("global.anthropic.") || - providerInfo.model.startsWith("databricks-claude-")) - if (providerInfo.format === opts.format && !providerInfo.payloadModifier && !specialAnthropic) { - const patched = body.replace(/"model"\s*:\s*"[^"]+"/, `"model":${JSON.stringify(providerInfo.model)}`) - if (providerInfo.format !== "oa-compat" || !isStream) return patched - return patched.replace(/}\s*$/, ',"stream_options":{"include_usage":true}}') - } - return undefined + const reqUrl = providerInfo.modifyUrl(providerInfo.api, googleStream ?? false) + const specialAnthropic = + providerInfo.format === "anthropic" && + (providerInfo.model.startsWith("arn:aws:bedrock:") || + providerInfo.model.startsWith("global.anthropic.") || + providerInfo.model.startsWith("databricks-claude-")) + if (providerInfo.format !== opts.format) throw new Error("Zen provider format must match request format") + if (specialAnthropic) throw new Error("Anthropic provider body modifiers are incompatible with streaming") + const prepared = requestBody + + const reqBody = (() => { + if (opts.format === "google") return body + if (!prepared) throw new Error("Missing prepared request body") + return prepared.stream(providerInfo.model, providerInfo.format === "oa-compat") })() - const reqBody = - directBody ?? - JSON.stringify( - providerInfo.modifyBody({ - ...createBodyConverter(opts.format, providerInfo.format)(JSON.parse(body)), - model: providerInfo.model, - ...(() => { - const replacer = (obj: Record): Record => - Object.fromEntries( - Object.entries(obj).flatMap(([k, v]) => { - if (Array.isArray(v)) return [[k, v]] - if (typeof v === "object") return [[k, replacer(v)]] - if (typeof v === "string") { - if (v === "$workspace") return authInfo?.workspaceID ? [[k, authInfo.workspaceID]] : [] - if (v === "$org") - return authInfo?.workspaceID ? [[k, authInfo.workspaceID.replace("wrk_", "org_")]] : [] - if (v === "$user") return stickyId ? [[k, stickyId]] : [] - if (v.startsWith("$header.")) { - const headerValue = input.request.headers.get(v.slice(8)) - return headerValue ? [[k, headerValue]] : [] - } - } - return [[k, v]] - }), - ) - return replacer(providerInfo.payloadModifier ?? {}) - })(), - }), - ) logger.debug("REQUEST URL: " + reqUrl) - logger.debug("REQUEST: " + reqBody.substring(0, 300) + "...") + logger.debug("REQUEST: " + (requestBody?.preview ?? "") + "...") const isNewInference = providerInfo.id.startsWith("console.") || providerInfo.id.startsWith("console-go.") || providerInfo.id.startsWith("inf.") || providerInfo.id.startsWith("inf-go.") - const res = await fetchWithRetryableStatus( - reqUrl, - { - method: "POST", - headers: (() => { - const headers = new Headers(input.request.headers) - providerInfo.modifyHeaders(headers, providerInfo.apiKey, stickyId) - Object.entries(providerInfo.headerModifier ?? {}).forEach(([k, v]) => { - if (v === "$ip") return headers.set(k, ip) - if (v === "$caller") return headers.set(k, stickyId) - if (v === "$session") return headers.set(k, sessionId) - if (v === "$model") return headers.set(k, model) - if (v === "$request") return headers.set(k, requestId) - if (v === "$project") return headers.set(k, projectId) - if (v === "$workspace") { - if (authInfo?.workspaceID) headers.set(k, authInfo.workspaceID) - return - } - if (v === "$org") { - if (authInfo?.workspaceID) headers.set(k, authInfo.workspaceID.replace("wrk_", "org_")) - return - } - headers.set(k, v) - }) - headers.delete("host") - headers.delete("content-length") - headers.delete("x-opencode-request") - if (!isNewInference) headers.delete("x-opencode-session") - headers.delete("x-opencode-project") - headers.delete("x-opencode-client") - return headers - })(), - body: reqBody, - // Propagate caller disconnects to the upstream provider request so - // abandoned Console requests do not leave orphaned inference work open. - signal: input.request.signal, - }, - { count: isNewInference ? MAX_RETRYABLE_STATUS_RETRIES : 0 }, - ) + const res = await fetch(reqUrl, { + method: "POST", + headers: (() => { + const headers = new Headers(input.request.headers) + providerInfo.modifyHeaders(headers, providerInfo.apiKey, stickyId) + Object.entries(providerInfo.headerModifier ?? {}).forEach(([k, v]) => { + if (v === "$ip") return headers.set(k, ip) + if (v === "$caller") return headers.set(k, stickyId) + if (v === "$session") return headers.set(k, sessionId) + if (v === "$model") return headers.set(k, model) + if (v === "$request") return headers.set(k, requestId) + if (v === "$project") return headers.set(k, projectId) + if (v === "$workspace") { + if (authInfo?.workspaceID) headers.set(k, authInfo.workspaceID) + return + } + if (v === "$org") { + if (authInfo?.workspaceID) headers.set(k, authInfo.workspaceID.replace("wrk_", "org_")) + return + } + headers.set(k, v) + }) + headers.delete("host") + headers.delete("content-length") + headers.delete("x-opencode-request") + if (!isNewInference) headers.delete("x-opencode-session") + headers.delete("x-opencode-project") + headers.delete("x-opencode-client") + return headers + })(), + body: reqBody, + // Propagate caller disconnects to the upstream provider request so + // abandoned Console requests do not leave orphaned inference work open. + signal: input.request.signal, + }) + const isStream = res.headers.get("content-type")?.toLowerCase().includes("text/event-stream") ?? false + logger.metric({ is_stream: isStream }) if (isNewInference) { const resEndpointId = res.headers.get("x-opencode-endpoint-id") @@ -305,29 +271,10 @@ export async function handler( }) } - // Try another provider => stop retrying if using fallback provider - if ( - //!isNewInference && - res.status !== 200 && - // ie. 400 error is usually provider error like malformed request - res.status !== 400 && - // ie. openai 404 error: Item with id 'msg_0ead8b004a3b165d0069436a6b6834819896da85b63b196a3f' not found. - !(modelInfo.id.startsWith("gpt-") && res.status === 404) && - // ie. cannot change codex model providers mid-session - modelInfo.stickyProvider !== "strict" && - modelInfo.fallbackProvider && - providerInfo.id !== modelInfo.fallbackProvider - ) { - return retriableRequest({ - excludeProviders: [...retry.excludeProviders, providerInfo.id], - retryCount: retry.retryCount + 1, - }) - } - - return { providerInfo, reqBody, res, startTimestamp } + return { providerInfo, res, startTimestamp, isStream } } - const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest() + const { providerInfo, res, startTimestamp, isStream } = await providerRequest() // Store sticky provider if (res.status === 200) await stickyTracker?.set(providerInfo.id) @@ -483,6 +430,8 @@ export async function handler( headers: resHeaders, }) } catch (error: any) { + if (requestBody) void requestBody.cancel().catch(() => {}) + else void input.request.body?.cancel().catch(() => {}) // The caller disconnected before we finished. Because the outbound provider // request shares input.request.signal, an aborted caller surfaces here as an // AbortError. There is no client left to receive a body, so skip the error @@ -607,7 +556,6 @@ export async function handler( modelInfo: ModelInfo, stickyId: string, trialProviders: string[] | undefined, - retry: RetryOptions, stickyProviderId: string | undefined, modelTpmLimits: Record | undefined, modelTpsLimits: Record | undefined, @@ -634,14 +582,11 @@ export async function handler( })) } - // Use fallback provider if max retries reached const fallbackProvider = allProviders.find((provider) => provider.id === modelInfo.fallbackProvider) - if (retry.retryCount === MAX_FAILOVER_RETRIES) return fallbackProvider let topPriority = Infinity const providers = allProviders .filter((provider) => provider.weight !== 0) - .filter((provider) => !retry.excludeProviders.includes(provider.id)) .filter((provider) => { if (provider.budgetPriority === undefined) return true if (!providerBudget) return true @@ -1049,15 +994,6 @@ export async function handler( providerInfo.apiKey = authInfo.provider.credentials } - async function fetchWithRetryableStatus(url: string, options: RequestInit, retry = { count: 0 }) { - const res = await fetch(url, options) - if ([429, 529].includes(res.status) && retry.count < MAX_RETRYABLE_STATUS_RETRIES) { - await new Promise((resolve) => setTimeout(resolve, Math.pow(2, retry.count) * 500)) - return fetchWithRetryableStatus(url, options, { count: retry.count + 1 }) - } - return res - } - function calculateCost(modelInfo: ModelInfo, usageInfo: UsageInfo) { const { inputTokens, outputTokens, reasoningTokens, cacheReadTokens, cacheWrite5mTokens, cacheWrite1hTokens } = usageInfo diff --git a/packages/console/app/src/routes/zen/util/requestBody.ts b/packages/console/app/src/routes/zen/util/requestBody.ts new file mode 100644 index 00000000000..458faf553e7 --- /dev/null +++ b/packages/console/app/src/routes/zen/util/requestBody.ts @@ -0,0 +1,191 @@ +const TAIL_LIMIT = 4 * 1024 +const encoder = new TextEncoder() + +export async function prepareRequestBody(body: ReadableStream) { + const reader = body.getReader() + const chunks: Uint8Array[] = [] + const decoder = new TextDecoder() + let text = "" + let done = false + let searchFrom = 0 + let bom = 0 + let match: RegExpExecArray | null = null + const pattern = /("model"\s*:\s*")([^"]+)"/g + + while (!done && !match) { + const next = await reader.read() + done = next.done + if (!next.value) continue + if (!chunks.length && next.value[0] === 0xef && next.value[1] === 0xbb && next.value[2] === 0xbf) bom = 3 + chunks.push(next.value) + text += decoder.decode(next.value, { stream: true }) + pattern.lastIndex = searchFrom + match = pattern.exec(text) + searchFrom = Math.max(0, text.length - 256) + } + if (done) { + text += decoder.decode() + if (!match) { + pattern.lastIndex = searchFrom + match = pattern.exec(text) + } + } + + const found = (() => { + if (!match) return + const start = bom + utf8Length(text, match.index + match[1].length) + return { model: match[2], start, end: start + utf8Length(match[2], match[2].length) } + })() + const preview = text.substring(0, 300) + text = "" + match = null + let used = false + + return { + model: found?.model ?? "", + preview, + cancel: () => reader.cancel(), + stream(providerModel: string, includeUsage: boolean) { + if (used) throw new Error("Request body stream already consumed") + if (!found) throw new Error("Missing model field") + used = true + + const initial = replace(chunks, found.start, found.end, providerModel) + chunks.length = 0 + const output = passthrough(initial, reader, done) + if (!includeUsage) return output + return appendUsage(output) + }, + } +} + +function utf8Length(value: string, end: number) { + let length = 0 + for (let i = 0; i < end; i++) { + const code = value.charCodeAt(i) + if (code <= 0x7f) length++ + else if (code <= 0x7ff) length += 2 + else if (code >= 0xd800 && code <= 0xdbff && i + 1 < end && value.charCodeAt(i + 1) >= 0xdc00) { + length += 4 + i++ + } else length += 3 + } + return length +} + +function replace(chunks: Uint8Array[], start: number, end: number, value: string) { + let offset = 0 + let inserted = false + return chunks.flatMap((chunk) => { + const chunkStart = offset + const chunkEnd = offset + chunk.length + offset = chunkEnd + if (chunkEnd <= start || chunkStart >= end) return [chunk] + + const parts = [chunk.subarray(0, Math.max(0, start - chunkStart))] + if (!inserted) { + parts.push(encoder.encode(value)) + inserted = true + } + parts.push(chunk.subarray(Math.min(chunk.length, end - chunkStart))) + return parts.filter((part) => part.length) + }) +} + +function passthrough( + initial: Array, + reader: ReadableStreamDefaultReader, + sourceDone: boolean, +) { + let done = sourceDone + let index = 0 + return new ReadableStream({ + async pull(controller) { + const chunk = initial[index] + if (chunk) { + initial[index++] = undefined + controller.enqueue(chunk) + return + } + initial.length = 0 + if (done) { + controller.close() + return + } + const next = await reader.read() + done = next.done + if (next.value) controller.enqueue(next.value) + if (done) controller.close() + }, + cancel(reason) { + initial.length = 0 + return reader.cancel(reason) + }, + }) +} + +function appendUsage(body: ReadableStream) { + const reader = body.getReader() + const decoder = new TextDecoder() + let tail = new Uint8Array() + let streamText = "" + let isStream = false + const inspect = (chunk?: Uint8Array) => { + streamText += chunk ? decoder.decode(chunk, { stream: true }) : decoder.decode() + for (const match of streamText.matchAll(/"stream"\s*:\s*(true|false)/g)) isStream = match[1] === "true" + streamText = streamText.slice(-64) + } + return new ReadableStream({ + async pull(controller) { + while (true) { + const next = await reader.read() + if (next.done) { + inspect() + if (!isStream) { + if (tail.length) controller.enqueue(tail) + controller.close() + return + } + const close = tail.lastIndexOf(125) + if (close < 0) { + controller.error(new Error("Invalid JSON request body")) + return + } + if (close) controller.enqueue(tail.subarray(0, close)) + controller.enqueue(encoder.encode(',"stream_options":{"include_usage":true}}')) + if (close + 1 < tail.length) controller.enqueue(tail.subarray(close + 1)) + controller.close() + return + } + + const chunk = next.value + inspect(chunk) + if (tail.length + chunk.length <= TAIL_LIMIT) { + const combined = new Uint8Array(tail.length + chunk.length) + combined.set(tail) + combined.set(chunk, tail.length) + tail = combined + continue + } + + const emit = tail.length + chunk.length - TAIL_LIMIT + if (emit <= tail.length) { + controller.enqueue(tail.subarray(0, emit)) + const combined = new Uint8Array(TAIL_LIMIT) + combined.set(tail.subarray(emit)) + combined.set(chunk, tail.length - emit) + tail = combined + return + } + + if (tail.length) controller.enqueue(tail) + controller.enqueue(chunk.subarray(0, emit - tail.length)) + tail = chunk.slice(emit - tail.length) + return + } + }, + cancel(reason) { + return reader.cancel(reason) + }, + }) +} diff --git a/packages/console/app/test/requestBody.test.ts b/packages/console/app/test/requestBody.test.ts new file mode 100644 index 00000000000..52d86297b4e --- /dev/null +++ b/packages/console/app/test/requestBody.test.ts @@ -0,0 +1,122 @@ +import { describe, expect, test } from "bun:test" +import { prepareRequestBody } from "../src/routes/zen/util/requestBody" + +describe("Zen request body streaming", () => { + test("patches the leading model without buffering the remaining body", async () => { + let reads = 0 + const body = new ReadableStream( + { + pull(controller) { + const chunks = [ + '{"model":"client-model","stream":true,"messages":[', + JSON.stringify({ role: "user", content: "large payload" }), + "]}", + ] + const chunk = chunks[reads++] + if (chunk) controller.enqueue(new TextEncoder().encode(chunk)) + else controller.close() + }, + }, + { highWaterMark: 0 }, + ) + + const request = await prepareRequestBody(body) + expect(request.model).toBe("client-model") + expect(reads).toBe(1) + + const output = await new Response(request.stream("provider-model", false)).text() + expect(JSON.parse(output)).toEqual({ + model: "provider-model", + stream: true, + messages: [{ role: "user", content: "large payload" }], + }) + }) + + test("appends stream usage options at the end of the request", async () => { + const body = new Blob(['{"model":"client-model","stream":true,"messages":[]} ']).stream() + const request = await prepareRequestBody(body) + const output = await new Response(request.stream("provider-model", true)).text() + + expect(JSON.parse(output)).toEqual({ + model: "provider-model", + stream: true, + messages: [], + stream_options: { include_usage: true }, + }) + expect(output.endsWith(" ")).toBe(true) + }) + + test("detects streaming after a large message while forwarding", async () => { + const content = "x".repeat(128 * 1024) + let reads = 0 + const chunks = [ + '{"model":"client-model","messages":[', + JSON.stringify({ role: "user", content }), + '],"stream":true}', + ] + const body = new ReadableStream( + { + pull(controller) { + const chunk = chunks[reads++] + if (chunk) controller.enqueue(new TextEncoder().encode(chunk)) + else controller.close() + }, + }, + { highWaterMark: 0 }, + ) + const request = await prepareRequestBody(body) + expect(reads).toBe(1) + const output = await new Response(request.stream("provider-model", true)).text() + + expect(JSON.parse(output)).toEqual({ + model: "provider-model", + messages: [{ role: "user", content }], + stream: true, + stream_options: { include_usage: true }, + }) + }) + + test("buffers through a late model field and then streams the rest", async () => { + const content = "こんにちは".repeat(32 * 1024) + let reads = 0 + const chunks = [ + '{"messages":[', + JSON.stringify({ role: "user", content }), + '],"model":"client-model","stream":true,"extra":"after-model"}', + ] + const body = new ReadableStream( + { + pull(controller) { + const chunk = chunks[reads++] + if (chunk) controller.enqueue(new TextEncoder().encode(chunk)) + else controller.close() + }, + }, + { highWaterMark: 0 }, + ) + const request = await prepareRequestBody(body) + + expect(request.model).toBe("client-model") + expect(reads).toBe(3) + expect(JSON.parse(await new Response(request.stream("provider-model", true)).text())).toEqual({ + messages: [{ role: "user", content }], + model: "provider-model", + stream: true, + extra: "after-model", + stream_options: { include_usage: true }, + }) + }) + + test("preserves a UTF-8 BOM while patching the model", async () => { + const body = new Blob(['\uFEFF{"messages":[],"model":"client-model","stream":false}']).stream() + const request = await prepareRequestBody(body) + const output = new Uint8Array(await new Response(request.stream("provider-model", false)).arrayBuffer()) + + expect([...output.subarray(0, 3)]).toEqual([0xef, 0xbb, 0xbf]) + expect(JSON.parse(new TextDecoder().decode(output))).toEqual({ + messages: [], + model: "provider-model", + stream: false, + }) + }) +})