MathLab/task-board/backend/data/tree-pruner.js

"use strict";

function buildMaps(nodes, edges) {
  const byId = new Map(nodes.map((node) => [node.id, node]));
  const parent = new Map();
  const children = new Map();

  for (const edge of edges) {
    if (edge.type !== "hierarchy") continue;
    parent.set(edge.to, edge.from);
    if (!children.has(edge.from)) children.set(edge.from, []);
    children.get(edge.from).push(edge.to);
  }

  return { byId, parent, children };
}

function getAncestors(focusId, maps) {
  const path = [];
  let cursor = focusId;
  const seen = new Set();
  while (cursor && !seen.has(cursor)) {
    seen.add(cursor);
    const node = maps.byId.get(cursor);
    if (!node) break;
    path.push(node);
    cursor = maps.parent.get(cursor) || null;
  }
  return path.reverse();
}

function summarizeNode(node) {
  return `[${node.status}] ${node.name}`;
}

function pruneConversationContext(graph, focusNodeId) {
  const nodes = Array.isArray(graph.nodes) ? graph.nodes : [];
  const edges = Array.isArray(graph.edges) ? graph.edges : [];
  const maps = buildMaps(nodes, edges);
  const focus = maps.byId.get(focusNodeId);
  if (!focus) {
    return {
      focusNodeId,
      markdown: "未找到焦点节点，无法生成上下文。",
      segments: {
        ancestorPath: [],
        directChildren: [],
        siblingSummaries: [],
        foldedSummary: "无",
      },
    };
  }

  const ancestorPath = getAncestors(focusNodeId, maps);
  const directChildren = (maps.children.get(focusNodeId) || []).map((id) => maps.byId.get(id)).filter(Boolean);

  const parentId = maps.parent.get(focusNodeId) || null;
  const siblingSummaries = parentId
    ? (maps.children.get(parentId) || [])
        .filter((id) => id !== focusNodeId)
        .map((id) => maps.byId.get(id))
        .filter(Boolean)
        .map(summarizeNode)
    : [];

  const keepIds = new Set([
    ...ancestorPath.map((node) => node.id),
    focusNodeId,
    ...directChildren.map((node) => node.id),
  ]);
  const foldedCount = nodes.filter((node) => !keepIds.has(node.id) && node.status === "completed").length;
  const foldedSummary = foldedCount > 0 ? `已折叠 ${foldedCount} 个已完成远端节点` : "无";

  const markdownLines = [];
  markdownLines.push("## Focused Context");
  markdownLines.push("");
  markdownLines.push("### Ancestor Path");
  if (ancestorPath.length === 0) {
    markdownLines.push("- (empty)");
  } else {
    for (const node of ancestorPath) markdownLines.push(`- ${summarizeNode(node)}`);
  }
  markdownLines.push("");
  markdownLines.push("### Direct Children");
  if (directChildren.length === 0) {
    markdownLines.push("- (empty)");
  } else {
    for (const node of directChildren) markdownLines.push(`- ${summarizeNode(node)}`);
  }
  markdownLines.push("");
  markdownLines.push("### Sibling Branches");
  if (siblingSummaries.length === 0) {
    markdownLines.push("- (empty)");
  } else {
    for (const summary of siblingSummaries) markdownLines.push(`- ${summary}`);
  }
  markdownLines.push("");
  markdownLines.push("### Folded History");
  markdownLines.push(`- ${foldedSummary}`);

  return {
    focusNodeId,
    markdown: markdownLines.join("\n"),
    segments: {
      ancestorPath,
      directChildren,
      siblingSummaries,
      foldedSummary,
    },
  };
}

module.exports = {
  pruneConversationContext,
};