diff --git a/apps/memos-local-plugin/bridge.cts b/apps/memos-local-plugin/bridge.cts index f348b8c12..354064128 100644 --- a/apps/memos-local-plugin/bridge.cts +++ b/apps/memos-local-plugin/bridge.cts @@ -197,6 +197,9 @@ async function main(): Promise { const { startStdioServer, waitForShutdown } = (await importEsm( runtimeModule("bridge/stdio.ts", "dist/bridge/stdio.js") )) as typeof import("./bridge/stdio.js"); + const { createStdioHostLlmBridge } = (await importEsm( + runtimeModule("bridge/host-llm.ts", "dist/bridge/host-llm.js") + )) as typeof import("./bridge/host-llm.js"); const { memoryBuffer, rootLogger } = (await importEsm( runtimeModule("core/logger/index.ts", "dist/core/logger/index.js") )) as typeof import("./core/logger/index.js"); @@ -237,47 +240,10 @@ async function main(): Promise { // independent `currentBridge` slots. Registering inside bootstrap // forces both ends to share the same module instance. let stdio: import("./bridge/stdio.js").StdioServerHandle | null = null; - const lazyHostLlmBridge: import("./core/llm/host-bridge.js").HostLlmBridge = - { - id: `stdio.host.${args.agent}.v1`, - async complete(input) { - if (!stdio) { - throw new Error( - "host LLM bridge invoked before stdio server was ready", - ); - } - const result = (await stdio.serverRequest( - "host.llm.complete", - { - messages: input.messages, - model: input.model, - temperature: input.temperature, - maxTokens: input.maxTokens, - timeoutMs: input.timeoutMs, - }, - { timeoutMs: (input.timeoutMs ?? 60_000) + 5_000 }, - )) as { - text?: string; - model?: string; - usage?: { - promptTokens?: number; - completionTokens?: number; - totalTokens?: number; - }; - durationMs?: number; - }; - return { - text: typeof result?.text === "string" ? result.text : "", - model: - typeof result?.model === "string" - ? result.model - : input.model ?? "", - usage: result?.usage, - durationMs: - typeof result?.durationMs === "number" ? result.durationMs : 0, - }; - }, - }; + const lazyHostLlmBridge = createStdioHostLlmBridge( + `stdio.host.${args.agent}.v1`, + () => stdio, + ); const { Telemetry } = (await importEsm( runtimeModule("core/telemetry/index.ts", "dist/core/telemetry/index.js") diff --git a/apps/memos-local-plugin/bridge.mts b/apps/memos-local-plugin/bridge.mts index 75f7e7c37..69d7c50a8 100644 --- a/apps/memos-local-plugin/bridge.mts +++ b/apps/memos-local-plugin/bridge.mts @@ -180,6 +180,7 @@ async function main(): Promise { // from source via `tsx bridge.mts`. const { bootstrapMemoryCoreFull } = await import("./core/pipeline/index.js"); const { startStdioServer, waitForShutdown } = await import("./bridge/stdio.js"); + const { createStdioHostLlmBridge } = await import("./bridge/host-llm.js"); const { memoryBuffer, rootLogger } = await import("./core/logger/index.js"); const { startHttpServer } = await import("./server/http.js"); @@ -206,47 +207,10 @@ async function main(): Promise { // independent `currentBridge` slots. Registering inside bootstrap // forces both ends to share the same module instance. let stdio: import("./bridge/stdio.js").StdioServerHandle | null = null; - const lazyHostLlmBridge: import("./core/llm/host-bridge.js").HostLlmBridge = - { - id: `stdio.host.${args.agent}.v1`, - async complete(input) { - if (!stdio) { - throw new Error( - "host LLM bridge invoked before stdio server was ready", - ); - } - const result = (await stdio.serverRequest( - "host.llm.complete", - { - messages: input.messages, - model: input.model, - temperature: input.temperature, - maxTokens: input.maxTokens, - timeoutMs: input.timeoutMs, - }, - { timeoutMs: (input.timeoutMs ?? 60_000) + 5_000 }, - )) as { - text?: string; - model?: string; - usage?: { - promptTokens?: number; - completionTokens?: number; - totalTokens?: number; - }; - durationMs?: number; - }; - return { - text: typeof result?.text === "string" ? result.text : "", - model: - typeof result?.model === "string" - ? result.model - : input.model ?? "", - usage: result?.usage, - durationMs: - typeof result?.durationMs === "number" ? result.durationMs : 0, - }; - }, - }; + const lazyHostLlmBridge = createStdioHostLlmBridge( + `stdio.host.${args.agent}.v1`, + () => stdio, + ); const { Telemetry } = await import("./core/telemetry/index.js"); diff --git a/apps/memos-local-plugin/bridge/host-llm.ts b/apps/memos-local-plugin/bridge/host-llm.ts new file mode 100644 index 000000000..f1a54c4de --- /dev/null +++ b/apps/memos-local-plugin/bridge/host-llm.ts @@ -0,0 +1,54 @@ +import type { + HostLlmBridge, + HostLlmCompletion, +} from "../core/llm/host-bridge.js"; +import type { StdioServerHandle } from "./stdio.js"; + +type ServerRequestTransport = Pick; + +export function createStdioHostLlmBridge( + id: string, + getTransport: () => ServerRequestTransport | null, +): HostLlmBridge { + return { + id, + async complete(input) { + const transport = getTransport(); + if (!transport) { + throw new Error("host LLM bridge invoked before stdio server was ready"); + } + + const configuredTimeoutMs = input.timeoutMs ?? 60_000; + const remainingDeadlineMs = + input.deadlineAt === undefined + ? configuredTimeoutMs + : input.deadlineAt - Date.now(); + if (remainingDeadlineMs <= 0) { + throw new Error("host LLM request deadline has passed"); + } + const timeoutMs = Math.min(configuredTimeoutMs, remainingDeadlineMs); + + const result = await transport.serverRequest>( + "host.llm.complete", + { + messages: input.messages, + model: input.model, + temperature: input.temperature, + maxTokens: input.maxTokens, + timeoutMs, + }, + { timeoutMs, signal: input.signal }, + ); + return { + text: typeof result?.text === "string" ? result.text : "", + model: + typeof result?.model === "string" + ? result.model + : input.model ?? "", + usage: result?.usage, + durationMs: + typeof result?.durationMs === "number" ? result.durationMs : 0, + }; + }, + }; +} diff --git a/apps/memos-local-plugin/bridge/stdio.ts b/apps/memos-local-plugin/bridge/stdio.ts index a7ecd8e50..6296e8702 100644 --- a/apps/memos-local-plugin/bridge/stdio.ts +++ b/apps/memos-local-plugin/bridge/stdio.ts @@ -64,7 +64,7 @@ export interface StdioServerHandle { serverRequest( method: string, params?: unknown, - options?: { timeoutMs?: number }, + options?: { timeoutMs?: number; signal?: AbortSignal }, ): Promise; } @@ -91,9 +91,22 @@ export function startStdioServer(options: StdioServerOptions): StdioServerHandle resolve: (value: unknown) => void; reject: (err: unknown) => void; timer: ReturnType | null; + signal?: AbortSignal; + abortHandler?: () => void; } >(); + function takeServerPending(id: string) { + const entry = serverPending.get(id); + if (!entry) return undefined; + serverPending.delete(id); + if (entry.timer) clearTimeout(entry.timer); + if (entry.signal && entry.abortHandler) { + entry.signal.removeEventListener("abort", entry.abortHandler); + } + return entry; + } + const eventsUnsubscribe = options.core.subscribeEvents((e) => { writeNotification(RPC_METHODS.EVENTS_NOTIFY, e); }); @@ -114,10 +127,8 @@ export function startStdioServer(options: StdioServerOptions): StdioServerHandle } catch { /* ignore */ } - for (const [id, entry] of serverPending) { - if (entry.timer) clearTimeout(entry.timer); - entry.reject(err ?? new Error("stdio bridge closed")); - serverPending.delete(id); + for (const id of serverPending.keys()) { + takeServerPending(id)?.reject(err ?? new Error("stdio bridge closed")); } } @@ -179,10 +190,8 @@ export function startStdioServer(options: StdioServerOptions): StdioServerHandle (raw.result !== undefined || raw.error !== undefined) ) { const id = raw.id as string; - const pending = serverPending.get(id); + const pending = takeServerPending(id); if (pending) { - serverPending.delete(id); - if (pending.timer) clearTimeout(pending.timer); if (raw.error != null) { pending.reject(raw.error); } else { @@ -260,8 +269,12 @@ export function startStdioServer(options: StdioServerOptions): StdioServerHandle function serverRequest( method: string, params?: unknown, - options?: { timeoutMs?: number }, + options?: { timeoutMs?: number; signal?: AbortSignal }, ): Promise { + const signal = options?.signal; + if (signal?.aborted) { + return Promise.reject(signal.reason ?? new Error(`serverRequest ${method} aborted`)); + } const id = `srv-${++serverRequestSeq}`; const timeoutMs = options?.timeoutMs ?? 60_000; return new Promise((resolve, reject) => { @@ -270,15 +283,29 @@ export function startStdioServer(options: StdioServerOptions): StdioServerHandle return; } const timer = setTimeout(() => { - if (serverPending.delete(id)) { - reject(new Error(`serverRequest ${method} timed out after ${timeoutMs}ms`)); - } + takeServerPending(id)?.reject( + new Error(`serverRequest ${method} timed out after ${timeoutMs}ms`), + ); }, timeoutMs); + const abortHandler = signal + ? () => { + takeServerPending(id)?.reject( + signal.reason ?? new Error(`serverRequest ${method} aborted`), + ); + } + : undefined; serverPending.set(id, { resolve: resolve as (value: unknown) => void, reject, timer, + signal, + abortHandler, }); + if (signal && abortHandler) { + signal.addEventListener("abort", abortHandler, { once: true }); + } + if (signal?.aborted) abortHandler?.(); + if (!serverPending.has(id)) return; writeLine({ jsonrpc: "2.0", id, method, params }); }); } diff --git a/apps/memos-local-plugin/core/llm/README.md b/apps/memos-local-plugin/core/llm/README.md index 2ffea8e06..6aa7ae8e6 100644 --- a/apps/memos-local-plugin/core/llm/README.md +++ b/apps/memos-local-plugin/core/llm/README.md @@ -105,6 +105,11 @@ we call the host once as a best-effort fallback. The resulting completion's `servedBy` is `"host_fallback"` so downstream dashboards can distinguish host wins from primary wins. +Fallback is skipped once the caller's abort signal fires or its absolute +deadline passes. When fallback is still allowed, both cancellation and the +remaining deadline propagate to the host bridge so the host request cannot +outlive the caller's budget. + Registering from an adapter: ```ts diff --git a/apps/memos-local-plugin/core/llm/client.ts b/apps/memos-local-plugin/core/llm/client.ts index 7e19f25ca..762e0bea4 100644 --- a/apps/memos-local-plugin/core/llm/client.ts +++ b/apps/memos-local-plugin/core/llm/client.ts @@ -182,11 +182,19 @@ export function createLlmClientWithProvider( ); } - function canUseHostFallback(): boolean { + function callerBudgetIsGone(opts: LlmCallOptions | undefined): boolean { + return ( + opts?.signal?.aborted === true || + (opts?.deadlineAt !== undefined && Date.now() >= opts.deadlineAt) + ); + } + + function canUseHostFallback(opts?: LlmCallOptions): boolean { return ( config.fallbackToHost === true && provider.name !== "host" && - getHostLlmBridge() !== null + getHostLlmBridge() !== null && + !callerBudgetIsGone(opts) ); } @@ -311,7 +319,7 @@ export function createLlmClientWithProvider( // Overview can surface that suppression is happening. if (breakerIsOpen()) { maybeEmitCircuitOpenStatus(opts, op); - if (canUseHostFallback()) { + if (canUseHostFallback(opts)) { return callHostFallback(makeBreakerOpenError(), messages, input, opts, op, { keepBreakerOpen: true, notifyError: false, @@ -347,7 +355,7 @@ export function createLlmClientWithProvider( }); return { completion }; } catch (err) { - if (shouldFallback(err, config, provider.name)) { + if (!callerBudgetIsGone(opts) && shouldFallback(err, config, provider.name)) { const primaryTerminal = breakerIsTerminal(err); if (primaryTerminal) breakerTrip(err); try { diff --git a/apps/memos-local-plugin/core/llm/host-bridge.ts b/apps/memos-local-plugin/core/llm/host-bridge.ts index 51d9ca82c..325f76341 100644 --- a/apps/memos-local-plugin/core/llm/host-bridge.ts +++ b/apps/memos-local-plugin/core/llm/host-bridge.ts @@ -18,6 +18,8 @@ export interface HostLlmCompleteInput { temperature?: number; maxTokens?: number; timeoutMs?: number; + /** Absolute caller deadline; adapters must not outlive it. */ + deadlineAt?: number; signal?: AbortSignal; } diff --git a/apps/memos-local-plugin/core/llm/providers/host.ts b/apps/memos-local-plugin/core/llm/providers/host.ts index 002c06888..eae9fdac3 100644 --- a/apps/memos-local-plugin/core/llm/providers/host.ts +++ b/apps/memos-local-plugin/core/llm/providers/host.ts @@ -43,6 +43,7 @@ export class HostLlmProvider implements LlmProvider { temperature: opts.temperature, maxTokens: opts.maxTokens, timeoutMs: ctx.config.timeoutMs, + deadlineAt: ctx.deadlineAt, signal: ctx.signal, }); return { diff --git a/apps/memos-local-plugin/tests/unit/bridge/host-llm-budget.test.ts b/apps/memos-local-plugin/tests/unit/bridge/host-llm-budget.test.ts new file mode 100644 index 000000000..5fcc56818 --- /dev/null +++ b/apps/memos-local-plugin/tests/unit/bridge/host-llm-budget.test.ts @@ -0,0 +1,54 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; + +import { createStdioHostLlmBridge } from "../../../bridge/host-llm.js"; + +describe("stdio host LLM bridge", () => { + afterEach(() => vi.useRealTimers()); + + it("caps reverse RPC at the remaining deadline and forwards cancellation", async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date("2026-01-01T00:00:00.000Z")); + const controller = new AbortController(); + const serverRequest = vi.fn(async () => ({ + text: "host response", + model: "host-model", + durationMs: 10, + })); + const bridge = createStdioHostLlmBridge("test.host.v1", () => ({ serverRequest })); + + const result = await bridge.complete({ + messages: [{ role: "user", content: "hello" }], + timeoutMs: 10_000, + deadlineAt: Date.now() + 2_000, + signal: controller.signal, + }); + + expect(result.text).toBe("host response"); + expect(serverRequest).toHaveBeenCalledWith( + "host.llm.complete", + expect.objectContaining({ timeoutMs: 2_000 }), + { timeoutMs: 2_000, signal: controller.signal }, + ); + }); + + it("rejects an expired deadline without sending reverse RPC", async () => { + const serverRequest = vi.fn(); + const bridge = createStdioHostLlmBridge("test.host.v1", () => ({ serverRequest })); + + await expect( + bridge.complete({ + messages: [{ role: "user", content: "hello" }], + deadlineAt: Date.now() - 1, + }), + ).rejects.toThrow("deadline has passed"); + expect(serverRequest).not.toHaveBeenCalled(); + }); + + it("rejects before stdio is ready", async () => { + const bridge = createStdioHostLlmBridge("test.host.v1", () => null); + + await expect( + bridge.complete({ messages: [{ role: "user", content: "hello" }] }), + ).rejects.toThrow("before stdio server was ready"); + }); +}); diff --git a/apps/memos-local-plugin/tests/unit/bridge/stdio.test.ts b/apps/memos-local-plugin/tests/unit/bridge/stdio.test.ts index 28cdd9a01..41e2a6013 100644 --- a/apps/memos-local-plugin/tests/unit/bridge/stdio.test.ts +++ b/apps/memos-local-plugin/tests/unit/bridge/stdio.test.ts @@ -7,7 +7,7 @@ * • request → error response with stable `code`. * • notifications (events + logs) stream back unordered. */ -import { describe, expect, it, vi } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { PassThrough } from "node:stream"; import { @@ -100,10 +100,12 @@ function wire() { logToStderr: false, }); const client = createStdioClient(serverOut, clientOut); - return { core, server, client, clientOut }; + return { core, server, client, clientOut, serverOut }; } describe("stdio transport", () => { + afterEach(() => vi.useRealTimers()); + it("round-trips a successful request", async () => { const { client } = wire(); const res = await client.request<{ ok: boolean }>("core.init", {}); @@ -159,4 +161,54 @@ describe("stdio transport", () => { expect(server.connected).toBe(false); expect(core.shutdown).toHaveBeenCalled(); }); + + it("cancels a pending server request when its signal aborts", async () => { + const { server } = wire(); + const controller = new AbortController(); + const request = server.serverRequest( + "host.llm.complete", + {}, + { timeoutMs: 1_000, signal: controller.signal }, + ); + + controller.abort(new Error("caller cancelled")); + + await expect(request).rejects.toThrow("caller cancelled"); + }); + + it("does not send a server request when its signal is already aborted", async () => { + const { server, serverOut } = wire(); + const controller = new AbortController(); + controller.abort(new Error("deadline passed")); + + await expect( + server.serverRequest( + "host.llm.complete", + {}, + { timeoutMs: 1_000, signal: controller.signal }, + ), + ).rejects.toThrow("deadline passed"); + expect(serverOut.read()).toBeNull(); + }); + + it("cleans up concurrent cancelled server requests", async () => { + vi.useFakeTimers(); + const { server } = wire(); + const controllers = Array.from({ length: 128 }, () => new AbortController()); + const requests = controllers.map((controller) => + server.serverRequest("host.llm.complete", {}, { + timeoutMs: 10_000, + signal: controller.signal, + }), + ); + + for (const controller of controllers) { + controller.abort(new Error("caller cancelled")); + } + + const results = await Promise.allSettled(requests); + expect(results).toHaveLength(128); + expect(results.every((result) => result.status === "rejected")).toBe(true); + expect(vi.getTimerCount()).toBe(0); + }); }); diff --git a/apps/memos-local-plugin/tests/unit/llm/client.test.ts b/apps/memos-local-plugin/tests/unit/llm/client.test.ts index 51051ab00..d4effe804 100644 --- a/apps/memos-local-plugin/tests/unit/llm/client.test.ts +++ b/apps/memos-local-plugin/tests/unit/llm/client.test.ts @@ -244,6 +244,86 @@ describe("llm/client", () => { expect(client.stats().hostFallbacks).toBe(1); }); + it("does NOT fall back after the caller aborts the primary request", async () => { + const controller = new AbortController(); + const primaryError = new MemosError( + ERROR_CODES.LLM_TIMEOUT, + "openai_compatible request was cancelled", + { cancelled: true }, + ); + const provider: LlmProvider = { + name: "openai_compatible", + async complete(_messages, _input, ctx) { + expect(ctx.signal?.aborted).toBe(false); + controller.abort(new Error("caller cancelled")); + throw primaryError; + }, + }; + let hostCalls = 0; + registerHostLlmBridge({ + id: "test.host.v1", + async complete() { + hostCalls++; + return { text: "too late", model: "host-m", durationMs: 1 }; + }, + }); + const client = createLlmClientWithProvider(cfg({ fallbackToHost: true }), provider); + + await expect(client.complete("ping", { signal: controller.signal })).rejects.toBe( + primaryError, + ); + expect(hostCalls).toBe(0); + expect(client.stats().hostFallbacks).toBe(0); + }); + + it("does NOT fall back after the caller deadline has passed", async () => { + const primaryError = new MemosError( + ERROR_CODES.LLM_TIMEOUT, + "openai_compatible request timed out", + ); + const thrower = new ThrowingProvider(primaryError); + let hostCalls = 0; + registerHostLlmBridge({ + id: "test.host.v1", + async complete() { + hostCalls++; + return { text: "too late", model: "host-m", durationMs: 1 }; + }, + }); + const client = createLlmClientWithProvider(cfg({ fallbackToHost: true }), thrower); + + await expect(client.complete("ping", { deadlineAt: Date.now() - 1 })).rejects.toBe( + primaryError, + ); + expect(hostCalls).toBe(0); + expect(client.stats().hostFallbacks).toBe(0); + }); + + it("still falls back on a provider timeout while the caller budget remains", async () => { + const thrower = new ThrowingProvider( + new MemosError(ERROR_CODES.LLM_TIMEOUT, "provider timed out"), + ); + const controller = new AbortController(); + const deadlineAt = Date.now() + 10_000; + let seen: { signal?: AbortSignal; deadlineAt?: number } | undefined; + registerHostLlmBridge({ + id: "test.host.v1", + async complete(input) { + seen = input; + return { text: "host response", model: "host-m", durationMs: 1 }; + }, + }); + const client = createLlmClientWithProvider(cfg({ fallbackToHost: true }), thrower); + + const result = await client.complete("ping", { + signal: controller.signal, + deadlineAt, + }); + + expect(result.servedBy).toBe("host_fallback"); + expect(seen).toMatchObject({ signal: controller.signal, deadlineAt }); + }); + it("does NOT fall back when primary throws a non-transient error", async () => { const thrower = new ThrowingProvider( new MemosError(ERROR_CODES.INVALID_ARGUMENT, "bad payload"),