diff --git a/README.md b/README.md index 989c1ec..a521c3e 100644 --- a/README.md +++ b/README.md @@ -271,6 +271,8 @@ Built-in providers today: - `pi` via `LOBSTER_PI_LLM_ADAPTER_URL` (typically supplied by the Pi extension) - `http` via `LOBSTER_LLM_ADAPTER_URL` +A host embedding Lobster can supply its own adapters through `ctx.llmAdapters`. Step `timeout_ms` and workflow cancellation reach an adapter as `ctx.signal`: Lobster stops waiting as soon as that signal aborts, so the step fails or retries on time either way, but it cannot cancel work an adapter has already started. An injected adapter should observe `ctx.signal` and abort its own request — otherwise a timed-out step can leave a model call running, and billed, in the background. + Workflow `_meta.cost` and `cost_limit` use a static pricing table plus optional overrides from `LOBSTER_LLM_PRICING_JSON`, for example `{"my-model":{"input":1.0,"output":2.0}}` in USD per million tokens. Unknown or missing model IDs still record token counts with zero estimated cost, but Lobster warns on stderr so stale or missing pricing does not fail silently. A cached or replayed answer is not billed again: a model call is counted once, in the run that made it, however many later steps re-emit its answer. This holds while the answer stays inside Lobster: through pipelines, renderers, projections, run state, `workflow:` steps and a resume. It does not survive a stage that hands the items to an external process and reads them back — `exec --stdin json --json ...` — because what comes back is whatever that process printed, and Lobster cannot tell a faithful copy of a replay from a fresh claim to have made the call. Such a step is billed as a call, which is what earlier versions did everywhere. A `workflow:` step counts what its sub-workflow spent, and a replay the sub-workflow returned is not billed a second time by the run that composed it. Spend also survives a pause — a workflow that stops at an approval or `input` gate keeps what it has recorded, so `_meta.cost` covers the whole run after a resume and `cost_limit` applies to the whole run rather than to the steps after the last gate. diff --git a/src/commands/stdlib/llm_invoke.ts b/src/commands/stdlib/llm_invoke.ts index ff8bcb7..1ffa3f2 100644 --- a/src/commands/stdlib/llm_invoke.ts +++ b/src/commands/stdlib/llm_invoke.ts @@ -739,6 +739,10 @@ async function runLlmInvoke({ config: CommandConfig; }) { const env = ctx.env ?? process.env; + const signal: AbortSignal | undefined = ctx?.signal; + // Run-state and cache hits return before any adapter call, so a cancelled run + // would otherwise still finish as a success. + throwIfCancelled(signal); const provider = resolveProvider(args, env, config.defaultProvider, ctx); const adapter = resolveAdapter({ provider, env, args, config, ctx }); const prompt = extractPrompt(args); @@ -785,6 +789,9 @@ async function runLlmInvoke({ const inputArtifacts: any[] = []; for await (const item of input) inputArtifacts.push(item); + // Draining pipeline input waits on the upstream step, so a timeout can fire + // here. Re-check before the reuse lookups below can answer with a success. + throwIfCancelled(signal); const normalizedArtifacts = [...inputArtifacts, ...providedArtifacts].map(normalizeArtifact); const artifactHashes = normalizedArtifacts.map(hashArtifact); @@ -801,6 +808,9 @@ async function runLlmInvoke({ if (stateKey && !forceRefresh) { const stored = await readReusableLlmState(env, stateKey, ctx.signal); + // Reading run state is I/O of unbounded duration; a signal that aborted + // during it must not be overtaken by the replay below. + throwIfCancelled(signal); const reused = pickReusableState(stored, cacheKey, config.stateType); if (reused) { const replay: LlmProvenance = { cacheKey, replayed: true }; @@ -816,6 +826,7 @@ async function runLlmInvoke({ if (!disableCache && !forceRefresh) { const cache = await readCacheEntry(env, cacheKey, config.cacheNamespace, ctx.signal); + throwIfCancelled(signal); if (cache) { const replay: LlmProvenance = { cacheKey, replayed: true }; return { @@ -849,6 +860,7 @@ async function runLlmInvoke({ let lastValidationErrors: string[] = []; while (true) { + throwIfCancelled(signal); attempt += 1; if (attempt > 1) { payload.retryContext = { @@ -862,9 +874,16 @@ async function runLlmInvoke({ let responseEnvelope: LlmResponseEnvelope; try { ctx.signal?.throwIfAborted(); - responseEnvelope = await adapter.invoke({ env, args, payload, signal: ctx.signal }); + responseEnvelope = await abortable( + adapter.invoke({ env, args, payload, signal: ctx.signal }), + signal, + ); ctx.signal?.throwIfAborted(); } catch (err: any) { + // Cancellation is the caller's error, not an adapter failure: surface it + // as an abort so workflow timeout and abort handling still recognizes it, + // rather than wrapping it in a "request failed" adapter error. + if (signal?.aborted) throw asCancellation(err, signal); throw new Error(`${config.name} request failed: ${err?.message ?? String(err)}`); } @@ -938,6 +957,66 @@ async function runLlmInvoke({ } } +/** + * Build the error a cancelled run rejects with. + * + * The workflow runner only treats an error named `AbortError` or coded + * `ABORT_ERR` as cancellation; everything else follows the step's retry and + * `on_error` policy. A host may abort with any reason it likes, so passing + * `signal.reason` straight through means `controller.abort(new Error("stop"))` + * leaves a cancelled step looking like an ordinary failure -- retried, or + * swallowed by `on_error: continue` and reported as a successful run. Keep the + * host's message and hang its reason off `cause`, but mark the rejection so the + * runner recognizes it. + */ +function cancellationError(signal: AbortSignal): unknown { + const reason: any = signal.reason; + if (reason === undefined || reason === null) { + return new DOMException("The operation was aborted.", "AbortError"); + } + if (reason?.name === "AbortError" || reason?.code === "ABORT_ERR") return reason; + const message = + typeof reason?.message === "string" && reason.message ? reason.message : String(reason); + const error: any = new Error(message, { cause: reason }); + error.name = "AbortError"; + error.code = "ABORT_ERR"; + return error; +} + +/** Rethrow `err` when it already reads as cancellation, else the run's reason. */ +function asCancellation(err: any, signal: AbortSignal): unknown { + if (err?.name === "AbortError" || err?.code === "ABORT_ERR") return err; + return cancellationError(signal); +} + +function throwIfCancelled(signal?: AbortSignal): void { + if (signal?.aborted) throw cancellationError(signal); +} + +/** + * Reject as soon as the run is cancelled instead of waiting for `promise`. + * HTTP adapters are cancelled at the socket, but an injected `ctx.llmAdapters` + * adapter may ignore `ctx.signal` entirely; without this, one of those keeps a + * timed-out step waiting for as long as it likes. + */ +function abortable(promise: Promise, signal?: AbortSignal): Promise { + if (!signal) return promise; + return new Promise((resolve, reject) => { + const onAbort = () => reject(cancellationError(signal)); + // Observe `promise` before anything can settle the wrapper, including the + // already-aborted case below. An adapter can cancel the run from inside its own + // `invoke` and reject afterwards; rejecting here without watching that promise + // leaves the rejection unhandled, which ends the process under Node's default + // handling -- long after this step was cancelled cleanly. + promise.then(resolve, reject).finally(() => signal.removeEventListener("abort", onAbort)); + if (signal.aborted) { + onAbort(); + return; + } + signal.addEventListener("abort", onAbort, { once: true }); + }); +} + function resolveProvider( args: any, env: any, @@ -983,6 +1062,7 @@ function resolveAdapter({ config: CommandConfig; ctx: any; }): Adapter { + const signal: AbortSignal | undefined = ctx?.signal; const direct = getDirectAdapter(ctx, provider); if (direct) { const invoke = typeof direct === "function" ? direct : direct.invoke; diff --git a/test/llm_invoke.test.ts b/test/llm_invoke.test.ts index 98892b7..1a85515 100644 --- a/test/llm_invoke.test.ts +++ b/test/llm_invoke.test.ts @@ -422,6 +422,380 @@ test("llm.invoke uses Pi adapter over local HTTP bridge", async () => { } }); +test( + "llm.invoke aborts the in-flight adapter request when ctx.signal aborts", + { timeout: 20_000 }, + async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + + const stalled = createStalledAdapter(); + await stalled.listen(); + + const controller = new AbortController(); + try { + const pending = cmd.run({ + input: streamOf([]), + args: { _: [], provider: "http", prompt: "Summarize", "disable-cache": true }, + ctx: { + ...baseCtx( + { LOBSTER_LLM_ADAPTER_URL: `http://127.0.0.1:${stalled.port}/invoke` }, + registry, + ), + signal: controller.signal, + }, + } as any); + + await stalled.requestReceived; + controller.abort(); + + const err = await pending.then( + () => null, + (e: any) => e, + ); + assert.ok(err, "llm.invoke should reject once the run is aborted"); + assert.ok( + err.name === "AbortError" || err.code === "ABORT_ERR", + `expected an abort error, got ${err.name}: ${err.message}`, + ); + // The abort must stay recognizable; wrapping it hides the cancellation from + // workflow timeout and abort handling. + assert.doesNotMatch(String(err.message), /request failed/); + await stalled.requestClosed; + } finally { + controller.abort(); + await stalled.close(); + } + }, +); + +test( + "llm.invoke stops waiting for a direct adapter that ignores ctx.signal", + { timeout: 20_000 }, + async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + + let invoked = () => {}; + const adapterInvoked = new Promise((resolve) => (invoked = resolve)); + // A supported ctx.llmAdapters adapter that never resolves and never looks + // at ctx.signal. + const stubborn = { + invoke() { + invoked(); + return new Promise(() => {}); + }, + }; + + const controller = new AbortController(); + const pending = cmd.run({ + input: streamOf([]), + args: { _: [], provider: "stubborn", prompt: "Summarize", "disable-cache": true }, + ctx: { + ...baseCtx({}, registry), + llmAdapters: { stubborn }, + signal: controller.signal, + }, + } as any); + + await adapterInvoked; + controller.abort(); + + const err = await pending.then( + () => null, + (e: any) => e, + ); + assert.ok(err, "llm.invoke should reject once the run is aborted"); + assert.ok( + err.name === "AbortError" || err.code === "ABORT_ERR", + `expected an abort error, got ${err.name}: ${err.message}`, + ); + assert.doesNotMatch(String(err.message), /request failed/); + }, +); + +test("llm.invoke observes a direct adapter that rejects after aborting the run", async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + + const controller = new AbortController(); + // An adapter that cancels the run from inside its own invoke and then fails: + // the shape of an SDK client that tears itself down on a fatal error. + const selfaborting = { + invoke() { + controller.abort(); + return Promise.reject(new Error("adapter tore down its client")); + }, + }; + + const unhandled: any[] = []; + const onUnhandled = (reason: any) => unhandled.push(reason); + process.on("unhandledRejection", onUnhandled); + try { + const err = await cmd + .run({ + input: streamOf([]), + args: { _: [], provider: "selfaborting", prompt: "Summarize", "disable-cache": true }, + ctx: { + ...baseCtx({}, registry), + llmAdapters: { selfaborting }, + signal: controller.signal, + }, + } as any) + .then( + () => null, + (e: any) => e, + ); + assert.ok( + err?.name === "AbortError" || err?.code === "ABORT_ERR", + `expected an abort error, got ${err?.name}: ${err?.message}`, + ); + // The adapter's own rejection must not outlive the cancelled step. Nothing is + // waiting on it once the abort has won the race, and an unobserved rejection + // ends the process under Node's default handling. + await new Promise((resolve) => setTimeout(resolve, 50)); + assert.equal( + unhandled.length, + 0, + `adapter rejection went unhandled: ${unhandled[0]?.message ?? unhandled[0]}`, + ); + } finally { + process.off("unhandledRejection", onUnhandled); + } +}); + +test("llm.invoke does not complete from cache when ctx.signal is already aborted", async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + const cacheDir = await mkdtemp(path.join(tmpdir(), "lobster-cache-")); + await mkdir(path.join(cacheDir, "llm.invoke"), { recursive: true }); + + let requests = 0; + const server = http.createServer((req, res) => { + requests += 1; + req.resume(); + req.on("end", () => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + ok: true, + result: { runId: "cached_1", output: { text: "hello", data: null } }, + }), + ); + }); + }); + + await new Promise((resolve) => server.listen(0, resolve)); + const addr = server.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + const env = { + LOBSTER_LLM_ADAPTER_URL: `http://127.0.0.1:${port}/invoke`, + LOBSTER_CACHE_DIR: cacheDir, + }; + const args = { _: [], provider: "http", prompt: "Summarize" }; + + try { + // Populate the cache, then repeat the same call on a cancelled run. + const warm = await cmd.run({ + input: streamOf([]), + args, + ctx: baseCtx(env, registry), + } as any); + assert.equal((await collect(warm.output!))[0].cached, false); + assert.equal(requests, 1); + + const controller = new AbortController(); + controller.abort(); + await assert.rejects( + () => + cmd.run({ + input: streamOf([]), + args, + ctx: { ...baseCtx(env, registry), signal: controller.signal }, + } as any), + (err: any) => err?.name === "AbortError" || err?.code === "ABORT_ERR", + ); + assert.equal(requests, 1, "the cached run must not reach the adapter either"); + } finally { + await rm(cacheDir, { recursive: true, force: true }); + await closeServer(server); + } +}); + +test("llm.invoke does not replay a cache hit when ctx.signal aborts while input is drained", async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + const cacheDir = await mkdtemp(path.join(tmpdir(), "lobster-cache-")); + await mkdir(path.join(cacheDir, "llm.invoke"), { recursive: true }); + + let requests = 0; + const server = http.createServer((req, res) => { + requests += 1; + req.resume(); + req.on("end", () => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + ok: true, + result: { runId: "cached_1", output: { text: "hello", data: null } }, + }), + ); + }); + }); + + await new Promise((resolve) => server.listen(0, resolve)); + const addr = server.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + const env = { + LOBSTER_LLM_ADAPTER_URL: `http://127.0.0.1:${port}/invoke`, + LOBSTER_CACHE_DIR: cacheDir, + }; + const args = { _: [], provider: "http", prompt: "Summarize" }; + + try { + const warm = await cmd.run({ + input: streamOf([]), + args, + ctx: baseCtx(env, registry), + } as any); + assert.equal((await collect(warm.output!))[0].cached, false); + assert.equal(requests, 1); + + // The step timeout fires while the upstream step is still producing input, + // which is after the entry check and before the cache lookup. + const controller = new AbortController(); + const abortingInput = (async function* () { + await new Promise((resolve) => setTimeout(resolve, 5)); + controller.abort(); + })(); + + await assert.rejects( + () => + cmd.run({ + input: abortingInput, + args, + ctx: { ...baseCtx(env, registry), signal: controller.signal }, + } as any), + (err: any) => err?.name === "AbortError" || err?.code === "ABORT_ERR", + ); + assert.equal(requests, 1, "a cancelled run must not reach the adapter either"); + } finally { + await rm(cacheDir, { recursive: true, force: true }); + await closeServer(server); + } +}); + +test("llm.invoke does not replay run state when ctx.signal aborts while input is drained", async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + const stateDir = await mkdtemp(path.join(tmpdir(), "lobster-state-")); + + let requests = 0; + const server = http.createServer((req, res) => { + requests += 1; + req.resume(); + req.on("end", () => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + ok: true, + result: { runId: "state_1", output: { text: "hello", data: null } }, + }), + ); + }); + }); + + await new Promise((resolve) => server.listen(0, resolve)); + const addr = server.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + const env = { + LOBSTER_LLM_ADAPTER_URL: `http://127.0.0.1:${port}/invoke`, + LOBSTER_STATE_DIR: stateDir, + }; + const args = { + _: [], + provider: "http", + prompt: "Summarize", + "state-key": "step-1", + "disable-cache": true, + }; + + try { + const warm = await cmd.run({ + input: streamOf([]), + args, + ctx: baseCtx(env, registry), + } as any); + assert.equal((await collect(warm.output!))[0].cached, false); + assert.equal(requests, 1); + + const controller = new AbortController(); + const abortingInput = (async function* () { + await new Promise((resolve) => setTimeout(resolve, 5)); + controller.abort(); + })(); + + await assert.rejects( + () => + cmd.run({ + input: abortingInput, + args, + ctx: { ...baseCtx(env, registry), signal: controller.signal }, + } as any), + (err: any) => err?.name === "AbortError" || err?.code === "ABORT_ERR", + ); + assert.equal(requests, 1, "a cancelled run must not reach the adapter either"); + } finally { + await rm(stateDir, { recursive: true, force: true }); + await closeServer(server); + } +}); + +test("llm.invoke does not call the adapter when ctx.signal is already aborted", async () => { + const registry = createDefaultRegistry(); + const cmd = registry.get("llm.invoke"); + assert.ok(cmd); + + let requests = 0; + const server = http.createServer((req, res) => { + requests += 1; + req.resume(); + res.writeHead(200, { "content-type": "application/json" }); + res.end(JSON.stringify({ ok: true, result: { runId: "r1", output: { data: {} } } })); + }); + + await new Promise((resolve) => server.listen(0, resolve)); + const addr = server.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + + try { + const controller = new AbortController(); + controller.abort(); + + await assert.rejects( + () => + cmd.run({ + input: streamOf([]), + args: { _: [], provider: "http", prompt: "Summarize", "disable-cache": true }, + ctx: { + ...baseCtx({ LOBSTER_LLM_ADAPTER_URL: `http://127.0.0.1:${port}/invoke` }, registry), + signal: controller.signal, + }, + } as any), + (err: any) => err?.name === "AbortError" || err?.code === "ABORT_ERR", + ); + assert.equal(requests, 0, "an already-cancelled run must not reach the adapter"); + } finally { + await closeServer(server); + } +}); + test("llm.invoke does not retry schema validation after adapter cancellation", async () => { const registry = createDefaultRegistry(); const cmd = registry.get("llm.invoke"); @@ -878,6 +1252,44 @@ function baseCtx(envOverrides: Record, registry?: any) { }; } +// An adapter that accepts the request and never answers, so the only thing that +// can end the call is the caller cancelling it. +function createStalledAdapter() { + const sockets = new Set(); + let markReceived = () => {}; + let markClosed = () => {}; + const requestReceived = new Promise((resolve) => (markReceived = resolve)); + const requestClosed = new Promise((resolve) => (markClosed = resolve)); + + const server = http.createServer((req, res) => { + req.resume(); + req.on("end", () => { + res.on("close", () => markClosed()); + markReceived(); + }); + }); + server.on("connection", (socket) => { + sockets.add(socket); + socket.on("close", () => sockets.delete(socket)); + }); + + return { + requestReceived, + requestClosed, + port: 0, + async listen() { + await new Promise((resolve) => server.listen(0, resolve)); + const addr = server.address(); + this.port = typeof addr === "object" && addr ? addr.port : 0; + }, + async close() { + for (const socket of sockets) socket.destroy(); + if (!server.listening) return; + await new Promise((resolve) => server.close(() => resolve())); + }, + }; +} + async function closeServer(server: http.Server) { if (!server.listening) return; await new Promise((resolve) => server.close(() => resolve())); diff --git a/test/step_timeout.test.ts b/test/step_timeout.test.ts index 1a36d7d..b418097 100644 --- a/test/step_timeout.test.ts +++ b/test/step_timeout.test.ts @@ -1,6 +1,7 @@ import test from "node:test"; import assert from "node:assert/strict"; import { promises as fsp } from "node:fs"; +import http from "node:http"; import os from "node:os"; import path from "node:path"; import { PassThrough } from "node:stream"; @@ -13,6 +14,8 @@ async function runWorkflow( opts?: { signal?: AbortSignal; dryRun?: boolean; + env?: Record; + llmAdapters?: Record; }, ) { const tmpDir = await fsp.mkdtemp(path.join(os.tmpdir(), "lobster-step-timeout-")); @@ -30,10 +33,11 @@ async function runWorkflow( stdin: process.stdin, stdout: process.stdout, stderr, - env: { ...process.env, LOBSTER_STATE_DIR: stateDir }, + env: { ...process.env, LOBSTER_STATE_DIR: stateDir, ...opts?.env }, mode: "tool", signal: opts?.signal, dryRun: opts?.dryRun, + llmAdapters: opts?.llmAdapters, registry: createDefaultRegistry(), }, }); @@ -188,6 +192,51 @@ test("external abort still propagates when timeout is configured", async () => { ); }); +test("timed-out llm.invoke step stops waiting for the adapter", { timeout: 20_000 }, async () => { + const sockets = new Set(); + let markClosed = () => {}; + const requestClosed = new Promise((resolve) => (markClosed = resolve)); + + // An adapter that accepts the request and never answers. + const server = http.createServer((req, res) => { + req.resume(); + req.on("end", () => res.on("close", () => markClosed())); + }); + server.on("connection", (socket) => { + sockets.add(socket); + socket.on("close", () => sockets.delete(socket)); + }); + + await new Promise((resolve) => server.listen(0, resolve)); + const addr = server.address(); + const port = typeof addr === "object" && addr ? addr.port : 0; + + try { + const started = Date.now(); + await assert.rejects( + () => + runWorkflow( + { + steps: [ + { + id: "ask", + pipeline: 'llm.invoke --prompt "Summarize" --disable-cache', + timeout_ms: 300, + }, + ], + }, + { env: { LOBSTER_LLM_ADAPTER_URL: `http://127.0.0.1:${port}/invoke` } }, + ), + /timed out after 300ms/, + ); + assert.ok(Date.now() - started < 10_000, "the step must not wait for the adapter"); + await requestClosed; + } finally { + for (const socket of sockets) socket.destroy(); + await new Promise((resolve) => server.close(() => resolve())); + } +}); + test("dry-run renders timeout and on_error details", async () => { const { stderrOutput } = await runWorkflow( { @@ -205,3 +254,46 @@ test("dry-run renders timeout and on_error details", async () => { assert.match(stderrOutput, /timeout: 5000ms/); assert.match(stderrOutput, /on_error: continue/); }); + +test( + "external abort with a custom reason still stops the workflow", + { timeout: 20_000 }, + async () => { + let invoked = () => {}; + const adapterInvoked = new Promise((resolve) => (invoked = resolve)); + // Ignores ctx.signal, so what the runner classifies is llm.invoke's own + // abort rejection rather than a cancelled socket. + const stubborn = { + invoke() { + invoked(); + return new Promise(() => {}); + }, + }; + + const controller = new AbortController(); + // A host may abort with any reason. A plain Error must still read as + // cancellation, not as an ordinary step failure. + void adapterInvoked.then(() => controller.abort(new Error("cancelled by host"))); + + await assert.rejects( + () => + runWorkflow( + { + steps: [ + { + id: "ask", + // on_error: continue is what makes a misclassified abort + // visible -- the run would otherwise report success. + pipeline: 'llm.invoke --provider stubborn --prompt "Summarize" --disable-cache', + on_error: "continue", + }, + ], + }, + { signal: controller.signal, llmAdapters: { stubborn } }, + ), + (err: any) => + (err?.name === "AbortError" || err?.code === "ABORT_ERR") && + /cancelled by host/.test(String(err?.message)), + ); + }, +);