cms-onpageeditor/test/assistant.test.ts

import { describe, expect, it, vi } from 'vitest'
import { parseToolCall, runAgent, visibleText } from '../src/ai/agent'
import type { Engine } from '../src/ai/engine'
import { applyProposal, createTools } from '../src/ai/tools'

const schema = { type: 'object', properties: { title: { type: 'string', maxLength: 20 }, links: { type: 'array', items: { type: 'object', properties: { text: { type: 'string' } } } } } }
const node = { id: 'n1', dataType: 'hero', version: 3, data: { title: 'Hello wrold', links: [{ text: 'a' }] } }

function fakeCtrl() {
  const setData = vi.fn(async () => ({}))
  return {
    setData,
    ctrl: {
      store: { get: () => ({ selected: 'n1', blocks: [{ id: 'n1', type: 'hero', depth: 0 }] }) },
      block: (id: string) => (id === 'n1' ? { id } : undefined),
      select: vi.fn(),
      schemaOf: () => schema,
      node: async () => node,
      schema: async () => schema,
      labelFor: () => 'Hero',
      setData,
    } as any,
  }
}
const scripted = (replies: string[]): Engine => ({ generate: async () => replies.shift() ?? 'done' })

describe('parseToolCall', () => {
  it('reads tagged, unclosed and fenced calls; plain text is not a call', () => {
    expect(parseToolCall('<tool_call>{"name":"a","arguments":{"x":1}}</tool_call>')).toEqual({ name: 'a', arguments: { x: 1 } })
    expect(parseToolCall('<tool_call>{"name":"a"}')).toEqual({ name: 'a', arguments: {} })
    expect(parseToolCall('```json\n{"name":"b","arguments":{}}\n```')?.name).toBe('b')
    expect(parseToolCall('Just a sentence.')).toBeNull()
    expect(parseToolCall('<tool_call>{oops</tool_call>')).toBeNull()
  })
})

describe('runAgent + tools', () => {
  it('reads the node, queues a validated proposal, and finishes with text', async () => {
    const { ctrl } = fakeCtrl()
    const proposals: any[] = []
    const tools = createTools({ ctrl, propose: (p) => proposals.push(p) })
    const out = await runAgent(
      scripted([
        '<tool_call>{"name":"get_selected_node","arguments":{}}</tool_call>',
        '<tool_call>{"name":"propose_set","arguments":{"path":"title","value":"Hello world"}}</tool_call>',
        'Fixed the typo.',
      ]),
      'sys',
      [{ role: 'user', content: 'fix spelling' }],
      tools,
    )
    expect(out).toBe('Fixed the typo.')
    expect(proposals).toEqual([{ kind: 'set', nodeId: 'n1', label: 'Hero', path: 'title', before: 'Hello wrold', after: 'Hello world' }])
  })
  it('rejects schema-invalid proposals and unknown fields, feeding the error back', async () => {
    const { ctrl } = fakeCtrl()
    const proposals: any[] = []
    const tools = createTools({ ctrl, propose: (p) => proposals.push(p) })
    const seen: string[] = []
    const engine: Engine = {
      generate: async (m) => {
        seen.push(m.at(-1)!.content)
        return seen.length === 1
          ? '<tool_call>{"name":"propose_set","arguments":{"path":"title","value":"this title is far too long for the limit"}}</tool_call>'
          : seen.length === 2
            ? '<tool_call>{"name":"propose_set","arguments":{"path":"nope","value":"x"}}</tool_call>'
            : 'ok'
      },
    }
    await runAgent(engine, 's', [{ role: 'user', content: 'q' }], tools)
    expect(proposals).toHaveLength(0)
    expect(seen[1]).toContain('At most 20 characters')
    expect(seen[2]).toContain('Unknown field')
  })
  it('stops after maxSteps', async () => {
    const { ctrl } = fakeCtrl()
    const loop: Engine = { generate: async () => '<tool_call>{"name":"get_page_outline","arguments":{}}</tool_call>' }
    expect(await runAgent(loop, 's', [], createTools({ ctrl, propose() {} }), { maxSteps: 2 })).toContain('too many steps')
  })
})

describe('applyProposal', () => {
  it('rebuilds the nested patch from the fresh node and saves with its version', async () => {
    const { ctrl, setData } = fakeCtrl()
    await applyProposal(ctrl, { id: 1, kind: 'set', nodeId: 'n1', label: 'Hero', path: 'links.0.text', before: 'a', after: 'b' })
    expect(setData).toHaveBeenCalledWith('n1', { links: [{ text: 'b' }] }, { version: 3 })
  })
})

describe('more tools', () => {
  it('find_text matches case-insensitively with the path', async () => {
    const { ctrl } = fakeCtrl()
    const t = createTools({ ctrl, propose() {} }).find((x) => x.name === 'find_text')!
    expect(await t.run({ query: 'WROLD' })).toEqual([{ id: 'n1', type: 'hero', path: 'title', value: 'Hello wrold' }])
    await expect(t.run({})).rejects.toThrow('query')
  })
  it('select_block selects known blocks and refuses unknown ones', async () => {
    const { ctrl } = fakeCtrl()
    const t = createTools({ ctrl, propose() {} }).find((x) => x.name === 'select_block')!
    await t.run({ id: 'n1' })
    expect(ctrl.select).toHaveBeenCalledWith('n1', { scroll: true })
    await expect(t.run({ id: 'zzz' })).rejects.toThrow('No block')
  })
  it('propose_remove queues a removal; applyProposal removes', async () => {
    const { ctrl } = fakeCtrl()
    ctrl.remove = vi.fn(async () => ({}))
    const got: any[] = []
    await createTools({ ctrl, propose: (p) => got.push(p) }).find((x) => x.name === 'propose_remove')!.run({ id: 'n1' })
    expect(got[0]).toMatchObject({ kind: 'remove', nodeId: 'n1', summary: 'Hello wrold' })
    await applyProposal(ctrl, { ...got[0], id: 1 })
    expect(ctrl.remove).toHaveBeenCalledWith('n1')
  })
})

describe('streaming', () => {
  it('forwards tokens and hides tool-call markup', async () => {
    const { ctrl } = fakeCtrl()
    const seen: string[] = []
    const engine: Engine = {
      generate: async (_m, o) => {
        for (const t of ['Let me ', 'look.<tool', '_call>{"name":"get_page_outline","arguments":{}}</tool_call>']) o?.onToken?.(t)
        return 'Let me look.<tool_call>{"name":"get_page_outline","arguments":{}}</tool_call>'
      },
    }
    let n = 0
    const wrapped: Engine = { generate: async (m, o) => (n++ ? 'Done.' : engine.generate(m, o)) }
    await runAgent(wrapped, 's', [], createTools({ ctrl, propose() {} }), { onText: (p) => seen.push(visibleText(p)) })
    expect(seen).toContain('Let me')
    expect(seen).toContain('Let me look.')
    expect(seen.every((s) => !s.includes('<tool'))).toBe(true)
  })
})