[AISAGE-735] feat(ai-gateway): send explicit session ID and affinity to Vercel

packages/web-shared/src/lib/ai-gateway/providerHash.test.ts

@@ -79,6 +79,30 @@ describe('session tracking identifiers', () => {
     expect(trackedSession).toBeUndefined();
   });
 
+  it.each(requests)('uses the shared session hash for Vercel $kind tracking', request => {
+    const trackedRequest = structuredClone(request);
+    applyTrackingIds(trackedRequest, VERCEL_AI_GATEWAY, userId, sessionId);
+
+    const trackedSession =
+      trackedRequest.kind === 'messages'
+        ? trackedRequest.body.session_id
+        : trackedRequest.body.prompt_cache_key;
+    expect(trackedSession).toBe(
+      generateProviderSpecificSessionHash(userId, sessionId, VERCEL_AI_GATEWAY)
+    );
+  });
+
+  it('sets the Vercel user hash on Vercel messages requests', () => {
+    const request: GatewayRequest = {
+      kind: 'messages',
+      body: { model: 'test-model', messages: [], max_tokens: 100 },
+    };
+    applyTrackingIds(request, VERCEL_AI_GATEWAY, userId, null);
+
+    expect(request.body).toHaveProperty('user', generateVercelDownstreamSafetyIdentifier(userId));
+    expect(request.body).not.toHaveProperty('session_id');
+  });
+
   const promptCacheKeyRequests = (promptCacheKey: string): GatewayRequest[] => [
     {
       kind: 'chat_completions',

packages/web-shared/src/lib/ai-gateway/providerHash.ts

@@ -69,7 +69,7 @@ export function applyTrackingIds(
   const taskHash = taskId ? generateProviderSpecificSessionHash(userId, taskId, provider) : '';
   if (request.kind === 'messages') {
     request.body.metadata = { ...request.body.metadata, user_id: userHash };
-    if (provider.id === 'openrouter') {
+    if (provider.id === 'openrouter' || provider.id === 'vercel') {
       request.body.user = userHash;
       if (taskHash) {
         request.body.session_id = taskHash;

packages/web-shared/src/lib/ai-gateway/providers/openrouter/types.ts

@@ -35,6 +35,9 @@ export type VercelInferenceProviderConfig =
 export type VercelProviderConfig = {
   gateway?: GatewayProviderOptions & {
     byok?: Record<string, VercelInferenceProviderConfig[]>;
+    // https://vercel.com/docs/ai-gateway/models-and-providers/session-affinity
+    sessionId?: string;
+    sessionAffinity?: boolean;
   };
   anthropic?: AnthropicProviderOptions;
 };

packages/web-shared/src/lib/ai-gateway/providers/vercel/index.test.ts

@@ -434,6 +434,7 @@ describe('applyVercelSettings BYOK pinning', () => {
         anthropic: [{ apiKey: 'sk-anthropic' }],
         bedrock: [{ accessKeyId: 'AKIAEXAMPLE', secretAccessKey: 'secret', region: 'us-east-1' }],
       },
+      sessionAffinity: true,
     });
     expect(request.body.provider).toBeUndefined();
   });
@@ -706,6 +707,96 @@ describe('applyVercelSettings managed requests', () => {
   });
 });
 
+describe('applyVercelSettings session affinity', () => {
+  const conversationId = 'conversation-hash';
+
+  function requestWithConversationId(
+    kind: GatewayRequest['kind'],
+    id: string | undefined
+  ): GatewayRequest {
+    const model = 'anthropic/claude-sonnet-4.5';
+    const provider = { order: ['google-vertex', 'amazon-bedrock'] };
+    if (kind === 'responses') {
+      return { kind, body: { model, input: 'hello', provider, prompt_cache_key: id } };
+    }
+    if (kind === 'messages') {
+      return {
+        kind,
+        body: {
+          model,
+          max_tokens: 128,
+          messages: [{ role: 'user', content: 'hello' }],
+          provider,
+          session_id: id,
+        },
+      };
+    }
+    return {
+      kind,
+      body: {
+        model,
+        messages: [{ role: 'user', content: 'hello' }],
+        provider,
+        prompt_cache_key: id,
+      },
+    };
+  }
+
+  async function loadApplyVercelSettings() {
+    jest.resetModules();
+    jest.doMock('@kilocode/web-shared/lib/ai-gateway/providers/gateway-models-cache', () => ({
+      getCachedVercelInferenceProviderIdsForModel: jest.fn(async () => ['vertex', 'bedrock']),
+      getVercelModelsFromDatabase: jest.fn(),
+    }));
+    return (await import('@kilocode/web-shared/lib/ai-gateway/providers/vercel'))
+      .applyVercelSettings;
+  }
+
+  const kinds = ['chat_completions', 'messages', 'responses'] as const;
+
+  it.each(kinds)('sends the conversation ID for managed %s requests', async kind => {
+    const applySettings = await loadApplyVercelSettings();
+    const request = requestWithConversationId(kind, conversationId);
+
+    await applySettings('anthropic/claude-sonnet-4.5', request, null);
+
+    expect(request.body.providerOptions?.gateway).toMatchObject({
+      sessionId: conversationId,
+      sessionAffinity: true,
+      order: ['vertex', 'bedrock'],
+    });
+  });
+
+  it.each(kinds)('sends the conversation ID for BYOK %s requests', async kind => {
+    const request = requestWithConversationId(kind, conversationId);
+
+    await applyVercelSettings('anthropic/claude-sonnet-4.5', request, [
+      { decryptedAPIKey: 'sk-anthropic', providerId: 'anthropic' },
+    ]);
+
+    expect(request.body.providerOptions?.gateway).toMatchObject({
+      sessionId: conversationId,
+      sessionAffinity: true,
+    });
+  });
+
+  it.each(kinds)('leaves sessionId unset without a conversation ID for %s', async kind => {
+    const applySettings = await loadApplyVercelSettings();
+    const managed = requestWithConversationId(kind, undefined);
+    const byok = requestWithConversationId(kind, '');
+
+    await applySettings('anthropic/claude-sonnet-4.5', managed, null);
+    await applySettings('anthropic/claude-sonnet-4.5', byok, [
+      { decryptedAPIKey: 'sk-anthropic', providerId: 'anthropic' },
+    ]);
+
+    for (const request of [managed, byok]) {
+      expect(request.body.providerOptions?.gateway?.sessionAffinity).toBe(true);
+      expect(request.body.providerOptions?.gateway).not.toHaveProperty('sessionId');
+    }
+  });
+});
+
 describe('isVercelRoutingOptOut', () => {
   it('only opts out exact model ID matches', () => {
     const optOutModels = new Set(['moonshotai/kimi-k3']);

packages/web-shared/src/lib/ai-gateway/providers/vercel/index.ts

@@ -121,6 +121,21 @@ export async function shouldRouteToVercel(
   return true;
 }
 
+function getConversationId(request: GatewayRequest): string | undefined {
+  const id = request.kind === 'messages' ? request.body.session_id : request.body.prompt_cache_key;
+  return typeof id === 'string' && id.length > 0 ? id : undefined;
+}
+
+// An explicit session ID makes Vercel pin the exact credential that served the
+// session, keeping its prompt cache warm after a fallback.
+// https://vercel.com/docs/ai-gateway/models-and-providers/session-affinity
+function getSessionAffinityOptions(
+  request: GatewayRequest
+): Pick<NonNullable<VercelProviderConfig['gateway']>, 'sessionId' | 'sessionAffinity'> {
+  const sessionId = getConversationId(request);
+  return sessionId ? { sessionId, sessionAffinity: true } : { sessionAffinity: true };
+}
+
 function convertProviderSort(
   sort: OpenRouterProviderConfig['sort']
 ): GatewayProviderOptions['sort'] {
@@ -172,6 +187,7 @@ export function convertProviderOptions(
       disallowPromptTraining: provider?.data_collection === 'deny' || undefined,
       models: requestToMutate.body.models,
       serviceTier: serviceTier === 'flex' || serviceTier === 'priority' ? serviceTier : undefined,
+      ...getSessionAffinityOptions(requestToMutate),
     },
   };
 }
@@ -285,6 +301,7 @@ export async function applyVercelSettings(
         sort: convertProviderSort(requestToMutate.body.provider?.sort),
         byok: byokProviders,
         models: requestToMutate.body.models,
+        ...getSessionAffinityOptions(requestToMutate),
       },
     };
   } else {