pi-agent-extensions 0.1.0

package/extensions/handoff/extraction.ts

@@ -0,0 +1,153 @@
+import type { ExtractionOutput, HandoffConfig, ParseResult } from "./types.js";
+import { parseExtractionResponse, normalizeExtraction } from "./parser.js";
+
+/**
+ * System prompt for the extraction LLM call.
+ * Instructs the model to extract structured context from the conversation.
+ */
+export const EXTRACTION_SYSTEM_PROMPT = `You are a context extraction assistant. Your task is to analyze a conversation and extract the most relevant context for continuing work in a new thread.
+
+Given the conversation history and the user's goal for the next thread, extract:
+1. **relevantFiles**: Files that were ACTUALLY MENTIONED in the conversation that are relevant to the goal. Include a brief reason for each.
+2. **relevantCommands**: Commands that were run and may need to be run again
+3. **relevantInformation**: Key context for accomplishing the goal
+4. **decisions**: Important decisions made during the conversation
+5. **openQuestions**: Unresolved questions, risks, or blockers
+
+## Output Format
+
+Respond with ONLY a valid JSON object in this exact format:
+
+\`\`\`json
+{
+  "relevantFiles": [
+    { "path": "path/to/file.ts", "reason": "Brief reason why this file matters" }
+  ],
+  "relevantCommands": ["npm test", "git status"],
+  "relevantInformation": [
+    "Key fact or context point",
+    "Another important detail"
+  ],
+  "decisions": [
+    "Decision that was made and why"
+  ],
+  "openQuestions": [
+    "Question that remains unanswered"
+  ]
+}
+\`\`\`
+
+## What to Extract
+
+**relevantInformation** - Focus on:
+- Project conventions learned (e.g., "Use TypeBox, not Zod for schemas")
+- Runtime behaviors discovered (e.g., "Extensions hot-reload with /reload")
+- Gotchas that could trip up the next agent (e.g., "Must use .js extension in imports")
+- Technical constraints or requirements
+- Key findings from exploration
+
+**relevantFiles** - Only include files that:
+- Were EXPLICITLY MENTIONED in the conversation (by path or filename)
+- Are directly related to accomplishing the goal
+- Contain patterns to follow or will need to be modified
+
+## What NOT to Extract
+
+- Completed tasks or work history ("We implemented X, then Y, then Z")
+- Obvious actions the agent will do anyway (running tests, building, linting)
+- Generic observations that don't help the specific goal
+- Files that were NOT mentioned in the conversation (do not invent paths)
+
+## Guidelines
+
+- Be GOAL-FOCUSED: extract what helps accomplish the user's stated goal
+- Be FUTURE-ORIENTED: what does the NEXT agent need to know?
+- Be CONCISE: one line per entry, no fluff
+- ONLY include files that were actually discussed - never invent file paths
+- If a category has no relevant items, use an empty array
+- Do NOT include any text outside the JSON block
+- Do NOT explain your reasoning - just output the JSON`;
+
+/**
+ * Stricter prompt for retry after parse failure
+ */
+export const EXTRACTION_RETRY_PROMPT = `Your previous response was not valid JSON. Please output ONLY a valid JSON object with this structure:
+
+{
+  "relevantFiles": [{ "path": "string", "reason": "string" }],
+  "relevantCommands": ["string"],
+  "relevantInformation": ["string"],
+  "decisions": ["string"],
+  "openQuestions": ["string"]
+}
+
+No explanations, no markdown, no text before or after - ONLY the JSON object.`;
+
+/**
+ * Builds the user message for the extraction call
+ */
+export function buildExtractionUserMessage(
+  conversationText: string,
+  goal: string,
+): string {
+  return `## Conversation History
+
+${conversationText}
+
+## User's Goal for New Thread
+
+${goal}
+
+Extract the relevant context for this goal and output ONLY the JSON object.`;
+}
+
+/**
+ * Result of an extraction attempt
+ */
+export interface ExtractionResult {
+  success: boolean;
+  extraction?: ExtractionOutput;
+  error?: string;
+  retried?: boolean;
+}
+
+/**
+ * Processes the LLM response and handles retry logic.
+ *
+ * @param responseText - The text response from the LLM
+ * @param config - Handoff configuration for normalization
+ * @param conversationText - Optional conversation text for file validation
+ * @returns Extraction result with normalized data or error
+ */
+export function processExtractionResponse(
+  responseText: string,
+  config: HandoffConfig,
+  conversationText?: string,
+): ParseResult & { normalized?: ExtractionOutput } {
+  const parseResult = parseExtractionResponse(responseText);
+
+  if (!parseResult.success || !parseResult.data) {
+    return parseResult;
+  }
+
+  // Normalize the extraction (dedupe, cap limits, strip @ prefixes, validate files)
+  const normalized = normalizeExtraction(parseResult.data, config, conversationText);
+
+  return {
+    success: true,
+    data: parseResult.data,
+    normalized,
+  };
+}
+
+/**
+ * Extracts text content from an assistant message
+ */
+export function extractTextFromAssistantMessage(
+  content: Array<{ type: string; text?: string }>,
+): string {
+  return content
+    .filter((c): c is { type: "text"; text: string } => c.type === "text")
+    .map((c) => c.text)
+    .join("\n");
+}

package/extensions/handoff/index.ts

@@ -0,0 +1,534 @@
+/**
+ * Handoff Extension
+ *
+ * Provides a `/handoff` command that generates a high-quality "new thread prompt"
+ * from the current session, then starts a new session with that prompt pre-filled
+ * for the user to review and send.
+ *
+ * Usage:
+ *   /handoff implement team-level handoff with proper tests
+ *   /handoff fix the authentication bug in login flow
+ *   /handoff add unit tests for the parser module
+ */
+
+import { complete, getModel, type Message, type Model } from "@mariozechner/pi-ai";
+import type {
+  ExtensionAPI,
+  ExtensionCommandContext,
+} from "@mariozechner/pi-coding-agent";
+import {
+  BorderedLoader,
+  convertToLlm,
+  serializeConversation,
+} from "@mariozechner/pi-coding-agent";
+
+import { loadConfig, validateGoal } from "./config.js";
+import { ProgressLoader, EXTRACTION_PHASES } from "./progress.js";
+import {
+  EXTRACTION_SYSTEM_PROMPT,
+  EXTRACTION_RETRY_PROMPT,
+  buildExtractionUserMessage,
+  processExtractionResponse,
+  extractTextFromAssistantMessage,
+} from "./extraction.js";
+import { collectSessionMetadata } from "./metadata.js";
+import { assembleHandoffPrompt } from "./prompt.js";
+import {
+  SKILL_ENTRY_TYPE,
+  type HandoffConfig,
+  type SkillEntry,
+} from "./types.js";
+
+/**
+ * Resolves the model to use for extraction based on config
+ */
+function resolveExtractionModel(
+  ctx: ExtensionCommandContext,
+  config: HandoffConfig,
+): Model<any> | undefined {
+  // Use current model if configured to do so or no override specified
+  if (config.useCurrentModel || !config.model) {
+    return ctx.model;
+  }
+
+  // Try to get the override model
+  const [provider, ...modelParts] = config.model.split("/");
+  const modelId = modelParts.join("/");
+
+  if (!provider || !modelId) {
+    // Invalid format, fall back to current
+    return ctx.model;
+  }
+
+  const overrideModel = getModel(provider, modelId);
+  if (!overrideModel) {
+    // Model not found, fall back to current
+    console.warn(`Handoff: Model ${config.model} not found, using current model`);
+    return ctx.model;
+  }
+
+  return overrideModel;
+}
+
+/**
+ * Main handoff command handler
+ */
+async function runHandoffCommand(
+  args: string | undefined,
+  ctx: ExtensionCommandContext,
+  pi: ExtensionAPI,
+  lastSkill: string | undefined,
+): Promise<void> {
+  // Load config from .pi/settings.json
+  const cwd = ctx.sessionManager.getCwd();
+  const config = loadConfig(cwd);
+
+  // Validate goal
+  const goal = args?.trim() ?? "";
+  const goalValidation = validateGoal(goal, config.minGoalLength);
+
+  if (!goalValidation.valid) {
+    if (ctx.hasUI) {
+      ctx.ui.notify(goalValidation.error!, "error");
+    } else {
+      console.error(goalValidation.error);
+    }
+    return;
+  }
+
+  // Check for model
+  if (!ctx.model) {
+    const errorMsg = "No model selected. Use /model to select a model first.";
+    if (ctx.hasUI) {
+      ctx.ui.notify(errorMsg, "error");
+    } else {
+      console.error(errorMsg);
+    }
+    return;
+  }
+
+  // Get conversation context
+  const sessionContext = ctx.sessionManager.buildSessionContext();
+  const messages = sessionContext.messages;
+
+  if (messages.length === 0) {
+    const errorMsg = "No conversation to hand off.";
+    if (ctx.hasUI) {
+      ctx.ui.notify(errorMsg, "error");
+    } else {
+      console.error(errorMsg);
+    }
+    return;
+  }
+
+  // Convert messages to LLM format and serialize
+  const llmMessages = convertToLlm(messages);
+  const conversationText = serializeConversation(llmMessages);
+  const currentSessionFile = ctx.sessionManager.getSessionFile();
+
+  // Collect metadata
+  const activeTools = pi.getActiveTools();
+  const sessionName = ctx.sessionManager.getSessionName();
+  const thinkingLevel = pi.getThinkingLevel();
+
+  const metadata = await collectSessionMetadata({
+    model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined,
+    thinkingLevel: thinkingLevel !== "off" ? thinkingLevel : undefined,
+    tools: activeTools,
+    sessionName: sessionName ?? undefined,
+    lastSkill,
+    exec: (cmd, args, opts) => pi.exec(cmd, args, opts),
+  });
+
+  // Resolve which model to use for extraction
+  const extractionModel = resolveExtractionModel(ctx, config);
+  if (!extractionModel) {
+    const errorMsg = "No model available for extraction.";
+    if (ctx.hasUI) {
+      ctx.ui.notify(errorMsg, "error");
+    } else {
+      console.error(errorMsg);
+    }
+    return;
+  }
+
+  // Generate extraction via LLM
+  const extractionResult = await generateExtraction(
+    conversationText,
+    goal,
+    config,
+    ctx,
+    extractionModel,
+  );
+
+  if (!extractionResult.success || !extractionResult.extraction) {
+    if (ctx.hasUI) {
+      ctx.ui.notify(
+        extractionResult.error ?? "Failed to generate handoff context",
+        "error",
+      );
+    } else {
+      console.error(extractionResult.error ?? "Failed to generate handoff context");
+    }
+    return;
+  }
+
+  // Assemble the handoff prompt
+  const handoffPrompt = assembleHandoffPrompt(
+    extractionResult.extraction,
+    goal,
+    metadata,
+    config,
+  );
+
+  // Non-UI mode: just print the prompt
+  if (!ctx.hasUI) {
+    console.log(handoffPrompt);
+    return;
+  }
+
+  // Interactive mode: let user edit the prompt
+  const editedPrompt = await ctx.ui.editor("Edit handoff prompt", handoffPrompt);
+
+  if (editedPrompt === undefined) {
+    ctx.ui.notify("Handoff cancelled", "info");
+    return;
+  }
+
+  // Create new session with parent tracking
+  const newSessionResult = await ctx.newSession({
+    parentSession: currentSessionFile,
+  });
+
+  if (newSessionResult.cancelled) {
+    ctx.ui.notify("New session cancelled", "info");
+    return;
+  }
+
+  // Set the edited prompt in the editor for submission
+  ctx.ui.setEditorText(editedPrompt);
+  ctx.ui.notify("Handoff ready. Press Enter to send.", "info");
+}
+
+/**
+ * Extraction result type
+ */
+interface ExtractionResult {
+  success: boolean;
+  extraction?: ReturnType<typeof processExtractionResponse>["normalized"];
+  error?: string;
+  completionMessage?: string;
+}
+
+/**
+ * Generates the extraction by calling the LLM with retry on parse failure
+ */
+async function generateExtraction(
+  conversationText: string,
+  goal: string,
+  config: HandoffConfig,
+  ctx: ExtensionCommandContext,
+  model: Model<any>,
+): Promise<ExtractionResult> {
+  if (!ctx.hasUI) {
+    // Non-UI mode: direct call without loader
+    return await doExtraction(conversationText, goal, config, ctx, model);
+  }
+
+  // Interactive mode: show loader during extraction
+  if (config.showProgressPhases) {
+    // Use phase-based progress loader
+    return await ctx.ui.custom<ExtractionResult>((tui, theme, _kb, done) => {
+      const loader = new ProgressLoader(tui, theme, EXTRACTION_PHASES[0]);
+      loader.onAbort = () => {
+        loader.dispose();
+        done({ success: false, error: "Cancelled" });
+      };
+
+      doExtractionWithPhases(conversationText, goal, config, ctx, model, loader.signal, (phase) => {
+        loader.setPhase(phase);
+      })
+        .then((result) => {
+          const completionMessage = loader.getCompletionMessage();
+          loader.dispose();
+          done({ ...result, completionMessage });
+        })
+        .catch((err) => {
+          loader.dispose();
+          console.error("Handoff extraction failed:", err);
+          done({ success: false, error: err.message ?? "Unknown error" });
+        });
+
+      return loader;
+    });
+  } else {
+    // Use simple bordered loader
+    return await ctx.ui.custom<ExtractionResult>((tui, theme, _kb, done) => {
+      const loader = new BorderedLoader(tui, theme, "Generating handoff context...");
+      loader.onAbort = () => done({ success: false, error: "Cancelled" });
+
+      doExtraction(conversationText, goal, config, ctx, model, loader.signal)
+        .then(done)
+        .catch((err) => {
+          console.error("Handoff extraction failed:", err);
+          done({ success: false, error: err.message ?? "Unknown error" });
+        });
+
+      return loader;
+    });
+  }
+}
+
+/**
+ * Performs the actual LLM extraction call with retry
+ */
+async function doExtraction(
+  conversationText: string,
+  goal: string,
+  config: HandoffConfig,
+  ctx: ExtensionCommandContext,
+  model: Model<any>,
+  signal?: AbortSignal,
+): Promise<ExtractionResult> {
+  const apiKey = await ctx.modelRegistry.getApiKey(model);
+
+  // Build user message
+  const userMessage: Message = {
+    role: "user",
+    content: [
+      { type: "text", text: buildExtractionUserMessage(conversationText, goal) },
+    ],
+    timestamp: Date.now(),
+  };
+
+  // First attempt
+  const response = await complete(
+    model,
+    { systemPrompt: EXTRACTION_SYSTEM_PROMPT, messages: [userMessage] },
+    { apiKey, signal },
+  );
+
+  if (response.stopReason === "aborted") {
+    return { success: false, error: "Cancelled" };
+  }
+
+  if (response.stopReason === "error") {
+    return { success: false, error: response.errorMessage ?? "LLM error" };
+  }
+
+  const responseText = extractTextFromAssistantMessage(response.content);
+  const result = processExtractionResponse(responseText, config, conversationText);
+
+  if (result.success && result.normalized) {
+    return { success: true, extraction: result.normalized };
+  }
+
+  // Retry with stricter prompt
+  const retryMessage: Message = {
+    role: "user",
+    content: [{ type: "text", text: EXTRACTION_RETRY_PROMPT }],
+    timestamp: Date.now(),
+  };
+
+  const assistantMessage: Message = {
+    role: "assistant",
+    content: response.content,
+    api: response.api,
+    provider: response.provider,
+    model: response.model,
+    usage: response.usage,
+    stopReason: response.stopReason,
+    timestamp: response.timestamp,
+  };
+
+  const retryResponse = await complete(
+    model,
+    {
+      systemPrompt: EXTRACTION_SYSTEM_PROMPT,
+      messages: [userMessage, assistantMessage, retryMessage],
+    },
+    { apiKey, signal },
+  );
+
+  if (retryResponse.stopReason === "aborted") {
+    return { success: false, error: "Cancelled" };
+  }
+
+  if (retryResponse.stopReason === "error") {
+    return { success: false, error: retryResponse.errorMessage ?? "LLM error on retry" };
+  }
+
+  const retryText = extractTextFromAssistantMessage(retryResponse.content);
+  const retryResult = processExtractionResponse(retryText, config, conversationText);
+
+  if (retryResult.success && retryResult.normalized) {
+    return { success: true, extraction: retryResult.normalized };
+  }
+
+  return {
+    success: false,
+    error: `Failed to parse extraction after retry: ${retryResult.error}`,
+  };
+}
+
+/**
+ * Performs extraction with phase updates for progress UI
+ */
+async function doExtractionWithPhases(
+  conversationText: string,
+  goal: string,
+  config: HandoffConfig,
+  ctx: ExtensionCommandContext,
+  model: Model<any>,
+  signal: AbortSignal,
+  onPhase: (phase: string) => void,
+): Promise<ExtractionResult> {
+  const apiKey = await ctx.modelRegistry.getApiKey(model);
+
+  // Phase 1: Analyzing conversation
+  onPhase(EXTRACTION_PHASES[0]);
+
+  // Build user message
+  const userMessage: Message = {
+    role: "user",
+    content: [
+      { type: "text", text: buildExtractionUserMessage(conversationText, goal) },
+    ],
+    timestamp: Date.now(),
+  };
+
+  // Phase 2: Extracting context (LLM call)
+  onPhase(EXTRACTION_PHASES[1]);
+
+  // First attempt
+  const response = await complete(
+    model,
+    { systemPrompt: EXTRACTION_SYSTEM_PROMPT, messages: [userMessage] },
+    { apiKey, signal },
+  );
+
+  if (response.stopReason === "aborted") {
+    return { success: false, error: "Cancelled" };
+  }
+
+  if (response.stopReason === "error") {
+    return { success: false, error: response.errorMessage ?? "LLM error" };
+  }
+
+  // Phase 3: Assembling prompt
+  onPhase(EXTRACTION_PHASES[2]);
+
+  const responseText = extractTextFromAssistantMessage(response.content);
+  const result = processExtractionResponse(responseText, config, conversationText);
+
+  if (result.success && result.normalized) {
+    return { success: true, extraction: result.normalized };
+  }
+
+  // Retry needed - stay on phase 2
+  onPhase("Retrying extraction...");
+
+  const retryMessage: Message = {
+    role: "user",
+    content: [{ type: "text", text: EXTRACTION_RETRY_PROMPT }],
+    timestamp: Date.now(),
+  };
+
+  const assistantMessage: Message = {
+    role: "assistant",
+    content: response.content,
+    api: response.api,
+    provider: response.provider,
+    model: response.model,
+    usage: response.usage,
+    stopReason: response.stopReason,
+    timestamp: response.timestamp,
+  };
+
+  const retryResponse = await complete(
+    model,
+    {
+      systemPrompt: EXTRACTION_SYSTEM_PROMPT,
+      messages: [userMessage, assistantMessage, retryMessage],
+    },
+    { apiKey, signal },
+  );
+
+  if (retryResponse.stopReason === "aborted") {
+    return { success: false, error: "Cancelled" };
+  }
+
+  if (retryResponse.stopReason === "error") {
+    return { success: false, error: retryResponse.errorMessage ?? "LLM error on retry" };
+  }
+
+  // Back to phase 3
+  onPhase(EXTRACTION_PHASES[2]);
+
+  const retryText = extractTextFromAssistantMessage(retryResponse.content);
+  const retryResult = processExtractionResponse(retryText, config, conversationText);
+
+  if (retryResult.success && retryResult.normalized) {
+    return { success: true, extraction: retryResult.normalized };
+  }
+
+  return {
+    success: false,
+    error: `Failed to parse extraction after retry: ${retryResult.error}`,
+  };
+}
+
+/**
+ * Main extension entry point
+ */
+export default function handoffExtension(pi: ExtensionAPI) {
+  // Track last used skill
+  let lastSkill: string | undefined;
+
+  // Restore last skill from session on startup
+  pi.on("session_start", async (_event, ctx) => {
+    lastSkill = undefined;
+    for (const entry of ctx.sessionManager.getEntries()) {
+      if (
+        entry.type === "custom" &&
+        (entry as any).customType === SKILL_ENTRY_TYPE
+      ) {
+        const data = (entry as any).data as SkillEntry | undefined;
+        if (data?.skillName) {
+          lastSkill = data.skillName;
+        }
+      }
+    }
+  });
+
+  // Track skill usage via input event
+  pi.on("input", async (event, _ctx) => {
+    const text = event.text.trim();
+
+    // Check if this is a skill command
+    if (text.startsWith("/skill:")) {
+      const skillMatch = text.match(/^\/skill:([^\s]+)/);
+      if (skillMatch) {
+        const skillName = skillMatch[1];
+        lastSkill = skillName;
+
+        // Persist to session
+        pi.appendEntry(SKILL_ENTRY_TYPE, {
+          skillName,
+          timestamp: Date.now(),
+        } as SkillEntry);
+      }
+    }
+
+    // Let the input continue processing
+    return { action: "continue" };
+  });
+
+  // Register the /handoff command
+  pi.registerCommand("handoff", {
+    description: "Transfer context to a new focused session",
+    handler: async (args, ctx) => {
+      await runHandoffCommand(args, ctx, pi, lastSkill);
+    },
+  });
+}