mnote/packages/pi-mnote/extensions/mnote-bridge.ts

import fs from "node:fs";
import path from "node:path";
import { fileURLToPath } from "node:url";

type ExtensionContext = {
  hasUI?: boolean;
  ui?: {
    confirm?: (title: string, message: string, options?: Record<string, unknown>) => Promise<boolean>;
  };
};

type ExtensionAPI = {
  registerTool: (spec: Record<string, unknown>) => void;
  on?: (
    eventName: string,
    handler: (event: Record<string, unknown>) => Record<string, unknown> | void,
  ) => void;
};

type MnoteToolManifestEntry = {
  piName: string;
  mnoteName: string;
  label: string;
  description: string;
};

const DEFAULT_TOOLS: MnoteToolManifestEntry[] = [
  {
    piName: "mnote_current_page_read",
    mnoteName: "mnote.current_page.read",
    label: "MNote current page read",
    description: "Read the current MNote page through MNote access scope.",
  },
  {
    piName: "mnote_selection_read",
    mnoteName: "mnote.selection.read",
    label: "MNote selection read",
    description: "Read the current MNote editor selection snapshot supplied by MNote.",
  },
  {
    piName: "mnote_allowed_roots_describe",
    mnoteName: "mnote.allowed_roots.describe",
    label: "MNote allowed roots describe",
    description: "Describe MNote allowed roots and disabled raw tools.",
  },
  {
    piName: "mnote_local_file_read",
    mnoteName: "mnote.local_file.read",
    label: "MNote local file read",
    description: "Read a file only through MNote allowed roots.",
  },
  {
    piName: "mnote_local_file_patch",
    mnoteName: "mnote.local_file.patch",
    label: "MNote local file patch",
    description: "Patch a file only through MNote allowed roots and watcher refresh.",
  },
  {
    piName: "mnote_knowledge_rag_status",
    mnoteName: "mnote.knowledge_rag.status",
    label: "MNote LightRAG status",
    description: "Check the MNote LightRAG knowledge provider status, indexed source registry, dashboard URL, and sync state. Use before knowledge-base questions when availability is uncertain. Params: optional workspaceId/rootUri; MNote fills the current Pi session context when omitted.",
  },
  {
    piName: "mnote_knowledge_rag_query",
    mnoteName: "mnote.knowledge_rag.query",
    label: "MNote LightRAG query",
    description: "Ask the MNote LightRAG knowledge library across indexed books, papers, Office files, PDFs, images, and attachments. Use this for knowledge-base questions and answers that require sources, citations, or evidence. For book or long-document questions, pass query, mode='naive' or 'mix', topK, chunkTopK, includeChunkContent=true, and includeDocumentStructureIndex=true. Use returned references/citations quotes as evidence; do not invent page numbers or hand-write /documents/mnote:// links.",
  },
  {
    piName: "mnote_knowledge_rag_section_context",
    mnoteName: "mnote.knowledge_rag.section_context",
    label: "MNote LightRAG section context",
    description: "Read bounded section blocks/chunks from a LightRAG sidecar using documentStructureIndex ranges returned by mnote_knowledge_rag_query. Use this for second-pass reading of large books or long documents when query references are not enough. Params include sourcePath/sourceId/lightRagDocId/filePath/sectionId, block or paragraph ordinal range, contextBefore/contextAfter, maxBlocks, maxChars.",
  },
  {
    piName: "mnote_knowledge_rag_open_reference",
    mnoteName: "mnote.knowledge_rag.open_reference",
    label: "MNote LightRAG open reference",
    description: "Convert a LightRAG reference, filePath, or chunkId returned by mnote_knowledge_rag_query into a MNote clickable local resource locator. Use when the user asks to open or verify a cited source.",
  },
  {
    piName: "mnote_reference_open",
    mnoteName: "mnote.reference.open",
    label: "MNote reference open",
    description: "Legacy alias for opening a citation/reference through MNote mapping. Prefer mnote_knowledge_rag_open_reference for LightRAG references.",
  },
  {
    piName: "mnote_codex_rescue_request",
    mnoteName: "mnote.codex_rescue.request",
    label: "MNote Codex rescue",
    description: "Ask local Codex to rescue hard MNote/Pi problems before escalating to the user. Use when tools, skills, MCP, LightRAG, environment, or local repo behavior looks broken and normal Pi troubleshooting is insufficient. This tool is admin-gated, approval-gated by default, runs codex exec with workspace-write sandbox and a timeout, and returns Codex's final answer plus stdout/stderr snippets. Provide issue, evidence/logs, attempted steps, and desired outcome. Call at most once per unresolved incident; if Codex cannot fix it, summarize the blocker to the user.",
  },
  {
    piName: "mnote_tool_receipt_write",
    mnoteName: "mnote.tool_receipt.write",
    label: "MNote tool receipt write",
    description: "Write a provider-neutral MNote tool receipt.",
  },
];

const EXTENSION_DIR = path.dirname(fileURLToPath(import.meta.url));
const DEFAULT_CONTEXT_FILE = path.join(EXTENSION_DIR, "mnote-context.json");
const CONTEXT_FILE = env("PI_MNOTE_CONTEXT_FILE")
  || (fs.existsSync(DEFAULT_CONTEXT_FILE) ? DEFAULT_CONTEXT_FILE : "");
const RUNTIME_IMPL = env("PI_MNOTE_RUNTIME_IMPL")
  || env("MNOTE_PI_RUNTIME_IMPL")
  || (CONTEXT_FILE ? "pi-rust" : "");
const BASE_URL = env("PI_MNOTE_BRIDGE_BASE_URL") || env("MNOTE_PI_BRIDGE_BASE_URL") || env("MNOTE_PI_LAB_BASE_URL") || "http://127.0.0.1:3000";
const SESSION_ID = env("PI_MNOTE_BRIDGE_SESSION_ID") || env("MNOTE_PI_BRIDGE_SESSION_ID") || env("MNOTE_PI_LAB_SESSION_ID") || "";
const BRIDGE_TOKEN = env("PI_MNOTE_BRIDGE_TOKEN") || env("MNOTE_PI_BRIDGE_TOKEN") || env("MNOTE_PI_LAB_BRIDGE_TOKEN") || "";
const HTTP_BRIDGE_AVAILABLE = env("PI_MNOTE_HTTP_BRIDGE_AVAILABLE") === "1";
const TOOL_POLICIES = parseJsonRecord(env("PI_MNOTE_BRIDGE_TOOL_POLICIES") || env("MNOTE_PI_BRIDGE_TOOL_POLICIES"), {});
const TOOLS = normalizeTools(parseJsonUnknown(env("PI_MNOTE_BRIDGE_TOOLS") || env("MNOTE_PI_BRIDGE_TOOLS")) ?? DEFAULT_TOOLS);
const CONTEXT_PREFIX = "[[MNOTE_PI_CONTEXT_V1:";
let liveContext: Record<string, unknown> | undefined;

function env(name: string): string {
  return (process.env[name] || "").trim();
}

function parseJsonUnknown(raw: string): unknown {
  if (!raw) return undefined;
  try {
    return JSON.parse(raw);
  } catch {
    return undefined;
  }
}

function parseJsonRecord(raw: string, fallback: Record<string, unknown>): Record<string, unknown> {
  const parsed = parseJsonUnknown(raw);
  return parsed && typeof parsed === "object" && !Array.isArray(parsed)
    ? parsed as Record<string, unknown>
    : fallback;
}

function normalizeTools(value: unknown): MnoteToolManifestEntry[] {
  if (!Array.isArray(value)) return DEFAULT_TOOLS;
  const tools = value
    .map((entry) => {
      if (!entry || typeof entry !== "object") return undefined;
      const record = entry as Record<string, unknown>;
      const piName = stringField(record, "piName") || stringField(record, "pi_name");
      const mnoteName = stringField(record, "mnoteName") || stringField(record, "mnote_name");
      const label = stringField(record, "label") || piName;
      const description = stringField(record, "description") || label;
      if (!piName || !mnoteName) return undefined;
      return { piName, mnoteName, label, description };
    })
    .filter((entry): entry is MnoteToolManifestEntry => Boolean(entry));
  return tools.length ? tools : DEFAULT_TOOLS;
}

function stringField(record: Record<string, unknown>, key: string): string {
  const value = record[key];
  return typeof value === "string" ? value.trim() : "";
}

function toolResult(payload: Record<string, unknown>, isError = false) {
  return {
    content: [{ type: "text", text: JSON.stringify(payload, null, 2) }],
    details: payload,
    isError,
  };
}

function decodeHexJson(value: string): Record<string, unknown> {
  if (!value || value.length % 2 !== 0 || !/^[0-9a-f]+$/i.test(value)) {
    throw new Error("MNote Pi context hex 无效");
  }
  let jsonText = "";
  for (let index = 0; index < value.length; index += 2) {
    jsonText += String.fromCharCode(Number.parseInt(value.slice(index, index + 2), 16));
  }
  const parsed = JSON.parse(decodeURIComponent(escape(jsonText)));
  if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
    throw new Error("MNote Pi context payload 格式无效");
  }
  return parsed as Record<string, unknown>;
}

function captureInputContext(event: Record<string, unknown>) {
  const text = typeof event.text === "string" ? event.text : "";
  const start = text.indexOf(CONTEXT_PREFIX);
  if (start < 0) return { action: "continue" };
  const end = text.indexOf("]]\n", start);
  if (end < 0) return { action: "continue" };
  const encoded = text.slice(start + CONTEXT_PREFIX.length, end);
  try {
    liveContext = decodeHexJson(encoded);
  } catch (error) {
    liveContext = {
      contextError: error instanceof Error ? error.message : String(error),
    };
  }
  return {
    action: "transform",
    text: `${text.slice(0, start)}${text.slice(end + 3)}`,
    images: event.images,
  };
}

function readContextSnapshot(): Record<string, unknown> {
  if (liveContext) return liveContext;
  if (!CONTEXT_FILE) {
    throw new Error("PI_MNOTE_CONTEXT_FILE 未配置");
  }
  const raw = fs.readFileSync(CONTEXT_FILE, "utf8");
  const parsed = JSON.parse(raw);
  if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
    throw new Error("MNote Pi 上下文文件格式无效");
  }
  return parsed as Record<string, unknown>;
}

function contextAllowedRoots(context: Record<string, unknown>): Record<string, unknown>[] {
  const snapshot = context.allowedRoots;
  if (!snapshot || typeof snapshot !== "object" || Array.isArray(snapshot)) return [];
  const roots = (snapshot as Record<string, unknown>).roots;
  if (!Array.isArray(roots)) return [];
  return roots.filter((root): root is Record<string, unknown> => (
    Boolean(root) && typeof root === "object" && !Array.isArray(root)
  ));
}

function pathIsInside(target: string, root: string): boolean {
  const relative = path.relative(root, target);
  return relative === "" || (!relative.startsWith("..") && !path.isAbsolute(relative));
}

function resolveAllowedTarget(
  context: Record<string, unknown>,
  rootUri: string,
  relativePath: string,
): { target: string; readPath: string } {
  const roots = contextAllowedRoots(context);
  const matchedRoot = roots.find((root) => stringField(root, "rootUri") === rootUri);
  const rootPath = matchedRoot
    ? stringField(matchedRoot, "rootPath")
    : stringField(context, "primaryRootPath") || process.cwd();
  if (!rootPath) throw new Error("MNote allowed root 缺少 rootPath");

  const canonicalRoot = fs.realpathSync(rootPath);
  const requestedTarget = path.isAbsolute(relativePath)
    ? path.resolve(relativePath)
    : path.resolve(canonicalRoot, relativePath);
  const canonicalTarget = fs.realpathSync(requestedTarget);
  const allowedRoots = roots
    .map((root) => stringField(root, "rootPath"))
    .filter(Boolean)
    .map((root) => fs.realpathSync(root));
  const candidates = allowedRoots.length ? allowedRoots : [canonicalRoot];
  if (!candidates.some((root) => pathIsInside(canonicalTarget, root))) {
    throw new Error(`路径超出 MNote allowed roots: ${relativePath}`);
  }
  const processRoot = fs.realpathSync(process.cwd());
  const readPath = pathIsInside(canonicalTarget, processRoot)
    ? path.relative(processRoot, canonicalTarget) || "."
    : canonicalTarget;
  return { target: canonicalTarget, readPath };
}

function executeNativeCurrentPageRead(params: unknown) {
  try {
    const input = params && typeof params === "object" && !Array.isArray(params)
      ? params as Record<string, unknown>
      : {};
    const context = readContextSnapshot();
    const rootUri = stringField(context, "rootUri")
      || stringField(input, "rootUri")
      || stringField(input, "root_uri");
    const pagePath = stringField(context, "pagePath")
      || stringField(input, "pagePath")
      || stringField(input, "page_path")
      || stringField(input, "path");
    if (!rootUri) throw new Error("读取当前页缺少 rootUri");
    if (!pagePath) throw new Error("读取当前页缺少 pagePath");

    const resolved = resolveAllowedTarget(context, rootUri, pagePath);
    const content = fs.readFileSync(resolved.readPath, "utf8");
    const stat = fs.statSync(resolved.readPath);
    return toolResult({
      ok: true,
      rootUri,
      pagePath,
      path: resolved.target,
      content,
      contentLength: content.length,
      format: "markdown",
      fileVersion: `pi-rust-native-${Math.trunc(stat.mtimeMs)}-${stat.size}`,
      transport: "pi-rust-native-fs",
      contextFile: CONTEXT_FILE,
    });
  } catch (error) {
    return toolResult({
      ok: false,
      code: "mnote_pi_rust_current_page_read_failed",
      message: error instanceof Error ? error.message : String(error),
      transport: "pi-rust-native-fs",
      contextFile: CONTEXT_FILE,
    }, true);
  }
}

function executeNativeSelectionRead() {
  try {
    const context = readContextSnapshot();
    return toolResult({
      ok: true,
      selection: context.selectedContext ?? null,
      selectionSource: "mnote_sidebar_host_snapshot",
      rootUri: context.rootUri ?? null,
      pagePath: context.pagePath ?? null,
      contextRefs: context.contextRefs ?? [],
      transport: "pi-rust-native-context-file",
    });
  } catch (error) {
    return toolResult({
      ok: false,
      code: "mnote_pi_rust_selection_read_failed",
      message: error instanceof Error ? error.message : String(error),
      transport: "pi-rust-native-context-file",
    }, true);
  }
}

function executeNativeAllowedRootsDescribe() {
  try {
    const context = readContextSnapshot();
    return toolResult({
      ok: true,
      allowedRoots: context.allowedRoots ?? { roots: [] },
      rootUri: context.rootUri ?? null,
      pagePath: context.pagePath ?? null,
      primaryRootPath: context.primaryRootPath ?? process.cwd(),
      transport: "pi-rust-native-context-file",
    });
  } catch (error) {
    return toolResult({
      ok: false,
      code: "mnote_pi_rust_allowed_roots_read_failed",
      message: error instanceof Error ? error.message : String(error),
      transport: "pi-rust-native-context-file",
    }, true);
  }
}

function executeNativeTool(tool: MnoteToolManifestEntry, params: unknown) {
  if (runtimeImplementation() !== "pi-rust" || toolPolicy(tool.mnoteName) !== "allow") return undefined;
  if (tool.mnoteName === "mnote.current_page.read") return executeNativeCurrentPageRead(params);
  if (tool.mnoteName === "mnote.selection.read") return executeNativeSelectionRead();
  if (tool.mnoteName === "mnote.allowed_roots.describe") return executeNativeAllowedRootsDescribe();
  return undefined;
}

async function callMnote(toolName: string, params: unknown) {
  const response = await fetch(`${BASE_URL}/api/page-ai/pi/tool-call-bridge`, {
    method: "POST",
    headers: {
      "content-type": "application/json",
      "x-mnote-pi-lab-bridge-token": BRIDGE_TOKEN,
    },
    body: JSON.stringify({ sessionId: SESSION_ID, toolName, params: params || {} }),
  });
  const payload = await response.json().catch(() => ({ ok: false, code: "bad_json" }));
  const text = JSON.stringify((payload as Record<string, unknown>).result || payload, null, 2);
  return {
    content: [{ type: "text", text }],
    details: payload,
  };
}

function stableJson(value: unknown): string {
  if (value === null || typeof value !== "object") return JSON.stringify(value);
  if (Array.isArray(value)) return `[${value.map(stableJson).join(",")}]`;
  const record = value as Record<string, unknown>;
  return `{${Object.keys(record).sort().map((key) => `${JSON.stringify(key)}:${stableJson(record[key])}`).join(",")}}`;
}

function stableHash(value: string): string {
  let hash = 5381 >>> 0;
  for (let index = 0; index < value.length; index += 1) {
    hash = (((hash << 5) + hash) + value.charCodeAt(index)) >>> 0;
  }
  return hash.toString(16);
}

function paramsHash(params: unknown): string {
  const copy = { ...((params || {}) as Record<string, unknown>) };
  delete copy.mnoteApproval;
  delete copy.mnote_approval;
  return stableHash(stableJson(copy));
}

function toolPolicy(toolName: string): string {
  const contextPolicies = liveContext?.toolPolicies;
  const contextValue = contextPolicies && typeof contextPolicies === "object" && !Array.isArray(contextPolicies)
    ? (contextPolicies as Record<string, unknown>)[toolName]
    : undefined;
  const value = contextValue ?? TOOL_POLICIES[toolName];
  return typeof value === "string" && value.trim() ? value.trim() : "allow";
}

function runtimeImplementation(): string {
  return liveContext ? stringField(liveContext, "runtimeImplementation") || RUNTIME_IMPL : RUNTIME_IMPL;
}

async function requestMnoteApprovalViaBridge(request: Record<string, unknown>): Promise<Record<string, unknown>> {
  const response = await fetch(`${BASE_URL}/api/page-ai/pi/ui-request-bridge`, {
    method: "POST",
    headers: {
      "content-type": "application/json",
      "x-mnote-pi-lab-bridge-token": BRIDGE_TOKEN,
    },
    body: JSON.stringify({ sessionId: SESSION_ID, ...request }),
  });
  return await response.json().catch(() => ({ ok: false, cancelled: true, code: "bad_json" }));
}

async function requestMnoteApproval(
  ctx: ExtensionContext | undefined,
  toolCallId: string,
  toolName: string,
  label: string,
  params: unknown,
): Promise<null | Record<string, unknown>> {
  if (toolPolicy(toolName) !== "ask") return null;
  const approvalId = toolCallId || `approval_${Date.now()}_${Math.random().toString(16).slice(2)}`;
  const approval = { approvalId, toolName, paramsHash: paramsHash(params) };
  const summary = JSON.stringify(params || {}, null, 2).slice(0, 1200);
  const message = `${label}\n${toolName}\n\n${summary}`;

  if (ctx?.hasUI && ctx.ui && typeof ctx.ui.confirm === "function") {
    const result = await ctx.ui.confirm("审批 Pi 工具调用", message, { timeout: 60000 });
    if (result === true) return approval;
    return { ...approval, denied: true };
  }

  const result = await requestMnoteApprovalViaBridge({
    id: approvalId,
    method: "confirm",
    title: "审批 Pi 工具调用",
    message,
    timeoutMs: 60000,
    mnoteApproval: approval,
  });
  if (result && result.confirmed === true) return approval;
  return { ...approval, denied: true };
}

function register(pi: ExtensionAPI, tool: MnoteToolManifestEntry) {
  pi.registerTool({
    name: tool.piName,
    label: tool.label,
    description: tool.description,
    promptSnippet: `${tool.label}: ${tool.description}`,
    parameters: { type: "object", additionalProperties: true },
    execute(toolCallId: unknown, params: unknown, _signal: unknown, _onUpdate: unknown, ctx: ExtensionContext | undefined) {
      if (toolPolicy(tool.mnoteName) === "deny") {
        return toolResult({
          ok: false,
          code: "mnote_tool_denied_by_permission_mode",
          message: `当前 Pi 模式禁止调用工具 ${tool.mnoteName}`,
          toolName: tool.mnoteName,
        }, true);
      }
      const nativeResult = executeNativeTool(tool, params);
      if (nativeResult) return nativeResult;
      if (runtimeImplementation() === "pi-rust" && !HTTP_BRIDGE_AVAILABLE) {
        return toolResult({
          ok: false,
          code: "mnote_pi_rust_service_bridge_unavailable",
          message: `${tool.mnoteName} 仍依赖旧 HTTP bridge；Pi Rust 当前应通过原生文件工具或专用 MCP 扩展调用。`,
          toolName: tool.mnoteName,
          runtimeImplementation: runtimeImplementation(),
        }, true);
      }
      return executeBridgeTool(toolCallId, params, ctx, tool);
    },
  });
}

async function executeBridgeTool(
  toolCallId: unknown,
  params: unknown,
  ctx: ExtensionContext | undefined,
  tool: MnoteToolManifestEntry,
) {
      const nextParams = { ...((params || {}) as Record<string, unknown>) };
      const approval = await requestMnoteApproval(ctx, String(toolCallId || ""), tool.mnoteName, tool.label, nextParams);
      if (approval && approval.denied) {
        return {
          content: [{ type: "text", text: JSON.stringify({ ok: false, code: "mnote_tool_approval_cancelled", message: "用户拒绝或未完成 MNote 工具审批", toolName: tool.mnoteName }, null, 2) }],
          details: { ok: false, code: "mnote_tool_approval_cancelled", toolName: tool.mnoteName },
        };
      }
      if (toolPolicy(tool.mnoteName) === "ask") {
        nextParams.mnoteApproval = {
          ...approval,
          confirmed: true,
          method: "extension_ui_confirm",
          toolCallId: String(toolCallId || ""),
          toolName: tool.mnoteName,
          approvedAt: new Date().toISOString(),
        };
      }
      return callMnote(tool.mnoteName, nextParams);
}

export default function mnotePi(pi: ExtensionAPI) {
  pi.on?.("input", captureInputContext);
  for (const tool of TOOLS) register(pi, tool);
}