@dreki-gg/pi-subagent 0.6.0 → 0.8.0

package/extensions/subagent/index.ts

@@ -12,61 +12,36 @@
  * Uses JSON mode to capture structured output from subagents.
  */
 
-import { spawn } from 'node:child_process';
-import { existsSync } from 'node:fs';
-import * as fs from 'node:fs';
-import { copyFile, mkdir, readdir, unlink } from 'node:fs/promises';
 import * as os from 'node:os';
-import { join } from 'node:path';
-import * as path from 'node:path';
-import type { AgentToolResult } from '@mariozechner/pi-agent-core';
-import type { Message } from '@mariozechner/pi-ai';
-import { StringEnum } from '@mariozechner/pi-ai';
+import type { AgentToolResult } from '@earendil-works/pi-agent-core';
+import type { Message } from '@earendil-works/pi-ai';
+import { StringEnum } from '@earendil-works/pi-ai';
 import {
   type ExtensionAPI,
   type ExtensionCommandContext,
+  AuthStorage,
   DefaultPackageManager,
   getAgentDir,
   getMarkdownTheme,
+  ModelRegistry,
   SettingsManager,
-  withFileMutationQueue,
-} from '@mariozechner/pi-coding-agent';
-import type { ResolvedPaths } from '@mariozechner/pi-coding-agent';
-import { Box, Container, Markdown, Spacer, Text } from '@mariozechner/pi-tui';
-import { Type } from '@sinclair/typebox';
+} from '@earendil-works/pi-coding-agent';
+import type { ResolvedPaths } from '@earendil-works/pi-coding-agent';
+import { Box, Container, Markdown, Spacer, Text, type AutocompleteItem } from '@earendil-works/pi-tui';
+import { Type } from 'typebox';
 import {
   type AgentConfig,
   type AgentScope,
   type AgentSource,
-  bundledAgentsDir,
   discoverAgents,
-  discoverAgentsWithPackages,
 } from './agents.js';
-import {
-  formatDelegateUsage,
-  formatRunAgentUsage,
-  parseDelegateArgs,
-  parseRunAgentArgs,
-} from './delegate-args.js';
-import { runAgent, runParallel } from './delegate-executor.js';
-import { buildSynthesisPrompt, extractRecentConversation } from './synthesis.js';
-import type {
-  AgentResult,
-  DelegateState,
-  PhaseDefinition,
-  PhaseResult,
-  UsageStats as DelegateUsageStats,
-} from './delegate-types.js';
-import {
-  aggregateUsage,
-  confirmPlan,
-  confirmSynthesis,
-  formatFullSummary,
-  formatPhaseHeader,
-  getFinalText,
-  pickWorkflow,
-} from './delegate-ui.js';
-import { getWorkflow, suggestWorkflow } from './workflows.js';
+import { formatRunAgentUsage, parseRunAgentArgs } from './run-agent-args.js';
+import { runAgent } from './agent-runner.js';
+import { extractRecentConversation } from './synthesis.js';
+import type { AgentResult } from './agent-runner-types.js';
+import { getFinalText } from './agent-result-utils.js';
+import { buildHandoffFromResult, renderHandoffForPrompt } from './handoffs.js';
+import { emptyUsage, spawnPiAgent } from './spawn-utils.js';
 
 const MAX_PARALLEL_TASKS = 8;
 const MAX_CONCURRENCY = 4;
@@ -218,7 +193,7 @@ interface SingleResult {
 interface SubagentDetails {
   mode: 'single' | 'parallel' | 'chain';
   agentScope: AgentScope;
-  projectAgentsDir: string | null;
+  projectPromptsDir: string | null;
   results: SingleResult[];
 }
 
@@ -272,39 +247,6 @@ async function mapWithConcurrencyLimit<TIn, TOut>(
   return results;
 }
 
-async function writePromptToTempFile(
-  agentName: string,
-  prompt: string,
-): Promise<{ dir: string; filePath: string }> {
-  const tmpDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), 'pi-subagent-'));
-  const safeName = agentName.replace(/[^\w.-]+/g, '_');
-  const filePath = path.join(tmpDir, `prompt-${safeName}.md`);
-  await withFileMutationQueue(filePath, async () => {
-    await fs.promises.writeFile(filePath, prompt, { encoding: 'utf-8', mode: 0o600 });
-  });
-  return { dir: tmpDir, filePath };
-}
-
-function getPiInvocation(args: string[]): { command: string; args: string[] } {
-  const currentScript = process.argv[1];
-
-  // Bun standalone binaries set argv[1] to a virtual FS path (e.g. /$bunfs/root/pi)
-  // that only resolves inside the running Bun process. Skip it — the binary IS the entry point.
-  const isBunVirtualPath = currentScript?.startsWith('/$bunfs/');
-
-  if (currentScript && !isBunVirtualPath && fs.existsSync(currentScript)) {
-    return { command: process.execPath, args: [currentScript, ...args] };
-  }
-
-  const execName = path.basename(process.execPath).toLowerCase();
-  const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
-  if (!isGenericRuntime) {
-    // Standalone binary (e.g. pi compiled with Bun) — invoke directly
-    return { command: process.execPath, args };
-  }
-
-  return { command: 'pi', args };
-}
 
 type OnUpdateCallback = (partial: AgentToolResult<SubagentDetails>) => void;
 
@@ -314,6 +256,8 @@ async function runSingleAgent(
   agentName: string,
   task: string,
   cwd: string | undefined,
+  modelOverride: string | undefined,
+  thinkingOverride: string | undefined,
   step: number | undefined,
   signal: AbortSignal | undefined,
   onUpdate: OnUpdateCallback | undefined,
@@ -330,26 +274,13 @@ async function runSingleAgent(
       exitCode: 1,
       messages: [],
       stderr: `Unknown agent: "${agentName}". Available agents: ${available}.`,
-      usage: {
-        input: 0,
-        output: 0,
-        cacheRead: 0,
-        cacheWrite: 0,
-        cost: 0,
-        contextTokens: 0,
-        turns: 0,
-      },
+      usage: emptyUsage(),
       step,
     };
   }
 
-  const args: string[] = ['--mode', 'json', '-p', '--no-session'];
-  if (agent.model) args.push('--model', agent.model);
-  if (agent.thinking) args.push('--thinking', agent.thinking);
-  if (agent.tools && agent.tools.length > 0) args.push('--tools', agent.tools.join(','));
-
-  let tmpPromptDir: string | null = null;
-  let tmpPromptPath: string | null = null;
+  const selectedModel = modelOverride ?? agent.model;
+  const selectedThinking = thinkingOverride ?? agent.thinking;
 
   const currentResult: SingleResult = {
     agent: agentName,
@@ -358,16 +289,8 @@ async function runSingleAgent(
     exitCode: 0,
     messages: [],
     stderr: '',
-    usage: {
-      input: 0,
-      output: 0,
-      cacheRead: 0,
-      cacheWrite: 0,
-      cost: 0,
-      contextTokens: 0,
-      turns: 0,
-    },
-    model: agent.model,
+    usage: emptyUsage(),
+    model: selectedModel,
     step,
   };
 
@@ -380,124 +303,58 @@ async function runSingleAgent(
     }
   };
 
-  try {
-    if (agent.systemPrompt.trim()) {
-      const tmp = await writePromptToTempFile(agent.name, agent.systemPrompt);
-      tmpPromptDir = tmp.dir;
-      tmpPromptPath = tmp.filePath;
-      args.push('--append-system-prompt', tmpPromptPath);
-    }
-
-    args.push(`Task: ${task}`);
-    let wasAborted = false;
-
-    const exitCode = await new Promise<number>((resolve) => {
-      const invocation = getPiInvocation(args);
-      const proc = spawn(invocation.command, invocation.args, {
-        cwd: cwd ?? defaultCwd,
-        shell: false,
-        stdio: ['ignore', 'pipe', 'pipe'],
-      });
-      let buffer = '';
-
-      const processLine = (line: string) => {
-        if (!line.trim()) return;
-        let event: any;
-        try {
-          event = JSON.parse(line);
-        } catch {
-          return;
-        }
-
-        if (event.type === 'message_end' && event.message) {
-          const msg = event.message as Message;
-          currentResult.messages.push(msg);
-
-          if (msg.role === 'assistant') {
-            currentResult.usage.turns++;
-            const usage = msg.usage;
-            if (usage) {
-              currentResult.usage.input += usage.input || 0;
-              currentResult.usage.output += usage.output || 0;
-              currentResult.usage.cacheRead += usage.cacheRead || 0;
-              currentResult.usage.cacheWrite += usage.cacheWrite || 0;
-              currentResult.usage.cost += usage.cost?.total || 0;
-              currentResult.usage.contextTokens = usage.totalTokens || 0;
-            }
-            if (!currentResult.model && msg.model) currentResult.model = msg.model;
-            if (msg.stopReason) currentResult.stopReason = msg.stopReason;
-            if (msg.errorMessage) currentResult.errorMessage = msg.errorMessage;
-          }
-          emitUpdate();
-        }
-
-        if (event.type === 'tool_result_end' && event.message) {
-          currentResult.messages.push(event.message as Message);
-          emitUpdate();
-        }
-      };
-
-      proc.stdout.on('data', (data) => {
-        buffer += data.toString();
-        const lines = buffer.split('\n');
-        buffer = lines.pop() || '';
-        for (const line of lines) processLine(line);
-      });
-
-      proc.stderr.on('data', (data) => {
-        currentResult.stderr += data.toString();
-      });
-
-      proc.on('close', (code) => {
-        if (buffer.trim()) processLine(buffer);
-        resolve(code ?? 0);
-      });
-
-      proc.on('error', () => {
-        resolve(1);
-      });
+  const spawnResult = await spawnPiAgent({
+    cwd: cwd ?? defaultCwd,
+    agentName: agent.name,
+    task,
+    systemPrompt: agent.systemPrompt,
+    model: selectedModel,
+    thinking: selectedThinking,
+    tools: agent.tools,
+    signal,
+    onMessage: (msg) => {
+      currentResult.messages = spawnResult.messages;
+      currentResult.usage = spawnResult.usage;
+      currentResult.model = spawnResult.model ?? currentResult.model;
+      currentResult.stopReason = spawnResult.stopReason;
+      currentResult.errorMessage = spawnResult.errorMessage;
+      emitUpdate();
+    },
+    onToolResult: () => {
+      currentResult.messages = spawnResult.messages;
+      emitUpdate();
+    },
+  });
 
-      if (signal) {
-        const killProc = () => {
-          wasAborted = true;
-          proc.kill('SIGTERM');
-          setTimeout(() => {
-            if (!proc.killed) proc.kill('SIGKILL');
-          }, 5000);
-        };
-        if (signal.aborted) killProc();
-        else signal.addEventListener('abort', killProc, { once: true });
-      }
-    });
+  currentResult.exitCode = spawnResult.exitCode;
+  currentResult.messages = spawnResult.messages;
+  currentResult.stderr = spawnResult.stderr;
+  currentResult.usage = spawnResult.usage;
+  currentResult.model = spawnResult.model ?? currentResult.model;
+  currentResult.stopReason = spawnResult.stopReason;
+  currentResult.errorMessage = spawnResult.errorMessage;
 
-    currentResult.exitCode = exitCode;
-    if (wasAborted) throw new Error('Subagent was aborted');
-    return currentResult;
-  } finally {
-    if (tmpPromptPath)
-      try {
-        fs.unlinkSync(tmpPromptPath);
-      } catch {
-        /* ignore */
-      }
-    if (tmpPromptDir)
-      try {
-        fs.rmdirSync(tmpPromptDir);
-      } catch {
-        /* ignore */
-      }
-  }
+  if (spawnResult.wasAborted) throw new Error('Subagent was aborted');
+  return currentResult;
 }
 
 const TaskItem = Type.Object({
   agent: Type.String({ description: 'Name of the agent to invoke' }),
   task: Type.String({ description: 'Task to delegate to the agent' }),
+  model: Type.Optional(Type.String({ description: 'Optional model override for this task' })),
+  thinking: Type.Optional(
+    Type.String({ description: 'Optional reasoning level override for this task' }),
+  ),
   cwd: Type.Optional(Type.String({ description: 'Working directory for the agent process' })),
 });
 
 const ChainItem = Type.Object({
   agent: Type.String({ description: 'Name of the agent to invoke' }),
   task: Type.String({ description: 'Task with optional {previous} placeholder for prior output' }),
+  model: Type.Optional(Type.String({ description: 'Optional model override for this step' })),
+  thinking: Type.Optional(
+    Type.String({ description: 'Optional reasoning level override for this step' }),
+  ),
   cwd: Type.Optional(Type.String({ description: 'Working directory for the agent process' })),
 });
 
@@ -512,11 +369,17 @@ const SubagentParams = Type.Object({
     Type.String({ description: 'Name of the agent to invoke (for single mode)' }),
   ),
   task: Type.Optional(Type.String({ description: 'Task to delegate (for single mode)' })),
+  model: Type.Optional(
+    Type.String({ description: 'Optional default model override for the run or all steps/tasks' }),
+  ),
+  thinking: Type.Optional(
+    Type.String({ description: 'Optional default reasoning level override for the run or all steps/tasks' }),
+  ),
   tasks: Type.Optional(
-    Type.Array(TaskItem, { description: 'Array of {agent, task} for parallel execution' }),
+    Type.Array(TaskItem, { description: 'Array of {agent, task, model?, thinking?} for parallel execution' }),
   ),
   chain: Type.Optional(
-    Type.Array(ChainItem, { description: 'Array of {agent, task} for sequential execution' }),
+    Type.Array(ChainItem, { description: 'Array of {agent, task, model?, thinking?} for sequential execution' }),
   ),
   agentScope: Type.Optional(AgentScopeSchema),
   confirmProjectAgents: Type.Optional(
@@ -530,42 +393,120 @@ const SubagentParams = Type.Object({
   ),
 });
 
-/**
- * Try to resolve package paths including the agents resource.
- * Returns null if the pi version does not support ResolvedPaths.agents.
- */
-async function tryResolvePackagePaths(cwd: string): Promise<ResolvedPaths | undefined> {
+async function resolvePackagePaths(cwd: string): Promise<ResolvedPaths | undefined> {
   try {
     const agentDir = getAgentDir();
     const settingsManager = SettingsManager.create(cwd, agentDir);
     const packageManager = new DefaultPackageManager({ cwd, agentDir, settingsManager });
-    const resolved = await packageManager.resolve();
-    // Check if this pi version has agents support
-    if ('agents' in resolved && Array.isArray((resolved as Record<string, unknown>).agents)) {
-      return resolved;
-    }
+    return await packageManager.resolve();
   } catch {
-    // Incompatible pi version or missing API
+    return undefined;
   }
-  return undefined;
 }
 
 export default function (pi: ExtensionAPI) {
+  let autocompleteCwd = process.cwd();
+
+  function parseArgumentText(argumentText: string): {
+    completedTokens: string[];
+    currentToken: string;
+  } {
+    const endsWithSpace = argumentText.length > 0 && /\s$/.test(argumentText);
+    const tokens = argumentText.trim().length > 0 ? argumentText.trim().split(/\s+/) : [];
+    return {
+      completedTokens: endsWithSpace ? tokens : tokens.slice(0, -1),
+      currentToken: endsWithSpace ? '' : (tokens.at(-1) ?? ''),
+    };
+  }
+
+  function replaceCurrentToken(argumentText: string, replacement: string): string {
+    const trimmedEnd = argumentText.replace(/\s+$/, '');
+    if (!trimmedEnd) return replacement;
+    if (/\s$/.test(argumentText)) return `${trimmedEnd} ${replacement}`;
+    const lastSpace = trimmedEnd.lastIndexOf(' ');
+    return lastSpace === -1
+      ? replacement
+      : `${trimmedEnd.slice(0, lastSpace + 1)}${replacement}`;
+  }
+
+  function buildArgumentCompletions(
+    argumentText: string,
+    items: Array<{ value: string; label?: string; description?: string }>,
+  ): AutocompleteItem[] | null {
+    const { currentToken } = parseArgumentText(argumentText);
+    const query = currentToken.toLowerCase();
+    const filtered = items.filter((item) => {
+      if (!query) return true;
+      return (
+        item.value.toLowerCase().startsWith(query) ||
+        item.label?.toLowerCase().includes(query) ||
+        item.description?.toLowerCase().includes(query)
+      );
+    });
+    if (filtered.length === 0) return null;
+    return filtered.map((item) => ({
+      value: replaceCurrentToken(argumentText, item.value),
+      label: item.label ?? item.value,
+      description: item.description,
+    }));
+  }
+
+  async function getAutocompleteAgents(scope: AgentScope): Promise<AgentConfig[]> {
+    const resolvedPaths = await resolvePackagePaths(autocompleteCwd);
+    return discoverAgents(autocompleteCwd, scope, resolvedPaths).agents;
+  }
+
+  async function getAutocompleteModels(scope: AgentScope): Promise<
+    Array<{ value: string; label?: string; description?: string }>
+  > {
+    const items = new Map<string, { value: string; label?: string; description?: string }>();
+
+    try {
+      const authStorage = AuthStorage.create();
+      const modelRegistry = ModelRegistry.create(authStorage);
+      const available = await modelRegistry.getAvailable();
+      for (const model of available) {
+        const value = `${model.provider}/${model.id}`;
+        items.set(value, {
+          value,
+          label: value,
+          description: 'Configured model with available credentials',
+        });
+      }
+    } catch {
+      // Fall back to models referenced by discovered agents.
+    }
+
+    const agents = await getAutocompleteAgents(scope);
+    for (const agent of agents) {
+      if (!agent.model) continue;
+      if (items.has(agent.model)) continue;
+      items.set(agent.model, {
+        value: agent.model,
+        label: agent.model,
+        description: `Default model on agent ${agent.name}`,
+      });
+    }
+
+    return [...items.values()];
+  }
+
   pi.registerTool({
     name: 'subagent',
     label: 'Subagent',
     description: [
       'Delegate tasks to specialized subagents with isolated context.',
       'Modes: single (agent + task), parallel (tasks array), chain (sequential with {previous} placeholder).',
-      'Default agent scope is "user" (from ~/.pi/agent/agents).',
-      'To enable project-local agents in .pi/agents, set agentScope: "both" (or "project").',
+      'Optional per-run or per-step model and reasoning-level overrides are supported via model/thinking.',
+      'Default agent scope is "user" (from ~/.pi/agent/prompts).',
+      'To enable project-local agents in .pi/prompts, set agentScope: "both" (or "project").',
     ].join(' '),
     parameters: SubagentParams,
 
     async execute(_toolCallId, params, signal, onUpdate, ctx) {
       const agentScope: AgentScope = params.agentScope ?? 'user';
-      const resolvedPaths = await tryResolvePackagePaths(ctx.cwd);
-      const discovery = discoverAgentsWithPackages(ctx.cwd, agentScope, resolvedPaths);
+      const resolvedPaths = await resolvePackagePaths(ctx.cwd);
+      const discovery = discoverAgents(ctx.cwd, agentScope, resolvedPaths);
       const agents = discovery.agents;
       const confirmProjectAgents = params.confirmProjectAgents ?? true;
 
@@ -579,7 +520,7 @@ export default function (pi: ExtensionAPI) {
         (results: SingleResult[]): SubagentDetails => ({
           mode,
           agentScope,
-          projectAgentsDir: discovery.projectAgentsDir,
+          projectPromptsDir: discovery.projectPromptsDir,
           results,
         });
 
@@ -612,7 +553,7 @@ export default function (pi: ExtensionAPI) {
 
         if (projectAgentsRequested.length > 0) {
           const names = projectAgentsRequested.map((a) => a.name).join(', ');
-          const dir = discovery.projectAgentsDir ?? '(unknown)';
+          const dir = discovery.projectPromptsDir ?? '(unknown)';
           const ok = await ctx.ui.confirm(
             'Run project-local agents?',
             `Agents: ${names}\nSource: ${dir}\n\nProject agents are repo-controlled. Only continue for trusted repositories.`,
@@ -648,12 +589,18 @@ export default function (pi: ExtensionAPI) {
               }
             : undefined;
 
+          if (ctx.hasUI) {
+            ctx.ui.setWorkingMessage(`Chain step ${i + 1}/${params.chain.length}: ${step.agent}`);
+          }
+
           const result = await runSingleAgent(
             ctx.cwd,
             agents,
             step.agent,
             taskWithContext,
             step.cwd,
+            step.model ?? params.model,
+            step.thinking ?? params.thinking,
             i + 1,
             signal,
             chainUpdate,
@@ -682,7 +629,18 @@ export default function (pi: ExtensionAPI) {
               isError: true,
             };
           }
-          previousOutput = getFinalOutput(result.messages);
+          const finalOutput = getFinalOutput(result.messages) || '(no output)';
+          previousOutput = renderHandoffForPrompt(
+            buildHandoffFromResult({
+              agent: step.agent,
+              step: i + 1,
+              task: step.task,
+              output: finalOutput,
+            }),
+          );
+        }
+        if (ctx.hasUI) {
+          ctx.ui.setWorkingMessage();
         }
         return {
           content: [
@@ -747,6 +705,10 @@ export default function (pi: ExtensionAPI) {
           }
         };
 
+        if (ctx.hasUI) {
+          ctx.ui.setWorkingMessage(`Running ${params.tasks.length} agents in parallel`);
+        }
+
         const results = await mapWithConcurrencyLimit(
           params.tasks,
           MAX_CONCURRENCY,
@@ -757,6 +719,8 @@ export default function (pi: ExtensionAPI) {
               t.agent,
               t.task,
               t.cwd,
+              t.model ?? params.model,
+              t.thinking ?? params.thinking,
               undefined,
               signal,
               // Per-task update callback
@@ -774,6 +738,10 @@ export default function (pi: ExtensionAPI) {
           },
         );
 
+        if (ctx.hasUI) {
+          ctx.ui.setWorkingMessage();
+        }
+
         const successCount = results.filter((r) => r.exitCode === 0).length;
         const summaries = results.map((r) => {
           const output = getFinalOutput(r.messages);
@@ -791,17 +759,27 @@ export default function (pi: ExtensionAPI) {
       }
 
       if (params.agent && params.task) {
+        if (ctx.hasUI) {
+          ctx.ui.setWorkingMessage(`Running agent: ${params.agent}`);
+        }
+
         const result = await runSingleAgent(
           ctx.cwd,
           agents,
           params.agent,
           params.task,
           params.cwd,
+          params.model,
+          params.thinking,
           undefined,
           signal,
           onUpdate,
           makeDetails('single'),
         );
+        if (ctx.hasUI) {
+          ctx.ui.setWorkingMessage();
+        }
+
         const isError =
           result.exitCode !== 0 || result.stopReason === 'error' || result.stopReason === 'aborted';
         if (isError) {
@@ -1164,31 +1142,18 @@ export default function (pi: ExtensionAPI) {
     },
   });
 
-  // ── Delegate commands ──────────────────────────────────────────────────────
-
-  function emptyDelegateUsage(): DelegateUsageStats {
-    return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 };
-  }
-
-  function combineParallelOutputs(results: AgentResult[]): string {
-    return results
-      .map((r) => {
-        const output = getFinalText(r);
-        return `## ${r.agent}\n${output}`;
-      })
-      .join('\n\n');
-  }
+  // ── Command helpers ───────────────────────────────────────────────────────
 
   async function confirmProjectAgentsIfNeeded(
     ctx: ExtensionCommandContext,
     agents: AgentConfig[],
-    projectAgentsDir: string | null,
+    projectPromptsDir: string | null,
     confirmProjectAgents: boolean,
   ): Promise<boolean> {
     if (!ctx.hasUI || !confirmProjectAgents || agents.length === 0) return true;
 
     const names = Array.from(new Set(agents.map((agent) => agent.name))).join(', ');
-    const dir = projectAgentsDir ?? '(unknown)';
+    const dir = projectPromptsDir ?? '(unknown)';
     return ctx.ui.confirm(
       'Run project-local agents?',
       `Agents: ${names}\nSource: ${dir}\n\nProject agents are repo-controlled. Only continue for trusted repositories.`,
@@ -1303,192 +1268,109 @@ export default function (pi: ExtensionAPI) {
     return box;
   });
 
-  async function executePhase(
-    cwd: string,
-    phase: PhaseDefinition,
-    synthesis: string,
-    previousOutput: string,
-    ctx: ExtensionCommandContext,
-    agentScope: AgentScope,
-    resolvedPaths?: ResolvedPaths,
-  ): Promise<{ results: AgentResult[]; output: string; aborted: boolean }> {
-    const task = phase.taskTemplate
-      .replace(/\{synthesis\}/g, synthesis)
-      .replace(/\{previous\}/g, previousOutput);
-
-    ctx.ui.notify(
-      `${formatPhaseHeader(phase.name, 'running')} ${phase.kind === 'parallel' ? `(${phase.agents.length} agents)` : phase.agents[0]}`,
-      'info',
-    );
-
-    let results: AgentResult[];
-    if (phase.kind === 'parallel') {
-      results = await runParallel(cwd, phase.agents, task, {
-        agentScope,
-        onUpdate: (_phaseName, _agentName, _result) => {
-          // Streaming updates happen per-message
-        },
-        phaseName: phase.name,
-        resolvedPaths,
-      });
-    } else {
-      const result = await runAgent(cwd, phase.agents[0], task, {
-        agentScope,
-        phaseName: phase.name,
-        resolvedPaths,
-      });
-      results = [result];
-    }
-
-    const hasError = results.some((r) => r.exitCode !== 0 || r.stopReason === 'error');
-    if (hasError) {
-      const failedAgents = results
-        .filter((r) => r.exitCode !== 0 || r.stopReason === 'error')
-        .map((r) => `${r.agent}: ${r.errorMessage || r.stderr || '(unknown error)'}`)
-        .join('\n');
-
-      ctx.ui.notify(`${formatPhaseHeader(phase.name, 'failed')}\n${failedAgents}`, 'error');
-      return { results, output: combineParallelOutputs(results), aborted: true };
-    }
-
-    const output = combineParallelOutputs(results);
-    ctx.ui.notify(formatPhaseHeader(phase.name, 'done'), 'info');
-
-    if (phase.requiresConfirmation && ctx.hasUI) {
-      const planText = output;
-      const proceed = await confirmPlan(ctx, planText);
-      if (!proceed) {
-        ctx.ui.notify('Delegate aborted by user after plan review.', 'warning');
-        return { results, output, aborted: true };
-      }
-    }
-
-    return { results, output, aborted: false };
-  }
-
-  async function listMdFiles(dir: string): Promise<Set<string>> {
-    if (!existsSync(dir)) return new Set();
-    try {
-      const entries = await readdir(dir);
-      return new Set(entries.filter((f) => f.endsWith('.md')).map((f) => f.replace(/\.md$/, '')));
-    } catch {
-      return new Set();
-    }
-  }
+  pi.on('session_start', async (_event, ctx) => {
+    autocompleteCwd = ctx.cwd;
+  });
 
-  pi.registerCommand('delegate-agents', {
-    description: 'List, customize, or reset delegate agents',
-    handler: async (args, ctx) => {
-      const userDir = join(getAgentDir(), 'agents');
-      const parts = args?.trim().split(/\s+/) || [];
-      const action = parts[0] || 'list';
-
-      if (action === 'list') {
-        const bundled = await listMdFiles(bundledAgentsDir);
-        const user = await listMdFiles(userDir);
-        const allNames = new Set([...bundled, ...user]);
-
-        const lines: string[] = ['## Delegate Agents\n'];
-        for (const name of [...allNames].sort()) {
-          const isBundled = bundled.has(name);
-          const isUser = user.has(name);
-          if (isUser && isBundled) {
-            lines.push(`- **${name}** — user override (bundled version available, \`/delegate-agents reset ${name}\` to restore)`);
-          } else if (isUser) {
-            lines.push(`- **${name}** — user-only`);
-          } else {
-            lines.push(`- **${name}** — bundled`);
+  pi.registerCommand('run-agent', {
+    description:
+      'Run a single agent directly — honors agent sessionStrategy metadata such as fork-at',
+    getArgumentCompletions: async (argumentText) => {
+      const { completedTokens } = parseArgumentText(argumentText);
+      let agentScope: AgentScope = 'user';
+      let expectsScopeValue = false;
+      let expectsThinkingValue = false;
+      let expectsModelValue = false;
+      let agentName: string | undefined;
+
+      for (const token of completedTokens) {
+        if (expectsScopeValue) {
+          if (token === 'user' || token === 'project' || token === 'both') {
+            agentScope = token;
           }
+          expectsScopeValue = false;
+          continue;
         }
-        pi.sendMessage({ customType: 'delegate-agents-list', content: lines.join('\n'), display: true });
-        return;
-      }
-
-      if (action === 'reset') {
-        const name = parts[1];
-        if (!name) {
-          ctx.ui.notify('Usage: /delegate-agents reset <name|--all>', 'warning');
-          return;
+        if (expectsThinkingValue) {
+          expectsThinkingValue = false;
+          continue;
         }
-
-        if (name === '--all') {
-          const user = await listMdFiles(userDir);
-          const bundled = await listMdFiles(bundledAgentsDir);
-          let count = 0;
-          for (const n of user) {
-            if (bundled.has(n)) {
-              await unlink(join(userDir, `${n}.md`));
-              count++;
-            }
-          }
-          ctx.ui.notify(count > 0 ? `Reset ${count} agent(s) to bundled versions.` : 'No user overrides to reset.', 'info');
-          return;
+        if (expectsModelValue) {
+          expectsModelValue = false;
+          continue;
         }
-
-        const userFile = join(userDir, `${name}.md`);
-        const bundledFile = join(bundledAgentsDir, `${name}.md`);
-
-        if (!existsSync(userFile)) {
-          ctx.ui.notify(`No user override for "${name}" — already using bundled version.`, 'info');
-          return;
+        if (token === '--scope') {
+          expectsScopeValue = true;
+          continue;
         }
-        if (!existsSync(bundledFile)) {
-          ctx.ui.notify(`"${name}" is user-only (no bundled version). Delete manually if needed.`, 'warning');
-          return;
+        if (token === '--thinking' || token === '--reasoning-level') {
+          expectsThinkingValue = true;
+          continue;
         }
-
-        await unlink(userFile);
-        ctx.ui.notify(`Reset "${name}" — now using bundled version.`, 'info');
-        return;
-      }
-
-      if (action === 'edit') {
-        const name = parts[1];
-        if (!name) {
-          ctx.ui.notify('Usage: /delegate-agents edit <name>', 'warning');
-          return;
+        if (token === '--model') {
+          expectsModelValue = true;
+          continue;
         }
+        if (token === '--yes-project-agents') continue;
+        if (token.startsWith('--')) continue;
+        if (!agentName) agentName = token;
+      }
 
-        const bundledFile = join(bundledAgentsDir, `${name}.md`);
-        const userFile = join(userDir, `${name}.md`);
-
-        if (existsSync(userFile)) {
-          ctx.ui.notify(`User override already exists: ${userFile}`, 'info');
-          return;
-        }
-        if (!existsSync(bundledFile)) {
-          ctx.ui.notify(`No bundled agent named "${name}".`, 'warning');
-          return;
-        }
+      if (expectsScopeValue) {
+        return buildArgumentCompletions(argumentText, [
+          { value: 'user', description: 'Only user/global agents from ~/.pi/agent/prompts' },
+          { value: 'project', description: 'Only project-local agents from .pi/prompts' },
+          { value: 'both', description: 'User/global agents plus project-local agents' },
+        ]);
+      }
 
-        await mkdir(userDir, { recursive: true });
-        await copyFile(bundledFile, userFile);
-        ctx.ui.notify(`Copied "${name}" to ${userFile} — edit it there to customize.`, 'info');
-        return;
+      if (expectsThinkingValue) {
+        return buildArgumentCompletions(argumentText, [
+          { value: 'off', description: 'Disable reasoning if the model supports it' },
+          { value: 'minimal', description: 'Minimal reasoning effort' },
+          { value: 'low', description: 'Lower reasoning effort' },
+          { value: 'medium', description: 'Balanced reasoning effort' },
+          { value: 'high', description: 'Higher reasoning effort' },
+          { value: 'xhigh', description: 'Maximum reasoning effort on supported models' },
+        ]);
       }
 
-      ctx.ui.notify('Unknown action. Usage: /delegate-agents [list|reset <name|--all>|edit <name>]', 'warning');
+      if (expectsModelValue) {
+        const models = await getAutocompleteModels(agentScope);
+        return buildArgumentCompletions(argumentText, models);
+      }
+      if (agentName) return null;
+
+      const agents = await getAutocompleteAgents(agentScope);
+      return buildArgumentCompletions(argumentText, [
+        { value: '--scope', description: 'Choose which agent directories are searched' },
+        { value: '--model', description: 'Override the agent model for this run' },
+        { value: '--thinking', description: 'Override the reasoning level for this run' },
+        { value: '--reasoning-level', description: 'Alias for --thinking' },
+        { value: '--yes-project-agents', description: 'Skip the project-agent trust confirmation prompt' },
+        ...agents.map((agent) => ({
+          value: agent.name,
+          label: agent.name,
+          description: `${agent.source} — ${agent.description}`,
+        })),
+      ]);
     },
-  });
-
-  pi.registerCommand('run-agent', {
-    description:
-      'Run a single agent directly — honors agent sessionStrategy metadata such as fork-at',
     handler: async (args, ctx) => {
+      autocompleteCwd = ctx.cwd;
       const parsed = parseRunAgentArgs(args);
       if (!parsed.ok) {
         ctx.ui.notify(parsed.error, 'warning');
         return;
       }
 
-      const { agentName, explicitTask, agentScope, confirmProjectAgents } = parsed.options;
+      const { agentName, explicitTask, agentScope, confirmProjectAgents, model, thinking } = parsed.options;
       if (!agentName) {
         ctx.ui.notify(formatRunAgentUsage(), 'warning');
         return;
       }
 
-      const resolvedPaths = await tryResolvePackagePaths(ctx.cwd);
-      const discovery = discoverAgentsWithPackages(ctx.cwd, agentScope, resolvedPaths);
+      const resolvedPaths = await resolvePackagePaths(ctx.cwd);
+      const discovery = discoverAgents(ctx.cwd, agentScope, resolvedPaths);
       const agent = discovery.agents.find((candidate) => candidate.name === agentName);
 
       if (!agent) {
@@ -1501,7 +1383,7 @@ export default function (pi: ExtensionAPI) {
       const approved = await confirmProjectAgentsIfNeeded(
         ctx,
         requestedProjectAgents,
-        discovery.projectAgentsDir,
+        discovery.projectPromptsDir,
         confirmProjectAgents,
       );
       if (!approved) {
@@ -1516,7 +1398,65 @@ export default function (pi: ExtensionAPI) {
         return;
       }
 
-      let forked = false;
+      const runAndReport = async (
+        commandCwd: string,
+        commandUi: typeof ctx.ui,
+        sendSummary: (message: {
+          customType: 'run-agent-summary';
+          content: string;
+          display: true;
+          details: RunAgentSummaryDetails;
+        }) => Promise<void> | void,
+        forked: boolean,
+      ) => {
+        const selectedModel = model ?? agent.model;
+        const selectedThinking = thinking ?? agent.thinking;
+
+        commandUi.notify(
+          `Running agent: ${agent.name} [${agent.source}${forked ? ', forked' : ', inline'}${selectedModel ? `, model=${selectedModel}` : ''}${selectedThinking ? `, thinking=${selectedThinking}` : ''}]`,
+          'info',
+        );
+
+        const result = await runAgent(commandCwd, agent.name, task, {
+          agentScope,
+          resolvedPaths,
+          model,
+          thinking,
+        });
+        const failed = result.exitCode !== 0 || result.stopReason === 'error';
+        const output = getFinalText(result).trim();
+
+        commandUi.notify(
+          failed ? `Agent "${agent.name}" failed.` : `Agent "${agent.name}" complete.`,
+          failed ? 'error' : 'info',
+        );
+
+        await sendSummary({
+          customType: 'run-agent-summary',
+          content: formatRunAgentSummary({
+            agent,
+            agentScope,
+            task,
+            result,
+            forked,
+          }),
+          display: true,
+          details: {
+            agent: agent.name,
+            agentSource: agent.source,
+            agentScope,
+            sessionStrategy: agent.sessionStrategy ?? 'inline',
+            forked,
+            exitCode: result.exitCode,
+            stopReason: result.stopReason,
+            errorMessage: result.errorMessage,
+            usage: result.usage,
+            task,
+            output,
+          } satisfies RunAgentSummaryDetails,
+        });
+      };
+
       if (agent.sessionStrategy === 'fork-at') {
         const leafId = ctx.sessionManager.getLeafId?.();
         if (!leafId) {
@@ -1525,188 +1465,26 @@ export default function (pi: ExtensionAPI) {
             'warning',
           );
         } else {
-          const forkCompat = ctx.fork as unknown as (
-            entryId: string,
-            options?: { position: 'at' },
-          ) => Promise<{ cancelled: boolean }>;
-          const forkResult = await forkCompat(leafId, { position: 'at' });
+          const forkResult = await ctx.fork(leafId, {
+            position: 'at',
+            withSession: async (replacementCtx) => {
+              await runAndReport(
+                replacementCtx.cwd,
+                replacementCtx.ui,
+                (message) => replacementCtx.sendMessage(message),
+                true,
+              );
+            },
+          });
           if (forkResult.cancelled) {
             ctx.ui.notify(`Run cancelled — unable to fork session for "${agent.name}".`, 'info');
-            return;
           }
-          forked = true;
-        }
-      }
-
-      ctx.ui.notify(
-        `Running agent: ${agent.name} [${agent.source}${forked ? ', forked' : ', inline'}]`,
-        'info',
-      );
-
-      const result = await runAgent(ctx.cwd, agent.name, task, { agentScope, resolvedPaths });
-      const failed = result.exitCode !== 0 || result.stopReason === 'error';
-      const output = getFinalText(result).trim();
-
-      ctx.ui.notify(
-        failed ? `Agent "${agent.name}" failed.` : `Agent "${agent.name}" complete.`,
-        failed ? 'error' : 'info',
-      );
-
-      pi.sendMessage({
-        customType: 'run-agent-summary',
-        content: formatRunAgentSummary({
-          agent,
-          agentScope,
-          task,
-          result,
-          forked,
-        }),
-        display: true,
-        details: {
-          agent: agent.name,
-          agentSource: agent.source,
-          agentScope,
-          sessionStrategy: agent.sessionStrategy ?? 'inline',
-          forked,
-          exitCode: result.exitCode,
-          stopReason: result.stopReason,
-          errorMessage: result.errorMessage,
-          usage: result.usage,
-          task,
-          output,
-        } satisfies RunAgentSummaryDetails,
-      });
-    },
-  });
-
-  pi.registerCommand('delegate', {
-    description:
-      'Orchestrate subagent workflows — supports --scope, --workflow, and project-local agents',
-    handler: async (args, ctx) => {
-      const parsed = parseDelegateArgs(args);
-      if (!parsed.ok) {
-        ctx.ui.notify(parsed.error, 'warning');
-        return;
-      }
-
-      const { explicitTask, agentScope, workflowId, confirmProjectAgents } = parsed.options;
-      const resolvedPaths = await tryResolvePackagePaths(ctx.cwd);
-
-      // Step 1: Synthesize
-      const conversation = extractRecentConversation(ctx);
-      if (!conversation.trim() && !explicitTask) {
-        ctx.ui.notify(`No conversation context and no task specified.\n\n${formatDelegateUsage()}`, 'warning');
-        return;
-      }
-
-      const prompt = buildSynthesisPrompt(conversation, explicitTask);
-      ctx.ui.notify('⏳ Synthesizing task from conversation...', 'info');
-
-      const synthResult = await runAgent(ctx.cwd, 'planner', prompt, { agentScope, resolvedPaths });
-      const synthesis = getFinalText(synthResult);
-
-      if (!synthesis.trim()) {
-        ctx.ui.notify('Failed to synthesize task from conversation.', 'error');
-        return;
-      }
-
-      // Step 2: Confirm synthesis
-      if (ctx.hasUI) {
-        const synthOk = await confirmSynthesis(ctx, synthesis);
-        if (!synthOk) {
-          ctx.ui.notify('Delegate cancelled.', 'info');
-          return;
-        }
-      }
-
-      // Step 3: Pick workflow
-      const workflow = workflowId
-        ? getWorkflow(workflowId)
-        : ctx.hasUI
-          ? await pickWorkflow(ctx, synthesis)
-          : getWorkflow(suggestWorkflow(synthesis));
-      if (!workflow) {
-        ctx.ui.notify('Delegate cancelled — no workflow selected.', 'info');
-        return;
-      }
-
-      if (agentScope === 'project' || agentScope === 'both') {
-        const discovery = discoverAgentsWithPackages(ctx.cwd, agentScope, resolvedPaths);
-        const requestedProjectAgents = workflow.phases
-          .flatMap((phase) => phase.agents)
-          .map((name) => discovery.agents.find((agent) => agent.name === name))
-          .filter((agent): agent is AgentConfig => agent?.source === 'project');
-
-        const ok = await confirmProjectAgentsIfNeeded(
-          ctx,
-          requestedProjectAgents,
-          discovery.projectAgentsDir,
-          confirmProjectAgents,
-        );
-        if (!ok) {
-          ctx.ui.notify('Delegate cancelled — project-local agents not approved.', 'info');
           return;
         }
       }
 
-      // Step 4: Execute phases
-      const state: DelegateState = {
-        synthesis,
-        workflow,
-        phases: [],
-        totalUsage: emptyDelegateUsage(),
-      };
-
-      let previousOutput = '';
-
-      ctx.ui.notify(`Starting workflow: ${workflow.label} [${agentScope}]`, 'info');
-
-      for (const phaseDef of workflow.phases) {
-        const phaseResult = await executePhase(
-          ctx.cwd,
-          phaseDef,
-          synthesis,
-          previousOutput,
-          ctx,
-          agentScope,
-          resolvedPaths,
-        );
-
-        const phase: PhaseResult = {
-          name: phaseDef.name,
-          kind: phaseDef.kind,
-          agents: phaseResult.results,
-        };
-
-        state.phases.push(phase);
-        const phaseUsage = aggregateUsage(phaseResult.results);
-        state.totalUsage.input += phaseUsage.input;
-        state.totalUsage.output += phaseUsage.output;
-        state.totalUsage.cacheRead += phaseUsage.cacheRead;
-        state.totalUsage.cacheWrite += phaseUsage.cacheWrite;
-        state.totalUsage.cost += phaseUsage.cost;
-        state.totalUsage.turns += phaseUsage.turns;
-
-        if (phaseResult.aborted) break;
-
-        previousOutput = phaseResult.output;
-      }
-
-      // Step 5: Show full summary
-      const summary = formatFullSummary(state);
-      ctx.ui.notify(`Delegate complete — ${workflow.label}`, 'info');
-
-      pi.sendMessage({
-        customType: 'delegate-summary',
-        content: summary,
-        display: true,
-        details: {
-          workflow: workflow.id,
-          agentScope,
-          phaseCount: state.phases.length,
-          totalUsage: state.totalUsage,
-        },
-      });
+      await runAndReport(ctx.cwd, ctx.ui, (message) => pi.sendMessage(message), false);
     },
   });
+
 }