zen: update sticky session logic
This commit is contained in:
@@ -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] })],
|
||||||
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user