konnosDashboard
AIAssistant.utils.test.ts159 lines · main
1import type { UIMessage } from 'ai'
2import { describe, expect, test } from 'vitest'
3
4import {
5 hasPendingToolApproval,
6 isReadOnlySelect,
7 resolvePendingToolApprovalsAsDenied,
8} from './AIAssistant.utils'
9
10const createMessageWithPart = (
11 part: UIMessage['parts'][number],
12 role: UIMessage['role'] = 'assistant'
13) =>
14 [
15 {
16 id: `${role}-msg-1`,
17 role,
18 parts: [part],
19 },
20 ] as UIMessage[]
21
22const pendingSqlApprovalPart = {
23 type: 'tool-execute_sql',
24 toolCallId: 'call-1',
25 state: 'approval-requested',
26 input: { sql: 'select 1', label: 'Test query' },
27 approval: { id: 'approval-1' },
28} satisfies UIMessage['parts'][number]
29
30describe('AIAssistant.utils.ts:isReadOnlySelect', () => {
31 test('Should return true for SQL that only contains SELECT operation', () => {
32 const sql = 'select * from countries where id > 100 order by id asc;'
33 const result = isReadOnlySelect(sql)
34 expect(result).toBe(true)
35 })
36 test('Should return false for SQL that contains INSERT operation', () => {
37 const sql = `insert into countries (id, name) values (1, 'hello');`
38 const result = isReadOnlySelect(sql)
39 expect(result).toBe(false)
40 })
41 test('Should return false for SQL that contains UPDATE operation', () => {
42 const sql = `update countries set name = 'hello' where id = 2;`
43 const result = isReadOnlySelect(sql)
44 expect(result).toBe(false)
45 })
46 test('Should return false for SQL that contains DELETE operation', () => {
47 const sql = `delete from countries where id = 2;`
48 const result = isReadOnlySelect(sql)
49 expect(result).toBe(false)
50 })
51 test('Should return false for SQL that contains ALTER operation', () => {
52 const sql = `alter table countries drop column id if exists;`
53 const result = isReadOnlySelect(sql)
54 expect(result).toBe(false)
55 })
56 test('Should return false for SQL that contains DROP operation', () => {
57 const sql = `drop table if exists countries;`
58 const result = isReadOnlySelect(sql)
59 expect(result).toBe(false)
60 })
61 test('Should return false for SQL that contains CREATE operation', () => {
62 const sql = `create schema test_schema;`
63 const result = isReadOnlySelect(sql)
64 expect(result).toBe(false)
65 })
66 test('Should return false for SQL that contains REPLACE operation', () => {
67 const sql = `create or replace view test_view as select * from countries where id > 500;`
68 const result = isReadOnlySelect(sql)
69 expect(result).toBe(false)
70 })
71 test('Should return false for SQL that calls a function not whitelisted', () => {
72 const sql = `select create_new_user();`
73 const result = isReadOnlySelect(sql)
74 expect(result).toBe(false)
75 })
76 test('Should return true for SQL that calls a function that is whitelisted', () => {
77 const sql = `select count(select * from countries);`
78 const result = isReadOnlySelect(sql)
79 expect(result).toBe(true)
80 })
81 test('Should return false for SQL that contains a write operation with a read operation', () => {
82 const sql1 = `select count(select * from countries); create schema joshen;`
83 const result1 = isReadOnlySelect(sql1)
84 expect(result1).toBe(false)
85 const sql2 = `create schema joshen; select count(select * from countries);`
86 const result2 = isReadOnlySelect(sql2)
87 expect(result2).toBe(false)
88 })
89})
90
91describe('AIAssistant.utils.ts:hasPendingToolApproval', () => {
92 test('Should return true when an assistant message has a pending tool approval', () => {
93 const messages = createMessageWithPart(pendingSqlApprovalPart)
94
95 expect(hasPendingToolApproval(messages)).toBe(true)
96 })
97
98 test('Should return false when approval has already been answered', () => {
99 const messages = createMessageWithPart({
100 type: 'tool-execute_sql',
101 toolCallId: 'call-1',
102 state: 'approval-responded',
103 input: { sql: 'select 1', label: 'Test query' },
104 approval: { id: 'approval-1', approved: true },
105 })
106
107 expect(hasPendingToolApproval(messages)).toBe(false)
108 })
109
110 test('Should ignore non-assistant messages', () => {
111 const messages = createMessageWithPart(pendingSqlApprovalPart, 'user')
112
113 expect(hasPendingToolApproval(messages)).toBe(false)
114 })
115
116 test('Should return true for dynamic tool approvals', () => {
117 const messages = createMessageWithPart({
118 type: 'dynamic-tool',
119 toolName: 'execute_sql',
120 toolCallId: 'call-1',
121 state: 'approval-requested',
122 input: { sql: 'select 1', label: 'Test query' },
123 approval: { id: 'approval-1' },
124 })
125
126 expect(hasPendingToolApproval(messages)).toBe(true)
127 })
128})
129
130describe('AIAssistant.utils.ts:resolvePendingToolApprovalsAsDenied', () => {
131 test('Should convert pending approvals into denied outputs', () => {
132 const messages = createMessageWithPart(pendingSqlApprovalPart)
133
134 const resolvedMessages = resolvePendingToolApprovalsAsDenied(messages)
135 const [part] = resolvedMessages[0].parts
136
137 expect(part).toMatchObject({
138 state: 'output-denied',
139 approval: {
140 id: 'approval-1',
141 approved: false,
142 reason: 'Skipped because the user sent a follow-up message.',
143 },
144 })
145 })
146
147 test('Should leave completed approvals unchanged', () => {
148 const messages = createMessageWithPart({
149 type: 'tool-execute_sql',
150 toolCallId: 'call-1',
151 state: 'output-available',
152 input: { sql: 'select 1', label: 'Test query' },
153 output: [],
154 approval: { id: 'approval-1', approved: true },
155 })
156
157 expect(resolvePendingToolApprovalsAsDenied(messages)).toEqual(messages)
158 })
159})