فهرست منبع

zen: fix affiliated header

Frank 2 ماه پیش
والد
کامیت
6de584f25c

+ 0 - 45
packages/console/app/src/routes/zen/util/dataDumper.ts

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

+ 9 - 27
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<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]

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

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

+ 2 - 2
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<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 {

+ 1 - 1
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<string, any>, apiKey: string) => {
+  modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
     headers.set("authorization", `Bearer ${apiKey}`)
   },
   modifyBody: (body: Record<string, any>) => body,

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

+ 7 - 3
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 {