dotfiles/pi/files/agent/extensions/timestamps.ts

/**
 * Timestamps extension for Pi.
 *
 * - Shows elapsed session time in footer (updates every second)
 * - Shows how long the last turn took
 * - Injects timestamp markers without triggering extra turns
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Box, Text } from "@mariozechner/pi-tui";

// Track session time
let sessionStart = Date.now();
let timerHandle: ReturnType<typeof setInterval> | null = null;
let turnStartTime: number | null = null;
let lastTurnDuration: number | null = null;

function formatTime(date: Date): string {
	return date.toLocaleTimeString("en-US", {
		hour: "2-digit",
		minute: "2-digit",
		second: "2-digit",
		hour12: false,
	});
}

function formatElapsed(ms: number): string {
	const seconds = Math.floor(ms / 1000);
	const minutes = Math.floor(seconds / 60);
	const hours = Math.floor(minutes / 60);

	if (hours > 0) return `${hours}h ${minutes % 60}m`;
	if (minutes > 0) return `${minutes}m ${seconds % 60}s`;
	return `${seconds}s`;
}

function formatDuration(ms: number): string {
	if (ms < 1000) return `${ms}ms`;
	if (ms < 60000) return `${(ms / 1000).toFixed(1)}s`;
	return `${(ms / 60000).toFixed(1)}m`;
}

export default function (pi: ExtensionAPI) {
	const updateStatus = (ctx: { ui: { setStatus: (id: string, text: string | undefined) => void; theme: { fg: (color: string, text: string) => string } } }) => {
		const elapsed = Date.now() - sessionStart;
		let status = ctx.ui.theme.fg("dim", `⏱ ${formatElapsed(elapsed)}`);
		if (lastTurnDuration !== null) {
			status += ctx.ui.theme.fg("muted", ` | took ${formatDuration(lastTurnDuration)}`);
		}
		ctx.ui.setStatus("timestamps", status);
	};

	// Renderer for user timestamp
	pi.registerMessageRenderer("timestamp-prefix", (message, _options, theme) => {
		const details = message.details as { timestamp: number; elapsed: string } | undefined;
		if (!details) return new Text("");

		const timeStr = formatTime(new Date(details.timestamp));
		const line = theme.fg("dim", `◆ ${timeStr} (+${details.elapsed})`);
		const box = new Box(0, 0, undefined);
		box.addChild(new Text(line, 0, 0));
		return box;
	});

	// Renderer for assistant timestamp
	pi.registerMessageRenderer("timestamp-suffix", (message, _options, theme) => {
		const details = message.details as { timestamp: number; elapsed: string; duration?: number } | undefined;
		if (!details) return new Text("");

		const timeStr = formatTime(new Date(details.timestamp));
		let line = `○ ${timeStr} (+${details.elapsed})`;
		if (details.duration && details.duration > 1000) {
			line += ` • took ${formatDuration(details.duration)}`;
		}

		const box = new Box(0, 0, undefined);
		box.addChild(new Text(theme.fg("dim", line), 0, 0));
		return box;
	});

	// Start timer on session start
	pi.on("session_start", async (_event, ctx) => {
		sessionStart = Date.now();
		turnStartTime = null;
		lastTurnDuration = null;

		if (timerHandle) clearInterval(timerHandle);
		timerHandle = setInterval(() => updateStatus(ctx), 1000);
		updateStatus(ctx);
	});

	// Track turn timing
	pi.on("turn_start", async () => {
		turnStartTime = Date.now();
	});

	pi.on("turn_end", async (_event, ctx) => {
		if (turnStartTime !== null) {
			lastTurnDuration = Date.now() - turnStartTime;
			turnStartTime = null;
			updateStatus(ctx);
		}

		if (lastTurnDuration !== null) {
			const now = Date.now();
			const elapsed = formatElapsed(now - sessionStart);
			pi.sendMessage(
				{
					customType: "timestamp-suffix",
					content: "",
					display: true,
					details: { timestamp: now, elapsed, duration: lastTurnDuration },
				},
				{ deliverAs: "followUp", triggerTurn: false }
			);
		}
	});

	// Inject timestamp marker right before agent starts
	pi.on("before_agent_start", async (event) => {
		const now = Date.now();
		const elapsed = formatElapsed(now - sessionStart);

		return {
			message: {
				customType: "timestamp-prefix",
				content: "",
				display: true,
				details: { timestamp: now, elapsed },
			},
		};
	});

	// Clean up on shutdown
	pi.on("session_shutdown", async () => {
		if (timerHandle) {
			clearInterval(timerHandle);
			timerHandle = null;
		}
	});
}