Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/ephemlets/runtime.ts

Raw
import { randomUUID } from "node:crypto";
import { fileURLToPath } from "node:url";
import type { JsonValue } from "@earendil-works/pi-ai";
import type {
	AgentToolResult,
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { buildEphemlet } from "./src/builder.ts";
import { loadFactory } from "./src/loader.ts";
import { PanelHost } from "./src/panel.ts";
import {
	boundedText,
	cloneState,
	ENTRY_TYPE,
	EphemletStore,
	saveMessage,
} from "./src/store.ts";
import type { BuildResult, Ephemlet } from "./src/types.ts";

const authoringGuidanceUrl = new URL(
	"./guidance/authoring.md",
	import.meta.url,
);

export type EphemletParams = {
	action:
		| "put"
		| "build"
		| "list"
		| "get"
		| "state"
		| "delete"
		| "show"
		| "hide";
	id?: string;
	expectedRevision?: number;
	title?: string;
	purpose?: string;
	kind?: Ephemlet["kind"];
	source?: string;
	state?: unknown;
	include?: ("source" | "state")[];
	brief?: string;
	wait?: boolean;
	model?: string;
};

type Logger = typeof import("./src/debug.ts").dbg;

type Job = {
	controller: AbortController;
	promise: Promise<BuildResult | undefined>;
};
function required(value: string | undefined, name: string): string {
	if (!value?.trim()) throw new Error(`${name} is required`);
	return value;
}
function excerpt(value: string, limit: number): string {
	return value.length > limit ? `${value.slice(0, limit - 1)}…` : value;
}
function metadata(record: Ephemlet) {
	return {
		id: record.id,
		title: record.title,
		purpose: excerpt(record.purpose, 160),
		kind: record.kind,
		revision: record.revision,
		status: record.status,
		error: record.error ? excerpt(record.error, 256) : null,
		jobId: record.jobId,
	};
}
function result(data: unknown): AgentToolResult<unknown> {
	const text = JSON.stringify(data);
	const details: JsonValue = JSON.parse(text);
	return {
		content: [{ type: "text", text }],
		details,
		structuredContent: details,
	};
}

export class EphemletRuntime {
	readonly store: EphemletStore;
	private context: ExtensionContext | undefined;
	private sessionId: string | undefined;
	private panel: PanelHost | undefined;
	private readonly jobs = new Map<string, Job>();
	private generation = 0;

	constructor(
		private readonly pi: ExtensionAPI,
		private readonly log: Logger,
	) {
		this.store = new EphemletStore((event) =>
			pi.appendEntry(ENTRY_TYPE, event),
		);
	}
	setContext(ctx: ExtensionContext): void {
		this.context = ctx;
	}
	private ctx(): ExtensionContext {
		if (!this.context) throw new Error("No ephemlet session context");
		return this.context;
	}

	async restore(ctx: ExtensionContext): Promise<void> {
		this.dispose();
		this.context = ctx;
		this.sessionId = ctx.sessionManager.getSessionId();
		this.store.replay(ctx.sessionManager.getBranch());
		for (const record of this.store.records.values()) {
			if (record.status === "building")
				this.store.commit({
					...record,
					status: "failed",
					jobId: null,
					error:
						"Build interrupted by session change/restart; request build again.",
				});
		}
		if (this.store.visible) await this.refresh(false);
	}

	private async refresh(focus: boolean): Promise<void> {
		if (this.ctx().mode !== "tui") return;
		if (!this.panel) {
			this.panel ??= new PanelHost(this.pi, this.store, () => this.ctx(), {
				changed: () => this.panel?.refresh(true),
				close: (id) => {
					this.close(id);
					this.panel?.refresh(true);
				},
				save: (id) => {
					this.pi.sendUserMessage(saveMessage(this.store.get(id)), {
						deliverAs: "followUp",
						expandPromptTemplates: false,
					});
				},
				hide: () => {
					this.store.view(this.store.selected, false);
					this.panel?.refresh();
				},
			});
		}
		this.panel.refresh(focus);
	}

	async show(ctx: ExtensionContext): Promise<void> {
		await this.prepare(ctx);
		if (!this.store.records.size) {
			ctx.ui.notify("No ephemlets in this session", "info");
			return;
		}
		this.store.view(
			this.store.selected ?? this.store.records.keys().next().value ?? null,
			true,
		);
		await this.refresh(true);
	}
	async hide(ctx: ExtensionContext): Promise<void> {
		await this.prepare(ctx);
		this.store.view(this.store.selected, false);
		await this.refresh(false);
	}
	async toggle(ctx: ExtensionContext): Promise<void> {
		await this.prepare(ctx);
		if (this.store.visible) await this.hide(ctx);
		else await this.show(ctx);
	}
	private async prepare(ctx: ExtensionContext): Promise<void> {
		if (this.sessionId !== ctx.sessionManager.getSessionId())
			await this.restore(ctx);
		else this.context = ctx;
	}

	private close(id: string): void {
		this.jobs.get(id)?.controller.abort();
		this.jobs.delete(id);
		this.store.delete(id);
		this.log?.("ephemlet.delete", { count: this.store.records.size });
	}

	async execute(
		params: EphemletParams,
		ctx: ExtensionContext,
		signal?: AbortSignal,
	): Promise<AgentToolResult<unknown>> {
		await this.prepare(ctx);
		signal?.throwIfAborted();
		if (params.action === "list")
			return result({
				ephemlets: [...this.store.records.values()].map(metadata),
				guidance: fileURLToPath(authoringGuidanceUrl),
			});
		if (params.action === "get") {
			const record = this.store.get(required(params.id, "id"));
			return result({
				...metadata(record),
				...(params.include?.includes("source")
					? { source: record.source }
					: {}),
				...(params.include?.includes("state")
					? { state: cloneState(record.state) }
					: {}),
			});
		}
		if (params.action === "show" || params.action === "hide") {
			if (params.id) this.store.view(params.id, params.action === "show");
			if (params.action === "show") await this.show(ctx);
			else await this.hide(ctx);
			return result({
				visible: this.store.visible,
				selected: this.store.selected,
			});
		}
		if (params.action === "delete") {
			this.close(required(params.id, "id"));
			await this.refresh(false);
			return result({ deleted: params.id });
		}
		if (params.action === "state") {
			const record = this.store.get(required(params.id, "id"));
			if (params.expectedRevision !== record.revision)
				throw new Error(`Revision conflict: expected ${record.revision}`);
			this.store.setState(record.id, record.revision, cloneState(params.state));
			this.panel?.refresh();
			return result(metadata(record));
		}
		const previous = params.id ? this.store.get(params.id) : undefined;
		if (previous && params.expectedRevision !== previous.revision)
			throw new Error(`Revision conflict: expected ${previous.revision}`);
		const title = required(params.title ?? previous?.title, "title");
		const purpose = required(params.purpose ?? previous?.purpose, "purpose");
		if (params.action === "put") {
			const kind = params.kind ?? previous?.kind ?? "markdown";
			const source = required(params.source, "source");
			boundedText(source, "Source");
			if (kind === "interactive") await loadFactory(source);
			signal?.throwIfAborted();
			const record = this.store.put({
				...(params.id ? { id: params.id } : {}),
				...(params.expectedRevision
					? { expectedRevision: params.expectedRevision }
					: {}),
				title,
				purpose,
				kind,
				source,
				...(params.state !== undefined
					? { state: cloneState(params.state) }
					: {}),
			});
			this.jobs.get(record.id)?.controller.abort();
			this.jobs.delete(record.id);
			this.store.view(record.id, true);
			await this.refresh(false);
			this.log?.("ephemlet.put", { count: this.store.records.size, kind });
			return result(metadata(record));
		}
		const brief = boundedText(required(params.brief, "brief"), "Brief", 32768);
		if (this.jobs.size >= 4 && !previous?.jobId)
			throw new Error("At most four ephemlet builds may run concurrently");
		// Allocate stable identity and persist the request before returning to the main agent.
		const record = this.store.put({
			...(params.id ? { id: params.id } : {}),
			...(params.expectedRevision
				? { expectedRevision: params.expectedRevision }
				: {}),
			title,
			purpose,
			kind: previous?.kind ?? "interactive",
			source: previous?.source ?? "",
			...(params.state !== undefined
				? { state: cloneState(params.state) }
				: {}),
		});
		this.jobs.get(record.id)?.controller.abort();
		const controller = new AbortController();
		const buildSignal =
			params.wait && signal
				? AbortSignal.any([controller.signal, signal])
				: controller.signal;
		const jobId = randomUUID();
		const generation = this.generation;
		this.store.commit({ ...record, status: "building", jobId });
		this.store.view(record.id, true);
		const request = {
			id: record.id,
			title,
			purpose,
			brief,
			state: cloneState(record.state),
			source: record.source,
			...(params.model ? { model: params.model } : {}),
		};
		const promise = (async (): Promise<BuildResult | undefined> => {
			try {
				const built = await buildEphemlet(
					ctx,
					request,
					buildSignal,
					authoringGuidanceUrl,
				);
				buildSignal.throwIfAborted();
				await loadFactory(built.source);
				if (
					generation !== this.generation ||
					buildSignal.aborted ||
					this.store.records.get(record.id)?.jobId !== jobId
				)
					return undefined;
				const current = this.store.get(record.id);
				this.store.commit({
					...current,
					revision: current.revision + 1,
					kind: "interactive",
					source: built.source,
					...(current.source
						? {
								fallback: current.fallback ?? {
									kind: current.kind,
									source: current.source,
									state: cloneState(current.state),
								},
							}
						: {}),
					status: "ready",
					jobId: null,
					error: null,
				});
				this.pi.appendEntry("ephemlets.build-usage", {
					id: record.id,
					jobId,
					usage: built.usage,
				});
				if (!params.wait)
					this.pi.sendMessage(
						{
							customType: "ephemlet-build",
							content: `Ephemlet ready: ${title} (${record.id})`,
							display: false,
							details: {
								id: record.id,
								jobId,
								status: "ready",
								usage: built.usage,
							},
						},
						{ triggerTurn: false },
					);
				await this.refresh(false);
				this.log?.("build.finish", { outcome: "ready" });
				return built;
			} catch (error) {
				if (
					generation === this.generation &&
					this.store.records.get(record.id)?.jobId === jobId
				) {
					this.store.commit({
						...this.store.get(record.id),
						status: "failed",
						jobId: null,
						error: String(error).slice(0, 4096),
					});
					if (!params.wait) {
						this.pi.sendMessage(
							{
								customType: "ephemlet-build",
								content: `Ephemlet build failed (${record.id}); use ephemlet get for status.`,
								display: false,
								details: { id: record.id, jobId, status: "failed" },
							},
							{ triggerTurn: false },
						);
						ctx.ui.notify(
							`Ephemlet build failed: ${title}: ${String(error)}`,
							"error",
						);
					}
					await this.refresh(false);
					this.log?.("build.finish", { outcome: "failed" });
				}
				if (params.wait) throw error;
				return undefined;
			} finally {
				if (this.jobs.get(record.id)?.controller === controller)
					this.jobs.delete(record.id);
			}
		})();
		this.jobs.set(record.id, { controller, promise });
		await this.refresh(false);
		if (params.wait) {
			const built = await promise;
			return {
				...result(metadata(this.store.get(record.id))),
				...(built ? { usage: built.usage } : {}),
			};
		}
		return result({ ...metadata(this.store.get(record.id)), jobId });
	}

	dispose(): void {
		this.generation++;
		for (const job of this.jobs.values()) job.controller.abort();
		this.jobs.clear();
		this.panel?.dispose();
		this.panel = undefined;
		this.context = undefined;
		this.sessionId = undefined;
	}
}