zen: fix affiliated header
This commit is contained in:
@@ -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 }),
|
|
||||||
),
|
|
||||||
)
|
|
||||||
},
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -39,7 +39,6 @@ import { openaiHelper } from "./provider/openai"
|
|||||||
import { oaCompatHelper } from "./provider/openai-compatible"
|
import { oaCompatHelper } from "./provider/openai-compatible"
|
||||||
import { createRateLimiter as createIpRateLimiter } from "./ipRateLimiter"
|
import { createRateLimiter as createIpRateLimiter } from "./ipRateLimiter"
|
||||||
import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter"
|
import { createRateLimiter as createKeyRateLimiter } from "./keyRateLimiter"
|
||||||
import { createDataDumper } from "./dataDumper"
|
|
||||||
import { createTrialLimiter } from "./trialLimiter"
|
import { createTrialLimiter } from "./trialLimiter"
|
||||||
import { createStickyTracker } from "./stickyProviderTracker"
|
import { createStickyTracker } from "./stickyProviderTracker"
|
||||||
import { LiteData } from "@opencode-ai/console-core/lite.js"
|
import { LiteData } from "@opencode-ai/console-core/lite.js"
|
||||||
@@ -103,7 +102,6 @@ export async function handler(
|
|||||||
const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey
|
const zenApiKey = rawZenApiKey === "public" ? undefined : rawZenApiKey
|
||||||
const sessionId = input.request.headers.get("x-opencode-session") ?? ""
|
const sessionId = input.request.headers.get("x-opencode-session") ?? ""
|
||||||
const requestId = input.request.headers.get("x-opencode-request") ?? ""
|
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 ocClient = input.request.headers.get("x-opencode-client") ?? ""
|
||||||
const userAgent = input.request.headers.get("user-agent") ?? ""
|
const userAgent = input.request.headers.get("user-agent") ?? ""
|
||||||
logger.metric({
|
logger.metric({
|
||||||
@@ -116,16 +114,16 @@ export async function handler(
|
|||||||
})
|
})
|
||||||
const zenData = ZenData.list(opts.modelList)
|
const zenData = ZenData.list(opts.modelList)
|
||||||
const modelInfo = validateModel(zenData, model)
|
const modelInfo = validateModel(zenData, model)
|
||||||
const dataDumper = createDataDumper(sessionId, requestId, projectId)
|
|
||||||
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
|
const trialLimiter = createTrialLimiter(modelInfo.trialProvider, ip)
|
||||||
const trialProviders = await trialLimiter?.check()
|
const trialProviders = await trialLimiter?.check()
|
||||||
const rateLimiter = modelInfo.allowAnonymous
|
const rateLimiter = modelInfo.allowAnonymous
|
||||||
? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
|
? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
|
||||||
: createKeyRateLimiter(modelInfo.id, modelInfo.rateLimit, zenApiKey, input.request)
|
: createKeyRateLimiter(modelInfo.id, modelInfo.rateLimit, zenApiKey, input.request)
|
||||||
await rateLimiter?.check()
|
await rateLimiter?.check()
|
||||||
const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, sessionId)
|
|
||||||
const stickyProvider = await stickyTracker?.get()
|
|
||||||
const authInfo = await authenticate(modelInfo, zenApiKey)
|
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)
|
const billingSource = validateBilling(authInfo, modelInfo)
|
||||||
logger.metric({ source: billingSource })
|
logger.metric({ source: billingSource })
|
||||||
const modelTpmLimiter = createModelTpmLimiter(modelInfo.providers)
|
const modelTpmLimiter = createModelTpmLimiter(modelInfo.providers)
|
||||||
@@ -139,8 +137,7 @@ export async function handler(
|
|||||||
zenData,
|
zenData,
|
||||||
authInfo,
|
authInfo,
|
||||||
modelInfo,
|
modelInfo,
|
||||||
ip,
|
stickyId,
|
||||||
sessionId,
|
|
||||||
trialProviders,
|
trialProviders,
|
||||||
retry,
|
retry,
|
||||||
stickyProvider,
|
stickyProvider,
|
||||||
@@ -167,13 +164,8 @@ export async function handler(
|
|||||||
if (Array.isArray(v)) return [[k, v]]
|
if (Array.isArray(v)) return [[k, v]]
|
||||||
if (typeof v === "object") return [[k, replacer(v)]]
|
if (typeof v === "object") return [[k, replacer(v)]]
|
||||||
if (typeof v === "string") {
|
if (typeof v === "string") {
|
||||||
if (v === "$ip") return [[k, ip]]
|
|
||||||
if (v === "$workspace") return authInfo?.workspaceID ? [[k, authInfo?.workspaceID]] : []
|
if (v === "$workspace") return authInfo?.workspaceID ? [[k, authInfo?.workspaceID]] : []
|
||||||
if (v === "$session") return sessionId ? [[k, sessionId]] : []
|
if (v === "$user") return stickyId ? [[k, stickyId]] : []
|
||||||
if (v === "$user") {
|
|
||||||
const user = sessionId ?? authInfo?.workspaceID ?? ip
|
|
||||||
return user ? [[k, user]] : []
|
|
||||||
}
|
|
||||||
if (v.startsWith("$header.")) {
|
if (v.startsWith("$header.")) {
|
||||||
const headerValue = input.request.headers.get(v.slice(8))
|
const headerValue = input.request.headers.get(v.slice(8))
|
||||||
return headerValue ? [[k, headerValue]] : []
|
return headerValue ? [[k, headerValue]] : []
|
||||||
@@ -192,7 +184,7 @@ export async function handler(
|
|||||||
method: "POST",
|
method: "POST",
|
||||||
headers: (() => {
|
headers: (() => {
|
||||||
const headers = new Headers(input.request.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]) => {
|
Object.entries(providerInfo.headerMappings ?? {}).forEach(([k, v]) => {
|
||||||
headers.set(k, headers.get(v)!)
|
headers.set(k, headers.get(v)!)
|
||||||
})
|
})
|
||||||
@@ -237,10 +229,6 @@ export async function handler(
|
|||||||
|
|
||||||
const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest()
|
const { providerInfo, reqBody, res, startTimestamp } = await retriableRequest()
|
||||||
|
|
||||||
// Store model request
|
|
||||||
dataDumper?.provideModel(providerInfo.storeModel)
|
|
||||||
dataDumper?.provideRequest(reqBody)
|
|
||||||
|
|
||||||
// Store sticky provider
|
// Store sticky provider
|
||||||
if (res.status === 200) await stickyTracker?.set(providerInfo.id)
|
if (res.status === 200) await stickyTracker?.set(providerInfo.id)
|
||||||
|
|
||||||
@@ -281,8 +269,6 @@ export async function handler(
|
|||||||
const body = JSON.stringify(responseConverter(json))
|
const body = JSON.stringify(responseConverter(json))
|
||||||
logger.metric({ response_length: body.length })
|
logger.metric({ response_length: body.length })
|
||||||
logger.debug("RESPONSE: " + body)
|
logger.debug("RESPONSE: " + body)
|
||||||
dataDumper?.provideResponse(body)
|
|
||||||
dataDumper?.flush()
|
|
||||||
return new Response(body, {
|
return new Response(body, {
|
||||||
status: resStatus,
|
status: resStatus,
|
||||||
statusText: res.statusText,
|
statusText: res.statusText,
|
||||||
@@ -313,7 +299,6 @@ export async function handler(
|
|||||||
response_length: responseLength,
|
response_length: responseLength,
|
||||||
"timestamp.last_byte": timestampLastByte,
|
"timestamp.last_byte": timestampLastByte,
|
||||||
})
|
})
|
||||||
dataDumper?.flush()
|
|
||||||
await rateLimiter?.track()
|
await rateLimiter?.track()
|
||||||
const usage = usageParser.retrieve()
|
const usage = usageParser.retrieve()
|
||||||
if (usage) {
|
if (usage) {
|
||||||
@@ -351,7 +336,6 @@ export async function handler(
|
|||||||
|
|
||||||
responseLength += value.length
|
responseLength += value.length
|
||||||
buffer += decoder.decode(value, { stream: true })
|
buffer += decoder.decode(value, { stream: true })
|
||||||
dataDumper?.provideStream(buffer)
|
|
||||||
|
|
||||||
const parts = buffer.split(providerInfo.streamSeparator)
|
const parts = buffer.split(providerInfo.streamSeparator)
|
||||||
buffer = parts.pop() ?? ""
|
buffer = parts.pop() ?? ""
|
||||||
@@ -490,8 +474,7 @@ export async function handler(
|
|||||||
zenData: ZenData,
|
zenData: ZenData,
|
||||||
authInfo: AuthInfo,
|
authInfo: AuthInfo,
|
||||||
modelInfo: ModelInfo,
|
modelInfo: ModelInfo,
|
||||||
ip: string,
|
stickyId: string,
|
||||||
sessionId: string,
|
|
||||||
trialProviders: string[] | undefined,
|
trialProviders: string[] | undefined,
|
||||||
retry: RetryOptions,
|
retry: RetryOptions,
|
||||||
stickyProvider: string | undefined,
|
stickyProvider: string | undefined,
|
||||||
@@ -541,11 +524,10 @@ export async function handler(
|
|||||||
.flatMap((provider) => Array<typeof provider>(provider.weight).fill(provider))
|
.flatMap((provider) => Array<typeof provider>(provider.weight).fill(provider))
|
||||||
|
|
||||||
// Use the last 4 characters of session ID to select a provider
|
// Use the last 4 characters of session ID to select a provider
|
||||||
const identifier = sessionId.length ? sessionId : ip
|
|
||||||
let h = 0
|
let h = 0
|
||||||
const l = identifier.length
|
const l = stickyId.length
|
||||||
for (let i = l - 4; i < l; i++) {
|
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 index = (h >>> 0) % providers.length // make unsigned + range 0..length-1
|
||||||
const provider = providers[index || 0]
|
const provider = providers[index || 0]
|
||||||
|
|||||||
@@ -28,7 +28,7 @@ export const anthropicHelper: ProviderHelper = ({ reqModel, providerModel }) =>
|
|||||||
isBedrock
|
isBedrock
|
||||||
? `${providerApi}/model/${isBedrockModelArn ? encodeURIComponent(providerModel) : providerModel}/${isStream ? "invoke-with-response-stream" : "invoke"}`
|
? `${providerApi}/model/${isBedrockModelArn ? encodeURIComponent(providerModel) : providerModel}/${isStream ? "invoke-with-response-stream" : "invoke"}`
|
||||||
: providerApi + "/messages",
|
: providerApi + "/messages",
|
||||||
modifyHeaders: (headers: Headers, body: Record<string, any>, apiKey: string) => {
|
modifyHeaders: (headers: Headers, apiKey: string, _stickyId: string) => {
|
||||||
if (isBedrock || isDatabricks) {
|
if (isBedrock || isDatabricks) {
|
||||||
headers.set("Authorization", `Bearer ${apiKey}`)
|
headers.set("Authorization", `Bearer ${apiKey}`)
|
||||||
} else {
|
} else {
|
||||||
|
|||||||
@@ -30,7 +30,7 @@ export const googleHelper: ProviderHelper = ({ providerModel }) => ({
|
|||||||
format: "google",
|
format: "google",
|
||||||
modifyUrl: (providerApi: string, isStream?: boolean) =>
|
modifyUrl: (providerApi: string, isStream?: boolean) =>
|
||||||
`${providerApi}/models/${providerModel}:${isStream ? "streamGenerateContent?alt=sse" : "generateContent"}`,
|
`${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)
|
headers.set("x-goog-api-key", apiKey)
|
||||||
},
|
},
|
||||||
modifyBody: (body: Record<string, any>) => {
|
modifyBody: (body: Record<string, any>) => {
|
||||||
|
|||||||
@@ -26,9 +26,9 @@ type Usage = {
|
|||||||
export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
|
export const oaCompatHelper: ProviderHelper = ({ adjustCacheUsage }) => ({
|
||||||
format: "oa-compat",
|
format: "oa-compat",
|
||||||
modifyUrl: (providerApi: string) => providerApi + "/chat/completions",
|
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("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) => {
|
modifyBody: (body: Record<string, any>, _workspaceID?: string) => {
|
||||||
return {
|
return {
|
||||||
|
|||||||
@@ -15,7 +15,7 @@ type Usage = {
|
|||||||
export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
|
export const openaiHelper: ProviderHelper = ({ workspaceID }) => ({
|
||||||
format: "openai",
|
format: "openai",
|
||||||
modifyUrl: (providerApi: string) => providerApi + "/responses",
|
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}`)
|
headers.set("authorization", `Bearer ${apiKey}`)
|
||||||
},
|
},
|
||||||
modifyBody: (body: Record<string, any>) => body,
|
modifyBody: (body: Record<string, any>) => body,
|
||||||
|
|||||||
@@ -41,7 +41,7 @@ export type ProviderHelper = (input: {
|
|||||||
}) => {
|
}) => {
|
||||||
format: ZenData.Format
|
format: ZenData.Format
|
||||||
modifyUrl: (providerApi: string, isStream?: boolean) => string
|
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>
|
modifyBody: (body: Record<string, any>) => Record<string, any>
|
||||||
createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined
|
createBinaryStreamDecoder: () => ((chunk: Uint8Array) => Uint8Array | undefined) | undefined
|
||||||
streamSeparator: string
|
streamSeparator: string
|
||||||
|
|||||||
@@ -1,10 +1,14 @@
|
|||||||
import { Database, eq } from "@opencode-ai/console-core/drizzle/index.js"
|
import { Database, eq } from "@opencode-ai/console-core/drizzle/index.js"
|
||||||
import { ModelStickyProviderTable } from "@opencode-ai/console-core/schema/ip.sql.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 (!stickyProvider) return
|
||||||
if (!session) return
|
if (!stickyId) return
|
||||||
const id = `${modelId}/${session}`
|
const id = `${modelId}/${stickyId}`
|
||||||
let _providerId: string | undefined
|
let _providerId: string | undefined
|
||||||
|
|
||||||
return {
|
return {
|
||||||
|
|||||||
Reference in New Issue
Block a user