index.ts

  1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
  2// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
  3//
  4// SPDX-License-Identifier: MIT
  5
  6import { complete, type Message, type Tool } from "@mariozechner/pi-ai";
  7import type {
  8	ExtensionAPI,
  9	ExtensionCommandContext,
 10	ExtensionContext,
 11	SessionEntry,
 12} from "@mariozechner/pi-coding-agent";
 13import { BorderedLoader, SessionManager, convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
 14import { Key, matchesKey } from "@mariozechner/pi-tui";
 15import { Type } from "@sinclair/typebox";
 16import * as fs from "node:fs";
 17import * as os from "node:os";
 18import * as path from "node:path";
 19
 20const STATUS_KEY = "handoff";
 21const COUNTDOWN_SECONDS = 10;
 22
 23/**
 24 * Model used for handoff extraction calls. Set PI_HANDOFF_MODEL env var
 25 * as "provider/modelId" (e.g. "anthropic/claude-haiku-4-5") to use a
 26 * different model for extraction than the session's current model.
 27 */
 28const HANDOFF_MODEL_OVERRIDE = process.env.PI_HANDOFF_MODEL;
 29
 30const SYSTEM_PROMPT = `You're helping transfer context between coding sessions. The next session starts fresh with no memory of this conversation, so extract what matters.
 31
 32Consider these questions:
 33- What was just done or implemented?
 34- What decisions were made and why?
 35- What technical details were discovered (APIs, methods, patterns)?
 36- What constraints, limitations, or caveats were found?
 37- What patterns or approaches are being followed?
 38- What is still unfinished or unresolved?
 39- What open questions or risks are worth flagging?
 40
 41Rules:
 42- Be concrete: prefer file paths, specific decisions, and actual commands over vague summaries.
 43- Only include files from the provided candidate list.
 44- Only include skills from the provided loaded skills list.
 45- If something wasn't explicitly discussed, don't include it.`;
 46
 47const QUERY_SYSTEM_PROMPT = `You answer questions about a prior pi session.
 48
 49Rules:
 50- Use only facts from the provided conversation.
 51- Prefer concrete outputs: file paths, decisions, TODOs, errors.
 52- If not present, say explicitly: "Not found in provided session.".
 53- Keep answer concise.`;
 54
 55/**
 56 * Tool definition passed to the extraction LLM call. Models produce
 57 * structured tool-call arguments far more reliably than free-form JSON
 58 * in a text response, so we use a single tool instead of responseFormat.
 59 */
 60const HANDOFF_EXTRACTION_TOOL: Tool = {
 61	name: "extract_handoff_context",
 62	description: "Extract handoff context from the conversation for a new session.",
 63	parameters: Type.Object({
 64		relevantFiles: Type.Array(Type.String(), {
 65			description: "File paths relevant to continuing this work, chosen from the provided candidate list",
 66		}),
 67		skillsInUse: Type.Array(Type.String(), {
 68			description: "Skills relevant to the goal, chosen from the provided list of loaded skills",
 69		}),
 70		context: Type.String({
 71			description:
 72				"Key context: what was done, decisions made, technical details discovered, constraints, and patterns being followed",
 73		}),
 74		openItems: Type.Array(Type.String(), {
 75			description: "Unfinished work, open questions, known risks, and things to watch out for",
 76		}),
 77	}),
 78};
 79
 80type PendingAutoSubmit = {
 81	ctx: ExtensionContext;
 82	sessionFile: string | undefined;
 83	interval: ReturnType<typeof setInterval>;
 84	unsubscribeInput: () => void;
 85};
 86
 87type PendingHandoff = {
 88	prompt: string;
 89	parentSession: string | undefined;
 90	newModel: string | undefined;
 91};
 92
 93function isEditableInput(data: string): boolean {
 94	if (!data) return false;
 95	if (data.length === 1) {
 96		const code = data.charCodeAt(0);
 97		if (code >= 32 && code !== 127) return true;
 98		if (code === 8 || code === 13) return true;
 99	}
100
101	if (data === "\n" || data === "\r") return true;
102	if (data === "\x7f") return true;
103
104	if (data.length > 1 && !data.startsWith("\x1b")) return true;
105
106	return false;
107}
108
109function statusLine(ctx: ExtensionContext, seconds: number): string {
110	const accent = ctx.ui.theme.fg("accent", `handoff auto-submit in ${seconds}s`);
111	const hint = ctx.ui.theme.fg("dim", "(type to edit, Esc to cancel)");
112	return `${accent} ${hint}`;
113}
114
115function getSessionsRoot(sessionFile: string | undefined): string | undefined {
116	if (!sessionFile) return undefined;
117	const normalized = sessionFile.replace(/\\/g, "/");
118	const marker = "/sessions/";
119	const idx = normalized.indexOf(marker);
120	if (idx === -1) {
121		return path.dirname(path.resolve(sessionFile));
122	}
123	return normalized.slice(0, idx + marker.length - 1);
124}
125
126function getFallbackSessionsRoot(): string | undefined {
127	const configuredDir = process.env.PI_CODING_AGENT_DIR;
128	const candidate = configuredDir
129		? path.resolve(configuredDir, "sessions")
130		: path.resolve(os.homedir(), ".pi", "agent", "sessions");
131	return fs.existsSync(candidate) ? candidate : undefined;
132}
133
134function normalizeSessionPath(sessionPath: string, sessionsRoot: string | undefined): string {
135	if (path.isAbsolute(sessionPath)) return path.resolve(sessionPath);
136	if (sessionsRoot) return path.resolve(sessionsRoot, sessionPath);
137	return path.resolve(sessionPath);
138}
139
140function sessionPathAllowed(candidate: string, sessionsRoot: string | undefined): boolean {
141	if (!sessionsRoot) return true;
142	const root = path.resolve(sessionsRoot);
143	const resolved = path.resolve(candidate);
144	return resolved === root || resolved.startsWith(`${root}${path.sep}`);
145}
146
147/**
148 * Build a candidate file set from two sources:
149 *   1. Primary: actual tool calls (read, write, edit, create) in the session
150 *   2. Secondary: file-like patterns in the conversation text (catches files
151 *      that were discussed but never opened)
152 */
153function extractCandidateFiles(entries: SessionEntry[], conversationText: string): Set<string> {
154	const files = new Set<string>();
155	const fileToolNames = new Set(["read", "write", "edit", "create"]);
156
157	// Primary: files from actual tool calls
158	for (const entry of entries) {
159		if (entry.type !== "message") continue;
160		const msg = entry.message;
161		if (msg.role !== "assistant") continue;
162
163		for (const block of msg.content) {
164			if (typeof block !== "object" || block === null || block.type !== "toolCall") continue;
165			if (!fileToolNames.has(block.name)) continue;
166
167			const args = block.arguments as Record<string, unknown>;
168			const filePath =
169				typeof args.path === "string" ? args.path : typeof args.file === "string" ? args.file : undefined;
170			if (!filePath) continue;
171			if (filePath.endsWith("/SKILL.md")) continue;
172
173			files.add(filePath);
174		}
175	}
176
177	// Secondary: file-like patterns from conversation text
178	const filePattern = /(?:^|\s)([a-zA-Z0-9._\-/]+\.[a-zA-Z0-9]+)(?:\s|$|[,;:\)])/gm;
179	let match;
180	while ((match = filePattern.exec(conversationText)) !== null) {
181		const candidate = match[1];
182		if (candidate && !candidate.startsWith(".") && candidate.length > 2) {
183			files.add(candidate);
184		}
185	}
186
187	return files;
188}
189
190/**
191 * Extract skill names that were actually loaded during the conversation.
192 * Looks for read() tool calls targeting SKILL.md files and derives the
193 * skill name from the parent directory (the convention for pi skills).
194 */
195function extractLoadedSkills(entries: SessionEntry[]): string[] {
196	const skills = new Set<string>();
197	for (const entry of entries) {
198		if (entry.type !== "message") continue;
199		const msg = entry.message;
200		if (msg.role !== "assistant") continue;
201
202		for (const block of msg.content) {
203			if (typeof block !== "object" || block === null || block.type !== "toolCall") continue;
204
205			// read() calls where the path ends in SKILL.md
206			if (block.name !== "read") continue;
207			const args = block.arguments as Record<string, unknown>;
208			const filePath = typeof args.path === "string" ? args.path : undefined;
209			if (!filePath?.endsWith("/SKILL.md")) continue;
210
211			// Skill name is the parent directory name:
212			//   .../skills/backing-up-with-keld/SKILL.md → backing-up-with-keld
213			const parent = path.basename(path.dirname(filePath));
214			if (parent && parent !== "skills") {
215				skills.add(parent);
216			}
217		}
218	}
219	return [...skills].sort();
220}
221
222/**
223 * Resolve the model to use for handoff extraction calls. Uses the
224 * PI_HANDOFF_MODEL env var if set, otherwise falls back to the session model.
225 */
226function resolveExtractionModel(ctx: {
227	model: ExtensionContext["model"];
228	modelRegistry: ExtensionContext["modelRegistry"];
229}): ExtensionContext["model"] {
230	if (!HANDOFF_MODEL_OVERRIDE) return ctx.model;
231	const slashIdx = HANDOFF_MODEL_OVERRIDE.indexOf("/");
232	if (slashIdx <= 0) return ctx.model;
233	const provider = HANDOFF_MODEL_OVERRIDE.slice(0, slashIdx);
234	const modelId = HANDOFF_MODEL_OVERRIDE.slice(slashIdx + 1);
235	return ctx.modelRegistry.find(provider, modelId) ?? ctx.model;
236}
237
238type HandoffExtraction = {
239	files: string[];
240	skills: string[];
241	context: string;
242	openItems: string[];
243};
244
245/**
246 * Run the extraction LLM call and return structured context, or null if
247 * aborted or the model didn't produce a valid tool call.
248 */
249async function extractHandoffContext(
250	model: NonNullable<ExtensionContext["model"]>,
251	apiKey: string,
252	headers: Record<string, string> | undefined,
253	conversationText: string,
254	goal: string,
255	candidateFiles: string[],
256	loadedSkills: string[],
257	signal?: AbortSignal,
258): Promise<HandoffExtraction | null> {
259	const filesContext =
260		candidateFiles.length > 0
261			? `\n\n## Candidate Files\n\nThese files were touched or mentioned during the session. Return only the ones relevant to the goal.\n\n${candidateFiles.map((f) => `- ${f}`).join("\n")}`
262			: "";
263
264	const skillsContext =
265		loadedSkills.length > 0
266			? `\n\n## Skills Loaded During This Session\n\n${loadedSkills.map((s) => `- ${s}`).join("\n")}\n\nReturn only the skills from this list that are relevant to the goal.`
267			: "";
268
269	const userMessage: Message = {
270		role: "user",
271		content: [
272			{
273				type: "text",
274				text: `## Conversation\n\n${conversationText}\n\n## Goal for Next Session\n\n${goal}${filesContext}${skillsContext}`,
275			},
276		],
277		timestamp: Date.now(),
278	};
279
280	const response = await complete(
281		model,
282		{
283			systemPrompt: SYSTEM_PROMPT,
284			messages: [userMessage],
285			tools: [HANDOFF_EXTRACTION_TOOL],
286		},
287		{ apiKey, headers, signal },
288	);
289
290	if (response.stopReason === "aborted") return null;
291
292	const toolCall = response.content.find((c) => c.type === "toolCall" && c.name === "extract_handoff_context");
293
294	if (!toolCall || toolCall.type !== "toolCall") {
295		console.error("Model did not call extract_handoff_context:", response.content);
296		return null;
297	}
298
299	const args = toolCall.arguments as Record<string, unknown>;
300
301	if (
302		!Array.isArray(args.relevantFiles) ||
303		!Array.isArray(args.skillsInUse) ||
304		typeof args.context !== "string" ||
305		!Array.isArray(args.openItems)
306	) {
307		console.error("Unexpected tool call arguments shape:", args);
308		return null;
309	}
310
311	const candidateSet = new Set(candidateFiles);
312	const loadedSkillSet = new Set(loadedSkills);
313
314	return {
315		files: (args.relevantFiles as string[]).filter((f) => typeof f === "string" && candidateSet.has(f)),
316		skills: (args.skillsInUse as string[]).filter((s) => typeof s === "string" && loadedSkillSet.has(s)),
317		context: args.context as string,
318		openItems: (args.openItems as string[]).filter((item) => typeof item === "string"),
319	};
320}
321
322/**
323 * Assemble the handoff draft. The user's goal goes last so it has the
324 * most weight in the new session (recency bias).
325 */
326function assembleHandoffDraft(result: HandoffExtraction, goal: string, parentSessionFile: string | undefined): string {
327	const parentBlock = parentSessionFile
328		? `**Parent session:** \`${parentSessionFile}\`\n\nUse the \`session_query\` tool with this path if you need details from the prior thread.\n\n`
329		: "";
330
331	const filesSection = result.files.length > 0 ? `## Files\n\n${result.files.map((f) => `- ${f}`).join("\n")}\n\n` : "";
332
333	const skillsSection =
334		result.skills.length > 0 ? `## Skills in Use\n\n${result.skills.map((s) => `- ${s}`).join("\n")}\n\n` : "";
335
336	const contextSection = `## Context\n\n${result.context}\n\n`;
337
338	const openItemsSection =
339		result.openItems.length > 0 ? `## Open Items\n\n${result.openItems.map((item) => `- ${item}`).join("\n")}\n\n` : "";
340
341	return `${parentBlock}${filesSection}${skillsSection}${contextSection}${openItemsSection}${goal}`.trim();
342}
343
344export default function (pi: ExtensionAPI) {
345	let pending: PendingAutoSubmit | null = null;
346	let pendingHandoff: PendingHandoff | null = null;
347	let handoffTimestamp: number | null = null;
348
349	const clearPending = (ctx?: ExtensionContext, notify?: string) => {
350		if (!pending) return;
351
352		clearInterval(pending.interval);
353		pending.unsubscribeInput();
354		pending.ctx.ui.setStatus(STATUS_KEY, undefined);
355
356		const local = pending;
357		pending = null;
358
359		if (notify && ctx) {
360			ctx.ui.notify(notify, "info");
361		} else if (notify) {
362			local.ctx.ui.notify(notify, "info");
363		}
364	};
365
366	const autoSubmitDraft = () => {
367		if (!pending) return;
368
369		const active = pending;
370		const currentSession = active.ctx.sessionManager.getSessionFile();
371		if (active.sessionFile && currentSession !== active.sessionFile) {
372			clearPending(undefined);
373			return;
374		}
375
376		const draft = active.ctx.ui.getEditorText().trim();
377		clearPending(undefined);
378
379		if (!draft) {
380			active.ctx.ui.notify("Handoff draft is empty", "warning");
381			return;
382		}
383
384		active.ctx.ui.setEditorText("");
385
386		try {
387			if (active.ctx.isIdle()) {
388				pi.sendUserMessage(draft);
389			} else {
390				pi.sendUserMessage(draft, { deliverAs: "followUp" });
391			}
392		} catch {
393			pi.sendUserMessage(draft);
394		}
395	};
396
397	const startCountdown = (ctx: ExtensionContext) => {
398		clearPending(ctx);
399
400		let seconds = COUNTDOWN_SECONDS;
401		ctx.ui.setStatus(STATUS_KEY, statusLine(ctx, seconds));
402
403		const unsubscribeInput = ctx.ui.onTerminalInput((data) => {
404			if (matchesKey(data, Key.escape)) {
405				clearPending(ctx, "Handoff auto-submit cancelled");
406				return { consume: true };
407			}
408
409			if (isEditableInput(data)) {
410				clearPending(ctx, "Handoff auto-submit stopped (editing)");
411			}
412
413			return undefined;
414		});
415
416		const interval = setInterval(() => {
417			if (!pending) return;
418
419			seconds -= 1;
420			if (seconds <= 0) {
421				autoSubmitDraft();
422				return;
423			}
424
425			ctx.ui.setStatus(STATUS_KEY, statusLine(ctx, seconds));
426		}, 1000);
427
428		pending = {
429			ctx,
430			sessionFile: ctx.sessionManager.getSessionFile(),
431			interval,
432			unsubscribeInput,
433		};
434	};
435
436	pi.on("session_before_switch", (_event, ctx) => {
437		if (pending) clearPending(ctx);
438	});
439
440	pi.on("session_switch", (_event, ctx) => {
441		if (pending) clearPending(ctx);
442		// A proper session switch (e.g. /new) fully resets agent state,
443		// so clear the context filter to avoid hiding new messages.
444		handoffTimestamp = null;
445	});
446
447	pi.on("session_before_fork", (_event, ctx) => {
448		if (pending) clearPending(ctx);
449	});
450
451	pi.on("session_fork", (_event, ctx) => {
452		if (pending) clearPending(ctx);
453	});
454
455	pi.on("session_before_tree", (_event, ctx) => {
456		if (pending) clearPending(ctx);
457	});
458
459	pi.on("session_tree", (_event, ctx) => {
460		if (pending) clearPending(ctx);
461	});
462
463	pi.on("session_shutdown", (_event, ctx) => {
464		if (pending) clearPending(ctx);
465	});
466
467	// --- Tool-path handoff coordination ---
468	//
469	// The /handoff command has ExtensionCommandContext with ctx.newSession()
470	// which does a full agent reset. The handoff tool only gets
471	// ExtensionContext, which lacks newSession(). So the tool stores a
472	// pending handoff and these handlers complete it:
473	//
474	// 1. agent_end: after the agent loop finishes, switch sessions and
475	//    send the handoff prompt in the next macrotask
476	// 2. context: filter pre-handoff messages since the low-level session
477	//    switch doesn't clear agent.state.messages
478	// 3. session_switch (above): clears the filter on proper switches
479
480	pi.on("agent_end", (_event, ctx) => {
481		if (!pendingHandoff) return;
482
483		const { prompt, parentSession, newModel } = pendingHandoff;
484		pendingHandoff = null;
485
486		handoffTimestamp = Date.now();
487		(ctx.sessionManager as any).newSession({ parentSession });
488
489		setTimeout(async () => {
490			if (newModel) {
491				const slashIdx = newModel.indexOf("/");
492				if (slashIdx > 0) {
493					const provider = newModel.slice(0, slashIdx);
494					const modelId = newModel.slice(slashIdx + 1);
495					const model = ctx.modelRegistry.find(provider, modelId);
496					if (model) await pi.setModel(model);
497				}
498			}
499			pi.sendUserMessage(prompt);
500		}, 0);
501	});
502
503	pi.on("context", (event) => {
504		if (handoffTimestamp === null) return;
505
506		const newMessages = (event as any).messages.filter((m: any) => m.timestamp >= handoffTimestamp);
507		if (newMessages.length > 0) {
508			return { messages: newMessages };
509		}
510	});
511
512	pi.registerTool({
513		name: "session_query",
514		label: (params) => `Session Query: ${params.question}`,
515		description:
516			"Query a prior pi session file. Use when handoff prompt references a parent session and you need details.",
517		parameters: Type.Object({
518			sessionPath: Type.String({
519				description:
520					"Session .jsonl path. Absolute path, or relative to sessions root (e.g. 2026-02-16/foo/session.jsonl)",
521			}),
522			question: Type.String({ description: "Question about that session" }),
523		}),
524		async execute(_toolCallId, params, signal, onUpdate, ctx) {
525			const currentSessionFile = ctx.sessionManager.getSessionFile();
526			const sessionsRoot = getSessionsRoot(currentSessionFile) ?? getFallbackSessionsRoot();
527			const resolvedPath = normalizeSessionPath(params.sessionPath, sessionsRoot);
528
529			const error = (text: string) => ({
530				content: [{ type: "text" as const, text }],
531				details: { error: true },
532			});
533
534			const cancelled = () => ({
535				content: [{ type: "text" as const, text: "Session query cancelled." }],
536				details: { cancelled: true },
537			});
538
539			if (signal.aborted) {
540				return cancelled();
541			}
542
543			if (!resolvedPath.endsWith(".jsonl")) {
544				return error(`Invalid session path (expected .jsonl): ${params.sessionPath}`);
545			}
546
547			if (!sessionPathAllowed(resolvedPath, sessionsRoot)) {
548				return error(`Session path outside allowed sessions directory: ${params.sessionPath}`);
549			}
550
551			if (!fs.existsSync(resolvedPath)) {
552				return error(`Session file not found: ${resolvedPath}`);
553			}
554
555			let fileStats: fs.Stats;
556			try {
557				fileStats = fs.statSync(resolvedPath);
558			} catch (err) {
559				return error(`Failed to stat session file: ${String(err)}`);
560			}
561
562			if (!fileStats.isFile()) {
563				return error(`Session path is not a file: ${resolvedPath}`);
564			}
565
566			onUpdate?.({
567				content: [{ type: "text", text: `Querying: ${resolvedPath}` }],
568				details: { status: "loading", sessionPath: resolvedPath },
569			});
570
571			let sessionManager: SessionManager;
572			try {
573				sessionManager = SessionManager.open(resolvedPath);
574			} catch (err) {
575				return error(`Failed to open session: ${String(err)}`);
576			}
577
578			const branch = sessionManager.getBranch();
579			const messages = branch
580				.filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
581				.map((entry) => entry.message);
582
583			if (messages.length === 0) {
584				return {
585					content: [{ type: "text" as const, text: "Session has no messages." }],
586					details: { empty: true, sessionPath: resolvedPath },
587				};
588			}
589
590			if (!ctx.model) {
591				return error("No model selected for session query.");
592			}
593
594			const conversationText = serializeConversation(convertToLlm(messages));
595			try {
596				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
597				if (!auth.ok) {
598					return error(`Failed to get API key: ${auth.error}`);
599				}
600				const userMessage: Message = {
601					role: "user",
602					content: [
603						{
604							type: "text",
605							text: `## Session\n\n${conversationText}\n\n## Question\n\n${params.question}`,
606						},
607					],
608					timestamp: Date.now(),
609				};
610
611				const response = await complete(
612					ctx.model,
613					{ systemPrompt: QUERY_SYSTEM_PROMPT, messages: [userMessage] },
614					{ apiKey: auth.apiKey, headers: auth.headers, signal },
615				);
616
617				if (response.stopReason === "aborted") {
618					return cancelled();
619				}
620
621				const answer = response.content
622					.filter((c): c is { type: "text"; text: string } => c.type === "text")
623					.map((c) => c.text)
624					.join("\n")
625					.trim();
626
627				return {
628					content: [{ type: "text" as const, text: answer || "No answer generated." }],
629					details: {
630						sessionPath: resolvedPath,
631						question: params.question,
632						messageCount: messages.length,
633					},
634				};
635			} catch (err) {
636				if (signal.aborted) {
637					return cancelled();
638				}
639				if (err instanceof Error && err.name === "AbortError") {
640					return cancelled();
641				}
642				return error(`Session query failed: ${String(err)}`);
643			}
644		},
645	});
646
647	pi.registerTool({
648		name: "handoff",
649		label: "Handoff",
650		description:
651			"Transfer context to a new session. Use when the user explicitly asks for a handoff or when the context window is nearly full. Provide a goal describing what the new session should focus on.",
652		parameters: Type.Object({
653			goal: Type.String({
654				description: "The goal/task for the new session",
655			}),
656			model: Type.Optional(
657				Type.String({
658					description: "Model for the new session as provider/modelId (e.g. 'anthropic/claude-haiku-4-5')",
659				}),
660			),
661		}),
662		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
663			if (!ctx.model) {
664				return {
665					content: [{ type: "text" as const, text: "No model selected." }],
666				};
667			}
668
669			const branch = ctx.sessionManager.getBranch();
670			const messages = branch
671				.filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
672				.map((entry) => entry.message);
673
674			if (messages.length === 0) {
675				return {
676					content: [
677						{
678							type: "text" as const,
679							text: "No conversation to hand off.",
680						},
681					],
682				};
683			}
684
685			const llmMessages = convertToLlm(messages);
686			const conversationText = serializeConversation(llmMessages);
687			const candidateFiles = [...extractCandidateFiles(branch, conversationText)];
688			const loadedSkills = extractLoadedSkills(branch);
689
690			const extractionModel = resolveExtractionModel(ctx) ?? ctx.model;
691			const auth = await ctx.modelRegistry.getApiKeyAndHeaders(extractionModel);
692			if (!auth.ok) {
693				return {
694					content: [
695						{
696							type: "text" as const,
697							text: `Failed to get API key: ${auth.error}`,
698						},
699					],
700				};
701			}
702
703			let result: HandoffExtraction | null;
704			try {
705				result = await extractHandoffContext(
706					extractionModel,
707					auth.apiKey,
708					auth.headers,
709					conversationText,
710					params.goal,
711					candidateFiles,
712					loadedSkills,
713					signal,
714				);
715			} catch (err) {
716				if (signal.aborted) {
717					return {
718						content: [
719							{
720								type: "text" as const,
721								text: "Handoff cancelled.",
722							},
723						],
724					};
725				}
726				return {
727					content: [
728						{
729							type: "text" as const,
730							text: `Handoff extraction failed: ${String(err)}`,
731						},
732					],
733				};
734			}
735
736			if (!result) {
737				return {
738					content: [
739						{
740							type: "text" as const,
741							text: "Handoff extraction failed or was cancelled.",
742						},
743					],
744				};
745			}
746
747			const currentSessionFile = ctx.sessionManager.getSessionFile();
748			const prompt = assembleHandoffDraft(result, params.goal, currentSessionFile);
749
750			pendingHandoff = {
751				prompt,
752				parentSession: currentSessionFile,
753				newModel: params.model,
754			};
755
756			return {
757				content: [
758					{
759						type: "text" as const,
760						text: "Handoff prepared. The session will switch after this turn completes.",
761					},
762				],
763			};
764		},
765	});
766
767	pi.registerCommand("handoff", {
768		description: "Transfer context to a new session (-model provider/modelId)",
769		handler: async (args: string, ctx: ExtensionCommandContext) => {
770			if (!ctx.hasUI) {
771				ctx.ui.notify("/handoff requires interactive mode", "error");
772				return;
773			}
774
775			if (!ctx.model) {
776				ctx.ui.notify("No model selected", "error");
777				return;
778			}
779
780			// Parse optional -model flag from args
781			let remaining = args;
782			let newSessionModel: string | undefined;
783
784			const modelMatch = remaining.match(/(?:^|\s)-model\s+(\S+)/);
785			if (modelMatch) {
786				newSessionModel = modelMatch[1];
787				remaining = remaining.replace(modelMatch[0], " ");
788			}
789
790			let goal = remaining.trim();
791			if (!goal) {
792				const entered = await ctx.ui.input("handoff goal", "What should the new thread do?");
793				if (!entered?.trim()) {
794					ctx.ui.notify("Handoff cancelled", "info");
795					return;
796				}
797				goal = entered.trim();
798			}
799
800			const branch = ctx.sessionManager.getBranch();
801			const messages = branch
802				.filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
803				.map((entry) => entry.message);
804
805			if (messages.length === 0) {
806				ctx.ui.notify("No conversation to hand off", "warning");
807				return;
808			}
809
810			const llmMessages = convertToLlm(messages);
811			const conversationText = serializeConversation(llmMessages);
812			const currentSessionFile = ctx.sessionManager.getSessionFile();
813			const candidateFiles = [...extractCandidateFiles(branch, conversationText)];
814			const loadedSkills = extractLoadedSkills(branch);
815			const extractionModel = resolveExtractionModel(ctx) ?? ctx.model;
816
817			const result = await ctx.ui.custom<HandoffExtraction | null>((tui, theme, _kb, done) => {
818				const loader = new BorderedLoader(tui, theme, "Extracting handoff context...");
819				loader.onAbort = () => done(null);
820
821				const run = async () => {
822					const auth = await ctx.modelRegistry.getApiKeyAndHeaders(extractionModel);
823					if (!auth.ok) {
824						throw new Error(`Failed to get API key: ${auth.error}`);
825					}
826
827					return extractHandoffContext(
828						extractionModel,
829						auth.apiKey,
830						auth.headers,
831						conversationText,
832						goal,
833						candidateFiles,
834						loadedSkills,
835						loader.signal,
836					);
837				};
838
839				run()
840					.then(done)
841					.catch((err) => {
842						console.error("handoff generation failed", err);
843						done(null);
844					});
845
846				return loader;
847			});
848
849			if (!result) {
850				ctx.ui.notify("Handoff cancelled", "info");
851				return;
852			}
853
854			const prefillDraft = assembleHandoffDraft(result, goal, currentSessionFile);
855
856			const editedPrompt = await ctx.ui.editor("Edit handoff draft", prefillDraft);
857			if (editedPrompt === undefined) {
858				ctx.ui.notify("Handoff cancelled", "info");
859				return;
860			}
861
862			const next = await ctx.newSession({
863				parentSession: currentSessionFile,
864			});
865
866			if (next.cancelled) {
867				ctx.ui.notify("New session cancelled", "info");
868				return;
869			}
870
871			// Apply -model if specified
872			if (newSessionModel) {
873				const slashIdx = newSessionModel.indexOf("/");
874				if (slashIdx > 0) {
875					const provider = newSessionModel.slice(0, slashIdx);
876					const modelId = newSessionModel.slice(slashIdx + 1);
877					const model = ctx.modelRegistry.find(provider, modelId);
878					if (model) {
879						await pi.setModel(model);
880					} else {
881						ctx.ui.notify(`Unknown model: ${newSessionModel}`, "warning");
882					}
883				} else {
884					ctx.ui.notify(`Invalid model format "${newSessionModel}", expected provider/modelId`, "warning");
885				}
886			}
887
888			const newSessionFile = ctx.sessionManager.getSessionFile();
889			if (newSessionFile) {
890				ctx.ui.notify(`Switched to new session: ${newSessionFile}`, "info");
891			}
892
893			ctx.ui.setEditorText(editedPrompt);
894			startCountdown(ctx);
895		},
896	});
897}