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