From 3e75f1aafa5f356e426f06583d281ad2c23c1f92 Mon Sep 17 00:00:00 2001 From: YasserYG8 Date: Sun, 28 Jun 2026 06:41:29 +0100 Subject: [PATCH] fix: validate API graph shape and add cycle guards + error boundary around diagram --- src/components/ui/CodeWorkspace.tsx | 8 +++- src/components/ui/Diagram.test.ts | 37 +++++++++++++++ src/components/ui/Diagram.tsx | 46 +++++++++++++++--- src/components/ui/ErrorBoundary.tsx | 72 +++++++++++++++++++++++++++++ src/test/mocks/styleMock.css | 1 + src/test/mocks/styleMock.js | 2 + vitest.config.ts | 5 +- 7 files changed, 162 insertions(+), 9 deletions(-) create mode 100644 src/components/ui/Diagram.test.ts create mode 100644 src/components/ui/ErrorBoundary.tsx create mode 100644 src/test/mocks/styleMock.css create mode 100644 src/test/mocks/styleMock.js diff --git a/src/components/ui/CodeWorkspace.tsx b/src/components/ui/CodeWorkspace.tsx index 7bc23bb..42dd6be 100644 --- a/src/components/ui/CodeWorkspace.tsx +++ b/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); diff --git a/src/components/ui/Diagram.test.ts b/src/components/ui/Diagram.test.ts new file mode 100644 index 0000000..bae19ce --- /dev/null +++ b/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); + }); +}); diff --git a/src/components/ui/Diagram.tsx b/src/components/ui/Diagram.tsx index 1af36ff..be968ae 100644 --- a/src/components/ui/Diagram.tsx +++ b/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(); 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(); 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(); + // 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(); 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(); 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 ( - - - + + + + + ); } diff --git a/src/components/ui/ErrorBoundary.tsx b/src/components/ui/ErrorBoundary.tsx new file mode 100644 index 0000000..1a65111 --- /dev/null +++ b/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 { + 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 ( +
+
+
+ + + + + +
+

+ Failed to render diagram +

+

+ The diagram contains malformed or cyclic data and could not be displayed correctly. +

+ +
+
+ ); + } + + return this.props.children; + } +} diff --git a/src/test/mocks/styleMock.css b/src/test/mocks/styleMock.css new file mode 100644 index 0000000..f52b664 --- /dev/null +++ b/src/test/mocks/styleMock.css @@ -0,0 +1 @@ +/* Mock CSS file for tests */ diff --git a/src/test/mocks/styleMock.js b/src/test/mocks/styleMock.js new file mode 100644 index 0000000..71cb127 --- /dev/null +++ b/src/test/mocks/styleMock.js @@ -0,0 +1,2 @@ +// Mock module for CSS files +export default {}; diff --git a/vitest.config.ts b/vitest.config.ts index 93ec33a..1a57e25 100644 --- a/vitest.config.ts +++ b/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)), + }, }, });