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, ExtensionContext } 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: ExtensionContext) => {
    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 legacy timestamp-suffix messages (from old sessions)
  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);
      const timeStr = formatTime(new Date(now));
      let line = `○ ${timeStr} (+${elapsed})`;
      if (lastTurnDuration > 1000) {
        line += ` • took ${formatDuration(lastTurnDuration)}`;
      }
      ctx.ui.notify(line, "info");
    }
  });

  // Strip old timestamp-suffix messages from LLM context
  pi.on("context", async (event) => {
    const messages = event.messages.filter(
      (m: any) =>
        m.role !== "custom" ||
        (m.customType !== "timestamp-prefix" &&
          m.customType !== "timestamp-suffix"),
    );
    return { messages };
  });

  // Inject timestamp marker right before agent starts (inline, persistent)
  pi.on("before_agent_start", async () => {
    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;
    }
  });
}