@remnic/core 1.1.31 → 9.3.516

package/src/embedding-fallback.ts

@@ -1,5 +1,5 @@
 import path from "node:path";
-import { mkdir, readFile, writeFile } from "node:fs/promises";
+import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises";
 import { log } from "./logger.js";
 import { readEnvVar } from "./runtime/env.js";
 import type { PluginConfig } from "./types.js";
@@ -132,6 +132,7 @@ export type EmbedMode = "lookup" | "index";
 export class EmbeddingFallback {
   private readonly indexPath: string;
   private loaded: EmbeddingIndexFile | null = null;
+  private mutationQueue: Promise<void> = Promise.resolve();
 
   constructor(private readonly config: PluginConfig) {
     this.indexPath = path.join(config.memoryDir, "state", "embeddings.json");
@@ -291,23 +292,36 @@ export class EmbeddingFallback {
     const vector = await this.embed(content, provider, { mode: "index" });
     if (!vector) return;
 
-    const index = await this.loadIndex(provider);
-    const relPath = toMemoryRelativePath(this.config.memoryDir, filePath);
-    index.entries[memoryId] = {
-      vector,
-      path: relPath,
-    };
-    await this.saveIndex(index);
+    await this.enqueueIndexMutation(async () => {
+      const index = await this.loadIndex(provider);
+      const relPath = toMemoryRelativePath(this.config.memoryDir, filePath);
+      index.entries[memoryId] = {
+        vector,
+        path: relPath,
+      };
+      await this.saveIndex(index);
+    });
   }
 
   async removeFromIndex(memoryId: string): Promise<void> {
     const provider = await this.resolveProvider();
     if (!provider) return;
 
-    const index = await this.loadIndex(provider);
-    if (!index.entries[memoryId]) return;
-    delete index.entries[memoryId];
-    await this.saveIndex(index);
+    await this.enqueueIndexMutation(async () => {
+      const index = await this.loadIndex(provider);
+      if (!index.entries[memoryId]) return;
+      delete index.entries[memoryId];
+      await this.saveIndex(index);
+    });
+  }
+
+  private enqueueIndexMutation<T>(mutation: () => Promise<T>): Promise<T> {
+    const run = this.mutationQueue.catch(() => undefined).then(mutation);
+    this.mutationQueue = run.then(
+      () => undefined,
+      () => undefined,
+    );
+    return run;
   }
 
   private async resolveProvider(): Promise<ProviderConfig | null> {
@@ -473,7 +487,14 @@ export class EmbeddingFallback {
     try {
       const raw = await readFile(this.indexPath, "utf-8");
       const parsed = JSON.parse(raw) as EmbeddingIndexFile;
-      if (parsed && parsed.version === 1 && parsed.entries && typeof parsed.entries === "object") {
+      if (
+        parsed &&
+        parsed.version === 1 &&
+        parsed.provider === provider.type &&
+        parsed.model === provider.model &&
+        parsed.entries &&
+        typeof parsed.entries === "object"
+      ) {
         this.loaded = {
           version: 1,
           provider: provider.type,
@@ -496,8 +517,19 @@ export class EmbeddingFallback {
   }
 
   private async saveIndex(index: EmbeddingIndexFile): Promise<void> {
-    await mkdir(path.dirname(this.indexPath), { recursive: true });
-    await writeFile(this.indexPath, JSON.stringify(index), "utf-8");
+    const dir = path.dirname(this.indexPath);
+    await mkdir(dir, { recursive: true });
+    const tempPath = path.join(
+      dir,
+      `.embeddings.${process.pid}.${Date.now()}.${Math.random().toString(16).slice(2)}.tmp`,
+    );
+    try {
+      await writeFile(tempPath, JSON.stringify(index), "utf-8");
+      await rename(tempPath, this.indexPath);
+    } catch (err) {
+      await rm(tempPath, { force: true }).catch(() => undefined);
+      throw err;
+    }
     this.loaded = index;
   }
 }

package/src/enrichment/audit.ts

@@ -61,6 +61,10 @@ export async function readAuditLog(
 ): Promise<EnrichmentAuditEntry[]> {
   const filePath = auditFilePath(auditDir);
   if (!existsSync(filePath)) return [];
+  const sinceMs = since === undefined ? undefined : Date.parse(since);
+  if (since !== undefined && !Number.isFinite(sinceMs)) {
+    throw new Error(`Invalid enrichment audit since timestamp: ${since}`);
+  }
 
   const raw = await readFile(filePath, "utf-8");
   const entries: EnrichmentAuditEntry[] = [];
@@ -77,7 +81,10 @@ export async function readAuditLog(
         "entityName" in parsed
       ) {
         const entry = parsed as EnrichmentAuditEntry;
-        if (since && entry.timestamp < since) continue;
+        if (typeof entry.timestamp !== "string") continue;
+        const entryMs = Date.parse(entry.timestamp);
+        if (!Number.isFinite(entryMs)) continue;
+        if (sinceMs !== undefined && entryMs < sinceMs) continue;
         entries.push(entry);
       }
     } catch {

package/src/enrichment/pipeline.ts

@@ -30,16 +30,17 @@ interface RateLimitBucket {
   dayReset: number;
 }
 
+const rateBuckets = new Map<string, RateLimitBucket>();
+
 function isRateLimited(
   provider: EnrichmentProvider,
   config: EnrichmentPipelineConfig,
-  buckets: Map<string, RateLimitBucket>,
 ): boolean {
   const providerCfg = config.providers.find((p) => p.id === provider.id);
   if (!providerCfg?.rateLimit) return false;
 
   const now = Date.now();
-  let bucket = buckets.get(provider.id);
+  let bucket = rateBuckets.get(provider.id);
   if (!bucket) {
     bucket = {
       minuteCount: 0,
@@ -47,7 +48,7 @@ function isRateLimited(
       dayCount: 0,
       dayReset: now + 86_400_000,
     };
-    buckets.set(provider.id, bucket);
+    rateBuckets.set(provider.id, bucket);
   }
 
   // Reset windows if expired
@@ -66,9 +67,8 @@ function isRateLimited(
 
 function recordCall(
   providerId: string,
-  buckets: Map<string, RateLimitBucket>,
 ): void {
-  const bucket = buckets.get(providerId);
+  const bucket = rateBuckets.get(providerId);
   if (bucket) {
     bucket.minuteCount += 1;
     bucket.dayCount += 1;
@@ -88,13 +88,21 @@ export async function runEnrichmentPipeline(
   if (!config.enabled) return [];
   if (entities.length === 0) return [];
 
-  const rateBuckets = new Map<string, RateLimitBucket>();
   const results: EnrichmentResult[] = [];
 
   for (const entity of entities) {
     const providers = registry.getForImportance(entity.importanceLevel, config);
+    const maxCandidates = config.maxCandidatesPerEntity;
+    const hasPositiveCandidateBudget = maxCandidates > 0;
+    let remainingCandidateBudget = hasPositiveCandidateBudget
+      ? maxCandidates
+      : Number.POSITIVE_INFINITY;
 
     for (const provider of providers) {
+      if (hasPositiveCandidateBudget && remainingCandidateBudget <= 0) {
+        break;
+      }
+
       const start = Date.now();
 
       // Check availability
@@ -122,7 +130,7 @@ export async function runEnrichmentPipeline(
       }
 
       // Check rate limit
-      if (isRateLimited(provider, config, rateBuckets)) {
+      if (isRateLimited(provider, config)) {
         log.debug?.(
           `enrichment: skipping provider ${provider.id} for ${entity.name} — rate limited`,
         );
@@ -146,7 +154,7 @@ export async function runEnrichmentPipeline(
       try {
         candidates = await provider.enrich(entity);
       } catch (err) {
-        recordCall(provider.id, rateBuckets);
+        recordCall(provider.id);
         log.error?.(
           `enrichment: provider ${provider.id} failed for ${entity.name}: ${err instanceof Error ? err.message : String(err)}`,
         );
@@ -161,21 +169,21 @@ export async function runEnrichmentPipeline(
         });
         continue;
       }
-      recordCall(provider.id, rateBuckets);
+      recordCall(provider.id);
 
       // Tag each candidate with provider id
       for (const candidate of candidates) {
         candidate.source = provider.id;
       }
 
-      // Cap at maxCandidatesPerEntity.
+      // Cap at maxCandidatesPerEntity across all providers for this entity.
       // 0 means "accept none"; undefined/negative means "no cap".
-      const maxCandidates = config.maxCandidatesPerEntity;
       let accepted: EnrichmentCandidate[];
       if (maxCandidates === 0) {
         accepted = [];
-      } else if (maxCandidates > 0 && candidates.length > maxCandidates) {
-        accepted = candidates.slice(0, maxCandidates);
+      } else if (hasPositiveCandidateBudget) {
+        accepted = candidates.slice(0, remainingCandidateBudget);
+        remainingCandidateBudget -= accepted.length;
       } else {
         accepted = candidates;
       }

package/src/enrichment/web-search-provider.ts

@@ -42,12 +42,7 @@ export class WebSearchProvider implements EnrichmentProvider {
     if (!this.searchFn) return [];
 
     const query = `${entity.name} ${entity.type}`;
-    let snippets: string[];
-    try {
-      snippets = await this.searchFn(query);
-    } catch {
-      return [];
-    }
+    const snippets = await this.searchFn(query);
 
     return snippets
       .filter((s) => typeof s === "string" && s.trim().length > 0)

package/src/entity-retrieval.ts

@@ -67,6 +67,7 @@ type EntityHintSnippet = {
 export interface BuildEntityRecallSectionOptions {
   config: PluginConfig;
   storage: StorageManager;
+  namespaceStorage?: (namespace: string) => Promise<StorageManager>;
   query: string;
   recallNamespaces?: string[];
   recentTurns: number;
@@ -278,6 +279,38 @@ async function readNativeChunks(
   }).catch(() => []);
 }
 
+async function resolveEntityIndexStorages(
+  storage: StorageManager,
+  config: PluginConfig,
+  recallNamespaces?: string[],
+  namespaceStorage?: (namespace: string) => Promise<StorageManager>,
+): Promise<StorageManager[]> {
+  if (
+    !config.namespacesEnabled ||
+    !namespaceStorage ||
+    !recallNamespaces ||
+    recallNamespaces.length === 0
+  ) {
+    return [storage];
+  }
+
+  const storages: StorageManager[] = [];
+  const seenDirs = new Set<string>();
+  for (const namespace of uniqueStrings(recallNamespaces)) {
+    try {
+      const scopedStorage = await namespaceStorage(namespace);
+      const storageDir = path.resolve(scopedStorage.dir);
+      if (seenDirs.has(storageDir)) continue;
+      seenDirs.add(storageDir);
+      storages.push(scopedStorage);
+    } catch {
+      continue;
+    }
+  }
+
+  return storages.length > 0 ? storages : [storage];
+}
+
 function entityIndexStatePath(storage: StorageManager): string {
   return path.join(storage.dir, "state", "entity-mention-index.json");
 }
@@ -353,13 +386,24 @@ async function buildEntityMentionIndex(
   storage: StorageManager,
   config: PluginConfig,
   recallNamespaces?: string[],
+  namespaceStorage?: (namespace: string) => Promise<StorageManager>,
 ): Promise<EntityMentionIndex> {
-  const [previousIndex, entityFiles, memories, nativeChunks] = await Promise.all([
-    readEntityIndexState(storage),
-    storage.readAllEntityFiles(),
-    storage.readAllMemories(),
+  const storages = await resolveEntityIndexStorages(
+    storage,
+    config,
+    recallNamespaces,
+    namespaceStorage,
+  );
+  const shouldPersistIndex =
+    storages.length === 1 && path.resolve(storages[0]!.dir) === path.resolve(storage.dir);
+  const [previousIndex, entityFileSets, memorySets, nativeChunks] = await Promise.all([
+    shouldPersistIndex ? readEntityIndexState(storage) : Promise.resolve(null),
+    Promise.all(storages.map((scopedStorage) => scopedStorage.readAllEntityFiles())),
+    Promise.all(storages.map((scopedStorage) => scopedStorage.readAllMemories())),
     readNativeChunks(config, recallNamespaces),
   ]);
+  const entityFiles = entityFileSets.flat();
+  const memories = memorySets.flat();
 
   const entities = new Map<string, EntityMentionIndexEntry>();
   for (const entity of entityFiles) {
@@ -436,7 +480,9 @@ async function buildEntityMentionIndex(
         : new Date().toISOString(),
     entities: sortedEntities,
   };
-  await writeEntityIndexState(storage, index);
+  if (shouldPersistIndex) {
+    await writeEntityIndexState(storage, index);
+  }
   return index;
 }
 
@@ -715,7 +761,12 @@ export async function buildEntityRecallSection(options: BuildEntityRecallSection
   const mode = detectEntityQueryMode(options.query);
   if (!mode) return null;
 
-  const index = await buildEntityMentionIndex(options.storage, options.config, options.recallNamespaces);
+  const index = await buildEntityMentionIndex(
+    options.storage,
+    options.config,
+    options.recallNamespaces,
+    options.namespaceStorage,
+  );
   if (index.entities.length === 0) return null;
 
   const explicitCandidates = resolveExplicitCandidates(index, options.query);

package/src/evals.ts

@@ -295,6 +295,14 @@ function assertIsoTimestamp(value: unknown, field: string): string {
   return normalized;
 }
 
+function assertNonNegativeInteger(value: unknown, field: string): number {
+  const parsed = Number(value);
+  if (!Number.isSafeInteger(parsed) || parsed < 0) {
+    throw new Error(`${field} must be a non-negative integer`);
+  }
+  return parsed;
+}
+
 function assertPathWithin(rootDir: string, targetPath: string, field: string): void {
   const resolvedRoot = path.resolve(rootDir);
   const resolvedTarget = path.resolve(targetPath);
@@ -374,25 +382,26 @@ export function validateEvalRunSummary(raw: unknown): EvalRunSummary {
     throw new Error("status must be one of running|completed|failed|partial");
   }
 
-  const totalCases = Number(raw.totalCases);
-  const passedCases = Number(raw.passedCases);
-  const failedCases = Number(raw.failedCases);
-  if (!Number.isFinite(totalCases) || totalCases < 0) throw new Error("totalCases must be a non-negative number");
-  if (!Number.isFinite(passedCases) || passedCases < 0) throw new Error("passedCases must be a non-negative number");
-  if (!Number.isFinite(failedCases) || failedCases < 0) throw new Error("failedCases must be a non-negative number");
+  const totalCases = assertNonNegativeInteger(raw.totalCases, "totalCases");
+  const passedCases = assertNonNegativeInteger(raw.passedCases, "passedCases");
+  const failedCases = assertNonNegativeInteger(raw.failedCases, "failedCases");
+  if (passedCases + failedCases > totalCases) {
+    throw new Error("passedCases + failedCases must be less than or equal to totalCases");
+  }
 
   const metrics = parseOptionalEvalRunMetrics(raw.metrics);
+  const completedAt =
+    typeof raw.completedAt === "string" && raw.completedAt.trim().length > 0
+      ? assertIsoTimestamp(raw.completedAt, "completedAt")
+      : undefined;
 
   return {
     schemaVersion: 1,
     runId: assertString(raw.runId, "runId"),
     benchmarkId: assertString(raw.benchmarkId, "benchmarkId"),
     status: status as EvalRunStatus,
-    startedAt: assertString(raw.startedAt, "startedAt"),
-    completedAt:
-      typeof raw.completedAt === "string" && raw.completedAt.trim().length > 0
-        ? raw.completedAt.trim()
-        : undefined,
+    startedAt: assertIsoTimestamp(raw.startedAt, "startedAt"),
+    completedAt,
     totalCases,
     passedCases,
     failedCases,
@@ -421,7 +430,7 @@ export function validateEvalBaselineSnapshot(raw: unknown): EvalBaselineSnapshot
       runId: assertString(item.runId, `benchmarks[${index}].runId`),
       completedAt:
         typeof item.completedAt === "string" && item.completedAt.trim().length > 0
-          ? item.completedAt.trim()
+          ? assertIsoTimestamp(item.completedAt, `benchmarks[${index}].completedAt`)
           : undefined,
       gitRef:
         typeof item.gitRef === "string" && item.gitRef.trim().length > 0
@@ -443,7 +452,7 @@ export function validateEvalBaselineSnapshot(raw: unknown): EvalBaselineSnapshot
   return {
     schemaVersion: 1,
     snapshotId: assertString(raw.snapshotId, "snapshotId"),
-    createdAt: assertString(raw.createdAt, "createdAt"),
+    createdAt: assertIsoTimestamp(raw.createdAt, "createdAt"),
     sourceRootDir: assertString(raw.sourceRootDir, "sourceRootDir"),
     benchmarkCount,
     benchmarks,

package/src/explicit-capture.test.ts

@@ -0,0 +1,40 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+
+import {
+  parseInlineExplicitCaptureNotes,
+  validateExplicitCaptureInput,
+} from "./explicit-capture.js";
+
+function parseSingleNote(confidenceLine: string): ReturnType<typeof parseInlineExplicitCaptureNotes>[number] {
+  const notes = parseInlineExplicitCaptureNotes(`
+<memory_note>
+content: Inline explicit capture content for testing
+category: fact
+${confidenceLine}
+</memory_note>
+`);
+  assert.equal(notes.length, 1);
+  return notes[0]!;
+}
+
+test("inline explicit capture rejects malformed confidence values", () => {
+  for (const confidenceLine of ["confidence: abc", "confidence: 0.5x"]) {
+    const note = parseSingleNote(confidenceLine);
+    assert.throws(
+      () => validateExplicitCaptureInput(note),
+      /confidence must be a finite number/,
+      `${confidenceLine} should be rejected`,
+    );
+  }
+});
+
+test("inline explicit capture preserves valid confidence values", () => {
+  const validated = validateExplicitCaptureInput(parseSingleNote("confidence: 0.5"));
+  assert.equal(validated.confidence, 0.5);
+});
+
+test("inline explicit capture defaults omitted confidence", () => {
+  const validated = validateExplicitCaptureInput(parseSingleNote(""));
+  assert.equal(validated.confidence, 0.95);
+});

package/src/explicit-capture.ts

@@ -205,6 +205,15 @@ function parseExplicitCaptureTtl(ttl: string | undefined): string | undefined {
   return new Date(Date.now() + amount * multiplier).toISOString();
 }
 
+function parseInlineConfidence(value: string): number {
+  const trimmed = value.trim();
+  if (!/^[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:e[+-]?\d+)?$/i.test(trimmed)) {
+    return Number.NaN;
+  }
+  const parsed = Number(trimmed);
+  return Number.isFinite(parsed) ? parsed : Number.NaN;
+}
+
 function parseInlineNote(block: string): ExplicitCaptureInput | null {
   const lines = block.replace(/\r/g, "").split("\n");
   const note: Partial<ExplicitCaptureInput> = {};
@@ -248,7 +257,7 @@ function parseInlineNote(block: string): ExplicitCaptureInput | null {
         note.category = value;
         break;
       case "confidence":
-        note.confidence = Number.parseFloat(value);
+        note.confidence = parseInlineConfidence(value);
         break;
       case "namespace":
         note.namespace = value;
@@ -326,7 +335,10 @@ export function validateExplicitCaptureInput(
   assertNoSecretLikeMetadata("ttl", input.ttl);
   assertNoSecretLikeMetadataList("tags", input.tags);
 
-  const confidence = Number.isFinite(input.confidence) ? Number(input.confidence) : 0.95;
+  if (input.confidence !== undefined && !Number.isFinite(input.confidence)) {
+    throw new Error("confidence must be a finite number");
+  }
+  const confidence = input.confidence === undefined ? 0.95 : Number(input.confidence);
   if (confidence < 0 || confidence > 1) {
     throw new Error("confidence must be between 0 and 1");
   }

package/src/extraction.ts

@@ -2,7 +2,7 @@ import OpenAI from "openai";
 import { log } from "./logger.js";
 import { delinearize } from "./delinearize.js";
 import { LocalLlmClient } from "./local-llm.js";
-import { FallbackLlmClient } from "./fallback-llm.js";
+import { FallbackLlmClient, fallbackLlmRuntimeContextFromConfig } from "./fallback-llm.js";
 import {
   ExtractionResultSchema,
   ConsolidationResultSchema,
@@ -115,9 +115,10 @@ export class ExtractionEngine {
       log.warn("no OpenAI API key — direct OpenAI client disabled; local and gateway fallback paths remain available");
     }
     this.localLlm = localLlm ?? new LocalLlmClient(config, modelRegistry);
-    this.fallbackLlm = new FallbackLlmClient(gatewayConfig, {
-      workspaceDir: config.workspaceDir,
-    });
+    this.fallbackLlm = new FallbackLlmClient(
+      gatewayConfig,
+      fallbackLlmRuntimeContextFromConfig(config),
+    );
     this.modelRegistry = modelRegistry ?? new ModelRegistry(config.memoryDir);
     if (config.modelSource === "gateway") {
       log.debug(
@@ -201,6 +202,20 @@ export class ExtractionEngine {
       || (result.relationships?.length ?? 0) > 0;
   }
 
+  private looksLikeExtractionResultPayload(parsed: any): boolean {
+    return !!parsed &&
+      typeof parsed === "object" &&
+      !Array.isArray(parsed) &&
+      (
+        "facts" in parsed ||
+        "entities" in parsed ||
+        "profileUpdates" in parsed ||
+        "questions" in parsed ||
+        "relationships" in parsed ||
+        "identityReflection" in parsed
+      );
+  }
+
   private normalizeExtractionResultPayload(parsed: any): ExtractionResult {
     const entities = Array.isArray(parsed?.entities)
       ? parsed.entities
@@ -1299,39 +1314,36 @@ ${truncatedConversation}`;
     // Avoid logging model output content by default (may contain user data).
     log.debug(`extractWithLocalLlm: got response content, length=${content.length}`);
 
-    try {
-      for (const candidate of extractJsonCandidates(content)) {
-        try {
-          log.debug(`extractWithLocalLlm: attempting JSON parse, candidate length=${candidate.length}`);
-          const parsed = JSON.parse(candidate);
+    for (const candidate of extractJsonCandidates(content)) {
+      try {
+        log.debug(`extractWithLocalLlm: attempting JSON parse, candidate length=${candidate.length}`);
+        const parsed = JSON.parse(candidate);
+        if (!this.looksLikeExtractionResultPayload(parsed)) {
+          continue;
+        }
 
-          const result: ExtractionResult = this.normalizeExtractionResultPayload(parsed);
+        const result: ExtractionResult = this.normalizeExtractionResultPayload(parsed);
 
-          log.debug(
-            `extractWithLocalLlm: successfully parsed response, facts=${result.facts.length}, entities=${result.entities.length}, profileUpdates=${result.profileUpdates.length}, questions=${result.questions.length}`,
-          );
-          return result;
-        } catch {
-          // keep trying candidates
-        }
-      }
-      return null;
-    } catch (err) {
-      // Try to extract partial facts from truncated JSON
-      log.debug("extractWithLocalLlm: JSON parse failed, attempting partial extraction...");
-      const partial = this.extractPartialFacts(content);
-      if (partial.facts.length > 0 || partial.entities.length > 0) {
         log.debug(
-          `extractWithLocalLlm: extracted ${partial.facts.length} partial facts from truncated JSON`,
+          `extractWithLocalLlm: successfully parsed response, facts=${result.facts.length}, entities=${result.entities.length}, profileUpdates=${result.profileUpdates.length}, questions=${result.questions.length}`,
         );
-        return partial;
+        return result;
+      } catch {
+        // keep trying candidates
       }
+    }
 
-      // Could not extract anything
-      const errMsg = err instanceof Error ? err.message : String(err);
-      log.debug(`extractWithLocalLlm: JSON parse error: ${errMsg}`);
-      return null;
+    // Try to extract partial facts from truncated JSON after all complete JSON
+    // candidates fail to parse.
+    log.debug("extractWithLocalLlm: JSON parse failed, attempting partial extraction...");
+    const partial = this.extractPartialFacts(content);
+    if (partial.facts.length > 0 || partial.entities.length > 0) {
+      log.debug(
+        `extractWithLocalLlm: extracted ${partial.facts.length} partial facts from truncated JSON`,
+      );
+      return partial;
     }
+    return null;
   }
 
   /**

package/src/fallback-llm.ts

@@ -7,7 +7,12 @@ import {
   buildChatCompletionTokenLimit,
   shouldAssumeOpenAiChatCompletions,
 } from "./openai-chat-compat.js";
-import { resolveProviderApiKey, getGatewayRuntimeAuthForModel } from "./resolve-provider-secret.js";
+import {
+  resolveProviderApiKey,
+  getGatewayRuntimeAuthForModel,
+  type GetRuntimeAuthForModelFn,
+  type ResolveApiKeyFn,
+} from "./resolve-provider-secret.js";
 import { loadModelsJsonProviders } from "./models-json.js";
 import { callCodexCliFallback } from "./codex-cli-fallback.js";
 import { resolveHomeDir } from "./runtime/env.js";
@@ -36,9 +41,32 @@ export interface FallbackLlmResponse {
 
 export interface FallbackLlmRuntimeContext {
   agentDir?: string;
+  getRuntimeAuthForModel?: GetRuntimeAuthForModelFn | null;
+  resolveApiKeyForProvider?: ResolveApiKeyFn | null;
   workspaceDir?: string;
 }
 
+export function fallbackLlmRuntimeContextFromConfig(
+  config: Pick<
+    GatewayBackedRuntimeConfig,
+    "providerApiKeyResolver" | "runtimeAuthForModelResolver" | "workspaceDir"
+  >,
+  overrides: FallbackLlmRuntimeContext = {},
+): FallbackLlmRuntimeContext {
+  return {
+    workspaceDir: config.workspaceDir,
+    resolveApiKeyForProvider: config.providerApiKeyResolver,
+    getRuntimeAuthForModel: config.runtimeAuthForModelResolver,
+    ...overrides,
+  };
+}
+
+type GatewayBackedRuntimeConfig = {
+  providerApiKeyResolver?: ResolveApiKeyFn | null;
+  runtimeAuthForModelResolver?: GetRuntimeAuthForModelFn | null;
+  workspaceDir?: string;
+};
+
 interface ModelRef {
   providerId: string;
   modelId: string;
@@ -488,7 +516,9 @@ export class FallbackLlmClient {
     model: ModelRef,
   ): Promise<{ apiKey?: string; baseUrl?: string } | null> {
     try {
-      const getRuntimeAuth = await getGatewayRuntimeAuthForModel();
+      const getRuntimeAuth = await getGatewayRuntimeAuthForModel({
+        getRuntimeAuthForModel: this.runtimeContext.getRuntimeAuthForModel,
+      });
       if (!getRuntimeAuth) return null;
 
       const result = await getRuntimeAuth({
@@ -526,6 +556,9 @@ export class FallbackLlmClient {
       model.providerConfig.apiKey,
       this.gatewayConfig,
       this.runtimeContext.agentDir,
+      {
+        resolveApiKeyForProvider: this.runtimeContext.resolveApiKeyForProvider,
+      },
     );
   }