From b3bad6b5818f6f354812f337fef626a38017cb2b Mon Sep 17 00:00:00 2001 From: Dax Date: Sun, 23 Aug 2026 13:03:20 -0400 Subject: [PATCH] fix(console): revert scanned zen request streaming (#44470) --- .../app/src/routes/zen/util/handler.ts | 204 ++++++++++++------ .../app/src/routes/zen/util/requestBody.ts | 191 ---------------- packages/console/app/test/requestBody.test.ts | 122 ----------- 3 files changed, 136 insertions(+), 381 deletions(-) delete mode 100644 packages/console/app/src/routes/zen/util/requestBody.ts delete 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 b0b502b9bb9..92c5515a646 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -28,7 +28,13 @@ import { GoUsageLimitError, BlackUsageLimitError, } from "./error" -import { buildCostChunk, createStreamPartConverter, createResponseConverter, UsageInfo } from "./provider/provider" +import { + buildCostChunk, + createBodyConverter, + createStreamPartConverter, + createResponseConverter, + UsageInfo, +} from "./provider/provider" import { anthropicHelper } from "./provider/anthropic" import { googleHelper } from "./provider/google" import { openaiHelper } from "./provider/openai" @@ -47,10 +53,12 @@ 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 PreparedBody = Awaited> +type RetryOptions = { + excludeProviders: string[] + retryCount: number +} type BillingSource = "anonymous" | "free" | "byok" | "subscription" | "lite" | "balance" function resolve(text: string, params?: Record) { @@ -78,6 +86,8 @@ 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 = [ @@ -86,14 +96,12 @@ export async function handler( "wrk_01KKZDKDWCS1VTJF8QTX62DD50", // contributors ] - let requestBody: PreparedBody | undefined try { const url = input.request.url - 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) : (requestBody?.model ?? "") - const googleStream = opts.format === "google" ? opts.parseIsStream(url, undefined) : undefined + const body = await input.request.text() + 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) 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) @@ -104,6 +112,7 @@ 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, @@ -165,7 +174,7 @@ export async function handler( ) const providerBudget = await providerBudgetTracker?.check() - const providerRequest = async () => { + const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => { const providerInfo = selectProvider( model, zenData, @@ -173,6 +182,7 @@ export async function handler( modelInfo, stickyId, trialProviders, + retry, stickyProvider, modelTpmLimits, modelTpsLimits, @@ -188,67 +198,95 @@ export async function handler( }) const startTimestamp = Date.now() - 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 (providerInfo.payloadModifier) - throw new Error("Zen provider payload modifiers are incompatible with streaming") - 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 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 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: " + (requestBody?.preview ?? "") + "...") + logger.debug("REQUEST: " + reqBody.substring(0, 300) + "...") const isNewInference = providerInfo.id.startsWith("console.") || providerInfo.id.startsWith("console-go.") || providerInfo.id.startsWith("inf.") || providerInfo.id.startsWith("inf-go.") - 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 }) + 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 }, + ) if (isNewInference) { const resEndpointId = res.headers.get("x-opencode-endpoint-id") @@ -267,10 +305,29 @@ export async function handler( }) } - return { providerInfo, res, startTimestamp, isStream } + // 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 } } - const { providerInfo, res, startTimestamp, isStream } = await providerRequest() + const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest() // Store sticky provider if (res.status === 200) await stickyTracker?.set(providerInfo.id) @@ -426,8 +483,6 @@ 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 @@ -552,6 +607,7 @@ export async function handler( modelInfo: ModelInfo, stickyId: string, trialProviders: string[] | undefined, + retry: RetryOptions, stickyProviderId: string | undefined, modelTpmLimits: Record | undefined, modelTpsLimits: Record | undefined, @@ -578,11 +634,14 @@ 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 @@ -990,6 +1049,15 @@ 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 deleted file mode 100644 index 458faf553e7..00000000000 --- a/packages/console/app/src/routes/zen/util/requestBody.ts +++ /dev/null @@ -1,191 +0,0 @@ -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 deleted file mode 100644 index 52d86297b4e..00000000000 --- a/packages/console/app/test/requestBody.test.ts +++ /dev/null @@ -1,122 +0,0 @@ -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, - }) - }) -})