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

import { describe, test, expect, vi } from 'vitest';
import { sanitizeAiTree, __test } from './apply-ai-response';

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);
  });
});

// ---------------------------------------------------------------------------
// 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();
  });
});