From 6de584f25cf577992e18236a52bb092f52f75993 Mon Sep 17 00:00:00 2001 From: Frank Date: Thu, 21 May 2026 02:36:08 -0400 Subject: [PATCH] zen: fix affiliated header --- .../app/src/routes/zen/util/dataDumper.ts | 45 ------------------- .../app/src/routes/zen/util/handler.ts | 36 ++++----------- .../src/routes/zen/util/provider/anthropic.ts | 2 +- .../src/routes/zen/util/provider/google.ts | 2 +- .../zen/util/provider/openai-compatible.ts | 4 +- .../src/routes/zen/util/provider/openai.ts | 2 +- .../src/routes/zen/util/provider/provider.ts | 2 +- .../routes/zen/util/stickyProviderTracker.ts | 10 +++-- 8 files changed, 22 insertions(+), 81 deletions(-) delete mode 100644 packages/console/app/src/routes/zen/util/dataDumper.ts diff --git a/packages/console/app/src/routes/zen/util/dataDumper.ts b/packages/console/app/src/routes/zen/util/dataDumper.ts deleted file mode 100644 index 4027e23341..0000000000 --- a/packages/console/app/src/routes/zen/util/dataDumper.ts +++ /dev/null @@ -1,45 +0,0 @@ -import { Resource, waitUntil } from "@opencode-ai/console-resource" - -export function createDataDumper(sessionId: string, requestId: string, projectId: string) { - return - if (Resource.App.stage !== "production") return - if (sessionId === "") return - - let data: Record = { sessionId, requestId, projectId } - let metadata: Record = { sessionId, requestId, projectId } - - return { - provideModel: (model?: string) => { - data.modelName = model - metadata.modelName = model - }, - provideRequest: (request: string) => (data.request = request), - provideResponse: (response: string) => (data.response = response), - provideStream: (chunk: string) => (data.response = (data.response ?? "") + chunk), - flush: () => { - if (!data.modelName) return - - const timestamp = new Date().toISOString().replace(/[^0-9]/g, "") - const year = timestamp.substring(0, 4) - const month = timestamp.substring(4, 6) - const day = timestamp.substring(6, 8) - const hour = timestamp.substring(8, 10) - const minute = timestamp.substring(10, 12) - const second = timestamp.substring(12, 14) - - void waitUntil( - Resource.ZenDataNew.put( - `data/${data.modelName}/${year}/${month}/${day}/${hour}/${minute}/${second}/${requestId}.json`, - JSON.stringify({ timestamp, ...data }), - ), - ) - - void waitUntil( - Resource.ZenDataNew.put( - `meta/${data.modelName}/${sessionId}/${requestId}.json`, - JSON.stringify({ timestamp, ...metadata }), - ), - ) - }, - } -} diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 6d11a94730..e4b42d741e 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -39,7 +39,6 @@ import { openaiHelper } from "./provider/openai" import { oaCompatHelper } from "./provider/openai-compatible" import { createRateLimiter as createIpRateLimiter } from "./ipRateLimiter" import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter" -import { createDataDumper } from "./dataDumper" import { createTrialLimiter } from "./trialLimiter" import { createStickyTracker } from "./stickyProviderTracker" import { LiteData } from "@opencode-ai/console-core/lite.js" @@ -103,7 +102,6 @@ export async function handler( const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey const sessionId = input.request.headers.get("x-opencode-session") ?? "" const requestId = input.request.headers.get("x-opencode-request") ?? "" - const projectId = input.request.headers.get("x-opencode-project") ?? "" const ocClient = input.request.headers.get("x-opencode-client") ?? "" const userAgent = input.request.headers.get("user-agent") ?? "" logger.metric({ @@ -116,16 +114,16 @@ export async function handler( }) const zenData = ZenData.list(opts.modelList) const modelInfo = validateModel(zenData, model) - const dataDumper = createDataDumper(sessionId, requestId, projectId) const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip) const trialProviders = await trialLimiter?.check() const rateLimiter = modelInfo.allowAnonymous ? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request) : createKeyRateLimiter(modelInfo.id, modelInfo.rateLimit, zenApiKey, input.request) await rateLimiter?.check() - const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, sessionId) - const stickyProvider = await stickyTracker?.get() const authInfo = await authenticate(modelInfo, zenApiKey) + const stickyId = sessionId ? sessionId : (authInfo?.workspaceID ?? ip) + const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, stickyId) + const stickyProvider = await stickyTracker?.get() const billingSource = validateBilling(authInfo, modelInfo) logger.metric({ source: billingSource }) const modelTpmLimiter = createModelTpmLimiter(modelInfo.providers) @@ -139,8 +137,7 @@ export async function handler( zenData, authInfo, modelInfo, - ip, - sessionId, + stickyId, trialProviders, retry, stickyProvider, @@ -167,13 +164,8 @@ export async function handler( if (Array.isArray(v)) return [[k, v]] if (typeof v === "object") return [[k, replacer(v)]] if (typeof v === "string") { - if (v === "$ip") return [[k, ip]] if (v === "$workspace") return authInfo?.workspaceID ? [[k, authInfo?.workspaceID]] : [] - if (v === "$session") return sessionId ? [[k, sessionId]] : [] - if (v === "$user") { - const user = sessionId ?? authInfo?.workspaceID ?? ip - return user ? [[k, user]] : [] - } + if (v === "$user") return stickyId ? [[k, stickyId]] : [] if (v.startsWith("$header.")) { const headerValue = input.request.headers.get(v.slice(8)) return headerValue ? [[k, headerValue]] : [] @@ -192,7 +184,7 @@ export async function handler( method: "POST", headers: (() => { const headers = new Headers(input.request.headers) - providerInfo.modifyHeaders(headers, body, providerInfo.apiKey) + providerInfo.modifyHeaders(headers, providerInfo.apiKey, stickyId) Object.entries(providerInfo.headerMappings ?? {}).forEach(([k, v]) => { headers.set(k, headers.get(v)!) }) @@ -237,10 +229,6 @@ export async function handler( const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest() - // Store model request - dataDumper?.provideModel(providerInfo.storeModel) - dataDumper?.provideRequest(reqBody) - // Store sticky provider if (res.status === 200) await stickyTracker?.set(providerInfo.id) @@ -281,8 +269,6 @@ export async function handler( const body = JSON.stringify(responseConverter(json)) logger.metric({ response_length: body.length }) logger.debug("RESPONSE: " + body) - dataDumper?.provideResponse(body) - dataDumper?.flush() return new Response(body, { status: resStatus, statusText: res.statusText, @@ -313,7 +299,6 @@ export async function handler( response_length: responseLength, "timestamp.last_byte": timestampLastByte, }) - dataDumper?.flush() await rateLimiter?.track() const usage = usageParser.retrieve() if (usage) { @@ -351,7 +336,6 @@ export async function handler( responseLength += value.length buffer += decoder.decode(value, { stream: true }) - dataDumper?.provideStream(buffer) const parts = buffer.split(providerInfo.streamSeparator) buffer = parts.pop() ?? "" @@ -490,8 +474,7 @@ export async function handler( zenData: ZenData, authInfo: AuthInfo, modelInfo: ModelInfo, - ip: string, - sessionId: string, + stickyId: string, trialProviders: string[] | undefined, retry: RetryOptions, stickyProvider: string | undefined, @@ -541,11 +524,10 @@ export async function handler( .flatMap((provider) => Array(provider.weight).fill(provider)) // Use the last 4 characters of session ID to select a provider - const identifier = sessionId.length ? sessionId : ip let h = 0 - const l = identifier.length + const l = stickyId.length for (let i = l - 4; i < l; i++) { - h = (h * 31 + identifier.charCodeAt(i)) | 0 // 32-bit int + h = (h * 31 + stickyId.charCodeAt(i)) | 0 // 32-bit int } const index = (h >>> 0) % providers.length // make unsigned + range 0..length-1 const provider = providers[index || 0] diff --git a/packages/console/app/src/routes/zen/util/provider/anthropic.ts b/packages/console/app/src/routes/zen/util/provider/anthropic.ts index d93bc58d82..8c394ee3e1 100644 --- a/packages/console/app/src/routes/zen/util/provider/anthropic.ts +++ b/packages/console/app/src/routes/zen/util/provider/anthropic.ts @@ -28,7 +28,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) => isBedrock ? `${providerApi}/model/${isBedrockModelArn ? encodeURIComponent(providerModel) : providerModel}/${isStream ? "invoke-with-response-stream" : "invoke"}` : providerApi + "/messages", - modifyHeaders: (headers: Headers, body: Record, apiKey: string) => { + modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => { if (isBedrock || isDatabricks) { headers.set("Authorization", `Bearer ${apiKey}`) } else { diff --git a/packages/console/app/src/routes/zen/util/provider/google.ts b/packages/console/app/src/routes/zen/util/provider/google.ts index ef7937c358..2954024e20 100644 --- a/packages/console/app/src/routes/zen/util/provider/google.ts +++ b/packages/console/app/src/routes/zen/util/provider/google.ts @@ -30,7 +30,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({ format: "google", modifyUrl: (providerApi: string, isStream?: boolean) => `${providerApi}/models/${providerModel}:${isStream ? "streamGenerateContent?alt=sse" : "generateContent"}`, - modifyHeaders: (headers: Headers, body: Record, apiKey: string) => { + modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => { headers.set("x-goog-api-key", apiKey) }, modifyBody: (body: Record) => { diff --git a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts index e6dedb1a4b..912c89092a 100644 --- a/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts +++ b/packages/console/app/src/routes/zen/util/provider/openai-compatible.ts @@ -26,9 +26,9 @@ type Usage = { export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({ format: "oa-compat", modifyUrl: (providerApi: string) => providerApi + "/chat/completions", - modifyHeaders: (headers: Headers, body: Record, apiKey: string) => { + modifyHeaders: (headers: Headers, apiKey: string, stickyId: string) => { headers.set("authorization", `Bearer ${apiKey}`) - headers.set("x-session-affinity", headers.get("x-opencode-session") ?? "") + headers.set("x-session-affinity", stickyId) }, modifyBody: (body: Record, _workspaceID?: string) => { return { diff --git a/packages/console/app/src/routes/zen/util/provider/openai.ts b/packages/console/app/src/routes/zen/util/provider/openai.ts index 1c5cbdb3c9..4b39407d44 100644 --- a/packages/console/app/src/routes/zen/util/provider/openai.ts +++ b/packages/console/app/src/routes/zen/util/provider/openai.ts @@ -15,7 +15,7 @@ type Usage = { export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({ format: "openai", modifyUrl: (providerApi: string) => providerApi + "/responses", - modifyHeaders: (headers: Headers, body: Record, apiKey: string) => { + modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => { headers.set("authorization", `Bearer ${apiKey}`) }, modifyBody: (body: Record) => body, diff --git a/packages/console/app/src/routes/zen/util/provider/provider.ts b/packages/console/app/src/routes/zen/util/provider/provider.ts index 86446bfd85..319f8fdca3 100644 --- a/packages/console/app/src/routes/zen/util/provider/provider.ts +++ b/packages/console/app/src/routes/zen/util/provider/provider.ts @@ -41,7 +41,7 @@ export type ProviderHelper = (input: { }) => { format: ZenData.Format modifyUrl: (providerApi: string, isStream?: boolean) => string - modifyHeaders: (headers: Headers, body: Record, apiKey: string) => void + modifyHeaders: (headers: Headers, apiKey: string, stickyId: string) => void modifyBody: (body: Record) => Record createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined streamSeparator: string diff --git a/packages/console/app/src/routes/zen/util/stickyProviderTracker.ts b/packages/console/app/src/routes/zen/util/stickyProviderTracker.ts index fae0cdf03c..67fe3ed405 100644 --- a/packages/console/app/src/routes/zen/util/stickyProviderTracker.ts +++ b/packages/console/app/src/routes/zen/util/stickyProviderTracker.ts @@ -1,10 +1,14 @@ import { Database, eq } from "@opencode-ai/console-core/drizzle/index.js" import { ModelStickyProviderTable } from "@opencode-ai/console-core/schema/ip.sql.js" -export function createStickyTracker(modelId: string, stickyProvider: "strict" | "prefer" | undefined, session: string) { +export function createStickyTracker( + modelId: string, + stickyProvider: "strict" | "prefer" | undefined, + stickyId: string, +) { if (!stickyProvider) return - if (!session) return - const id = `${modelId}/${session}` + if (!stickyId) return + const id = `${modelId}/${stickyId}` let _providerId: string | undefined return {