zen: update sticky session logic

This commit is contained in:
Frank
2026-05-13 12:45:11 -04:00
parent 8ad3a4b217
commit fa077b92b1
5 changed files with 2871 additions and 7 deletions
@@ -123,7 +123,7 @@ export async function handler(
? 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.stickyProvider, sessionId) const stickyTracker = createStickyTracker(modelInfo.id, modelInfo.stickyProvider, sessionId)
const stickyProvider = await stickyTracker?.get() const stickyProvider = await stickyTracker?.get()
const authInfo = await authenticate(modelInfo, zenApiKey) const authInfo = await authenticate(modelInfo, zenApiKey)
const billingSource = validateBilling(authInfo, modelInfo) const billingSource = validateBilling(authInfo, modelInfo)
@@ -238,7 +238,7 @@ export async function handler(
dataDumper?.provideRequest(reqBody) dataDumper?.provideRequest(reqBody)
// Store sticky provider // 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 // Temporarily change 404 to 400 status code b/c solid start automatically override 404 response
const resStatus = res.status === 404 ? 400 : res.status const resStatus = res.status === 404 ? 400 : res.status
@@ -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 (!stickyProvider) return
if (!session) return if (!session) return
const key = `sticky:${session}` const id = `${modelId}/${session}`
let _providerId: string | undefined
return { return {
get: async () => { 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) => { 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,
},
}),
)
}, },
} }
} }
@@ -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
@@ -51,3 +51,13 @@ export const ModelTpsRateLimitTable = mysqlTable(
}, },
(table) => [primaryKey({ columns: [table.id, table.interval] })], (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] })],
)