dotfiles/pi/files/agent/extensions/session-name.ts

/**
 * Enhanced session naming with AI-powered auto-naming.
 *
 * Features:
 * - Manual naming: /session-name [name]
 * - Auto-naming command: /session-name --auto
 * - Automatic naming: triggers after 3 messages if no name set
 *
 * Auto-naming analyzes the conversation history and generates a concise,
 * descriptive name for the session using a cheap model (MiniMax).
 */

import { complete, type Message } from "@mariozechner/pi-ai";
import { getModel } from "@mariozechner/pi-ai";
import type { ExtensionAPI, SessionEntry } from "@mariozechner/pi-coding-agent";
import {
  BorderedLoader,
  convertToLlm,
  serializeConversation,
} from "@mariozechner/pi-coding-agent";
import * as fs from "node:fs";
import * as path from "node:path";
import * as os from "node:os";

const SYSTEM_PROMPT = `You are a session naming assistant. Given a conversation history, generate a short, descriptive session name (2-5 words) that captures the main topic or task.

Guidelines:
- Be concise but specific
- Use kebab-case or natural language
- Focus on the core task/question
- Avoid generic names like "discussion" or "conversation"
- No quotes, no punctuation at the end

Examples:
- "fix auth bug" -> "fix-auth-bug" or "authentication fix"
- "how do I deploy to vercel" -> "vercel deployment"
- "explain react hooks" -> "react hooks explanation"
- "optimize database queries" -> "db query optimization"

Output ONLY the session name, nothing else.`;

// Cheap model for auto-naming: Minimax from OpenCode Go
const AUTO_NAME_MODEL = getModel("opencode-go", "minimax-m2.5");

// Number of messages before auto-naming kicks in
const AUTO_NAME_THRESHOLD = 2;

// Debug log file
const LOG_FILE = path.join(os.homedir(), ".pi", "session-name-debug.log");

function log(message: string) {
  const timestamp = new Date().toISOString();
  const entry = `[${timestamp}] ${message}\n`;
  try {
    fs.appendFileSync(LOG_FILE, entry);
  } catch (e) {
    // ignore write errors
  }
}

export default function (pi: ExtensionAPI) {
  // Track if we've already attempted auto-naming for this session
  let autoNamedAttempted = false;

  // Listen for agent_end to auto-name sessions (non-blocking)
  pi.on("agent_end", (_event, ctx) => {
    log("=== agent_end triggered ===");
    log(`hasUI: ${ctx.hasUI}`);
    log(`current session name: ${pi.getSessionName()}`);
    log(`autoNamedAttempted: ${autoNamedAttempted}`);

    // Skip if already has a name or already attempted
    if (pi.getSessionName() || autoNamedAttempted) {
      log("Skipping: already has name or already attempted");
      return;
    }

    // Count user messages in the branch
    const branch = ctx.sessionManager.getBranch();
    const userMessages = branch.filter(
      (entry): entry is SessionEntry & { type: "message" } =>
        entry.type === "message" && entry.message.role === "user",
    );

    log(`Total entries in branch: ${branch.length}`);
    log(`User messages: ${userMessages.length}`);
    log(`Threshold: ${AUTO_NAME_THRESHOLD}`);

    // Only auto-name after threshold is reached
    if (userMessages.length < AUTO_NAME_THRESHOLD) {
      log("Skipping: below threshold");
      return;
    }

    // Mark as attempted so we don't try again
    autoNamedAttempted = true;
    log("Threshold reached, attempting auto-name");

    // Only auto-name in interactive mode
    if (!ctx.hasUI) {
      log("Skipping: no UI (non-interactive mode)");
      return;
    }

    // Gather conversation context
    const messages = branch
      .filter(
        (entry): entry is SessionEntry & { type: "message" } =>
          entry.type === "message",
      )
      .map((entry) => entry.message);

    log(`Total messages to analyze: ${messages.length}`);

    if (messages.length === 0) {
      log("No messages found, aborting");
      return;
    }

    // Convert to LLM format and serialize
    const llmMessages = convertToLlm(messages);
    const conversationText = serializeConversation(llmMessages);

    log(`Conversation text length: ${conversationText.length}`);

    // Truncate if too long (keep costs low)
    const maxChars = 4000;
    const truncatedText =
      conversationText.length > maxChars
        ? conversationText.slice(0, maxChars) + "\n..."
        : conversationText;

    log(`Truncated text length: ${truncatedText.length}`);
    log("Starting background auto-name...");

    // Fire-and-forget: run auto-naming in background without blocking
    const doAutoName = async () => {
      const apiKey = await ctx.modelRegistry.getApiKey(AUTO_NAME_MODEL);
      log(`Got API key: ${apiKey ? "yes" : "no"}`);

      if (!apiKey) {
        log("No API key available, aborting");
        return;
      }

      const userMessage: Message = {
        role: "user",
        content: [
          {
            type: "text",
            text: `## Conversation History\n\n${truncatedText}\n\nGenerate a concise session name for this conversation.`,
          },
        ],
        timestamp: Date.now(),
      };

      const response = await complete(
        AUTO_NAME_MODEL,
        { systemPrompt: SYSTEM_PROMPT, messages: [userMessage] },
        { apiKey },
      );

      log(`Response received, stopReason: ${response.stopReason}`);

      if (response.stopReason === "aborted") {
        log("Request was aborted");
        return;
      }

      const name = response.content
        .filter((c): c is { type: "text"; text: string } => c.type === "text")
        .map((c) => c.text.trim())
        .join(" ")
        .replace(/^[\"']|[\"']$/g, ""); // Remove surrounding quotes

      log(`Generated name: "${name}"`);

      // Clean up the generated name
      const cleanName = name
        .replace(/\n/g, " ")
        .replace(/\s+/g, " ")
        .trim()
        .slice(0, 50); // Max 50 chars

      log(`Cleaned name: "${cleanName}"`);

      if (cleanName) {
        pi.setSessionName(cleanName);
        ctx.ui.notify(`Auto-named: ${cleanName}`, "info");
        log(`Successfully set session name to: ${cleanName}`);
      } else {
        log("Cleaned name was empty, not setting");
      }
    };

    // Run in background without awaiting - don't block the agent
    doAutoName().catch((err) => {
      log(`ERROR: ${err}`);
      console.error("Auto-naming failed:", err);
    });
  });

  // Reset flag on new session
  pi.on("session_start", () => {
    log("=== session_start ===");
    autoNamedAttempted = false;
    log("Reset autoNamedAttempted to false");
  });

  pi.on("session_switch", () => {
    log("=== session_switch ===");
    autoNamedAttempted = false;
    log("Reset autoNamedAttempted to false");
  });

  // Manual command for setting/getting session name
  pi.registerCommand("session-name", {
    description:
      "Set, show, or auto-generate session name (usage: /session-name [name] or /session-name --auto)",
    handler: async (args, ctx) => {
      const trimmedArgs = args.trim();

      // Show current name if no args
      if (!trimmedArgs) {
        const current = pi.getSessionName();
        ctx.ui.notify(
          current ? `Session: ${current}` : "No session name set",
          "info",
        );
        return;
      }

      // Auto-generate name using AI
      if (trimmedArgs === "--auto" || trimmedArgs === "-a") {
        if (!ctx.hasUI) {
          ctx.ui.notify("Auto-naming requires interactive mode", "error");
          return;
        }

        // Gather conversation context
        const branch = ctx.sessionManager.getBranch();
        const messages = branch
          .filter(
            (entry): entry is SessionEntry & { type: "message" } =>
              entry.type === "message",
          )
          .map((entry) => entry.message);

        if (messages.length === 0) {
          ctx.ui.notify("No conversation to analyze", "error");
          return;
        }

        // Convert to LLM format and serialize
        const llmMessages = convertToLlm(messages);
        const conversationText = serializeConversation(llmMessages);

        // Truncate if too long (keep costs low)
        const maxChars = 4000;
        const truncatedText =
          conversationText.length > maxChars
            ? conversationText.slice(0, maxChars) + "\n..."
            : conversationText;

        // Generate name with loader UI
        const result = await ctx.ui.custom<string | null>(
          (tui, theme, _kb, done) => {
            const loader = new BorderedLoader(
              tui,
              theme,
              "Generating session name...",
            );
            loader.onAbort = () => done(null);

            const doGenerate = async () => {
              const apiKey = await ctx.modelRegistry.getApiKey(AUTO_NAME_MODEL);

              const userMessage: Message = {
                role: "user",
                content: [
                  {
                    type: "text",
                    text: `## Conversation History\n\n${truncatedText}\n\nGenerate a concise session name for this conversation.`,
                  },
                ],
                timestamp: Date.now(),
              };

              const response = await complete(
                AUTO_NAME_MODEL,
                { systemPrompt: SYSTEM_PROMPT, messages: [userMessage] },
                { apiKey, signal: loader.signal },
              );

              if (response.stopReason === "aborted") {
                return null;
              }

              const name = response.content
                .filter(
                  (c): c is { type: "text"; text: string } => c.type === "text",
                )
                .map((c) => c.text.trim())
                .join(" ")
                .replace(/^[\"']|[\"']$/g, ""); // Remove surrounding quotes

              return name;
            };

            doGenerate()
              .then(done)
              .catch((err) => {
                console.error("Auto-naming failed:", err);
                done(null);
              });

            return loader;
          },
        );

        if (result === null) {
          ctx.ui.notify("Auto-naming cancelled", "info");
          return;
        }

        // Clean up the generated name
        const cleanName = result
          .replace(/\n/g, " ")
          .replace(/\s+/g, " ")
          .trim()
          .slice(0, 50); // Max 50 chars

        if (!cleanName) {
          ctx.ui.notify("Failed to generate name", "error");
          return;
        }

        pi.setSessionName(cleanName);
        ctx.ui.notify(`Session auto-named: ${cleanName}`, "info");
        return;
      }

      // Manual naming
      pi.setSessionName(trimmedArgs);
      ctx.ui.notify(`Session named: ${trimmedArgs}`, "info");
    },
  });
}