MailUI4Agents/plugins/pi-mail-bridge/extension/index.ts

/**
 * pi 交互式会话的邮件工具扩展。
 *
 * # 与 src/index.mjs（常驻守护进程）的分工
 *
 *   - 守护进程负责**收信**：订阅 SSE、fork worker、跑模型、自动回信
 *   - 本扩展只给**交互式 TUI 会话**装上同一套邮件工具
 *
 * 两者是同一条 AgentMail 身份（agent `pi`）的两个入口，与 DSH 的
 * 「TUI + 邮箱是同一个 Agent」完全一致。人可以在 TUI 里直接收发邮件，
 * 也可以在邮箱里给 pi 发信 —— 两边看到的是同一条会话流。
 *
 * # 为什么必须单独做这个扩展
 *
 * 守护进程用 `noExtensions: true` 起会话（见 src/session-pool.mjs），
 * 所以它的邮件工具**不会**出现在人的交互式 pi 里。而平台的建设者与维护者
 * 恰恰是在交互式 pi 里工作的：没有 send_mail / read_inbox，他既看不到
 * 自己刚发出的那封信，也无法回信，只能绕到 curl + 密钥直连 Gateway ——
 * 一个「邮件驱动」的平台，维护者自己收不到邮件。
 *
 * # 密钥从哪来
 *
 * 交互式 pi 的进程环境里通常**没有** AGENTMAIL_* 变量（守护进程的
 * EnvironmentFile 只注入给它自己）。因此按以下顺序解析：
 *
 *   1. 进程环境（`AGENTMAIL_AGENT_KEY` / `AGENTMAIL_GATEWAY_URL` / `AGENTMAIL_AGENT_NAME`）
 *   2. `AGENTMAIL_ENV_FILE`（默认 `/etc/agentmail/pi.env`）—— 部署时的权威来源，
 *      与守护进程用的是同一把密钥，因此身份一致
 *   3. `$AGENTMAIL_CONFIG_DIR/agent.key` 或 `~/.agentmail/agent.key`
 *      （插件的历史约定；接受 `key` 与 `key_token` 两种字段名）
 *
 * 找不到密钥时**不注册任何工具**并明确告知 —— 挂上一组永远 401 的工具比没有更糟：
 * 模型会以为自己能发信，人却只看到一串认证失败。
 *
 * # 为什么不注册 connect_to_server
 *
 * 那个工具会重写 Gateway 坐标并重新登记密钥。在交互式会话里调用它会**动到守护
 * 进程的配置**（两者共用同一把密钥与同一个 Agent 名），而守护进程才是真正长期
 * 持有 SSE 长连的一方。坐标变更属于部署动作，不该由一次 TUI 对话触发。
 */

import { existsSync, readFileSync } from 'node:fs';
import { homedir } from 'node:os';
import { join } from 'node:path';
import type { ExtensionAPI } from '@earendil-works/pi-coding-agent';

const DEFAULT_ENV_FILE = '/etc/agentmail/pi.env';
const DEFAULT_GATEWAY = 'http://127.0.0.1:8180';

/** 守护进程专用的生命周期工具：交互式会话里不注册（理由见文件头注释）。 */
const EXCLUDED_TOOLS = new Set(['connect_to_server']);

/**
 * registerTool 期望的完整定义形状。
 *
 * 用 SDK 自己的类型取，而不是复制一份：SDK 升级时形状变化会在这里变成编译错误，
 * 而不是运行时的静默错配。
 */
type MailTool = Parameters<ExtensionAPI['registerTool']>[0];

/** 极简 .env 解析：只认 `KEY=value`，忽略注释与空行，不处理引号与转义。 */
function parseEnvFile(path: string): Record<string, string> {
  const out: Record<string, string> = {};
  if (!existsSync(path)) return out;
  let raw: string;
  try {
    raw = readFileSync(path, 'utf8');
  } catch {
    return out;
  }
  for (const line of raw.split('\n')) {
    const trimmed = line.trim();
    if (!trimmed || trimmed.startsWith('#')) continue;
    const eq = trimmed.indexOf('=');
    if (eq <= 0) continue;
    const key = trimmed.slice(0, eq).trim();
    // 值里可能有 `=`（例如 base64），只切第一个
    out[key] = trimmed.slice(eq + 1).trim();
  }
  return out;
}

/** 读插件约定的本地密钥文件，兼容 `key` 与 `key_token` 两种字段名。 */
function readLocalKey(configDir?: string): string {
  const candidates = [
    join(configDir || join(homedir(), '.agentmail'), 'agent.key'),
    join(homedir(), '.agentmail', 'agent.key'),
  ];
  for (const path of candidates) {
    try {
      if (!existsSync(path)) continue;
      const raw = JSON.parse(readFileSync(path, 'utf8')) as
        { key?: unknown; key_token?: unknown };
      // 历史文件用的是 key_token（见部署实况），只认 `key` 会静默读不到。
      const key = raw?.key ?? raw?.key_token;
      if (typeof key === 'string' && key.trim()) return key.trim();
    } catch {
      // 读不了就试下一个
    }
  }
  return '';
}

interface MailConfig {
  url: string;
  agentName: string;
  agentKey: string;
}

/** @returns null 表示拿不到密钥（调用方据此不注册任何工具） */
function resolveConfig(): MailConfig | null {
  let url = process.env.AGENTMAIL_GATEWAY_URL || '';
  let agentName = process.env.AGENTMAIL_AGENT_NAME || '';
  let agentKey = process.env.AGENTMAIL_AGENT_KEY || '';

  if (!agentKey) {
    const envFile = process.env.AGENTMAIL_ENV_FILE || DEFAULT_ENV_FILE;
    const parsed = parseEnvFile(envFile);
    agentKey = parsed.AGENTMAIL_AGENT_KEY || '';
    url = url || parsed.AGENTMAIL_GATEWAY_URL || '';
    agentName = agentName || parsed.AGENTMAIL_AGENT_NAME || '';
  }

  if (!agentKey) agentKey = readLocalKey(process.env.AGENTMAIL_CONFIG_DIR);
  if (!agentKey) return null;

  return {
    url: (url || DEFAULT_GATEWAY).replace(/\/+$/, ''),
    agentName: agentName || 'pi',
    agentKey,
  };
}

export default async function (pi: ExtensionAPI): Promise<void> {
  const cfg = resolveConfig();

  // 动态 import 仓库里的 .mjs：它们没有 .d.ts（也不该造一份 —— 守护进程用同一份
  // 对象经 customTools 注册，额外声明只会两边不同步），因此在边界上收窄成我们
  // 真正依赖的最小形状。工具对象的 `parameters` 是普通 JSON Schema，而
  // registerTool 的类型要求 TypeBox 的 TSchema；两者在运行期是同一套 JSON Schema
  // 校验（守护进程已长期验证），差异只在静态类型上，故此处一次收窄并注明。
  let registered = 0;
  let notice = '';

  if (!cfg) {
    notice = 'pi-mail-bridge：未找到 AgentMail 密钥（AGENTMAIL_AGENT_KEY / '
      + `${process.env.AGENTMAIL_ENV_FILE || DEFAULT_ENV_FILE} / ~/.agentmail/agent.key），`
      + '邮件工具未注册。';
  } else {
    const gatewayMod = await import(new URL('../src/gateway.mjs', import.meta.url).href);
    const toolsMod = await import(new URL('../src/tools.mjs', import.meta.url).href);
    const GatewayClient = gatewayMod.GatewayClient as new (opts: {
      url: string; agentName: string; agentKey: string; agentSecret: string;
    }) => unknown;
    const createMailTools = toolsMod.createMailTools as (opts: {
      client: unknown; agentName: string; log: (msg: string) => void;
    }) => MailTool[];

    const tools = createMailTools({
      client: new GatewayClient({
        url: cfg.url,
        agentName: cfg.agentName,
        agentKey: cfg.agentKey,
        agentSecret: '',
      }),
      agentName: cfg.agentName,
      // 出错只进 stderr（不污染 TUI）；正常路径本身不吵闹。
      log: (msg: string) => console.error(`[pi-mail-tools] ${msg}`),
    });

    for (const tool of tools) {
      if (!tool?.name || EXCLUDED_TOOLS.has(tool.name)) continue;
      pi.registerTool(tool);
      registered += 1;
    }
  }

  // 只在 session_start 里报：扩展工厂可能在「不会开会话」的调用里执行
  // （例如 pi --list-models），那里没有 ctx，也没有界面可显示。
  pi.on('session_start', async (_event, ctx) => {
    if (notice) {
      ctx.ui.notify(notice, 'warning');
      return;
    }
    ctx.ui.setStatus('pi-mail', `邮件已接入（${cfg?.agentName}@${cfg?.url}，${registered} 个工具）`);
  });
}