craft/src/state/PageContext.treeToState.test.ts

@@ -0,0 +1,117 @@
+import { describe, test, expect, vi } from 'vitest';
+import type { SerializedTreeNode } from '../types/sitesmith';
+import { treeToCraftState } from './PageContext';
+
+/**
+ * E3: treeToState (here, the extracted pure `treeToCraftState`) previously had
+ * NO resolvedName allowlist — an AI `replace` response (scope site/page,
+ * reaches this function via `actions.deserialize`) containing an unknown
+ * component would produce a state that could throw at render/deserialize
+ * time. This mirrors the guard `buildNodeTree` (apply-ai-response.ts)
+ * already had.
+ *
+ * Tested directly against the exported pure function (no React/Craft.js
+ * Editor needed) per the brief's guidance to extract+test the core
+ * transform rather than exercise this through the PageProvider.
+ */
+describe('treeToCraftState resolvedName guard', () => {
+  test('an unknown resolvedName at the root does not throw and falls back to a valid empty canvas', () => {
+    const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const tree: SerializedTreeNode = {
+      type: { resolvedName: 'NotARealComponent' },
+      props: { node_id: 'bad-root' },
+      nodes: [],
+    };
+    let state = '';
+    expect(() => {
+      state = treeToCraftState(tree);
+    }).not.toThrow();
+
+    const parsed = JSON.parse(state);
+    expect(parsed.ROOT).toBeDefined();
+    expect(parsed.ROOT.type.resolvedName).toBe('Container');
+    expect(parsed.ROOT.nodes).toEqual([]);
+    expect(warnSpy).toHaveBeenCalled();
+    warnSpy.mockRestore();
+  });
+
+  test('an unknown resolvedName on a nested child is soft-skipped; valid siblings still render', () => {
+    const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const tree: SerializedTreeNode = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 'sec-1' },
+      nodes: [
+        { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] },
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'h-1', text: 'Welcome' }, nodes: [] },
+      ],
+    };
+    const state = treeToCraftState(tree);
+    const parsed = JSON.parse(state);
+
+    // The unknown node must not appear anywhere in the produced state.
+    for (const id of Object.keys(parsed)) {
+      expect(parsed[id].type.resolvedName).not.toBe('NotARealComponent');
+    }
+    // The Heading sibling survived (wrapped in Section's section-inner).
+    expect(parsed['h-1']).toBeDefined();
+    expect(parsed['h-1'].props.text).toBe('Welcome');
+    expect(warnSpy).toHaveBeenCalled();
+    warnSpy.mockRestore();
+  });
+
+  test('a fully-valid tree still produces a working ROOT-keyed state (regression)', () => {
+    const tree: SerializedTreeNode = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 'sec-1' },
+      nodes: [
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'h-1', text: 'Hi' }, nodes: [] },
+      ],
+    };
+    const state = treeToCraftState(tree);
+    const parsed = JSON.parse(state);
+
+    expect(parsed.ROOT).toBeDefined();
+    expect(parsed.ROOT.type.resolvedName).toBe('Section');
+    expect(parsed.ROOT.isCanvas).toBe(true);
+    expect(parsed.ROOT.parent).toBeNull();
+    const innerId = parsed.ROOT.linkedNodes['section-inner'];
+    expect(innerId).toBeDefined();
+    expect(parsed[innerId].nodes).toEqual(['h-1']);
+    expect(parsed['h-1'].parent).toBe(innerId);
+    expect(parsed['h-1'].props.text).toBe('Hi');
+  });
+
+  test('a ColumnLayout root still flattens children into col-N linkedNodes (regression)', () => {
+    const tree: SerializedTreeNode = {
+      type: { resolvedName: 'ColumnLayout' },
+      props: { node_id: 'cols-1' },
+      nodes: [
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'a' }, nodes: [] },
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'b' }, nodes: [] },
+      ],
+    };
+    const state = treeToCraftState(tree);
+    const parsed = JSON.parse(state);
+
+    expect(parsed.ROOT.type.resolvedName).toBe('ColumnLayout');
+    expect(parsed.ROOT.linkedNodes).toEqual({ 'col-0': 'a', 'col-1': 'b' });
+    expect(parsed.ROOT.props.columns).toBe(2);
+    // NB: the ROOT-aliasing step only reassigns `nodes[]` children's parent,
+    // not `linkedNodes` children's — a pre-existing quirk (not introduced by
+    // this refactor, and out of scope for E3) that leaves column children's
+    // `parent` pointing at the tree's original (pre-alias) root id.
+    expect(parsed['a'].parent).toBe('cols-1');
+    expect(parsed['a'].isCanvas).toBe(true);
+  });
+
+  test('style: [] normalizes to {} (regression)', () => {
+    const tree: SerializedTreeNode = {
+      type: { resolvedName: 'Heading' },
+      props: { node_id: 'h-1', style: [] as unknown as Record<string, unknown> },
+      nodes: [],
+    };
+    const state = treeToCraftState(tree);
+    const parsed = JSON.parse(state);
+    expect(parsed.ROOT.props.style).toEqual({});
+  });
+});

craft/src/state/PageContext.tsx

@@ -3,26 +3,7 @@ import { useEditor } from '@craftjs/core';
 import { PageData } from '../types';
 import { SerializedTreeNode } from '../types/sitesmith';
 import { useSiteDesign, SiteDesign } from './SiteDesignContext';
-
+import { sanitizeAiTree, flattenTreeForCraft, FlatCraftNode } from '../utils/craft-tree';
-/** Only `Container` instances are "real" canvases in serialized state — they
- * directly render whatever is in node.data.nodes. Layout-shell components
- * (Section, HeroSimple, FeaturesGrid, ColumnLayout, CTASection, etc) use
- * Craft.js <Element canvas id="…"> linkedNodes internally; their own
- * isCanvas must be FALSE or Craft.js's toNodeTree walker trips an Invariant
- * because the shell claims to be a canvas but its render ignores `nodes`. */
-const CANVAS_TYPES = new Set<string>(['Container']);
-
-/** Shells that wrap their content in a single <Element id="<key>" is={Container}>.
- * When the AI puts content directly under one of these, the children end up
- * orphaned (the shell ignores data.nodes — it renders via the linkedNode) and
- * Craft.js auto-creates the linkedNode at render time with a botched type
- * field, which then crashes toNodeTree. Pre-create the linkedNode ourselves
- * to keep the state shape Craft.js expects. */
-const SHELL_INNER: Record<string, string> = {
-  Section: 'section-inner',
-  BackgroundSection: 'bg-section-inner',
-  FormContainer: 'form-inner',
-};
 
 interface PageContextValue {
   pages: PageData[];
@@ -60,12 +41,50 @@ const EMPTY_HEADER =
 const EMPTY_FOOTER =
   '{"ROOT":{"type":{"resolvedName":"Container"},"isCanvas":true,"props":{"style":{"minHeight":"60px","backgroundColor":"#0f172a","color":"#94a3b8","padding":"40px 24px","textAlign":"center"},"tag":"footer"},"displayName":"Container","custom":{},"hidden":false,"nodes":[],"linkedNodes":{}}}';
 
+/**
+ * Flatten a `SerializedTreeNode` (as produced by the AI, a template, etc.)
+ * into a Craft.js `SerializedNodes` JSON string ready for
+ * `actions.deserialize()`.
+ *
+ * Untrusted input is validated the same way `apply-ai-response.ts`'s
+ * `buildNodeTree` validates AI `patch`/section-replace trees: any node whose
+ * `type.resolvedName` isn't a registered component is dropped (subtree and
+ * all) via `sanitizeAiTree`, and a bad/colliding/`'ROOT'` id is regenerated
+ * — never left as-is. This matters here specifically because `replace`
+ * scope `site`/`page` responses reach this function via `actions.deserialize`
+ * with no further validation downstream, unlike the `buildNodeTree` path
+ * which also gets Craft.js's own `parseFreshNode` as a second line of
+ * defense. If the ROOT node itself is invalid, sanitizeAiTree returns null
+ * and we fall back to an empty canvas rather than handing deserialize()
+ * something that could throw.
+ *
+ * Exported standalone (not a hook) so it's directly unit-testable without
+ * mounting a Craft.js `<Editor>`.
+ */
+export function treeToCraftState(tree: SerializedTreeNode): string {
+  const sanitized = sanitizeAiTree(tree, new Set());
+  if (!sanitized) return EMPTY_CANVAS;
+
+  const { rootNodeId, nodes: flatNodes } = flattenTreeForCraft(sanitized);
+  const nodes: Record<string, FlatCraftNode> = flatNodes;
+
+  // Craft.js deserialize requires the root node keyed as 'ROOT'.
+  if (rootNodeId !== 'ROOT') {
+    nodes['ROOT'] = { ...nodes[rootNodeId], parent: null, isCanvas: true };
+    delete nodes[rootNodeId];
+    for (const childId of nodes['ROOT'].nodes) {
+      if (nodes[childId]) nodes[childId].parent = 'ROOT';
+    }
+  }
+  return JSON.stringify(nodes);
+}
+
 // Default header seed: a ROOT header Container holding a Navbar with default
 // links. New sites previously opened with an EMPTY header (just a bare
 // Container), so there was no menu to edit and the empty zone rendered as a
 // stray band above the page. Seeding a real Navbar gives every new site an
 // editable menu-with-links out of the box (and removes the empty-header gap).
-// Node shape matches serializeTreeForCraft() / Craft's actions.deserialize().
+// Node shape matches treeToCraftState() / Craft's actions.deserialize().
 export const DEFAULT_HEADER_STATE = JSON.stringify({
   ROOT: {
     type: { resolvedName: 'Container' },
@@ -379,101 +398,6 @@ export const PageProvider: React.FC<{ children: ReactNode }> = ({ children }) =>
     })));
   }, []);
 
-  /** Flatten a SerializedTreeNode into a Craft.js SerializedNodes JSON string */
-  const treeToState = useCallback((tree: SerializedTreeNode): string => {
-    let counter = 0;
-    const nodes: Record<string, unknown> = {};
-    const walk = (node: SerializedTreeNode, parent: string | null): string => {
-      const id = (node.props.node_id as string | undefined) || `ai-auto-${counter++}`;
-      const childIds: string[] = [];
-      const typeName = node.type?.resolvedName;
-      // Normalize props: the AI sometimes emits `style: []` instead of `{}`.
-      // React/Craft.js choke when a CSSProperties slot is an array — normalize it.
-      const rawProps = node.props ?? {};
-      const props: Record<string, unknown> = { ...rawProps };
-      if (Array.isArray(props.style)) props.style = {};
-      nodes[id] = {
-        type: node.type,
-        // isCanvas must match the component's craft.rules — only layout
-        // wrappers accept children. Setting it true on leaf components
-        // (Heading, TextBlock, ButtonLink, etc) makes Craft.js render them
-        // as empty drop-canvas wrappers and the actual content disappears.
-        isCanvas: typeName ? CANVAS_TYPES.has(typeName) : false,
-        props,
-        displayName: typeName,
-        custom: {},
-        hidden: false,
-        parent,
-        nodes: childIds,
-        linkedNodes: {},
-      };
-      for (const child of node.nodes ?? []) {
-        childIds.push(walk(child, id));
-      }
-      // ColumnLayout uses Craft.js linkedNodes with fixed ids (col-0, col-1, ...).
-      // The AI emits children as direct `nodes`, but ColumnLayout's render ignores
-      // them and creates fresh column Elements — the AI's children become orphans
-      // and any subsequent toNodeTree walk hits an Invariant. Move direct children
-      // into linkedNodes so they render in the columns the user actually sees.
-      if (typeName === 'ColumnLayout' && childIds.length > 0) {
-        const linked: Record<string, string> = {};
-        childIds.forEach((cid, i) => {
-          linked[`col-${i}`] = cid;
-          if (nodes[cid]) (nodes[cid] as any).isCanvas = true; // columns are canvases
-        });
-        (nodes[id] as any).nodes = [];
-        (nodes[id] as any).linkedNodes = linked;
-        // Reflect the actual column count on the component so its render matches.
-        const cur = (nodes[id] as any).props || {};
-        if (!cur.columns || cur.columns !== childIds.length) cur.columns = childIds.length;
-        (nodes[id] as any).props = cur;
-      }
-      // Section/BackgroundSection/FormContainer each render a single
-      // <Element id="<key>" is={Container} canvas> ... </Element>. If the AI
-      // nests content as direct children, Craft.js will auto-create the
-      // linkedNode on first render — and store its type as the Container
-      // component class rather than {resolvedName:'Container'}, which then
-      // crashes toNodeTree with "type (undefined) does not exist in resolver".
-      // Pre-create the linkedNode ourselves with the correct serialized type
-      // so Craft.js never has to materialize it.
-      const innerKey = SHELL_INNER[typeName ?? ''];
-      if (innerKey && childIds.length > 0) {
-        const innerId = `${id}__${innerKey}`;
-        nodes[innerId] = {
-          type: { resolvedName: 'Container' },
-          isCanvas: true,
-          props: { tag: 'div' },
-          displayName: 'Container',
-          custom: {},
-          hidden: false,
-          parent: id,
-          nodes: [...childIds],
-          linkedNodes: {},
-        };
-        for (const cid of childIds) {
-          if (nodes[cid]) (nodes[cid] as any).parent = innerId;
-        }
-        (nodes[id] as any).nodes = [];
-        (nodes[id] as any).linkedNodes = { [innerKey]: innerId };
-      }
-      return id;
-    };
-    const rootId = walk(tree, null);
-    // Craft.js deserialize requires the root node keyed as 'ROOT'
-    if (rootId !== 'ROOT') {
-      nodes['ROOT'] = nodes[rootId];
-      (nodes['ROOT'] as any).parent = null;
-      // ROOT must be a canvas regardless of component type so children render.
-      (nodes['ROOT'] as any).isCanvas = true;
-      delete nodes[rootId];
-      // Fix up parent references from ROOT's children
-      for (const childId of (nodes['ROOT'] as any).nodes) {
-        if (nodes[childId]) (nodes[childId] as any).parent = 'ROOT';
-      }
-    }
-    return JSON.stringify(nodes);
-  }, []);
-
   /**
    * AI helper: replace all pages with newly generated trees.
    * Stores each page's serialized state without touching the live canvas
@@ -494,7 +418,7 @@ export const PageProvider: React.FC<{ children: ReactNode }> = ({ children }) =>
         id: i === 0 ? 'home' : `page_${Date.now()}_${i}`,
         name: p.name,
         slug,
-        craftState: treeToState(p.tree),
+        craftState: treeToCraftState(p.tree),
       };
     });
     setPages(built);
@@ -503,14 +427,14 @@ export const PageProvider: React.FC<{ children: ReactNode }> = ({ children }) =>
     setActivePageId(built[0].id);
     activePageIdRef.current = built[0].id;
     loadState(firstState, EMPTY_CANVAS);
-  }, [treeToState, loadState]);
+  }, [loadState]);
 
   /**
    * AI helper: replace the current page's tree.
    * Deserializes the new tree into the live Craft.js canvas and persists it.
    */
   const replaceCurrentPage = useCallback((page: { name: string; tree: SerializedTreeNode }) => {
-    const craftState = treeToState(page.tree);
+    const craftState = treeToCraftState(page.tree);
     const currentId = activePageIdRef.current;
     if (currentId === HEADER_ID) {
       setHeaderPage((prev) => ({ ...prev, name: page.name, craftState }));
@@ -522,33 +446,33 @@ export const PageProvider: React.FC<{ children: ReactNode }> = ({ children }) =>
       );
     }
     loadState(craftState, EMPTY_CANVAS);
-  }, [treeToState, loadState]);
+  }, [loadState]);
 
   /**
    * AI helper: replace the shared header tree.
    * Updates stored state; does NOT switch the canvas to header view.
    */
   const setHeader = useCallback((tree: SerializedTreeNode) => {
-    const craftState = treeToState(tree);
+    const craftState = treeToCraftState(tree);
     setHeaderPage((prev) => ({ ...prev, craftState }));
     // If the canvas is currently showing the header, refresh it live
     if (activePageIdRef.current === HEADER_ID) {
       loadState(craftState, EMPTY_HEADER);
     }
-  }, [treeToState, loadState]);
+  }, [loadState]);
 
   /**
    * AI helper: replace the shared footer tree.
    * Updates stored state; does NOT switch the canvas to footer view.
    */
   const setFooter = useCallback((tree: SerializedTreeNode) => {
-    const craftState = treeToState(tree);
+    const craftState = treeToCraftState(tree);
     setFooterPage((prev) => ({ ...prev, craftState }));
     // If the canvas is currently showing the footer, refresh it live
     if (activePageIdRef.current === FOOTER_ID) {
       loadState(craftState, EMPTY_FOOTER);
     }
-  }, [treeToState, loadState]);
+  }, [loadState]);
 
   return (
     <PageContext.Provider