From fd97380a9be5cc42a77eb2c6340caa8f01f598c1 Mon Sep 17 00:00:00 2001
From: Sawyer Cutler <sawyer@dirtroad.dev>
Date: Wed, 5 Aug 2026 18:23:23 -0700
Subject: [PATCH] feat: configurable embed/rerank client timeouts via env vars

Add EMBED_TIMEOUT_MS and RERANK_TIMEOUT_MS so operators can raise the
fixed 10s HTTP timeout per deployment. Ollama embed calls have been
observed taking 6-9s in production, leaving little margin against the
hardcoded default.
---
 src/config.ts            |  6 ++++
 src/memory.test.ts       |  4 +++
 src/mount-config.test.ts | 59 ++++++++++++++++++++++++++++++++++++++++
 src/mount-config.ts      |  2 ++
 src/services/search.ts   |  4 +++
 5 files changed, 75 insertions(+)
 create mode 100644 src/mount-config.test.ts

diff --git a/src/config.ts b/src/config.ts
index b99a11e..f34223d 100644
--- a/src/config.ts
+++ b/src/config.ts
@@ -26,6 +26,9 @@ export type EngineConfig = {
     model: string;
     apiStyle: string;
     apiKey: string | undefined;
+    // Per-request timeout override (EMBED_TIMEOUT_MS). `undefined` means
+    // "use embed-client.ts's own default" — must NOT be defaulted here.
+    timeoutMs: number | undefined;
   };
   rerank: {
     baseUrl: string | undefined;
@@ -40,5 +43,8 @@ export type EngineConfig = {
     // single constant is exactly how a small-model budget got applied to
     // the default model's much larger context window.
     maxDocChars: number | undefined;
+    // Per-request timeout override (RERANK_TIMEOUT_MS). `undefined` means
+    // "use rerank-client.ts's own default" — must NOT be defaulted here.
+    timeoutMs: number | undefined;
   };
 };
diff --git a/src/memory.test.ts b/src/memory.test.ts
index c4957cd..11fc1c4 100644
--- a/src/memory.test.ts
+++ b/src/memory.test.ts
@@ -113,12 +113,14 @@ const wiringConfig: MemoryConfig = {
       model: "m",
       apiStyle: "openai",
       apiKey: undefined,
+      timeoutMs: undefined,
     },
     rerank: {
       baseUrl: undefined,
       model: undefined,
       apiKey: undefined,
       maxDocChars: undefined,
+      timeoutMs: undefined,
     },
   },
 };
@@ -145,6 +147,7 @@ function baseConfig(
         model: "m",
         apiStyle: "openai",
         apiKey: undefined,
+        timeoutMs: undefined,
       },
       rerank,
     },
@@ -163,6 +166,7 @@ describe("createMemory — construction validation", () => {
           model: "bge-reranker-base",
           apiKey: undefined,
           maxDocChars: 5_000,
+          timeoutMs: undefined,
         }),
       }),
     ).toThrow(RerankConfigError);
diff --git a/src/mount-config.test.ts b/src/mount-config.test.ts
new file mode 100644
index 0000000..84bfa33
--- /dev/null
+++ b/src/mount-config.test.ts
@@ -0,0 +1,59 @@
+import { afterEach, beforeEach, describe, expect, it } from "bun:test";
+import { loadMemoryConfig } from "./mount-config.ts";
+
+const REQUIRED_ENV = {
+  KNOWLEDGE_DATABASE_URL: "postgres://localhost:5432/test",
+  EMBED_BASE_URL: "http://embed.example",
+  EMBED_MODEL: "test-model",
+};
+
+const ENV_KEYS = [
+  ...Object.keys(REQUIRED_ENV),
+  "EMBED_TIMEOUT_MS",
+  "RERANK_TIMEOUT_MS",
+];
+
+let savedEnv: Record<string, string | undefined>;
+
+beforeEach(() => {
+  savedEnv = {};
+  for (const key of ENV_KEYS) {
+    savedEnv[key] = process.env[key];
+    delete process.env[key];
+  }
+  for (const [key, value] of Object.entries(REQUIRED_ENV)) {
+    process.env[key] = value;
+  }
+});
+
+afterEach(() => {
+  for (const key of ENV_KEYS) {
+    if (savedEnv[key] === undefined) delete process.env[key];
+    else process.env[key] = savedEnv[key];
+  }
+});
+
+describe("loadMemoryConfig — EMBED_TIMEOUT_MS / RERANK_TIMEOUT_MS", () => {
+  it("leaves embed.timeoutMs and rerank.timeoutMs undefined when unset, so the clients' own defaults apply", () => {
+    const config = loadMemoryConfig();
+    expect(config.memory.embed.timeoutMs).toBeUndefined();
+    expect(config.memory.rerank.timeoutMs).toBeUndefined();
+  });
+
+  it("flows EMBED_TIMEOUT_MS through to embed.timeoutMs", () => {
+    process.env.EMBED_TIMEOUT_MS = "20000";
+    const config = loadMemoryConfig();
+    expect(config.memory.embed.timeoutMs).toBe(20_000);
+  });
+
+  it("flows RERANK_TIMEOUT_MS through to rerank.timeoutMs", () => {
+    process.env.RERANK_TIMEOUT_MS = "15000";
+    const config = loadMemoryConfig();
+    expect(config.memory.rerank.timeoutMs).toBe(15_000);
+  });
+
+  it("rejects a non-positive-integer EMBED_TIMEOUT_MS", () => {
+    process.env.EMBED_TIMEOUT_MS = "not-a-number";
+    expect(() => loadMemoryConfig()).toThrow("EMBED_TIMEOUT_MS must be a positive integer");
+  });
+});
diff --git a/src/mount-config.ts b/src/mount-config.ts
index 8ee7d17..fc593cd 100644
--- a/src/mount-config.ts
+++ b/src/mount-config.ts
@@ -64,12 +64,14 @@ export function loadMemoryConfig(): MemoryConfig {
         model: requireEnv("EMBED_MODEL"),
         apiStyle: optionalEnv("EMBED_API_STYLE") ?? "openai",
         apiKey: optionalEnv("EMBED_API_KEY"),
+        timeoutMs: optionalIntEnv("EMBED_TIMEOUT_MS"),
       },
       rerank: {
         baseUrl: optionalEnv("RERANK_BASE_URL"),
         model: optionalEnv("RERANK_MODEL"),
         apiKey: optionalEnv("RERANK_API_KEY"),
         maxDocChars: optionalIntEnv("RERANK_MAX_DOC_CHARS"),
+        timeoutMs: optionalIntEnv("RERANK_TIMEOUT_MS"),
       },
     },
   };
diff --git a/src/services/search.ts b/src/services/search.ts
index 4d76d9e..9205c81 100644
--- a/src/services/search.ts
+++ b/src/services/search.ts
@@ -672,6 +672,7 @@ function toEmbedClientConfig(embed: EngineConfig["embed"]): EmbedClientConfig {
     modelId: embed.model,
     apiStyle: embed.apiStyle as EmbedClientConfig["apiStyle"],
     ...(embed.apiKey !== undefined ? { apiKey: embed.apiKey } : {}),
+    ...(embed.timeoutMs !== undefined ? { timeoutMs: embed.timeoutMs } : {}),
   };
 }
 
@@ -694,6 +695,9 @@ export function toRerankClientConfig(
     ...(rerank.maxDocChars !== undefined
       ? { maxDocChars: rerank.maxDocChars }
       : {}),
+    ...(rerank.timeoutMs !== undefined
+      ? { timeoutMs: rerank.timeoutMs }
+      : {}),
   };
 }