"use client"; import { useEffect, useMemo, useRef, useState } from "react"; import { QueryClient, QueryClientProvider, useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { AnimatePresence, motion } from "framer-motion"; import { FileJson, Loader2, Save, X } from "lucide-react"; import { Button } from "@/components/ui/Button"; import { Input } from "@/components/ui/Input"; import { Textarea } from "@/components/ui/Textarea"; import { CompactTreeNavigator, type CompactNavigatorFocusMode } from "@/components/tree-visualizer/CompactTreeNavigator"; import { TreeCanvas, type TreeCanvasFocusRequest } from "@/components/tree-visualizer/TreeCanvas"; import { TreeCommandBar } from "@/components/tree-visualizer/TreeCommandBar"; import { TreeLibraryRail } from "@/components/tree-visualizer/TreeLibraryRail"; import { CreateTreeDialog } from "@/components/tree-visualizer/CreateTreeDialog"; import { ModelStreamPanel } from "@/components/tree-visualizer/ModelStreamPanel"; import { NodeInspector } from "@/components/tree-visualizer/NodeInspector"; import type { TreeNodeViewMode } from "@/components/tree-visualizer/viewMode"; import { createTreeFromRaw, deleteTree, deleteTrees, fetchTree, fetchTrees, updateTree, validateTreePayload } from "@/lib/api"; import { getTreeNavigationState } from "@/lib/tree/navigation"; import { parseTrajectoryText } from "@/lib/tree/parser"; import { visibleNodeIdsForStep } from "@/lib/tree/playback"; import type { TrajectoryTree, TrajectoryTreeNode, ValidationResult } from "@/lib/tree/types"; import { cn } from "@/lib/utils"; const defaultRailWidth = 300; const defaultInspectorWidth = 420; const railWidthKey = "tree-visualizer.layout.railWidth"; const inspectorWidthKey = "tree-visualizer.layout.inspectorWidth"; const nodeViewModeKey = "tree-visualizer.nodeViewMode"; function clamp(value: number, min: number, max: number) { return Math.min(max, Math.max(min, value)); } function storedWidth(key: string, fallback: number, min: number, max: number) { if (typeof window === "undefined") return fallback; const value = Number(window.localStorage.getItem(key)); return Number.isFinite(value) ? clamp(value, min, max) : fallback; } function storedNodeViewMode() { if (typeof window === "undefined") return "full" as const; return window.localStorage.getItem(nodeViewModeKey) === "atom" ? "atom" : "full"; } function stableTreeJson(tree: TrajectoryTree | null) { if (!tree) return ""; return JSON.stringify({ title: tree.title, raw_text: tree.raw_text, nodes: tree.nodes }); } function safeFilePart(value: string) { return value .toLowerCase() .replace(/[^a-z0-9_-]+/g, "_") .replace(/^_+|_+$/g, "") .slice(0, 80); } function downloadText(filename: string, content: string, type = "text/plain;charset=utf-8") { const blob = new Blob([content], { type }); const url = URL.createObjectURL(blob); const link = document.createElement("a"); link.href = url; link.download = filename; document.body.appendChild(link); link.click(); link.remove(); URL.revokeObjectURL(url); } function GeneratedTreeSaveDialog({ open, tree, isSaving, error, onClose, onSave }: { open: boolean; tree: TrajectoryTree | null; isSaving: boolean; error: Error | null; onClose: () => void; onSave: (title: string) => void; }) { const [title, setTitle] = useState(""); useEffect(() => { if (!open) return; setTitle(tree?.title || "Model Stream"); }, [open, tree?.title]); const preview = useMemo( () => JSON.stringify( { title: title.trim() || tree?.title || "Model Stream", rawText: tree?.raw_text ?? "", nodes: tree?.nodes ?? [] }, null, 2 ), [title, tree] ); return ( {open && tree && ( Save generated tree {tree.nodes.length} generated nodes Tree name setTitle(event.target.value)} placeholder="Generated reasoning tree" disabled={isSaving} /> {error && ( {error.message} )} Raw JSON to save Cancel onSave(title)} disabled={isSaving || tree.nodes.length === 0}> {isSaving ? : } Save tree )} ); } function TreeVisualizerInner() { const queryClient = useQueryClient(); const [selectedTreeId, setSelectedTreeId] = useState(null); const [tree, setTree] = useState(null); const [validation, setValidation] = useState(null); const [selectedNodeId, setSelectedNodeId] = useState(null); const [createOpen, setCreateOpen] = useState(false); const [streamOpen, setStreamOpen] = useState(false); const [generatedSaveOpen, setGeneratedSaveOpen] = useState(false); const [liveStreamTreeId, setLiveStreamTreeId] = useState(null); const liveStreamTreeIdRef = useRef(null); const [lastSavedJson, setLastSavedJson] = useState(""); const [revealCount, setRevealCount] = useState(0); const [isPlaying, setIsPlaying] = useState(false); const [fitNonce, setFitNonce] = useState(0); const [focusRequest, setFocusRequest] = useState(null); const [autoReadEnabled, setAutoReadEnabled] = useState(false); const [createdReplayTreeId, setCreatedReplayTreeId] = useState(null); const [darkMode, setDarkMode] = useState(false); const [railWidth, setRailWidth] = useState(defaultRailWidth); const [inspectorWidth, setInspectorWidth] = useState(defaultInspectorWidth); const [resizingPanel, setResizingPanel] = useState<"rail" | "inspector" | null>(null); const [nodeViewMode, setNodeViewMode] = useState("full"); const treesQuery = useQuery({ queryKey: ["trees"], queryFn: fetchTrees }); useEffect(() => { setDarkMode(window.localStorage.getItem("tree-visualizer.theme") === "dark"); setRailWidth(storedWidth(railWidthKey, defaultRailWidth, 240, 460)); setInspectorWidth(storedWidth(inspectorWidthKey, defaultInspectorWidth, 340, 620)); setNodeViewMode(storedNodeViewMode()); }, []); useEffect(() => { window.localStorage.setItem("tree-visualizer.theme", darkMode ? "dark" : "light"); }, [darkMode]); useEffect(() => { window.localStorage.setItem(railWidthKey, String(railWidth)); }, [railWidth]); useEffect(() => { window.localStorage.setItem(inspectorWidthKey, String(inspectorWidth)); }, [inspectorWidth]); useEffect(() => { window.localStorage.setItem(nodeViewModeKey, nodeViewMode); }, [nodeViewMode]); useEffect(() => { if (!resizingPanel) return; const previousCursor = document.body.style.cursor; const previousUserSelect = document.body.style.userSelect; document.body.style.cursor = "col-resize"; document.body.style.userSelect = "none"; const handlePointerMove = (event: PointerEvent) => { if (resizingPanel === "rail") { setRailWidth(clamp(event.clientX, 240, 460)); return; } setInspectorWidth(clamp(window.innerWidth - event.clientX, 340, 620)); }; const stopResize = () => setResizingPanel(null); window.addEventListener("pointermove", handlePointerMove); window.addEventListener("pointerup", stopResize, { once: true }); window.addEventListener("pointercancel", stopResize, { once: true }); return () => { document.body.style.cursor = previousCursor; document.body.style.userSelect = previousUserSelect; window.removeEventListener("pointermove", handlePointerMove); window.removeEventListener("pointerup", stopResize); window.removeEventListener("pointercancel", stopResize); }; }, [resizingPanel]); useEffect(() => { if (selectedTreeId || !treesQuery.data?.trees.length) return; setSelectedTreeId(treesQuery.data.trees[0].id); }, [selectedTreeId, treesQuery.data?.trees]); const treeQuery = useQuery({ queryKey: ["tree", selectedTreeId], queryFn: () => fetchTree(selectedTreeId as string), enabled: Boolean(selectedTreeId) }); useEffect(() => { if (!treeQuery.data?.tree) return; if (liveStreamTreeIdRef.current) return; const loadedTree = treeQuery.data.tree; setTree(loadedTree); setValidation(treeQuery.data.validation); setLastSavedJson(stableTreeJson(loadedTree)); liveStreamTreeIdRef.current = null; setLiveStreamTreeId(null); setSelectedNodeId(null); setAutoReadEnabled(false); if (createdReplayTreeId === loadedTree.id) { setRevealCount(0); setIsPlaying(true); setCreatedReplayTreeId(null); } else { setRevealCount(loadedTree.nodes.length); setIsPlaying(false); } }, [createdReplayTreeId, treeQuery.data]); useEffect(() => { if (!tree || !isPlaying) return; if (revealCount >= tree.nodes.length) { setIsPlaying(false); return; } const timer = window.setTimeout(() => { setRevealCount((current) => Math.min(tree.nodes.length, current + 1)); }, 380); return () => window.clearTimeout(timer); }, [isPlaying, revealCount, tree]); useEffect(() => { if (!tree || !autoReadEnabled) return; const visibleIds = visibleNodeIdsForStep(tree.nodes, revealCount); const navigation = getTreeNavigationState(tree.nodes, visibleIds, selectedNodeId); const target = selectedNodeId ? navigation.nextNode : navigation.currentNode; if (!target) { setAutoReadEnabled(false); return; } const timer = window.setTimeout(() => { setSelectedNodeId(target.id); setFocusRequest({ nodeId: target.id, mode: "node", nonce: Date.now() }); }, selectedNodeId ? 2600 : 120); return () => window.clearTimeout(timer); }, [autoReadEnabled, revealCount, selectedNodeId, tree]); const saveMutation = useMutation({ mutationFn: async (input?: { title?: string }) => { if (!tree) throw new Error("No tree loaded."); if (liveStreamTreeId && tree.id === liveStreamTreeId) { return createTreeFromRaw({ title: input?.title?.trim() || tree.title, rawText: tree.raw_text, nodes: tree.nodes }); } return updateTree(tree.id, { title: tree.title, raw_text: tree.raw_text, nodes: tree.nodes }); }, onSuccess: ({ tree: savedTree, validation: nextValidation }) => { setTree(savedTree); setValidation(nextValidation); setLastSavedJson(stableTreeJson(savedTree)); liveStreamTreeIdRef.current = null; setLiveStreamTreeId(null); setSelectedTreeId(savedTree.id); setGeneratedSaveOpen(false); queryClient.invalidateQueries({ queryKey: ["trees"] }); queryClient.invalidateQueries({ queryKey: ["tree", savedTree.id] }); } }); const validateMutation = useMutation({ mutationFn: async () => { if (!tree) throw new Error("No tree loaded."); return validateTreePayload({ tree }); }, onSuccess: ({ validation: nextValidation }) => { setValidation(nextValidation); } }); const deleteMutation = useMutation({ mutationFn: deleteTree, onSuccess: (_result, deletedId) => { queryClient.invalidateQueries({ queryKey: ["trees"] }); queryClient.removeQueries({ queryKey: ["tree", deletedId] }); if (selectedTreeId === deletedId) { setSelectedTreeId(null); setTree(null); setValidation(null); setSelectedNodeId(null); setAutoReadEnabled(false); } } }); const bulkDeleteMutation = useMutation({ mutationFn: deleteTrees, onSuccess: ({ deletedIds }) => { const deleted = new Set(deletedIds); queryClient.invalidateQueries({ queryKey: ["trees"] }); for (const deletedId of deleted) { queryClient.removeQueries({ queryKey: ["tree", deletedId] }); } if (selectedTreeId && deleted.has(selectedTreeId)) { setSelectedTreeId(null); setTree(null); setValidation(null); setSelectedNodeId(null); setAutoReadEnabled(false); } } }); const isDirty = Boolean(tree && stableTreeJson(tree) !== lastSavedJson); const selectedNode = useMemo( () => tree?.nodes.find((node) => node.id === selectedNodeId) ?? null, [selectedNodeId, tree] ); const isLiveStreamTree = Boolean(tree && liveStreamTreeId && tree.id === liveStreamTreeId); const hasActionableTree = Boolean(tree && (!isLiveStreamTree || tree.nodes.length > 0)); const updateNodes = (nodes: TrajectoryTreeNode[]) => { setTree((current) => (current ? { ...current, nodes } : current)); }; const updateNode = (nextNode: TrajectoryTreeNode) => { setTree((current) => current ? { ...current, nodes: current.nodes.map((node) => (node.id === nextNode.id ? nextNode : node)) } : current ); }; const handleCreated = (createdTree: TrajectoryTree) => { setCreatedReplayTreeId(createdTree.id); setSelectedTreeId(createdTree.id); setTree(createdTree); setValidation(null); setSelectedNodeId(null); setAutoReadEnabled(false); liveStreamTreeIdRef.current = null; setLiveStreamTreeId(null); setRevealCount(0); setIsPlaying(true); setLastSavedJson(stableTreeJson(createdTree)); queryClient.invalidateQueries({ queryKey: ["trees"] }); queryClient.invalidateQueries({ queryKey: ["tree", createdTree.id] }); }; const handleModelStreamStart = () => { const nextId = `model_stream_${Date.now()}`; const now = new Date().toISOString(); const nextTree: TrajectoryTree = { version: 1, id: nextId, title: "Model Stream", createdAt: now, updatedAt: now, raw_text: "", nodes: [] }; if (selectedTreeId) { void queryClient.cancelQueries({ queryKey: ["tree", selectedTreeId] }); } liveStreamTreeIdRef.current = nextId; setLiveStreamTreeId(nextId); setSelectedTreeId(null); setTree(nextTree); setValidation(null); setSelectedNodeId(null); setAutoReadEnabled(false); setRevealCount(0); setIsPlaying(false); setLastSavedJson(""); }; const handleModelStreamRawText = (rawText: string, done: boolean) => { const parse = parseTrajectoryText(rawText); const now = new Date().toISOString(); const streamTreeId = liveStreamTreeIdRef.current ?? liveStreamTreeId ?? `model_stream_${Date.now()}`; if (!liveStreamTreeIdRef.current) liveStreamTreeIdRef.current = streamTreeId; if (!liveStreamTreeId) setLiveStreamTreeId(streamTreeId); if (parse.nodes.length === 0) { setTree((current) => current && current.id === streamTreeId ? { ...current, raw_text: rawText, updatedAt: now } : current ); setValidation(parse.validation); return; } const previousCount = tree?.id === streamTreeId ? tree.nodes.length : 0; const nextTree: TrajectoryTree = { version: 1, id: streamTreeId, title: parse.title || "Model Stream", createdAt: tree?.id === streamTreeId ? tree.createdAt : now, updatedAt: now, raw_text: rawText, nodes: parse.nodes }; setTree(nextTree); setValidation(parse.validation); setRevealCount(parse.nodes.length); setIsPlaying(false); if (parse.nodes.length > previousCount || done) { const newest = parse.nodes[parse.nodes.length - 1]; if (newest) { setSelectedNodeId(newest.id); setFocusRequest({ nodeId: newest.id, mode: "node", nonce: Date.now() }); } } }; const handleModelStreamError = (message: string) => { setValidation((current) => ({ ok: false, issues: [ { level: "error", code: "model_stream_error", message }, ...(current?.issues ?? []) ], rootIds: current?.rootIds ?? [], finalNodeId: current?.finalNodeId ?? null })); }; const handleSave = () => { if (!tree) return; if (isLiveStreamTree) { setGeneratedSaveOpen(true); return; } saveMutation.mutate(undefined); }; const replay = () => { if (!tree) return; setSelectedNodeId(null); setAutoReadEnabled(false); setRevealCount(0); setIsPlaying(true); }; const step = () => { if (!tree) return; setAutoReadEnabled(false); setIsPlaying(false); setRevealCount((current) => Math.min(tree.nodes.length, current + 1)); }; const exportRawMessage = () => { if (!tree) return; const filenameBase = safeFilePart(tree.title || tree.id) || tree.id || "trajectory_tree"; const rawText = tree.raw_text?.trim() ? tree.raw_text : JSON.stringify( { title: tree.title, nodes: tree.nodes }, null, 2 ); downloadText(`${filenameBase}_raw_trajectory.txt`, rawText); }; const navigateToNode = (nodeId: string, mode: CompactNavigatorFocusMode) => { setAutoReadEnabled(false); setIsPlaying(false); setSelectedNodeId(nodeId); setFocusRequest({ nodeId, mode, nonce: Date.now() }); }; const toggleAutoRead = () => { if (!tree) return; setIsPlaying(false); setRevealCount(tree.nodes.length); setAutoReadEnabled((current) => !current); }; return ( { event.preventDefault(); setResizingPanel("rail"); }} onToggleDarkMode={() => setDarkMode((current) => !current)} onCreate={() => setCreateOpen(true)} onSelect={(treeId) => { liveStreamTreeIdRef.current = null; setLiveStreamTreeId(null); setSelectedTreeId(treeId); setSelectedNodeId(null); setAutoReadEnabled(false); }} onDelete={(treeId) => { if (window.confirm("Delete this saved tree?")) deleteMutation.mutate(treeId); }} onBulkDelete={(treeIds) => bulkDeleteMutation.mutate({ ids: treeIds })} onDeleteAll={() => bulkDeleteMutation.mutate({ all: true })} /> setCreateOpen(true)} onImport={() => setCreateOpen(true)} onModelStream={() => setStreamOpen(true)} onSave={handleSave} onExportRaw={exportRawMessage} onValidate={() => validateMutation.mutate()} onReplay={replay} onPause={() => setIsPlaying(false)} onStep={step} onFitView={() => setFitNonce((current) => current + 1)} onNodeViewModeChange={setNodeViewMode} /> setCreateOpen(true)} onImport={() => setCreateOpen(true)} /> setStreamOpen(false)} onStreamStart={handleModelStreamStart} onStreamRawText={handleModelStreamRawText} onStreamError={handleModelStreamError} /> setGeneratedSaveOpen(false)} onSave={(title) => saveMutation.mutate({ title })} /> {(treeQuery.isFetching || deleteMutation.isPending || bulkDeleteMutation.isPending) && ( Syncing tree )} { event.preventDefault(); setResizingPanel("inspector"); }} onUpdateNode={updateNode} onSaveTree={handleSave} onClose={() => setSelectedNodeId(null)} /> setCreateOpen(false)} onCreated={handleCreated} /> ); } export function TreeVisualizerPage() { const [client] = useState(() => new QueryClient()); return ( ); }
{tree.nodes.length} generated nodes