diff --git a/src/lib/db/sessionAccountAffinity.ts b/src/lib/db/sessionAccountAffinity.ts new file mode 100644 index 0000000000..80ba0a729d --- /dev/null +++ b/src/lib/db/sessionAccountAffinity.ts @@ -0,0 +1,31 @@ +// Stubbed functions for session account affinity (PR 1887 pending) + +export function getSessionAccountAffinity(sessionKey: string, provider: string): any { + return null; +} + +export function upsertSessionAccountAffinity( + sessionKey: string, + provider: string, + connectionId: string, + now: number = Date.now() +): void {} + +export function touchSessionAccountAffinity( + sessionKey: string, + provider: string, + now: number = Date.now() +): void {} + +export function deleteSessionAccountAffinity(sessionKey: string, provider: string): void {} + +export function cleanupStaleSessionAccountAffinities( + ttlMs: number = 30 * 60 * 1000, + now: number = Date.now() +): number { + return 0; +} + +export function startSessionAccountAffinityCleanup(): void {} + +export function stopSessionAccountAffinityCleanupForTests(): void {} diff --git a/src/sse/services/auth.ts b/src/sse/services/auth.ts index 84d3fc5c86..863204ed92 100644 --- a/src/sse/services/auth.ts +++ b/src/sse/services/auth.ts @@ -1,4 +1,4 @@ -import { randomUUID } from "crypto"; +import { randomUUID, createHash } from "crypto"; import { getProviderConnections, validateApiKey, @@ -78,6 +78,7 @@ interface CredentialSelectionOptions { bypassQuotaPolicy?: boolean; forcedConnectionId?: string | null; excludeConnectionIds?: string[] | null; + sessionKey?: string | null; } interface CooldownInspectionState { @@ -145,6 +146,94 @@ function toBooleanOrDefault(value: unknown, fallback: boolean): boolean { return typeof value === "boolean" ? value : fallback; } +function readHeaderValue( + headers: Headers | { get?: (name: string) => string | null } | null | undefined, + name: string +): string | null { + if (!headers || typeof headers.get !== "function") return null; + const value = headers.get(name); + return typeof value === "string" && value.trim().length > 0 ? value.trim() : null; +} + +function normalizeSessionKey(value: unknown, prefix: string): string | null { + if (typeof value !== "string" || value.trim().length === 0) return null; + const trimmed = value.trim(); + if (trimmed.length <= 180 && /^[A-Za-z0-9._:-]+$/.test(trimmed)) { + return `${prefix}:${trimmed}`; + } + return `${prefix}:sha256:${createHash("sha256").update(trimmed).digest("hex")}`; +} + +function extractTextForSessionHash(value: unknown): string | null { + if (typeof value === "string") return value; + if (Array.isArray(value)) { + const parts = value + .map((item) => { + if (typeof item === "string") return item; + const record = asRecord(item); + if (typeof record.text === "string") return record.text; + if (typeof record.content === "string") return record.content; + return null; + }) + .filter(Boolean) as string[]; + return parts.length > 0 ? parts.join("\n") : JSON.stringify(value); + } + if (value && typeof value === "object") return JSON.stringify(value); + return null; +} + +function getFirstInputText(body: unknown): string | null { + const record = asRecord(body); + if (record.input !== undefined) { + if (typeof record.input === "string") return record.input; + if (Array.isArray(record.input)) { + for (const item of record.input) { + const itemRecord = asRecord(item); + const text = extractTextForSessionHash(itemRecord.content ?? item); + if (text && text.trim().length > 0) return text; + } + } + const text = extractTextForSessionHash(record.input); + if (text && text.trim().length > 0) return text; + } + + if (Array.isArray(record.messages)) { + const userMessage = record.messages.find((message) => asRecord(message).role === "user"); + const firstMessage = userMessage ?? record.messages[0]; + const text = extractTextForSessionHash(asRecord(firstMessage).content ?? firstMessage); + if (text && text.trim().length > 0) return text; + } + + return null; +} + +export function extractSessionAffinityKey( + body: unknown, + headers?: Headers | { get?: (name: string) => string | null } | null +): string | null { + const headerKey = normalizeSessionKey( + readHeaderValue(headers, "x-codex-session-id") ?? + readHeaderValue(headers, "x-session-id") ?? + readHeaderValue(headers, "x-omniroute-session"), + "header" + ); + if (headerKey) return headerKey; + + const record = asRecord(body); + const metadata = asRecord(record.metadata); + const explicitKey = + normalizeSessionKey(metadata.session_id, "metadata") ?? + normalizeSessionKey(metadata.sessionId, "metadata") ?? + normalizeSessionKey(record.conversation_id, "conversation") ?? + normalizeSessionKey(record.session_id, "session") ?? + normalizeSessionKey(record.prompt_cache_key, "prompt-cache"); + if (explicitKey) return explicitKey; + + const inputText = getFirstInputText(body); + if (!inputText || inputText.trim().length === 0) return null; + return `input:sha256:${createHash("sha256").update(inputText.slice(0, 4096)).digest("hex")}`; +} + function getCodexLimitPolicy(providerSpecificData: JsonRecord): { use5h: boolean; useWeekly: boolean;