fix(llama-cpp): dispose runtime on plugin stop

extensions/llama-cpp/index.test.ts

@@ -1,6 +1,7 @@
 import os from "node:os";
 import path from "node:path";
 import { expectDefined } from "@openclaw/normalization-core";
+import type { OpenClawPluginService } from "openclaw/plugin-sdk/plugin-entry";
 import { createTestPluginApi } from "openclaw/plugin-sdk/plugin-test-api";
 import {
   createPluginRegistryFixture,
@@ -116,6 +117,30 @@ afterEach(() => {
 });
 
 describe("llama.cpp provider plugin", () => {
+  it("registers process-owned inference cleanup as a plugin service", async () => {
+    const services: OpenClawPluginService[] = [];
+    llamaCppPlugin.register(
+      createTestPluginApi({
+        id: "llama-cpp",
+        name: "llama.cpp Provider",
+        source: "test",
+        config: {},
+        pluginConfig: {},
+        runtime: {} as never,
+        registerService: (service) => services.push(service),
+      }),
+    );
+
+    expect(services).toEqual([
+      expect.objectContaining({
+        id: "llama-cpp-inference-runtime",
+        start: expect.any(Function),
+        stop: expect.any(Function),
+      }),
+    ]);
+    await services[0]?.stop?.({} as never);
+  });
+
   it("registers the local text-inference provider", () => {
     expect(registerLlamaCppTextProvider()).toEqual(
       expect.objectContaining({

extensions/llama-cpp/index.ts

@@ -8,7 +8,7 @@ import {
   resolveLlamaCppSyntheticApiKey,
 } from "./src/defaults.js";
 import { llamaCppEmbeddingProviderAdapter } from "./src/embedding-provider.js";
-import { createLlamaCppStreamFn } from "./src/inference-provider.js";
+import { createLlamaCppInferenceRuntime } from "./src/inference-provider.js";
 import { detectLlamaCppSetup, prepareLlamaCppSetup, runLlamaCppSetup } from "./src/setup.js";
 
 export default definePluginEntry({
@@ -16,6 +16,12 @@ export default definePluginEntry({
   name: "llama.cpp Provider",
   description: "Local GGUF text inference and embeddings through node-llama-cpp",
   register(api: OpenClawPluginApi) {
+    const inferenceRuntime = createLlamaCppInferenceRuntime();
+    api.registerService({
+      id: "llama-cpp-inference-runtime",
+      start: () => undefined,
+      stop: () => inferenceRuntime.dispose(),
+    });
     api.registerEmbeddingProvider(llamaCppEmbeddingProviderAdapter);
     api.registerProvider({
       id: LLAMA_CPP_PROVIDER_ID,
@@ -51,7 +57,7 @@ export default definePluginEntry({
         if (model.baseUrl !== LLAMA_CPP_LOCAL_BASE_URL) {
           return undefined;
         }
-        return createLlamaCppStreamFn({
+        return inferenceRuntime.createStreamFn({
           providerConfig: config?.models?.providers?.[provider],
         });
       },

extensions/llama-cpp/src/inference-provider.test.ts

@@ -54,19 +54,17 @@ vi.mock("node-llama-cpp", () => ({
   },
 }));
 
-import { createLlamaCppStreamFn } from "./inference-provider.js";
+import {
+  createLlamaCppInferenceRuntime,
+  llamaCppInferenceTestApi,
+  type LlamaCppInferenceRuntime,
+} from "./inference-provider.js";
 
-const {
-  clearLlamaCppInferenceCacheForTests,
-  mapContextToLlamaChatHistory,
-  mapToolsToLlamaFunctions,
-} = (globalThis as Record<PropertyKey, unknown>)[
-  Symbol.for("openclaw.llamaCppInferenceTestApi")
-] as {
-  clearLlamaCppInferenceCacheForTests: () => Promise<void>;
-  mapContextToLlamaChatHistory: (context: Context) => unknown[];
-  mapToolsToLlamaFunctions: (context: Context) => Record<string, unknown> | undefined;
-};
+if (!llamaCppInferenceTestApi) {
+  throw new Error("expected llama.cpp inference test API");
+}
+const { mapContextToLlamaChatHistory, mapToolsToLlamaFunctions } = llamaCppInferenceTestApi;
+let inferenceRuntime: LlamaCppInferenceRuntime;
 
 const model: Model = {
   id: "test.gguf",
@@ -113,11 +111,11 @@ type TestStreamParams = {
   selectedModel?: Model;
   prompt?: string;
   tools?: Context["tools"];
-  options?: Parameters<ReturnType<typeof createLlamaCppStreamFn>>[2];
+  options?: Parameters<ReturnType<LlamaCppInferenceRuntime["createStreamFn"]>>[2];
 };
 
 async function createTestStream(params: TestStreamParams = {}) {
-  return await createLlamaCppStreamFn({})(
+  return await inferenceRuntime.createStreamFn({})(
     params.selectedModel ?? model,
     {
       messages: [{ role: "user", content: params.prompt ?? "Hi", timestamp: 1 }],
@@ -131,8 +129,24 @@ async function collectTestEvents(params: TestStreamParams = {}) {
   return await collectEvents(await createTestStream(params));
 }
 
-beforeEach(async () => {
-  await clearLlamaCppInferenceCacheForTests();
+function deferGeneration() {
+  let finishGeneration: (() => void) | undefined;
+  mocks.generateResponse.mockImplementationOnce(
+    async () =>
+      await new Promise((resolve) => {
+        finishGeneration = () =>
+          resolve({
+            response: "",
+            functionCalls: undefined,
+            metadata: { stopReason: "eogToken" },
+          });
+      }),
+  );
+  return () => finishGeneration?.();
+}
+
+beforeEach(() => {
+  inferenceRuntime = createLlamaCppInferenceRuntime();
   vi.clearAllMocks();
   mocks.generateResponse.mockResolvedValue({
     response: "",
@@ -142,7 +156,7 @@ beforeEach(async () => {
 });
 
 afterEach(async () => {
-  await clearLlamaCppInferenceCacheForTests();
+  await inferenceRuntime.dispose();
 });
 
 describe("llama.cpp inference provider", () => {
@@ -836,7 +850,7 @@ describe("llama.cpp inference provider", () => {
   });
 
   it("disposes the previous model and context when the model changes", async () => {
-    const streamFn = createLlamaCppStreamFn({});
+    const streamFn = inferenceRuntime.createStreamFn({});
     await collectEvents(
       await streamFn(model, { messages: [{ role: "user", content: "one", timestamp: 1 }] }),
     );
@@ -853,7 +867,7 @@ describe("llama.cpp inference provider", () => {
   });
 
   it("reuses one context sequence across serialized requests for the same model", async () => {
-    const streamFn = createLlamaCppStreamFn({});
+    const streamFn = inferenceRuntime.createStreamFn({});
     await collectEvents(
       await streamFn(model, { messages: [{ role: "user", content: "one", timestamp: 1 }] }),
     );
@@ -865,6 +879,84 @@ describe("llama.cpp inference provider", () => {
     expect(mocks.llama.loadModel).toHaveBeenCalledTimes(1);
   });
 
+  it("disposes the context, model, and native runtime in ownership order", async () => {
+    await collectTestEvents();
+
+    await inferenceRuntime.dispose();
+
+    expect(mocks.contextDispose).toHaveBeenCalledOnce();
+    expect(mocks.modelDispose).toHaveBeenCalledOnce();
+    expect(mocks.llamaDispose).toHaveBeenCalledOnce();
+    expect(mocks.contextDispose.mock.invocationCallOrder[0]).toBeLessThan(
+      mocks.modelDispose.mock.invocationCallOrder[0] ?? 0,
+    );
+    expect(mocks.modelDispose.mock.invocationCallOrder[0]).toBeLessThan(
+      mocks.llamaDispose.mock.invocationCallOrder[0] ?? 0,
+    );
+  });
+
+  it("waits for admitted inference before disposing the runtime", async () => {
+    const finishGeneration = deferGeneration();
+    const stream = await createTestStream();
+    await vi.waitFor(() => expect(mocks.generateResponse).toHaveBeenCalledOnce());
+
+    const disposing = inferenceRuntime.dispose();
+    await Promise.resolve();
+    expect(mocks.contextDispose).not.toHaveBeenCalled();
+
+    finishGeneration();
+    await stream.result();
+    await disposing;
+
+    expect(mocks.contextDispose).toHaveBeenCalledOnce();
+    expect(mocks.modelDispose).toHaveBeenCalledOnce();
+    expect(mocks.llamaDispose).toHaveBeenCalledOnce();
+  });
+
+  it("rejects new inference once runtime disposal begins", async () => {
+    const finishGeneration = deferGeneration();
+    const activeStream = await createTestStream();
+    await vi.waitFor(() => expect(mocks.generateResponse).toHaveBeenCalledOnce());
+
+    const disposing = inferenceRuntime.dispose();
+    const rejectedStream = await createTestStream({ prompt: "too late" });
+
+    await expect(rejectedStream.result()).resolves.toMatchObject({
+      stopReason: "error",
+      errorMessage: "llama.cpp runtime is stopping",
+    });
+    expect(mocks.generateResponse).toHaveBeenCalledOnce();
+
+    finishGeneration();
+    await activeStream.result();
+    await disposing;
+  });
+
+  it("shares concurrent runtime disposal and performs cleanup once", async () => {
+    await collectTestEvents();
+
+    const disposals = [inferenceRuntime.dispose(), inferenceRuntime.dispose()];
+    expect(disposals[1]).toBe(disposals[0]);
+    await Promise.all(disposals);
+    expect(mocks.contextDispose).toHaveBeenCalledOnce();
+    expect(mocks.modelDispose).toHaveBeenCalledOnce();
+    expect(mocks.llamaDispose).toHaveBeenCalledOnce();
+  });
+
+  it("retains failed cleanup ownership so disposal can be retried", async () => {
+    await collectTestEvents();
+    mocks.contextDispose.mockRejectedValueOnce(new Error("context cleanup failed"));
+
+    await expect(inferenceRuntime.dispose()).rejects.toThrow("context cleanup failed");
+    expect([mocks.modelDispose.mock.calls.length, mocks.llamaDispose.mock.calls.length]).toEqual([
+      0, 0,
+    ]);
+    await inferenceRuntime.dispose();
+    expect(mocks.contextDispose).toHaveBeenCalledTimes(2);
+    expect(mocks.modelDispose).toHaveBeenCalledOnce();
+    expect(mocks.llamaDispose).toHaveBeenCalledOnce();
+  });
+
   it.each([
     {
       scenario: "a smaller advertised model window",
@@ -964,7 +1056,7 @@ describe("llama.cpp inference provider", () => {
           resolveFirst = resolve;
         }),
     );
-    const streamFn = createLlamaCppStreamFn({});
+    const streamFn = inferenceRuntime.createStreamFn({});
     const firstStream = await streamFn(model, {
       messages: [{ role: "user", content: "first", timestamp: 1 }],
     });

extensions/llama-cpp/src/inference-provider.ts

@@ -41,11 +41,18 @@ type LoadedModel = {
 
 type LlamaJsonSchemaInput = Parameters<Llama["createGrammarForJsonSchema"]>[0];
 
-// Process-owned, single-slot cache. A model/context pair lives until another
-// model replaces it or the process exits, bounding resident model memory.
-let loadedModel: LoadedModel | undefined;
-let llamaInstance: Llama | undefined;
-let operationQueue: Promise<void> = Promise.resolve();
+type LlamaCppInferenceRuntimeState = {
+  loadedModel?: LoadedModel;
+  llamaInstance?: Llama;
+  operationQueue: Promise<void>;
+  lifecycle: "open" | "closing" | "closed";
+  disposePromise?: Promise<void>;
+};
+
+export type LlamaCppInferenceRuntime = {
+  createStreamFn: (params: { providerConfig?: ModelProviderConfig }) => StreamFn;
+  dispose: () => Promise<void>;
+};
 
 function zeroCostUsage(input = 0, output = 0): Usage {
   return {
@@ -235,17 +242,20 @@ function resolveContextSize(
   return { max: modelCap };
 }
 
-async function disposeLoadedModel(): Promise<void> {
-  if (!loadedModel) {
+async function disposeLoadedModel(state: LlamaCppInferenceRuntimeState): Promise<void> {
+  if (!state.loadedModel) {
     return;
   }
-  const previous = loadedModel;
-  loadedModel = undefined;
+  const previous = state.loadedModel;
   await previous.context.dispose();
   await previous.model.dispose();
+  if (state.loadedModel === previous) {
+    state.loadedModel = undefined;
+  }
 }
 
 async function getLoadedModel(params: {
+  state: LlamaCppInferenceRuntimeState;
   runtime: NodeLlamaCppModule;
   model: Parameters<StreamFn>[0];
   providerConfig?: ModelProviderConfig;
@@ -258,12 +268,12 @@ async function getLoadedModel(params: {
   });
   const contextSize = resolveContextSize(params.model, params.providerConfig);
   const key = `${modelPath}\0${JSON.stringify(contextSize)}`;
-  if (loadedModel?.key === key) {
-    return loadedModel;
+  if (params.state.loadedModel?.key === key) {
+    return params.state.loadedModel;
   }
-  await disposeLoadedModel();
-  const llama = llamaInstance ?? (await params.runtime.getLlama());
-  llamaInstance = llama;
+  await disposeLoadedModel(params.state);
+  const llama = params.state.llamaInstance ?? (await params.runtime.getLlama());
+  params.state.llamaInstance = llama;
   const fitContextSize = typeof contextSize === "number" ? contextSize : contextSize.max;
   const model = await llama.loadModel({
     modelPath,
@@ -276,8 +286,8 @@ async function getLoadedModel(params: {
     // Serialized requests reuse this one sequence. Disposing/reallocating it per
     // turn races node-llama-cpp's asynchronous sequence-id reclamation.
     const sequence = context.getSequence();
-    loadedModel = { key, llama, model, context, sequence };
-    return loadedModel;
+    params.state.loadedModel = { key, llama, model, context, sequence };
+    return params.state.loadedModel;
   } catch (error) {
     await context?.dispose();
     await model.dispose();
@@ -285,25 +295,66 @@ async function getLoadedModel(params: {
   }
 }
 
-async function serialize(operation: () => Promise<void>): Promise<void> {
-  const current = operationQueue.then(operation, operation);
-  operationQueue = current.catch(() => undefined);
+async function serialize(
+  state: LlamaCppInferenceRuntimeState,
+  operation: () => Promise<void>,
+): Promise<void> {
+  const current = state.operationQueue.then(operation, operation);
+  state.operationQueue = current.catch(() => undefined);
   await current;
 }
 
-async function clearLlamaCppInferenceCacheForTests(): Promise<void> {
-  await serialize(async () => {
-    await disposeLoadedModel();
-    if (llamaInstance) {
-      await llamaInstance.dispose();
-      llamaInstance = undefined;
+function disposeLlamaCppInferenceRuntime(state: LlamaCppInferenceRuntimeState): Promise<void> {
+  if (state.lifecycle === "closed") {
+    return Promise.resolve();
+  }
+  if (state.disposePromise) {
+    return state.disposePromise;
+  }
+  state.lifecycle = "closing";
+  const attempt = serialize(state, async () => {
+    await disposeLoadedModel(state);
+    if (state.llamaInstance) {
+      const previous = state.llamaInstance;
+      await previous.dispose();
+      if (state.llamaInstance === previous) {
+        state.llamaInstance = undefined;
+      }
     }
   });
+  state.disposePromise = attempt.then(
+    () => {
+      state.lifecycle = "closed";
+      state.disposePromise = undefined;
+    },
+    (error: unknown) => {
+      state.disposePromise = undefined;
+      throw error;
+    },
+  );
+  return state.disposePromise;
 }
 
-export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderConfig }): StreamFn {
+function createLlamaCppStreamFnForRuntime(
+  state: LlamaCppInferenceRuntimeState,
+  params: { providerConfig?: ModelProviderConfig },
+): StreamFn {
   return createPlainTextToolCallCompatWrapper((model, context, options) => {
     const stream = createAssistantMessageEventStream();
+    if (state.lifecycle !== "open") {
+      stream.push({
+        type: "error",
+        reason: "error",
+        error: buildMessage({
+          model,
+          content: [],
+          stopReason: "error",
+          errorMessage: "llama.cpp runtime is stopping",
+        }),
+      });
+      stream.end();
+      return stream;
+    }
     let streamedText = "";
     const streamedContent: AssistantMessage["content"] = [];
     let generationAborted = false;
@@ -340,6 +391,7 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
       try {
         const runtime = await importNodeLlamaCpp();
         const loaded = await getLoadedModel({
+          state,
           runtime,
           model,
           providerConfig: params.providerConfig,
@@ -458,10 +510,10 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
         const appendFunctionCallParamsChunk = (chunk: LlamaChatResponseFunctionCallParamsChunk) => {
           closeThinkingBlock();
           closeTextBlock();
-          let state = streamedToolCalls.get(chunk.callIndex);
-          if (!state) {
+          let callState = streamedToolCalls.get(chunk.callIndex);
+          if (!callState) {
             ensureStreamStarted();
-            state = {
+            callState = {
               toolCall: {
                 type: "toolCall",
                 id: `llama_cpp_call_${randomUUID()}`,
@@ -471,26 +523,26 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
               contentIndex: streamedContent.length,
               partialArgs: "",
             };
-            streamedToolCalls.set(chunk.callIndex, state);
-            streamedContent.push(state.toolCall);
+            streamedToolCalls.set(chunk.callIndex, callState);
+            streamedContent.push(callState.toolCall);
             stream.push({
               type: "toolcall_start",
-              contentIndex: state.contentIndex,
+              contentIndex: callState.contentIndex,
               partial: partial(),
             });
           }
           if (chunk.paramsChunk) {
-            state.partialArgs += chunk.paramsChunk;
+            callState.partialArgs += chunk.paramsChunk;
             // Replace the block so already queued partial snapshots retain the
             // exact argument state they exposed before this streamed delta.
-            state.toolCall = {
-              ...state.toolCall,
-              arguments: parseStreamingJson(state.partialArgs),
+            callState.toolCall = {
+              ...callState.toolCall,
+              arguments: parseStreamingJson(callState.partialArgs),
             };
-            streamedContent[state.contentIndex] = state.toolCall;
+            streamedContent[callState.contentIndex] = callState.toolCall;
             stream.push({
               type: "toolcall_delta",
-              contentIndex: state.contentIndex,
+              contentIndex: callState.contentIndex,
               delta: chunk.paramsChunk,
               partial: partial(),
             });
@@ -542,35 +594,35 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
           const confirmedCalls =
             result.metadata.stopReason === "maxTokens" ? [] : (result.functionCalls ?? []);
           const toolCalls: ToolCall[] = confirmedCalls.map((call, callIndex) => {
-            let state = streamedToolCalls.get(callIndex);
+            let callState = streamedToolCalls.get(callIndex);
             const argumentsObject = normalizeArguments(call.params);
-            if (!state) {
+            if (!callState) {
               appendFunctionCallParamsChunk({
                 callIndex,
                 functionName: call.functionName,
                 paramsChunk: JSON.stringify(argumentsObject),
                 done: true,
               });
-              state = streamedToolCalls.get(callIndex);
+              callState = streamedToolCalls.get(callIndex);
             }
-            if (!state) {
+            if (!callState) {
               throw new Error("llama.cpp native tool call stream state is missing");
             }
-            state.toolCall = {
-              ...state.toolCall,
+            callState.toolCall = {
+              ...callState.toolCall,
               name: call.functionName,
               arguments: argumentsObject,
             };
-            streamedContent[state.contentIndex] = state.toolCall;
+            streamedContent[callState.contentIndex] = callState.toolCall;
             // The dependency reports its final argument chunk before checking the
             // token budget; only this authoritative result can complete a call.
             stream.push({
               type: "toolcall_end",
-              contentIndex: state.contentIndex,
-              toolCall: state.toolCall,
+              contentIndex: callState.contentIndex,
+              toolCall: callState.toolCall,
               partial: partial(),
             });
-            return state.toolCall;
+            return callState.toolCall;
           });
           const confirmedToolCallIds = new Set(toolCalls.map((toolCall) => toolCall.id));
           const content = streamedContent.filter(
@@ -612,16 +664,27 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC
       }
     };
     if (!ended) {
-      queueMicrotask(() => void serialize(run));
+      void serialize(state, run);
     }
     return stream;
   });
 }
 
-if (process.env.VITEST || process.env.NODE_ENV === "test") {
-  (globalThis as Record<PropertyKey, unknown>)[Symbol.for("openclaw.llamaCppInferenceTestApi")] = {
-    mapContextToLlamaChatHistory,
-    mapToolsToLlamaFunctions,
-    clearLlamaCppInferenceCacheForTests,
+export function createLlamaCppInferenceRuntime(): LlamaCppInferenceRuntime {
+  const state: LlamaCppInferenceRuntimeState = {
+    operationQueue: Promise.resolve(),
+    lifecycle: "open",
+  };
+  return {
+    createStreamFn: (params) => createLlamaCppStreamFnForRuntime(state, params),
+    dispose: () => disposeLlamaCppInferenceRuntime(state),
   };
 }
+
+export const llamaCppInferenceTestApi =
+  process.env.VITEST || process.env.NODE_ENV === "test"
+    ? {
+        mapContextToLlamaChatHistory,
+        mapToolsToLlamaFunctions,
+      }
+    : undefined;