400 lines
15 KiB
TypeScript
400 lines
15 KiB
TypeScript
/**
|
|
* Extension loaded into sub-agents.
|
|
* - Shows agent identity + available tools as a styled widget above the editor (toggle with Ctrl+Alt+O)
|
|
* - Provides an `ask_question` tool for asking the parent orchestrator a question
|
|
*
|
|
* Subagents do NOT self-terminate via a tool. Auto-exit agents shut down
|
|
* automatically when their agent loop ends (see the `agent_end` handler);
|
|
* interactive agents end when the human exits the pane.
|
|
*
|
|
* `ask_question` keeps the session OPEN: it writes a `${sessionFile}.ask`
|
|
* signal the parent's watcher picks up, parks the session in a "waiting" state
|
|
* (auto-exit is suppressed for that turn via `awaitingAnswer`), and the parent
|
|
* replies with subagent_message — which lands as the subagent's next turn.
|
|
*/
|
|
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
|
|
import { Box, Text } from "@earendil-works/pi-tui";
|
|
import { Type } from "typebox";
|
|
import { writeFileSync } from "node:fs";
|
|
import { createSubagentActivityRecorder } from "./activity.ts";
|
|
|
|
export function shouldMarkUserTookOver(agentStarted: boolean): boolean {
|
|
return agentStarted;
|
|
}
|
|
|
|
/**
|
|
* Number of child subagents this session itself still has in flight.
|
|
*
|
|
* When this extension is loaded inside a subagent that can spawn its own
|
|
* children (e.g. a worker delegating to scout/researcher), `index.ts` runs in
|
|
* the same process and publishes a live count through a shared process-global
|
|
* symbol. A subagent that spawns children and then writes a "waiting for
|
|
* results" message would otherwise auto-exit the instant that turn ends —
|
|
* killing the session before its children report back. Reading this count lets
|
|
* `agent_end` keep the session open until every child has finished and its
|
|
* result has been delivered.
|
|
*
|
|
* Returns 0 when the spawning tools aren't loaded (scout/researcher, or a
|
|
* standalone session), so those agents auto-exit exactly as before.
|
|
*/
|
|
export function runningChildrenCount(): number {
|
|
const fn = (globalThis as any)[Symbol.for("pi-subagents/running-children-count")];
|
|
if (typeof fn !== "function") return 0;
|
|
try {
|
|
const n = fn();
|
|
return typeof n === "number" && n > 0 ? n : 0;
|
|
} catch {
|
|
return 0;
|
|
}
|
|
}
|
|
|
|
export function shouldAutoExitOnAgentEnd(
|
|
_userTookOver: boolean,
|
|
messages: any[] | undefined,
|
|
): boolean {
|
|
// Manual input should not strand an auto-exit subagent. If the latest agent
|
|
// turn completed normally, close the session. Escape/abort still leaves it
|
|
// open for inspection or another prompt.
|
|
//
|
|
// stopReason: "error" (e.g. exhausted retries on a provider overload) also
|
|
// returns true — we want to shut down so the parent is woken up — but we
|
|
// pair this with findLatestAssistantError() so the parent learns it was an
|
|
// error, not a clean completion.
|
|
if (messages) {
|
|
for (let i = messages.length - 1; i >= 0; i--) {
|
|
const msg = messages[i];
|
|
if (msg?.role === "assistant") {
|
|
return msg.stopReason !== "aborted";
|
|
}
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
export interface SubagentErrorInfo {
|
|
errorMessage: string;
|
|
stopReason: "error";
|
|
}
|
|
|
|
/**
|
|
* If the last assistant message in the turn ended with `stopReason: "error"`
|
|
* (typically auto-retry exhausted on an overload / rate limit / server error),
|
|
* return its error info so the parent orchestrator can surface a clear
|
|
* failure instead of silently treating the run as completed.
|
|
*
|
|
* Returns `null` when the latest assistant turn completed normally or was
|
|
* aborted by the user (handled separately by shouldAutoExitOnAgentEnd).
|
|
*/
|
|
export function findLatestAssistantError(
|
|
messages: any[] | undefined,
|
|
): SubagentErrorInfo | null {
|
|
if (!messages) return null;
|
|
for (let i = messages.length - 1; i >= 0; i--) {
|
|
const msg = messages[i];
|
|
if (msg?.role !== "assistant") continue;
|
|
if (msg.stopReason !== "error") return null;
|
|
const raw = typeof msg.errorMessage === "string" ? msg.errorMessage.trim() : "";
|
|
return {
|
|
errorMessage: raw || "Subagent agent loop ended with stopReason=error (no errorMessage field).",
|
|
stopReason: "error",
|
|
};
|
|
}
|
|
return null;
|
|
}
|
|
|
|
export function parseDeniedTools(rawValue: string | undefined): string[] {
|
|
return (rawValue ?? "")
|
|
.split(",")
|
|
.map((value) => value.trim())
|
|
.filter(Boolean);
|
|
}
|
|
|
|
export default function (pi: ExtensionAPI) {
|
|
let toolNames: string[] = [];
|
|
let denied: string[] = [];
|
|
let expanded = false;
|
|
|
|
// Read subagent identity from env vars (set by parent orchestrator)
|
|
const subagentName = process.env.PI_SUBAGENT_NAME ?? "";
|
|
const subagentAgent = process.env.PI_SUBAGENT_AGENT ?? "";
|
|
const deniedToolsValue = process.env.PI_DENY_TOOLS;
|
|
const autoExit = process.env.PI_SUBAGENT_AUTO_EXIT === "1";
|
|
const recorder = createSubagentActivityRecorder({
|
|
runningChildId: process.env.PI_SUBAGENT_ID,
|
|
activityFile: process.env.PI_SUBAGENT_ACTIVITY_FILE,
|
|
});
|
|
|
|
function renderWidget(ctx: { ui: { setWidget: Function } }, _theme: any) {
|
|
ctx.ui.setWidget(
|
|
"subagent-tools",
|
|
(_tui: any, theme: any) => {
|
|
const box = new Box(1, 0, (text: string) => theme.bg("toolSuccessBg", text));
|
|
|
|
const label = subagentAgent || subagentName;
|
|
const agentTag = label ? theme.bold(theme.fg("accent", `[${label}]`)) : "";
|
|
|
|
if (expanded) {
|
|
// Expanded: full tool list + denied
|
|
const countInfo = theme.fg("dim", ` — ${toolNames.length} available`);
|
|
const hint = theme.fg("muted", " (Ctrl+Alt+O to collapse)");
|
|
|
|
const toolList = toolNames
|
|
.map((name: string) => theme.fg("dim", name))
|
|
.join(theme.fg("muted", ", "));
|
|
|
|
let deniedLine = "";
|
|
if (denied.length > 0) {
|
|
const deniedList = denied
|
|
.map((name: string) => theme.fg("error", name))
|
|
.join(theme.fg("muted", ", "));
|
|
deniedLine = "\n" + theme.fg("muted", "denied: ") + deniedList;
|
|
}
|
|
|
|
const content = new Text(
|
|
`${agentTag}${countInfo}${hint}\n${toolList}${deniedLine}`,
|
|
0,
|
|
0,
|
|
);
|
|
box.addChild(content);
|
|
} else {
|
|
// Collapsed: one-line summary
|
|
const countInfo = theme.fg("dim", ` — ${toolNames.length} tools`);
|
|
const deniedInfo =
|
|
denied.length > 0
|
|
? theme.fg("dim", " · ") + theme.fg("error", `${denied.length} denied`)
|
|
: "";
|
|
const hint = theme.fg("muted", " (Ctrl+Alt+O to expand)");
|
|
|
|
const content = new Text(`${agentTag}${countInfo}${deniedInfo}${hint}`, 0, 0);
|
|
box.addChild(content);
|
|
}
|
|
|
|
return box;
|
|
},
|
|
{ placement: "aboveEditor" },
|
|
);
|
|
}
|
|
|
|
let userTookOver = false;
|
|
let agentStarted = false;
|
|
// Set when ask_question is called; suppresses auto-exit so the session stays
|
|
// open while it waits for the orchestrator's reply. Cleared when the reply
|
|
// lands — on `input` (covers a reply steered into the current run) and on
|
|
// `agent_start` (covers a reply that starts a fresh turn after parking).
|
|
let awaitingAnswer = false;
|
|
|
|
// Show widget + status bar on session start
|
|
pi.on("session_start", (_event, ctx) => {
|
|
recorder.sessionStart();
|
|
const tools = pi.getAllTools();
|
|
toolNames = tools.map((t) => t.name).sort();
|
|
denied = parseDeniedTools(deniedToolsValue);
|
|
|
|
renderWidget(ctx, null);
|
|
});
|
|
|
|
pi.on("input", () => {
|
|
recorder.input();
|
|
// A submitted message is the orchestrator's (or a human's) reply — the
|
|
// pending ask_question has been answered, however it was delivered. Clear
|
|
// here, not only on agent_start, because a reply steered in *mid-run* is
|
|
// absorbed into the current run (pi's `steer` behavior injects it before
|
|
// the next LLM call): no new agent_start fires, so without this the flag
|
|
// would stay set and agent_end would park the session as `waiting` even
|
|
// though the answer already arrived and was consumed. (The `input` event
|
|
// fires for mid-run steers because prompt() emits it before queueing.)
|
|
awaitingAnswer = false;
|
|
// Ignore the initial task message that starts an autonomous subagent.
|
|
// Only inputs after the first agent run has started count as user takeover.
|
|
if (!shouldMarkUserTookOver(agentStarted)) return;
|
|
userTookOver = true;
|
|
});
|
|
|
|
pi.on("before_agent_start", () => {
|
|
recorder.beforeAgentStart();
|
|
});
|
|
|
|
pi.on("agent_start", () => {
|
|
agentStarted = true;
|
|
// A new turn is starting — any pending ask_question has now been answered
|
|
// (or superseded), so let auto-exit resume normally when this turn ends.
|
|
awaitingAnswer = false;
|
|
recorder.agentStart();
|
|
});
|
|
|
|
pi.on("agent_end", (event, ctx) => {
|
|
const messages = (event as any).messages as any[] | undefined;
|
|
// Never shut down while this session still has work in flight:
|
|
// - awaitingAnswer: an ask_question is pending the orchestrator's reply.
|
|
// - runningChildrenCount(): this subagent spawned its own children and is
|
|
// waiting for their results (delivered as steered turns). Exiting now
|
|
// would strand those children and drop their results.
|
|
// In both cases the session parks as `waiting` and resumes when the next
|
|
// turn lands.
|
|
const hasPendingChildren = runningChildrenCount() > 0;
|
|
const shouldExit =
|
|
!awaitingAnswer &&
|
|
!hasPendingChildren &&
|
|
autoExit &&
|
|
shouldAutoExitOnAgentEnd(userTookOver, messages);
|
|
|
|
if (shouldExit) {
|
|
// Surface stopReason: "error" turns (auto-retry exhausted, provider
|
|
// overload, etc.) to the parent via the .exit sidecar so the watcher
|
|
// can report a clear failure with the underlying error message.
|
|
// Without this the parent would only see exit code 0 and a stale
|
|
// assistant message, mistaking the crash for a successful completion.
|
|
const errorInfo = findLatestAssistantError(messages);
|
|
const sessionFile = process.env.PI_SUBAGENT_SESSION;
|
|
if (errorInfo && sessionFile) {
|
|
try {
|
|
writeFileSync(
|
|
`${sessionFile}.exit`,
|
|
JSON.stringify({
|
|
type: "error",
|
|
errorMessage: errorInfo.errorMessage,
|
|
stopReason: errorInfo.stopReason,
|
|
}),
|
|
);
|
|
} catch {
|
|
// Best effort — even without the sidecar, watcher's session-file
|
|
// fallback can still recover the errorMessage.
|
|
}
|
|
}
|
|
|
|
recorder.agentEndDone();
|
|
ctx.shutdown();
|
|
return;
|
|
}
|
|
|
|
recorder.agentEndWaiting();
|
|
if (autoExit) {
|
|
// Reset any recorded manual input marker. Auto-exit is decided by whether
|
|
// the latest agent turn completed normally, not by who initiated it.
|
|
userTookOver = false;
|
|
}
|
|
});
|
|
|
|
pi.on("turn_start", (event) => {
|
|
recorder.turnStart((event as any).turnIndex);
|
|
});
|
|
|
|
pi.on("turn_end", (event) => {
|
|
recorder.turnEnd((event as any).turnIndex);
|
|
});
|
|
|
|
pi.on("before_provider_request", () => {
|
|
recorder.beforeProviderRequest();
|
|
});
|
|
|
|
pi.on("after_provider_response", () => {
|
|
recorder.afterProviderResponse();
|
|
});
|
|
|
|
pi.on("message_update", (event) => {
|
|
recorder.messageUpdate((event as any).assistantMessageEvent?.type);
|
|
});
|
|
|
|
pi.on("tool_execution_start", (event) => {
|
|
recorder.toolExecutionStart((event as any).toolCallId, (event as any).toolName);
|
|
});
|
|
|
|
pi.on("tool_call", (event) => {
|
|
recorder.toolCall((event as any).toolCallId, (event as any).toolName);
|
|
});
|
|
|
|
pi.on("tool_execution_update", (event) => {
|
|
recorder.toolExecutionUpdate((event as any).toolCallId, (event as any).toolName);
|
|
});
|
|
|
|
pi.on("tool_result", (event) => {
|
|
recorder.toolResult((event as any).toolCallId, (event as any).toolName);
|
|
});
|
|
|
|
pi.on("tool_execution_end", (event) => {
|
|
recorder.toolExecutionEnd((event as any).toolCallId, (event as any).toolName);
|
|
});
|
|
|
|
pi.on("session_shutdown", (event) => {
|
|
recorder.sessionShutdown((event as any).reason);
|
|
});
|
|
|
|
// Toggle expand/collapse with Ctrl+Alt+O
|
|
pi.registerShortcut("ctrl+alt+o", {
|
|
description: "Toggle subagent tools widget",
|
|
handler: (ctx) => {
|
|
expanded = !expanded;
|
|
renderWidget(ctx, null);
|
|
},
|
|
});
|
|
|
|
pi.registerTool({
|
|
name: "ask_question",
|
|
label: "ask_question",
|
|
description:
|
|
"Ask the orchestrator (the parent agent that spawned you) a single question and pause until they reply. " +
|
|
"Use this when requirements are ambiguous, a decision would materially affect your work, you're blocked, " +
|
|
"or you need information or confirmation only the orchestrator has. Prefer asking over guessing. " +
|
|
"Your session stays open while you wait — the answer arrives as your next message, then you continue. " +
|
|
"Ask exactly one question per call; make separate calls for unrelated questions.",
|
|
promptSnippet:
|
|
"Use this tool to ask the orchestrator one clarifying, missing-requirement, preference, or decision question before continuing — instead of guessing.",
|
|
promptGuidelines: [
|
|
"Ask exactly one question per tool call.",
|
|
"If you need answers to multiple things, make separate ask_question calls instead of bundling them.",
|
|
"Prefer this tool over guessing when requirements, preferences, or implementation choices are unclear.",
|
|
"Use it when multiple valid paths exist and the right one depends on the orchestrator's intent.",
|
|
"Give enough context in the question that the orchestrator can answer without re-reading your whole task.",
|
|
"After asking, stop and wait — the reply will arrive as your next message.",
|
|
],
|
|
parameters: Type.Object({
|
|
question: Type.String({
|
|
description:
|
|
"The single freeform question to ask the orchestrator. Include enough context to answer it directly.",
|
|
}),
|
|
}),
|
|
async execute(_toolCallId, params, _signal, _onUpdate, _ctx) {
|
|
const sessionFile = process.env.PI_SUBAGENT_SESSION;
|
|
if (!sessionFile) {
|
|
throw new Error(
|
|
"ask_question is only available in subagent contexts. " +
|
|
"PI_SUBAGENT_SESSION environment variable is not set.",
|
|
);
|
|
}
|
|
|
|
// Keep the session open: suppress auto-exit for this turn and park in the
|
|
// "waiting" phase. The parent's watcher picks up the `.ask` signal and
|
|
// notifies the orchestrator, who replies via subagent_message.
|
|
awaitingAnswer = true;
|
|
recorder.askQuestion();
|
|
const askData = {
|
|
name: process.env.PI_SUBAGENT_NAME ?? "subagent",
|
|
agent: process.env.PI_SUBAGENT_AGENT ?? "",
|
|
question: params.question,
|
|
};
|
|
writeFileSync(`${sessionFile}.ask`, JSON.stringify(askData));
|
|
|
|
return {
|
|
content: [
|
|
{
|
|
type: "text",
|
|
text:
|
|
"Question sent to the orchestrator. Stop here and wait — do not continue working or " +
|
|
"assume an answer. Their reply will arrive as your next message.",
|
|
},
|
|
],
|
|
details: { question: params.question },
|
|
};
|
|
},
|
|
|
|
renderCall(args, theme) {
|
|
const text =
|
|
theme.fg("toolTitle", theme.bold("ask_question ")) +
|
|
theme.fg("muted", String((args as any).question ?? ""));
|
|
return new Text(text, 0, 0);
|
|
},
|
|
});
|
|
|
|
}
|