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-messages.ts b/extensions/llama-cpp/src/inference-messages.ts new file mode 100644 index 000000000000..f7116bd34e1d --- /dev/null +++ b/extensions/llama-cpp/src/inference-messages.ts @@ -0,0 +1,51 @@ +import type { StreamFn } from "openclaw/plugin-sdk/agent-core"; +import type { AssistantMessage, StopReason, Usage } from "openclaw/plugin-sdk/llm"; + +export function zeroCostUsage(input = 0, output = 0): Usage { + return { + input, + output, + cacheRead: 0, + cacheWrite: 0, + totalTokens: input + output, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; +} + +export function buildMessage(params: { + model: Parameters[0]; + content: AssistantMessage["content"]; + stopReason: StopReason; + usage?: Usage; + errorMessage?: string; +}): AssistantMessage { + return { + role: "assistant", + content: params.content, + api: params.model.api, + provider: params.model.provider, + model: params.model.id, + stopReason: params.stopReason, + usage: params.usage ?? zeroCostUsage(), + timestamp: Date.now(), + ...(params.errorMessage ? { errorMessage: params.errorMessage } : {}), + }; +} + +export function runtimeUnavailableErrorMessage(restartRequired: boolean): string { + return restartRequired + ? "llama.cpp runtime stopped after native cleanup failed. Fully stop the managed Gateway service or foreground Gateway process, then start it again. An in-process restart cannot recover native resources." + : "llama.cpp runtime is stopping"; +} + +export function runtimeUnavailableMessage( + model: Parameters[0], + restartRequired: boolean, +): AssistantMessage { + return buildMessage({ + model, + content: [], + stopReason: "error", + errorMessage: runtimeUnavailableErrorMessage(restartRequired), + }); +} diff --git a/extensions/llama-cpp/src/inference-provider.mapping.test.ts b/extensions/llama-cpp/src/inference-provider.mapping.test.ts new file mode 100644 index 000000000000..c70dd6b3b401 --- /dev/null +++ b/extensions/llama-cpp/src/inference-provider.mapping.test.ts @@ -0,0 +1,101 @@ +import type { Context } from "openclaw/plugin-sdk/llm"; +import { describe, expect, it } from "vitest"; +import "./inference-provider.js"; + +const { mapContextToLlamaChatHistory, mapToolsToLlamaFunctions } = ( + globalThis as Record +)[Symbol.for("openclaw.llamaCppInferenceTestApi")] as { + mapContextToLlamaChatHistory: (context: Context) => unknown[]; + mapToolsToLlamaFunctions: (context: Context) => Record | undefined; +}; + +describe("llama.cpp inference mappings", () => { + it("maps OpenClaw history and tool results into the model chat template history", () => { + const context = { + systemPrompt: "Be concise.", + messages: [ + { role: "user" as const, content: "weather?", timestamp: 1 }, + { + role: "assistant" as const, + api: "openai-completions", + provider: "test", + model: "test", + stopReason: "toolUse" as const, + usage: { + input: 1, + output: 1, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 2, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: 2, + content: [ + { type: "text" as const, text: "Checking." }, + { + type: "toolCall" as const, + id: "call-1", + name: "weather", + arguments: { city: "Berlin" }, + }, + ], + }, + { + role: "toolResult" as const, + toolCallId: "call-1", + toolName: "weather", + content: [{ type: "text" as const, text: "Sunny" }], + isError: false, + timestamp: 3, + }, + { role: "user" as const, content: "thanks", timestamp: 4 }, + ], + }; + + expect(mapContextToLlamaChatHistory(context)).toEqual([ + { type: "system", text: "Be concise." }, + { type: "user", text: "weather?" }, + { + type: "model", + response: [ + "Checking.", + { + type: "functionCall", + name: "weather", + params: { city: "Berlin" }, + result: "Sunny", + }, + ], + }, + { type: "user", text: "thanks" }, + ]); + }); + + it("maps JSON-schema tools to native node-llama-cpp function definitions", () => { + expect( + mapToolsToLlamaFunctions({ + messages: [], + tools: [ + { + name: "weather", + description: "Get weather", + parameters: { + type: "object", + properties: { city: { type: "string" } }, + required: ["city"], + }, + }, + ], + }), + ).toEqual({ + weather: { + description: "Get weather", + params: { + type: "object", + properties: { city: { type: "string" } }, + required: ["city"], + }, + }, + }); + }); +}); diff --git a/extensions/llama-cpp/src/inference-provider.test.ts b/extensions/llama-cpp/src/inference-provider.test.ts index e645b1fce79f..e395bb8116c0 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 } from "./inference-provider.js"; -const { - clearLlamaCppInferenceCacheForTests, - mapContextToLlamaChatHistory, - mapToolsToLlamaFunctions, -} = (globalThis as Record)[ +type LlamaCppInferenceRuntime = ReturnType; +let inferenceRuntime: LlamaCppInferenceRuntime; +const testApi = (globalThis as Record)[ Symbol.for("openclaw.llamaCppInferenceTestApi") ] as { - clearLlamaCppInferenceCacheForTests: () => Promise; - mapContextToLlamaChatHistory: (context: Context) => unknown[]; - mapToolsToLlamaFunctions: (context: Context) => Record | undefined; + resetInferenceRuntimeCoordinator: () => void; }; +const NATIVE_CLEANUP_RECOVERY_MESSAGE = + "llama.cpp runtime stopped after native cleanup failed. Fully stop the managed Gateway service or foreground Gateway process, then start it again. An in-process restart cannot recover native resources."; 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,31 @@ 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?.(); +} + +function expectDisposeCalls(contextCount: number, modelCount: number, llamaCount: number) { + expect(mocks.contextDispose).toHaveBeenCalledTimes(contextCount); + expect(mocks.modelDispose).toHaveBeenCalledTimes(modelCount); + expect(mocks.llamaDispose).toHaveBeenCalledTimes(llamaCount); +} + +beforeEach(() => { + testApi.resetInferenceRuntimeCoordinator(); + inferenceRuntime = createLlamaCppInferenceRuntime(); vi.clearAllMocks(); mocks.generateResponse.mockResolvedValue({ response: "", @@ -142,99 +163,10 @@ beforeEach(async () => { }); afterEach(async () => { - await clearLlamaCppInferenceCacheForTests(); + await inferenceRuntime.dispose().catch(() => undefined); }); describe("llama.cpp inference provider", () => { - it("maps OpenClaw history and tool results into the model chat template history", () => { - const context = { - systemPrompt: "Be concise.", - messages: [ - { role: "user" as const, content: "weather?", timestamp: 1 }, - { - role: "assistant" as const, - api: "openai-completions", - provider: "test", - model: "test", - stopReason: "toolUse" as const, - usage: { - input: 1, - output: 1, - cacheRead: 0, - cacheWrite: 0, - totalTokens: 2, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }, - timestamp: 2, - content: [ - { type: "text" as const, text: "Checking." }, - { - type: "toolCall" as const, - id: "call-1", - name: "weather", - arguments: { city: "Berlin" }, - }, - ], - }, - { - role: "toolResult" as const, - toolCallId: "call-1", - toolName: "weather", - content: [{ type: "text" as const, text: "Sunny" }], - isError: false, - timestamp: 3, - }, - { role: "user" as const, content: "thanks", timestamp: 4 }, - ], - }; - - expect(mapContextToLlamaChatHistory(context)).toEqual([ - { type: "system", text: "Be concise." }, - { type: "user", text: "weather?" }, - { - type: "model", - response: [ - "Checking.", - { - type: "functionCall", - name: "weather", - params: { city: "Berlin" }, - result: "Sunny", - }, - ], - }, - { type: "user", text: "thanks" }, - ]); - }); - - it("maps JSON-schema tools to native node-llama-cpp function definitions", () => { - expect( - mapToolsToLlamaFunctions({ - messages: [], - tools: [ - { - name: "weather", - description: "Get weather", - parameters: { - type: "object", - properties: { city: { type: "string" } }, - required: ["city"], - }, - }, - ], - }), - ).toEqual({ - weather: { - description: "Get weather", - params: { - type: "object", - properties: { city: { type: "string" } }, - required: ["city"], - }, - }, - }); - }); - it("streams text deltas and reports native token-meter usage", async () => { mocks.generateResponse.mockImplementationOnce(async (_history, options) => { options.onTextChunk("Hel"); @@ -835,36 +767,201 @@ describe("llama.cpp inference provider", () => { expect(mocks.generateResponse.mock.calls[0]?.[1]).not.toHaveProperty("grammar"); }); - it("disposes the previous model and context when the model changes", async () => { - const streamFn = createLlamaCppStreamFn({}); - await collectEvents( - await streamFn(model, { messages: [{ role: "user", content: "one", timestamp: 1 }] }), - ); - await collectEvents( - await streamFn( - { ...model, id: "other.gguf", params: { modelPath: "other.gguf" } }, - { messages: [{ role: "user", content: "two", timestamp: 2 }] }, - ), - ); - - expect(mocks.contextDispose).toHaveBeenCalledTimes(1); - expect(mocks.modelDispose).toHaveBeenCalledTimes(1); + it("makes failed changed-model cleanup terminal", async () => { + const otherModel = { ...model, id: "other.gguf", params: { modelPath: "other.gguf" } }; + await collectTestEvents({ prompt: "one" }); + await collectTestEvents({ selectedModel: otherModel, prompt: "two" }); + let rejectCleanup!: (error: Error) => void; + const cleanup = new Promise((_resolve, reject) => { + rejectCleanup = reject; + }); + mocks.contextDispose.mockImplementationOnce(async () => await cleanup); + const failedSwitch = await createTestStream({ prompt: "three" }); + await vi.waitFor(() => expect(mocks.contextDispose).toHaveBeenCalledTimes(2)); + const unavailable = await createTestStream({ selectedModel: otherModel, prompt: "four" }); + const disposing = inferenceRuntime.dispose(); + rejectCleanup(new Error("context cleanup failed")); + await expect(failedSwitch.result()).resolves.toMatchObject({ + errorMessage: NATIVE_CLEANUP_RECOVERY_MESSAGE, + }); + await expect(unavailable.result()).resolves.toMatchObject({ + errorMessage: NATIVE_CLEANUP_RECOVERY_MESSAGE, + }); + await expect(disposing).rejects.toThrow("context cleanup failed"); + expectDisposeCalls(2, 1, 0); expect(mocks.llama.loadModel).toHaveBeenCalledTimes(2); }); + it("records cleanup failure during partial model initialization", async () => { + mocks.model.createContext.mockRejectedValueOnce(new Error("context creation failed")); + mocks.modelDispose.mockRejectedValueOnce(new Error("model cleanup failed")); + const failedInitialization = await createTestStream(); + await expect(failedInitialization.result()).resolves.toMatchObject({ + errorMessage: NATIVE_CLEANUP_RECOVERY_MESSAGE, + }); + await expect(inferenceRuntime.dispose()).rejects.toThrow("model cleanup failed"); + expectDisposeCalls(0, 1, 0); + }); + it("reuses one context sequence across serialized requests for the same model", async () => { - const streamFn = createLlamaCppStreamFn({}); - await collectEvents( - await streamFn(model, { messages: [{ role: "user", content: "one", timestamp: 1 }] }), - ); - await collectEvents( - await streamFn(model, { messages: [{ role: "user", content: "two", timestamp: 2 }] }), - ); + await collectTestEvents({ prompt: "one" }); + await collectTestEvents({ prompt: "two" }); expect(mocks.context.getSequence).toHaveBeenCalledTimes(1); expect(mocks.llama.loadModel).toHaveBeenCalledTimes(1); }); + it("disposes the context, model, and native runtime in ownership order", async () => { + await collectTestEvents(); + + await inferenceRuntime.dispose(); + + expectDisposeCalls(1, 1, 1); + 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("blocks a published replacement until predecessor service stop finishes", async () => { + const retiringRuntime = inferenceRuntime; + await collectTestEvents(); + let finishCleanup!: () => void; + mocks.contextDispose.mockImplementationOnce( + async () => + await new Promise((resolve) => { + finishCleanup = resolve; + }), + ); + // Gateway publishes the replacement registry before stopping old services. + const disposing = retiringRuntime.dispose(); + await vi.waitFor(() => expect(mocks.contextDispose).toHaveBeenCalledOnce()); + + const replacementRuntime = createLlamaCppInferenceRuntime(); + const replacementEvents = collectEvents( + await replacementRuntime.createStreamFn({})(model, { + messages: [{ role: "user", content: "after reload", timestamp: 2 }], + }), + ); + await Promise.resolve(); + expect(mocks.llama.loadModel).toHaveBeenCalledOnce(); + + finishCleanup(); + await disposing; + await expect(replacementEvents).resolves.toEqual( + expect.arrayContaining([expect.objectContaining({ type: "done", reason: "stop" })]), + ); + expect(mocks.llama.loadModel).toHaveBeenCalledTimes(2); + inferenceRuntime = replacementRuntime; + }); + + it("keeps a published replacement blocked when predecessor service stop fails", async () => { + const retiringRuntime = inferenceRuntime; + await collectTestEvents(); + mocks.contextDispose.mockRejectedValueOnce(new Error("retiring cleanup failed")); + + const replacementRuntime = createLlamaCppInferenceRuntime(); + // Match reload ordering: replacement is reachable before old service stop settles. + const disposing = retiringRuntime.dispose(); + const replacementEvents = collectEvents( + await replacementRuntime.createStreamFn({})(model, { + messages: [{ role: "user", content: "after failed reload", timestamp: 2 }], + }), + ); + + await expect(disposing).rejects.toThrow("retiring cleanup failed"); + await expect(replacementEvents).resolves.toEqual( + expect.arrayContaining([ + expect.objectContaining({ + type: "error", + error: expect.objectContaining({ + errorMessage: NATIVE_CLEANUP_RECOVERY_MESSAGE, + }), + }), + ]), + ); + expect(mocks.llama.loadModel).toHaveBeenCalledOnce(); + await expect(replacementRuntime.dispose()).resolves.toBeUndefined(); + inferenceRuntime = replacementRuntime; + }); + + it("lets a superseded waiting replacement stop before its predecessor retires", async () => { + await collectTestEvents(); + const waitingRuntime = createLlamaCppInferenceRuntime(); + const waitingStream = await waitingRuntime.createStreamFn({})(model, { + messages: [{ role: "user", content: "superseded reload", timestamp: 2 }], + }); + await Promise.resolve(); + + const disposingWaitingRuntime = waitingRuntime.dispose(); + + await expect(waitingStream.result()).resolves.toMatchObject({ + stopReason: "error", + errorMessage: "llama.cpp runtime is stopping", + }); + await expect(disposingWaitingRuntime).resolves.toBeUndefined(); + expect(mocks.llama.loadModel).toHaveBeenCalledOnce(); + }); + + 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; + + expectDisposeCalls(1, 1, 1); + }); + + 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); + expectDisposeCalls(1, 1, 1); + }); + + it("keeps a failed native runtime cleanup terminal", async () => { + await collectTestEvents(); + mocks.llamaDispose.mockRejectedValueOnce(new Error("llama cleanup failed")); + const firstDisposal = inferenceRuntime.dispose(); + await expect(firstDisposal).rejects.toThrow("llama cleanup failed"); + expectDisposeCalls(1, 1, 1); + const unavailable = await createTestStream({ prompt: "after failed stop" }); + await expect(unavailable.result()).resolves.toMatchObject({ + errorMessage: NATIVE_CLEANUP_RECOVERY_MESSAGE, + }); + expect(inferenceRuntime.dispose()).toBe(firstDisposal); + }); + it.each([ { scenario: "a smaller advertised model window", @@ -964,7 +1061,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..22096b9e5559 100644 --- a/extensions/llama-cpp/src/inference-provider.ts +++ b/extensions/llama-cpp/src/inference-provider.ts @@ -10,13 +10,7 @@ import type { LlamaModel, } from "node-llama-cpp"; import type { StreamFn } from "openclaw/plugin-sdk/agent-core"; -import type { - AssistantMessage, - Context, - StopReason, - ToolCall, - Usage, -} from "openclaw/plugin-sdk/llm"; +import type { AssistantMessage, Context, StopReason, ToolCall } from "openclaw/plugin-sdk/llm"; import { createAssistantMessageEventStream, parseStreamingJson } from "openclaw/plugin-sdk/llm"; import type { ModelProviderConfig } from "openclaw/plugin-sdk/provider-model-shared"; import { createPlainTextToolCallCompatWrapper } from "openclaw/plugin-sdk/provider-stream-shared"; @@ -25,6 +19,16 @@ import { resolveLlamaCppModelCacheDir, resolveLlamaCppModelSource, } from "./defaults.js"; +import { + buildMessage, + runtimeUnavailableErrorMessage, + runtimeUnavailableMessage, + zeroCostUsage, +} from "./inference-messages.js"; +import { + createLlamaCppInferenceRuntimeToken, + type LlamaCppInferenceRuntimeToken, +} from "./inference-runtime-coordinator.js"; import { formatLlamaCppSetupError, importNodeLlamaCpp, @@ -41,41 +45,24 @@ 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 = { + admission: LlamaCppInferenceRuntimeToken; + loadedModel?: LoadedModel; + llamaInstance?: Llama; + operationQueue: Promise; + lifecycle: "open" | "closing" | "closed"; + cleanupFailure?: { error: Error }; + retiringRuntimeFailure?: boolean; + disposePromise?: Promise; +}; -function zeroCostUsage(input = 0, output = 0): Usage { - return { - input, - output, - cacheRead: 0, - cacheWrite: 0, - totalTokens: input + output, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }; -} +type LlamaCppInferenceRuntime = { + createStreamFn: (params: { providerConfig?: ModelProviderConfig }) => StreamFn; + dispose: () => Promise; +}; -function buildMessage(params: { - model: Parameters[0]; - content: AssistantMessage["content"]; - stopReason: StopReason; - usage?: Usage; - errorMessage?: string; -}): AssistantMessage { - return { - role: "assistant", - content: params.content, - api: params.model.api, - provider: params.model.provider, - model: params.model.id, - stopReason: params.stopReason, - usage: params.usage ?? zeroCostUsage(), - timestamp: Date.now(), - ...(params.errorMessage ? { errorMessage: params.errorMessage } : {}), - }; +function runtimeRequiresRestart(state: LlamaCppInferenceRuntimeState): boolean { + return Boolean(state.cleanupFailure || state.retiringRuntimeFailure); } function extractText(content: unknown): string { @@ -235,17 +222,34 @@ 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; - await previous.context.dispose(); - await previous.model.dispose(); + const previous = state.loadedModel; + state.loadedModel = undefined; + try { + await previous.context.dispose(); + await previous.model.dispose(); + } catch (error) { + recordCleanupFailure(state, error); + throw error; + } +} + +function recordCleanupFailure(state: LlamaCppInferenceRuntimeState, error: unknown): void { + state.cleanupFailure ??= { error: error instanceof Error ? error : new Error(String(error)) }; + state.lifecycle = "closed"; + state.admission.fail(state.cleanupFailure.error); +} + +function recordRetiringRuntimeFailure(state: LlamaCppInferenceRuntimeState): void { + state.retiringRuntimeFailure = true; + state.lifecycle = "closed"; } async function getLoadedModel(params: { + state: LlamaCppInferenceRuntimeState; runtime: NodeLlamaCppModule; model: Parameters[0]; providerConfig?: ModelProviderConfig; @@ -258,12 +262,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,34 +280,81 @@ 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(); + try { + await context?.dispose(); + await model.dispose(); + } catch (cleanupError) { + recordCleanupFailure(params.state, cleanupError); + throw cleanupError; + } throw error; } } -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.disposePromise) { + return state.disposePromise; + } + if (state.cleanupFailure) { + state.disposePromise = Promise.reject(state.cleanupFailure.error); + return state.disposePromise; + } + state.lifecycle = "closing"; + state.admission.close(); + // node-llama-cpp disposers are one-shot and child cleanup releases the + // parent's disposal guard. Do not force parent cleanup after a child rejects: + // the retained guard can make that parent disposer wait forever. + state.disposePromise = serialize(state, async () => { + if (state.cleanupFailure) { + throw state.cleanupFailure.error; } - }); + await disposeLoadedModel(state); + if (state.llamaInstance) { + const previous = state.llamaInstance; + await previous.dispose(); + if (state.llamaInstance === previous) { + state.llamaInstance = undefined; + } + } + state.admission.release(); + }) + .catch((error: unknown) => { + recordCleanupFailure(state, error); + throw error; + }) + .finally(() => { + state.lifecycle = "closed"; + }); + 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: runtimeUnavailableMessage(model, runtimeRequiresRestart(state)), + }); + stream.end(); + return stream; + } let streamedText = ""; const streamedContent: AssistantMessage["content"] = []; let generationAborted = false; @@ -338,8 +389,25 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC started = true; signal?.removeEventListener("abort", abortWhileQueued); try { + if (state.lifecycle !== "open") { + stream.push({ + type: "error", + reason: "error", + error: runtimeUnavailableMessage(model, runtimeRequiresRestart(state)), + }); + return; + } + await state.admission.acquire({ + signal, + isLive: () => state.lifecycle === "open", + onRestartRequired: () => recordRetiringRuntimeFailure(state), + }); + if (state.lifecycle !== "open") { + throw new Error("llama.cpp runtime is stopping"); + } const runtime = await importNodeLlamaCpp(); const loaded = await getLoadedModel({ + state, runtime, model, providerConfig: params.providerConfig, @@ -458,10 +526,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 +539,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 +610,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( @@ -595,7 +663,11 @@ export function createLlamaCppStreamFn(params: { providerConfig?: ModelProviderC } catch (error) { const aborted = generationAborted || options?.signal?.aborted === true; const reason = aborted ? "aborted" : "error"; - const errorMessage = aborted ? "Request was aborted" : formatLlamaCppSetupError(error); + const errorMessage = aborted + ? "Request was aborted" + : state.lifecycle !== "open" + ? runtimeUnavailableErrorMessage(runtimeRequiresRestart(state)) + : formatLlamaCppSetupError(error); stream.push({ type: "error", reason, @@ -612,16 +684,29 @@ 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 = { + admission: createLlamaCppInferenceRuntimeToken(), + operationQueue: Promise.resolve(), + lifecycle: "open", + }; + return { + createStreamFn: (params) => createLlamaCppStreamFnForRuntime(state, params), + dispose: () => disposeLlamaCppInferenceRuntime(state), }; } + +if (process.env.VITEST || process.env.NODE_ENV === "test") { + const globalStore = globalThis as Record; + const testApiKey = Symbol.for("openclaw.llamaCppInferenceTestApi"); + Object.assign((globalStore[testApiKey] ??= {}), { + mapContextToLlamaChatHistory, + mapToolsToLlamaFunctions, + }); +} diff --git a/extensions/llama-cpp/src/inference-runtime-coordinator.test.ts b/extensions/llama-cpp/src/inference-runtime-coordinator.test.ts new file mode 100644 index 000000000000..040e9a98f928 --- /dev/null +++ b/extensions/llama-cpp/src/inference-runtime-coordinator.test.ts @@ -0,0 +1,114 @@ +import { beforeEach, describe, expect, it } from "vitest"; +import { createLlamaCppInferenceRuntimeToken } from "./inference-runtime-coordinator.js"; + +const testApi = (globalThis as Record)[ + Symbol.for("openclaw.llamaCppInferenceTestApi") +] as { + resetInferenceRuntimeCoordinator: () => void; +}; + +function createLiveToken() { + let live = true; + const token = createLlamaCppInferenceRuntimeToken(); + return { + acquire: (signal?: AbortSignal) => + token.acquire({ signal, isLive: () => live, onRestartRequired: () => undefined }), + dispose: () => { + live = false; + token.close(); + }, + fail: token.fail, + release: token.release, + }; +} + +beforeEach(() => { + testApi.resetInferenceRuntimeCoordinator(); +}); + +describe("llama.cpp inference runtime coordinator", () => { + it("does not reserve native ownership until a runtime acquires it", async () => { + createLiveToken(); + const replacement = createLiveToken(); + + await expect(replacement.acquire()).resolves.toBeUndefined(); + }); + + it("hands native ownership to rapid contenders one generation at a time", async () => { + const first = createLiveToken(); + const second = createLiveToken(); + const third = createLiveToken(); + await first.acquire(); + let secondAcquired = false; + let thirdAcquired = false; + const secondAcquisition = second.acquire().then(() => { + secondAcquired = true; + }); + const thirdAcquisition = third.acquire().then(() => { + thirdAcquired = true; + }); + + first.release(); + await secondAcquisition; + expect(secondAcquired).toBe(true); + expect(thirdAcquired).toBe(false); + + second.release(); + await thirdAcquisition; + expect(thirdAcquired).toBe(true); + }); + + it("skips a disposed waiter before handing ownership to the next runtime", async () => { + const first = createLiveToken(); + const disposed = createLiveToken(); + const replacement = createLiveToken(); + await first.acquire(); + const disposedAcquisition = disposed.acquire(); + let replacementAcquired = false; + const replacementAcquisition = replacement.acquire().then(() => { + replacementAcquired = true; + }); + + disposed.dispose(); + + await expect(disposedAcquisition).rejects.toThrow("runtime is stopping"); + expect(replacementAcquired).toBe(false); + + first.release(); + await expect(replacementAcquisition).resolves.toBeUndefined(); + }); + + it("removes an aborted waiter without letting it claim ownership", async () => { + const first = createLiveToken(); + const aborted = createLiveToken(); + const replacement = createLiveToken(); + const abortController = new AbortController(); + await first.acquire(); + const abortedAcquisition = aborted.acquire(abortController.signal); + const replacementAcquisition = replacement.acquire(); + + abortController.abort(); + first.release(); + + await expect(abortedAcquisition).rejects.toThrow(); + await expect(replacementAcquisition).resolves.toBeUndefined(); + }); + + it("latches cleanup failure for current and future waiters", async () => { + const first = createLiveToken(); + const waiting = createLiveToken(); + await first.acquire(); + const waitingAcquisition = waiting.acquire(); + + first.fail(new Error("native cleanup failed")); + + await expect(waitingAcquisition).rejects.toMatchObject({ + name: "LlamaCppInferenceRestartRequiredError", + code: "LLAMA_CPP_INFERENCE_RESTART_REQUIRED", + }); + await expect(createLiveToken().acquire()).rejects.toMatchObject({ + name: "LlamaCppInferenceRestartRequiredError", + code: "LLAMA_CPP_INFERENCE_RESTART_REQUIRED", + }); + }); +}); diff --git a/extensions/llama-cpp/src/inference-runtime-coordinator.ts b/extensions/llama-cpp/src/inference-runtime-coordinator.ts new file mode 100644 index 000000000000..8d0626cece71 --- /dev/null +++ b/extensions/llama-cpp/src/inference-runtime-coordinator.ts @@ -0,0 +1,174 @@ +import { resolveGlobalSingleton } from "openclaw/plugin-sdk/global-singleton"; + +const COORDINATOR_KEY = Symbol.for("openclaw.llamaCppInferenceRuntimeCoordinator"); +const TEST_API_KEY = Symbol.for("openclaw.llamaCppInferenceTestApi"); +const RESTART_REQUIRED_CODE = "LLAMA_CPP_INFERENCE_RESTART_REQUIRED"; + +type Completion = { + promise: Promise; + complete: () => void; +}; + +type CoordinatorState = { + owner?: Completion; + restartRequired?: Error; +}; + +type RuntimeToken = { + closing: Completion; + owner?: Completion; + released: boolean; +}; + +export type LlamaCppInferenceRuntimeToken = { + acquire: (params: { + signal?: AbortSignal; + isLive: () => boolean; + onRestartRequired: () => void; + }) => Promise; + close: () => void; + fail: (error: unknown) => void; + release: () => void; +}; + +class LlamaCppInferenceRestartRequiredError extends Error { + readonly code = RESTART_REQUIRED_CODE; + + constructor(cause: Error) { + super("A previous llama.cpp runtime failed to release native resources", { cause }); + this.name = "LlamaCppInferenceRestartRequiredError"; + } +} + +function getCoordinatorState(): CoordinatorState { + return resolveGlobalSingleton(COORDINATOR_KEY, () => ({})); +} + +function createCompletion(): Completion { + let complete!: () => void; + const promise = new Promise((resolve) => { + complete = resolve; + }); + return { promise, complete }; +} + +function abortedError(signal: AbortSignal): Error { + return signal.reason instanceof Error ? signal.reason : new Error("Request was aborted"); +} + +function waitForTurn(params: { + closing: Promise; + retired: Promise; + signal?: AbortSignal; +}): Promise<"closed" | "retired"> { + const signal = params.signal; + if (signal?.aborted) { + return Promise.reject(abortedError(signal)); + } + return new Promise((resolve, reject) => { + const cleanup = () => signal?.removeEventListener("abort", abort); + const abort = () => { + cleanup(); + reject(signal ? abortedError(signal) : new Error("Request was aborted")); + }; + const finish = (result: "closed" | "retired") => { + cleanup(); + resolve(result); + }; + signal?.addEventListener("abort", abort, { once: true }); + void params.closing.then(() => finish("closed")); + void params.retired.then(() => finish("retired")); + }); +} + +function isRestartRequiredError(error: unknown): error is LlamaCppInferenceRestartRequiredError { + return error instanceof Error && "code" in error && error.code === RESTART_REQUIRED_CODE; +} + +async function acquireRuntime( + token: RuntimeToken, + params: { signal?: AbortSignal; isLive: () => boolean }, +): Promise { + if (token.owner) { + return; + } + while (true) { + if (token.released || !params.isLive()) { + throw new Error("llama.cpp runtime is stopping"); + } + if (params.signal?.aborted) { + throw abortedError(params.signal); + } + const state = getCoordinatorState(); + if (state.restartRequired) { + throw new LlamaCppInferenceRestartRequiredError(state.restartRequired); + } + if (!state.owner) { + const owner = createCompletion(); + state.owner = owner; + token.owner = owner; + return; + } + if ( + (await waitForTurn({ + closing: token.closing.promise, + retired: state.owner.promise, + signal: params.signal, + })) === "closed" + ) { + throw new Error("llama.cpp runtime is stopping"); + } + } +} + +function releaseRuntime(token: RuntimeToken): void { + token.released = true; + if (!token.owner) { + return; + } + const state = getCoordinatorState(); + if (state.owner === token.owner) { + state.owner = undefined; + } + token.owner.complete(); + token.owner = undefined; +} + +function failRuntime(token: RuntimeToken, error: unknown): void { + const state = getCoordinatorState(); + state.restartRequired ??= error instanceof Error ? error : new Error(String(error)); + token.owner?.complete(); +} + +export function createLlamaCppInferenceRuntimeToken(): LlamaCppInferenceRuntimeToken { + // Registration creates only a local token; native ownership remains lazy. + const token: RuntimeToken = { closing: createCompletion(), released: false }; + return { + acquire: async ({ onRestartRequired, ...params }) => { + try { + await acquireRuntime(token, params); + } catch (error) { + // Reloaded plugin chunks have distinct class identities. + if (isRestartRequiredError(error)) { + onRestartRequired(); + } + throw error; + } + }, + close: () => token.closing.complete(), + fail: (error) => failRuntime(token, error), + release: () => releaseRuntime(token), + }; +} + +if (process.env.VITEST || process.env.NODE_ENV === "test") { + const globalStore = globalThis as Record; + const testApi = (globalStore[TEST_API_KEY] ?? {}) as Record; + testApi.resetInferenceRuntimeCoordinator = () => { + const state = getCoordinatorState(); + state.owner?.complete(); + state.owner = undefined; + state.restartRequired = undefined; + }; + globalStore[TEST_API_KEY] = testApi; +}