mirror of
https://github.com/anomalyco/opencode.git
synced 2026-05-22 11:25:15 +00:00
zen: fix affiliated header
This commit is contained in:
Frank
parent
8c72edb3ff
commit
6de584f25c
8 changed files with 22 additions and 81 deletions
|
|
@ -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 { 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]
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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>) => {
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue