co-mono/packages/coding-agent/src/modes/interactive/components/tree-selector.ts

770 lines
25 KiB
TypeScript
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import {
type Component,
Container,
Input,
isArrowDown,
isArrowUp,
isBackspace,
isCtrlC,
isCtrlO,
isEnter,
isEscape,
Spacer,
Text,
TruncatedText,
truncateToWidth,
} from "@mariozechner/pi-tui";
import type { SessionTreeNode } from "../../../core/session-manager.js";
import { theme } from "../theme/theme.js";
import { DynamicBorder } from "./dynamic-border.js";
/** Flattened tree node for navigation */
interface FlatNode {
node: SessionTreeNode;
/** Indentation level (each level = 2 spaces) */
indent: number;
/** Whether to show connector (├─ or └─) - true if parent has multiple children */
showConnector: boolean;
/** If showConnector, true = last sibling (└─), false = not last (├─) */
isLast: boolean;
/** For each ancestor branch point, true = show │ (more siblings below), false = show space */
gutters: boolean[];
/** True if this node is a root under a virtual branching root (multiple roots) */
isVirtualRootChild: boolean;
}
/** Filter mode for tree display */
type FilterMode = "default" | "user-only" | "labeled-only" | "all";
/**
* Tree list component with selection and ASCII art visualization
*/
/** Tool call info for lookup */
interface ToolCallInfo {
name: string;
arguments: Record<string, unknown>;
}
class TreeList implements Component {
private flatNodes: FlatNode[] = [];
private filteredNodes: FlatNode[] = [];
private selectedIndex = 0;
private currentLeafId: string | null;
private maxVisibleLines: number;
private filterMode: FilterMode = "default";
private searchQuery = "";
private toolCallMap: Map<string, ToolCallInfo> = new Map();
private multipleRoots = false;
private activePathIds: Set<string> = new Set();
public onSelect?: (entryId: string) => void;
public onCancel?: () => void;
public onLabelEdit?: (entryId: string, currentLabel: string | undefined) => void;
constructor(tree: SessionTreeNode[], currentLeafId: string | null, maxVisibleLines: number) {
this.currentLeafId = currentLeafId;
this.maxVisibleLines = maxVisibleLines;
this.multipleRoots = tree.length > 1;
this.flatNodes = this.flattenTree(tree);
this.buildActivePath();
this.applyFilter();
// Start with current leaf selected
const leafIndex = this.filteredNodes.findIndex((n) => n.node.entry.id === currentLeafId);
if (leafIndex !== -1) {
this.selectedIndex = leafIndex;
} else {
this.selectedIndex = Math.max(0, this.filteredNodes.length - 1);
}
}
/** Build the set of entry IDs on the path from root to current leaf */
private buildActivePath(): void {
this.activePathIds.clear();
if (!this.currentLeafId) return;
// Build a map of id -> entry for parent lookup
const entryMap = new Map<string, FlatNode>();
for (const flatNode of this.flatNodes) {
entryMap.set(flatNode.node.entry.id, flatNode);
}
// Walk from leaf to root
let currentId: string | null = this.currentLeafId;
while (currentId) {
this.activePathIds.add(currentId);
const node = entryMap.get(currentId);
if (!node) break;
currentId = node.node.entry.parentId ?? null;
}
}
private flattenTree(roots: SessionTreeNode[]): FlatNode[] {
const result: FlatNode[] = [];
this.toolCallMap.clear();
// Indentation rules:
// - At indent 0: stay at 0 unless parent has >1 children (then +1)
// - At indent 1: children always go to indent 2 (visual grouping of subtree)
// - At indent 2+: stay flat for single-child chains, +1 only if parent branches
// Stack items: [node, indent, justBranched, showConnector, isLast, gutters, isVirtualRootChild]
type StackItem = [SessionTreeNode, number, boolean, boolean, boolean, boolean[], boolean];
const stack: StackItem[] = [];
// Determine which subtrees contain the active leaf (to sort current branch first)
const containsActive = new Map<SessionTreeNode, boolean>();
const leafId = this.currentLeafId;
const markContains = (node: SessionTreeNode): boolean => {
let has = leafId !== null && node.entry.id === leafId;
for (const child of node.children) {
if (markContains(child)) {
has = true;
}
}
containsActive.set(node, has);
return has;
};
for (const root of roots) {
markContains(root);
}
// Add roots in reverse order, prioritizing the one containing the active leaf
// If multiple roots, treat them as children of a virtual root that branches
const multipleRoots = roots.length > 1;
const orderedRoots = [...roots].sort((a, b) => Number(containsActive.get(b)) - Number(containsActive.get(a)));
for (let i = orderedRoots.length - 1; i >= 0; i--) {
const isLast = i === orderedRoots.length - 1;
stack.push([orderedRoots[i], multipleRoots ? 1 : 0, multipleRoots, multipleRoots, isLast, [], multipleRoots]);
}
while (stack.length > 0) {
const [node, indent, justBranched, showConnector, isLast, gutters, isVirtualRootChild] = stack.pop()!;
// Extract tool calls from assistant messages for later lookup
const entry = node.entry;
if (entry.type === "message" && entry.message.role === "assistant") {
const content = (entry.message as { content?: unknown }).content;
if (Array.isArray(content)) {
for (const block of content) {
if (typeof block === "object" && block !== null && "type" in block && block.type === "toolCall") {
const tc = block as { id: string; name: string; arguments: Record<string, unknown> };
this.toolCallMap.set(tc.id, { name: tc.name, arguments: tc.arguments });
}
}
}
}
result.push({ node, indent, showConnector, isLast, gutters, isVirtualRootChild });
const children = node.children;
const multipleChildren = children.length > 1;
// Order children so the branch containing the active leaf comes first
const orderedChildren = (() => {
const prioritized: SessionTreeNode[] = [];
const rest: SessionTreeNode[] = [];
for (const child of children) {
if (containsActive.get(child)) {
prioritized.push(child);
} else {
rest.push(child);
}
}
return [...prioritized, ...rest];
})();
// Calculate child indent
let childIndent: number;
if (multipleChildren) {
// Parent branches: children get +1
childIndent = indent + 1;
} else if (justBranched && indent > 0) {
// First generation after a branch: +1 for visual grouping
childIndent = indent + 1;
} else {
// Single-child chain: stay flat
childIndent = indent;
}
// Build gutters for children
// If this node showed a connector, add a gutter entry for descendants
const childGutters = showConnector ? [...gutters, !isLast] : gutters;
// Add children in reverse order
for (let i = orderedChildren.length - 1; i >= 0; i--) {
const childIsLast = i === orderedChildren.length - 1;
stack.push([
orderedChildren[i],
childIndent,
multipleChildren,
multipleChildren,
childIsLast,
childGutters,
false,
]);
}
}
return result;
}
private applyFilter(): void {
// Remember currently selected node to preserve cursor position
const previouslySelectedId = this.filteredNodes[this.selectedIndex]?.node.entry.id;
const searchTokens = this.searchQuery.toLowerCase().split(/\s+/).filter(Boolean);
this.filteredNodes = this.flatNodes.filter((flatNode) => {
const entry = flatNode.node.entry;
const isCurrentLeaf = entry.id === this.currentLeafId;
// Skip assistant messages with only tool calls (no text) unless error/aborted
// Always show current leaf so active position is visible
if (entry.type === "message" && entry.message.role === "assistant" && !isCurrentLeaf) {
const msg = entry.message as { stopReason?: string; content?: unknown };
const hasText = this.hasTextContent(msg.content);
const isErrorOrAborted = msg.stopReason && msg.stopReason !== "stop" && msg.stopReason !== "toolUse";
// Only hide if no text AND not an error/aborted message
if (!hasText && !isErrorOrAborted) {
return false;
}
}
// Apply filter mode
let passesFilter = true;
if (this.filterMode === "user-only") {
passesFilter =
(entry.type === "message" && entry.message.role === "user") ||
(entry.type === "custom_message" && entry.display);
} else if (this.filterMode === "labeled-only") {
passesFilter = flatNode.node.label !== undefined;
} else if (this.filterMode !== "all") {
// Default mode: hide label and custom entries
passesFilter = entry.type !== "label" && entry.type !== "custom";
}
if (!passesFilter) return false;
// Apply search filter
if (searchTokens.length > 0) {
const nodeText = this.getSearchableText(flatNode.node).toLowerCase();
return searchTokens.every((token) => nodeText.includes(token));
}
return true;
});
// Try to preserve cursor on the same node after filtering
if (previouslySelectedId) {
const newIndex = this.filteredNodes.findIndex((n) => n.node.entry.id === previouslySelectedId);
if (newIndex !== -1) {
this.selectedIndex = newIndex;
return;
}
}
// Fall back: clamp index if out of bounds
if (this.selectedIndex >= this.filteredNodes.length) {
this.selectedIndex = Math.max(0, this.filteredNodes.length - 1);
}
}
/** Get searchable text content from a node */
private getSearchableText(node: SessionTreeNode): string {
const entry = node.entry;
const parts: string[] = [];
if (node.label) {
parts.push(node.label);
}
switch (entry.type) {
case "message": {
const msg = entry.message;
parts.push(msg.role);
if ("content" in msg && msg.content) {
parts.push(this.extractContent(msg.content));
}
if (msg.role === "bashExecution") {
const bashMsg = msg as { command?: string };
if (bashMsg.command) parts.push(bashMsg.command);
}
break;
}
case "custom_message": {
parts.push(entry.customType);
if (typeof entry.content === "string") {
parts.push(entry.content);
} else {
parts.push(this.extractContent(entry.content));
}
break;
}
case "compaction":
parts.push("compaction");
break;
case "branch_summary":
parts.push("branch summary", entry.summary);
break;
case "model_change":
parts.push("model", entry.modelId);
break;
case "thinking_level_change":
parts.push("thinking", entry.thinkingLevel);
break;
case "custom":
parts.push("custom", entry.customType);
break;
case "label":
parts.push("label", entry.label ?? "");
break;
}
return parts.join(" ");
}
invalidate(): void {}
getSearchQuery(): string {
return this.searchQuery;
}
getSelectedNode(): SessionTreeNode | undefined {
return this.filteredNodes[this.selectedIndex]?.node;
}
updateNodeLabel(entryId: string, label: string | undefined): void {
for (const flatNode of this.flatNodes) {
if (flatNode.node.entry.id === entryId) {
flatNode.node.label = label;
break;
}
}
}
private getFilterLabel(): string {
switch (this.filterMode) {
case "user-only":
return " [user]";
case "labeled-only":
return " [labeled]";
case "all":
return " [all]";
default:
return "";
}
}
render(width: number): string[] {
const lines: string[] = [];
if (this.filteredNodes.length === 0) {
lines.push(truncateToWidth(theme.fg("muted", " No entries found"), width));
lines.push(truncateToWidth(theme.fg("muted", ` (0/0)${this.getFilterLabel()}`), width));
return lines;
}
const startIndex = Math.max(
0,
Math.min(
this.selectedIndex - Math.floor(this.maxVisibleLines / 2),
this.filteredNodes.length - this.maxVisibleLines,
),
);
const endIndex = Math.min(startIndex + this.maxVisibleLines, this.filteredNodes.length);
for (let i = startIndex; i < endIndex; i++) {
const flatNode = this.filteredNodes[i];
const entry = flatNode.node.entry;
const isSelected = i === this.selectedIndex;
const isCurrentLeaf = entry.id === this.currentLeafId;
// Build line: cursor + gutters + connector + extra indent + label + content + suffix
const cursor = isSelected ? theme.fg("accent", " ") : " ";
// If multiple roots, shift display (roots at 0, not 1)
const displayIndent = this.multipleRoots ? Math.max(0, flatNode.indent - 1) : flatNode.indent;
// Build prefix: gutters + connector + extra spaces
const gutterStr = flatNode.gutters.map((g) => (g ? "│ " : " ")).join("");
const connector =
flatNode.showConnector && !flatNode.isVirtualRootChild ? (flatNode.isLast ? "└─ " : "├─ ") : "";
// Extra indent for visual grouping beyond gutters/connector
const prefixLevels = flatNode.gutters.length + (connector ? 1 : 0);
const extraIndent = " ".repeat(Math.max(0, displayIndent - prefixLevels));
const prefix = gutterStr + connector + extraIndent;
// Active path marker - shown right before the entry text
const isOnActivePath = this.activePathIds.has(entry.id);
const pathMarker = isOnActivePath ? theme.fg("accent", "• ") : "";
const label = flatNode.node.label ? theme.fg("warning", `[${flatNode.node.label}] `) : "";
const content = this.getEntryDisplayText(flatNode.node, isSelected);
const suffix = isCurrentLeaf ? theme.fg("accent", " *") : "";
const line = cursor + theme.fg("dim", prefix) + pathMarker + label + content + suffix;
lines.push(truncateToWidth(line, width));
}
lines.push(
truncateToWidth(
theme.fg("muted", ` (${this.selectedIndex + 1}/${this.filteredNodes.length})${this.getFilterLabel()}`),
width,
),
);
return lines;
}
private getEntryDisplayText(node: SessionTreeNode, isSelected: boolean): string {
const entry = node.entry;
let result: string;
const normalize = (s: string) => s.replace(/[\n\t]/g, " ").trim();
switch (entry.type) {
case "message": {
const msg = entry.message;
const role = msg.role;
if (role === "user") {
const msgWithContent = msg as { content?: unknown };
const content = normalize(this.extractContent(msgWithContent.content));
result = theme.fg("accent", "user: ") + content;
} else if (role === "assistant") {
const msgWithContent = msg as { content?: unknown; stopReason?: string; errorMessage?: string };
const textContent = normalize(this.extractContent(msgWithContent.content));
if (textContent) {
result = theme.fg("success", "assistant: ") + textContent;
} else if (msgWithContent.stopReason === "aborted") {
result = theme.fg("warning", "assistant: ") + theme.fg("muted", "(aborted)");
} else if (msgWithContent.errorMessage) {
const errMsg = normalize(msgWithContent.errorMessage).slice(0, 80);
result = theme.fg("error", "assistant: ") + errMsg;
} else {
result = theme.fg("muted", "assistant: (no content)");
}
} else if (role === "toolResult") {
const toolMsg = msg as { toolCallId?: string; toolName?: string };
const toolCall = toolMsg.toolCallId ? this.toolCallMap.get(toolMsg.toolCallId) : undefined;
if (toolCall) {
result = theme.fg("muted", this.formatToolCall(toolCall.name, toolCall.arguments));
} else {
result = theme.fg("muted", `[${toolMsg.toolName ?? "tool"}]`);
}
} else if (role === "bashExecution") {
const bashMsg = msg as { command?: string };
result = theme.fg("dim", `[bash]: ${normalize(bashMsg.command ?? "")}`);
} else {
result = theme.fg("dim", `[${role}]`);
}
break;
}
case "custom_message": {
const content =
typeof entry.content === "string"
? entry.content
: entry.content
.filter((c): c is { type: "text"; text: string } => c.type === "text")
.map((c) => c.text)
.join("");
result = theme.fg("customMessageLabel", `[${entry.customType}]: `) + normalize(content);
break;
}
case "compaction": {
const tokens = Math.round(entry.tokensBefore / 1000);
result = theme.fg("borderAccent", `[compaction: ${tokens}k tokens]`);
break;
}
case "branch_summary":
result = theme.fg("warning", `[branch summary]: `) + normalize(entry.summary);
break;
case "model_change":
result = theme.fg("dim", `[model: ${entry.modelId}]`);
break;
case "thinking_level_change":
result = theme.fg("dim", `[thinking: ${entry.thinkingLevel}]`);
break;
case "custom":
result = theme.fg("dim", `[custom: ${entry.customType}]`);
break;
case "label":
result = theme.fg("dim", `[label: ${entry.label ?? "(cleared)"}]`);
break;
default:
result = "";
}
return isSelected ? theme.bold(result) : result;
}
private extractContent(content: unknown): string {
const maxLen = 200;
if (typeof content === "string") return content.slice(0, maxLen);
if (Array.isArray(content)) {
let result = "";
for (const c of content) {
if (typeof c === "object" && c !== null && "type" in c && c.type === "text") {
result += (c as { text: string }).text;
if (result.length >= maxLen) return result.slice(0, maxLen);
}
}
return result;
}
return "";
}
private hasTextContent(content: unknown): boolean {
if (typeof content === "string") return content.trim().length > 0;
if (Array.isArray(content)) {
for (const c of content) {
if (typeof c === "object" && c !== null && "type" in c && c.type === "text") {
const text = (c as { text?: string }).text;
if (text && text.trim().length > 0) return true;
}
}
}
return false;
}
private formatToolCall(name: string, args: Record<string, unknown>): string {
const shortenPath = (p: string): string => {
const home = process.env.HOME || process.env.USERPROFILE || "";
if (home && p.startsWith(home)) return `~${p.slice(home.length)}`;
return p;
};
switch (name) {
case "read": {
const path = shortenPath(String(args.path || args.file_path || ""));
const offset = args.offset as number | undefined;
const limit = args.limit as number | undefined;
let display = path;
if (offset !== undefined || limit !== undefined) {
const start = offset ?? 1;
const end = limit !== undefined ? start + limit - 1 : "";
display += `:${start}${end ? `-${end}` : ""}`;
}
return `[read: ${display}]`;
}
case "write": {
const path = shortenPath(String(args.path || args.file_path || ""));
return `[write: ${path}]`;
}
case "edit": {
const path = shortenPath(String(args.path || args.file_path || ""));
return `[edit: ${path}]`;
}
case "bash": {
const cmd = String(args.command || "").slice(0, 50);
return `[bash: ${cmd}${(args.command as string)?.length > 50 ? "..." : ""}]`;
}
case "grep": {
const pattern = String(args.pattern || "");
const path = shortenPath(String(args.path || "."));
return `[grep: /${pattern}/ in ${path}]`;
}
case "find": {
const pattern = String(args.pattern || "");
const path = shortenPath(String(args.path || "."));
return `[find: ${pattern} in ${path}]`;
}
case "ls": {
const path = shortenPath(String(args.path || "."));
return `[ls: ${path}]`;
}
default: {
// Custom tool - show name and truncated JSON args
const argsStr = JSON.stringify(args).slice(0, 40);
return `[${name}: ${argsStr}${JSON.stringify(args).length > 40 ? "..." : ""}]`;
}
}
}
handleInput(keyData: string): void {
if (isArrowUp(keyData)) {
this.selectedIndex = this.selectedIndex === 0 ? this.filteredNodes.length - 1 : this.selectedIndex - 1;
} else if (isArrowDown(keyData)) {
this.selectedIndex = this.selectedIndex === this.filteredNodes.length - 1 ? 0 : this.selectedIndex + 1;
} else if (isEnter(keyData)) {
const selected = this.filteredNodes[this.selectedIndex];
if (selected && this.onSelect) {
this.onSelect(selected.node.entry.id);
}
} else if (isEscape(keyData)) {
if (this.searchQuery) {
this.searchQuery = "";
this.applyFilter();
} else {
this.onCancel?.();
}
} else if (isCtrlC(keyData)) {
this.onCancel?.();
} else if (isCtrlO(keyData)) {
// Cycle filter: default → user-only → labeled-only → all → default
const modes: FilterMode[] = ["default", "user-only", "labeled-only", "all"];
const currentIndex = modes.indexOf(this.filterMode);
this.filterMode = modes[(currentIndex + 1) % modes.length];
this.applyFilter();
} else if (isBackspace(keyData)) {
if (this.searchQuery.length > 0) {
this.searchQuery = this.searchQuery.slice(0, -1);
this.applyFilter();
}
} else if (keyData === "l" && !this.searchQuery) {
const selected = this.filteredNodes[this.selectedIndex];
if (selected && this.onLabelEdit) {
this.onLabelEdit(selected.node.entry.id, selected.node.label);
}
} else {
const hasControlChars = [...keyData].some((ch) => {
const code = ch.charCodeAt(0);
return code < 32 || code === 0x7f || (code >= 0x80 && code <= 0x9f);
});
if (!hasControlChars && keyData.length > 0) {
this.searchQuery += keyData;
this.applyFilter();
}
}
}
}
/** Component that displays the current search query */
class SearchLine implements Component {
constructor(private treeList: TreeList) {}
invalidate(): void {}
render(width: number): string[] {
const query = this.treeList.getSearchQuery();
if (query) {
return [truncateToWidth(` ${theme.fg("muted", "Search:")} ${theme.fg("accent", query)}`, width)];
}
return [truncateToWidth(` ${theme.fg("muted", "Search:")}`, width)];
}
handleInput(_keyData: string): void {}
}
/** Label input component shown when editing a label */
class LabelInput implements Component {
private input: Input;
private entryId: string;
public onSubmit?: (entryId: string, label: string | undefined) => void;
public onCancel?: () => void;
constructor(entryId: string, currentLabel: string | undefined) {
this.entryId = entryId;
this.input = new Input();
if (currentLabel) {
this.input.setValue(currentLabel);
}
}
invalidate(): void {}
render(width: number): string[] {
const lines: string[] = [];
const indent = " ";
const availableWidth = width - indent.length;
lines.push(truncateToWidth(`${indent}${theme.fg("muted", "Label (empty to remove):")}`, width));
lines.push(...this.input.render(availableWidth).map((line) => truncateToWidth(`${indent}${line}`, width)));
lines.push(truncateToWidth(`${indent}${theme.fg("dim", "enter: save esc: cancel")}`, width));
return lines;
}
handleInput(keyData: string): void {
if (isEnter(keyData)) {
const value = this.input.getValue().trim();
this.onSubmit?.(this.entryId, value || undefined);
} else if (isEscape(keyData)) {
this.onCancel?.();
} else {
this.input.handleInput(keyData);
}
}
}
/**
* Component that renders a session tree selector for navigation
*/
export class TreeSelectorComponent extends Container {
private treeList: TreeList;
private labelInput: LabelInput | null = null;
private labelInputContainer: Container;
private treeContainer: Container;
private onLabelChangeCallback?: (entryId: string, label: string | undefined) => void;
constructor(
tree: SessionTreeNode[],
currentLeafId: string | null,
terminalHeight: number,
onSelect: (entryId: string) => void,
onCancel: () => void,
onLabelChange?: (entryId: string, label: string | undefined) => void,
) {
super();
this.onLabelChangeCallback = onLabelChange;
const maxVisibleLines = Math.max(5, Math.floor(terminalHeight / 2));
this.treeList = new TreeList(tree, currentLeafId, maxVisibleLines);
this.treeList.onSelect = onSelect;
this.treeList.onCancel = onCancel;
this.treeList.onLabelEdit = (entryId, currentLabel) => this.showLabelInput(entryId, currentLabel);
this.treeContainer = new Container();
this.treeContainer.addChild(this.treeList);
this.labelInputContainer = new Container();
this.addChild(new Spacer(1));
this.addChild(new Text(theme.bold("Session Tree"), 1, 0));
this.addChild(new DynamicBorder());
this.addChild(new TruncatedText(theme.fg("muted", " Type to search. l: label. ^O: cycle filter"), 0, 0));
this.addChild(new SearchLine(this.treeList));
this.addChild(new DynamicBorder());
this.addChild(new Spacer(1));
this.addChild(this.treeContainer);
this.addChild(this.labelInputContainer);
this.addChild(new Spacer(1));
this.addChild(new DynamicBorder());
if (tree.length === 0) {
setTimeout(() => onCancel(), 100);
}
}
private showLabelInput(entryId: string, currentLabel: string | undefined): void {
this.labelInput = new LabelInput(entryId, currentLabel);
this.labelInput.onSubmit = (id, label) => {
this.treeList.updateNodeLabel(id, label);
this.onLabelChangeCallback?.(id, label);
this.hideLabelInput();
};
this.labelInput.onCancel = () => this.hideLabelInput();
this.treeContainer.clear();
this.labelInputContainer.clear();
this.labelInputContainer.addChild(this.labelInput);
}
private hideLabelInput(): void {
this.labelInput = null;
this.labelInputContainer.clear();
this.treeContainer.clear();
this.treeContainer.addChild(this.treeList);
}
handleInput(keyData: string): void {
if (this.labelInput) {
this.labelInput.handleInput(keyData);
} else {
this.treeList.handleInput(keyData);
}
}
getTreeList(): TreeList {
return this.treeList;
}
}