@gmickel/gno 1.46.0 → 2.0.0

package/src/llm/nodeLlamaCpp/simulator-session.ts

@@ -0,0 +1,240 @@
+/** Guarded simulator section from node-llama-cpp PR636 (MIT, Gilad S.).
+ * See simulator-handle.ts and THIRD_PARTY_NOTICES.md for source attribution.
+ * Retained on 3.20.0 for failed-load cleanup and joined session disposal,
+ * beyond upstream PR636. Estimation inputs and selection remain unchanged.
+ */
+import type {
+  SimulatorBackend,
+  SimulatorCache,
+  SimulatorContextOptions,
+  SimulatorDependencies,
+  SimulatorDisposable,
+  SimulatorMemory,
+  SimulatorModelOptions,
+} from "./simulator-types";
+
+import { SimulatorModelHandle } from "./simulator-handle";
+
+function defined(options: Record<string, unknown>): Record<string, unknown> {
+  return Object.fromEntries(
+    Object.entries(options).filter(([, value]) => value !== undefined)
+  );
+}
+
+export class GuardedSimulatorSession {
+  private readonly models: SimulatorCache<Promise<SimulatorModelHandle>>;
+  private isDisposed = false;
+  private disposal?: Promise<void>;
+
+  constructor(
+    private readonly llama: SimulatorBackend,
+    private readonly dependencies: SimulatorDependencies,
+    lruCacheSize = 10,
+    private readonly loadModelError?: Error
+  ) {
+    this.models = new dependencies.LruCache(lruCacheSize, {
+      async onDelete(_key, value) {
+        try {
+          const handle = await value;
+          // Give queued consumers the same acquisition turn as upstream PR636.
+          await Promise.resolve();
+          await handle.dispose();
+        } catch {
+          /* Failed loads have no live cached handle. */
+        }
+      },
+    });
+  }
+
+  async estimateModelResources({
+    modelSource,
+    gpuLayers,
+    useMmap = false,
+  }: SimulatorModelOptions): Promise<SimulatorMemory> {
+    const handle = await this.getModel({ modelSource, gpuLayers, useMmap });
+    const lease = this.acquire(handle);
+    try {
+      const result = handle.model.getMemoryBreakdown();
+      if (this.llama._shouldLog(this.dependencies.debug))
+        this.llama._log(
+          this.dependencies.debug,
+          `Simulating model resource usage. gpuLayers=${gpuLayers} useMmap=${useMmap} memoryBreakdownCpuRam=${this.dependencies.bytes(result.cpuRam)} memoryBreakdownGpuVram=${this.dependencies.bytes(result.gpuVram)}`
+        );
+      return result;
+    } finally {
+      await lease.dispose();
+    }
+  }
+
+  async estimateContextResources({
+    modelSource,
+    gpuLayers,
+    contextSize,
+    batchSize,
+    sequences,
+    isEmbeddingContext = false,
+    flashAttention = "auto",
+    swaFullCache = false,
+    useMmap = false,
+    kvCacheKeyType = this.dependencies.f16,
+    kvCacheValueType = this.dependencies.f16,
+  }: SimulatorContextOptions): Promise<SimulatorMemory> {
+    const handle = await this.getModel({ modelSource, gpuLayers, useMmap });
+    const lease = this.acquire(handle);
+    try {
+      const context = new this.llama._bindings.AddonContext(
+        handle.model,
+        defined({
+          contextSize,
+          batchSize,
+          sequences,
+          embeddings: isEmbeddingContext,
+          flashAttention,
+          kvCacheKeyType,
+          kvCacheValueType,
+          swaFullCache,
+        })
+      );
+      try {
+        const loadingLock = this.dependencies.needInitLock(this.llama.gpu)
+          ? await this.dependencies.acquireLock([
+              this.llama._memoryLock,
+              this.dependencies.addonInit,
+            ])
+          : undefined;
+        const restoreLog = this.llama._createLogLevelOverride(
+          this.dependencies.error
+        );
+        try {
+          if (!(await context.init()))
+            throw new Error("Failed to create context");
+        } finally {
+          restoreLog();
+          await loadingLock?.dispose();
+        }
+        const result = context.getMemoryBreakdown();
+        if (this.llama._shouldLog(this.dependencies.debug))
+          this.llama._log(
+            this.dependencies.debug,
+            `Simulating context resource usage. gpuLayers=${gpuLayers} contextSize=${contextSize.toLocaleString("en-US", { notation: "compact" })} batchSize=${batchSize} sequences=${sequences} isEmbeddingContext=${isEmbeddingContext} flashAttention=${flashAttention} swaFullCache=${swaFullCache} kvCacheKeyType=${kvCacheKeyType} kvCacheValueType=${kvCacheValueType} useMmap=${useMmap} memoryBreakdownCpuRam=${this.dependencies.bytes(result.cpuRam)} memoryBreakdownGpuVram=${this.dependencies.bytes(result.gpuVram)}`
+          );
+        return result;
+      } finally {
+        await context.dispose();
+      }
+    } finally {
+      await lease.dispose();
+    }
+  }
+
+  get disposed(): boolean {
+    return this.isDisposed;
+  }
+  [Symbol.asyncDispose](): Promise<void> {
+    return this.dispose();
+  }
+
+  dispose(): Promise<void> {
+    this.disposal ??= this.disposeInternal();
+    return this.disposal;
+  }
+
+  private async disposeInternal(): Promise<void> {
+    this.isDisposed = true;
+    const pending = [...this.models.values()];
+    this.models.clear();
+    const handles = await Promise.all(
+      pending.map((value) => value.catch(() => undefined))
+    );
+    await Promise.all(handles.map((handle) => handle?.dispose()));
+  }
+
+  private acquire(handle: SimulatorModelHandle): SimulatorDisposable {
+    try {
+      return handle.disposeGuard.createPreventDisposalHandle();
+    } catch {
+      throw new Error("Model is disposed");
+    }
+  }
+
+  private async getModel({
+    modelSource,
+    gpuLayers,
+    useMmap = this.llama.supportsMmap,
+  }: SimulatorModelOptions): Promise<SimulatorModelHandle> {
+    if (this.isDisposed) throw new Error("simulator session is disposed");
+    if (this.loadModelError) throw this.loadModelError;
+    let backendHandle: SimulatorDisposable;
+    try {
+      backendHandle =
+        this.llama._backendDisposeGuard.createPreventDisposalHandle();
+    } catch {
+      throw new Error("Llama instance is disposed");
+    }
+    try {
+      const key = `${gpuLayers}:${useMmap}`;
+      const existing = this.models.get(key);
+      if (existing) return await existing;
+      if (this.llama._shouldLog(this.dependencies.debug))
+        this.llama._log(
+          this.dependencies.debug,
+          `Loading model for simulator session. gpuLayers=${gpuLayers} useMmap=${useMmap}`
+        );
+      const pending = this.loadModel({ modelSource, gpuLayers, useMmap });
+      this.models.set(key, pending);
+      try {
+        return await pending;
+      } catch (error) {
+        this.models.delete(key);
+        throw error;
+      }
+    } finally {
+      await backendHandle.dispose();
+    }
+  }
+
+  private async loadModel({
+    modelSource,
+    gpuLayers,
+    useMmap = false,
+  }: SimulatorModelOptions): Promise<SimulatorModelHandle> {
+    const model = new this.llama._bindings.AddonModel(
+      typeof modelSource === "string" ? modelSource : "",
+      {
+        gpuLayers,
+        noAlloc: true,
+        useMmap,
+        useMlock: false,
+      }
+    );
+    const loadingLock = this.dependencies.needInitLock(this.llama.gpu)
+      ? await this.dependencies.acquireLock([
+          this.llama._memoryLock,
+          this.dependencies.addonInit,
+        ])
+      : undefined;
+    const restoreLog = this.llama._createLogLevelOverride(
+      this.dependencies.error
+    );
+    try {
+      if (
+        !(await (typeof modelSource === "string"
+          ? model.init()
+          : model.init(modelSource)))
+      )
+        throw new Error("Failed to load model");
+    } catch (error) {
+      await model.dispose().catch(() => {});
+      throw error;
+    } finally {
+      restoreLog();
+      await loadingLock?.dispose();
+    }
+    try {
+      return new SimulatorModelHandle(this.llama, model, this.dependencies);
+    } catch (error) {
+      await model.dispose().catch(() => {});
+      throw error;
+    }
+  }
+}

package/src/llm/nodeLlamaCpp/simulator-types.ts

@@ -0,0 +1,80 @@
+/** Private compatibility shapes pinned to node-llama-cpp 3.20.0. */
+export interface SimulatorDisposable {
+  dispose(): void | Promise<void>;
+}
+export interface SimulatorGuard {
+  createPreventDisposalHandle(): SimulatorDisposable;
+  acquireDisposeLock(): Promise<void>;
+}
+export interface SimulatorMemory {
+  cpuRam: number;
+  gpuVram: number;
+}
+export interface SimulatorModel {
+  init(source?: unknown): Promise<boolean>;
+  getMemoryBreakdown(): SimulatorMemory;
+  dispose(): Promise<void>;
+}
+export interface SimulatorContext {
+  init(): Promise<boolean>;
+  getMemoryBreakdown(): SimulatorMemory;
+  dispose(): Promise<void>;
+}
+export interface SimulatorBackend {
+  gpu: string | false;
+  supportsMmap: boolean;
+  _memoryLock: object;
+  _backendDisposeGuard: SimulatorGuard;
+  _bindings: {
+    AddonModel: new (
+      source: string,
+      options: Record<string, unknown>
+    ) => SimulatorModel;
+    AddonContext: new (
+      model: SimulatorModel,
+      options: Record<string, unknown>
+    ) => SimulatorContext;
+  };
+  onDispose: {
+    createListener(callback: () => void | Promise<void>): SimulatorDisposable;
+  };
+  _createLogLevelOverride(level: string): () => void;
+  _shouldLog(level: string): boolean;
+  _log(level: string, message: string): void;
+}
+export interface SimulatorCache<T> {
+  get(key: string): T | undefined;
+  set(key: string, value: T): unknown;
+  delete(key: string): unknown;
+  clear(): void;
+  values(): IterableIterator<T>;
+}
+export interface SimulatorDependencies {
+  DisposeGuard: new (parents: SimulatorGuard[]) => SimulatorGuard;
+  LruCache: new <T>(
+    size: number,
+    options?: { onDelete(key: string, value: T): Promise<void> }
+  ) => SimulatorCache<T>;
+  acquireLock(keys: unknown[]): Promise<SimulatorDisposable>;
+  needInitLock(gpu: string | false): boolean;
+  addonInit: string;
+  debug: string;
+  error: string;
+  f16: number;
+  bytes(value: number): string;
+}
+export interface SimulatorModelOptions {
+  modelSource: unknown;
+  gpuLayers: number;
+  useMmap?: boolean;
+}
+export interface SimulatorContextOptions extends SimulatorModelOptions {
+  contextSize: number;
+  batchSize: number;
+  sequences: number;
+  isEmbeddingContext?: boolean;
+  flashAttention?: boolean | "auto";
+  swaFullCache?: boolean;
+  kvCacheKeyType?: number;
+  kvCacheValueType?: number;
+}

package/src/llm/types.ts

@@ -15,6 +15,15 @@ export type LlmResult<T> =
   | { ok: true; value: T }
   | { ok: false; error: LlmError };
 
+/** Optional caller lifetime; distinct from the configured native execution timeout.
+ * Adapters must retain native ownership until work settles, even after abort.
+ */
+export interface InferenceOptions {
+  signal?: AbortSignal;
+  /** Absolute Unix epoch milliseconds; includes admission, queue and model load. */
+  deadlineAt?: number;
+}
+
 // ─────────────────────────────────────────────────────────────────────────────
 // Model Types
 // ─────────────────────────────────────────────────────────────────────────────
@@ -81,14 +90,27 @@ export interface RerankScore {
 // Port Interfaces
 // ─────────────────────────────────────────────────────────────────────────────
 
+export interface EmbeddingIdentity {
+  contextSize: number;
+  /** Versioned policy including the effective token limit. */
+  truncationPolicy: string;
+  modelFingerprint: string;
+  runtimeFingerprint: string;
+}
+
 export interface EmbeddingPort {
   readonly modelUri: string;
   /** Initialize the embedding context (loads model). Call before dimensions(). */
-  init(): Promise<LlmResult<void>>;
-  embed(text: string): Promise<LlmResult<number[]>>;
-  embedBatch(texts: string[]): Promise<LlmResult<number[][]>>;
+  init(options?: InferenceOptions): Promise<LlmResult<void>>;
+  embed(text: string, options?: InferenceOptions): Promise<LlmResult<number[]>>;
+  embedBatch(
+    texts: string[],
+    options?: InferenceOptions
+  ): Promise<LlmResult<number[][]>>;
   /** Returns embedding dimensions. Must call init() first. */
   dimensions(): number;
+  /** Verified native identity after init; absent for unverified/HTTP backends. */
+  getIdentity?(): EmbeddingIdentity | undefined;
   dispose(): Promise<void>;
 }
 
@@ -96,13 +118,21 @@ export interface GenerationPort {
   readonly modelUri: string;
   /** Undefined is treated as unsupported for backwards-compatible ports. */
   readonly structuredOutput?: StructuredOutputCapability;
-  generate(prompt: string, params?: GenParams): Promise<LlmResult<string>>;
+  generate(
+    prompt: string,
+    params?: GenParams,
+    options?: InferenceOptions
+  ): Promise<LlmResult<string>>;
   dispose(): Promise<void>;
 }
 
 export interface RerankPort {
   readonly modelUri: string;
-  rerank(query: string, documents: string[]): Promise<LlmResult<RerankScore[]>>;
+  rerank(
+    query: string,
+    documents: string[],
+    options?: InferenceOptions
+  ): Promise<LlmResult<RerankScore[]>>;
   dispose(): Promise<void>;
 }
 

package/src/mcp/context.ts

@@ -11,6 +11,7 @@ import type {
 } from "../core/egress-policy";
 import type { EgressLineage } from "../core/egress-provenance";
 import type { JobManager } from "../core/job-manager";
+import type { NativeDisposeOptions } from "../llm/native-worker/owned-exit";
 import type { ModelLease } from "../llm/nodeLlamaCpp/lifecycle";
 import type { ResidentStatus } from "../serve/status-model";
 import type { SqliteAdapter } from "../store/sqlite/adapter";
@@ -18,6 +19,7 @@ import type { StoreResult } from "../store/types";
 import type { McpToolProfile } from "./tool-profile";
 
 import { MCP_SERVER_NAME, VERSION } from "../app/constants";
+import { LlmAdapter } from "../llm/nodeLlamaCpp/adapter";
 import { createStandaloneResidentStatus } from "../serve/resident-status";
 import { registerResources } from "./resources/index";
 import { registerTools } from "./tools/index";
@@ -86,6 +88,8 @@ export interface ToolContext {
   toolProfile?: McpToolProfile;
   isShuttingDown: () => boolean;
   getResidentStatus?: () => ResidentStatus;
+  getModelAdapter?: () => LlmAdapter;
+  disposeModels?: (options?: NativeDisposeOptions) => Promise<void>;
   acquireModelLease?: () => ModelLease;
   markContentMutation?: () => void;
   markIndexMutation?: () => void;
@@ -125,6 +129,8 @@ export interface CreateToolContextOptions {
   toolProfile?: McpToolProfile;
   isShuttingDown: () => boolean;
   getResidentStatus?: () => ResidentStatus;
+  getModelAdapter?: (config: Config) => LlmAdapter | undefined;
+  disposeModels?: (options?: NativeDisposeOptions) => Promise<void>;
   acquireModelLease?: () => ModelLease;
   markContentMutation?: () => void;
   markIndexMutation?: () => void;
@@ -141,6 +147,8 @@ export interface CreateToolContextOptions {
 export function createToolContext(
   options: CreateToolContextOptions
 ): ToolContext {
+  const adapters = new Map<Config, LlmAdapter>();
+  let modelsClosed = false;
   const requestSnapshot = new AsyncLocalStorage<ToolContextSnapshot>();
   const currentSnapshot = (): ToolContextSnapshot =>
     requestSnapshot.getStore() ??
@@ -176,6 +184,25 @@ export function createToolContext(
     getResidentStatus:
       options.getResidentStatus ??
       (() => createStandaloneResidentStatus("stdio")),
+    getModelAdapter: () => {
+      if (modelsClosed) throw new Error("Resident runtime is shutting down");
+      const config = currentSnapshot().config;
+      const shared = options.getModelAdapter?.(config);
+      if (shared) return shared;
+      let adapter = adapters.get(config);
+      if (!adapter) {
+        adapter = new LlmAdapter(config);
+        adapters.set(config, adapter);
+      }
+      return adapter;
+    },
+    disposeModels: async (disposeOptions) => {
+      modelsClosed = true;
+      await Promise.all(
+        [...adapters.values()].map((adapter) => adapter.dispose(disposeOptions))
+      );
+      adapters.clear();
+    },
     acquireModelLease: options.acquireModelLease,
     markContentMutation: options.markContentMutation,
     markIndexMutation: options.markIndexMutation,

package/src/mcp/http-transport.ts

@@ -1,9 +1,8 @@
-/** Web Standard Streamable HTTP request routing for the resident MCP runtime. */
-
 import {
   isInitializeRequest,
   type McpServer,
 } from "@modelcontextprotocol/server";
+/** Web Standard Streamable HTTP request routing for the resident MCP runtime. */
 
 import type { DestinationClassification } from "../core/destination-classifier";
 import type { ResidentRequestHandle } from "../serve/resident-runtime";
@@ -16,6 +15,7 @@ import type {
 
 import { MCP_SERVER_NAME, VERSION } from "../app/constants";
 import { EgressDeniedError } from "../core/egress-enforcement";
+import { withInferenceScope } from "../llm/inference-scope";
 import { createMcpServerSurface, type ToolContext } from "./context";
 import {
   enforceHttpMcpEgress,
@@ -416,14 +416,16 @@ export class HttpMcpTransport {
         throw error;
       }
       const handle = () =>
-        transport
-          ? transport.handleRequest(
-              request,
-              requestBody === undefined
-                ? undefined
-                : { parsedBody: requestBody }
-            )
-          : this.#modern.fetch(request, requestBody);
+        withInferenceScope({ signal: admission.signal }, async () =>
+          transport
+            ? transport.handleRequest(
+                request,
+                requestBody === undefined
+                  ? undefined
+                  : { parsedBody: requestBody }
+              )
+            : this.#modern.fetch(request, requestBody)
+        );
       const destinationZone =
         context.peerClassification?.zone ??
         (context.identity === "loopback" ? "loopback" : "remote");

package/src/mcp/server.ts

@@ -164,6 +164,8 @@ export async function startMcpServer(options: McpServerOptions): Promise<void> {
       // Best-effort - server may already be closed
     }
 
+    await ctx.disposeModels?.();
+
     // 4. Close DB (safe now - no tool or job is running)
     await store.close();
 
@@ -214,4 +216,5 @@ export async function startMcpServer(options: McpServerOptions): Promise<void> {
     });
     // Also resolve on SIGTERM/SIGINT (already handled by shutdown())
   });
+  await shutdown();
 }

package/src/mcp/tool-profile.ts

@@ -1,3 +1,11 @@
+import type {
+  Icon,
+  McpServer,
+  StandardSchemaWithJSON,
+  ToolAnnotations,
+  ToolCallback,
+  ServerContext,
+} from "@modelcontextprotocol/server";
 /**
  * MCP tool profiles.
  *
@@ -9,13 +17,7 @@
  * @module src/mcp/tool-profile
  */
 
-import type {
-  Icon,
-  McpServer,
-  StandardSchemaWithJSON,
-  ToolAnnotations,
-  ToolCallback,
-} from "@modelcontextprotocol/server";
+import { withInferenceScope } from "../llm/inference-scope";
 
 export const MCP_TOOL_PROFILES = ["core", "full"] as const;
 
@@ -105,8 +107,28 @@ export function createProfileToolRegistrar(
   profile: McpToolProfile
 ): ProfileToolRegistrar {
   const allowlist = mcpToolProfileAllowlist(profile);
+  const connection = new AbortController();
+  const previousClose = server.server.onclose;
+  server.server.onclose = () => {
+    connection.abort();
+    previousClose?.();
+  };
   return (name, config, cb) => {
     if (allowlist !== null && !allowlist.has(name)) return;
-    server.registerTool(name, config, cb);
+    // The SDK callback is conditional on schema presence: the last argument
+    // is always ServerContext, while a schema-bearing call also has input.
+    const wrapped = (...args: unknown[]) => {
+      const request = args.at(-1) as ServerContext;
+      return withInferenceScope(
+        {
+          signal: AbortSignal.any([connection.signal, request.mcpReq.signal]),
+        },
+        async () =>
+          args.length === 1
+            ? (cb as ToolCallback<undefined>)(request)
+            : (cb as ToolCallback<StandardSchemaWithJSON>)(args[0], request)
+      );
+    };
+    server.registerTool(name, config, wrapped as unknown as typeof cb);
   };
 }

package/src/mcp/tools/context.ts

@@ -37,7 +37,7 @@ import {
 import { LlmAdapter } from "../../llm/nodeLlamaCpp/adapter";
 import { resolveDownloadPolicy } from "../../llm/policy";
 import { resolveModelUri } from "../../llm/registry";
-import { createVectorIndexPort } from "../../store/vector";
+import { createLazyVectorIndex } from "../../store/vector/lazy";
 
 interface ContextToolResultData {
   structuredContent: Record<string, unknown>;
@@ -123,7 +123,7 @@ export const createMcpModelPorts = async (
   factoryOverride?: McpModelPortFactory,
   options: { generation?: boolean } = {}
 ): Promise<McpModelPorts> => {
-  const llm = new LlmAdapter(context.config);
+  const llm = context.getModelAdapter?.() ?? new LlmAdapter(context.config);
   const factory = factoryOverride ?? llm;
   const lease = factory.acquireModelLease?.();
   const policy = resolveDownloadPolicy(process.env, {});
@@ -152,15 +152,12 @@ export const createMcpModelPorts = async (
     if (embedResult.ok) {
       // Take ownership before init: init failures must not leak the port.
       ownedEmbedPort = embedResult.value;
-      const initialized = await ownedEmbedPort.init();
-      if (initialized.ok) {
-        embedPort = ownedEmbedPort;
-        const vectorResult = await createVectorIndexPort(
-          context.store.getRawDb(),
-          { model: embedUri, dimensions: ownedEmbedPort.dimensions() }
-        );
-        if (vectorResult.ok) vectorIndex = vectorResult.value;
-      }
+      embedPort = ownedEmbedPort;
+      vectorIndex = await createLazyVectorIndex(
+        context.store.getRawDb(),
+        embedUri,
+        ownedEmbedPort
+      );
     }
     const rerankResult = await factory.createRerankPort(
       resolveModelUri(context.config, "rerank", undefined, collection),

package/src/mcp/tools/embed.ts

@@ -85,7 +85,7 @@ export function handleEmbed(
             const lease = ctx.acquireModelLease?.();
             try {
               // Create LLM adapter with offline policy (fail-fast, no download)
-              const llm = new LlmAdapter(ctx.config);
+              const llm = ctx.getModelAdapter?.() ?? new LlmAdapter(ctx.config);
               const embedResult = await llm.createEmbeddingPort(modelUri, {
                 egressCollections: collection ? [collection.name] : "all",
                 policy: { offline: true, allowDownload: false },

package/src/mcp/tools/index-cmd.ts

@@ -148,7 +148,7 @@ export function handleIndex(
               recordContentMutation(syncResult, ctx.markContentMutation);
 
               // Phase 2: Embed
-              const llm = new LlmAdapter(ctx.config);
+              const llm = ctx.getModelAdapter?.() ?? new LlmAdapter(ctx.config);
               const embedResult = await llm.createEmbeddingPort(modelUri, {
                 egressCollections: collections,
                 policy: { offline: true, allowDownload: false },