Merge pull request #41 from off-grid-ai/fix/review-followup-proxyres-dims

fix: review follow-up — proxy stream guard, all-branch dim clamp, safer scripts

Author: alichherawalla

scripts/smoke-api.sh

@@ -10,14 +10,21 @@
 # (slower). Exits non-zero on the first failure so it works as a CI/pre-release gate.
 set -uo pipefail
 GW="${OFFGRID_GATEWAY_URL:-http://127.0.0.1:7878}"
+
+# Per-run temp files via mktemp (not fixed /tmp paths) - avoids a TOCTOU/symlink hijack on a
+# predictable name and lets concurrent runs coexist. Cleaned up on any exit.
+MODELS_JSON="$(mktemp -t smoke_models.XXXXXX)"
+STREAM_TXT="$(mktemp -t smoke_stream.XXXXXX)"
+trap 'rm -f "$MODELS_JSON" "$STREAM_TXT"' EXIT
+
 fail() { echo "  FAIL: $1"; exit 1; }
 ok() { echo "  ok: $1"; }
 
 echo "[smoke] gateway = $GW"
 
 # 0. Gateway reachable + model catalog (models-manager path)
-curl -s -m 5 "$GW/v1/models" >/tmp/smoke_models.json 2>&1 || fail "gateway not reachable - start the app first (npm run dev)"
-node -e 'const j=require("/tmp/smoke_models.json");if(!j.data||!j.data.length)process.exit(1);console.log("  models:",j.data.map(m=>m.id).join(", "))' \
+curl -s -m 5 "$GW/v1/models" >"$MODELS_JSON" 2>&1 || fail "gateway not reachable - start the app first (npm run dev)"
+MODELS_JSON="$MODELS_JSON" node -e 'const j=require(process.env.MODELS_JSON);if(!j.data||!j.data.length)process.exit(1);console.log("  models:",j.data.map(m=>m.id).join(", "))' \
   < /dev/null 2>/dev/null || fail "/v1/models returned no models"
 ok "/v1/models"
 
@@ -29,8 +36,8 @@ ok "chat non-stream"
 
 # 2. Chat stream (SSE proxy + retry refactor) - count content OR reasoning deltas
 curl -sN -m 120 "$GW/v1/chat/completions" -H 'Content-Type: application/json' \
-  -d '{"messages":[{"role":"user","content":"Count one to five."}],"max_tokens":48,"stream":true}' > /tmp/smoke_stream.txt 2>&1
-node -e 'const fs=require("fs");let n=0;for(const l of fs.readFileSync("/tmp/smoke_stream.txt","utf8").split("\n")){if(!l.startsWith("data:"))continue;const p=l.slice(5).trim();if(p==="[DONE]")continue;try{const d=JSON.parse(p).choices?.[0]?.delta||{};if(d.content||d.reasoning_content)n++;}catch(e){}}if(n<1)process.exit(1);console.log("  stream deltas:",n)' || fail "chat stream produced no deltas"
+  -d '{"messages":[{"role":"user","content":"Count one to five."}],"max_tokens":48,"stream":true}' > "$STREAM_TXT" 2>&1
+STREAM_TXT="$STREAM_TXT" node -e 'const fs=require("fs");let n=0;for(const l of fs.readFileSync(process.env.STREAM_TXT,"utf8").split("\n")){if(!l.startsWith("data:"))continue;const p=l.slice(5).trim();if(p==="[DONE]")continue;try{const d=JSON.parse(p).choices?.[0]?.delta||{};if(d.content||d.reasoning_content)n++;}catch(e){}}if(n<1)process.exit(1);console.log("  stream deltas:",n)' || fail "chat stream produced no deltas"
 ok "chat stream"
 
 # 3. Embeddings

scripts/test-db.sh

@@ -8,7 +8,9 @@
 # Kept OUT of the default `npm test` (those files are *.dbtest.ts, not *.test.ts) precisely
 # because of this ABI swap. Run it explicitly: `npm run test:db`. A CI job that runs it must
 # do so in an isolated step (the rebuild mutates node_modules).
-set -uo pipefail
+# -e: abort on any failure (e.g. `npm rebuild` failing) so we never run the tests against a stale
+# ABI. The EXIT trap below still fires on abort, so Electron's build is restored either way.
+set -euo pipefail
 
 restore() {
   echo "[test:db] restoring the Electron ABI build of better-sqlite3-multiple-ciphers..."

src/main/__tests__/stream-guards.test.ts

@@ -1,6 +1,6 @@
-import { describe, it, expect } from 'vitest';
+import { describe, it, expect, vi } from 'vitest';
 import { EventEmitter } from 'node:events';
-import { guardConsoleStreams } from '../stream-guards';
+import { guardConsoleStreams, guardProxyStreams } from '../stream-guards';
 
 // Fails-before / passes-after: a Node EventEmitter throws on emit('error') when there is NO
 // 'error' listener. That is exactly why an EPIPE on stdout crashed the main process. The guard
@@ -28,3 +28,54 @@ describe('guardConsoleStreams', () => {
     expect(() => b.emit('error', new Error('y'))).not.toThrow();
   });
 });
+
+// Regression for the proxyToLlama mid-stream crash: proxyRes.pipe(res) wired the pipe but neither
+// stream had an 'error' listener, so an upstream reset (llama-server aborting mid-stream) or a
+// client disconnect turned the unhandled 'error' event into an uncaught exception that took the
+// whole main process down. guardProxyStreams installs listeners on both ends.
+describe('guardProxyStreams', () => {
+  it('makes an upstream reset non-fatal and tears down the client response (would throw without the guard)', () => {
+    const upstream = new EventEmitter();
+    const client = Object.assign(new EventEmitter(), { destroy: vi.fn() });
+    const reset = Object.assign(new Error('read ECONNRESET'), { code: 'ECONNRESET' });
+
+    // Sanity: before guarding, an upstream 'error' with no listener throws (the crash we saw).
+    expect(() => upstream.emit('error', reset)).toThrow(/ECONNRESET/);
+
+    // After guarding, the same emit is swallowed and the client response is destroyed to free it.
+    const guarded = guardProxyStreams(upstream, client);
+    expect(guarded).toBe(2);
+    expect(() => upstream.emit('error', reset)).not.toThrow();
+    expect(client.destroy).toHaveBeenCalledTimes(1);
+  });
+
+  it('makes a client disconnect non-fatal AND tears down the upstream (stops reading llama-server)', () => {
+    const upstream = Object.assign(new EventEmitter(), { destroy: vi.fn() });
+    const client = new EventEmitter();
+    const guarded = guardProxyStreams(upstream, client);
+    expect(guarded).toBe(2);
+    // A client disconnect must not throw...
+    expect(() => client.emit('error', new Error('EPIPE'))).not.toThrow();
+    // ...and must destroy the upstream, else proxyRes.pipe(res) keeps draining llama-server
+    // after the client is gone (wasted local inference + a held socket).
+    expect(upstream.destroy).toHaveBeenCalledTimes(1);
+  });
+
+  it('a client disconnect with no upstream is still non-fatal', () => {
+    const client = new EventEmitter();
+    guardProxyStreams(undefined, client);
+    expect(() => client.emit('error', new Error('EPIPE'))).not.toThrow();
+  });
+
+  it('tolerates a client with no destroy method (does not throw when tearing down)', () => {
+    const upstream = new EventEmitter();
+    const client = new EventEmitter(); // no destroy()
+    expect(guardProxyStreams(upstream, client)).toBe(2);
+    expect(() => upstream.emit('error', new Error('reset'))).not.toThrow();
+  });
+
+  it('skips undefined streams and counts only the guarded ones', () => {
+    expect(guardProxyStreams(undefined, undefined)).toBe(0);
+    expect(guardProxyStreams(new EventEmitter(), undefined)).toBe(1);
+  });
+});

src/main/model-server.ts

@@ -37,6 +37,7 @@ import { llm, type LlmSettings } from './llm';
 import { LLAMA_SERVER_PORT, GATEWAY_PORT } from '../shared/ports';
 import { retryWithDeadline } from './lib/retry';
 import { resolveDims } from './model-server/dimensions';
+import { guardProxyStreams } from './stream-guards';
 import { classifyRef, decodeDataUrl, stripFileScheme, mimeFromExt, extForMime, toDataUrl } from './model-server/data-url';
 import { errBody, errMeta } from './model-server/errors';
 import { isAsync, matchPollRoute } from './model-server/async-request';
@@ -178,6 +179,11 @@ function proxyToLlama(
         { hostname: UPSTREAM_HOST, port: UPSTREAM_PORT, path: req.url, method: req.method, headers },
         (proxyRes) => {
           res.writeHead(proxyRes.statusCode || 502, proxyRes.headers);
+          // Guard both ends BEFORE piping: a mid-stream reset from llama-server (or a client
+          // disconnect) emits 'error' on these streams, and with no listener that becomes an
+          // uncaught exception that crashes the main process. Does not re-settle this promise —
+          // it has already resolved once piping begins.
+          guardProxyStreams(proxyRes, res);
           proxyRes.pipe(res);
           resolve();
         }

src/main/model-server/__tests__/dimensions.test.ts

@@ -26,6 +26,12 @@ describe('round64', () => {
     expect(round64(1024)).toBe(1024);
     expect(round64(512)).toBe(512);
   });
+
+  it('rejects non-finite input (NaN/Infinity clamp to the 256 floor, never leak NaN)', () => {
+    expect(round64(NaN)).toBe(256);
+    expect(round64(Infinity)).toBe(256);
+    expect(round64(-Infinity)).toBe(256);
+  });
 });
 
 describe('parseSize', () => {
@@ -55,19 +61,35 @@ describe('parseSize', () => {
 });
 
 describe('resolveDims', () => {
-  it('uses explicit numeric width/height as-is (no rounding)', () => {
-    expect(resolveDims({ width: 500, height: 700 })).toEqual({ width: 500, height: 700 });
+  it('rounds explicit numeric width/height to /64 (every branch is rounded)', () => {
+    // 500 -> 512, 700 -> 704 (700/64 = 10.9 -> 11 -> 704)
+    expect(resolveDims({ width: 500, height: 700 })).toEqual({ width: 512, height: 704 });
+  });
+
+  it('clamps out-of-range explicit width/height instead of passing raw values through', () => {
+    // Regression: the explicit width/height branch used to return raw values, so a caller could
+    // slip 10000 / 8 straight into image generation. Now clamped to 256-2048.
+    expect(resolveDims({ width: 10000, height: 8 })).toEqual({ width: 2048, height: 256 });
+  });
+
+  it('rejects NaN explicit width/height (clamps to the 256 floor, never leaks NaN)', () => {
+    expect(resolveDims({ width: NaN, height: NaN })).toEqual({ width: 256, height: 256 });
   });
 
   it('ignores explicit width/height when either is not a number', () => {
     // falls through to size
     expect(resolveDims({ width: 500, height: '700', size: '256x256' })).toEqual({ width: 256, height: 256 });
   });
 
-  it('parses an OpenAI size string when width/height absent', () => {
+  it('parses an OpenAI size string when width/height absent (and rounds it)', () => {
     expect(resolveDims({ size: '768x512' })).toEqual({ width: 768, height: 512 });
   });
 
+  it('rounds and clamps an out-of-range size string (the size branch is rounded too)', () => {
+    // Regression: the size-string branch used to return raw parsed values. 500 -> 512, 9000 -> 2048.
+    expect(resolveDims({ size: '500x9000' })).toEqual({ width: 512, height: 2048 });
+  });
+
   it('derives square dims from a 1:1 aspect ratio at the default 1K resolution', () => {
     expect(resolveDims({ aspect_ratio: '1:1' })).toEqual({ width: 1024, height: 1024 });
   });
@@ -100,10 +122,11 @@ describe('resolveDims', () => {
     expect(resolveDims({ size: 'nope', aspect_ratio: 'nope' })).toEqual({});
   });
 
-  it('prefers explicit dims over size over aspect_ratio', () => {
+  it('prefers explicit dims over size over aspect_ratio (still rounded/clamped)', () => {
+    // Explicit wins; 100 -> 256 (floor), 200 -> round64(200)=192 -> 256 (floor).
     expect(resolveDims({ width: 100, height: 200, size: '512x512', aspect_ratio: '1:1' })).toEqual({
-      width: 100,
-      height: 200,
+      width: 256,
+      height: 256,
     });
     expect(resolveDims({ size: '512x512', aspect_ratio: '1:1' })).toEqual({ width: 512, height: 512 });
   });

src/main/model-server/dimensions.ts

@@ -2,8 +2,11 @@
 // Extracted from model-server.ts so the branch logic (OpenAI size, OpenRouter
 // aspect_ratio + resolution, explicit width/height) is unit-testable.
 
-// Round to a multiple of 64 (diffusion models require it), clamped sane.
+// Round to a multiple of 64 (diffusion models require it), clamped sane. A non-finite input
+// (NaN/Infinity - e.g. a bad size string or a 0-denominator aspect ratio) clamps to the 256 floor
+// rather than leaking NaN downstream into image generation.
 export function round64(n: number): number {
+  if (!Number.isFinite(n)) return 256;
   return Math.max(256, Math.min(2048, Math.round(n / 64) * 64));
 }
 
@@ -16,17 +19,21 @@ export function parseSize(size: unknown): { width?: number; height?: number } {
 }
 
 // Resolve output dimensions from any of: explicit width/height, an OpenAI
-// "WIDTHxHEIGHT" size, or an OpenRouter aspect_ratio + resolution pair.
+// "WIDTHxHEIGHT" size, or an OpenRouter aspect_ratio + resolution pair. EVERY branch runs its
+// result through round64 (round to /64, clamp 256-2048, reject NaN) so no path can leak an
+// out-of-range or non-numeric dimension into image generation.
 export function resolveDims(p: {
   width?: unknown;
   height?: unknown;
   size?: unknown;
   aspect_ratio?: unknown;
   resolution?: unknown;
 }): { width?: number; height?: number } {
-  if (typeof p.width === 'number' && typeof p.height === 'number') return { width: p.width, height: p.height };
+  if (typeof p.width === 'number' && typeof p.height === 'number') {
+    return { width: round64(p.width), height: round64(p.height) };
+  }
   const fromSize = parseSize(p.size);
-  if (fromSize.width && fromSize.height) return fromSize;
+  if (fromSize.width && fromSize.height) return { width: round64(fromSize.width), height: round64(fromSize.height) };
   if (typeof p.aspect_ratio === 'string') {
     const m = /^(\d+)\s*[:x×]\s*(\d+)$/.exec(p.aspect_ratio.trim());
     if (m) {

src/main/stream-guards.ts

@@ -24,3 +24,33 @@ export function guardConsoleStreams(streams: Array<NodeJS.EventEmitter | undefin
   }
   return n;
 }
+
+/** Guard both ends of a proxied stream (upstream response + client response) against a mid-stream
+ *  reset. Piping `proxyRes` into `res` only wires the pipe; neither stream gets an 'error' listener,
+ *  so if llama-server aborts the connection or the client disconnects mid-stream, the unhandled
+ *  'error' event becomes an uncaught exception and crashes the whole main process (an EventEmitter
+ *  throws on 'error' only when there is NO listener). We swallow the error and destroy the client
+ *  response so its socket is released. This does NOT settle any outer promise — by the time piping
+ *  starts the proxy attempt has already resolved, and re-settling here would be a double-settle.
+ *  Returns the count guarded (for tests). */
+export function guardProxyStreams(
+  upstream: (NodeJS.EventEmitter & { destroy?: (err?: Error) => void }) | undefined,
+  client: (NodeJS.EventEmitter & { destroy?: (err?: Error) => void }) | undefined
+): number {
+  let n = 0;
+  const destroy = (s?: { destroy?: (err?: Error) => void }): void => {
+    try { s?.destroy?.(); } catch { /* already destroyed */ }
+  };
+  if (upstream && typeof upstream.on === 'function') {
+    // Upstream (llama-server) reset mid-stream: tear down the client response so its socket frees.
+    upstream.on('error', () => destroy(client));
+    n++;
+  }
+  if (client && typeof client.on === 'function') {
+    // Client disconnected mid-stream: stop reading from llama-server too, else proxyRes.pipe(res)
+    // keeps draining the upstream (wasted local inference + a held socket) after the client is gone.
+    client.on('error', () => destroy(upstream));
+    n++;
+  }
+  return n;
+}