diff --git a/packages/core/src/db.ts b/packages/core/src/db.ts index e65df7e..68ff831 100644 --- a/packages/core/src/db.ts +++ b/packages/core/src/db.ts @@ -1329,6 +1329,38 @@ export function loadSessionTracking(sessionID: string): LoadedSessionTracking | }; } +/** + * Load all persisted header → session ID mappings from the session_state table. + * + * Used on gateway startup (in initIfNeeded) to pre-populate the in-memory + * headerSessionIndex so Tier 1 session identification works immediately + * after a process restart — without this, the first post-restart request + * with a known session header would generate a new session ID and orphan + * the old session's persisted state. + */ +export function loadHeaderSessionIndex(): Array<{ + sessionId: string; + headerSessionId: string; + headerName: string; +}> { + const rows = db() + .query( + `SELECT session_id, header_session_id, header_name + FROM session_state + WHERE header_session_id IS NOT NULL AND header_name IS NOT NULL`, + ) + .all() as Array<{ + session_id: string; + header_session_id: string; + header_name: string; + }>; + return rows.map((row) => ({ + sessionId: row.session_id, + headerSessionId: row.header_session_id, + headerName: row.header_name, + })); +} + // --------------------------------------------------------------------------- // Key-value store (kv_meta table) // --------------------------------------------------------------------------- diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 7107636..17bbf5b 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -75,6 +75,7 @@ export { type SessionCostSnapshot, saveSessionTracking, loadSessionTracking, + loadHeaderSessionIndex, type SessionTrackingState, type LoadedSessionTracking, getKV, diff --git a/packages/core/test/db.test.ts b/packages/core/test/db.test.ts index 6dee43f..615b30c 100644 --- a/packages/core/test/db.test.ts +++ b/packages/core/test/db.test.ts @@ -1,5 +1,5 @@ import { describe, test, expect } from "bun:test"; -import { db, close, ensureProject, projectId, mergeProjectInternal, loadForceMinLayer, saveForceMinLayer, getMeta, setMeta, getInstanceId, saveSessionCosts, loadSessionCosts, loadAllSessionCosts, getLastImportAt, setLastImportAt, saveSessionTracking, loadSessionTracking, getKV, setKV } from "../src/db"; +import { db, close, ensureProject, projectId, mergeProjectInternal, loadForceMinLayer, saveForceMinLayer, getMeta, setMeta, getInstanceId, saveSessionCosts, loadSessionCosts, loadAllSessionCosts, getLastImportAt, setLastImportAt, saveSessionTracking, loadSessionTracking, loadHeaderSessionIndex, getKV, setKV } from "../src/db"; describe("db", () => { @@ -634,6 +634,39 @@ describe("db", () => { expect(loaded!.lastBustAt).toBe(0); }); + // ------------------------------------------------------------------------- + // loadHeaderSessionIndex + // ------------------------------------------------------------------------- + + test("loadHeaderSessionIndex only returns sessions with non-null headers", () => { + const sid1 = `test-hsi-1-${crypto.randomUUID()}`; + const sid2 = `test-hsi-2-${crypto.randomUUID()}`; + saveSessionTracking(sid1, { + headerSessionId: "uuid-aaa", + headerName: "x-claude-code-session-id", + }); + saveSessionTracking(sid2, { + headerSessionId: "uuid-bbb", + headerName: "x-session-affinity", + }); + // Session without headers should NOT appear + const sid3 = `test-hsi-3-${crypto.randomUUID()}`; + saveSessionTracking(sid3, { messageCount: 5 }); + + const entries = loadHeaderSessionIndex(); + const found1 = entries.find((e) => e.sessionId === sid1); + const found2 = entries.find((e) => e.sessionId === sid2); + const found3 = entries.find((e) => e.sessionId === sid3); + + expect(found1).toBeDefined(); + expect(found1!.headerSessionId).toBe("uuid-aaa"); + expect(found1!.headerName).toBe("x-claude-code-session-id"); + expect(found2).toBeDefined(); + expect(found2!.headerSessionId).toBe("uuid-bbb"); + expect(found2!.headerName).toBe("x-session-affinity"); + expect(found3).toBeUndefined(); + }); + // ------------------------------------------------------------------------- // KV helpers (kv_meta table) // ------------------------------------------------------------------------- diff --git a/packages/gateway/src/pipeline.ts b/packages/gateway/src/pipeline.ts index e21da43..b8c810b 100644 --- a/packages/gateway/src/pipeline.ts +++ b/packages/gateway/src/pipeline.ts @@ -47,6 +47,7 @@ import { embedding, saveSessionTracking, loadSessionTracking, + loadHeaderSessionIndex, } from "@loreai/core"; import type { @@ -195,6 +196,7 @@ export function setUpstreamInterceptor( export async function resetPipelineState(): Promise { initialized = false; sessions.clear(); + headerSessionIndex.clear(); ltmSessionCache.clear(); ltmPinnedText.clear(); // Shut down batch queue gracefully before clearing the client @@ -546,6 +548,23 @@ async function initIfNeeded(projectPath: string, config?: GatewayConfig): Promis log.error("lat-reader startup refresh error:", e); } + // Pre-populate headerSessionIndex from DB so Tier 1 session identification + // works immediately after process restart. Without this, the first request + // with a known session header generates a new session ID and orphans the + // old session's persisted state. + try { + const headerEntries = loadHeaderSessionIndex(); + for (const entry of headerEntries) { + const indexKey = `${entry.headerName}:${entry.headerSessionId}`; + headerSessionIndex.set(indexKey, entry.sessionId); + } + if (headerEntries.length > 0) { + log.info(`restored ${headerEntries.length} header→session mappings from DB`); + } + } catch (e) { + log.warn("header session index restore failed:", e); + } + // Pre-warm models.dev pricing/limits cache so synchronous lookups in the // request hot path (getModelSpec, emitCostMetric) resolve from memory. fetchModelData().catch((e) => log.warn("models.dev pre-warm failed:", e));