konnosDashboard
tool-sanitizer.test.ts176 lines · main
1import type { ToolUIPart } from 'ai'
2import { describe, expect, test } from 'vitest'
3
4// End of third-party imports
5
6import { prepareMessagesForAPI } from '../message-utils'
7import {
8 createAssistantMessageWithExecuteSqlTool,
9 createAssistantMessageWithMultipleTools,
10 createLongConversation,
11} from '../test-fixtures'
12import { NO_DATA_PERMISSIONS, sanitizeMessagePart } from './tool-sanitizer'
13
14describe('messages are sanitized based on opt-in level', () => {
15 test('messages are sanitized at disabled level', () => {
16 const messages = [
17 createAssistantMessageWithExecuteSqlTool('SELECT email FROM users', [
18 { email: 'test@example.com' },
19 ]),
20 ]
21
22 // Prepare messages as frontend would
23 const preparedMessages = prepareMessagesForAPI(messages)
24
25 // Sanitize messages as API endpoint would
26 const processedMessages = preparedMessages.map((msg) => {
27 if (msg.role === 'assistant' && msg.parts) {
28 const processedParts = msg.parts.map((part) => {
29 return sanitizeMessagePart(part, 'disabled')
30 })
31
32 return { ...msg, parts: processedParts }
33 }
34 return msg
35 })
36
37 const output = (processedMessages[0].parts[1] as ToolUIPart).output
38 expect(output).toMatch(NO_DATA_PERMISSIONS)
39 })
40
41 test('messages are sanitized at schema level', () => {
42 const messages = [
43 createAssistantMessageWithExecuteSqlTool('SELECT email FROM users', [
44 { email: 'test@example.com' },
45 ]),
46 ]
47
48 // Prepare messages as frontend would
49 const preparedMessages = prepareMessagesForAPI(messages)
50
51 // Sanitize messages as API endpoint would
52 const processedMessages = preparedMessages.map((msg) => {
53 if (msg.role === 'assistant' && msg.parts) {
54 const processedParts = msg.parts.map((part) => {
55 return sanitizeMessagePart(part, 'schema')
56 })
57
58 return { ...msg, parts: processedParts }
59 }
60 return msg
61 })
62
63 const output = (processedMessages[0].parts[1] as ToolUIPart).output
64 expect(output).toMatch(NO_DATA_PERMISSIONS)
65 })
66
67 test('messages are sanitized at schema and log level', () => {
68 const messages = [
69 createAssistantMessageWithExecuteSqlTool('SELECT email FROM users', [
70 { email: 'test@example.com' },
71 ]),
72 ]
73
74 // Prepare messages as frontend would
75 const preparedMessages = prepareMessagesForAPI(messages)
76
77 // Sanitize messages as API endpoint would
78 const processedMessages = preparedMessages.map((msg) => {
79 if (msg.role === 'assistant' && msg.parts) {
80 const processedParts = msg.parts.map((part) => {
81 return sanitizeMessagePart(part, 'schema_and_log')
82 })
83
84 return { ...msg, parts: processedParts }
85 }
86 return msg
87 })
88
89 const output = (processedMessages[0].parts[1] as ToolUIPart).output
90 expect(output).toMatch(NO_DATA_PERMISSIONS)
91 })
92
93 test('messages are not sanitized at data level', () => {
94 const messages = [
95 createAssistantMessageWithExecuteSqlTool('SELECT email FROM users', [
96 { email: 'test@example.com' },
97 ]),
98 ]
99
100 // Prepare messages as frontend would
101 const preparedMessages = prepareMessagesForAPI(messages)
102
103 // Sanitize messages as API endpoint would
104 const processedMessages = preparedMessages.map((msg) => {
105 if (msg.role === 'assistant' && msg.parts) {
106 const processedParts = msg.parts.map((part) => {
107 return sanitizeMessagePart(part, 'schema_and_log_and_data')
108 })
109
110 return { ...msg, parts: processedParts }
111 }
112 return msg
113 })
114
115 const output = (processedMessages[0].parts[1] as ToolUIPart).output
116 expect(output).toEqual([{ email: 'test@example.com' }])
117 })
118
119 test('multiple tool parts in message are sanitized', () => {
120 const messages = [createAssistantMessageWithMultipleTools()]
121
122 // Prepare messages as frontend would
123 const preparedMessages = prepareMessagesForAPI(messages)
124
125 // Sanitize messages as API endpoint would
126 const processedMessages = preparedMessages.map((msg) => {
127 if (msg.role === 'assistant' && msg.parts) {
128 const processedParts = msg.parts.map((part) => {
129 return sanitizeMessagePart(part, 'schema')
130 })
131
132 return { ...msg, parts: processedParts }
133 }
134 return msg
135 })
136
137 const parts = processedMessages[0].parts
138 parts.forEach((part) => {
139 if (part.type.startsWith('tool')) {
140 const tool = part as ToolUIPart
141 expect(tool.output).toMatch(NO_DATA_PERMISSIONS)
142 }
143 })
144 })
145
146 test('long message chain is sanitized', () => {
147 const messages = createLongConversation()
148
149 // Prepare messages as frontend would
150 const preparedMessages = prepareMessagesForAPI(messages)
151
152 // Sanitize messages as API endpoint would
153 const processedMessages = preparedMessages.map((msg) => {
154 if (msg.role === 'assistant' && msg.parts) {
155 const processedParts = msg.parts.map((part) => {
156 return sanitizeMessagePart(part, 'schema')
157 })
158
159 return { ...msg, parts: processedParts }
160 }
161 return msg
162 })
163
164 processedMessages.forEach((msg) => {
165 if (msg.role === 'assistant' && msg.parts) {
166 const parts = msg.parts
167 parts.forEach((part) => {
168 if (part.type.startsWith('tool')) {
169 const tool = part as ToolUIPart
170 expect(tool.output).toMatch(NO_DATA_PERMISSIONS)
171 }
172 })
173 }
174 })
175 })
176})