fix: pre-populate headerSessionIndex from DB on startup

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)
 // ---------------------------------------------------------------------------

packages/core/src/index.ts

@@ -75,6 +75,7 @@ export {
   type SessionCostSnapshot,
   saveSessionTracking,
   loadSessionTracking,
+  loadHeaderSessionIndex,
   type SessionTrackingState,
   type LoadedSessionTracking,
   getKV,

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)
   // -------------------------------------------------------------------------

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<void> {
   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));