1  import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@mariozechner/pi-coding-agent";
 2  
 3  import type { SessionPickerResult } from "./picker.ts";
 4  import { openSessionSwitchPicker } from "./picker.ts";
 5  import { executeStartupAction, resolveStartupAction, resolveStartupSessionTarget } from "./relaunch.ts";
 6  
 7  export function resolveCommandPickerAction(
 8  	result: SessionPickerResult,
 9  ): { kind: "noop" | "shutdown" } | { kind: "switch"; sessionPath: string } {
10  	if (result.kind === "dismissed") {
11  		return result.reason === "exit" ? { kind: "shutdown" } : { kind: "noop" };
12  	}
13  
14  	return { kind: "switch", sessionPath: result.sessionPath };
15  }
16  
17  async function runSessionSwitchCommand(pi: ExtensionAPI, ctx: ExtensionCommandContext): Promise<void> {
18  	const result = await openSessionSwitchPicker(pi, ctx);
19  	const action = resolveCommandPickerAction(result);
20  	if (action.kind === "noop") {
21  		return;
22  	}
23  	if (action.kind === "shutdown") {
24  		ctx.shutdown();
25  		return;
26  	}
27  
28  	const switchResult = await ctx.switchSession(action.sessionPath);
29  	if (switchResult.cancelled) {
30  		ctx.ui.notify("Session switch cancelled", "info");
31  	}
32  }
33  
34  async function runSessionSwitchStartup(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
35  	const result = await openSessionSwitchPicker(pi, ctx);
36  	if (result.kind === "selected") {
37  		const target = resolveStartupSessionTarget(result.sessionPath);
38  		if ("warning" in target) {
39  			ctx.ui.notify(target.warning, "warning");
40  			return;
41  		}
42  
43  		const action = resolveStartupAction(result, { cwd: target.cwd });
44  		executeStartupAction(ctx, action);
45  		return;
46  	}
47  
48  	const action = resolveStartupAction(result, { cwd: ctx.cwd });
49  	executeStartupAction(ctx, action);
50  }
51  
52  export default function sessionSwitchExtension(pi: ExtensionAPI) {
53  	pi.registerFlag("switch-session", {
54  		description: "Open the session-switch picker after startup, then relaunch into the selected session",
55  		type: "boolean",
56  		default: false,
57  	});
58  
59  	pi.on("session_start", async (event, ctx) => {
60  		if (event.reason !== "startup" || !ctx.hasUI || pi.getFlag("switch-session") !== true) {
61  			return;
62  		}
63  
64  		await runSessionSwitchStartup(pi, ctx);
65  	});
66  
67  	pi.registerCommand("switch-session", {
68  		description: "Session picker (mirrors /resume) with live preview",
69  		handler: async (_args, ctx) => {
70  			await runSessionSwitchCommand(pi, ctx);
71  		},
72  	});
73  }