initial ai assistance

test/assistant.test.ts

@@ -0,0 +1,132 @@
+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)
+  })
+})