zen: update sticky session logic

Author: fwang

packages/console/app/src/routes/zen/util/handler.ts

@@ -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

packages/console/app/src/routes/zen/util/stickyProviderTracker.ts

@@ -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,
+            },
+          }),
+      )
     },
   }
 }

packages/console/core/migrations/20260513163955_tearful_whistler/migration.sql

@@ -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
+);