refactor(console): stream zen bodies without modifiers (#44472)

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<ReturnType<typeof ZenData.list>>
-type RetryOptions = {
-  excludeProviders: string[]
-  retryCount: number
-}
+type PreparedBody = Awaited<ReturnType<typeof prepareRequestBody>>
 type BillingSource = "anonymous" | "free" | "byok" | "subscription" | "lite" | "balance"
 
 function resolve(text: string, params?: Record<string, string | number>) {
@@ -86,8 +83,6 @@ 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 = [
@@ -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<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: " + 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<string, number> | undefined,
     modelTpsLimits: Record<string, { qualify: number; unqualify: number }> | 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

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<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

@@ -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<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,
+    })
+  })
+})