Skip to content

Commit 6de584f

Browse files
fwangfwang
committed
zen: fix affiliated header
1 parent 8c72edb commit 6de584f

8 files changed

Lines changed: 22 additions & 81 deletions

File tree

‎packages/console/app/src/routes/zen/util/dataDumper.ts‎

Lines changed: 0 additions & 45 deletions
This file was deleted.

‎packages/console/app/src/routes/zen/util/handler.ts‎

Lines changed: 9 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ import { openaiHelper } from "./provider/openai"
3939
import { oaCompatHelper } from "./provider/openai-compatible"
4040
import { createRateLimiter as createIpRateLimiter } from "./ipRateLimiter"
4141
import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter"
42-
import { createDataDumper } from "./dataDumper"
4342
import { createTrialLimiter } from "./trialLimiter"
4443
import { createStickyTracker } from "./stickyProviderTracker"
4544
import { LiteData } from "@opencode-ai/console-core/lite.js"
@@ -103,7 +102,6 @@ export async function handler(
103102
const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey
104103
const sessionId = input.request.headers.get("x-opencode-session") ?? ""
105104
const requestId = input.request.headers.get("x-opencode-request") ?? ""
106-
const projectId = input.request.headers.get("x-opencode-project") ?? ""
107105
const ocClient = input.request.headers.get("x-opencode-client") ?? ""
108106
const userAgent = input.request.headers.get("user-agent") ?? ""
109107
logger.metric({
@@ -116,16 +114,16 @@ export async function handler(
116114
})
117115
const zenData = ZenData.list(opts.modelList)
118116
const modelInfo = validateModel(zenData, model)
119-
const dataDumper = createDataDumper(sessionId, requestId, projectId)
120117
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
121118
const trialProviders = await trialLimiter?.check()
122119
const rateLimiter = modelInfo.allowAnonymous
123120
? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
124121
: createKeyRateLimiter(modelInfo.id, modelInfo.rateLimit, zenApiKey, input.request)
125122
await rateLimiter?.check()
126-
const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, sessionId)
127-
const stickyProvider = await stickyTracker?.get()
128123
const authInfo = await authenticate(modelInfo, zenApiKey)
124+
const stickyId = sessionId ? sessionId : (authInfo?.workspaceID ?? ip)
125+
const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, stickyId)
126+
const stickyProvider = await stickyTracker?.get()
129127
const billingSource = validateBilling(authInfo, modelInfo)
130128
logger.metric({ source: billingSource })
131129
const modelTpmLimiter = createModelTpmLimiter(modelInfo.providers)
@@ -139,8 +137,7 @@ export async function handler(
139137
zenData,
140138
authInfo,
141139
modelInfo,
142-
ip,
143-
sessionId,
140+
stickyId,
144141
trialProviders,
145142
retry,
146143
stickyProvider,
@@ -167,13 +164,8 @@ export async function handler(
167164
if (Array.isArray(v)) return [[k, v]]
168165
if (typeof v === "object") return [[k, replacer(v)]]
169166
if (typeof v === "string") {
170-
if (v === "$ip") return [[k, ip]]
171167
if (v === "$workspace") return authInfo?.workspaceID ? [[k, authInfo?.workspaceID]] : []
172-
if (v === "$session") return sessionId ? [[k, sessionId]] : []
173-
if (v === "$user") {
174-
const user = sessionId ?? authInfo?.workspaceID ?? ip
175-
return user ? [[k, user]] : []
176-
}
168+
if (v === "$user") return stickyId ? [[k, stickyId]] : []
177169
if (v.startsWith("$header.")) {
178170
const headerValue = input.request.headers.get(v.slice(8))
179171
return headerValue ? [[k, headerValue]] : []
@@ -192,7 +184,7 @@ export async function handler(
192184
method: "POST",
193185
headers: (() => {
194186
const headers = new Headers(input.request.headers)
195-
providerInfo.modifyHeaders(headers, body, providerInfo.apiKey)
187+
providerInfo.modifyHeaders(headers, providerInfo.apiKey, stickyId)
196188
Object.entries(providerInfo.headerMappings ?? {}).forEach(([k, v]) => {
197189
headers.set(k, headers.get(v)!)
198190
})
@@ -237,10 +229,6 @@ export async function handler(
237229

238230
const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest()
239231

240-
// Store model request
241-
dataDumper?.provideModel(providerInfo.storeModel)
242-
dataDumper?.provideRequest(reqBody)
243-
244232
// Store sticky provider
245233
if (res.status === 200) await stickyTracker?.set(providerInfo.id)
246234

@@ -281,8 +269,6 @@ export async function handler(
281269
const body = JSON.stringify(responseConverter(json))
282270
logger.metric({ response_length: body.length })
283271
logger.debug("RESPONSE: " + body)
284-
dataDumper?.provideResponse(body)
285-
dataDumper?.flush()
286272
return new Response(body, {
287273
status: resStatus,
288274
statusText: res.statusText,
@@ -313,7 +299,6 @@ export async function handler(
313299
response_length: responseLength,
314300
"timestamp.last_byte": timestampLastByte,
315301
})
316-
dataDumper?.flush()
317302
await rateLimiter?.track()
318303
const usage = usageParser.retrieve()
319304
if (usage) {
@@ -351,7 +336,6 @@ export async function handler(
351336

352337
responseLength += value.length
353338
buffer += decoder.decode(value, { stream: true })
354-
dataDumper?.provideStream(buffer)
355339

356340
const parts = buffer.split(providerInfo.streamSeparator)
357341
buffer = parts.pop() ?? ""
@@ -490,8 +474,7 @@ export async function handler(
490474
zenData: ZenData,
491475
authInfo: AuthInfo,
492476
modelInfo: ModelInfo,
493-
ip: string,
494-
sessionId: string,
477+
stickyId: string,
495478
trialProviders: string[] | undefined,
496479
retry: RetryOptions,
497480
stickyProvider: string | undefined,
@@ -541,11 +524,10 @@ export async function handler(
541524
.flatMap((provider) => Array<typeof provider>(provider.weight).fill(provider))
542525

543526
// Use the last 4 characters of session ID to select a provider
544-
const identifier = sessionId.length ? sessionId : ip
545527
let h = 0
546-
const l = identifier.length
528+
const l = stickyId.length
547529
for (let i = l - 4; i < l; i++) {
548-
h = (h * 31 + identifier.charCodeAt(i)) | 0 // 32-bit int
530+
h = (h * 31 + stickyId.charCodeAt(i)) | 0 // 32-bit int
549531
}
550532
const index = (h >>> 0) % providers.length // make unsigned + range 0..length-1
551533
const provider = providers[index || 0]

‎packages/console/app/src/routes/zen/util/provider/anthropic.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) =>
2828
isBedrock
2929
? `${providerApi}/model/${isBedrockModelArn ? encodeURIComponent(providerModel) : providerModel}/${isStream ? "invoke-with-response-stream" : "invoke"}`
3030
: providerApi + "/messages",
31-
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
31+
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
3232
if (isBedrock || isDatabricks) {
3333
headers.set("Authorization", `Bearer ${apiKey}`)
3434
} else {

‎packages/console/app/src/routes/zen/util/provider/google.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({
3030
format: "google",
3131
modifyUrl: (providerApi: string, isStream?: boolean) =>
3232
`${providerApi}/models/${providerModel}:${isStream ? "streamGenerateContent?alt=sse" : "generateContent"}`,
33-
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
33+
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
3434
headers.set("x-goog-api-key", apiKey)
3535
},
3636
modifyBody: (body: Record<string, any>) => {

‎packages/console/app/src/routes/zen/util/provider/openai-compatible.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,9 @@ type Usage = {
2626
export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
2727
format: "oa-compat",
2828
modifyUrl: (providerApi: string) => providerApi + "/chat/completions",
29-
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
29+
modifyHeaders: (headers: Headers, apiKey: string, stickyId: string) => {
3030
headers.set("authorization", `Bearer ${apiKey}`)
31-
headers.set("x-session-affinity", headers.get("x-opencode-session") ?? "")
31+
headers.set("x-session-affinity", stickyId)
3232
},
3333
modifyBody: (body: Record<string, any>, _workspaceID?: string) => {
3434
return {

‎packages/console/app/src/routes/zen/util/provider/openai.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ type Usage = {
1515
export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
1616
format: "openai",
1717
modifyUrl: (providerApi: string) => providerApi + "/responses",
18-
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
18+
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
1919
headers.set("authorization", `Bearer ${apiKey}`)
2020
},
2121
modifyBody: (body: Record<string, any>) => body,

‎packages/console/app/src/routes/zen/util/provider/provider.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ export type ProviderHelper = (input: {
4141
}) => {
4242
format: ZenData.Format
4343
modifyUrl: (providerApi: string, isStream?: boolean) => string
44-
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => void
44+
modifyHeaders: (headers: Headers, apiKey: string, stickyId: string) => void
4545
modifyBody: (body: Record<string, any>) => Record<string, any>
4646
createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined
4747
streamSeparator: string

‎packages/console/app/src/routes/zen/util/stickyProviderTracker.ts‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,14 @@
11
import { Database, eq } from "@opencode-ai/console-core/drizzle/index.js"
22
import { ModelStickyProviderTable } from "@opencode-ai/console-core/schema/ip.sql.js"
33

4-
export function createStickyTracker(modelId: string, stickyProvider: "strict" | "prefer" | undefined, session: string) {
4+
export function createStickyTracker(
5+
modelId: string,
6+
stickyProvider: "strict" | "prefer" | undefined,
7+
stickyId: string,
8+
) {
59
if (!stickyProvider) return
6-
if (!session) return
7-
const id = `${modelId}/${session}`
10+
if (!stickyId) return
11+
const id = `${modelId}/${stickyId}`
812
let _providerId: string | undefined
913

1014
return {

0 commit comments

Comments
 (0)