1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Petr Baudis <pasky@ucw.cz>
3//
4// SPDX-License-Identifier: MIT
5
6import type { ExtensionAPI, SessionEntry } from "@mariozechner/pi-coding-agent";
7import { convertToLlm, serializeConversation } from "@mariozechner/pi-coding-agent";
8import { Type } from "@sinclair/typebox";
9import { extractCandidateFiles, extractLoadedSkills, resolveExtractionModel } from "./session-analysis.js";
10import { type HandoffExtraction, assembleHandoffDraft, extractHandoffContext } from "./handoff-extraction.js";
11
12export type PendingHandoff = {
13 prompt: string;
14 parentSession: string | undefined;
15 newModel: string | undefined;
16};
17
18export function registerHandoffTool(pi: ExtensionAPI, setPendingHandoff: (h: PendingHandoff) => void) {
19 pi.registerTool({
20 name: "handoff",
21 label: "Handoff",
22 description:
23 "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.",
24 parameters: Type.Object({
25 goal: Type.String({
26 description: "The goal/task for the new session",
27 }),
28 model: Type.Optional(
29 Type.String({
30 description: "Model for the new session as provider/modelId (e.g. 'anthropic/claude-haiku-4-5')",
31 }),
32 ),
33 }),
34 async execute(_toolCallId, params, signal, _onUpdate, ctx) {
35 if (!ctx.model) {
36 return {
37 content: [{ type: "text" as const, text: "No model selected." }],
38 details: undefined,
39 };
40 }
41
42 const branch = ctx.sessionManager.getBranch();
43 const messages = branch
44 .filter((entry): entry is SessionEntry & { type: "message" } => entry.type === "message")
45 .map((entry) => entry.message);
46
47 if (messages.length === 0) {
48 return {
49 content: [
50 {
51 type: "text" as const,
52 text: "No conversation to hand off.",
53 },
54 ],
55 details: undefined,
56 };
57 }
58
59 const llmMessages = convertToLlm(messages);
60 const conversationText = serializeConversation(llmMessages);
61 const candidateFiles = [...extractCandidateFiles(branch, conversationText)];
62 const loadedSkills = extractLoadedSkills(branch);
63
64 const extractionModel = resolveExtractionModel(ctx) ?? ctx.model;
65 const auth = await ctx.modelRegistry.getApiKeyAndHeaders(extractionModel);
66 if (!auth.ok) {
67 return {
68 content: [
69 {
70 type: "text" as const,
71 text: `Failed to get API key: ${auth.error}`,
72 },
73 ],
74 details: undefined,
75 };
76 }
77
78 let result: HandoffExtraction | null;
79 try {
80 result = await extractHandoffContext(
81 extractionModel,
82 auth.apiKey,
83 auth.headers,
84 conversationText,
85 params.goal,
86 candidateFiles,
87 loadedSkills,
88 signal,
89 );
90 } catch (err) {
91 if (signal?.aborted) {
92 return {
93 content: [
94 {
95 type: "text" as const,
96 text: "Handoff cancelled.",
97 },
98 ],
99 details: undefined,
100 };
101 }
102 return {
103 content: [
104 {
105 type: "text" as const,
106 text: `Handoff extraction failed: ${String(err)}`,
107 },
108 ],
109 details: undefined,
110 };
111 }
112
113 if (!result) {
114 return {
115 content: [
116 {
117 type: "text" as const,
118 text: "Handoff extraction failed or was cancelled.",
119 },
120 ],
121 details: undefined,
122 };
123 }
124
125 const currentSessionFile = ctx.sessionManager.getSessionFile();
126 const prompt = assembleHandoffDraft(result, params.goal, currentSessionFile);
127
128 setPendingHandoff({
129 prompt,
130 parentSession: currentSessionFile,
131 newModel: params.model,
132 });
133
134 return {
135 content: [
136 {
137 type: "text" as const,
138 text: "Handoff prepared. The session will switch after this turn completes.",
139 },
140 ],
141 details: undefined,
142 };
143 },
144 });
145}