import { useMemo, useCallback, useRef, useState, useEffect } from 'react'; import ForceGraph3D from 'react-force-graph-3d'; import { useConversation } from '../context/ConversationContext'; import { toGraphData } from '../utils/tree'; import BranchDialog from './BranchDialog'; import * as THREE from 'three'; export default function TreeView() { const { state, switchToBranch, setNodeAsMainBranch, } = useConversation(); const fgRef = useRef(null); const containerRef = useRef(null); const treePaneRef = useRef(null); const prevNodeCountRef = useRef(0); const [dimensions, setDimensions] = useState({ width: 420, height: 600 }); const [contextMenu, setContextMenu] = useState(null); const [isFullscreen, setIsFullscreen] = useState(false); const [branchNode, setBranchNode] = useState(null); const [showLegend, setShowLegend] = useState(true); const conv = state.activeConversationId ? state.conversations[state.activeConversationId] : null; // --- Stable graph structure fingerprint --- const structureFingerprint = useMemo(() => { if (!conv) return ''; const nodeIds = Object.keys(conv.nodes).sort().join(','); const links = Object.values(conv.nodes) .filter((n) => n.parentId) .map((n) => `${n.parentId}->${n.id}`) .sort() .join(','); const mainPath = conv.mainBranchPath.join(','); const activePath = conv.activeBranchPath.join(','); const statuses = Object.values(conv.nodes) .map((n) => `${n.id}:${n.status}`) .join(','); return `${nodeIds}|${links}|${mainPath}|${activePath}|${statuses}`; }, [conv]); const graphData = useMemo(() => { if (!conv || Object.keys(conv.nodes).length === 0) { return { nodes: [], links: [] }; } return toGraphData(conv.nodes, conv.mainBranchPath, conv.activeBranchPath); // eslint-disable-next-line react-hooks/exhaustive-deps }, [structureFingerprint]); // Auto-fit camera when new nodes are added useEffect(() => { const currentCount = graphData.nodes.length; if (currentCount > prevNodeCountRef.current && currentCount > 0 && fgRef.current) { setTimeout(() => { fgRef.current?.zoomToFit(600, 60); }, 300); } prevNodeCountRef.current = currentCount; }, [graphData.nodes.length]); // Resize observer useEffect(() => { if (!containerRef.current) return; const observer = new ResizeObserver((entries) => { const { width, height } = entries[0].contentRect; setDimensions({ width, height }); }); observer.observe(containerRef.current); return () => observer.disconnect(); }, []); // Close context menu on click outside useEffect(() => { if (!contextMenu) return; const handler = () => setContextMenu(null); window.addEventListener('click', handler); return () => window.removeEventListener('click', handler); }, [contextMenu]); // Handle Escape to exit fullscreen useEffect(() => { if (!isFullscreen) return; const handler = (e) => { if (e.key === 'Escape') setIsFullscreen(false); }; window.addEventListener('keydown', handler); return () => window.removeEventListener('keydown', handler); }, [isFullscreen]); const handleNodeClick = useCallback( (node) => { if (!conv) return; switchToBranch(state.activeConversationId, node.id); }, [conv, state.activeConversationId, switchToBranch] ); const handleNodeRightClick = useCallback( (node, event) => { event.preventDefault(); setContextMenu({ x: event.clientX, y: event.clientY, nodeId: node.id, userMessage: node.userMessage || '', }); }, [] ); const handleSetMainBranch = useCallback(() => { if (!contextMenu || !conv) return; setNodeAsMainBranch(state.activeConversationId, contextMenu.nodeId); setContextMenu(null); }, [contextMenu, conv, state.activeConversationId, setNodeAsMainBranch]); const handleBranchFromGraph = useCallback(() => { if (!contextMenu || !conv) return; const node = conv.nodes[contextMenu.nodeId]; if (!node) return; setBranchNode(node); setContextMenu(null); }, [contextMenu, conv]); const handleResetView = useCallback(() => { if (!fgRef.current) return; fgRef.current.zoomToFit(400, 40); }, []); const toggleFullscreen = useCallback(() => { setIsFullscreen((prev) => !prev); setTimeout(() => { if (fgRef.current) { fgRef.current.zoomToFit(400, 40); } }, 100); }, []); // --- Color helpers --- const getNodeColor = useCallback((node) => { if (node.status === 'generating') return '#f0c040'; if (node.isActive && node.isMain) return '#00ffcc'; if (node.isActive) return '#6c7aff'; if (node.isMain) return '#00d4aa'; return '#555566'; }, []); const linkColor = useCallback((link) => { if (link.isActive && link.isMain) return 'rgba(0, 255, 204, 0.5)'; if (link.isActive) return 'rgba(108, 122, 255, 0.4)'; if (link.isMain) return 'rgba(0, 212, 170, 0.35)'; return 'rgba(85, 85, 102, 0.2)'; }, []); const linkWidth = useCallback((link) => { if (link.isMain || link.isActive) return 2.5; return 1; }, []); // --- Rich hover tooltip --- const nodeLabel = useCallback((node) => { const userPreview = node.userMessage ? node.userMessage.slice(0, 120) + (node.userMessage.length > 120 ? '…' : '') : '(start)'; const aiPreview = node.assistantMessage ? node.assistantMessage.slice(0, 200) + (node.assistantMessage.length > 200 ? '…' : '') : node.status === 'generating' ? '⏳ Generating…' : '—'; const statusBadge = node.status === 'generating' ? 'GENERATING' : node.childCount > 1 ? `${node.childCount} branches` : ''; return `
Turn ${node.depth + 1} ${statusBadge}
👤 User
${escapeHtml(userPreview)}
🤖 Assistant
${escapeHtml(aiPreview)}
`; }, []); // --- Custom 3D node with text label --- const nodeThreeObject = useCallback((node) => { const group = new THREE.Group(); // Core sphere const isHighlighted = node.isMain || node.isActive; const size = isHighlighted ? 5 : 3; const color = getNodeColor(node); const geo = new THREE.SphereGeometry(size, 20, 20); const mat = new THREE.MeshPhongMaterial({ color, emissive: color, emissiveIntensity: node.isActive ? 0.6 : 0.25, transparent: true, opacity: node.status === 'generating' ? 0.65 : 1, shininess: 60, }); const mesh = new THREE.Mesh(geo, mat); group.add(mesh); // Outer glow ring for highlighted nodes if (isHighlighted) { const ringGeo = new THREE.RingGeometry(size * 1.4, size * 1.7, 32); const ringMat = new THREE.MeshBasicMaterial({ color, transparent: true, opacity: 0.18, side: THREE.DoubleSide, }); const ringMesh = new THREE.Mesh(ringGeo, ringMat); group.add(ringMesh); } // Branch indicator (diamond) for nodes with multiple children if (node.childCount > 1) { const diamondGeo = new THREE.OctahedronGeometry(2.5); const diamondMat = new THREE.MeshPhongMaterial({ color: '#ff6c8a', emissive: '#ff6c8a', emissiveIntensity: 0.3, transparent: true, opacity: 0.7, }); const diamond = new THREE.Mesh(diamondGeo, diamondMat); diamond.position.set(size + 4, 0, 0); diamond.rotation.set(0, 0, Math.PI / 4); group.add(diamond); } // Sprite label below the node (pure WebGL, no DOM elements) const label = createSpriteLabel(node, color); label.position.set(0, -(size + 8), 0); group.add(label); return group; }, [getNodeColor]); const nodeCount = graphData.nodes.length; return (
Graph View
{nodeCount} node{nodeCount !== 1 ? 's' : ''} {nodeCount > 0 && ( <> )}
{nodeCount === 0 ? (

Start chatting to see your conversation tree

) : ( (node.isMain || node.isActive ? 3 : 1.5)} nodeThreeObject={nodeThreeObject} nodeThreeObjectExtend={false} linkColor={linkColor} linkWidth={linkWidth} linkCurvature={0.15} linkCurveRotation={0} linkDirectionalArrowLength={4} linkDirectionalArrowRelPos={0.85} linkDirectionalArrowColor={linkColor} linkDirectionalParticles={(link) => link.isActive ? 3 : link.isMain ? 2 : 0 } linkDirectionalParticleWidth={(link) => link.isActive ? 2.5 : 1.5 } linkDirectionalParticleColor={(link) => link.isActive ? '#6c7aff' : '#00d4aa' } linkDirectionalParticleSpeed={0.004} linkOpacity={0.7} onNodeClick={handleNodeClick} onNodeRightClick={handleNodeRightClick} enableNodeDrag={false} cooldownTicks={100} warmupTicks={50} d3AlphaDecay={0.08} d3VelocityDecay={0.4} showNavInfo={false} /> )} {/* Color Legend */} {showLegend && nodeCount > 0 && (
Legend
Active + Main
Active branch
Main branch
Generating
Inactive
Branch point
)}
{contextMenu && (
e.stopPropagation()} >
)} {branchNode && ( setBranchNode(null)} onConfirm={(nodeId, newMessage) => { window.dispatchEvent( new CustomEvent('branchchat:branch', { detail: { nodeId, newMessage }, }) ); setBranchNode(null); }} /> )}
); } /** * Create a canvas-rendered sprite label for a graph node. * Uses THREE.Sprite + CanvasTexture instead of CSS2DObject to avoid * orphaned DOM elements (ghost text) when the graph re-renders. */ function createSpriteLabel(node, color) { const preview = node.userMessage ? node.userMessage.slice(0, 22) + (node.userMessage.length > 22 ? '…' : '') : 'Start'; const turnNum = node.depth + 1; const statusIcon = node.status === 'generating' ? ' ⏳' : ''; const turnText = `#${turnNum}${statusIcon}`; const canvas = document.createElement('canvas'); const ctx = canvas.getContext('2d'); // High-DPI canvas for crisp text const scale = 2; const canvasW = 200 * scale; const canvasH = 48 * scale; canvas.width = canvasW; canvas.height = canvasH; ctx.scale(scale, scale); // Turn number line ctx.font = 'bold 11px Inter, sans-serif'; ctx.fillStyle = color; ctx.textAlign = 'center'; ctx.textBaseline = 'top'; ctx.fillText(turnText, 100, 4); // Preview line ctx.font = '9px Inter, sans-serif'; ctx.fillStyle = 'rgba(232, 232, 240, 0.5)'; ctx.fillText(preview, 100, 22); const texture = new THREE.CanvasTexture(canvas); texture.minFilter = THREE.LinearFilter; const spriteMat = new THREE.SpriteMaterial({ map: texture, transparent: true, depthWrite: false, }); const sprite = new THREE.Sprite(spriteMat); // Scale sprite to a reasonable world-space size sprite.scale.set(20, 20 * (canvasH / canvasW), 1); return sprite; } function escapeHtml(str) { return str .replace(/&/g, '&') .replace(//g, '>') .replace(/"/g, '"'); }