diff --git a/CHANGELOG.md b/CHANGELOG.md index 389f26d66c81..5bd4777fec09 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,6 +64,7 @@ Docs: https://docs.openclaw.ai ### Fixes +- **llama.cpp runtime cleanup:** dispose in-process GGUF contexts, models, and native runtimes when the provider plugin stops or reloads so replaced registrations do not retain model memory. Fixes #119144. - **Guided onboarding skip-UI routing:** keep `openclaw onboard --skip-ui` and `openclaw setup --skip-ui` on guided onboarding while skipping both browser and terminal handoffs, instead of silently switching to the classic wizard. Thanks @shakkernerd. - **Telegram durable ingress:** preserve pre-identity control-lane ownership during replay and attempt each drain snapshot row only once per pass, preventing targeted commands from spinning the spool and blocking polling shutdown. - **Control UI operator session permissions:** honor Gateway-advertised operator scopes for new-thread creation, thread management, checkpoints, and sharing controls while preserving read-only navigation and legacy Gateway compatibility. Fixes #117786. Thanks @shakkernerd. diff --git a/extensions/llama-cpp/index.test.ts b/extensions/llama-cpp/index.test.ts index 15a1e5c4bd73..4483334ae9cb 100644 --- a/extensions/llama-cpp/index.test.ts +++ b/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({ diff --git a/extensions/llama-cpp/index.ts b/extensions/llama-cpp/index.ts index ed2a74b32656..dbbd42102e19 100644 --- a/extensions/llama-cpp/index.ts +++ b/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], }); }, diff --git a/extensions/llama-cpp/src/inference-provider.test.ts b/extensions/llama-cpp/src/inference-provider.test.ts index e645b1fce79f..214dacb61ede 100644 --- a/extensions/llama-cpp/src/inference-provider.test.ts +++ b/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)[ - Symbol.for("openclaw.llamaCppInferenceTestApi") -] as { - clearLlamaCppInferenceCacheForTests: () => Promise; - mapContextToLlamaChatHistory: (context: Context) => unknown[]; - mapToolsToLlamaFunctions: (context: Context) => Record | 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>[2]; + options?: Parameters>[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 }], }); diff --git a/extensions/llama-cpp/src/inference-provider.ts b/extensions/llama-cpp/src/inference-provider.ts index a6d04d927831..f55dee5bd4d2 100644 --- a/extensions/llama-cpp/src/inference-provider.ts +++ b/extensions/llama-cpp/src/inference-provider.ts @@ -41,11 +41,18 @@ type LoadedModel = { type LlamaJsonSchemaInput = Parameters[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 = Promise.resolve(); +type LlamaCppInferenceRuntimeState = { + loadedModel?: LoadedModel; + llamaInstance?: Llama; + operationQueue: Promise; + lifecycle: "open" | "closing" | "closed"; + disposePromise?: Promise; +}; + +export type LlamaCppInferenceRuntime = { + createStreamFn: (params: { providerConfig?: ModelProviderConfig }) => StreamFn; + dispose: () => Promise; +}; function zeroCostUsage(input = 0, output = 0): Usage { return { @@ -235,17 +242,20 @@ function resolveContextSize( return { max: modelCap }; } -async function disposeLoadedModel(): Promise { - if (!loadedModel) { +async function disposeLoadedModel(state: LlamaCppInferenceRuntimeState): Promise { + 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[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): Promise { - const current = operationQueue.then(operation, operation); - operationQueue = current.catch(() => undefined); +async function serialize( + state: LlamaCppInferenceRuntimeState, + operation: () => Promise, +): Promise { + const current = state.operationQueue.then(operation, operation); + state.operationQueue = current.catch(() => undefined); await current; } -async function clearLlamaCppInferenceCacheForTests(): Promise { - await serialize(async () => { - await disposeLoadedModel(); - if (llamaInstance) { - await llamaInstance.dispose(); - llamaInstance = undefined; +function disposeLlamaCppInferenceRuntime(state: LlamaCppInferenceRuntimeState): Promise { + 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)[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;