zen: sticky session

1 files changed, 9 insertions, 5 deletions

diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts
index 2e4993787..330bcc1cf 100644
--- a/packages/console/app/src/routes/zen/util/handler.ts
+++ b/packages/console/app/src/routes/zen/util/handler.ts
@@ -66,7 +66,7 @@ export async function handler(
     await rateLimiter?.check()
 
     const retriableRequest = async (retry: RetryOptions = { excludeProviders: [], retryCount: 0 }) => {
-      const providerInfo = selectProvider(zenData, modelInfo, ip, retry)
+      const providerInfo = selectProvider(zenData, modelInfo, sessionId, retry)
       const authInfo = await authenticate(modelInfo, providerInfo)
       validateBilling(authInfo, modelInfo)
       validateModelSettings(authInfo)
@@ -275,7 +275,7 @@ export async function handler(
     return { id: modelId, ...modelData }
   }
 
-  function selectProvider(zenData: ZenData, modelInfo: ModelInfo, ip: string, retry: RetryOptions) {
+  function selectProvider(zenData: ZenData, modelInfo: ModelInfo, sessionId: string, retry: RetryOptions) {
     const provider = (() => {
       if (retry.retryCount === MAX_RETRIES) {
         return modelInfo.providers.find((provider) => provider.id === modelInfo.fallbackProvider)
@@ -286,9 +286,13 @@ export async function handler(
         .filter((provider) => !retry.excludeProviders.includes(provider.id))
         .flatMap((provider) => Array<typeof provider>(provider.weight ?? 1).fill(provider))
 
-      // Use the last 2 characters of IP address to select a provider
-      const lastChars = ip.slice(-2)
-      const index = parseInt(lastChars, 16) % providers.length
+      // Use the last 4 characters of session ID to select a provider
+      let h = 0
+      const l = sessionId.length
+      for (let i = l - 4; i < l; i++) {
+        h = (h * 31 + sessionId.charCodeAt(i)) | 0 // 32-bit int
+      }
+      const index = (h >>> 0) % providers.length // make unsigned + range 0..length-1
       return providers[index || 0]
     })()