From a7bddc8bb01cfe5e847fafbc07bd7574c29ab5a7 Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Thu, 25 Sep 2025 15:23:59 +0800
Subject: [PATCH] Fix tests

---
 apps/studio/lib/ai/tools/tool-sanitizer.test.ts | 12 ++++++------
 apps/studio/lib/ai/tools/tool-sanitizer.ts      |  6 ++++--
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.test.ts b/apps/studio/lib/ai/tools/tool-sanitizer.test.ts
index b104900ee9b..c8c7600b7a1 100644
--- a/apps/studio/lib/ai/tools/tool-sanitizer.test.ts
+++ b/apps/studio/lib/ai/tools/tool-sanitizer.test.ts
@@ -8,7 +8,7 @@ import {
   createAssistantMessageWithMultipleTools,
   createLongConversation,
 } from '../test-fixtures'
-import { sanitizeMessagePart } from './tool-sanitizer'
+import { NO_DATA_PERMISSIONS, sanitizeMessagePart } from './tool-sanitizer'
 
 describe('messages are sanitized based on opt-in level', () => {
   test('messages are sanitized at disabled level', () => {
@@ -34,7 +34,7 @@ describe('messages are sanitized based on opt-in level', () => {
     })
 
     const output = (processedMessages[0].parts[1] as ToolUIPart).output
-    expect(output).toMatch(/result is not shared/)
+    expect(output).toMatch(NO_DATA_PERMISSIONS)
   })
 
   test('messages are sanitized at schema level', () => {
@@ -60,7 +60,7 @@ describe('messages are sanitized based on opt-in level', () => {
     })
 
     const output = (processedMessages[0].parts[1] as ToolUIPart).output
-    expect(output).toMatch(/result is not shared/)
+    expect(output).toMatch(NO_DATA_PERMISSIONS)
   })
 
   test('messages are sanitized at schema and log level', () => {
@@ -86,7 +86,7 @@ describe('messages are sanitized based on opt-in level', () => {
     })
 
     const output = (processedMessages[0].parts[1] as ToolUIPart).output
-    expect(output).toMatch(/result is not shared/)
+    expect(output).toMatch(NO_DATA_PERMISSIONS)
   })
 
   test('messages are not sanitized at data level', () => {
@@ -137,7 +137,7 @@ describe('messages are sanitized based on opt-in level', () => {
     parts.forEach((part) => {
       if (part.type.startsWith('tool')) {
         const tool = part as ToolUIPart
-        expect(tool.output).toMatch(/result is not shared/)
+        expect(tool.output).toMatch(NO_DATA_PERMISSIONS)
       }
     })
   })
@@ -166,7 +166,7 @@ describe('messages are sanitized based on opt-in level', () => {
         parts.forEach((part) => {
           if (part.type.startsWith('tool')) {
             const tool = part as ToolUIPart
-            expect(tool.output).toMatch(/result is not shared/)
+            expect(tool.output).toMatch(NO_DATA_PERMISSIONS)
           }
         })
       }
diff --git a/apps/studio/lib/ai/tools/tool-sanitizer.ts b/apps/studio/lib/ai/tools/tool-sanitizer.ts
index 8b6883e991a..9ea3d5c2559 100644
--- a/apps/studio/lib/ai/tools/tool-sanitizer.ts
+++ b/apps/studio/lib/ai/tools/tool-sanitizer.ts
@@ -9,6 +9,9 @@ interface ToolSanitizer {
   sanitize: <Tool extends ToolUIPart>(tool: Tool, optInLevel: AiOptInLevel) => Tool
 }
 
+export const NO_DATA_PERMISSIONS =
+  'The query was executed and the user has viewed the results but decided not to share in the conversation due to permission levels. Continue with your plan unless instructed to interpret the result.'
+
 const executeSqlSanitizer: ToolSanitizer = {
   toolName: 'execute_sql',
   sanitize: (tool, optInLevel) => {
@@ -17,8 +20,7 @@ const executeSqlSanitizer: ToolSanitizer = {
 
     if (optInLevel !== 'schema_and_log_and_data') {
       if (Array.isArray(output)) {
-        sanitizedOutput =
-          'The query was executed and the user has viewed the results but decided not to share in the conversation due to permission levels. Continue with your plan unless instructed to interpret the result.'
+        sanitizedOutput = NO_DATA_PERMISSIONS
       }
     } else {
       sanitizedOutput = output