Translate latest message using another model

tl.ts

import { mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import type { Api, Model, UserMessage } from "@earendil-works/pi-ai";
import type {
  ExtensionAPI,
  ExtensionCommandContext,
} from "@earendil-works/pi-coding-agent";
import { getAgentDir, getMarkdownTheme } from "@earendil-works/pi-coding-agent";
import { Container, Markdown, Text } from "@earendil-works/pi-tui";

const SYSTEM_PROMPT = [
  "You are a competent engineer, and your job is to rephrase a block of text, as if explaining it to a colleague.",
  "",
  "Rewrite the provided text in simple technical English. Remove dramatic framing, suspense-building, hype, and buzzy metaphors (e.g. 'load-bearing assumption', 'here's the kicker', 'the most instructive part', 'this changes everything'). Plain sentences, no reveals. Keep every technical fact, number, file path, and command exactly intact — only the style changes, not the substance, and do not shorten beyond what removing fluff removes. Do not edit the code in code blocks, but feel free to rewrite code comments. Preserve semantics, so the text stays accurate. You do not have to maintain the structure of the text — reorganize sections, section headers, lists, etc., as needed. Output only the rewritten text with no preamble or additional commentary.",
].join("\n");

interface TlConfig {
  model: string;
}

interface TlTimings {
  idleWaitMs: number;
  authMs: number;
  firstEventMs?: number;
  firstTokenMs?: number;
  totalMs: number;
  thinkingChars: number;
  inputTokens: number;
  outputTokens: number;
  reasoningTokens?: number;
}

interface TranslationData {
  model: string;
  translation: string;
  timestamp: number;
  timings?: TlTimings;
}

function formatMs(ms: number): string {
  return ms < 1000 ? `${Math.round(ms)}ms` : `${(ms / 1000).toFixed(2)}s`;
}

const CONFIG_PATH = join(getAgentDir(), "tl.json");

function loadConfig(): TlConfig | undefined {
  try {
    const parsed = JSON.parse(readFileSync(CONFIG_PATH, "utf8")) as TlConfig;
    return typeof parsed.model === "string" ? parsed : undefined;
  } catch {
    return undefined;
  }
}

function saveConfig(config: TlConfig): void {
  mkdirSync(getAgentDir(), { recursive: true });
  writeFileSync(CONFIG_PATH, `${JSON.stringify(config, null, "\t")}\n`, "utf8");
}

function resolveConfiguredModel(
  ctx: ExtensionCommandContext,
): Model<Api> | undefined {
  const config = loadConfig();
  if (!config) return undefined;
  const slash = config.model.indexOf("/");
  if (slash === -1) return undefined;
  const model = ctx.modelRegistry.find(
    config.model.slice(0, slash),
    config.model.slice(slash + 1),
  );
  if (!model || !ctx.modelRegistry.hasConfiguredAuth(model)) return undefined;
  return model;
}

async function configureModel(
  ctx: ExtensionCommandContext,
): Promise<Model<Api> | undefined> {
  if (!ctx.hasUI) return undefined;
  const models = ctx.modelRegistry
    .getAvailable()
    .filter((model) => ctx.modelRegistry.hasConfiguredAuth(model));
  if (models.length === 0) {
    ctx.ui.notify("No models with configured authentication found", "error");
    return undefined;
  }
  const options = [
    ...new Set(models.map((model) => `${model.provider}/${model.id}`)),
  ].sort();
  const current = loadConfig()?.model;
  const choice = await ctx.ui.select(
    current ? `Translator model (current: ${current})` : "Translator model",
    options,
  );
  if (!choice) return undefined;
  saveConfig({ model: choice });
  ctx.ui.notify(`Translator model set to ${choice}`, "info");
  return resolveConfiguredModel(ctx);
}

function extractText(content: unknown): string {
  if (typeof content === "string") return content.trim();
  if (!Array.isArray(content)) return "";
  const parts: string[] = [];
  for (const block of content) {
    if (
      block &&
      typeof block === "object" &&
      (block as { type?: string }).type === "text"
    ) {
      const text = (block as { text?: unknown }).text;
      if (typeof text === "string") parts.push(text);
    }
  }
  return parts.join("\n\n").trim();
}

function latestAssistantText(ctx: ExtensionCommandContext): string | undefined {
  const branch = ctx.sessionManager.getBranch();
  for (let i = branch.length - 1; i >= 0; i--) {
    const entry = branch[i];
    if (entry.type !== "message" || entry.message.role !== "assistant")
      continue;
    const text = extractText(entry.message.content);
    if (text) return text;
  }
  return undefined;
}

export default function (pi: ExtensionAPI) {
  pi.registerEntryRenderer<TranslationData>(
    "tl-translation",
    (entry, options, theme) => {
      const data = entry.data;
      if (!data) return undefined;
      const container = new Container();
      const t = data.timings;
      let header = `tl · ${data.model}`;
      if (t) {
        header += ` · ${formatMs(t.totalMs)}`;
        if (t.firstTokenMs !== undefined) {
          header += ` (ttft ${formatMs(t.firstTokenMs)})`;
        }
      }
      let headerText = theme.fg("dim", header);
      if (t && ((t.reasoningTokens ?? 0) > 0 || t.thinkingChars > 0)) {
        headerText += theme.fg("warning", " · thinking detected");
      }
      container.addChild(new Text(headerText, 1, 0));
      if (options.expanded && t) {
        const phases = [
          `idle wait ${formatMs(t.idleWaitMs)}`,
          `auth ${formatMs(t.authMs)}`,
        ];
        if (t.firstEventMs !== undefined) {
          phases.push(`first event ${formatMs(t.firstEventMs)}`);
        }
        if (t.firstTokenMs !== undefined) {
          phases.push(`first token ${formatMs(t.firstTokenMs)}`);
        }
        phases.push(`total ${formatMs(t.totalMs)}`);
        container.addChild(new Text(theme.fg("dim", phases.join(" · ")), 1, 0));
        const tokens = [`in ${t.inputTokens}`, `out ${t.outputTokens}`];
        if (t.reasoningTokens !== undefined) {
          tokens.push(`reasoning ${t.reasoningTokens}`);
        }
        if (t.thinkingChars > 0) {
          tokens.push(`thinking chars ${t.thinkingChars}`);
        }
        container.addChild(
          new Text(theme.fg("dim", `tokens: ${tokens.join(" · ")}`), 1, 0),
        );
      }
      container.addChild(
        new Markdown(data.translation, 1, 0, getMarkdownTheme()),
      );
      return container;
    },
  );

  pi.registerCommand("tl", {
    description:
      "Rephrase the latest assistant message (usage: /tl or /tl config)",
    getArgumentCompletions: (prefix) =>
      "config".startsWith(prefix.toLowerCase())
        ? [
            {
              value: "config",
              label: "config",
              description: "Select the translator model",
            },
          ]
        : null,
    handler: async (args, ctx) => {
      if (args.trim().toLowerCase() === "config") {
        await configureModel(ctx);
        return;
      }

      let model = resolveConfiguredModel(ctx);
      if (!model) {
        model = await configureModel(ctx);
      }
      if (!model) {
        if (ctx.hasUI) ctx.ui.notify("No translator model configured", "error");
        return;
      }

      const idleStart = performance.now();
      await ctx.waitForIdle();
      const idleWaitMs = performance.now() - idleStart;

      const source = latestAssistantText(ctx);
      if (!source) {
        ctx.ui.notify("No assistant message to translate", "warning");
        return;
      }

      const modelName = `${model.provider}/${model.id}`;
      ctx.ui.setStatus("tl", `translating via ${modelName}...`);
      try {
        const provider = ctx.modelRegistry.getProvider(model.provider);
        if (!provider) {
          ctx.ui.notify(`Provider not found: ${model.provider}`, "error");
          return;
        }
        const authStart = performance.now();
        const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
        const authMs = performance.now() - authStart;
        if (!auth.ok) {
          ctx.ui.notify(`Translation failed: ${auth.error}`, "error");
          return;
        }
        const message: UserMessage = {
          role: "user",
          content: [{ type: "text", text: source }],
          timestamp: Date.now(),
        };
        const requestStart = performance.now();
        const stream = provider.streamSimple(
          auth.baseUrl ? { ...model, baseUrl: auth.baseUrl } : model,
          {
            systemPrompt: SYSTEM_PROMPT,
            messages: [message],
          },
          {
            apiKey: auth.apiKey,
            headers: auth.headers,
            env: auth.env,
            reasoning:
              model.api === "openai-codex-responses" ? "minimal" : undefined,
          },
        );
        let firstEventMs: number | undefined;
        let firstTokenMs: number | undefined;
        let thinkingChars = 0;
        for await (const event of stream) {
          const elapsed = performance.now() - requestStart;
          if (firstEventMs === undefined) firstEventMs = elapsed;
          if (event.type === "text_delta" || event.type === "thinking_delta") {
            if (firstTokenMs === undefined) firstTokenMs = elapsed;
            if (event.type === "thinking_delta") {
              thinkingChars += event.delta.length;
            }
          }
        }
        const response = await stream.result();
        const totalMs = performance.now() - requestStart;
        if (
          response.stopReason === "error" ||
          response.stopReason === "aborted"
        ) {
          ctx.ui.notify(
            `Translation failed: ${response.errorMessage ?? response.stopReason}`,
            "error",
          );
          return;
        }
        const translation = extractText(response.content);
        if (!translation) {
          ctx.ui.notify("Translator returned no text", "warning");
          return;
        }
        pi.appendEntry<TranslationData>("tl-translation", {
          model: modelName,
          translation,
          timestamp: Date.now(),
          timings: {
            idleWaitMs,
            authMs,
            firstEventMs,
            firstTokenMs,
            totalMs,
            thinkingChars,
            inputTokens: response.usage.input,
            outputTokens: response.usage.output,
            reasoningTokens: response.usage.reasoning,
          },
        });
      } catch (error) {
        ctx.ui.notify(
          `Translation failed: ${error instanceof Error ? error.message : String(error)}`,
          "error",
        );
      } finally {
        ctx.ui.setStatus("tl", undefined);
      }
    },
  });
}