1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
3//
4// SPDX-License-Identifier: MIT
5
6import type {
7 ExtensionAPI,
8 ExtensionCommandContext,
9 ExtensionContext,
10 SessionEntry,
11} from "@mariozechner/pi-coding-agent";
12import { BorderedLoader, convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
13import { extractCandidateFiles, extractLoadedSkills, resolveExtractionModel } from "./session-analysis.js";
14import { type HandoffExtraction, assembleHandoffDraft, extractHandoffContext } from "./handoff-extraction.js";
15import { parseAndSetModelWithNotify } from "./model-utils.js";
16
17export function registerHandoffCommand(pi: ExtensionAPI, startCountdown: (ctx: ExtensionContext) => void) {
18 pi.registerCommand("handoff", {
19 description: "Transfer context to a new session (-model provider/modelId)",
20 handler: async (args: string, ctx: ExtensionCommandContext) => {
21 if (!ctx.hasUI) {
22 ctx.ui.notify("/handoff requires interactive mode", "error");
23 return;
24 }
25
26 if (!ctx.model) {
27 ctx.ui.notify("No model selected", "error");
28 return;
29 }
30
31 // Parse optional -model flag from args
32 let remaining = args;
33 let newSessionModel: string | undefined;
34
35 const modelMatch = remaining.match(/(?:^|\s)-model\s+(\S+)/);
36 if (modelMatch) {
37 newSessionModel = modelMatch[1];
38 remaining = remaining.replace(modelMatch[0], " ");
39 }
40
41 let goal = remaining.trim();
42 if (!goal) {
43 const entered = await ctx.ui.input("handoff goal", "What should the new thread do?");
44 if (!entered?.trim()) {
45 ctx.ui.notify("Handoff cancelled", "info");
46 return;
47 }
48 goal = entered.trim();
49 }
50
51 const branch = ctx.sessionManager.getBranch();
52 const messages = branch
53 .filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
54 .map((entry) => entry.message);
55
56 if (messages.length === 0) {
57 ctx.ui.notify("No conversation to hand off", "warning");
58 return;
59 }
60
61 const llmMessages = convertToLlm(messages);
62 const conversationText = serializeConversation(llmMessages);
63 const currentSessionFile = ctx.sessionManager.getSessionFile();
64 const candidateFiles = [...extractCandidateFiles(branch, conversationText)];
65 const loadedSkills = extractLoadedSkills(branch);
66 const extractionModel = resolveExtractionModel(ctx) ?? ctx.model;
67
68 const result = await ctx.ui.custom<HandoffExtraction | null>((tui, theme, _kb, done) => {
69 const loader = new BorderedLoader(tui, theme, "Extracting handoff context...");
70 loader.onAbort = () => done(null);
71
72 const run = async () => {
73 const auth = await ctx.modelRegistry.getApiKeyAndHeaders(extractionModel);
74 if (!auth.ok) {
75 throw new Error(`Failed to get API key: ${auth.error}`);
76 }
77
78 return extractHandoffContext(
79 extractionModel,
80 auth.apiKey,
81 auth.headers,
82 conversationText,
83 goal,
84 candidateFiles,
85 loadedSkills,
86 loader.signal,
87 );
88 };
89
90 run()
91 .then(done)
92 .catch((err) => {
93 console.error("handoff generation failed", err);
94 done(null);
95 });
96
97 return loader;
98 });
99
100 if (!result) {
101 ctx.ui.notify("Handoff cancelled", "info");
102 return;
103 }
104
105 const prefillDraft = assembleHandoffDraft(result, goal, currentSessionFile);
106
107 const editedPrompt = await ctx.ui.editor("Edit handoff draft", prefillDraft);
108 if (editedPrompt === undefined) {
109 ctx.ui.notify("Handoff cancelled", "info");
110 return;
111 }
112
113 const next = await ctx.newSession({
114 parentSession: currentSessionFile ?? undefined,
115 });
116
117 if (next.cancelled) {
118 ctx.ui.notify("New session cancelled", "info");
119 return;
120 }
121
122 // Apply -model if specified
123 if (newSessionModel) {
124 await parseAndSetModelWithNotify(newSessionModel, ctx, pi);
125 }
126
127 const newSessionFile = ctx.sessionManager.getSessionFile();
128 ctx.ui.notify(`Switched to new session: ${newSessionFile ?? "(unnamed)"}`, "info");
129
130 ctx.ui.setEditorText(editedPrompt);
131 startCountdown(ctx);
132 },
133 });
134}