dev_sandbox/.pi/agent/extensions/quick-question.ts

import type { ExtensionAPI, ExtensionCommandContext, SessionEntry } from "@earendil-works/pi-coding-agent";

const MARKER_TYPE = "quick-question-session";

type QuickSessionMarker = {
	returnSession: string;
	createdAt: string;
};

function getQuickSessionMarker(entries: SessionEntry[]): QuickSessionMarker | undefined {
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry.type !== "custom" || entry.customType !== MARKER_TYPE) continue;

		const data = entry.data as Partial<QuickSessionMarker> | undefined;
		if (typeof data?.returnSession === "string") {
			return {
				returnSession: data.returnSession,
				createdAt: typeof data.createdAt === "string" ? data.createdAt : "",
			};
		}
	}

	return undefined;
}

export default function (pi: ExtensionAPI) {
	pi.on("input", async (event, ctx) => {
		if (event.source === "extension") return { action: "continue" };

		const aliases: Record<string, string> = {
			"/quick-close": "/quickclose",
			"/quick-done": "/quickdone",
			"/quick-back": "/quickback",
		};
		const replacement = aliases[event.text.trim()];
		if (!replacement) return { action: "continue" };

		const marker = getQuickSessionMarker(ctx.sessionManager.getBranch());
		if (!marker) {
			ctx.ui.notify("This is not a quick question session.", "error");
			return { action: "handled" };
		}

		(ctx.sessionManager as unknown as { setSessionFile(path: string): void }).setSessionFile(marker.returnSession);
		ctx.ui.setStatus("quick-question", undefined);
		ctx.ui.setWidget("quick-question", undefined);
		ctx.ui.notify("Returned from quick session.", "info");
		return { action: "handled" };
	});

	pi.on("session_start", async (_event, ctx) => {
		const marker = getQuickSessionMarker(ctx.sessionManager.getBranch());

		if (!marker) {
			ctx.ui.setStatus("quick-question", undefined);
			ctx.ui.setWidget("quick-question", undefined);
			return;
		}

		ctx.ui.setStatus("quick-question", ctx.ui.theme.fg("accent", "quick"));
		ctx.ui.setWidget("quick-question", [
			ctx.ui.theme.fg(
				"dim",
				"Quick question session. Use /quickclose, /quickdone, or /quickback to return to the original session.",
			),
		]);
	});

	pi.registerCommand("quick", {
		description: "Start a temporary one-off question session; pass a question to ask immediately",
		handler: async (args, ctx) => {
			if (!ctx.hasUI) {
				ctx.ui.notify("quick requires interactive mode", "error");
				return;
			}

			await ctx.waitForIdle();

			const returnSession = ctx.sessionManager.getSessionFile();
			if (!returnSession) {
				ctx.ui.notify("Current session is ephemeral; cannot return to it later.", "error");
				return;
			}

			const question = args.trim();

			const result = await ctx.newSession({
				parentSession: returnSession,
				setup: async (session) => {
					session.appendCustomEntry(MARKER_TYPE, {
						returnSession,
						createdAt: new Date().toISOString(),
					} satisfies QuickSessionMarker);
					session.appendSessionInfo("Quick question");
				},
				withSession: async (replacementCtx) => {
					replacementCtx.ui.notify(
						"Quick session started. Use /quickclose, /quickdone, or /quickback to return.",
						"info",
					);

					if (question) {
						await replacementCtx.sendUserMessage(question);
					} else {
						replacementCtx.ui.setEditorText("");
					}
				},
			});

			if (result.cancelled) {
				ctx.ui.notify("Quick session cancelled", "info");
			}
		},
	});

	async function closeQuickSession(ctx: ExtensionCommandContext) {
		await ctx.waitForIdle();

		const marker = getQuickSessionMarker(ctx.sessionManager.getBranch());
		if (!marker) {
			ctx.ui.notify("This is not a quick question session.", "error");
			return;
		}

		const result = await ctx.switchSession(marker.returnSession, {
			withSession: async (replacementCtx) => {
				replacementCtx.ui.notify("Returned from quick session.", "info");
			},
		});

		if (result.cancelled) {
			ctx.ui.notify("Return to original session cancelled", "info");
		}
	}

	for (const name of ["quickclose", "quickdone", "quickback", "quick-close", "quick-done", "quick-back"] as const) {
		pi.registerCommand(name, {
			description: "Close the quick question session and return to the original session",
			handler: async (_args, ctx) => closeQuickSession(ctx),
		});
	}
}