pi-squire 0.10.0 → 0.12.0

package/src/resume.ts

@@ -0,0 +1,472 @@
+// Resuming an exited delegate: the eligibility checks, the launch that reopens
+// its own session file, and the `delegation_resume` tool. A resume is recorded
+// as a new delegation (`resumedFrom` the old id) that keeps the old session id
+// (`sessionId`), so mail, replies and usage-limit marks still find the session.
+// `index.ts` supplies the pieces it shares with a launch: model choice, trust,
+// the task mail and the window.
+import { randomUUID } from "node:crypto";
+import { statSync } from "node:fs";
+import { isAbsolute } from "node:path";
+
+import type { ExtensionContext, ModelRegistry } from "@earendil-works/pi-coding-agent";
+import { type TObject, Type } from "typebox";
+
+import { type Agent, type Roster, THINKING_LEVELS } from "./agents.ts";
+import { sessionHeaderCwd, sessionIdOf } from "./result-sessions.ts";
+import {
+	deriveState,
+	finishedModel,
+	isStarted,
+	locateSession,
+	type Results,
+	recordedName,
+	type StartedRecorded,
+} from "./results.ts";
+import { type OpenedWindow, requireTmuxVersion, type TmuxClient } from "./tmux.ts";
+import { listWindows, ownWindow, type ToolDef, toolResult, windowState } from "./tracking.ts";
+import type { WindowPlacement } from "./window-placement.ts";
+import type { Workspace } from "./workspace.ts";
+
+/** What a resume asks for; the fields after `task` are optional overrides. */
+export type ResumeRequest = {
+	id: string;
+	task: string;
+	model?: string;
+	thinking?: string;
+	autoExit?: boolean;
+	/** The `delegation_resume` tool call that asked for it; undefined from the command. */
+	toolCallId?: string;
+};
+
+/** A candidate model passed over, as `chooseModel` reports it. */
+type Skipped = { model: string; why: string; reason: string };
+
+/** What a resume reads from the launch path in `index.ts`. */
+export type ResumeDeps = {
+	results: Results;
+	tmux: TmuxClient;
+	/** The `delegate` tool's result schema, which a resume's result extends. */
+	delegateSchema: TObject;
+	/** The refusal text for a session outside tmux. */
+	notInTmux: string;
+	startPolling(): void;
+	/** The window placement for this call; throws on an invalid setting. */
+	placement(): WindowPlacement;
+	roster(ctx: ExtensionContext): Roster;
+	chooseModel(
+		explicit: string | undefined,
+		agent: Agent,
+		models: ModelRegistry,
+		preferred: readonly string[],
+	): Promise<{ model: string; skipped: Skipped[]; notices: string[]; rest: string[] }>;
+	thinkingProblem(value: string): string | undefined;
+	/** The workspace and trust for an existing directory, deciding trust again. */
+	prepareDirectory(
+		cwd: string,
+		ctx: ExtensionContext,
+		signal: AbortSignal | undefined,
+	): Promise<{ workspace: Workspace; trusted: boolean }>;
+	/** Mail the task to a session id; returns the request id. */
+	sendTask(to: string, task: string): string;
+	/** Open the window with `--session <path>`; returns its id and session. */
+	open(input: {
+		id: string;
+		name: string;
+		sessionPath: string;
+		workspace: Workspace;
+		trusted: boolean;
+		cwd: string;
+		sessionDir: string;
+		start: { agent: Agent; model: string; rest: string[]; thinking: string; autoExit: boolean };
+		toolCallId: string | undefined;
+		ctx: ExtensionContext;
+		placement: WindowPlacement;
+	}): Promise<OpenedWindow>;
+};
+
+/** A started resume: what the tool and the command report. */
+export type Resumed = {
+	id: string;
+	resumedFrom: string;
+	agent: string;
+	model: string;
+	skipped: { model: string; reason: string }[];
+	notices: string[];
+	thinking: string;
+	autoExit: boolean;
+	windowId: string;
+	windowName: string;
+	/** The session the window opened in, when tmux could say. */
+	windowSession?: string;
+	cwd: string;
+	sessionDir: string;
+	workspace: Workspace;
+	/** The sentence naming a fallback model, or "". */
+	fallbackNote: string;
+};
+
+/** The started records, in the order recorded. */
+function startedRecords(results: Results): StartedRecorded[] {
+	return results.list().filter(isStarted);
+}
+
+/** The record a `/delegate-resume` id can name: started, with no later resume of it. */
+export function resumableIds(results: Results): string[] {
+	const started = startedRecords(results);
+	return started.filter((record) => !started.some((item) => item.resumedFrom === record.id)).map(({ id }) => id);
+}
+
+/** Steps 1 and 2: the id names a started delegation that nothing resumed. */
+function newestStarted(results: Results, id: string): StartedRecorded {
+	const record = results.find(id);
+	if (record === undefined) {
+		const known = results.list().map((item) => item.id);
+		throw new Error(`no delegation ${JSON.stringify(id)}; known: ${known.join(", ") || "none"}`);
+	}
+	if (!isStarted(record)) {
+		throw new Error(
+			`delegation ${id} is a workspace attempt (${record.attempt.phase}), not a started delegate, so there is no session to resume`,
+		);
+	}
+	const started = startedRecords(results);
+	let newest = record;
+	for (let next = started.find((item) => item.resumedFrom === newest.id); next !== undefined; ) {
+		newest = next;
+		next = started.find((item) => item.resumedFrom === newest.id);
+	}
+	if (newest !== record) {
+		throw new Error(
+			`delegation ${id} was already resumed; resume its newest delegation ${newest.id} instead (it continues the same session)`,
+		);
+	}
+	return record;
+}
+
+/**
+ * Steps 3 and 4: no process of this delegate is left. A running delegate is
+ * mailed instead; tmux that cannot be asked proves nothing; a window with the
+ * record's id and name that still runs its program is refused, and one whose
+ * program exited is killed first.
+ */
+async function ensureExited(record: StartedRecorded, tmux: TmuxClient): Promise<void> {
+	const windows = await listWindows(tmux);
+	const state = deriveState(record, windowState(windows, record));
+	if (state === "running") {
+		throw new Error(`delegation ${record.id} is running; send it a message instead of resuming it`);
+	}
+	if (windows === undefined || state === "unknown") {
+		throw new Error(`could not ask tmux whether delegation ${record.id}'s window is open, so it is not resumed`);
+	}
+	const own = ownWindow(windows, record);
+	if (own === undefined) return;
+	if (!own.exited) {
+		throw new Error(
+			`delegation ${record.id}'s window ${record.windowName} (${record.windowId}) is still open; its Pi may still be running, so it is not resumed. Close it with delegation_close, or send it a message`,
+		);
+	}
+	try {
+		await tmux.kill(own.id);
+	} catch (err) {
+		throw new Error(
+			`could not kill the exited window ${record.windowName} (${record.windowId}): ${(err as Error).message}`,
+		);
+	}
+}
+
+/** Whether `path` is an existing directory. */
+function isDirectory(path: string): boolean {
+	try {
+		return statSync(path).isDirectory();
+	} catch {
+		return false;
+	}
+}
+
+/** A delegate's session file, the directory its header says Pi runs it in, and the session directory to pass. */
+type SessionFile = { path: string; cwd: string; dir: string };
+
+/** Step 5: the session file, and the directory its header says Pi will run it in. */
+function sessionFor(record: StartedRecorded, ctx: ExtensionContext): SessionFile {
+	const dir = record.sessionDir ?? ctx.sessionManager.getSessionDir();
+	const path = locateSession(record, ctx.sessionManager.getSessionDir());
+	if (path === undefined) {
+		throw new Error(
+			`no session file for session ${sessionIdOf(record)} was found in ${dir}, so delegation ${record.id} is not resumable`,
+		);
+	}
+	if (!isAbsolute(path) || !isAbsolute(dir)) {
+		throw new Error(
+			`the session file ${path} or its directory ${dir} is not an absolute path, so delegation ${record.id} is not resumable`,
+		);
+	}
+	const cwd = sessionHeaderCwd(path);
+	if (cwd === undefined) {
+		throw new Error(
+			`could not read a working directory from the header of ${path}, so delegation ${record.id} is not resumable`,
+		);
+	}
+	if (!isDirectory(cwd)) {
+		throw new Error(
+			`delegation ${record.id}'s session runs in ${cwd}, which no longer exists, so it is not resumable (a removed worktree cannot be resumed)`,
+		);
+	}
+	return { path, cwd, dir };
+}
+
+/** The model choices: an explicit model alone, else where the delegate ended, then the agent's fallbacks. */
+function preferredModels(record: StartedRecorded, agent: Agent, sessionPath: string): string[] {
+	const ended = finishedModel(sessionPath) ?? record.result?.finishedModel ?? record.model;
+	return [...new Set([ended, ...agent.fallback])];
+}
+
+/** The tool's result: the `delegate` tool's plus the delegation it resumed. */
+function resultSchema(delegateSchema: TObject): TObject {
+	return Type.Object({
+		...delegateSchema.properties,
+		resumedFrom: Type.String({ description: "The delegation whose session this one reopened." }),
+	});
+}
+
+/** The model, thinking and auto-exit a resume runs with, and the directory and workspace it opens in. */
+type Plan = {
+	start: { agent: Agent; model: string; rest: string[]; thinking: string; autoExit: boolean };
+	chosen: Awaited<ReturnType<ResumeDeps["chooseModel"]>>;
+	workspace: Workspace;
+	trusted: boolean;
+};
+
+/** Settle the request's overrides against the record, and decide trust again for the session's directory. */
+async function planResume(
+	deps: ResumeDeps,
+	input: { record: StartedRecorded; agent: Agent; session: SessionFile; request: ResumeRequest },
+	ctx: ExtensionContext,
+	signal: AbortSignal | undefined,
+): Promise<Plan> {
+	const { record, agent, session, request } = input;
+	const thinking = request.thinking ?? record.thinking;
+	const badThinking = deps.thinkingProblem(thinking);
+	if (badThinking !== undefined) throw new Error(badThinking);
+	const preferred = preferredModels(record, agent, session.path);
+	const chosen = await deps.chooseModel(request.model, agent, ctx.modelRegistry, preferred);
+	const autoExit = request.autoExit ?? record.autoExit === true;
+	const prepared = await deps.prepareDirectory(session.cwd, ctx, signal);
+	// A managed worktree keeps its launch facts while the session still runs in it.
+	const workspace = record.workspace?.cwd === prepared.workspace.cwd ? record.workspace : prepared.workspace;
+	return {
+		start: { agent, model: chosen.model, rest: chosen.rest, thinking, autoExit },
+		chosen,
+		workspace,
+		trusted: prepared.trusted,
+	};
+}
+
+/** Mail the task, open the window on the session file and record the new delegation. */
+/**
+ * Mail the task to a window already open. Unlike a launch, a resume mails
+ * after the window opens: the session's inbox outlives a failed attempt, so
+ * mail sent for a window that never opened would reach the next resume as a
+ * second task. The reopened session's provider still delivers mail that
+ * arrives after it starts. A failed send kills the window, which has no task.
+ */
+async function mailTask(deps: ResumeDeps, windowId: string, sessionId: string, task: string): Promise<string> {
+	try {
+		return deps.sendTask(sessionId, task);
+	} catch (err) {
+		await deps.tmux.kill(windowId).catch(() => {});
+		throw err;
+	}
+}
+
+async function startResume(
+	deps: ResumeDeps,
+	input: {
+		record: StartedRecorded;
+		session: SessionFile;
+		plan: Plan;
+		task: string;
+		toolCallId: string | undefined;
+		placement: WindowPlacement;
+	},
+	ctx: ExtensionContext,
+): Promise<Resumed> {
+	const { record, session, plan, task, toolCallId, placement } = input;
+	const { start, chosen, workspace } = plan;
+	const id = randomUUID();
+	const sessionId = sessionIdOf(record);
+	const name = recordedName(record);
+	const { id: windowId, session: windowSession } = await deps.open({
+		id,
+		name,
+		sessionPath: session.path,
+		workspace,
+		trusted: plan.trusted,
+		cwd: workspace.cwd,
+		sessionDir: session.dir,
+		start,
+		toolCallId,
+		ctx,
+		placement,
+	});
+	const requestId = await mailTask(deps, windowId, sessionId, task);
+	deps.results.recordStart(
+		{
+			id,
+			agent: start.agent.name,
+			model: start.model,
+			thinking: start.thinking,
+			name,
+			windowId,
+			windowName: name,
+			requestId,
+			autoExit: start.autoExit,
+			cwd: workspace.cwd,
+			sessionDir: session.dir,
+			workspace,
+			resumedFrom: record.id,
+			sessionId,
+		},
+		ctx,
+	);
+	deps.startPolling();
+	return {
+		id,
+		resumedFrom: record.id,
+		agent: start.agent.name,
+		model: start.model,
+		skipped: chosen.skipped.map(({ model, reason }) => ({ model, reason })),
+		notices: chosen.notices,
+		thinking: start.thinking,
+		autoExit: start.autoExit,
+		windowId,
+		windowName: name,
+		...(windowSession === undefined ? {} : { windowSession }),
+		cwd: workspace.cwd,
+		sessionDir: session.dir,
+		workspace,
+		fallbackNote:
+			chosen.skipped.length === 0
+				? ""
+				: `Used fallback ${chosen.model} because ${chosen.skipped.map((item) => `${item.model} is ${item.why}`).join("; ")}.`,
+	};
+}
+
+/** The delegation ids with a resume under way in this process. */
+const inFlight = new Set<string>();
+
+/** Resume `request.id`: every eligibility check in order, then the launch. */
+async function resume(
+	deps: ResumeDeps,
+	request: ResumeRequest,
+	ctx: ExtensionContext,
+	signal: AbortSignal | undefined,
+): Promise<Resumed> {
+	if (!deps.tmux.insideTmux()) throw new Error(deps.notInTmux);
+	await requireTmuxVersion(deps.tmux);
+	const placement = deps.placement(); // read once, before the exited window is killed or anything is mailed
+	if (request.task.trim() === "") throw new Error("a task is required");
+	const record = newestStarted(deps.results, request.id);
+	await ensureExited(record, deps.tmux);
+	const session = sessionFor(record, ctx);
+	const agent = deps.roster(ctx).agents.find((item) => item.name === record.agent);
+	if (agent === undefined) {
+		throw new Error(
+			`agent ${JSON.stringify(record.agent)} is no longer in the roster, so delegation ${record.id} is not resumed`,
+		);
+	}
+	if (inFlight.has(record.id)) throw new Error(`delegation ${record.id} is already being resumed`);
+	inFlight.add(record.id);
+	try {
+		const plan = await planResume(deps, { record, agent, session, request }, ctx, signal);
+		return await startResume(
+			deps,
+			{ record, session, plan, task: request.task, toolCallId: request.toolCallId, placement },
+			ctx,
+		);
+	} finally {
+		inFlight.delete(record.id);
+	}
+}
+
+/** The session a window opened in, for the start text. */
+const inSession = (session: string | undefined): string => (session === undefined ? "" : ` in session ${session}`);
+
+/** The one line the command and the tool tell the user. */
+export function startedLine(resumed: Resumed): string {
+	return `delegate: resumed ${resumed.resumedFrom} as ${resumed.id}: ${resumed.agent} in window ${resumed.windowName}${inSession(resumed.windowSession)}${resumed.fallbackNote === "" ? "" : `. ${resumed.fallbackNote}`}${resumed.notices.map((line) => `. ${line}`).join("")}`;
+}
+
+/** The `delegation_resume` tool. */
+function resumeTool(deps: ResumeDeps): ToolDef {
+	return {
+		name: "delegation_resume",
+		label: "Resume delegation",
+		description:
+			"Continue an exited delegate with a follow-up task: reopen its own session, so it keeps the context it built, in a new window. The follow-up is recorded as a new delegation, and its result arrives as a message on its own, like a delegate's. Use it when the next task needs what the delegate already read or decided; for unrelated work, use delegate. It refuses a delegation that is still running (send it a message instead), one whose window is still open, an older delegation of a chain that was already resumed (it names the newest), and one whose working directory is gone, a removed worktree included. A delegate whose window was killed mid-turn can be resumed.",
+		parameters: Type.Object({
+			id: Type.String({ description: "The delegation to resume, from a delegate result or delegation_status." }),
+			task: Type.String({
+				description: "The follow-up task. The delegate sees its earlier conversation plus this.",
+			}),
+			model: Type.Optional(
+				Type.String({
+					description:
+						"provider/id to resume on; used as given. Omit to continue on the model the delegate finished on, or its agent's next fallback when that one is unavailable.",
+				}),
+			),
+			thinking: Type.Optional(
+				Type.String({ description: `Thinking level (${THINKING_LEVELS.join(", ")}); defaults to the delegation's.` }),
+			),
+			auto_exit: Type.Optional(
+				Type.Boolean({ description: "Close the window once it finishes; defaults to the delegation's setting." }),
+			),
+		}),
+		outputSchema: resultSchema(deps.delegateSchema),
+		// Opens a window and removes nothing but an exited window of the old run; each call starts another run.
+		annotations: { readOnlyHint: false, destructiveHint: true, idempotentHint: false, openWorldHint: true },
+		async execute(
+			toolCallId: string,
+			params: { id: string; task: string; model?: string | null; thinking?: string | null; auto_exit?: boolean | null },
+			signal: AbortSignal | undefined,
+			_onUpdate: unknown,
+			ctx: ExtensionContext,
+		): Promise<ReturnType<typeof toolResult>> {
+			const resumed = await resume(deps, { ...requestFrom(params), toolCallId }, ctx, signal);
+			const { fallbackNote, ...details } = resumed;
+			ctx.ui.notify(startedLine(resumed), "info");
+			return toolResult(
+				`Resumed delegation ${resumed.resumedFrom} as ${resumed.id}: ${resumed.agent} (${resumed.model}, thinking ${resumed.thinking}) in window ${resumed.windowName}${inSession(resumed.windowSession)}. It runs in the background, and its result arrives as a message that starts your next turn on its own; delegation_close ends it.${fallbackNote === "" ? "" : ` ${fallbackNote}`}${resumed.notices.map((line) => ` ${line}`).join("")}`,
+				{ outcome: "started", ...details },
+			);
+		},
+	};
+}
+
+/** The request a tool call makes; some models fill every optional field with null or the string "null", and both mean unset. */
+function requestFrom(params: {
+	id: string;
+	task: string;
+	model?: string | null;
+	thinking?: string | null;
+	auto_exit?: boolean | null;
+}): ResumeRequest {
+	const given = (value: string | null | undefined): string | undefined =>
+		value == null || value === "null" ? undefined : value;
+	return {
+		id: params.id,
+		task: params.task ?? "",
+		model: given(params.model),
+		thinking: given(params.thinking),
+		autoExit: typeof params.auto_exit === "boolean" ? params.auto_exit : undefined,
+	};
+}
+
+/** Resuming as a function for the command, and the tool for the model. */
+export function createResume(deps: ResumeDeps): {
+	resume(request: ResumeRequest, ctx: ExtensionContext, signal?: AbortSignal): Promise<Resumed>;
+	tool(): ToolDef;
+} {
+	return {
+		resume: (request: ResumeRequest, ctx: ExtensionContext, signal?: AbortSignal) => resume(deps, request, ctx, signal),
+		tool: () => resumeTool(deps),
+	};
+}