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 = 1;

// 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 auth = await ctx.modelRegistry.getApiKeyAndHeaders(AUTO_NAME_MODEL);
			log(`Got API key: ${auth.ok ? "yes" : "no"}`);

			if (!auth.ok) {
				log(`No API key available, aborting: ${auth.error}`);
				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: auth.apiKey, headers: auth.headers },
			);

			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 auth = await ctx.modelRegistry.getApiKeyAndHeaders(AUTO_NAME_MODEL);
							if (!auth.ok) throw new Error(auth.error);

							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: auth.apiKey, headers: auth.headers, 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");
		},
	});
}