Preserve active subagents across reload
This commit is contained in:
+221
-161
@@ -66,32 +66,11 @@ import {
|
||||
/** Absolute path to `pi-extension/subagents`. https://github.com/nodejs/node/issues/37845 */
|
||||
const SUBAGENTS_DIR = dirname(fileURLToPath(import.meta.url));
|
||||
|
||||
// Survive /reload: clear timers and abort poll loops from the previous module load.
|
||||
// /reload re-imports this file, giving fresh module-level state, but closures from
|
||||
// the old module keep running. See https://github.com/HazAT/pi-interactive-subagents/issues/5
|
||||
const WIDGET_INTERVAL_KEY = Symbol.for("pi-subagents/widget-interval");
|
||||
const STATUS_INTERVAL_KEY = Symbol.for("pi-subagents/status-interval");
|
||||
const POLL_ABORT_KEY = Symbol.for("pi-subagents/poll-abort-controller");
|
||||
|
||||
{
|
||||
const prevInterval = (globalThis as any)[WIDGET_INTERVAL_KEY];
|
||||
if (prevInterval) {
|
||||
clearInterval(prevInterval);
|
||||
(globalThis as any)[WIDGET_INTERVAL_KEY] = null;
|
||||
}
|
||||
const prevStatusInterval = (globalThis as any)[STATUS_INTERVAL_KEY];
|
||||
if (prevStatusInterval) {
|
||||
clearInterval(prevStatusInterval);
|
||||
(globalThis as any)[STATUS_INTERVAL_KEY] = null;
|
||||
}
|
||||
const prevAbort = (globalThis as any)[POLL_ABORT_KEY] as AbortController | undefined;
|
||||
if (prevAbort) prevAbort.abort();
|
||||
(globalThis as any)[POLL_ABORT_KEY] = new AbortController();
|
||||
}
|
||||
|
||||
function getModuleAbortSignal(): AbortSignal {
|
||||
return ((globalThis as any)[POLL_ABORT_KEY] as AbortController).signal;
|
||||
}
|
||||
// Pi re-imports extension modules during /reload while callbacks from the old
|
||||
// module can still be running. Keep live subagent ownership process-global so a
|
||||
// fresh extension runtime can rebind to existing watchers instead of cancelling
|
||||
// their terminal sessions.
|
||||
const SHARED_RUNTIME_KEY = Symbol.for("pi-subagents/shared-runtime");
|
||||
|
||||
const SubagentParams = Type.Object({
|
||||
agent: Type.String({
|
||||
@@ -634,10 +613,41 @@ interface RunningSubagent {
|
||||
* subagent's pane (e.g. planner).
|
||||
*/
|
||||
interactive: boolean;
|
||||
/** Suppress terminal delivery when a real session shutdown cancels this run. */
|
||||
suppressDelivery?: boolean;
|
||||
}
|
||||
|
||||
/** All currently running subagents, keyed by id. */
|
||||
const runningSubagents = new Map<string, RunningSubagent>();
|
||||
type PendingDelivery = (pi: ExtensionAPI) => void;
|
||||
|
||||
interface SharedSubagentRuntime {
|
||||
runningSubagents: Map<string, RunningSubagent>;
|
||||
latestCtx: ExtensionContext | null;
|
||||
latestPi: ExtensionAPI | null;
|
||||
widgetInterval: ReturnType<typeof setInterval> | null;
|
||||
statusInterval: ReturnType<typeof setInterval> | null;
|
||||
pendingDeliveries: PendingDelivery[];
|
||||
deliveryFlushScheduled: boolean;
|
||||
activeRuntimeToken: symbol | null;
|
||||
}
|
||||
|
||||
function createSharedRuntime(): SharedSubagentRuntime {
|
||||
return {
|
||||
runningSubagents: new Map(),
|
||||
latestCtx: null,
|
||||
latestPi: null,
|
||||
widgetInterval: null,
|
||||
statusInterval: null,
|
||||
pendingDeliveries: [],
|
||||
deliveryFlushScheduled: false,
|
||||
activeRuntimeToken: null,
|
||||
};
|
||||
}
|
||||
|
||||
const sharedRuntime = ((globalThis as any)[SHARED_RUNTIME_KEY] ??=
|
||||
createSharedRuntime()) as SharedSubagentRuntime;
|
||||
|
||||
/** All currently running subagents, preserved across /reload. */
|
||||
const runningSubagents = sharedRuntime.runningSubagents;
|
||||
|
||||
// When this extension is loaded inside a subagent that itself spawns children
|
||||
// (e.g. a worker delegating to scout/researcher), `subagent-done.ts` runs in the
|
||||
@@ -648,18 +658,42 @@ const runningSubagents = new Map<string, RunningSubagent>();
|
||||
const RUNNING_CHILDREN_COUNT_KEY = Symbol.for("pi-subagents/running-children-count");
|
||||
(globalThis as any)[RUNNING_CHILDREN_COUNT_KEY] = () => runningSubagents.size;
|
||||
|
||||
// ── Widget management ──
|
||||
// ── Widget and runtime delivery management ──
|
||||
|
||||
/** Latest ExtensionContext from session_start, used for widget updates. */
|
||||
let latestCtx: ExtensionContext | null = null;
|
||||
/** Latest ExtensionAPI, used to deliver ask_question notifications from the watcher. */
|
||||
let latestPi: ExtensionAPI | null = null;
|
||||
function sendWhenRuntimeReady(delivery: PendingDelivery): void {
|
||||
const pi = sharedRuntime.latestPi;
|
||||
if (pi) {
|
||||
delivery(pi);
|
||||
return;
|
||||
}
|
||||
sharedRuntime.pendingDeliveries.push(delivery);
|
||||
}
|
||||
|
||||
/** Interval timer for widget re-renders. */
|
||||
let widgetInterval: ReturnType<typeof setInterval> | null = null;
|
||||
function flushPendingDeliveries(): void {
|
||||
sharedRuntime.deliveryFlushScheduled = false;
|
||||
const pi = sharedRuntime.latestPi;
|
||||
if (!pi || sharedRuntime.pendingDeliveries.length === 0) return;
|
||||
|
||||
/** Interval timer for status transition checks. */
|
||||
let statusInterval: ReturnType<typeof setInterval> | null = null;
|
||||
const pending = sharedRuntime.pendingDeliveries.splice(0);
|
||||
for (const delivery of pending) delivery(pi);
|
||||
}
|
||||
|
||||
function schedulePendingDeliveryFlush(): void {
|
||||
if (sharedRuntime.deliveryFlushScheduled) return;
|
||||
sharedRuntime.deliveryFlushScheduled = true;
|
||||
setImmediate(flushPendingDeliveries);
|
||||
}
|
||||
|
||||
function clearRefreshIntervals(): void {
|
||||
if (sharedRuntime.widgetInterval) {
|
||||
clearInterval(sharedRuntime.widgetInterval);
|
||||
sharedRuntime.widgetInterval = null;
|
||||
}
|
||||
if (sharedRuntime.statusInterval) {
|
||||
clearInterval(sharedRuntime.statusInterval);
|
||||
sharedRuntime.statusInterval = null;
|
||||
}
|
||||
}
|
||||
|
||||
function formatElapsedMMSS(startTime: number): string {
|
||||
const seconds = Math.floor((Date.now() - startTime) / 1000);
|
||||
@@ -755,19 +789,19 @@ function renderSubagentWidgetLines(agents: RunningSubagent[], width: number): st
|
||||
}
|
||||
|
||||
function updateWidget() {
|
||||
if (!latestCtx?.hasUI) return;
|
||||
const ctx = sharedRuntime.latestCtx;
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
if (runningSubagents.size === 0) {
|
||||
latestCtx.ui.setWidget("subagent-status", undefined);
|
||||
if (widgetInterval) {
|
||||
clearInterval(widgetInterval);
|
||||
widgetInterval = null;
|
||||
(globalThis as any)[WIDGET_INTERVAL_KEY] = null;
|
||||
ctx.ui.setWidget("subagent-status", undefined);
|
||||
if (sharedRuntime.widgetInterval) {
|
||||
clearInterval(sharedRuntime.widgetInterval);
|
||||
sharedRuntime.widgetInterval = null;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
latestCtx.ui.setWidget(
|
||||
ctx.ui.setWidget(
|
||||
"subagent-status",
|
||||
(_tui: any, _theme: any) => {
|
||||
return {
|
||||
@@ -1066,15 +1100,14 @@ function handleSubagentSteer(
|
||||
};
|
||||
}
|
||||
|
||||
function startStatusRefresh(pi: ExtensionAPI) {
|
||||
if (!statusConfig.enabled || statusInterval) return;
|
||||
function startStatusRefresh() {
|
||||
if (!statusConfig.enabled || sharedRuntime.statusInterval) return;
|
||||
|
||||
statusInterval = setInterval(() => {
|
||||
sharedRuntime.statusInterval = setInterval(() => {
|
||||
if (runningSubagents.size === 0) {
|
||||
if (statusInterval) {
|
||||
clearInterval(statusInterval);
|
||||
statusInterval = null;
|
||||
(globalThis as any)[STATUS_INTERVAL_KEY] = null;
|
||||
if (sharedRuntime.statusInterval) {
|
||||
clearInterval(sharedRuntime.statusInterval);
|
||||
sharedRuntime.statusInterval = null;
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -1104,19 +1137,19 @@ function startStatusRefresh(pi: ExtensionAPI) {
|
||||
|
||||
if (transitionLines.length > 0) {
|
||||
const capped = capStatusLines(transitionLines, statusConfig.lineLimit);
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_status",
|
||||
content: formatStatusAggregate(transitionLines, statusConfig.lineLimit),
|
||||
display: true,
|
||||
details: { lines: capped.visibleLines, overflow: capped.overflow },
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
sendWhenRuntimeReady((pi) => {
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_status",
|
||||
content: formatStatusAggregate(transitionLines, statusConfig.lineLimit),
|
||||
display: true,
|
||||
details: { lines: capped.visibleLines, overflow: capped.overflow },
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
});
|
||||
}
|
||||
}, 1000);
|
||||
|
||||
(globalThis as any)[STATUS_INTERVAL_KEY] = statusInterval;
|
||||
}
|
||||
|
||||
// Resuming a finished session is always autonomous: the relaunched agent runs
|
||||
@@ -1158,15 +1191,17 @@ export const __test__ = {
|
||||
contextWindowFor,
|
||||
formatUsageSegments,
|
||||
widgetIcon,
|
||||
sendWhenRuntimeReady,
|
||||
flushPendingDeliveries,
|
||||
sharedRuntime,
|
||||
};
|
||||
|
||||
function startWidgetRefresh() {
|
||||
if (widgetInterval) return;
|
||||
if (sharedRuntime.widgetInterval) return;
|
||||
updateWidget(); // immediate first render
|
||||
widgetInterval = setInterval(() => {
|
||||
sharedRuntime.widgetInterval = setInterval(() => {
|
||||
updateWidget();
|
||||
}, 1000);
|
||||
(globalThis as any)[WIDGET_INTERVAL_KEY] = widgetInterval;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -1531,20 +1566,22 @@ function deliverPendingQuestion(running: RunningSubagent): void {
|
||||
const elapsed = Math.floor((Date.now() - running.startTime) / 1000);
|
||||
const replyHint = `\n\nReply with subagent_message({ name: "${name}", message: "…" }) — the same name works whether it is still running or has since exited. It stays open until you reply.`;
|
||||
|
||||
latestPi?.sendMessage(
|
||||
{
|
||||
customType: "subagent_question",
|
||||
content: `Sub-agent "${name}" asks (${formatElapsed(elapsed)}):\n\n${payload.question}${replyHint}`,
|
||||
display: true,
|
||||
details: {
|
||||
name,
|
||||
agent: running.agent,
|
||||
question: payload.question,
|
||||
...(sessionId ? { sessionId } : {}),
|
||||
sendWhenRuntimeReady((pi) => {
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_question",
|
||||
content: `Sub-agent "${name}" asks (${formatElapsed(elapsed)}):\n\n${payload.question}${replyHint}`,
|
||||
display: true,
|
||||
details: {
|
||||
name,
|
||||
agent: running.agent,
|
||||
question: payload.question,
|
||||
...(sessionId ? { sessionId } : {}),
|
||||
},
|
||||
},
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
async function watchSubagent(
|
||||
@@ -1554,7 +1591,7 @@ async function watchSubagent(
|
||||
const { name, task, surface, startTime, sessionFile } = running;
|
||||
|
||||
try {
|
||||
const result = await pollForExit(surface, AbortSignal.any([signal, getModuleAbortSignal()]), {
|
||||
const result = await pollForExit(surface, signal, {
|
||||
interval: 1000,
|
||||
sessionFile,
|
||||
sentinelFile: running.sentinelFile,
|
||||
@@ -1668,35 +1705,45 @@ async function watchSubagent(
|
||||
}
|
||||
|
||||
export default function subagentsExtension(pi: ExtensionAPI) {
|
||||
latestPi = pi;
|
||||
// Capture the UI context for widget updates
|
||||
const runtimeToken = Symbol("pi-subagents/runtime");
|
||||
|
||||
// Bind only after Pi has installed the fresh runtime. During /reload the old
|
||||
// session_shutdown handler clears these references first, so watcher results
|
||||
// cannot enqueue a turn through stale extension APIs.
|
||||
pi.on("session_start", (_event, ctx) => {
|
||||
latestCtx = ctx;
|
||||
// pi runs multiple sessions in one process. A prior session's shutdown
|
||||
// aborts the shared module poll-abort controller; install a fresh one so
|
||||
// subagents spawned in this session aren't watched against a dead signal.
|
||||
// See https://github.com/HazAT/pi-interactive-subagents/issues/5
|
||||
const prevAbort = (globalThis as any)[POLL_ABORT_KEY] as AbortController | undefined;
|
||||
if (!prevAbort || prevAbort.signal.aborted) {
|
||||
(globalThis as any)[POLL_ABORT_KEY] = new AbortController();
|
||||
sharedRuntime.activeRuntimeToken = runtimeToken;
|
||||
sharedRuntime.latestPi = pi;
|
||||
sharedRuntime.latestCtx = ctx;
|
||||
|
||||
if (runningSubagents.size > 0) {
|
||||
startWidgetRefresh();
|
||||
startStatusRefresh();
|
||||
}
|
||||
// A completion can land in the short gap between old shutdown and this
|
||||
// fresh session_start. Flush on the next macrotask, after reload finishes.
|
||||
schedulePendingDeliveryFlush();
|
||||
});
|
||||
|
||||
// Clean up on session shutdown
|
||||
pi.on("session_shutdown", (_event, _ctx) => {
|
||||
if (widgetInterval) {
|
||||
clearInterval(widgetInterval);
|
||||
widgetInterval = null;
|
||||
(globalThis as any)[WIDGET_INTERVAL_KEY] = null;
|
||||
}
|
||||
if (statusInterval) {
|
||||
clearInterval(statusInterval);
|
||||
statusInterval = null;
|
||||
(globalThis as any)[STATUS_INTERVAL_KEY] = null;
|
||||
}
|
||||
const moduleAbort = (globalThis as any)[POLL_ABORT_KEY] as AbortController | undefined;
|
||||
if (moduleAbort) moduleAbort.abort();
|
||||
for (const [_id, agent] of runningSubagents) {
|
||||
pi.on("session_shutdown", (event, _ctx) => {
|
||||
// A delayed event from an invalidated extension instance must not tear down
|
||||
// state already rebound to the replacement runtime.
|
||||
if (sharedRuntime.activeRuntimeToken !== runtimeToken) return;
|
||||
|
||||
sharedRuntime.activeRuntimeToken = null;
|
||||
sharedRuntime.latestPi = null;
|
||||
sharedRuntime.latestCtx = null;
|
||||
clearRefreshIntervals();
|
||||
|
||||
// /reload keeps the same parent session and terminal surfaces. Existing
|
||||
// watchers continue and rebind to the fresh runtime in session_start.
|
||||
if (event.reason === "reload") return;
|
||||
|
||||
// Real session replacement/quit owns a different lifecycle. Cancel those
|
||||
// children and suppress the synthetic cancellation result; there is no
|
||||
// valid destination session to receive it.
|
||||
sharedRuntime.pendingDeliveries.length = 0;
|
||||
for (const agent of runningSubagents.values()) {
|
||||
agent.suppressDelivery = true;
|
||||
agent.abortController?.abort();
|
||||
}
|
||||
runningSubagents.clear();
|
||||
@@ -1848,47 +1895,54 @@ export default function subagentsExtension(pi: ExtensionAPI) {
|
||||
|
||||
// Start widget refresh and status supervision when the first agent launches
|
||||
startWidgetRefresh();
|
||||
startStatusRefresh(pi);
|
||||
startStatusRefresh();
|
||||
|
||||
// Fire-and-forget: start watching in background
|
||||
// Fire-and-forget: start watching in background. Delivery resolves the
|
||||
// current Pi runtime at completion time so /reload cannot retain a stale
|
||||
// ExtensionAPI closure.
|
||||
watchSubagent(running, watcherAbort.signal)
|
||||
.then((result) => {
|
||||
updateWidget(); // reflect removal from Map immediately
|
||||
if (running.suppressDelivery) return;
|
||||
|
||||
const presentation = resolveResultPresentation(result, running.name);
|
||||
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: presentation,
|
||||
display: true,
|
||||
details: {
|
||||
name: running.name,
|
||||
task: running.task,
|
||||
agent: running.agent,
|
||||
exitCode: result.exitCode,
|
||||
elapsed: result.elapsed,
|
||||
sessionFile: result.sessionFile,
|
||||
...(result.sessionId ? { sessionId: result.sessionId } : {}),
|
||||
...(result.errorMessage ? { errorMessage: result.errorMessage } : {}),
|
||||
...(result.claudeSessionId ? { claudeSessionId: result.claudeSessionId } : {}),
|
||||
...(result.stats ? { stats: result.stats } : {}),
|
||||
sendWhenRuntimeReady((currentPi) => {
|
||||
currentPi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: presentation,
|
||||
display: true,
|
||||
details: {
|
||||
name: running.name,
|
||||
task: running.task,
|
||||
agent: running.agent,
|
||||
exitCode: result.exitCode,
|
||||
elapsed: result.elapsed,
|
||||
sessionFile: result.sessionFile,
|
||||
...(result.sessionId ? { sessionId: result.sessionId } : {}),
|
||||
...(result.errorMessage ? { errorMessage: result.errorMessage } : {}),
|
||||
...(result.claudeSessionId ? { claudeSessionId: result.claudeSessionId } : {}),
|
||||
...(result.stats ? { stats: result.stats } : {}),
|
||||
},
|
||||
},
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
});
|
||||
})
|
||||
.catch((err) => {
|
||||
updateWidget();
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: `Sub-agent "${running.name}" error: ${err?.message ?? String(err)}`,
|
||||
display: true,
|
||||
details: { name: running.name, task: running.task, error: err?.message },
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
if (running.suppressDelivery) return;
|
||||
sendWhenRuntimeReady((currentPi) => {
|
||||
currentPi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: `Sub-agent "${running.name}" error: ${err?.message ?? String(err)}`,
|
||||
display: true,
|
||||
details: { name: running.name, task: running.task, error: err?.message },
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
// Return immediately
|
||||
@@ -2286,7 +2340,7 @@ export default function subagentsExtension(pi: ExtensionAPI) {
|
||||
};
|
||||
runningSubagents.set(id, running);
|
||||
startWidgetRefresh();
|
||||
startStatusRefresh(pi);
|
||||
startStatusRefresh();
|
||||
|
||||
// Fire-and-forget watcher
|
||||
const watcherAbort = new AbortController();
|
||||
@@ -2295,6 +2349,7 @@ export default function subagentsExtension(pi: ExtensionAPI) {
|
||||
watchSubagent(running, watcherAbort.signal)
|
||||
.then((result) => {
|
||||
updateWidget();
|
||||
if (running.suppressDelivery) return;
|
||||
|
||||
const allEntries = getNewEntries(sessionPath, entryCountBefore);
|
||||
const summary = findLastAssistantMessage(allEntries) ??
|
||||
@@ -2308,35 +2363,40 @@ export default function subagentsExtension(pi: ExtensionAPI) {
|
||||
name,
|
||||
);
|
||||
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: presentation,
|
||||
display: true,
|
||||
details: {
|
||||
name,
|
||||
task: message,
|
||||
exitCode: result.exitCode,
|
||||
elapsed: result.elapsed,
|
||||
sessionFile: sessionPath,
|
||||
sessionId: resumedSessionId,
|
||||
...(result.errorMessage ? { errorMessage: result.errorMessage } : {}),
|
||||
sendWhenRuntimeReady((currentPi) => {
|
||||
currentPi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: presentation,
|
||||
display: true,
|
||||
details: {
|
||||
name,
|
||||
task: message,
|
||||
exitCode: result.exitCode,
|
||||
elapsed: result.elapsed,
|
||||
sessionFile: sessionPath,
|
||||
sessionId: resumedSessionId,
|
||||
...(result.errorMessage ? { errorMessage: result.errorMessage } : {}),
|
||||
},
|
||||
},
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
});
|
||||
})
|
||||
.catch((err) => {
|
||||
updateWidget();
|
||||
pi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: `Resume error: ${err?.message ?? String(err)}`,
|
||||
display: true,
|
||||
details: { name, error: err?.message },
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
if (running.suppressDelivery) return;
|
||||
sendWhenRuntimeReady((currentPi) => {
|
||||
currentPi.sendMessage(
|
||||
{
|
||||
customType: "subagent_result",
|
||||
content: `Resume error: ${err?.message ?? String(err)}`,
|
||||
display: true,
|
||||
details: { name, error: err?.message },
|
||||
},
|
||||
{ triggerTurn: true, deliverAs: "steer" },
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
return {
|
||||
|
||||
Reference in New Issue
Block a user