From b46b6915a5f9e421552356eb9fdad217b0ee306e Mon Sep 17 00:00:00 2001 From: Josh Knapp Date: Sun, 12 Jul 2026 12:37:25 -0700 Subject: [PATCH] fix(builder): validate AI response before deserializing into craft state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- craft/src/utils/apply-ai-response.test.ts | 181 +++++++++++++++++++++- craft/src/utils/apply-ai-response.ts | 125 ++++++++++++++- 2 files changed, 298 insertions(+), 8 deletions(-) diff --git a/craft/src/utils/apply-ai-response.test.ts b/craft/src/utils/apply-ai-response.test.ts index e2e37fb..13bc8e7 100644 --- a/craft/src/utils/apply-ai-response.test.ts +++ b/craft/src/utils/apply-ai-response.test.ts @@ -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) { + const propsById: Record = { [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(); + }); +}); diff --git a/craft/src/utils/apply-ai-response.ts b/craft/src/utils/apply-ai-response.ts index c61f5d9..1a2894f 100644 --- a/craft/src/utils/apply-ai-response.ts +++ b/craft/src/utils/apply-ai-response.ts @@ -2,6 +2,7 @@ import { useEditor } from '@craftjs/core'; import type { NodeTree } from '@craftjs/core'; import { usePages } from '../state/PageContext'; import { SitesmithResponse, SerializedTreeNode } from '../types/sitesmith'; +import { componentResolver } from '../components/resolver'; /** Only Container is a "real" Craft.js canvas in serialized state. Layout * shells (Section/HeroSimple/ColumnLayout/etc) use linkedNodes @@ -16,6 +17,81 @@ const SHELL_INNER: Record = { FormContainer: 'form-inner', }; +/** + * update_props may never touch these keys, regardless of what the AI sends. + * `node_id` is the stable identifier other ops (and findNodeIdByAiNodeId) + * rely on to re-target a node in later turns of the conversation — letting + * the AI clobber it would silently orphan future patches. + */ +const PROTECTED_UPDATE_PROPS_KEYS = new Set(['node_id']); + +/** True if `name` is a registered component in the resolver (i.e. safe to + * hand to `actions.addNodeTree` without it throwing "does not exist in + * resolver" at render time). */ +function isKnownResolvedName(name: unknown): name is string { + return typeof name === 'string' && Object.prototype.hasOwnProperty.call(componentResolver, name); +} + +/** + * Validate + sanitize an AI-supplied tree before it is handed to + * `buildNodeTree`/`query.parseFreshNode`. Fails SOFT — never throws: + * + * - **resolvedName allowlist**: any node whose `type.resolvedName` is not a + * key of `componentResolver` is dropped (it and its subtree), with a + * `console.warn`. If the tree's ROOT is itself invalid, this returns + * `null` and the caller must skip the whole op/tree. + * - **id policy — REGENERATE (not skip)**: a node id that is `'ROOT'`, + * empty/non-string, or collides with an id already in `usedIds` is + * replaced with a fresh `ai-auto-N` id (and a warning is logged). The + * node itself is kept — only structurally invalid *types* are dropped; + * a merely-bad *id* is repaired so we never lose otherwise-valid AI + * content over an id collision. + * + * `usedIds` should be seeded with the live document's existing node ids + * (see `buildNodeTree`) so AI-supplied ids can never collide with content + * already on the canvas; it is also used internally to keep ids unique + * within the tree being sanitized. + */ +export function sanitizeAiTree( + tree: SerializedTreeNode, + usedIds: Set = new Set(), + idCounter: { n: number } = { n: 0 }, +): SerializedTreeNode | null { + const resolvedName = tree?.type?.resolvedName; + if (!isKnownResolvedName(resolvedName)) { + console.warn(`sitesmith: dropping node with unknown component "${String(resolvedName)}"`); + return null; + } + + const rawId = tree.props?.node_id; + let id = typeof rawId === 'string' ? rawId.trim() : ''; + if (!id) { + id = `ai-auto-${idCounter.n++}`; + } else if (id === 'ROOT') { + console.warn('sitesmith: rejecting AI-supplied node id "ROOT" (reserved); regenerating'); + id = `ai-auto-${idCounter.n++}`; + } else if (usedIds.has(id)) { + console.warn(`sitesmith: AI-supplied node id "${id}" collides with an existing id; regenerating`); + id = `ai-auto-${idCounter.n++}`; + } + // Belt-and-suspenders: keep drawing fresh ids in the (pathological) case + // the freshly generated id itself collides. + while (usedIds.has(id)) id = `ai-auto-${idCounter.n++}`; + usedIds.add(id); + + const children: SerializedTreeNode[] = []; + for (const child of tree.nodes ?? []) { + const sanitizedChild = sanitizeAiTree(child, usedIds, idCounter); + if (sanitizedChild) children.push(sanitizedChild); + } + + return { + type: tree.type, + props: { ...tree.props, node_id: id }, + nodes: children, + }; +} + /** * Flatten a SerializedTreeNode tree into a Craft.js node map ready for * `actions.deserialize()`. @@ -68,11 +144,25 @@ export function serializeTreeForCraft(tree: SerializedTreeNode): { rootNodeId: s * inserting/replacing sections or individual nodes. */ function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree { - const idCounter = { n: 0 }; + // Validate + sanitize at the AI boundary before touching Craft.js at all: + // unknown resolvedNames are dropped (never reach parseFreshNode/addNodeTree, + // which would throw), and ids are repaired ('ROOT'/empty/colliding → + // regenerated) against the live document's existing node ids so a new + // AI-supplied id can never clobber or duplicate one already on the canvas. + const existingIds = new Set( + typeof query?.getNodes === 'function' ? Object.keys(query.getNodes()) : [], + ); + const sanitized = sanitizeAiTree(tree, existingIds); + if (!sanitized) { + throw new Error('sitesmith: AI tree root has an unknown/invalid component type; nothing to build'); + } + const craftNodes: Record = {}; const walk = (node: SerializedTreeNode, parent: string | null): string => { - const id = (node.props.node_id as string | undefined) || `ai-auto-${idCounter.n++}`; + // sanitizeAiTree() has already guaranteed every node here has a valid, + // unique, non-'ROOT' node_id — no fallback/collision handling needed. + const id = node.props.node_id as string; const craftNode = (query.parseFreshNode({ id, data: { @@ -138,7 +228,7 @@ function buildNodeTree(query: any, tree: SerializedTreeNode): NodeTree { return id; }; - const rootId = walk(tree, null); + const rootId = walk(sanitized, null); return { rootNodeId: rootId, nodes: craftNodes }; } @@ -156,7 +246,7 @@ export function findNodeIdByAiNodeId(query: any, aiNodeId: string): string | nul } /** Exported for unit tests */ -export const __test = { findNodeIdByAiNodeId }; +export const __test = { findNodeIdByAiNodeId, buildNodeTree, applyPatch }; /** * React hook that returns an `apply` function. @@ -243,9 +333,32 @@ function applyPatch( } switch (op.op) { - case 'update_props': - actions.setProp(id, (p: any) => { Object.assign(p, op.props); }); + case 'update_props': { + const patchProps = op.props; + if (!patchProps || typeof patchProps !== 'object' || Array.isArray(patchProps)) { + console.warn('sitesmith patch: update_props with invalid props payload, skipping', op.node_id); + break; + } + actions.setProp(id, (p: any) => { + for (const [key, value] of Object.entries(patchProps)) { + if (PROTECTED_UPDATE_PROPS_KEYS.has(key)) { + console.warn(`sitesmith patch: refusing to overwrite protected prop "${key}"`); + continue; + } + // Merge `style` shallowly instead of blindly replacing it — the + // AI usually only means to change one or two style keys, and a + // blind Object.assign would clobber every other style the user + // (or a previous AI turn) had already set. + if (key === 'style' && value && typeof value === 'object' && !Array.isArray(value)) { + const existingStyle = p.style && typeof p.style === 'object' ? p.style : {}; + p.style = { ...existingStyle, ...(value as Record) }; + } else { + p[key] = value; + } + } + }); break; + } case 'replace_node': { try {