pi-recurse 0.1.0

package/index.ts

@@ -0,0 +1,546 @@
+/**
+ * Pi Recurse Extension
+ *
+ * Enables programmatic recursive subagent spawning with guardrails.
+ * The key capability: LLM makes ONE tool call, extension code handles
+ * parallel spawning and result aggregation without autoregressive steps.
+ *
+ * Usage:
+ *   recurse({ mode: "single", prompt: "Analyze src/auth.ts" })
+ *   recurse({ mode: "parallel", tasks: [...], concurrency: 4 })
+ *   recurse({ mode: "chain", chain: [...] })
+ */
+
+import type { ExtensionAPI, ExtensionContext } from '@earendil-works/pi-coding-agent';
+import { Type } from '@sinclair/typebox';
+import { Text } from '@earendil-works/pi-tui';
+import type {
+  RecurseParams,
+  RecurseResult,
+  SubagentResult,
+  SubagentProgress,
+  RecurseSingleParams,
+  RecurseParallelParams,
+  RecurseChainParams,
+} from './types.js';
+import {
+  getCurrentDepth,
+  getMaxDepth,
+  getTraceId,
+  checkDepthGuard,
+  checkCallGuard,
+  checkTimeoutGuard,
+  checkBudgetGuard,
+  spawnSubagent,
+  runParallel,
+  getRecursiveSystemPrompt,
+  loadAccumulatedCost,
+  saveAccumulatedCost,
+  DEFAULTS,
+} from './lib.js';
+import {
+  renderParallelStatus,
+  renderSubagentStatus,
+  formatDuration,
+  formatTokens,
+  renderRecurseTree,
+  buildRecurseTree,
+} from './formatters.js';
+import { formatAgentLabel } from './names.js';
+
+export default function piRecurseExtension(pi: ExtensionAPI) {
+  const currentDepth = getCurrentDepth();
+  const maxDepth = getMaxDepth();
+  const isDeep = currentDepth > 0;
+
+  // At deep depths, disable the recurse tool entirely
+  const disableToolAt = parseInt(
+    process.env.RLM_DISABLE_TOOL_AT || String(DEFAULTS.DISABLE_TOOL_AT_DEPTH),
+    10
+  );
+  const toolEnabled = currentDepth < disableToolAt;
+
+  pi.on('before_agent_start', async (event) => {
+    const modifiedPrompt = getRecursiveSystemPrompt(event.systemPrompt, currentDepth);
+    return { systemPrompt: modifiedPrompt };
+  });
+
+  pi.on('session_start', async (_event, ctx) => {
+    if (ctx.hasUI && toolEnabled) {
+      const budget = checkBudgetGuard();
+      const statusText =
+        budget.remaining !== Infinity
+          ? `∞ depth ${currentDepth}/${maxDepth} · $${budget.remaining.toFixed(2)}`
+          : `∞ depth ${currentDepth}/${maxDepth}`;
+      ctx.ui.setStatus('recurse', statusText);
+    }
+  });
+
+  pi.registerTool({
+    name: 'recurse',
+    label: 'Recurse',
+    description: toolEnabled
+      ? `Spawn subagents programmatically. Mode "single" for one task, "parallel" for concurrent batch processing, "chain" for sequential dependency chains. Returns aggregated results.`
+      : `[DISABLED at depth ${currentDepth}] Work directly instead of recursing.`,
+    promptSnippet: 'Delegate work to subagents in single/parallel/chain mode',
+    promptGuidelines: [
+      "Use recurse({ mode: 'single', prompt: '...' }) for one-off delegation.",
+      "Use recurse({ mode: 'parallel', tasks: [...] }) for independent batch work.",
+      "Use recurse({ mode: 'chain', chain: [...] }) when each step depends on the previous.",
+      'Subagents return compact results; aggregate and synthesize in parent.',
+      'Check result.stats before proceeding with expensive operations.',
+    ],
+
+    parameters: Type.Object({
+      mode: StringEnum(['single', 'parallel', 'chain'] as const, {
+        description: 'Execution mode: single task, parallel batch, or sequential chain',
+      }),
+
+      // Single mode params
+      prompt: Type.Optional(Type.String({ description: 'Prompt for single mode' })),
+      context: Type.Optional(Type.String({ description: 'Context data to pipe to subagent' })),
+      fork: Type.Optional(Type.Boolean({ description: 'Fork session history (default: false)' })),
+
+      // Parallel mode params
+      tasks: Type.Optional(
+        Type.Array(
+          Type.Object({
+            id: Type.String({ description: 'Task identifier' }),
+            prompt: Type.String({ description: 'Subagent prompt' }),
+            context: Type.Optional(Type.String({ description: 'Task-specific context' })),
+          }),
+          { description: 'Tasks for parallel execution' }
+        )
+      ),
+      concurrency: Type.Optional(
+        Type.Number({
+          description: 'Max concurrent subagents',
+          default: DEFAULTS.CONCURRENCY,
+        })
+      ),
+      timeoutPerTask: Type.Optional(
+        Type.Number({
+          description: 'Timeout per task in seconds',
+        })
+      ),
+
+      // Chain mode params
+      chain: Type.Optional(
+        Type.Array(
+          Type.Object({
+            id: Type.String({ description: 'Step identifier' }),
+            prompt: Type.String({ description: 'Step prompt (use {previous} for prior output)' }),
+          }),
+          { description: 'Sequential chain steps' }
+        )
+      ),
+    }),
+
+    async execute(toolCallId, rawParams, signal, onUpdate, ctx) {
+      const params = rawParams as RecurseParams;
+      const startTime = Date.now();
+
+      const depthCheck = checkDepthGuard();
+      if (!depthCheck.allowed) {
+        return {
+          content: [{ type: 'text', text: `Blocked: ${depthCheck.reason}` }],
+          details: { blocked: true, reason: depthCheck.reason },
+          isError: true,
+        };
+      }
+
+      const callCheck = checkCallGuard();
+      if (!callCheck.allowed) {
+        return {
+          content: [{ type: 'text', text: `Blocked: ${callCheck.reason}` }],
+          details: { blocked: true, reason: callCheck.reason },
+          isError: true,
+        };
+      }
+
+      const timeoutCheck = checkTimeoutGuard();
+      if (!timeoutCheck.allowed) {
+        return {
+          content: [{ type: 'text', text: `Blocked: ${timeoutCheck.reason}` }],
+          details: { blocked: true, reason: timeoutCheck.reason },
+          isError: true,
+        };
+      }
+
+      const budgetCheck = checkBudgetGuard();
+      if (!budgetCheck.allowed) {
+        return {
+          content: [
+            {
+              type: 'text',
+              text: `Blocked: Budget exceeded ($${loadAccumulatedCost().toFixed(4)} spent)`,
+            },
+          ],
+          details: { blocked: true, reason: 'budget' },
+          isError: true,
+        };
+      }
+
+      // Execute based on mode
+      let results: SubagentResult[];
+
+      switch (params.mode) {
+        case 'single': {
+          if (!params.prompt) {
+            return {
+              content: [{ type: 'text', text: "Missing required 'prompt' for single mode" }],
+              details: { error: 'missing prompt' },
+              isError: true,
+            };
+          }
+
+          let currentData: { output: string; progress?: SubagentProgress } = {
+            output: '',
+            progress: {
+              status: 'running',
+              recentOutput: [],
+              recentTools: [],
+              toolCount: 0,
+              tokens: 0,
+              durationMs: 0,
+            },
+          };
+          const displayName = formatAgentLabel('subagent', true);
+
+          const result = await spawnSubagent({
+            prompt: params.prompt,
+            context: params.context,
+            fork: params.fork,
+            onUpdate: onUpdate
+              ? (data) => {
+                  currentData = data as { output: string; progress?: SubagentProgress };
+                  const lines = renderSubagentStatus('subagent', currentData, 100, true);
+                  onUpdate({
+                    content: [{ type: 'text', text: lines.join('\n') }],
+                    details: { progress: data.progress },
+                  });
+                }
+              : undefined,
+          });
+          results = [result];
+          break;
+        }
+
+        case 'parallel': {
+          if (!params.tasks || params.tasks.length === 0) {
+            return {
+              content: [{ type: 'text', text: "Missing required 'tasks' array for parallel mode" }],
+              details: { error: 'missing tasks' },
+              isError: true,
+            };
+          }
+
+          const concurrency = params.concurrency || DEFAULTS.CONCURRENCY;
+          onUpdate?.({
+            content: [
+              {
+                type: 'text',
+                text: `Spawning ${params.tasks.length} subagents (max ${concurrency} concurrent)...`,
+              },
+            ],
+            details: {},
+          });
+
+          // Track progress for each task
+          const taskProgress = new Map<string, { output: string; progress?: SubagentProgress }>();
+
+          // Programmatic parallel spawning — NO LLM involvement between spawns
+          results = await runParallel(
+            params.tasks,
+            async (task) => {
+              taskProgress.set(task.id, {
+                output: '',
+                progress: {
+                  status: 'running',
+                  recentOutput: [],
+                  recentTools: [],
+                  toolCount: 0,
+                  tokens: 0,
+                  durationMs: 0,
+                },
+              });
+
+              const result = await spawnSubagent({
+                prompt: task.prompt,
+                context: task.context,
+                timeout: params.timeoutPerTask,
+                onUpdate: onUpdate
+                  ? (data) => {
+                      taskProgress.set(
+                        task.id,
+                        data as { output: string; progress?: SubagentProgress }
+                      );
+                      // Render full multi-line status like pi-subagents
+                      const statusText = renderParallelStatus(taskProgress);
+                      onUpdate({
+                        content: [{ type: 'text', text: statusText }],
+                        details: { taskProgress: Object.fromEntries(taskProgress) },
+                      });
+                    }
+                  : undefined,
+              });
+
+              // IMPORTANT: Update with final result so status shows completed/failed
+              if (onUpdate) {
+                taskProgress.set(task.id, {
+                  output: result.output,
+                  progress: {
+                    status: result.success ? 'completed' : 'failed',
+                    recentOutput: result.progress?.recentOutput || [],
+                    recentTools: result.progress?.recentTools || [],
+                    toolCount: result.progress?.toolCount || 0,
+                    tokens: (result.usage?.input || 0) + (result.usage?.output || 0),
+                    durationMs: result.durationMs,
+                  },
+                });
+                const statusText = renderParallelStatus(taskProgress);
+                onUpdate({
+                  content: [{ type: 'text', text: statusText }],
+                  details: { taskProgress: Object.fromEntries(taskProgress) },
+                });
+              }
+
+              return result;
+            },
+            concurrency
+          );
+          break;
+        }
+
+        case 'chain': {
+          if (!params.chain || params.chain.length === 0) {
+            return {
+              content: [{ type: 'text', text: "Missing required 'chain' array for chain mode" }],
+              details: { error: 'missing chain' },
+              isError: true,
+            };
+          }
+
+          results = [];
+          let previousOutput = '';
+
+          for (const step of params.chain) {
+            // Check cancellation
+            if (signal?.aborted) {
+              results.push({
+                id: step.id,
+                success: false,
+                output: '',
+                error: 'Cancelled',
+                durationMs: 0,
+              });
+              break;
+            }
+
+            // Substitute {previous} placeholder
+            const prompt = step.prompt.replace(/\{previous\}/g, previousOutput);
+
+            let stepData: { output: string; progress?: SubagentProgress } = {
+              output: '',
+              progress: {
+                status: 'running',
+                recentOutput: [],
+                recentTools: [],
+                toolCount: 0,
+                tokens: 0,
+                durationMs: 0,
+              },
+            };
+            const stepLabel = formatAgentLabel(step.id, true);
+
+            const result = await spawnSubagent({
+              prompt,
+              onUpdate: onUpdate
+                ? (data) => {
+                    stepData = data as { output: string; progress?: SubagentProgress };
+                    const lines = renderSubagentStatus(step.id, stepData, 100, true);
+                    // Show step in context of chain
+                    const header = `Step ${results.length + 1}/${params.chain!.length}: ${stepLabel}`;
+                    onUpdate({
+                      content: [{ type: 'text', text: `${header}\n${lines.join('\n')}` }],
+                      details: {
+                        stepId: step.id,
+                        stepIndex: results.length + 1,
+                        totalSteps: params.chain!.length,
+                        stepProgress: data.progress,
+                      },
+                    });
+                  }
+                : undefined,
+            });
+            results.push(result);
+
+            if (!result.success) {
+              break; // Stop chain on failure
+            }
+
+            previousOutput = result.output;
+          }
+          break;
+        }
+
+        default:
+          return {
+            content: [{ type: 'text', text: `Unknown mode: ${(params as any).mode}` }],
+            details: { error: 'unknown mode' },
+            isError: true,
+          };
+      }
+
+      // Calculate aggregated stats
+      const succeeded = results.filter((r) => r.success).length;
+      const totalCost = results.reduce((sum, r) => sum + (r.usage?.cost || 0), 0);
+
+      // Update accumulated cost
+      if (totalCost > 0) {
+        const currentCost = loadAccumulatedCost();
+        saveAccumulatedCost(currentCost + totalCost);
+      }
+
+      const totalDuration = Date.now() - startTime;
+      const invocationId = Math.random().toString(36).substring(2, 10);
+
+      const result: RecurseResult = {
+        results,
+        stats: {
+          total: results.length,
+          succeeded,
+          failed: results.length - succeeded,
+          totalDurationMs: totalDuration,
+          totalCost,
+        },
+        depth: currentDepth,
+        mode: params.mode,
+        invocationId,
+      };
+
+      // Format output
+      const lines: string[] = [
+        `## Recurse Results (depth ${currentDepth})`,
+        '',
+        `**Stats:** ${result.stats.succeeded}/${result.stats.total} succeeded · ${(result.stats.totalDurationMs / 1000).toFixed(1)}s${totalCost > 0 ? ` · $${totalCost.toFixed(4)}` : ''}`,
+        '',
+      ];
+
+      for (const r of results) {
+        const icon = r.success ? '✓' : '✗';
+        const duration = (r.durationMs / 1000).toFixed(1);
+        const stopReason = r.stopReason && r.stopReason !== 'completed' ? ` [${r.stopReason}]` : '';
+        lines.push(`### ${icon} ${r.id} (${duration}s)${stopReason}`);
+        if (r.error) {
+          lines.push(`**Error:** ${r.error}`);
+        }
+        lines.push(r.output || '*(no output)*');
+        lines.push('');
+      }
+
+      return {
+        content: [{ type: 'text', text: lines.join('\n') }],
+        details: result,
+      };
+    },
+
+    // Custom rendering
+    renderCall(args, theme) {
+      const mode = (args.mode as string) || 'single';
+      const modeLabel = theme.fg('accent', mode);
+
+      let detail = '';
+      if (mode === 'single' && args.prompt) {
+        const prompt = String(args.prompt).slice(0, 40);
+        detail = prompt.length < String(args.prompt).length ? `${prompt}...` : prompt;
+      } else if (mode === 'parallel' && args.tasks) {
+        detail = `${(args.tasks as any[]).length} tasks`;
+      } else if (mode === 'chain' && args.chain) {
+        detail = `${(args.chain as any[]).length} steps`;
+      }
+
+      const text =
+        theme.fg('toolTitle', 'recurse ') +
+        modeLabel +
+        (detail ? theme.fg('dim', ` "${detail}"`) : '');
+
+      return new Text(text, 0, 0);
+    },
+
+    renderResult(result, { expanded }, theme) {
+      const data = result.details as RecurseResult | undefined;
+      if (!data) {
+        return new Text(theme.fg('dim', 'No result data'), 0, 0);
+      }
+
+      const { stats, depth, mode } = data;
+      const icon = stats.failed === 0 ? theme.fg('success', '✓') : theme.fg('warning', '⚠');
+      const hasChildren = data.results.some((r) => r.children);
+
+      let text = `${icon} ${stats.succeeded}/${stats.total} at depth ${depth}${mode ? ` · ${mode}` : ''}`;
+
+      if (stats.totalCost && stats.totalCost > 0) {
+        text += theme.fg('dim', ` · $${stats.totalCost.toFixed(4)}`);
+      }
+
+      if (hasChildren) {
+        text += theme.fg('accent', ' [has children]');
+      }
+
+      if (expanded) {
+        if (hasChildren && mode) {
+          // Render tree view
+          const tree = buildRecurseTree(data, mode);
+          const treeLines = renderRecurseTree(tree, 100);
+          text += '\n' + treeLines.join('\n');
+        } else {
+          // Simple flat view
+          text += '\n';
+          for (const r of data.results) {
+            const status = r.success ? theme.fg('success', '✓') : theme.fg('error', '✗');
+            text += `  ${status} ${r.id}`;
+            if (r.children) {
+              text += theme.fg('accent', ` → ${r.children.stats.total} children`);
+            }
+            text += '\n';
+          }
+        }
+      }
+
+      return new Text(text, 0, 0);
+    },
+  });
+
+  pi.registerCommand('recurse-status', {
+    description: 'Show current recursion status and guardrails',
+    handler: async (_args, ctx) => {
+      const depth = getCurrentDepth();
+      const max = getMaxDepth();
+      const budget = checkBudgetGuard();
+
+      const lines = [
+        '## Recurse Status',
+        '',
+        `**Current depth:** ${depth} / ${max}`,
+        `**Tool enabled:** ${toolEnabled}`,
+        `**Budget:** ${budget.remaining === Infinity ? 'unlimited' : `$${budget.remaining.toFixed(2)} remaining`}`,
+        `**Trace ID:** ${getTraceId()}`,
+      ];
+
+      ctx.ui.notify(lines.join('\n'), 'info');
+    },
+  });
+
+  function StringEnum<T extends readonly string[]>(
+    values: T,
+    options?: { description?: string; default?: T[number] }
+  ) {
+    return Type.Unsafe<T[number]>({
+      type: 'string',
+      enum: [...values],
+      ...(options?.description && { description: options.description }),
+      ...(options?.default && { default: options.default }),
+    });
+  }
+}