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