konnosDashboard
studio-tools.test.ts233 lines · main
1import { safeSql } from '@supabase/pg-meta'
2import { beforeEach, describe, expect, it, vi } from 'vitest'
3
4import { getStudioTools } from './studio-tools'
5import { executeSql } from '@/data/sql/execute-sql-query'
6import { NO_DATA_PERMISSIONS } from '@/lib/ai/tools/tool-sanitizer'
7
8vi.mock('@/data/sql/execute-sql-query', () => ({
9 executeSql: vi.fn(),
10}))
11
12describe('ai/tools/studio-tools', () => {
13 beforeEach(() => {
14 vi.mocked(executeSql).mockReset()
15 })
16
17 describe('getStudioTools', () => {
18 it('should return an object with tool definitions', () => {
19 const tools = getStudioTools()
20
21 expect(tools).toBeDefined()
22 expect(typeof tools).toBe('object')
23 })
24
25 it('should include execute_sql tool', () => {
26 const tools = getStudioTools()
27
28 expect(tools.execute_sql).toBeDefined()
29 expect(tools.execute_sql.description).toContain('execute a SQL statement')
30 })
31
32 it('should include deploy_edge_function tool', () => {
33 const tools = getStudioTools()
34
35 expect(tools.deploy_edge_function).toBeDefined()
36 expect(tools.deploy_edge_function.description).toContain('deploy a Briven Edge Function')
37 })
38
39 it('should include rename_chat tool', () => {
40 const tools = getStudioTools()
41
42 expect(tools.rename_chat).toBeDefined()
43 expect(tools.rename_chat.description).toContain('Rename the current chat session')
44 })
45
46 it('should have exactly 4 tools', () => {
47 const tools = getStudioTools()
48 const toolNames = Object.keys(tools)
49
50 expect(toolNames).toHaveLength(4)
51 expect(toolNames).toContain('load_knowledge')
52 expect(toolNames).toContain('execute_sql')
53 expect(toolNames).toContain('deploy_edge_function')
54 expect(toolNames).toContain('rename_chat')
55 })
56
57 it('should have execute_sql with correct input schema fields', () => {
58 const tools = getStudioTools()
59 const executeSqlTool = tools.execute_sql
60
61 // Check that the tool has an input schema
62 expect(executeSqlTool.inputSchema).toBeDefined()
63
64 // Verify the schema exists and is a Zod object
65 const schema = executeSqlTool.inputSchema
66 expect(schema).toBeDefined()
67 expect((schema as any)._def.typeName).toBe('ZodObject')
68 })
69
70 it('should have deploy_edge_function with input schema', () => {
71 const tools = getStudioTools()
72 const deployTool = tools.deploy_edge_function
73
74 expect(deployTool.inputSchema).toBeDefined()
75
76 // Verify the schema exists and is a Zod object
77 expect(deployTool.inputSchema).toBeDefined()
78 expect((deployTool.inputSchema as any)._def.typeName).toBe('ZodObject')
79 })
80
81 it('should have rename_chat with execute function', async () => {
82 const tools = getStudioTools()
83 const renameTool = tools.rename_chat
84
85 expect(renameTool.execute).toBeDefined()
86 expect(typeof renameTool.execute).toBe('function')
87
88 // Test the execute function
89 if (!renameTool.execute) throw new Error('execute is undefined')
90 const result = await renameTool.execute(
91 { newName: 'Test Chat' },
92 { toolCallId: 'test', messages: [] }
93 )
94 expect(result).toEqual({ status: 'Chat request sent to client' })
95 })
96
97 it('should validate execute_sql input schema correctly', () => {
98 const tools = getStudioTools()
99 const schema = tools.execute_sql.inputSchema
100
101 // Check if schema is a Zod schema with safeParse
102 if ('safeParse' in schema) {
103 // Valid input
104 const validInput = {
105 sql: safeSql`SELECT * FROM users`,
106 label: 'Get users',
107 chartConfig: { view: 'table' as const },
108 isWriteQuery: false,
109 }
110 expect(schema.safeParse(validInput).success).toBe(true)
111
112 // Valid chart config
113 const validChartInput = {
114 sql: safeSql`SELECT count(*) FROM users`,
115 label: 'User count',
116 chartConfig: { view: 'chart' as const, xAxis: 'date', yAxis: 'count' },
117 isWriteQuery: false,
118 }
119 expect(schema.safeParse(validChartInput).success).toBe(true)
120
121 // Missing required field
122 const invalidInput = {
123 sql: safeSql`SELECT * FROM users`,
124 // missing label, chartConfig, isWriteQuery
125 }
126 expect(schema.safeParse(invalidInput).success).toBe(false)
127 } else {
128 // Skip test if schema doesn't have safeParse
129 expect(schema).toBeDefined()
130 }
131 })
132
133 it('should require approval for read and write SQL queries', () => {
134 const tools = getStudioTools()
135
136 expect(tools.execute_sql.needsApproval).toBe(true)
137 })
138
139 it('should return execute_sql rows to the UI and sanitize model output without data opt-in', async () => {
140 const rows = [{ email: 'test@example.com' }]
141 vi.mocked(executeSql).mockResolvedValue({ result: rows })
142
143 const tools = getStudioTools({
144 projectRef: 'test-project',
145 connectionString: 'encrypted-connection-string',
146 aiOptInLevel: 'schema',
147 })
148
149 if (!tools.execute_sql.execute) throw new Error('execute is undefined')
150 const result = await tools.execute_sql.execute(
151 {
152 sql: 'SELECT email FROM users',
153 label: 'Get emails',
154 chartConfig: { view: 'table' },
155 isWriteQuery: false,
156 },
157 { toolCallId: 'test', messages: [] }
158 )
159
160 expect(executeSql).toHaveBeenCalledWith(
161 {
162 projectRef: 'test-project',
163 connectionString: 'encrypted-connection-string',
164 sql: 'SELECT email FROM users',
165 },
166 undefined,
167 undefined
168 )
169 expect(result).toEqual(rows)
170 expect((tools.execute_sql as any).toModelOutput({ output: result })).toEqual({
171 type: 'text',
172 value: NO_DATA_PERMISSIONS,
173 })
174 })
175
176 it('should return execute_sql rows with data opt-in', async () => {
177 const rows = [{ email: 'test@example.com' }]
178 vi.mocked(executeSql).mockResolvedValue({ result: rows })
179
180 const tools = getStudioTools({
181 projectRef: 'test-project',
182 connectionString: 'encrypted-connection-string',
183 aiOptInLevel: 'schema_and_log_and_data',
184 })
185
186 if (!tools.execute_sql.execute) throw new Error('execute is undefined')
187 const result = await tools.execute_sql.execute(
188 {
189 sql: 'SELECT email FROM users',
190 label: 'Get emails',
191 chartConfig: { view: 'table' },
192 isWriteQuery: false,
193 },
194 { toolCallId: 'test', messages: [] }
195 )
196
197 expect(executeSql).toHaveBeenCalledWith(
198 {
199 projectRef: 'test-project',
200 connectionString: 'encrypted-connection-string',
201 sql: 'SELECT email FROM users',
202 },
203 undefined,
204 undefined
205 )
206 expect(result).toEqual(rows)
207 expect((tools.execute_sql as any).toModelOutput({ output: result })).toEqual({
208 type: 'json',
209 value: rows,
210 })
211 })
212
213 it('should validate rename_chat input schema correctly', () => {
214 const tools = getStudioTools()
215 const schema = tools.rename_chat.inputSchema
216
217 // Check if schema is a Zod schema with safeParse
218 if ('safeParse' in schema) {
219 // Valid input
220 expect(schema.safeParse({ newName: 'My Chat' }).success).toBe(true)
221
222 // Invalid input - missing newName
223 expect(schema.safeParse({}).success).toBe(false)
224
225 // Invalid input - wrong type
226 expect(schema.safeParse({ newName: 123 }).success).toBe(false)
227 } else {
228 // Skip test if schema doesn't have safeParse
229 expect(schema).toBeDefined()
230 }
231 })
232 })
233})