fix: validate API graph shape and add cycle guards + error boundary a…

src/components/ui/CodeWorkspace.tsx

@@ -4,6 +4,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import Diagram from "./Diagram";
 import { saveGraph } from "@/lib/graphs";
 import type { Graph } from "@/lib/analysis/types";
+import { graphSchema } from "@/lib/validation";
 
 const LANGUAGES = ["python", "javascript", "typescript", "go", "rust", "sql"];
 
@@ -297,7 +298,12 @@ export default function CodeWorkspace({
       });
       const data = await res.json();
       if (!res.ok) throw new Error(data.error ?? "Error");
-      setGraph(data as Graph);
+
+      const parsed = graphSchema.safeParse(data);
+      if (!parsed.success) {
+        throw new Error("Invalid graph data returned by the API.");
+      }
+      setGraph(parsed.data as Graph);
     } catch (err) {
       setError(err instanceof Error ? err.message : "Error");
       setGraph(null);

src/components/ui/Diagram.test.ts

@@ -0,0 +1,37 @@
+import { describe, expect, test } from "vitest";
+import { layout } from "./Diagram";
+import type { Graph } from "@/lib/analysis/types";
+
+describe("layout", () => {
+  test("handles basic linear parent hierarchy without crashing", () => {
+    const graph: Graph = {
+      nodes: [
+        { id: "root", label: "root", type: "module" },
+        { id: "child", label: "child", type: "function", parent: "root" },
+      ],
+      edges: [],
+    };
+    const result = layout(graph, false);
+    expect(result.nodes.length).toBe(2);
+  });
+
+  test("breaks out of parent reference cycles to avoid infinite loops and stack overflow", () => {
+    // Node A's parent is B, and Node B's parent is A
+    const graph: Graph = {
+      nodes: [
+        { id: "A", label: "Node A", type: "module", parent: "B" },
+        { id: "B", label: "Node B", type: "module", parent: "A" },
+      ],
+      edges: [],
+    };
+
+    // This should run quickly and complete without throwing "Maximum call stack size exceeded"
+    // or entering an infinite loop.
+    const startTime = Date.now();
+    const result = layout(graph, false);
+    const duration = Date.now() - startTime;
+
+    expect(duration).toBeLessThan(1000); // Must complete almost instantly
+    expect(result.nodes.length).toBe(2);
+  });
+});

src/components/ui/Diagram.tsx

@@ -22,6 +22,7 @@ import { toPng, toJpeg, toSvg } from "html-to-image";
 import dagre from "@dagrejs/dagre";
 import "reactflow/dist/style.css";
 import type { Graph, GraphNode, TableColumn } from "@/lib/analysis/types";
+import { ErrorBoundary } from "./ErrorBoundary";
 
 const NODE_W = 156;
 const NODE_H = 40;
@@ -140,7 +141,7 @@ const isContainer = (n: GraphNode) => n.type === "module" || n.type === "class";
 
 type Layout = { nodes: Node[]; edges: Edge[] };
 
-function layout(graph: Graph, dark: boolean): Layout {
+export function layout(graph: Graph, dark: boolean): Layout {
   const byId = new Map(graph.nodes.map((n) => [n.id, n]));
   const childrenOf = new Map<string, GraphNode[]>();
   const roots: GraphNode[] = [];
@@ -160,7 +161,12 @@ function layout(graph: Graph, dark: boolean): Layout {
   // Direct child of `container` that is, or contains, `nodeId`.
   function ancestorIn(nodeId: string, container: string): string | null {
     let cur: string | undefined = nodeId;
+    const visited = new Set<string>();
     while (cur) {
+      if (visited.has(cur)) {
+        return null;
+      }
+      visited.add(cur);
       const n = byId.get(cur);
       if (!n) return null;
       if (n.parent === container) return cur;
@@ -169,12 +175,25 @@ function layout(graph: Graph, dark: boolean): Layout {
     return null;
   }
 
+  const visitedSizeOf = new Set<string>();
+
   // Lay out a container's direct children (recursively) and return its size.
   function sizeOf(id: string): { w: number; h: number } {
-    const node = byId.get(id)!;
-    if (!isContainer(node)) return { w: NODE_W, h: NODE_H };
+    if (visitedSizeOf.has(id)) {
+      return { w: NODE_W, h: NODE_H };
+    }
+    visitedSizeOf.add(id);
+
+    const node = byId.get(id);
+    if (!node || !isContainer(node)) {
+      visitedSizeOf.delete(id);
+      return { w: NODE_W, h: NODE_H };
+    }
     const kids = childrenOf.get(id) ?? [];
-    if (kids.length === 0) return { w: 200, h: HEADER + 14 };
+    if (kids.length === 0) {
+      visitedSizeOf.delete(id);
+      return { w: 200, h: HEADER + 14 };
+    }
 
     const g = new dagre.graphlib.Graph();
     g.setGraph({ rankdir: "TB", nodesep: 24, ranksep: 40 });
@@ -207,6 +226,7 @@ function layout(graph: Graph, dark: boolean): Layout {
       maxX = Math.max(maxX, rx + s.w);
       maxY = Math.max(maxY, ry + s.h);
     }
+    visitedSizeOf.delete(id);
     return { w: maxX + PAD, h: maxY + PAD };
   }
 
@@ -222,7 +242,12 @@ function layout(graph: Graph, dark: boolean): Layout {
   // Top-level layout of modules (imports + cross-module calls/extends).
   const topAncestor = (nodeId: string): string | null => {
     let cur: string | undefined = nodeId;
+    const visited = new Set<string>();
     while (cur) {
+      if (visited.has(cur)) {
+        return null;
+      }
+      visited.add(cur);
       const n = byId.get(cur);
       if (!n) return null;
       if (!n.parent) return cur;
@@ -252,7 +277,12 @@ function layout(graph: Graph, dark: boolean): Layout {
   const depthOf = (n: GraphNode): number => {
     let d = 0;
     let cur = n.parent;
+    const visited = new Set<string>();
     while (cur) {
+      if (visited.has(cur)) {
+        break;
+      }
+      visited.add(cur);
       d++;
       cur = byId.get(cur)?.parent;
     }
@@ -647,8 +677,10 @@ function DiagramFlow({
 
 export default function Diagram(props: { graph: Graph; emptyLabel: string }) {
   return (
-    <ReactFlowProvider>
-      <DiagramFlow {...props} />
-    </ReactFlowProvider>
+    <ErrorBoundary>
+      <ReactFlowProvider>
+        <DiagramFlow {...props} />
+      </ReactFlowProvider>
+    </ErrorBoundary>
   );
 }

src/components/ui/ErrorBoundary.tsx

@@ -0,0 +1,72 @@
+"use client";
+
+import React, { Component, type ErrorInfo, type ReactNode } from "react";
+
+interface Props {
+  children: ReactNode;
+  fallback?: ReactNode;
+}
+
+interface State {
+  hasError: boolean;
+  error: Error | null;
+}
+
+export class ErrorBoundary extends Component<Props, State> {
+  public state: State = {
+    hasError: false,
+    error: null,
+  };
+
+  public static getDerivedStateFromError(error: Error): State {
+    return { hasError: true, error };
+  }
+
+  public componentDidCatch(error: Error, errorInfo: ErrorInfo) {
+    console.error("ErrorBoundary caught an error:", error, errorInfo);
+  }
+
+  public render() {
+    if (this.state.hasError) {
+      if (this.props.fallback) {
+        return this.props.fallback;
+      }
+      return (
+        <div className="flex h-full w-full flex-col items-center justify-center p-6 text-center bg-[#f6f7f9] dark:bg-[#0b0d12]">
+          <div className="max-w-md rounded-xl border border-red-100 dark:border-red-950 bg-white dark:bg-[#12151c] p-6 shadow-sm">
+            <div className="mx-auto mb-3 flex h-10 w-10 items-center justify-center rounded-full bg-red-100 dark:bg-red-950 text-red-600 dark:text-red-400">
+              <svg
+                width="20"
+                height="20"
+                viewBox="0 0 24 24"
+                fill="none"
+                stroke="currentColor"
+                strokeWidth="2"
+                strokeLinecap="round"
+                strokeLinejoin="round"
+              >
+                <circle cx="12" cy="12" r="10" />
+                <line x1="12" y1="8" x2="12" y2="12" />
+                <line x1="12" y1="16" x2="12.01" y2="16" />
+              </svg>
+            </div>
+            <h3 className="text-md font-semibold text-[#0f172a] dark:text-[#e6e9ef] mb-1.5">
+              Failed to render diagram
+            </h3>
+            <p className="text-[13px] leading-relaxed text-[#475569] dark:text-[#9aa6b8] mb-4">
+              The diagram contains malformed or cyclic data and could not be displayed correctly.
+            </p>
+            <button
+              onClick={() => this.setState({ hasError: false, error: null })}
+              className="rounded-lg bg-[#4f46e5] px-3.5 py-2 text-xs font-medium text-white hover:bg-[#4338ca] transition-colors"
+            >
+              Try Again
+            </button>
+          </div>
+        </div>
+      );
+    }
+
+    return this.props.children;
+  }
+}

src/test/mocks/styleMock.css

@@ -0,0 +1 @@
+/* Mock CSS file for tests */

src/test/mocks/styleMock.js

@@ -0,0 +1,2 @@
+// Mock module for CSS files
+export default {};

vitest.config.ts

@@ -7,6 +7,9 @@ export default defineConfig({
     include: ["src/**/*.test.ts"],
   },
   resolve: {
-    alias: { "@": fileURLToPath(new URL("./src", import.meta.url)) },
+    alias: {
+      "@": fileURLToPath(new URL("./src", import.meta.url)),
+      "reactflow/dist/style.css": fileURLToPath(new URL("./src/test/mocks/styleMock.js", import.meta.url)),
+    },
   },
 });