craft/src/state/PageContext.tsx

@@ -1,6 +1,7 @@
 import React, { createContext, useContext, useState, useCallback, useRef, ReactNode } from 'react';
 import { useEditor } from '@craftjs/core';
 import { PageData } from '../types';
+import { SerializedTreeNode } from '../types/sitesmith';
 import { useSiteDesign, SiteDesign } from './SiteDesignContext';
 
 interface PageContextValue {
@@ -19,6 +20,11 @@ interface PageContextValue {
   setHeaderCraftState: (craftState: string) => void;
   setFooterCraftState: (craftState: string) => void;
   setPagesCraftState: (pagesData: { id: string; name: string; slug: string; craftState: string | null }[]) => void;
+  /** AI helpers — replace entire site or page with a new tree */
+  replaceAllPages: (pages: { name: string; tree: SerializedTreeNode }[]) => void;
+  replaceCurrentPage: (page: { name: string; tree: SerializedTreeNode }) => void;
+  setHeader: (tree: SerializedTreeNode) => void;
+  setFooter: (tree: SerializedTreeNode) => void;
   siteDesign: SiteDesign;
 }
 
@@ -50,6 +56,10 @@ const PageContext = createContext<PageContextValue>({
   setHeaderCraftState: () => {},
   setFooterCraftState: () => {},
   setPagesCraftState: () => {},
+  replaceAllPages: () => {},
+  replaceCurrentPage: () => {},
+  setHeader: () => {},
+  setFooter: () => {},
   siteDesign: {} as SiteDesign,
 });
 
@@ -273,6 +283,110 @@ 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[] = [];
+      nodes[id] = {
+        type: node.type,
+        isCanvas: true,
+        props: node.props,
+        displayName: node.type.resolvedName,
+        custom: {},
+        hidden: false,
+        parent,
+        nodes: childIds,
+        linkedNodes: {},
+      };
+      for (const child of node.nodes ?? []) {
+        childIds.push(walk(child, id));
+      }
+      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;
+      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
+   * (the canvas still shows the currently active page — call switchPage() if needed).
+   */
+  const replaceAllPages = useCallback((newPages: { name: string; tree: SerializedTreeNode }[]) => {
+    if (newPages.length === 0) return;
+    const built = newPages.map((p, i) => ({
+      id: i === 0 ? 'home' : `page_${Date.now()}_${i}`,
+      name: p.name,
+      slug: slugify(p.name),
+      craftState: treeToState(p.tree),
+      headCode: '',
+    }));
+    setPages(built);
+    // Load the first page into the live canvas
+    const firstState = built[0].craftState;
+    setActivePageId(built[0].id);
+    activePageIdRef.current = built[0].id;
+    loadState(firstState, EMPTY_CANVAS);
+  }, [treeToState, 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 currentId = activePageIdRef.current;
+    if (currentId === HEADER_ID) {
+      setHeaderPage((prev) => ({ ...prev, name: page.name, craftState }));
+    } else if (currentId === FOOTER_ID) {
+      setFooterPage((prev) => ({ ...prev, name: page.name, craftState }));
+    } else {
+      setPages((prev) =>
+        prev.map((p) => (p.id === currentId ? { ...p, name: page.name, craftState } : p)),
+      );
+    }
+    loadState(craftState, EMPTY_CANVAS);
+  }, [treeToState, 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);
+    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]);
+
+  /**
+   * 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);
+    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]);
+
   return (
     <PageContext.Provider
       value={{
@@ -291,6 +405,10 @@ export const PageProvider: React.FC<{ children: ReactNode }> = ({ children }) =>
         setHeaderCraftState,
         setFooterCraftState,
         setPagesCraftState,
+        replaceAllPages,
+        replaceCurrentPage,
+        setHeader,
+        setFooter,
         siteDesign: design,
       }}
     >

craft/src/utils/apply-ai-response.test.ts

@@ -0,0 +1,90 @@
+import { describe, test, expect } from 'vitest';
+import { serializeTreeForCraft, __test } from './apply-ai-response';
+
+describe('serializeTreeForCraft', () => {
+  test('flattens nested tree', () => {
+    const tree = {
+      type: { resolvedName: 'Section' },
+      props: { aiName: 'Hero', node_id: 'ai-hero-1' },
+      nodes: [
+        {
+          type: { resolvedName: 'Heading' },
+          props: { aiName: 'Title', node_id: 'ai-h-1', text: 'Welcome' },
+          nodes: [],
+        },
+      ],
+    };
+    const out = serializeTreeForCraft(tree);
+    expect(out.rootNodeId).toBe('ai-hero-1');
+    expect((out.nodes['ai-hero-1'] as any).nodes).toEqual(['ai-h-1']);
+    expect((out.nodes['ai-h-1'] as any).parent).toBe('ROOT');
+  });
+
+  test('auto-generates ids when node_id is missing', () => {
+    const tree = { type: { resolvedName: 'Heading' }, props: {}, nodes: [] };
+    const out = serializeTreeForCraft(tree);
+    expect(typeof out.rootNodeId).toBe('string');
+    expect(out.nodes[out.rootNodeId]).toBeDefined();
+  });
+
+  test('sets isCanvas true for layout components', () => {
+    const tree = {
+      type: { resolvedName: 'Container' },
+      props: { node_id: 'c1' },
+      nodes: [],
+    };
+    const out = serializeTreeForCraft(tree);
+    expect((out.nodes['ROOT'] as any).isCanvas).toBe(true);
+  });
+
+  test('sets isCanvas false for leaf components', () => {
+    const tree = {
+      type: { resolvedName: 'Heading' },
+      props: { node_id: 'h1' },
+      nodes: [],
+    };
+    const out = serializeTreeForCraft(tree);
+    expect((out.nodes['ROOT'] as any).isCanvas).toBe(false);
+  });
+
+  test('aliases root node to ROOT key', () => {
+    const tree = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 'ai-section-1' },
+      nodes: [],
+    };
+    const out = serializeTreeForCraft(tree);
+    expect(out.nodes['ROOT']).toBeDefined();
+    expect((out.nodes['ROOT'] as any).parent).toBeNull();
+  });
+});
+
+describe('findNodeIdByAiNodeId', () => {
+  const query = {
+    getNodes: () => ({
+      'craft-id-1': { data: { props: { node_id: 'ai-hero-1' } } },
+      'craft-id-2': { data: { props: { node_id: 'ai-cta-1' } } },
+    }),
+  };
+
+  test('returns craft id for matching node_id prop', () => {
+    expect(__test.findNodeIdByAiNodeId(query, 'ai-hero-1')).toBe('craft-id-1');
+  });
+
+  test('returns craft id for second entry', () => {
+    expect(__test.findNodeIdByAiNodeId(query, 'ai-cta-1')).toBe('craft-id-2');
+  });
+
+  test('falls back to raw id match', () => {
+    const q = {
+      getNodes: () => ({
+        'exact-id': { data: { props: {} } },
+      }),
+    };
+    expect(__test.findNodeIdByAiNodeId(q, 'exact-id')).toBe('exact-id');
+  });
+
+  test('returns null when not found', () => {
+    expect(__test.findNodeIdByAiNodeId({ getNodes: () => ({}) }, 'missing')).toBe(null);
+  });
+});

craft/src/utils/apply-ai-response.ts

@@ -0,0 +1,219 @@
+import { useEditor } from '@craftjs/core';
+import type { NodeTree } from '@craftjs/core';
+import { usePages } from '../state/PageContext';
+import { SitesmithResponse, SerializedTreeNode } from '../types/sitesmith';
+
+/** Component types that act as drop targets (isCanvas = true) */
+const CANVAS_TYPES = new Set([
+  'Container', 'Section', 'ColumnLayout', 'BackgroundSection',
+  'HeroSimple', 'FeaturesGrid', 'CTASection',
+  'FormContainer', 'Navbar', 'Footer',
+]);
+
+/**
+ * Flatten a SerializedTreeNode tree into a Craft.js node map ready for
+ * `actions.deserialize()`.
+ *
+ * Returns `{ rootNodeId, nodes }` where `nodes` is a flat map keyed by node id.
+ * The root entry is also aliased under 'ROOT' so Craft.js can find it when
+ * calling `actions.deserialize(JSON.stringify(nodes))`.
+ */
+export function serializeTreeForCraft(tree: SerializedTreeNode): { rootNodeId: string; nodes: Record<string, unknown> } {
+  const idCounter = { n: 0 };
+  const nodes: Record<string, any> = {};
+
+  const walk = (node: SerializedTreeNode, parent: string | null): string => {
+    const id = (node.props.node_id as string | undefined) || `ai-auto-${idCounter.n++}`;
+    nodes[id] = {
+      type: node.type,
+      props: node.props,
+      displayName: node.type.resolvedName,
+      isCanvas: CANVAS_TYPES.has(node.type.resolvedName),
+      parent,
+      nodes: [] as string[],
+      hidden: false,
+      custom: {},
+      linkedNodes: {},
+    };
+    for (const child of node.nodes ?? []) {
+      const childId = walk(child, id);
+      nodes[id].nodes.push(childId);
+    }
+    return id;
+  };
+
+  const rootId = walk(tree, null);
+
+  // Craft.js frame expects a 'ROOT' key; alias it if the AI gave a different id
+  if (rootId !== 'ROOT') {
+    nodes['ROOT'] = { ...nodes[rootId], parent: null };
+    // Fix children's parent reference to 'ROOT'
+    for (const childId of nodes['ROOT'].nodes as string[]) {
+      if (nodes[childId]) nodes[childId].parent = 'ROOT';
+    }
+  }
+
+  return { rootNodeId: rootId, nodes };
+}
+
+/**
+ * Build a Craft.js `NodeTree` from a `SerializedTreeNode` using `query.parseFreshNode`.
+ * This is the correct way to construct a tree for `actions.addNodeTree()` when
+ * inserting/replacing sections or individual nodes.
+ */
+function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree {
+  const idCounter = { n: 0 };
+  const craftNodes: Record<string, any> = {};
+
+  const walk = (node: SerializedTreeNode, parent: string | null): string => {
+    const id = (node.props.node_id as string | undefined) || `ai-auto-${idCounter.n++}`;
+    const craftNode = (query.parseFreshNode({
+      id,
+      data: {
+        type: node.type,
+        props: node.props,
+        displayName: node.type.resolvedName,
+        isCanvas: CANVAS_TYPES.has(node.type.resolvedName),
+        parent,
+        nodes: [],
+        linkedNodes: {},
+        hidden: false,
+        custom: {},
+      },
+    }) as any).toNode() as any;
+    craftNodes[id] = craftNode;
+    for (const child of node.nodes ?? []) {
+      const childId = walk(child, id);
+      craftNodes[id].data.nodes.push(childId);
+    }
+    return id;
+  };
+
+  const rootId = walk(tree, null);
+  return { rootNodeId: rootId, nodes: craftNodes };
+}
+
+/**
+ * Find the Craft.js node id that corresponds to an AI node_id value.
+ * Checks `data.props.node_id` first, then falls back to raw id equality.
+ */
+export function findNodeIdByAiNodeId(query: any, aiNodeId: string): string | null {
+  const all = query.getNodes() as Record<string, any>;
+  for (const [id, n] of Object.entries(all)) {
+    if (n.data?.props?.node_id === aiNodeId) return id;
+    if (id === aiNodeId) return id;
+  }
+  return null;
+}
+
+/** Exported for unit tests */
+export const __test = { findNodeIdByAiNodeId };
+
+/**
+ * React hook that returns an `apply` function.
+ * Call `apply(response)` after a successful Sitesmith API call to materialize
+ * the AI's instructions into the editor.
+ */
+export function useApplyAiResponse() {
+  const { actions, query } = useEditor();
+  const pages = usePages();
+
+  return async function apply(resp: SitesmithResponse): Promise<{ ok: boolean; message?: string }> {
+    // 'ask' type = AI wants clarification, nothing to apply
+    if (resp.type === 'ask') return { ok: true };
+
+    if (resp.type === 'replace') {
+      if (resp.scope === 'site') {
+        pages.replaceAllPages(resp.pages.map((p) => ({ name: p.name, tree: p.tree })));
+        if (resp.header) pages.setHeader(resp.header.tree);
+        if (resp.footer) pages.setFooter(resp.footer.tree);
+        return { ok: true };
+      }
+
+      if (resp.scope === 'page') {
+        pages.replaceCurrentPage(resp.pages[0]);
+        return { ok: true };
+      }
+
+      if (resp.scope === 'section') {
+        // Insert each provided tree as a new node tree appended to ROOT
+        for (const p of resp.pages) {
+          try {
+            const nodeTree = buildNodeTree(query, p.tree);
+            actions.addNodeTree(nodeTree, 'ROOT');
+          } catch (e) {
+            console.warn('sitesmith: failed to add section tree', e);
+          }
+        }
+        return { ok: true };
+      }
+    }
+
+    if (resp.type === 'patch') {
+      return applyPatch(actions, query, resp.ops);
+    }
+
+    return { ok: false, message: 'Unknown response type' };
+  };
+}
+
+function applyPatch(
+  actions: any,
+  query: any,
+  ops: any[],
+): { ok: boolean; message?: string } {
+  for (const op of ops) {
+    const id = findNodeIdByAiNodeId(query, op.node_id);
+    if (!id) {
+      console.warn('sitesmith patch: node_id not found, skipping op:', op.node_id, op.op);
+      continue;
+    }
+
+    switch (op.op) {
+      case 'update_props':
+        actions.setProp(id, (p: any) => { Object.assign(p, op.props); });
+        break;
+
+      case 'replace_node': {
+        try {
+          const nodeTree = buildNodeTree(query, op.tree);
+          const parent: string = query.node(id).get().data.parent ?? 'ROOT';
+          const siblings: string[] = query.node(parent).childNodes();
+          const index = siblings.indexOf(id);
+          actions.delete(id);
+          actions.addNodeTree(nodeTree, parent, index);
+        } catch (e) {
+          console.warn('sitesmith patch: replace_node failed', e);
+        }
+        break;
+      }
+
+      case 'insert_after':
+      case 'insert_before': {
+        try {
+          const nodeTree = buildNodeTree(query, op.tree);
+          const parent: string = query.node(id).get().data.parent ?? 'ROOT';
+          const siblings: string[] = query.node(parent).childNodes();
+          const index = siblings.indexOf(id);
+          const at = op.op === 'insert_after' ? index + 1 : index;
+          actions.addNodeTree(nodeTree, parent, at);
+        } catch (e) {
+          console.warn(`sitesmith patch: ${op.op} failed`, e);
+        }
+        break;
+      }
+
+      case 'delete_node':
+        try {
+          actions.delete(id);
+        } catch (e) {
+          console.warn('sitesmith patch: delete_node failed', e);
+        }
+        break;
+
+      default:
+        console.warn('sitesmith patch: unknown op', (op as any).op);
+    }
+  }
+  return { ok: true };
+}