fix(builder): validate AI response before deserializing into craft state

Adds an AI-boundary validation/sanitization pass in apply-ai-response.ts so malformed or malicious Sitesmith output can never crash addNodeTree or corrupt live craft state: - resolvedName allowlist: any node whose type.resolvedName is not a registered component in componentResolver is dropped (with its subtree) and a console.warn is logged; buildNodeTree throws before touching Craft.js if the tree ROOT itself is invalid, which existing call sites already catch and warn on (fail soft, never throws into addNodeTree). - update_props protected-key list (node_id) — an AI-supplied node_id can no longer overwrite the target node's real id; style is merged shallowly instead of replaced wholesale so unrelated existing style keys survive. - id policy: regenerate, never skip. Any AI-supplied node id that is 'ROOT', empty/non-string, or collides with an existing/already-used id is replaced with a fresh ai-auto-N id; the node itself is kept. Extends apply-ai-response.test.ts with coverage for all of the above plus regression tests proving fully-valid AI responses still apply exactly as before. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-12 12:37:25 -07:00
parent 4b36ce0d6a
commit b46b6915a5
2 changed files with 298 additions and 8 deletions
@@ -1,5 +1,5 @@
-import { describe, test, expect } from 'vitest';
-import { serializeTreeForCraft, __test } from './apply-ai-response';
+import { describe, test, expect, vi } from 'vitest';
+import { serializeTreeForCraft, sanitizeAiTree, __test } from './apply-ai-response';

 describe('serializeTreeForCraft', () => {
  test('flattens nested tree', () => {
@@ -88,3 +88,180 @@ describe('findNodeIdByAiNodeId', () => {
    expect(__test.findNodeIdByAiNodeId({ getNodes: () => ({}) }, 'missing')).toBe(null);
  });
 });
+
+// ---------------------------------------------------------------------------
+// Helpers for the tests below: fake @craftjs/core `query` / `actions` objects
+// good enough to exercise buildNodeTree/applyPatch without a real Editor.
+// ---------------------------------------------------------------------------
+
+function makeFakeQuery(existingIds: string[] = ['ROOT']) {
+  return {
+    getNodes: () => Object.fromEntries(existingIds.map((id) => [id, { data: { props: {} } }])),
+    parseFreshNode: (input: any) => ({ toNode: () => input }),
+    node: (id: string) => ({
+      get: () => ({ data: { parent: 'ROOT' } }),
+      childNodes: () => [id],
+    }),
+  };
+}
+
+/** A fake query+actions pair backed by a single existing node, wired so
+ * findNodeIdByAiNodeId can resolve `aiNodeId` back to `craftId`. */
+function makeSingleNodeHarness(craftId: string, aiNodeId: string, initialProps: Record<string, unknown>) {
+  const propsById: Record<string, any> = { [craftId]: { node_id: aiNodeId, ...initialProps } };
+  const query = {
+    getNodes: () =>
+      Object.fromEntries(Object.entries(propsById).map(([id, p]) => [id, { data: { props: p } }])),
+    parseFreshNode: (input: any) => ({ toNode: () => input }),
+    node: (id: string) => ({
+      get: () => ({ data: { parent: 'ROOT' } }),
+      childNodes: () => [craftId],
+    }),
+  };
+  const actions = {
+    setProp: (id: string, fn: (p: any) => void) => {
+      if (propsById[id]) fn(propsById[id]);
+    },
+    delete: (id: string) => {
+      delete propsById[id];
+    },
+    addNodeTree: () => {},
+  };
+  return { query, actions, propsById };
+}
+
+describe('sanitizeAiTree', () => {
+  test('unknown resolvedName at the root is rejected (returns null)', () => {
+    const tree = { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] };
+    expect(sanitizeAiTree(tree, new Set())).toBeNull();
+  });
+
+  test('unknown resolvedName on a nested child drops only that subtree; valid siblings survive', () => {
+    const tree = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 's1' },
+      nodes: [
+        { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] },
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'h1' }, nodes: [] },
+      ],
+    };
+    const out = sanitizeAiTree(tree, new Set());
+    expect(out).not.toBeNull();
+    expect(out!.nodes!.length).toBe(1);
+    expect(out!.nodes![0].props.node_id).toBe('h1');
+  });
+
+  test('a node id of "ROOT" is regenerated, never left as ROOT', () => {
+    const tree = { type: { resolvedName: 'Heading' }, props: { node_id: 'ROOT' }, nodes: [] };
+    const out = sanitizeAiTree(tree, new Set(['ROOT']));
+    expect(out).not.toBeNull();
+    expect(out!.props.node_id).not.toBe('ROOT');
+  });
+
+  test('an empty/non-string id is regenerated to a non-empty string', () => {
+    const tree = { type: { resolvedName: 'Heading' }, props: { node_id: '' }, nodes: [] };
+    const out = sanitizeAiTree(tree, new Set());
+    expect(typeof out!.props.node_id).toBe('string');
+    expect(out!.props.node_id).toBeTruthy();
+  });
+
+  test('an id colliding with an existing node is regenerated, not reused', () => {
+    const tree = { type: { resolvedName: 'Heading' }, props: { node_id: 'existing-1' }, nodes: [] };
+    const out = sanitizeAiTree(tree, new Set(['existing-1']));
+    expect(out!.props.node_id).not.toBe('existing-1');
+  });
+
+  test('duplicate ids within the same tree are de-duplicated (no two nodes share an id)', () => {
+    const tree = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 'dup-1' },
+      nodes: [
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'dup-1' }, nodes: [] },
+      ],
+    };
+    const out = sanitizeAiTree(tree, new Set());
+    expect(out!.props.node_id).toBe('dup-1');
+    expect(out!.nodes![0].props.node_id).not.toBe('dup-1');
+  });
+
+  test('a fully-valid tree passes through with structure and ids preserved (regression)', () => {
+    const tree = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 'sec-1', aiName: 'Hero' },
+      nodes: [
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'h-1', text: 'Welcome' }, nodes: [] },
+      ],
+    };
+    const out = sanitizeAiTree(tree, new Set());
+    expect(out!.props.node_id).toBe('sec-1');
+    expect(out!.nodes!.length).toBe(1);
+    expect(out!.nodes![0].props.node_id).toBe('h-1');
+    expect(out!.nodes![0].props.text).toBe('Welcome');
+  });
+});
+
+describe('buildNodeTree', () => {
+  test('builds a NodeTree preserving structure for a fully valid AI tree (regression)', () => {
+    const tree = {
+      type: { resolvedName: 'Section' },
+      props: { node_id: 'sec-1' },
+      nodes: [
+        { type: { resolvedName: 'Heading' }, props: { node_id: 'h-1', text: 'Hi' }, nodes: [] },
+      ],
+    };
+    const result = __test.buildNodeTree(makeFakeQuery(), tree);
+    expect(result.rootNodeId).toBe('sec-1');
+    expect(result.nodes['sec-1']).toBeDefined();
+    // Section wraps direct children in a pre-created section-inner linkedNode.
+    const innerId = result.nodes['sec-1'].data.linkedNodes['section-inner'];
+    expect(innerId).toBeDefined();
+    expect(result.nodes[innerId].data.nodes).toEqual(['h-1']);
+    expect(result.nodes['h-1'].data.parent).toBe(innerId);
+  });
+
+  test('throws when the tree root has an unknown resolvedName (caught by callers, never reaches addNodeTree)', () => {
+    const tree = { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] };
+    expect(() => __test.buildNodeTree(makeFakeQuery(), tree)).toThrow();
+  });
+});
+
+describe('applyPatch', () => {
+  test('update_props: does not overwrite node_id, applies a legit prop, merges style shallowly', () => {
+    const { query, actions, propsById } = makeSingleNodeHarness('craft-1', 'ai-node-1', {
+      text: 'Old',
+      style: { color: 'red', fontSize: '16px' },
+    });
+    const ops = [
+      {
+        op: 'update_props',
+        node_id: 'ai-node-1',
+        props: { node_id: 'ai-hacked', text: 'New', style: { color: 'blue' } },
+      },
+    ];
+    const result = __test.applyPatch(actions, query, ops);
+    expect(result.ok).toBe(true);
+    expect(propsById['craft-1'].node_id).toBe('ai-node-1'); // protected, unchanged
+    expect(propsById['craft-1'].text).toBe('New'); // legit prop applied
+    // style merges shallowly: color overwritten, fontSize preserved
+    expect(propsById['craft-1'].style).toEqual({ color: 'blue', fontSize: '16px' });
+  });
+
+  test('an op with an unknown resolvedName tree is skipped without throwing; other ops in the batch still apply', () => {
+    const warnSpy = vi.spyOn(console, 'warn').mockImplementation(() => {});
+    const { query, actions, propsById } = makeSingleNodeHarness('craft-1', 'ai-node-1', {});
+    const ops = [
+      {
+        op: 'insert_after',
+        node_id: 'ai-node-1',
+        tree: { type: { resolvedName: 'NotARealComponent' }, props: {}, nodes: [] },
+      },
+      { op: 'delete_node', node_id: 'ai-node-1' },
+    ];
+    let result: { ok: boolean; message?: string } | undefined;
+    expect(() => { result = __test.applyPatch(actions, query, ops); }).not.toThrow();
+    expect(result!.ok).toBe(true);
+    expect(propsById['craft-1']).toBeUndefined(); // delete_node (op 2) still applied
+    expect(warnSpy).toHaveBeenCalled(); // op 1's failure was logged, not thrown
+    warnSpy.mockRestore();
+  });
+});