From 0061e9636464b36d2ec6dd178a546153898e453e Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Mon, 19 Oct 2026 10:59:38 +0200 Subject: [PATCH 1/2] feat(ai-gateway): send explicit session ID and session affinity to Vercel --- .../src/lib/ai-gateway/providerHash.test.ts | 23 +++++ .../src/lib/ai-gateway/providerHash.ts | 6 +- .../ai-gateway/providers/openrouter/types.ts | 3 + .../ai-gateway/providers/vercel/index.test.ts | 91 +++++++++++++++++++ .../lib/ai-gateway/providers/vercel/index.ts | 17 ++++ 5 files changed, 137 insertions(+), 3 deletions(-) diff --git a/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts b/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts index a20e2209479..3b94e9c937f 100644 --- a/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts @@ -79,6 +79,29 @@ 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('does not set an OpenRouter user on Vercel messages requests', () => { + const request: GatewayRequest = { + kind: 'messages', + body: { model: 'test-model', messages: [], max_tokens: 100 }, + }; + applyTrackingIds(request, VERCEL_AI_GATEWAY, userId, sessionId); + + expect(request.body).not.toHaveProperty('user'); + }); + const promptCacheKeyRequests = (promptCacheKey: string): GatewayRequest[] => [ { kind: 'chat_completions', diff --git a/packages/web-shared/src/lib/ai-gateway/providerHash.ts b/packages/web-shared/src/lib/ai-gateway/providerHash.ts index 495c6bdd232..8ac68802c84 100644 --- a/packages/web-shared/src/lib/ai-gateway/providerHash.ts +++ b/packages/web-shared/src/lib/ai-gateway/providerHash.ts @@ -71,9 +71,9 @@ export function applyTrackingIds( request.body.metadata = { ...request.body.metadata, user_id: userHash }; if (provider.id === 'openrouter') { request.body.user = userHash; - if (taskHash) { - request.body.session_id = taskHash; - } + } + if (taskHash && (provider.id === 'openrouter' || provider.id === 'vercel')) { + request.body.session_id = taskHash; } } else { const clientCacheKey = request.body.prompt_cache_key; diff --git a/packages/web-shared/src/lib/ai-gateway/providers/openrouter/types.ts b/packages/web-shared/src/lib/ai-gateway/providers/openrouter/types.ts index 4f31609c5ce..ac094f28da7 100644 --- a/packages/web-shared/src/lib/ai-gateway/providers/openrouter/types.ts +++ b/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; + // https://vercel.com/docs/ai-gateway/models-and-providers/session-affinity + sessionId?: string; + sessionAffinity?: boolean; }; anthropic?: AnthropicProviderOptions; }; diff --git a/packages/web-shared/src/lib/ai-gateway/providers/vercel/index.test.ts b/packages/web-shared/src/lib/ai-gateway/providers/vercel/index.test.ts index 020f8191094..36e354d4e90 100644 --- a/packages/web-shared/src/lib/ai-gateway/providers/vercel/index.test.ts +++ b/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']); diff --git a/packages/web-shared/src/lib/ai-gateway/providers/vercel/index.ts b/packages/web-shared/src/lib/ai-gateway/providers/vercel/index.ts index 178956d3c66..3ef65ccd6f5 100644 --- a/packages/web-shared/src/lib/ai-gateway/providers/vercel/index.ts +++ b/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, '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 { From f179bca53e941b02da188258b8c99d5df4ccb50c Mon Sep 17 00:00:00 2001 From: Christiaan Arnoldus Date: Mon, 19 Oct 2026 11:43:33 +0200 Subject: [PATCH 2/2] feat(ai-gateway): send hashed user on Vercel messages requests --- .../web-shared/src/lib/ai-gateway/providerHash.test.ts | 7 ++++--- packages/web-shared/src/lib/ai-gateway/providerHash.ts | 8 ++++---- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts b/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts index 3b94e9c937f..82bb6f8dca4 100644 --- a/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts +++ b/packages/web-shared/src/lib/ai-gateway/providerHash.test.ts @@ -92,14 +92,15 @@ describe('session tracking identifiers', () => { ); }); - it('does not set an OpenRouter user on Vercel messages requests', () => { + 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, sessionId); + applyTrackingIds(request, VERCEL_AI_GATEWAY, userId, null); - expect(request.body).not.toHaveProperty('user'); + expect(request.body).toHaveProperty('user', generateVercelDownstreamSafetyIdentifier(userId)); + expect(request.body).not.toHaveProperty('session_id'); }); const promptCacheKeyRequests = (promptCacheKey: string): GatewayRequest[] => [ diff --git a/packages/web-shared/src/lib/ai-gateway/providerHash.ts b/packages/web-shared/src/lib/ai-gateway/providerHash.ts index 8ac68802c84..0ca86e62e19 100644 --- a/packages/web-shared/src/lib/ai-gateway/providerHash.ts +++ b/packages/web-shared/src/lib/ai-gateway/providerHash.ts @@ -69,11 +69,11 @@ 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 && (provider.id === 'openrouter' || provider.id === 'vercel')) { - request.body.session_id = taskHash; + if (taskHash) { + request.body.session_id = taskHash; + } } } else { const clientCacheKey = request.body.prompt_cache_key;