ai-agent/frontend/src/tree.ts

import type { FileEntry, TreeNode } from "./types";

/** Build a nested folder/file tree from a flat list of files (by path). */
export function buildTree(files: FileEntry[]): TreeNode {
  const root: TreeNode = { name: "", path: "", dir: true, children: [] };

  for (const file of files) {
    const parts = file.path.split("/");
    let node = root;
    let acc = "";
    parts.forEach((part, i) => {
      acc = acc ? `${acc}/${part}` : part;
      const isLeaf = i === parts.length - 1;
      let child = node.children.find((c) => c.name === part);
      if (!child) {
        child = {
          name: part,
          path: acc,
          dir: !isLeaf,
          file: isLeaf ? file : undefined,
          children: [],
        };
        node.children.push(child);
      }
      node = child;
    });
  }

  sortNode(root);
  return root;
}

function sortNode(node: TreeNode) {
  node.children.sort((a, b) => {
    if (a.dir !== b.dir) return a.dir ? -1 : 1; // folders first
    return a.name.localeCompare(b.name);
  });
  node.children.forEach(sortNode);
}

/** Case-insensitive substring filter on the full path. */
export function filterFiles(files: FileEntry[], query: string): FileEntry[] {
  const q = query.trim().toLowerCase();
  if (!q) return files;
  return files.filter((f) => f.path.toLowerCase().includes(q));
}

/** Sum the token count + dollar cost over the files contained under a node. */
export function nodeStats(node: TreeNode): { tokens: number; cost: number } {
  if (node.file) {
    return { tokens: node.file.tokens, cost: node.file.cost };
  }
  return node.children.reduce(
    (acc, c) => {
      const t = nodeStats(c);
      return { tokens: acc.tokens + t.tokens, cost: acc.cost + t.cost };
    },
    { tokens: 0, cost: 0 },
  );
}