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, ExtensionContext } from "@mariozechner/pi-coding-agent";
7import type { SessionManager } from "@mariozechner/pi-coding-agent";
8import { Key, matchesKey } from "@mariozechner/pi-tui";
9import { registerHandoffCommand } from "./handoff-command.js";
10import { registerHandoffTool, type PendingHandoff } from "./handoff-tool.js";
11import { registerSessionQueryTool } from "./session-query-tool.js";
12
13const STATUS_KEY = "handoff";
14const COUNTDOWN_SECONDS = 10;
15
16type PendingAutoSubmit = {
17 ctx: ExtensionContext;
18 sessionFile: string | undefined;
19 interval: ReturnType<typeof setInterval>;
20 unsubscribeInput: () => void;
21};
22
23function isEditableInput(data: string): boolean {
24 if (!data) return false;
25 if (data.length === 1) {
26 const code = data.charCodeAt(0);
27 if (code >= 32 && code !== 127) return true;
28 if (code === 8 || code === 13) return true;
29 }
30
31 if (data === "\n" || data === "\r") return true;
32 if (data === "\x7f") return true;
33
34 if (data.length > 1 && !data.startsWith("\x1b")) return true;
35
36 return false;
37}
38
39function statusLine(ctx: ExtensionContext, seconds: number): string {
40 const accent = ctx.ui.theme.fg("accent", `handoff auto-submit in ${seconds}s`);
41 const hint = ctx.ui.theme.fg("dim", "(type to edit, Esc to cancel)");
42 return `${accent} ${hint}`;
43}
44
45export default function (pi: ExtensionAPI) {
46 let pending: PendingAutoSubmit | null = null;
47 let pendingHandoff: PendingHandoff | null = null;
48 let handoffTimestamp: number | null = null;
49
50 const clearPending = (ctx?: ExtensionContext, notify?: string) => {
51 if (!pending) return;
52
53 clearInterval(pending.interval);
54 pending.unsubscribeInput();
55 pending.ctx.ui.setStatus(STATUS_KEY, undefined);
56
57 const local = pending;
58 pending = null;
59
60 if (notify && ctx) {
61 ctx.ui.notify(notify, "info");
62 } else if (notify) {
63 local.ctx.ui.notify(notify, "info");
64 }
65 };
66
67 const autoSubmitDraft = () => {
68 if (!pending) return;
69
70 const active = pending;
71 const currentSession = active.ctx.sessionManager.getSessionFile();
72 if (active.sessionFile && currentSession !== active.sessionFile) {
73 clearPending(undefined);
74 return;
75 }
76
77 const draft = active.ctx.ui.getEditorText().trim();
78 clearPending(undefined);
79
80 if (!draft) {
81 active.ctx.ui.notify("Handoff draft is empty", "warning");
82 return;
83 }
84
85 active.ctx.ui.setEditorText("");
86
87 try {
88 if (active.ctx.isIdle()) {
89 pi.sendUserMessage(draft);
90 } else {
91 pi.sendUserMessage(draft, { deliverAs: "followUp" });
92 }
93 } catch {
94 pi.sendUserMessage(draft);
95 }
96 };
97
98 const startCountdown = (ctx: ExtensionContext) => {
99 clearPending(ctx);
100
101 let seconds = COUNTDOWN_SECONDS;
102 ctx.ui.setStatus(STATUS_KEY, statusLine(ctx, seconds));
103
104 const unsubscribeInput = ctx.ui.onTerminalInput((data) => {
105 if (matchesKey(data, Key.escape)) {
106 clearPending(ctx, "Handoff auto-submit cancelled");
107 return { consume: true };
108 }
109
110 if (isEditableInput(data)) {
111 clearPending(ctx, "Handoff auto-submit stopped (editing)");
112 }
113
114 return undefined;
115 });
116
117 const interval = setInterval(() => {
118 if (!pending) return;
119
120 seconds -= 1;
121 if (seconds <= 0) {
122 autoSubmitDraft();
123 return;
124 }
125
126 ctx.ui.setStatus(STATUS_KEY, statusLine(ctx, seconds));
127 }, 1000);
128
129 pending = {
130 ctx,
131 sessionFile: ctx.sessionManager.getSessionFile(),
132 interval,
133 unsubscribeInput,
134 };
135 };
136
137 // --- Event handlers ---
138
139 pi.on("session_before_switch", (_event, ctx) => {
140 if (pending) clearPending(ctx);
141 });
142
143 pi.on("session_start", (event, ctx) => {
144 if (pending) clearPending(ctx);
145 // A proper session switch (e.g. /new) or fork fully resets agent state,
146 // so clear the context filter to avoid hiding new messages.
147 if (event.reason === "new" || event.reason === "resume" || event.reason === "fork") {
148 handoffTimestamp = null;
149 }
150 });
151
152 pi.on("session_before_fork", (_event, ctx) => {
153 if (pending) clearPending(ctx);
154 });
155
156 pi.on("session_before_tree", (_event, ctx) => {
157 if (pending) clearPending(ctx);
158 });
159
160 pi.on("session_tree", (_event, ctx) => {
161 if (pending) clearPending(ctx);
162 });
163
164 pi.on("session_shutdown", (_event, ctx) => {
165 if (pending) clearPending(ctx);
166 });
167
168 // --- Tool-path handoff coordination ---
169 //
170 // The /handoff command has ExtensionCommandContext with ctx.newSession()
171 // which does a full agent reset. The handoff tool only gets
172 // ExtensionContext, which lacks newSession(). So the tool stores a
173 // pending handoff and these handlers complete it:
174 //
175 // 1. agent_end: after the agent loop finishes, switch sessions and
176 // send the handoff prompt in the next macrotask
177 // 2. context: filter pre-handoff messages since the low-level session
178 // switch doesn't clear agent.state.messages
179 // 3. session_start (above): clears the filter on proper switches
180
181 pi.on("agent_end", (_event, ctx) => {
182 if (!pendingHandoff) return;
183
184 const { prompt, parentSession, newModel } = pendingHandoff;
185 pendingHandoff = null;
186
187 // Create new session first — this triggers session_start which clears
188 // any stale handoffTimestamp. Then set the timestamp for the current
189 // handoff so the context filter is active for the new session's first turn.
190 (ctx.sessionManager as SessionManager).newSession({ parentSession });
191 handoffTimestamp = Date.now();
192
193 setTimeout(async () => {
194 if (newModel) {
195 const slashIdx = newModel.indexOf("/");
196 if (slashIdx > 0) {
197 const provider = newModel.slice(0, slashIdx);
198 const modelId = newModel.slice(slashIdx + 1);
199 const model = ctx.modelRegistry.find(provider, modelId);
200 if (model) await pi.setModel(model);
201 }
202 }
203 pi.sendUserMessage(prompt);
204 }, 0);
205 });
206
207 pi.on("context", (event, _ctx) => {
208 if (handoffTimestamp === null) return;
209
210 const cutoff = handoffTimestamp;
211 const newMessages = event.messages.filter((m) => m.timestamp >= cutoff);
212 if (newMessages.length > 0) {
213 return { messages: newMessages };
214 }
215 });
216
217 // --- Register tools and commands ---
218
219 registerSessionQueryTool(pi);
220 registerHandoffTool(pi, (h) => {
221 pendingHandoff = h;
222 });
223 registerHandoffCommand(pi, startCountdown);
224}