fix(console): revert scanned zen request streaming (#44470)

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<ReturnType<typeof ZenData.list>>
-type PreparedBody = Awaited<ReturnType<typeof prepareRequestBody>>
+type RetryOptions = {
+  excludeProviders: string[]
+  retryCount: number
+}
 type BillingSource = "anonymous" | "free" | "byok" | "subscription" | "lite" | "balance"
 
 function resolve(text: string, params?: Record<string, string | number>) {
@@ -78,6 +86,8 @@ export async function handler(
   type ProviderInfo = Awaited<ReturnType<typeof selectProvider>>
   type CostInfo = ReturnType<typeof calculateCost>
 
+  const MAX_FAILOVER_RETRIES = 3
+  const MAX_RETRYABLE_STATUS_RETRIES = 3
   const dict = i18n(localeFromRequest(input.request))
   const t = (key: Key, params?: Record<string, string | number>) => 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<string, any>): Record<string, any> =>
+                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<string, number> | undefined,
     modelTpsLimits: Record<string, { qualify: number; unqualify: number }> | 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

packages/console/app/src/routes/zen/util/requestBody.ts

@@ -1,191 +0,0 @@
-const TAIL_LIMIT = 4 * 1024
-const encoder = new TextEncoder()
-
-export async function prepareRequestBody(body: ReadableStream<Uint8Array>) {
-  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<Uint8Array | undefined>,
-  reader: ReadableStreamDefaultReader<Uint8Array>,
-  sourceDone: boolean,
-) {
-  let done = sourceDone
-  let index = 0
-  return new ReadableStream<Uint8Array>({
-    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<Uint8Array>) {
-  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<Uint8Array>({
-    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)
-    },
-  })
-}

packages/console/app/test/requestBody.test.ts

@@ -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<Uint8Array>(
-      {
-        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<Uint8Array>(
-      {
-        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<Uint8Array>(
-      {
-        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,
-    })
-  })
-})