From 5a033b948ba44c92143dc149402e82daea0736b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elberte=20Pl=C3=ADnio?= Date: Sun, 19 Jul 2026 03:00:42 -0300 Subject: [PATCH 1/2] feat(pi-kit): non-blocking lanes_spawn with lanes_status, lanes_wait, lanes_abandon lanes_spawn now returns immediately; lanes run in the background and survive ESC/turn cancellation. lanes_status answers questions about the run or a single lane mid-flight, lanes_wait blocks for final results (abort detaches without stopping lanes), lanes_abandon is the only way to stop lanes. On settlement an idle session gets a one-shot nudge to collect results. --- packages/pi-kit/extensions/lanes.ts | 287 +++++++++++++++---- packages/pi-kit/test/lanes-extension.test.ts | 189 ++++++++++++ 2 files changed, 421 insertions(+), 55 deletions(-) create mode 100644 packages/pi-kit/test/lanes-extension.test.ts diff --git a/packages/pi-kit/extensions/lanes.ts b/packages/pi-kit/extensions/lanes.ts index b5dfa34..9d775cb 100644 --- a/packages/pi-kit/extensions/lanes.ts +++ b/packages/pi-kit/extensions/lanes.ts @@ -10,6 +10,14 @@ interface LastRun { runner: LaneRunner; projection: RunProjection; hidden: boolean; + /** Resolves when every lane settles. */ + settled: Promise; + /** True once run_end was journaled and the summary is available. */ + ended: boolean; + /** True once the model has consumed the final results (wait or nudge). */ + reported: boolean; + /** Number of lanes_wait calls currently attached; suppresses the settle nudge. */ + waiters: number; } interface SpawnLaneInput { @@ -136,7 +144,7 @@ export default function lanesExtension(pi: ExtensionAPI): void { pi.registerTool({ name: "lanes_spawn", label: "Spawn lanes", - description: `Run independent Pi lanes concurrently. You MUST state an explicit model and effort for every lane. Any effort off|minimal|low|medium|high|xhigh is allowed per lane (the value in parentheses is only that model's starting prior). Models: ${MODEL_TABLE.map((row) => `${row.selector} (prior ${row.prior})`).join(", ")}.`, + description: `Spawn independent Pi lanes that run concurrently in the background (non-blocking: this tool returns immediately; lanes survive turn cancellation). Collect results with lanes_wait, inspect progress or a specific lane with lanes_status, stop lanes only with lanes_abandon. While lanes run you can keep working with the user. You MUST state an explicit model and effort for every lane. Any effort off|minimal|low|medium|high|xhigh is allowed per lane (the value in parentheses is only that model's starting prior). Models: ${MODEL_TABLE.map((row) => `${row.selector} (prior ${row.prior})`).join(", ")}.`, parameters: Type.Object({ lanes: Type.Array( Type.Object({ @@ -171,71 +179,240 @@ export default function lanesExtension(pi: ExtensionAPI): void { }; } + if (lastRun && !lastRun.ended) { + return { + content: [ + { + type: "text" as const, + text: `Run ${lastRun.id} is still active. Use lanes_status to inspect it, lanes_wait to collect it, or lanes_abandon to stop it before spawning a new run.`, + }, + ], + details: { activeRun: lastRun.id }, + isError: true, + }; + } + const runId = newRunId(); const startedAt = Date.now(); - let view: LastRun | undefined; - let timer: NodeJS.Timeout | undefined; - let runner: LaneRunner | undefined; - const onAbort = () => runner?.abandonAll("aborted by parent"); - try { - appendEvent({ - v: 1, - t: new Date().toISOString(), - run: runId, - type: "run_created", - lanes: specs.length, - origin: "lanes_spawn", - }); + appendEvent({ + v: 1, + t: new Date().toISOString(), + run: runId, + type: "run_created", + lanes: specs.length, + origin: "lanes_spawn", + }); - const activeRunner = new LaneRunner({ - runId, - append: appendEvent, - onUpdate() { - if (view && lastRun === view) renderWidget(ctx, view); - }, - }); - runner = activeRunner; - view = { id: runId, runner: activeRunner, projection: activeRunner.projection(), hidden: false }; - lastRun = view; - renderWidget(ctx, view); + let view: LastRun; + const activeRunner = new LaneRunner({ + runId, + append: appendEvent, + onUpdate() { + if (view && lastRun === view) renderWidget(ctx, view); + }, + }); - signal?.addEventListener("abort", onAbort, { once: true }); - timer = ctx.hasUI ? setInterval(() => renderWidget(ctx, view!), 1_000) : undefined; - timer?.unref(); - const projections = await activeRunner.dispatch(specs); + // Deliberately NOT wired to the tool-call abort signal: lanes survive + // ESC / turn cancellation. Stopping a run is explicit via lanes_abandon + // or /lanes abandon. + const timer = ctx.hasUI ? setInterval(() => renderWidget(ctx, view), 1_000) : undefined; + timer?.unref(); - const ok = projections.every((lane) => lane.state === "done"); - appendEvent({ - v: 1, - t: new Date().toISOString(), - run: runId, - type: "run_end", - ok, - durationMs: Date.now() - startedAt, + const settled = activeRunner + .dispatch(specs) + .then((projections) => { + const ok = projections.every((lane) => lane.state === "done"); + appendEvent({ + v: 1, + t: new Date().toISOString(), + run: runId, + type: "run_end", + ok, + durationMs: Date.now() - startedAt, + }); + return projections; + }) + .catch((error: unknown) => { + activeRunner.abandonAll("lanes run error"); + appendEvent({ + v: 1, + t: new Date().toISOString(), + run: runId, + type: "run_end", + ok: false, + durationMs: Date.now() - startedAt, + }); + return [...activeRunner.projection().lanes.values()]; + }) + .finally(() => { + clearInterval(timer); + view.ended = true; + // Widget stays visible after completion; only /lanes hide or a new run replaces it. + if (lastRun === view) renderWidget(ctx, view); + // Nudge the model when it is idle and never collected the results. + if (lastRun === view && !view.reported && view.waiters === 0) { + try { + if (ctx.isIdle()) { + view.reported = true; + pi.sendUserMessage( + `[lanes] run ${runId} settled. Call lanes_status for the results, then continue where we left off.`, + ); + } else { + pi.sendUserMessage( + `[lanes] run ${runId} settled — results available via lanes_status/lanes_wait when convenient.`, + { deliverAs: "followUp" }, + ); + } + } catch { + // Nudge is best-effort; lanes_status still works. + } + } }); - // Widget stays visible after completion; only /lanes hide or a new run replaces it. - renderWidget(ctx, view); - const text = projections - .map((lane) => { - const answer = lane.answer?.replace(/\s+/g, " ").trim().slice(0, 400) ?? ""; - const failure = lane.state !== "done" ? ` · ${lane.abandonReason ?? "failed"}` : ""; - return `${lane.spec.lane}: ${lane.state} · ${lane.spec.model}:${lane.spec.effort} · tok ${lane.tokensIn}/${lane.tokensOut} · $${lane.cost.toFixed(4)}${failure}${answer ? ` · ${answer}` : ""}`; - }) - .join("\n"); - return { content: [{ type: "text" as const, text }], details: projections, isError: !ok }; - } catch (error) { - runner?.abandonAll("lanes_spawn error"); - const message = error instanceof Error ? error.message : String(error); + view = { + id: runId, + runner: activeRunner, + projection: activeRunner.projection(), + hidden: false, + settled, + ended: false, + reported: false, + waiters: 0, + }; + lastRun = view; + renderWidget(ctx, view); + + const text = [ + `Run ${runId} spawned with ${specs.length} lane${specs.length === 1 ? "" : "s"} (non-blocking):`, + ...specs.map((spec) => ` ${spec.lane}: ${spec.model}:${spec.effort}`), + "Lanes run in the background and survive turn cancellation.", + "Use lanes_status to check progress or answer questions about a lane, lanes_wait to block for final results, lanes_abandon to stop lanes.", + "You can keep working with the user meanwhile — do not poll lanes_status in a loop.", + ].join("\n"); + return { content: [{ type: "text" as const, text }], details: { run: runId, lanes: specs.length } }; + }, + }); + + function summarizeLane(lane: LaneProjection, verbose: boolean): string { + const answerLimit = verbose ? 4_000 : 400; + const answer = lane.answer?.replace(/\s+/g, " ").trim().slice(0, answerLimit) ?? ""; + const failure = lane.state === "failed" || lane.state === "abandoned" ? ` · ${lane.abandonReason ?? "failed"}` : ""; + const activity = + lane.state === "running" ? ` · ${(lane.currentTool ?? lane.lastStatus ?? "working").slice(0, verbose ? 200 : 60)}` : ""; + const duration = lane.durationMs !== undefined ? ` · ${fmtDuration(lane.durationMs)}` : ""; + return `${lane.spec.lane}: ${lane.state} · ${lane.spec.model}:${lane.spec.effort} · tok ${lane.tokensIn}/${lane.tokensOut} · $${lane.cost.toFixed(4)}${duration}${failure}${activity}${answer ? ` · ${answer}` : ""}`; + } + + function summarizeRun(run: LastRun, laneFilter?: string): { text: string; lanes: LaneProjection[]; missing?: string } { + const all = [...run.projection.lanes.values()]; + if (laneFilter) { + const lane = run.projection.lanes.get(laneFilter); + if (!lane) return { text: "", lanes: [], missing: laneFilter }; + return { text: `run ${run.id} (${run.ended ? "ended" : "active"})\n${summarizeLane(lane, true)}`, lanes: [lane] }; + } + const header = `run ${run.id} (${run.ended ? "ended" : "active"}) · $${run.projection.totalCost.toFixed(4)} · tok ${run.projection.totalTokensIn}/${run.projection.totalTokensOut}`; + return { text: [header, ...all.map((lane) => summarizeLane(lane, run.ended))].join("\n"), lanes: all }; + } + + pi.registerTool({ + name: "lanes_status", + label: "Lane status", + description: + "Report live status of the current lane run without blocking: per-lane state, current activity, tokens, cost, and (for settled lanes) answers. Pass `lane` for a detailed view of one lane — use this when the user asks about a specific lane. Do not call in a polling loop; prefer lanes_wait when you only need the final results.", + parameters: Type.Object({ + lane: Type.Optional(Type.String()), + }), + async execute(_toolCallId, params) { + if (!lastRun) { + return { content: [{ type: "text" as const, text: "No lane run has been spawned this session." }], details: undefined }; + } + const { text, missing } = summarizeRun(lastRun, params.lane); + if (missing) { + const known = [...lastRun.projection.lanes.keys()].join(", "); return { - content: [{ type: "text" as const, text: `lanes_spawn failed: ${message}` }], - details: { error: message }, + content: [{ type: "text" as const, text: `Unknown lane "${missing}". Lanes in run ${lastRun.id}: ${known}` }], + details: undefined, isError: true, }; - } finally { - clearInterval(timer); - signal?.removeEventListener("abort", onAbort); } + if (lastRun.ended) lastRun.reported = true; + return { content: [{ type: "text" as const, text }], details: undefined }; + }, + }); + + pi.registerTool({ + name: "lanes_wait", + label: "Wait for lanes", + description: + "Block until the current lane run settles and return the final per-lane results. Call this when you have nothing else to do for the user, or when the user asks for the results. Cancelling this wait does NOT stop the lanes — they keep running and lanes_status/lanes_wait remain available.", + parameters: Type.Object({}), + async execute(_toolCallId, _params, signal) { + if (!lastRun) { + return { content: [{ type: "text" as const, text: "No lane run has been spawned this session." }], details: undefined }; + } + const run = lastRun; + if (!run.ended) { + // Wait for settlement or the tool call being aborted — abort only + // detaches the wait, it never touches the lanes themselves. + run.waiters++; + try { + await Promise.race([ + run.settled, + new Promise((resolve) => { + if (!signal) return; + if (signal.aborted) resolve(); + else signal.addEventListener("abort", () => resolve(), { once: true }); + }), + ]); + } finally { + run.waiters--; + } + if (!run.ended) { + return { + content: [ + { type: "text" as const, text: `Wait detached; run ${run.id} keeps running. Use lanes_status or lanes_wait again.` }, + ], + details: undefined, + }; + } + } + run.reported = true; + const { text, lanes } = summarizeRun(run); + const ok = lanes.every((lane) => lane.state === "done"); + return { content: [{ type: "text" as const, text }], details: lanes, isError: !ok }; + }, + }); + + pi.registerTool({ + name: "lanes_abandon", + label: "Abandon lanes", + description: + "Explicitly stop lanes of the current run. Pass `lane` to stop one lane, omit it to stop all. This is the only way to stop lanes — they are not cancelled by ESC or turn cancellation.", + parameters: Type.Object({ + lane: Type.Optional(Type.String()), + reason: Type.Optional(Type.String()), + }), + async execute(_toolCallId, params, _signal, _onUpdate, ctx) { + if (!lastRun) { + return { content: [{ type: "text" as const, text: "No lane run has been spawned this session." }], details: undefined }; + } + const reason = params.reason?.trim() || "abandoned by model"; + if (params.lane) { + if (!lastRun.projection.lanes.has(params.lane)) { + const known = [...lastRun.projection.lanes.keys()].join(", "); + return { + content: [{ type: "text" as const, text: `Unknown lane "${params.lane}". Lanes: ${known}` }], + details: undefined, + isError: true, + }; + } + lastRun.runner.abandon(params.lane, reason); + } else { + lastRun.runner.abandonAll(reason); + } + renderWidget(ctx, lastRun); + const { text } = summarizeRun(lastRun); + return { content: [{ type: "text" as const, text }], details: undefined }; }, }); diff --git a/packages/pi-kit/test/lanes-extension.test.ts b/packages/pi-kit/test/lanes-extension.test.ts new file mode 100644 index 0000000..4890a1f --- /dev/null +++ b/packages/pi-kit/test/lanes-extension.test.ts @@ -0,0 +1,189 @@ +import { chmod, mkdtemp, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; + +import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; +import lanesExtension from "../extensions/lanes.ts"; + +const streamFixture = fileURLToPath(new URL("./fixtures/lane-stream.jsonl", import.meta.url)); +const originalDataDir = process.env.PIKIT_DATA_DIR; +const originalPath = process.env.PATH; + +interface ToolResult { + content: Array<{ type: string; text: string }>; + details?: unknown; + isError?: boolean; +} + +type ToolExecute = ( + toolCallId: string, + params: Record, + signal: AbortSignal | undefined, + onUpdate: undefined, + ctx: ExtensionContext, +) => Promise; + +function makeHarness(opts: { idle?: boolean } = {}) { + const tools = new Map(); + const sent: Array<{ content: string; options?: unknown }> = []; + const pi = { + registerTool: (tool: { name: string; execute: ToolExecute }) => { + tools.set(tool.name, tool); + }, + registerCommand: () => {}, + sendUserMessage: (content: string, options?: unknown) => { + sent.push({ content, options }); + }, + } as unknown as ExtensionAPI; + const ctx = { + hasUI: false, + isIdle: () => opts.idle ?? true, + ui: {}, + } as unknown as ExtensionContext; + lanesExtension(pi); + const call = (name: string, params: Record = {}, signal?: AbortSignal) => { + const tool = tools.get(name); + if (!tool) throw new Error(`tool ${name} not registered`); + return tool.execute("tc", params, signal, undefined, ctx); + }; + return { call, sent, tools }; +} + +let dataDir: string; + +async function fixtureBinary(name: string, source: string): Promise { + const path = join(dataDir, name); + await writeFile(path, source, "utf8"); + await chmod(path, 0o755); +} + +/** Fake `pi` on PATH that replays the fixture stream after a short delay. */ +async function installSlowPi(delayMs: number): Promise { + await fixtureBinary( + "pi", + `#!/usr/bin/env node\nconst { readFileSync } = require("node:fs");\nconst lines = readFileSync(${JSON.stringify(streamFixture)}, "utf8").trimEnd().split("\\n");\nsetTimeout(() => { for (const line of lines) console.log(line); }, ${delayMs});\n`, + ); + process.env.PATH = `${dataDir}:${originalPath}`; +} + +const spec = { lane: "worker", task: "fixture task", model: "openai-codex/gpt-5.6-sol", effort: "medium" }; + +async function waitFor(check: () => boolean, timeoutMs = 5_000): Promise { + const deadline = Date.now() + timeoutMs; + while (!check()) { + if (Date.now() > deadline) throw new Error("waitFor timeout"); + await new Promise((resolve) => setTimeout(resolve, 25)); + } +} + +beforeEach(async () => { + dataDir = await mkdtemp(join(tmpdir(), "pi-kit-lanes-ext-")); + process.env.PIKIT_DATA_DIR = dataDir; +}); + +afterEach(async () => { + if (originalDataDir === undefined) delete process.env.PIKIT_DATA_DIR; + else process.env.PIKIT_DATA_DIR = originalDataDir; + process.env.PATH = originalPath; + await rm(dataDir, { recursive: true, force: true }); +}); + +describe("async lanes extension", () => { + it("lanes_spawn returns immediately while lanes keep running", async () => { + await installSlowPi(300); + const { call } = makeHarness(); + const started = Date.now(); + const result = await call("lanes_spawn", { lanes: [spec] }); + expect(Date.now() - started).toBeLessThan(250); + expect(result.isError).toBeFalsy(); + expect(result.content[0]!.text).toContain("non-blocking"); + expect(result.content[0]!.text).toContain("worker"); + + const status = await call("lanes_status", {}); + expect(status.content[0]!.text).toMatch(/worker: (queued|running)/); + + const final = await call("lanes_wait", {}); + expect(final.isError).toBeFalsy(); + expect(final.content[0]!.text).toContain("worker: done"); + expect(final.content[0]!.text).toContain("Hello from lane."); + }); + + it("lanes_status with a lane filter returns the detailed single-lane view", async () => { + await installSlowPi(10); + const { call } = makeHarness(); + await call("lanes_spawn", { lanes: [spec] }); + await call("lanes_wait", {}); + + const detail = await call("lanes_status", { lane: "worker" }); + expect(detail.content[0]!.text).toContain("worker: done"); + + const unknown = await call("lanes_status", { lane: "nope" }); + expect(unknown.isError).toBe(true); + expect(unknown.content[0]!.text).toContain('Unknown lane "nope"'); + expect(unknown.content[0]!.text).toContain("worker"); + }); + + it("rejects a second spawn while a run is active, allows it after settlement", async () => { + await installSlowPi(300); + const { call } = makeHarness(); + await call("lanes_spawn", { lanes: [spec] }); + const second = await call("lanes_spawn", { lanes: [spec] }); + expect(second.isError).toBe(true); + expect(second.content[0]!.text).toContain("still active"); + + await call("lanes_wait", {}); + const third = await call("lanes_spawn", { lanes: [{ ...spec, lane: "again" }] }); + expect(third.isError).toBeFalsy(); + await call("lanes_wait", {}); + }); + + it("aborting lanes_wait detaches the wait without stopping lanes", async () => { + await installSlowPi(400); + const { call } = makeHarness(); + await call("lanes_spawn", { lanes: [spec] }); + + const controller = new AbortController(); + const waiting = call("lanes_wait", {}, controller.signal); + controller.abort(); + const detached = await waiting; + expect(detached.content[0]!.text).toContain("keeps running"); + + const final = await call("lanes_wait", {}); + expect(final.content[0]!.text).toContain("worker: done"); + }); + + it("lanes_abandon stops a named lane and reports unknown lanes", async () => { + await installSlowPi(2_000); + const { call } = makeHarness(); + await call("lanes_spawn", { lanes: [spec] }); + + const unknown = await call("lanes_abandon", { lane: "nope" }); + expect(unknown.isError).toBe(true); + + const result = await call("lanes_abandon", { lane: "worker", reason: "test stop" }); + expect(result.content[0]!.text).toContain("worker: abandoned"); + const final = await call("lanes_wait", {}); + expect(final.isError).toBe(true); + expect(final.content[0]!.text).toContain("test stop"); + }); + + it("nudges an idle session once when the run settles unobserved", async () => { + await installSlowPi(50); + const { call, sent } = makeHarness({ idle: true }); + await call("lanes_spawn", { lanes: [spec] }); + await waitFor(() => sent.length > 0); + expect(sent[0]!.content).toContain("settled"); + expect(sent[0]!.content).toContain("lanes_status"); + }); + + it("does not nudge when the results were already collected via lanes_wait", async () => { + await installSlowPi(50); + const { call, sent } = makeHarness({ idle: true }); + await call("lanes_spawn", { lanes: [spec] }); + await call("lanes_wait", {}); + await new Promise((resolve) => setTimeout(resolve, 150)); + expect(sent).toEqual([]); + }); +}); From 03f5a0ce9f8c78de551ae0c3aaea3152652da0f8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Elberte=20Pl=C3=ADnio?= Date: Sun, 19 Jul 2026 03:11:52 -0300 Subject: [PATCH 2/2] fix(pi-kit): abandon active lanes on session_shutdown Lanes survive ESC/turn cancellation by design, but must not outlive their owning session: on shutdown any active run is abandoned and the settle nudge is suppressed. Matches the shutdown pattern already used by the delegation and journal extensions. --- packages/pi-kit/extensions/lanes.ts | 10 ++++++++ packages/pi-kit/test/lanes-extension.test.ts | 26 ++++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/packages/pi-kit/extensions/lanes.ts b/packages/pi-kit/extensions/lanes.ts index 9d775cb..0b26ac8 100644 --- a/packages/pi-kit/extensions/lanes.ts +++ b/packages/pi-kit/extensions/lanes.ts @@ -141,6 +141,16 @@ function renderWidgetUnsafe(ctx: WidgetContext, run: LastRun): void { export default function lanesExtension(pi: ExtensionAPI): void { let lastRun: LastRun | undefined; + // Lanes survive ESC/turn cancellation, but not the session itself: on + // shutdown (/new, reload, exit) any still-active run is abandoned so child + // processes cannot outlive their owning session. + pi.on("session_shutdown", () => { + if (lastRun && !lastRun.ended) { + lastRun.reported = true; // never nudge a session that is going away + lastRun.runner.abandonAll("session ended"); + } + }); + pi.registerTool({ name: "lanes_spawn", label: "Spawn lanes", diff --git a/packages/pi-kit/test/lanes-extension.test.ts b/packages/pi-kit/test/lanes-extension.test.ts index 4890a1f..d50b0ca 100644 --- a/packages/pi-kit/test/lanes-extension.test.ts +++ b/packages/pi-kit/test/lanes-extension.test.ts @@ -33,6 +33,7 @@ function makeHarness(opts: { idle?: boolean } = {}) { tools.set(tool.name, tool); }, registerCommand: () => {}, + on: () => {}, sendUserMessage: (content: string, options?: unknown) => { sent.push({ content, options }); }, @@ -187,3 +188,28 @@ describe("async lanes extension", () => { expect(sent).toEqual([]); }); }); + +describe("session shutdown", () => { + it("abandons an active run and suppresses the nudge on session_shutdown", async () => { + await installSlowPi(2_000); + const handlers = new Map void>(); + const tools = new Map(); + const sent: unknown[] = []; + const pi = { + registerTool: (tool: { name: string; execute: ToolExecute }) => tools.set(tool.name, tool), + registerCommand: () => {}, + sendUserMessage: (content: unknown) => sent.push(content), + on: (name: string, handler: () => void) => handlers.set(name, handler), + } as unknown as ExtensionAPI; + const ctx = { hasUI: false, isIdle: () => true, ui: {} } as unknown as ExtensionContext; + lanesExtension(pi); + await tools.get("lanes_spawn")!.execute("tc", { lanes: [spec] }, undefined, undefined, ctx); + + handlers.get("session_shutdown")!(); + + const status = await tools.get("lanes_status")!.execute("tc", {}, undefined, undefined, ctx); + expect(status.content[0]!.text).toContain("worker: abandoned"); + await new Promise((resolve) => setTimeout(resolve, 200)); + expect(sent).toEqual([]); + }); +});