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