zen: update sticky session logic

This commit is contained in:
Frank
2026-05-13 12:45:08 -04:00
parent 8ad3a4b217
commit fa077b92b1
5 changed files with 2871 additions and 7 deletions

View File

@@ -123,7 +123,7 @@ export async function handler(
? createIpRateLimiter(modelInfo.id, modelInfo.rateLimit, ip, input.request)
: createKeyRateLimiter(modelInfo.id, modelInfo.rateLimit, zenApiKey, input.request)
await rateLimiter?.check()
const stickyTracker = createStickyTracker(modelInfo.stickyProvider, sessionId)
const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, sessionId)
const stickyProvider = await stickyTracker?.get()
const authInfo = await authenticate(modelInfo, zenApiKey)
const billingSource = validateBilling(authInfo, modelInfo)
@@ -238,7 +238,7 @@ export async function handler(
dataDumper?.provideRequest(reqBody)
// Store sticky provider
await stickyTracker?.set(providerInfo.id)
if (res.status === 200) await stickyTracker?.set(providerInfo.id)
// Temporarily change 404 to 400 status code b/c solid start automatically override 404 response
const resStatus = res.status === 404 ? 400 : res.status

View File

@@ -1,16 +1,42 @@
import { Resource } from "@opencode-ai/console-resource"
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(stickyProvider: "strict" | "prefer" | undefined, session: string) {
export function createStickyTracker(modelId: string, stickyProvider: "strict" | "prefer" | undefined, session: string) {
if (!stickyProvider) return
if (!session) return
const key = `sticky:${session}`
const id = `${modelId}/${session}`
let _providerId: string | undefined
return {
get: async () => {
return await Resource.GatewayKv.get(key)
const data = await Database.use((tx) =>
tx
.select({
providerId: ModelStickyProviderTable.providerId,
})
.from(ModelStickyProviderTable)
.where(eq(ModelStickyProviderTable.id, id))
.limit(1),
)
_providerId = data[0]?.providerId
return _providerId
},
set: async (providerId: string) => {
await Resource.GatewayKv.put(key, providerId, { expirationTtl: 86400 })
if (_providerId === providerId) return
await Database.use((tx) =>
tx
.insert(ModelStickyProviderTable)
.values({
id,
providerId,
})
.onDuplicateKeyUpdate({
set: {
providerId,
},
}),
)
},
}
}

View File

@@ -0,0 +1,7 @@
CREATE TABLE `model_sticky_provider` (
`id` varchar(255) PRIMARY KEY,
`time_created` timestamp(3) NOT NULL DEFAULT (now()),
`time_updated` timestamp(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
`time_deleted` timestamp(3),
`provider_id` varchar(255) NOT NULL
);

File diff suppressed because it is too large Load Diff

View File

@@ -51,3 +51,13 @@ export const ModelTpsRateLimitTable = mysqlTable(
},
(table) => [primaryKey({ columns: [table.id, table.interval] })],
)
export const ModelStickyProviderTable = mysqlTable(
"model_sticky_provider",
{
id: varchar("id", { length: 255 }).notNull(),
...timestamps,
providerId: varchar("provider_id", { length: 255 }).notNull(),
},
(table) => [primaryKey({ columns: [table.id] })],
)