@lincoln504/pi-research 1.1.2 → 1.3.0

package/src/infrastructure/embedding/embedding-client.ts

@@ -8,6 +8,7 @@
 import * as http from 'node:http';
 import { logger } from '../../logger.ts';
 import type { IEmbedder } from '../../core/interfaces/knowledge-interfaces.ts';
+import { Utf8Body } from '../../utils/http-body.ts';
 
 export class EmbeddingClient implements IEmbedder {
   // Public so store.ts can set the dimension from existing table schema
@@ -17,10 +18,21 @@ export class EmbeddingClient implements IEmbedder {
   constructor(
     private readonly port: number,
     private _device: string = 'unknown',
+    /**
+     * Shared secret published by the leader in state. Undefined when the leader
+     * predates auth (it does not check a header either), so the request simply
+     * carries none — an older leader keeps working rather than 403-ing.
+     */
+    private readonly authSecret?: string,
   ) {
     logger.log(`[EmbeddingClient] Connecting to embedding server at http://127.0.0.1:${port}`);
   }
 
+  /** Auth header for the leader, or nothing when connecting to a pre-auth leader. */
+  private authHeaders(): Record<string, string> {
+    return this.authSecret ? { 'x-embedding-auth': this.authSecret } : {};
+  }
+
   // ---- IEmbedder ----
 
   getDevice(): string {
@@ -93,25 +105,26 @@ export class EmbeddingClient implements IEmbedder {
           port: this.port,
           path,
           method: 'GET',
+          headers: this.authHeaders(),
         },
         (res) => {
           // Keep the timeout armed through the body read — clearing it on headers
           // would let a stalled body hang for the full process lifetime.
-          let body = '';
-          res.on('data', (chunk: Buffer) => { body += chunk; });
+          const collected = new Utf8Body();
+          res.on('data', (chunk: Buffer) => { collected.push(chunk); });
           res.on('end', () => {
             if (resolved) return;
             resolved = true;
             clearTimeout(timer);
             try {
-              const parsed = JSON.parse(body) as T;
+              const parsed = JSON.parse(collected.toString()) as T;
               if (res.statusCode !== 200) {
                 reject(new Error(`[EmbeddingClient] GET ${path} returned HTTP ${res.statusCode}`));
               } else {
                 resolve(parsed);
               }
             } catch (_e) {
-              reject(new Error(`[EmbeddingClient] Failed to parse GET ${path} response: ${body}`));
+              reject(new Error(`[EmbeddingClient] Failed to parse GET ${path} response: ${collected.toString()}`));
             }
           });
           res.on('error', (err) => {
@@ -160,18 +173,19 @@ export class EmbeddingClient implements IEmbedder {
           headers: {
             'Content-Type': 'application/json',
             'Content-Length': Buffer.byteLength(body),
+            ...this.authHeaders(),
           },
         },
         (res) => {
           // Keep the timeout armed through the body read (see getRequest).
-          let responseBody = '';
-          res.on('data', (chunk: Buffer) => { responseBody += chunk; });
+          const collected = new Utf8Body();
+          res.on('data', (chunk: Buffer) => { collected.push(chunk); });
           res.on('end', () => {
             if (resolved) return;
             resolved = true;
             clearTimeout(timer);
             try {
-              const parsed = JSON.parse(responseBody) as T;
+              const parsed = JSON.parse(collected.toString()) as T;
               if (res.statusCode !== 200) {
                 const msg = (parsed as Record<string, unknown>)?.['error'];
                 reject(new Error(typeof msg === 'string' ? msg : `HTTP ${res.statusCode}`));
@@ -179,7 +193,7 @@ export class EmbeddingClient implements IEmbedder {
                 resolve(parsed);
               }
             } catch (_e) {
-              reject(new Error(`[EmbeddingClient] Failed to parse POST ${path} response: ${responseBody}`));
+              reject(new Error(`[EmbeddingClient] Failed to parse POST ${path} response: ${collected.toString()}`));
             }
           });
           res.on('error', (err) => {

package/src/infrastructure/embedding/embedding-factory.ts

@@ -19,7 +19,7 @@ import type { IEmbedder } from '../../core/interfaces/knowledge-interfaces.ts';
 import { Embedder } from '../../knowledge/embedder.ts';
 import { getModelEmbedderConfig } from '../../knowledge/model-config.ts';
 
-import { EmbeddingServer } from './embedding-server.ts';
+import { EmbeddingServer, getEmbeddingServerAuthSecret } from './embedding-server.ts';
 import { EmbeddingClient } from './embedding-client.ts';
 
 // ---------------------------------------------------------------------------
@@ -76,6 +76,13 @@ export async function isEmbeddingLeaderAlive(
 
 let _embeddingInstance: IEmbedder | null = null;
 let _embeddingInitPromise: Promise<IEmbedder> | null = null;
+// Config the in-flight init was started for. The join below must compare against
+// these — the cached-instance path checks model/device before reuse, and the
+// adoption paths refuse mismatched leaders, but a caller joining a promise built
+// for a different config would silently receive (and stamp the cache with) the
+// wrong-model embedder, defeating both guards.
+let _initModel: string | null = null;
+let _initDevice: string | null = null;
 let _cachedModel: string | null = null;
 let _cachedDevice: string | null = null;
 
@@ -119,11 +126,29 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
       return _embeddingInstance;
     }
     logger.info(`[EmbeddingFactory] Configuration change detected (${_cachedModel} on ${_cachedDevice} -> ${cfg.EMBEDDING_MODEL} on ${cfg.EMBEDDING_DEVICE}). Disposing stale instance.`);
-    await _embeddingInstance.dispose?.();
+    // Claim the instance SYNCHRONOUSLY, before the first await. Two callers that
+    // resolve different EMBEDDING_MODEL/EMBEDDING_DEVICE values would otherwise both
+    // observe it, both await dispose() on the same live EmbeddingServer (double
+    // close(), double embedder dispose) and both proceed to init().
+    const stale = _embeddingInstance;
     _embeddingInstance = null;
+    await stale.dispose?.();
   }
 
-  if (_embeddingInitPromise) return _embeddingInitPromise;
+  // An internal retry (_attempt > 0) is re-entering from INSIDE the promise stored
+  // here, so handing that promise back would make it await itself. Only a fresh
+  // external call may join an in-flight init — and only one started for the SAME
+  // config: a mismatched joiner would get the wrong-model embedder (same-dimension
+  // models cross-embed silently). On mismatch, wait the in-flight init out (never
+  // race a second election in this process), then re-resolve for this config — the
+  // cached-instance check above disposes the mismatched instance.
+  if (_embeddingInitPromise && _attempt === 0) {
+    if (_initModel === cfg.EMBEDDING_MODEL && _initDevice === cfg.EMBEDDING_DEVICE) {
+      return _embeddingInitPromise;
+    }
+    await _embeddingInitPromise.catch(() => {});
+    return getEmbedder(cfg, _attempt);
+  }
 
   let p: Promise<IEmbedder>;
 
@@ -166,6 +191,9 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
         pid: process.pid,
         serverId,
         model: cfg.EMBEDDING_MODEL,
+        // Published with the claim, not with the port, so a follower that adopts
+        // this leader always has the secret available whenever the port is.
+        authSecret: getEmbeddingServerAuthSecret(),
         ...(myStartTime !== null ? { startTime: myStartTime } : {}),
       };
       iAmCandidate = true;
@@ -211,7 +239,12 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
         if (!(await isEmbeddingLeaderAlive(info, processLifecycle))) {
           // Candidate process died without cleaning up (or its PID was recycled by
           // an unrelated process — the start-time pairing catches that too).
-          await stateManager.clearEmbeddingServer().catch((err) => logger.debug('Swallowed clear embedding server error:', err));
+          // serverId-scoped: several waiters typically detect the death in the same
+          // poll tick; the first one's clear+claim can commit before this one's clear
+          // runs, and an unconditional delete here would deregister that NEW claimant
+          // (whose Phase-3 port publish is CAS-guarded and silently no-ops) — leaving
+          // two concurrent model inits and an invisible leader.
+          await stateManager.clearEmbeddingServer({ serverId: info.serverId }).catch((err) => logger.debug('Swallowed clear embedding server error:', err));
           break;
         }
         // Re-check model identity every poll: the leader may have changed while
@@ -223,7 +256,7 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
           const portOk = await isPortListening(info.port);
           if (portOk) {
             logger.info(`[EmbeddingFactory] Connecting to embedding server on port ${info.port}`);
-            const client = new EmbeddingClient(info.port);
+            const client = new EmbeddingClient(info.port, cfg.EMBEDDING_DEVICE, info.authSecret);
             await client.fetchHealth();
             _embeddingInstance = client;
             _cachedModel = cfg.EMBEDDING_MODEL;
@@ -235,14 +268,21 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
           const secondCheck = await isPortListening(info.port);
           if (!secondCheck) {
             logger.warn(`[EmbeddingFactory] Registered port ${info.port} is unreachable after two checks — clearing stale state`);
-            await stateManager.clearEmbeddingServer().catch((err) => logger.debug('Swallowed clear embedding server error:', err));
+            // serverId-scoped: this snapshot is up to ~2.5s old (two probes + sleep);
+            // the entry may already belong to a new claimant. See the dead-leader
+            // clear above for the failure this prevents.
+            await stateManager.clearEmbeddingServer({ serverId: info.serverId }).catch((err) => logger.debug('Swallowed clear embedding server error:', err));
             break;
           }
         }
       }
 
-      // Timed out or candidate died — attempt to claim candidacy ourselves (tail-call)
-      _embeddingInitPromise = null;
+      // Timed out or candidate died — attempt to claim candidacy ourselves (tail-call).
+      // The in-flight marker is deliberately LEFT in place: clearing it here, while
+      // this very promise is still running, let a concurrent getEmbedder() see no
+      // instance and no init in flight and start a second, competing election in the
+      // same process. The retry re-enters with _attempt > 0, which skips the join
+      // above, so leaving it set cannot deadlock.
       if (_attempt + 1 >= MAX_LEADER_WAIT_RETRIES) {
         throw new Error(
           `[EmbeddingFactory] Could not obtain an embedding server after ${MAX_LEADER_WAIT_RETRIES} attempts ` +
@@ -267,8 +307,10 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
     try {
       port = await server.startServer();
     } catch (err) {
-      // Init failed — clear our placeholder so others can try
-      await stateManager.clearEmbeddingServer().catch((err) => logger.debug('Swallowed clear embedding server error:', err));
+      // Init failed — clear our placeholder so others can try. serverId-scoped so a
+      // concurrent steal-and-reclaim (a waiter that judged us dead and claimed) is
+      // never deregistered by our own failure path.
+      await stateManager.clearEmbeddingServer({ serverId }).catch((err) => logger.debug('Swallowed clear embedding server error:', err));
       throw err;
     }
 
@@ -292,16 +334,20 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
 
   p = init();
   _embeddingInitPromise = p;
+  _initModel = cfg.EMBEDDING_MODEL;
+  _initDevice = cfg.EMBEDDING_DEVICE;
   // init() already assigns _embeddingInstance before it resolves, so do NOT
   // re-assign it here. If clearEmbeddingInstance() ran (and disposed the server)
   // between init resolving and this microtask, re-assigning would resurrect the
   // disposed instance and hand it to the next getEmbedder() caller. Only clear
   // the in-flight promise marker.
-  p.then(() => {
-    _embeddingInitPromise = null;
-  }).catch(() => {
-    _embeddingInitPromise = null;
-  });
+  // Clear the marker only if it is STILL ours. An unconditional clear let a settling
+  // promise wipe a newer caller's in-flight marker — admitting yet another election,
+  // and cascading. Same guard as scheduler-factory's initialization promise.
+  const clearIfCurrent = (): void => {
+    if (_embeddingInitPromise === p) _embeddingInitPromise = null;
+  };
+  p.then(clearIfCurrent).catch(clearIfCurrent);
   return p;
 }
 

package/src/infrastructure/embedding/embedding-server.ts

@@ -6,14 +6,67 @@
  */
 
 import * as http from 'node:http';
+import * as crypto from 'node:crypto';
 import * as path from 'node:path';
 import { logger } from '../../logger.ts';
 import { captureStdio } from '../../utils/stdio-capture.ts';
 import { buildDefaultDebugLogPath } from '../../utils/log-utils.ts';
 import { DiskSpaceChecker } from '../../utils/disk-space-checker.ts';
+import { Utf8Body } from '../../utils/http-body.ts';
 import type { IEmbedder } from '../../core/interfaces/knowledge-interfaces.ts';
 import type { IStateManager } from '../../core/interfaces/state-manager-interfaces.ts';
 import type { Embedder } from '../../knowledge/embedder.ts';
+import { markWebGpuFallback } from '../../knowledge/embedder-utils.ts';
+
+/**
+ * Per-process shared secret every embed request must present.
+ *
+ * The server binds 127.0.0.1, but loopback is not an authorization boundary —
+ * any local process could otherwise enqueue embedding work. That matters more
+ * than raw GPU cost: the serial queue steps the leader down when it poisons, so
+ * sustained abuse denies the knowledge store cluster-wide. The browser server has
+ * had this since #21; the embedding server was the asymmetric gap.
+ *
+ * Generated lazily and held in memory only — it is published to followers through
+ * the state file, which is already 0o600 and already carries the browser secret.
+ */
+let _embeddingAuthSecret: string | null = null;
+
+export function getEmbeddingServerAuthSecret(): string {
+  if (!_embeddingAuthSecret) {
+    _embeddingAuthSecret = crypto.randomBytes(32).toString('hex');
+  }
+  return _embeddingAuthSecret;
+}
+
+/**
+ * Constant-time comparison of a presented header against the expected secret.
+ * Length is compared first because timingSafeEqual throws on a length mismatch.
+ */
+function isAuthorized(presented: string | string[] | undefined): boolean {
+  const header = Array.isArray(presented) ? presented[0] : presented;
+  const expected = Buffer.from(getEmbeddingServerAuthSecret(), 'utf8');
+  const actual = Buffer.from(header ?? '', 'utf8');
+  return actual.length === expected.length && crypto.timingSafeEqual(actual, expected);
+}
+
+/**
+ * Refuse to serialize a vector containing NaN/±Infinity. JSON.stringify emits
+ * `null` for non-finite numbers, and the client's `new Float32Array(resp.data)`
+ * coerces null back to 0 — a silently part-zeroed vector that passes every
+ * downstream width/finite gate (0 IS finite) and lands in the ANN index. Failing
+ * the request instead surfaces it as an ordinary embed failure on the client.
+ */
+function assertFiniteVector(vec: ArrayLike<number>, vectorIndex?: number): void {
+  for (let i = 0; i < vec.length; i++) {
+    if (!Number.isFinite(vec[i]!)) {
+      const where = vectorIndex === undefined ? `position ${i}` : `vector ${vectorIndex}, position ${i}`;
+      throw new Error(
+        `[EmbeddingServer] embedder produced a non-finite value (${vec[i]}) at ${where} — refusing to serialize (JSON would launder it to null/0)`,
+      );
+    }
+  }
+}
 
 // ---------------------------------------------------------------------------
 // SerialQueue — ensures embed/embedMany never run concurrently on the GPU
@@ -240,6 +293,15 @@ export class EmbeddingServer implements IEmbedder {
       'Stepping down as leader so a fresh process can re-elect with a clean GPU context. ' +
       err.message,
     );
+    // INVARIANT: after a poisoned (permanently hung) inference, THIS process must
+    // never touch WebGPU again. The wedged native call still holds the same-PID
+    // re-entrant GPU lock, so if this process wins the next election and
+    // re-initializes WebGPU it would run a second inference beside the stuck one —
+    // the concurrent-GPU-inference segfault class the SerialQueue exists to
+    // prevent. Latch the process-wide fallback so any embedder built in this
+    // process from here on comes up on CPU; only a FRESH process may re-elect
+    // with a clean GPU context.
+    markWebGpuFallback();
     // Step down WITHOUT disposing the embedder: pipeline.dispose() would join the
     // stuck native thread (itself hanging) or touch a device-lost context. Clearing
     // the state registration + closing the HTTP server is enough for re-election;
@@ -305,14 +367,14 @@ export class EmbeddingServer implements IEmbedder {
       this.leadershipTimer = null;
     }
 
-    // Clear our registration if we still own it
+    // Clear our registration if we still own it. The serverId CAS runs inside the
+    // state lock, so a successor that claimed the slot between our last check and
+    // this shutdown can never be deregistered by us (the old read-then-clear had
+    // exactly that window).
     try {
-      const serverInfo = await this.stateManager.getEmbeddingServer();
-      if (serverInfo?.serverId === this.serverId) {
-        await this.stateManager.clearEmbeddingServer().catch((err) => {
-          logger.warn('[EmbeddingServer] Failed to clear embedding server state:', err);
-        });
-      }
+      await this.stateManager.clearEmbeddingServer({ serverId: this.serverId }).catch((err) => {
+        logger.warn('[EmbeddingServer] Failed to clear embedding server state:', err);
+      });
     } catch (err) {
       const msg = err instanceof Error ? err.message : String(err);
       // State manager is torn down before the embedding server during process exit —
@@ -428,13 +490,27 @@ export class EmbeddingServer implements IEmbedder {
   // ---- HTTP request handler ----
 
   private async handleRequest(req: http.IncomingMessage, res: http.ServerResponse): Promise<void> {
+    // Authorize BEFORE reading the body, so an unauthenticated caller cannot make
+    // us buffer up to MAX_BODY_SIZE, and before routing, so it cannot probe
+    // leadership/health state either. Mirrors the browser server.
+    if (!isAuthorized(req.headers['x-embedding-auth'])) {
+      res.writeHead(403, { 'Content-Type': 'application/json' });
+      res.end(JSON.stringify({ error: 'Unauthorized' }));
+      req.resume(); // drain so the socket is released rather than left half-read
+      return;
+    }
+
     const MAX_BODY_SIZE = 50 * 1024 * 1024; // 50 MB
 
-    let body = '';
+    // Buffered, then decoded once: `body += chunk` would split multi-byte UTF-8
+    // across chunk boundaries and silently replace both halves with U+FFFD — and
+    // the text arriving here is exactly what gets embedded, so that corruption
+    // would be baked into the vectors. See Utf8Body.
+    const collected = new Utf8Body();
     let oversized = false;
     req.on('data', (chunk: Buffer) => {
-      body += chunk;
-      if (body.length > MAX_BODY_SIZE) {
+      collected.push(chunk);
+      if (collected.byteLength > MAX_BODY_SIZE) {
         oversized = true;
         req.destroy(new Error('Payload too large'));
       }
@@ -459,6 +535,19 @@ export class EmbeddingServer implements IEmbedder {
         finish();
       });
 
+      // A client that goes away mid-request emits 'aborted'/'close' and does NOT
+      // reliably emit 'error', so neither handler above would fire and this promise
+      // would dangle forever — retaining req, res and the buffered body for the rest
+      // of the process's life. That is not an exotic case: EmbeddingClient destroys
+      // the request on its own 120s timeout, which is the same as this server's
+      // SerialQueue soft timeout, so a slow embed reliably produces one. handleRequest
+      // is invoked with `void`, so nothing else would ever observe the leak.
+      // On a normal request this fires after 'end' and finish() is already a no-op;
+      // nothing follows this await (the caller invokes handleRequest with `void`), so
+      // an early resolve is inert either way — the 'end' handler still writes the
+      // response on its own.
+      req.on('close', finish);
+
       req.on('end', async () => {
         try {
           if (req.method === 'GET' && req.url === '/health') {
@@ -480,7 +569,7 @@ export class EmbeddingServer implements IEmbedder {
             return;
           }
 
-          const data = JSON.parse(body);
+          const data = JSON.parse(collected.toString());
 
           switch (req.url) {
             case '/embed': {
@@ -491,6 +580,9 @@ export class EmbeddingServer implements IEmbedder {
               // message drives the client to reconnect to the new leader.
               this.assertServing();
               const result = await this.queue.enqueue(() => this.embedder.embed(data.text as string));
+              // Scan BEFORE serialization: a device-lost mid-inference yields NaN
+              // entries that JSON would launder into null → 0 on the client.
+              assertFiniteVector(result);
               res.writeHead(200, { 'Content-Type': 'application/json' });
               res.end(JSON.stringify({ data: Array.from(result) }));
               break;
@@ -498,6 +590,8 @@ export class EmbeddingServer implements IEmbedder {
             case '/embedMany': {
               this.assertServing();
               const results = await this.queue.enqueue(() => this.embedder.embedMany(data.texts as string[]));
+              // Same non-finite scan as /embed, per vector (see assertFiniteVector).
+              for (let i = 0; i < results.length; i++) assertFiniteVector(results[i]!, i);
               res.writeHead(200, { 'Content-Type': 'application/json' });
               res.end(JSON.stringify({ results: results.map(r => Array.from(r)) }));
               break;