pi-ext

/**
 * session-rename — auto-generate and refine short Pi session names
 *
 * Policy:
 *   - User/manual names are never overwritten automatically.
 *   - Extension-generated names are continuously refined with model "keep or rename" judgement.
 *   - Rename cadence slows as session grows to reduce churn.
 *   - Names are 1–3 words.
 */

import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg, span } from "./src/debug.ts";
import {
	MAX_TIMER_MS,
	parseNumberSetting,
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";

type ContentBlock = {
	type?: string;
	text?: string;
	name?: string;
	arguments?: Record<string, unknown>;
};
type Entry = {
	id?: string;
	type: string;
	customType?: string;
	data?: unknown;
	summary?: string;
	message?: { role?: string; content?: unknown; toolName?: string };
};
type Model = NonNullable<ExtensionContext["model"]>;
type RenameReason = "load" | "turn" | "manual";

type MarkerData = {
	kind?: "auto-title" | "user-lock";
	name?: string;
	userTurnCount?: number;
	revision?: number;
	previousNames?: string[];
};

type TitleState =
	| { mode: "none"; latestMarker?: MarkerData }
	| { mode: "auto"; name: string; latestMarker: MarkerData }
	| { mode: "manual"; name?: string; latestMarker?: MarkerData };

const CUSTOM_TYPE = "session-rename";
const STATUS_KEY = "session-rename";
const IDLE_DELAY_SETTING: SettingDeclaration<number> = {
	key: "session-rename.idleDelayMs",
	env: "PI_SESSION_RENAME_IDLE_DELAY_MS",
	default: 10_000,
	parse: (raw, source) => {
		const value = parseNumberSetting(raw, source);
		if (value === undefined) return undefined;
		if (!Number.isInteger(value) || value < 0 || value > MAX_TIMER_MS)
			throw new Error(`expected an integer from 0 to ${MAX_TIMER_MS}`);
		return value;
	},
};
const LOAD_DELAY_MS = 700;
const MAX_TRANSCRIPT_CHARS = 8000;
const MAX_NAME_WORDS = 3;
const SMALL_EDIT_DISTANCE = 2;
const GENERIC_SINGLE_WORDS = new Set([
	"task",
	"work",
	"coding",
	"chat",
	"conversation",
	"session",
]);

// ── Text/session helpers ──────────────────────────────────────────────────

function extractText(content: unknown): string {
	if (typeof content === "string") return content;
	if (!Array.isArray(content)) return "";
	return (content as ContentBlock[])
		.flatMap((p) =>
			p?.type === "text" && typeof p.text === "string" ? [p.text] : [],
		)
		.join("\n");
}

function extractToolCalls(content: unknown): string[] {
	if (!Array.isArray(content)) return [];
	return (content as ContentBlock[])
		.filter((p) => p?.type === "toolCall" && typeof p.name === "string")
		.map((p) => {
			const summary = JSON.stringify(p.arguments ?? {}).slice(0, 240);
			return `Tool: ${p.name}(${summary})`;
		});
}

function countUserTurns(entries: Entry[]): number {
	let count = 0;
	for (const e of entries) {
		if (
			e.type === "message" &&
			e.message?.role === "user" &&
			extractText(e.message.content).trim()
		)
			count++;
	}
	return count;
}

function hasMeaningfulActivity(entries: Entry[]): boolean {
	let users = 0;
	let assistants = 0;
	let toolCalls = 0;
	for (const e of entries) {
		if (e.type !== "message" || !e.message?.role) continue;
		if (e.message.role === "user" && extractText(e.message.content).trim())
			users++;
		if (e.message.role === "assistant") {
			if (extractText(e.message.content).trim()) assistants++;
			toolCalls += extractToolCalls(e.message.content).length;
		}
	}
	return users > 0 && (assistants > 0 || toolCalls > 0);
}

function buildTranscript(
	entries: Entry[],
	oldName: string | undefined,
): string {
	const firstUsers: string[] = [];
	const summaries: string[] = [];
	const messageLines: string[] = [];
	const messageEntries = entries.filter(
		(e) => e.type === "message" && e.message?.role,
	);

	for (const e of entries) {
		if (
			(e.type === "compaction" || e.type === "branch_summary") &&
			typeof e.summary === "string" &&
			e.summary.trim()
		) {
			summaries.push(`${e.type}: ${e.summary.trim().slice(0, 1200)}`);
		}
		if (
			e.type === "message" &&
			e.message?.role === "user" &&
			firstUsers.length < 3
		) {
			const text = extractText(e.message.content).trim();
			if (text)
				firstUsers.push(
					`User goal ${firstUsers.length + 1}: ${text.slice(0, 1000)}`,
				);
		}
	}

	for (const e of messageEntries.slice(-10)) {
		const role = e.message?.role;
		if (role === "user") {
			const text = extractText(e.message?.content).trim();
			if (text) messageLines.push(`Recent user: ${text.slice(0, 900)}`);
		} else if (role === "assistant") {
			const text = extractText(e.message?.content).trim();
			if (text) messageLines.push(`Recent assistant: ${text.slice(0, 900)}`);
			messageLines.push(...extractToolCalls(e.message?.content).slice(0, 8));
		} else if (role === "toolResult") {
			const text = extractText(e.message?.content).trim();
			if (text)
				messageLines.push(
					`Result(${e.message?.toolName ?? "tool"}): ${text.slice(0, 300)}`,
				);
		}
	}

	const sections = [
		oldName ? `Current name: ${oldName}` : "Current name: <none>",
		firstUsers.length ? `Initial goals:\n${firstUsers.join("\n")}` : "",
		summaries.length ? `Summaries:\n${summaries.slice(-2).join("\n")}` : "",
		messageLines.length ? `Recent activity:\n${messageLines.join("\n")}` : "",
	].filter(Boolean);

	return sections.join("\n\n").slice(0, MAX_TRANSCRIPT_CHARS);
}

function currentSessionFile(ctx: ExtensionContext): string | undefined {
	try {
		return ctx.sessionManager.getSessionFile();
	} catch {
		return undefined;
	}
}

function currentName(ctx: ExtensionContext): string | undefined {
	try {
		return ctx.sessionManager.getSessionName();
	} catch {
		return undefined;
	}
}

function allEntries(ctx: ExtensionContext): Entry[] {
	try {
		return ctx.sessionManager.getEntries() as Entry[];
	} catch {
		return [];
	}
}

function branchEntries(ctx: ExtensionContext): Entry[] {
	try {
		return ctx.sessionManager.getBranch() as Entry[];
	} catch {
		return [];
	}
}

// ── Marker/manual-vs-auto state ────────────────────────────────────────────

function isMarkerData(data: unknown): data is MarkerData {
	return (
		typeof data === "object" &&
		data !== null &&
		((data as MarkerData).kind === "auto-title" ||
			(data as MarkerData).kind === "user-lock")
	);
}

function getTitleState(ctx: ExtensionContext): TitleState {
	const entries = allEntries(ctx);
	const name = currentName(ctx);
	let latestSessionInfoIndex = -1;
	let latestSessionInfoName: string | undefined;
	let latestMarkerIndex = -1;
	let latestMarker: MarkerData | undefined;

	for (let i = entries.length - 1; i >= 0; i--) {
		const e = entries[i] as Entry & { name?: string };
		if (latestSessionInfoIndex < 0 && e.type === "session_info") {
			latestSessionInfoIndex = i;
			latestSessionInfoName =
				typeof e.name === "string" ? e.name.trim() : undefined;
		}
		if (
			latestMarkerIndex < 0 &&
			e.type === "custom" &&
			e.customType === CUSTOM_TYPE &&
			isMarkerData(e.data)
		) {
			latestMarkerIndex = i;
			latestMarker = e.data;
		}
		if (latestSessionInfoIndex >= 0 && latestMarkerIndex >= 0) break;
	}

	if (
		latestMarker?.kind === "user-lock" &&
		latestMarkerIndex > latestSessionInfoIndex
	) {
		return { mode: "manual", name, latestMarker };
	}

	if (name) {
		const markerMatches =
			latestMarker?.kind === "auto-title" &&
			canonicalName(latestMarker.name ?? "") === canonicalName(name);
		if (
			latestMarker &&
			markerMatches &&
			latestMarkerIndex > latestSessionInfoIndex
		) {
			return { mode: "auto", name, latestMarker };
		}
		return { mode: "manual", name, latestMarker };
	}

	// Explicitly-cleared title counts as manual; no surprise re-title after user clears it.
	if (
		latestSessionInfoIndex >= 0 &&
		!latestSessionInfoName &&
		latestMarkerIndex < latestSessionInfoIndex
	) {
		return { mode: "manual", latestMarker };
	}

	return { mode: "none", latestMarker };
}

function appendAutoMarker(
	pi: ExtensionAPI,
	name: string,
	turnCount: number,
	previous?: MarkerData,
) {
	const priorNames = [
		...(previous?.previousNames ?? []),
		previous?.name,
	].filter((n): n is string => typeof n === "string" && n.trim().length > 0);

	const uniquePrior = Array.from(
		new Map(priorNames.map((n) => [canonicalName(n), n])).values(),
	).slice(-8);
	pi.appendEntry<MarkerData>(CUSTOM_TYPE, {
		kind: "auto-title",
		name,
		userTurnCount: turnCount,
		revision: (previous?.revision ?? 0) + 1,
		previousNames: uniquePrior,
	});
}

function appendUserLock(pi: ExtensionAPI, name: string, turnCount: number) {
	pi.appendEntry<MarkerData>(CUSTOM_TYPE, {
		kind: "user-lock",
		name,
		userTurnCount: turnCount,
	});
}

// ── Name generation/normalization ──────────────────────────────────────────

async function generateName(
	transcript: string,
	oldName: string | undefined,
	ctx: ExtensionContext,
	signal: AbortSignal,
): Promise<string | undefined> {
	const model: Model | undefined = ctx.model;
	if (!model) return undefined;

	const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
	if (!auth.ok) return undefined;

	const prompt =
		"Generate or refine a short display name for a Pi coding-agent session.\n\n" +
		`Current auto-generated name: ${oldName ? `"${oldName}"` : "<none>"}\n\n` +
		"Rules:\n" +
		"- Output exactly one line.\n" +
		"- Name must be 1-3 words.\n" +
		"- No quotes, markdown, emoji, punctuation, or trailing period.\n" +
		"- Prefer concrete nouns: project, file, feature, bug, domain.\n" +
		"- If current name is still accurate enough, output it exactly.\n" +
		"- Rename only if durable topic changed or current name is misleading/vague.\n" +
		"- Avoid synonyms, grammar tweaks, or style-only changes.\n" +
		"- Do not chase the latest incidental tool/action.\n" +
		"- Avoid generic one-word names like task, work, coding, chat, conversation, session.\n\n" +
		"<session>\n" +
		transcript +
		"\n</session>";

	const finish = span?.("name.generate");
	try {
		const response = await ctx.modelRegistry.complete(
			model,
			{
				systemPrompt:
					"You name coding-agent sessions tersely and conservatively.",
				messages: [
					{
						role: "user",
						content: [{ type: "text", text: prompt }],
						timestamp: Date.now(),
					},
				],
			},
			{
				signal,
				cacheRetention: "none",
				...(model.reasoning ? { reasoning: "medium" as const } : {}),
			},
		);

		const text = response.content
			.filter((c): c is { type: "text"; text: string } => c.type === "text")
			.map((c) => c.text)
			.join("\n")
			.trim();
		finish?.();
		return normalizeGeneratedName(text);
	} catch (error) {
		finish?.("error");
		throw error;
	}
}

function stripMarkdownWrapping(s: string): string {
	return s.replace(/^(```|`)+|(```|`)+$/g, "").trim();
}

function stripQuotes(s: string): string {
	return s.replace(/^['“”‘’"]+|['“”‘’"]+$/g, "").trim();
}

function stripLabelPrefix(s: string): string {
	return s.replace(/^(title|name|session)\s*[:=-]\s*/i, "");
}

function normalizeToAsciiWords(s: string): string {
	return s
		.normalize("NFKC")
		.toLowerCase()
		.replace(/[^\p{L}\p{N}\s]/gu, " ")
		.replace(/\s+/g, " ")
		.trim();
}

function normalizeGeneratedName(raw: string | undefined): string | undefined {
	if (!raw) return undefined;
	let s = raw.split(/\r?\n/, 1)[0] ?? "";
	s = stripMarkdownWrapping(s);
	s = stripQuotes(s);
	s = stripLabelPrefix(s);
	s = normalizeToAsciiWords(s);
	if (!s) return undefined;

	const words = s.split(" ").filter(Boolean).slice(0, MAX_NAME_WORDS);
	if (words.length === 0) return undefined;
	if (words.length === 1 && GENERIC_SINGLE_WORDS.has(words[0]))
		return undefined;
	return words.join(" ");
}

function canonicalName(name: string): string {
	return (
		normalizeGeneratedName(name) ??
		name.trim().toLowerCase().replace(/\s+/g, " ")
	);
}

function levenshtein(a: string, b: string): number {
	if (a === b) return 0;
	if (a.length === 0) return b.length;
	if (b.length === 0) return a.length;

	let prev = Array.from({ length: b.length + 1 }, (_, i) => i);
	let curr = new Array<number>(b.length + 1);
	for (let i = 1; i <= a.length; i++) {
		curr[0] = i;
		for (let j = 1; j <= b.length; j++) {
			const cost = a[i - 1] === b[j - 1] ? 0 : 1;
			curr[j] = Math.min(curr[j - 1] + 1, prev[j] + 1, prev[j - 1] + cost);
		}
		[prev, curr] = [curr, prev];
	}
	return prev[b.length];
}

function shouldApplyCandidate(
	candidate: string,
	oldName: string | undefined,
	marker: MarkerData | undefined,
): boolean {
	if (!oldName) return true;
	const oldCanonical = canonicalName(oldName);
	const nextCanonical = canonicalName(candidate);
	if (!nextCanonical || nextCanonical === oldCanonical) return false;
	if (levenshtein(oldCanonical, nextCanonical) <= SMALL_EDIT_DISTANCE)
		return false;
	const previous = marker?.previousNames ?? [];
	if (previous.some((name) => canonicalName(name) === nextCanonical))
		return false;
	return true;
}

function cooldownTurns(turnCount: number): number {
	if (turnCount <= 3) return 1;
	if (turnCount <= 10) return 2;
	return 5;
}

// ── Extension ──────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
	let idleTimer: ReturnType<typeof setTimeout> | undefined;
	let activeController: AbortController | undefined;
	let _currentCtx: ExtensionContext | undefined;
	let idleDelayMs = IDLE_DELAY_SETTING.default;

	const clearTimer = () => {
		if (idleTimer) {
			clearTimeout(idleTimer);
			idleTimer = undefined;
		}
	};

	const cancelActive = (ctx?: ExtensionContext) => {
		if (activeController) {
			activeController.abort();
			activeController = undefined;
		}
		if (ctx?.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
	};

	const scheduleRename = (
		ctx: ExtensionContext,
		reason: RenameReason,
		delayMs = idleDelayMs,
	) => {
		clearTimer();
		if (!currentSessionFile(ctx)) return;
		idleTimer = setTimeout(() => {
			idleTimer = undefined;
			void generateAndMaybeApply(ctx, reason);
		}, delayMs);
	};

	const generateAndMaybeApply = async (
		ctx: ExtensionContext,
		reason: RenameReason,
		force = false,
	) => {
		const sessionFile = currentSessionFile(ctx);
		if (!sessionFile) {
			dbg?.("name.outcome", { outcome: "no_session" });
			return;
		}

		const branch = branchEntries(ctx);
		if (!hasMeaningfulActivity(branch)) {
			dbg?.("name.outcome", { outcome: "no_activity" });
			return;
		}

		const turnCount = countUserTurns(branch);
		const state = getTitleState(ctx);
		if (state.mode === "manual" && !force) {
			dbg?.("name.outcome", { outcome: "manual_lock" });
			return;
		}

		const oldName = state.mode === "auto" ? state.name : undefined;
		if (state.mode === "auto" && !force) {
			const lastTurn = state.latestMarker.userTurnCount ?? 0;
			if (turnCount - lastTurn < cooldownTurns(turnCount)) {
				dbg?.("name.outcome", { outcome: "cooldown" });
				return;
			}
		}

		const startName = currentName(ctx);
		const transcript = buildTranscript(branch, oldName);
		if (!transcript.trim()) {
			dbg?.("name.outcome", { outcome: "no_transcript" });
			return;
		}

		cancelActive(ctx);
		const controller = new AbortController();
		activeController = controller;
		if (ctx.hasUI)
			ctx.ui.setStatus(STATUS_KEY, ctx.ui.theme.fg("dim", "󰑕 naming session…"));

		try {
			const candidate = await generateName(
				transcript,
				oldName,
				ctx,
				controller.signal,
			);
			if (!candidate) {
				dbg?.("name.outcome", { outcome: "no_candidate" });
				return;
			}
			if (controller.signal.aborted) {
				dbg?.("name.outcome", { outcome: "cancelled" });
				return;
			}
			if (currentSessionFile(ctx) !== sessionFile) {
				dbg?.("name.outcome", { outcome: "stale_session" });
				return;
			}
			if (currentName(ctx) !== startName) {
				dbg?.("name.outcome", { outcome: "name_changed" });
				return;
			}
			if (!shouldApplyCandidate(candidate, oldName, state.latestMarker)) {
				dbg?.("name.outcome", { outcome: "unchanged" });
				return;
			}

			pi.setSessionName(candidate);
			appendAutoMarker(
				pi,
				candidate,
				turnCount,
				state.mode === "auto" ? state.latestMarker : undefined,
			);
			if (ctx.hasUI && reason === "manual")
				ctx.ui.notify(`session → ${candidate}`, "info");
			dbg?.("name.outcome", { outcome: "renamed" });
		} catch (err) {
			dbg?.("name.outcome", {
				outcome: controller.signal.aborted ? "cancelled" : "failed",
			});
			if (!controller.signal.aborted && ctx.hasUI)
				ctx.ui.notify(
					`session-rename failed: ${err instanceof Error ? err.message : String(err)}`,
					"error",
				);
		} finally {
			if (activeController === controller) {
				activeController = undefined;
				if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
			}
		}
	};

	pi.on("session_start", async (_event, ctx) => {
		dbg?.("session.start");
		_currentCtx = ctx;
		const idleDelay = resolveSetting(pi, ctx, IDLE_DELAY_SETTING);
		idleDelayMs = idleDelay.ok ? idleDelay.value : IDLE_DELAY_SETTING.default;
		if (!idleDelay.ok)
			ctx.ui.notify(
				`${idleDelay.error}; using ${IDLE_DELAY_SETTING.default} ms`,
				"warning",
			);
		scheduleRename(ctx, "load", LOAD_DELAY_MS);
	});

	pi.on("turn_end", async (_event, ctx) => {
		_currentCtx = ctx;
		scheduleRename(ctx, "turn");
	});

	pi.on("turn_start", async (_event, ctx) => {
		clearTimer();
		cancelActive(ctx);
	});

	pi.on("input", async (_event, ctx) => {
		clearTimer();
		cancelActive(ctx);
	});

	pi.on("agent_start", async (_event, ctx) => {
		clearTimer();
		cancelActive(ctx);
	});

	pi.on("session_shutdown", async (_event, ctx) => {
		dbg?.("session.shutdown");
		clearTimer();
		cancelActive(ctx);
		_currentCtx = undefined;
		closeDebug();
	});

	pi.registerCommand("session-rename", {
		description: "Generate/refine session name, or set exact name with args",
		handler: async (args, ctx) => {
			const exact = args.trim();
			if (exact) {
				clearTimer();
				cancelActive(ctx);
				const name = exact.replace(/\s+/g, " ").trim();
				pi.setSessionName(name);
				appendUserLock(pi, name, countUserTurns(branchEntries(ctx)));
				if (ctx.hasUI) ctx.ui.notify(`session → ${name}`, "info");
				return;
			}
			await generateAndMaybeApply(ctx, "manual", true);
		},
	});
}