@gmickel/gno 1.45.1 → 2.0.0

package/src/llm/nodeLlamaCpp/lifecycle.ts

@@ -1,33 +1,29 @@
-/**
- * Model lifecycle manager.
- * Handles lazy loading, caching, and disposal of LLM models.
- *
- * @module src/llm/nodeLlamaCpp/lifecycle
- */
-
-import type { LlamaOptions } from "node-llama-cpp";
-
-import { platform } from "node:os";
+import type { LlamaInitOptions } from "./lifecycle-options";
+
+import {
+  resolveLlamaGpuMode,
+  resolveLlamaBuildMode,
+  resolveLlamaBackendInitTimeoutMs,
+  shouldRetryLlamaWithCpu,
+} from "./lifecycle-options";
+export {
+  resolveLlamaGpuMode,
+  resolveLlamaBuildMode,
+  resolveLlamaBackendInitTimeoutMs,
+  shouldRetryLlamaWithCpu,
+} from "./lifecycle-options";
+export type { LlamaGpuMode, LlamaBuildMode } from "./lifecycle-options";
 
 import type { ModelConfig } from "../../config/types";
 import type { LlmResult, LoadedModel, ModelType } from "../types";
 
 import { loadFailedError, outOfMemoryError, timeoutError } from "../errors";
+import { installSimulatorLifetimeGuard } from "./simulator-install";
 
-// ─────────────────────────────────────────────────────────────────────────────
 // Types
-// ─────────────────────────────────────────────────────────────────────────────
 
 type Llama = Awaited<ReturnType<typeof import("node-llama-cpp").getLlama>>;
 type LlamaModel = Awaited<ReturnType<Llama["loadModel"]>>;
-export type LlamaGpuMode = "auto" | "metal" | "vulkan" | "cuda" | false;
-export type LlamaBuildMode = "never" | "autoAttempt";
-
-type LlamaInitOptions = LlamaOptions & {
-  build: LlamaBuildMode;
-  gpu: LlamaGpuMode;
-};
-
 interface CachedModel {
   uri: string;
   type: ModelType;
@@ -50,100 +46,14 @@ export interface ModelLifecycleStats {
   inflightLoads: number;
 }
 
-let invalidGpuModeWarned = false;
-let invalidBuildModeWarned = false;
 let gpuFallbackWarned = false;
 let backendTimeoutWarned = false;
 
-const DEFAULT_BACKEND_INIT_TIMEOUT_MS = 30_000;
-
-export function resolveLlamaGpuMode(
-  env: NodeJS.ProcessEnv = process.env
-): LlamaGpuMode {
-  const raw = (env.GNO_LLAMA_GPU ?? env.NODE_LLAMA_CPP_GPU ?? "auto")
-    .trim()
-    .toLowerCase();
-  if (!raw || raw === "auto") {
-    return "auto";
-  }
-  if (raw === "metal" || raw === "vulkan" || raw === "cuda") {
-    return raw;
-  }
-  if (
-    raw === "false" ||
-    raw === "off" ||
-    raw === "none" ||
-    raw === "disable" ||
-    raw === "disabled" ||
-    raw === "0"
-  ) {
-    return false;
-  }
-  if (!invalidGpuModeWarned) {
-    invalidGpuModeWarned = true;
-    console.warn(
-      `[llama] Invalid GNO_LLAMA_GPU/NODE_LLAMA_CPP_GPU value "${raw}", using auto`
-    );
-  }
-  return "auto";
-}
-
-export function resolveLlamaBuildMode(
-  env: NodeJS.ProcessEnv = process.env
-): LlamaBuildMode {
-  const raw = (env.GNO_LLAMA_BUILD ?? "never").trim().toLowerCase();
-  if (
-    !raw ||
-    raw === "never" ||
-    raw === "prebuilt" ||
-    raw === "prebuilt-only"
-  ) {
-    return "never";
-  }
-  if (
-    raw === "autoattempt" ||
-    raw === "auto-attempt" ||
-    raw === "source" ||
-    raw === "build"
-  ) {
-    return "autoAttempt";
-  }
-  if (!invalidBuildModeWarned) {
-    invalidBuildModeWarned = true;
-    console.warn(`[llama] Invalid GNO_LLAMA_BUILD value "${raw}", using never`);
-  }
-  return "never";
-}
-
-export function resolveLlamaBackendInitTimeoutMs(
-  env: NodeJS.ProcessEnv = process.env
-): number {
-  const raw = env.GNO_LLAMA_INIT_TIMEOUT_MS;
-  if (!raw) {
-    return DEFAULT_BACKEND_INIT_TIMEOUT_MS;
-  }
-  const parsed = Number.parseInt(raw, 10);
-  return Number.isFinite(parsed) && parsed > 0
-    ? parsed
-    : DEFAULT_BACKEND_INIT_TIMEOUT_MS;
-}
-
-export function shouldRetryLlamaWithCpu(
-  gpu: LlamaGpuMode,
-  platformName = platform()
-): boolean {
-  if (gpu === false) {
-    return false;
-  }
-  return gpu !== "auto" || platformName === "win32";
-}
-
-// ─────────────────────────────────────────────────────────────────────────────
-// ModelManager
-// ─────────────────────────────────────────────────────────────────────────────
-
 export class ModelManager {
   private llama: Llama | null = null;
+  private llamaInit: Promise<Llama> | null = null;
+  private closing: Promise<void> | null = null;
+  private readonly lateCleanup = new Set<Promise<void>>();
   private readonly models: Map<string, CachedModel> = new Map();
   private readonly disposalTimers: Map<string, ReturnType<typeof setTimeout>> =
     new Map();
@@ -152,22 +62,38 @@ export class ModelManager {
   private readonly leaseDrainWaiters = new Set<() => void>();
   private readonly config: ModelConfig;
   private activeLeases = 0;
+  private broadLeases = 0;
+  private readonly modelLeases = new Map<string, number>();
+  private readonly expiresAt = new Map<string, number>();
+  private readonly retiringModels = new Map<string, Promise<void>>();
   private leaseAcquisitions = 0;
   private leaseReleases = 0;
   private loadAttempts = 0;
   private loadSuccesses = 0;
   private loadFailures = 0;
 
-  constructor(config: ModelConfig) {
+  constructor(
+    config: ModelConfig,
+    private readonly awaitNativeLoadSettlement = false
+  ) {
     this.config = config;
   }
 
-  /**
-   * Get or initialize the Llama instance.
-   * Uses lazy loading - only imports node-llama-cpp on first use.
-   */
   async getLlama(): Promise<Llama> {
+    if (this.closing) throw new Error("Model manager is disposing");
+    if (this.llama) return this.llama;
+    this.llamaInit ??= this.initializeLlama().finally(() => {
+      this.llamaInit = null;
+    });
+    return this.llamaInit;
+  }
+
+  private async initializeLlama(): Promise<Llama> {
+    // A timed-out native constructor still owns work until its loser is cleaned.
+    await Promise.allSettled(this.lateCleanup);
+    if (this.closing) throw new Error("Model manager is disposing");
     if (!this.llama) {
+      await installSimulatorLifetimeGuard();
       const { getLlama, LlamaLogLevel } = await import("node-llama-cpp");
       const gpu = resolveLlamaGpuMode();
       const build = resolveLlamaBuildMode();
@@ -184,7 +110,11 @@ export class ModelManager {
           timeoutMs
         );
       } catch (error) {
-        if (!shouldRetryLlamaWithCpu(gpu)) {
+        if (
+          this.closing ||
+          (error instanceof Error && error.name === "TimeoutError") ||
+          !shouldRetryLlamaWithCpu(gpu)
+        ) {
           throw error;
         }
         if (!gpuFallbackWarned) {
@@ -206,6 +136,12 @@ export class ModelManager {
         );
       }
     }
+    if (this.closing) {
+      const llama = this.llama;
+      this.llama = null;
+      await llama.dispose();
+      throw new Error("Model manager disposed during initialization");
+    }
     return this.llama;
   }
 
@@ -215,21 +151,46 @@ export class ModelManager {
     timeoutMs: number
   ): Promise<Llama> {
     let timeoutId: ReturnType<typeof setTimeout> | null = null;
+    let expired = false;
+    const initializationDeadline = performance.now() + timeoutMs;
+    const initialization = getLlama(options).then(async (llama) => {
+      if (expired || performance.now() >= initializationDeadline) {
+        await llama.dispose();
+        throw new DOMException(
+          "Backend initialization expired",
+          "TimeoutError"
+        );
+      }
+      return llama;
+    });
+    const cleanup = initialization.then(
+      () => {},
+      () => {}
+    );
+    this.lateCleanup.add(cleanup);
+    void cleanup.finally(() => this.lateCleanup.delete(cleanup));
     try {
-      return await Promise.race([
-        getLlama(options),
-        new Promise<never>((_, reject) => {
-          timeoutId = setTimeout(() => {
-            if (!backendTimeoutWarned) {
-              backendTimeoutWarned = true;
-              console.warn(
-                `[llama] Backend initialization timed out after ${timeoutMs}ms`
-              );
-            }
-            reject(new Error(`Backend init timeout after ${timeoutMs}ms`));
-          }, timeoutMs);
-        }),
-      ]);
+      const deadline = new Promise<never>((_, reject) => {
+        timeoutId = setTimeout(() => {
+          expired = true;
+          if (!backendTimeoutWarned) {
+            backendTimeoutWarned = true;
+            console.warn(
+              `[llama] Backend initialization timed out after ${timeoutMs}ms`
+            );
+          }
+          const error = new Error(`Backend init timeout after ${timeoutMs}ms`);
+          error.name = "TimeoutError";
+          reject(error);
+        }, timeoutMs);
+      });
+      // Native child completion must wait for actual backend settlement even
+      // after the init budget expires. Expiry still rejects and forbids fallback.
+      if (this.awaitNativeLoadSettlement) {
+        void deadline.catch(() => {});
+        return await initialization;
+      }
+      return await Promise.race([initialization, deadline]);
     } finally {
       if (timeoutId) {
         clearTimeout(timeoutId);
@@ -237,19 +198,23 @@ export class ModelManager {
     }
   }
 
-  /**
-   * Load a model by path.
-   * Uses caching, inflight deduplication, and TTL-based disposal.
-   */
   loadModel(
     modelPath: string,
     uri: string,
-    type: ModelType
+    type: ModelType,
+    signal?: AbortSignal
   ): Promise<LlmResult<LoadedModel>> {
+    const retiring = this.retiringModels.get(uri);
+    if (retiring)
+      return retiring.then(() => this.loadModel(modelPath, uri, type, signal));
+    if (this.closing)
+      return Promise.resolve({
+        ok: false,
+        error: loadFailedError(uri, new Error("Model manager is disposing")),
+      });
     // Check cache first
     const cached = this.models.get(uri);
     if (cached) {
-      this.resetDisposalTimer(uri);
       return Promise.resolve({
         ok: true as const,
         value: {
@@ -268,50 +233,76 @@ export class ModelManager {
     }
 
     // Start new load with cleanup
-    const loadPromise = this.loadModelInternal(modelPath, uri, type).finally(
-      () => {
-        this.inflightLoads.delete(uri);
-      }
-    );
+    const loadPromise = this.loadModelInternal(
+      modelPath,
+      uri,
+      type,
+      signal
+    ).finally(() => {
+      this.inflightLoads.delete(uri);
+    });
     this.inflightLoads.set(uri, loadPromise);
     return loadPromise;
   }
 
-  /**
-   * Internal model loading with timeout handling.
-   */
   private async loadModelInternal(
     modelPath: string,
     uri: string,
-    type: ModelType
+    type: ModelType,
+    signal?: AbortSignal
   ): Promise<LlmResult<LoadedModel>> {
     this.loadAttempts += 1;
     const timeoutMs = this.config.loadTimeout;
     let timeoutId: ReturnType<typeof setTimeout> | null = null;
     let timedOut = false;
+    const controller = new AbortController();
+    const loadSignal = signal
+      ? AbortSignal.any([signal, controller.signal])
+      : controller.signal;
 
     // Capture loadPromise outside try block so we can dispose it on timeout
     let loadPromise: Promise<LlamaModel> | null = null;
 
     try {
       const llama = await this.getLlama();
-      loadPromise = llama.loadModel({ modelPath });
+      loadSignal.throwIfAborted();
+      const loadDeadline = performance.now() + timeoutMs;
+      loadPromise = llama.loadModel({ modelPath, loadSignal });
 
       // Create timeout with proper cleanup
       const timeoutPromise = new Promise<never>((_, reject) => {
         timeoutId = setTimeout(() => {
           timedOut = true;
+          controller.abort(
+            new DOMException("Model load timed out", "TimeoutError")
+          );
           reject(new Error(`Load timeout after ${timeoutMs}ms`));
         }, timeoutMs);
       });
 
-      const model = await Promise.race([loadPromise, timeoutPromise]);
+      // The isolated child cannot report completion while a noncooperative
+      // loader still owns native allocations. Parent cancellation delivers early.
+      if (this.awaitNativeLoadSettlement) void timeoutPromise.catch(() => {});
+      const model = this.awaitNativeLoadSettlement
+        ? await loadPromise
+        : await Promise.race([loadPromise, timeoutPromise]);
+      if (this.awaitNativeLoadSettlement && performance.now() >= loadDeadline)
+        timedOut = true;
+      if (this.awaitNativeLoadSettlement && (timedOut || loadSignal.aborted)) {
+        await model.dispose();
+        if (timedOut) throw new Error("Model load timeout");
+        loadSignal.throwIfAborted();
+      }
 
       // Clear timeout on success
       if (timeoutId) {
         clearTimeout(timeoutId);
       }
 
+      if (this.closing) {
+        await model.dispose();
+        throw new Error("Model manager disposed during model loading");
+      }
       const now = Date.now();
       const cachedModel: CachedModel = {
         uri,
@@ -321,6 +312,7 @@ export class ModelManager {
       };
 
       this.models.set(uri, cachedModel);
+      this.expiresAt.set(uri, Date.now() + this.config.warmModelTtl);
       this.setDisposalTimer(uri);
       this.loadSuccesses += 1;
 
@@ -341,11 +333,11 @@ export class ModelManager {
       }
 
       // Dispose late-arriving model after timeout to prevent memory leak
-      if (timedOut && loadPromise) {
-        loadPromise.then(
+      if (timedOut && loadPromise && !this.awaitNativeLoadSettlement) {
+        const cleanup = loadPromise.then(
           (model) => {
             // Dispose model that arrived after timeout
-            model.dispose().catch(() => {
+            return model.dispose().catch(() => {
               // Ignore dispose errors
             });
           },
@@ -353,10 +345,16 @@ export class ModelManager {
             // Ignore load errors after timeout
           }
         );
+        this.lateCleanup.add(cleanup);
+        void cleanup.finally(() => this.lateCleanup.delete(cleanup));
       }
 
       if (e instanceof Error) {
-        if (e.message.includes("timeout")) {
+        if (
+          timedOut ||
+          e.name === "TimeoutError" ||
+          e.message.includes("timeout")
+        ) {
           return {
             ok: false,
             error: timeoutError(uri, "load", this.config.loadTimeout),
@@ -370,23 +368,23 @@ export class ModelManager {
     }
   }
 
-  /**
-   * Get a loaded model by URI (no loading).
-   */
   getLoadedModel(uri: string): CachedModel | undefined {
-    const model = this.models.get(uri);
-    if (model) {
-      this.resetDisposalTimer(uri);
-    }
-    return model;
+    if (this.closing) return;
+    return this.models.get(uri);
   }
 
-  /** Keep warm models alive while one request owns model-backed ports. */
-  acquireLease(): ModelLease {
+  acquireLease(uri?: string, activity = true): ModelLease {
+    if (this.closing) throw new Error("Model manager is disposing");
     this.activeLeases += 1;
     this.leaseAcquisitions += 1;
-    for (const timer of this.disposalTimers.values()) clearTimeout(timer);
-    this.disposalTimers.clear();
+    if (uri === undefined) this.broadLeases += 1;
+    else this.modelLeases.set(uri, (this.modelLeases.get(uri) ?? 0) + 1);
+    for (const [key, timer] of this.disposalTimers) {
+      if (uri === undefined || key === uri) {
+        clearTimeout(timer);
+        this.disposalTimers.delete(key);
+      }
+    }
 
     let released = false;
     return {
@@ -395,10 +393,22 @@ export class ModelManager {
         released = true;
         this.activeLeases = Math.max(0, this.activeLeases - 1);
         this.leaseReleases += 1;
+        if (uri === undefined) this.broadLeases -= 1;
+        else {
+          const count = (this.modelLeases.get(uri) ?? 1) - 1;
+          if (count) this.modelLeases.set(uri, count);
+          else this.modelLeases.delete(uri);
+        }
+        for (const key of this.models.keys()) {
+          if (uri === undefined || key === uri) {
+            if (activity)
+              this.expiresAt.set(key, Date.now() + this.config.warmModelTtl);
+            this.resetDisposalTimer(key);
+          }
+        }
         if (this.activeLeases === 0) {
           for (const resolve of this.leaseDrainWaiters) resolve();
           this.leaseDrainWaiters.clear();
-          for (const uri of this.models.keys()) this.setDisposalTimer(uri);
         }
       },
     };
@@ -417,21 +427,14 @@ export class ModelManager {
     };
   }
 
-  /**
-   * Check if a model is loaded.
-   */
   isLoaded(uri: string): boolean {
     return this.models.has(uri);
   }
 
-  /**
-   * Dispose a specific model.
-   */
   async dispose(uri: string): Promise<void> {
-    if (this.activeLeases > 0) {
-      this.resetDisposalTimer(uri);
-      return;
-    }
+    if (this.broadLeases > 0 || this.modelLeases.has(uri)) return;
+    const retiring = this.retiringModels.get(uri);
+    if (retiring) return retiring;
     const cached = this.models.get(uri);
     if (!cached) {
       return;
@@ -444,20 +447,29 @@ export class ModelManager {
       this.disposalTimers.delete(uri);
     }
 
-    // Dispose the model
-    try {
-      await cached.model.dispose();
-    } catch {
-      // Ignore disposal errors
-    }
-
+    // Remove ownership before yielding so acquisitions cannot see a retiring model.
     this.models.delete(uri);
+    this.expiresAt.delete(uri);
+    const cleanup = cached.model
+      .dispose()
+      .catch(() => {})
+      .finally(() => {
+        this.retiringModels.delete(uri);
+      });
+    this.retiringModels.set(uri, cleanup);
+    this.lateCleanup.add(cleanup);
+    await cleanup;
+    this.lateCleanup.delete(cleanup);
+  }
+
+  disposeAll(): Promise<void> {
+    this.closing ??= this.disposeAllInternal().finally(() => {
+      this.closing = null;
+    });
+    return this.closing;
   }
 
-  /**
-   * Dispose all loaded models.
-   */
-  async disposeAll(): Promise<void> {
+  private async disposeAllInternal(): Promise<void> {
     if (this.activeLeases > 0) {
       await new Promise<void>((resolve) => this.leaseDrainWaiters.add(resolve));
     }
@@ -479,13 +491,13 @@ export class ModelManager {
       this.models.delete(uri);
     }
 
-    // Clear llama instance
+    await this.llamaInit?.catch(() => {});
+    await Promise.allSettled(this.lateCleanup);
+    const llama = this.llama;
     this.llama = null;
+    await llama?.dispose();
   }
 
-  /**
-   * Get list of loaded models.
-   */
   getLoadedModels(): Array<{ uri: string; type: ModelType; loadedAt: number }> {
     return Array.from(this.models.values()).map((m) => ({
       uri: m.uri,
@@ -494,17 +506,18 @@ export class ModelManager {
     }));
   }
 
-  // ───────────────────────────────────────────────────────────────────────────
   // Private
-  // ───────────────────────────────────────────────────────────────────────────
 
   private setDisposalTimer(uri: string): void {
-    if (this.activeLeases > 0) return;
-    const timer = setTimeout(() => {
-      this.dispose(uri).catch(() => {
-        // Ignore disposal errors in timer callback
-      });
-    }, this.config.warmModelTtl);
+    if (this.broadLeases > 0 || this.modelLeases.has(uri)) return;
+    const timer = setTimeout(
+      () => {
+        this.dispose(uri).catch(() => {
+          // Ignore disposal errors in timer callback
+        });
+      },
+      Math.max(0, (this.expiresAt.get(uri) ?? Date.now()) - Date.now())
+    );
 
     // Allow CLI processes to exit without waiting for TTL timer
     if (typeof timer.unref === "function") {
@@ -523,15 +536,10 @@ export class ModelManager {
   }
 }
 
-// ─────────────────────────────────────────────────────────────────────────────
 // Singleton
-// ─────────────────────────────────────────────────────────────────────────────
 
 let defaultManager: ModelManager | null = null;
 
-/**
- * Get the default ModelManager singleton.
- */
 export function getModelManager(config: ModelConfig): ModelManager {
   if (!defaultManager) {
     defaultManager = new ModelManager(config);
@@ -539,9 +547,6 @@ export function getModelManager(config: ModelConfig): ModelManager {
   return defaultManager;
 }
 
-/**
- * Reset the default manager (for testing).
- */
 export async function resetModelManager(): Promise<void> {
   if (defaultManager) {
     await defaultManager.disposeAll();