zen: fix affiliated header

This commit is contained in:
Frank
2026-05-21 02:36:08 -04:00
parent 8c72edb3ff
commit 6de584f25c
8 changed files with 22 additions and 81 deletions

View File

@@ -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<string, any> = { sessionId, requestId, projectId }
let metadata: Record<string, any> = { 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 }),
),
)
},
}
}

View File

@@ -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<typeof provider>(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]

View File

@@ -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<string, any>, apiKey: string) => {
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
if (isBedrock || isDatabricks) {
headers.set("Authorization", `Bearer ${apiKey}`)
} else {

View File

@@ -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<string, any>, apiKey: string) => {
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
headers.set("x-goog-api-key", apiKey)
},
modifyBody: (body: Record<string, any>) => {

View File

@@ -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<string, any>, 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<string, any>, _workspaceID?: string) => {
return {

View File

@@ -15,7 +15,7 @@ type Usage = {
export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
format: "openai",
modifyUrl: (providerApi: string) => providerApi + "/responses",
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
headers.set("authorization", `Bearer ${apiKey}`)
},
modifyBody: (body: Record<string, any>) => body,

View File

@@ -41,7 +41,7 @@ export type ProviderHelper = (input: {
}) => {
format: ZenData.Format
modifyUrl: (providerApi: string, isStream?: boolean) => string
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => void
modifyHeaders: (headers: Headers, apiKey: string, stickyId: string) => void
modifyBody: (body: Record<string, any>) => Record<string, any>
createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined
streamSeparator: string

View File

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