Agent-JAE/default-extensions/replay.ts

import type { ExtensionAPI } from "@jaeswift/jae-coding-agent";
import { writeFileSync } from "node:fs";
import { join } from "node:path";

interface ReplayEvent {
  time: number;
  type: string;
  text: string;
}

export default function (pi: ExtensionAPI) {
  let recording = false;
  let events: ReplayEvent[] = [];
  let startTime = 0;

  pi.on("message_end", async (event, ctx) => {
    if (!recording) return;
    const msg = event.message;
    if (!msg) return;

    let text = "";
    if ("content" in msg && typeof msg.content === "string") {
      text = msg.content;
    } else if ("content" in msg && Array.isArray(msg.content)) {
      text = (msg.content as any[])
        .filter((b: any) => b.type === "text")
        .map((b: any) => b.text)
        .join("");
    }

    if (text) {
      events.push({
        time: (Date.now() - startTime) / 1000,
        type: (msg as any).role || "unknown",
        text,
      });
    }
  });

  pi.on("tool_execution_end", async (event, _ctx) => {
    if (!recording) return;
    const resultText =
      typeof event.result === "string"
        ? event.result
        : JSON.stringify(event.result).substring(0, 500);
    events.push({
      time: (Date.now() - startTime) / 1000,
      type: `tool:${event.toolName}`,
      text: resultText,
    });
  });

  pi.registerCommand("replay", {
    description: "Session replay: /replay start | /replay stop | /replay export [file.cast]",
    handler: async (args, ctx) => {
      const parts = args.trim().split(/\s+/);
      const sub = parts[0] || "status";

      if (sub === "start") {
        if (recording) {
          ctx.ui.notify("Already recording.", "warning");
          return;
        }
        recording = true;
        events = [];
        startTime = Date.now();
        ctx.ui.setStatus("replay", "\u{1F534} REC");
        ctx.ui.notify("\u{1F3AC} Replay recording started.", "info");
        return;
      }

      if (sub === "stop") {
        if (!recording) {
          ctx.ui.notify("Not recording.", "warning");
          return;
        }
        recording = false;
        ctx.ui.setStatus("replay", undefined);
        ctx.ui.notify(`\u{1F3AC} Recording stopped. ${events.length} events captured. Use /replay export to save.`, "info");
        return;
      }

      if (sub === "export") {
        if (events.length === 0) {
          ctx.ui.notify("No events recorded. Start a recording first.", "warning");
          return;
        }
        const filename = parts[1] || `jae-replay-${Date.now()}.cast`;
        const outPath = join(ctx.cwd, filename);

        // asciicast v2 format
        const header = JSON.stringify({
          version: 2,
          width: 120,
          height: 40,
          timestamp: Math.floor(startTime / 1000),
          title: "JAE Session Replay",
          env: { TERM: "xterm-256color" },
        });

        const lines = [header];
        for (const ev of events) {
          const prefix = ev.type === "assistant" ? "\x1b[36m" : ev.type.startsWith("tool:") ? "\x1b[33m" : "\x1b[37m";
          const label = ev.type === "assistant" ? "JAE" : ev.type === "user" ? "YOU" : ev.type;
          const displayText = `${prefix}[${label}]\x1b[0m ${ev.text}\r\n`;
          lines.push(JSON.stringify([ev.time, "o", displayText]));
        }

        writeFileSync(outPath, lines.join("\n") + "\n", "utf-8");
        ctx.ui.notify(`\u{1F3AC} Exported ${events.length} events to ${outPath}\nPlay with: asciinema play ${filename}`, "info");
        return;
      }

      if (sub === "status") {
        ctx.ui.notify(
          recording
            ? `\u{1F534} Recording in progress: ${events.length} events (${((Date.now() - startTime) / 1000).toFixed(0)}s)`
            : `\u{1F3AC} Not recording. ${events.length} events in buffer.`,
          "info",
        );
        return;
      }

      ctx.ui.notify("Usage: /replay start | stop | export [file.cast] | status", "info");
    },
  });
}