import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; import { Text } from "@earendil-works/pi-tui"; import { Type } from "typebox"; import type { AgentRegistry } from "../agents/registry"; import type { AgentTask } from "../agents/types"; import type { AgentManager } from "../runtime/agentManager"; import { buildAgentTask } from "../runtime/buildAgentTask"; import { createLifecycleWidgetEmitter, createParentSessionWriter, getParentSessionReference, } from "../runtime/sessionRelationship"; const agentToolSchema = Type.Object({ agent: Type.String({ description: "Name of the agent definition to run" }), prompt: Type.String({ description: "Task prompt for the selected agent" }), background: Type.Optional(Type.Boolean({ description: "Run in background and return an agent id immediately" })), }); const AGENT_TOOL_DESCRIPTION = "Delegate independent research or implementation work to one named sub-agent using the shared agent manager."; const AGENT_TOOL_PROMPT_SNIPPET = "Use Agent for focused sub-tasks, especially codebase exploration, specialist review, or isolated code changes."; const AGENT_TOOL_PROMPT_GUIDELINES = [ "Use Agent when a task is independent enough to hand off to one specialized sub-agent; do not use it for a single file read, quick grep, or another trivial local check.", "When you decide to use Agent, call it directly without narrating the delegation plan to the user first.", "In the prompt, state whether the sub-agent should do research or modify code, point it to the most relevant files or directories, and specify what result it should return.", "Once you delegate work, do not duplicate the same exploration or edits yourself; wait for the result or continue only with non-overlapping tasks.", "Use Agent with background=true only when the work can continue asynchronously while you do something else unrelated.", ] as const; interface AgentToolDetails { agentId: string; sessionId?: string; sessionPath?: string; status?: string; mode: "background" | "foreground"; } /** * Registers the main-agent `Agent` tool backed by the shared manager. */ export function registerAgentTool(pi: ExtensionAPI, registry: AgentRegistry, manager: AgentManager): void { pi.registerTool({ name: "Agent", label: "Agent", description: AGENT_TOOL_DESCRIPTION, promptSnippet: AGENT_TOOL_PROMPT_SNIPPET, promptGuidelines: [...AGENT_TOOL_PROMPT_GUIDELINES], parameters: agentToolSchema, async execute(_toolCallId, params, signal, _onUpdate, ctx) { registry.reload(); const agentDefinition = registry.getByName(params.agent); if (!agentDefinition) { throw new Error(`Unknown agent: ${params.agent}`); } const task: AgentTask = buildAgentTask({ agentName: agentDefinition.name, agentDefinition, prompt: params.prompt, cwd: ctx.cwd, mode: params.background ? "background" : "foreground", parentModel: ctx.model, parentThinking: pi.getThinkingLevel(), }); task.parentToolCallId = _toolCallId; task.parentSession = getParentSessionReference(ctx.sessionManager, ctx.cwd); task.parentSessionWriter = createParentSessionWriter(ctx.sessionManager, { emitLifecycle: createLifecycleWidgetEmitter(ctx.ui, ctx.mode), }); if (params.background) { const agentId = await manager.runInBackground(task, { signal }); return { content: [{ type: "text", text: `Started ${agentDefinition.name} as ${agentId}` }], details: { agentId, mode: "background" } satisfies AgentToolDetails, }; } const result = await manager.run(task, { signal }); const payload = result.responseText ?? result.error ?? "(no output)"; return { content: [{ type: "text", text: payload }], details: { agentId: result.agentId, sessionId: result.sessionId, sessionPath: result.sessionPath, status: result.status, mode: "foreground", } satisfies AgentToolDetails, }; }, renderCall(args, theme, _context) { const modeLabel = args.background ? theme.fg("muted", " [bg]") : ""; let text = theme.fg("toolTitle", theme.bold("Agent ")) + theme.fg("accent", args.agent) + modeLabel; text += `\n${theme.fg("dim", truncateSingleLine(args.prompt, 80))}`; return new Text(text, 0, 0); }, renderResult(result, { expanded, isPartial }, theme, _context) { if (isPartial) { return new Text(theme.fg("warning", "Delegating to sub-agent..."), 0, 0); } const details = result.details as AgentToolDetails | undefined; const payload = getToolTextContent(result); if (!details) { return new Text(payload, 0, 0); } if (details.mode === "background") { return new Text(theme.fg("success", `Started in background as ${details.agentId}`), 0, 0); } const statusText = details.status ?? "completed"; const statusColor = statusText === "completed" ? "success" : statusText === "failed" ? "error" : "warning"; let text = theme.fg(statusColor, statusText); text += theme.fg("muted", ` ยท ${details.agentId}`); if (!expanded) { text += `\n${theme.fg("muted", "(expand for full sub-agent output)")}`; return new Text(text, 0, 0); } if (details.sessionPath) { text += `\n${theme.fg("dim", `session: ${details.sessionPath}`)}`; } text += `\n\n${payload}`; return new Text(text, 0, 0); }, }); } function getToolTextContent(result: { content: Array<{ type: string; text?: string }> }): string { return ( result.content .filter((item): item is { type: "text"; text: string } => item.type === "text" && typeof item.text === "string") .map((item) => item.text) .join("\n") .trim() || "(no output)" ); } function truncateSingleLine(text: string, maxLength: number): string { const normalizedText = text.replaceAll(/\s+/g, " ").trim(); if (normalizedText.length <= maxLength) { return normalizedText; } return `${normalizedText.slice(0, maxLength - 3)}...`; }