mirror of
https://github.com/anomalyco/opencode.git
synced 2026-05-15 09:02:35 +00:00
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)
|
||||
: 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
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
}),
|
||||
)
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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] })],
|
||||
)
|
||||
|
||||
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