diff --git a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts
index a77d7be3e..dda75f9a9 100644
--- a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts
+++ b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts
@@ -13,6 +13,10 @@ describe('Think and Prep Agent Integration Tests', () => {
         messageId: 'test-message-123',
         sql_dialect_guidance: 'postgresql',
         chatId: 'test-chat-123',
+        organizationId: 'test-organization-123',
+        userId: 'test-user-123',
+        dataSourceId: 'test-data-source-123',
+        dataSourceSyntax: 'postgresql',
       });
 
       const streamResult = await thinkAndPrepAgent.stream({
diff --git a/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts b/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts
index 35b5e3586..329a548d3 100644
--- a/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts
+++ b/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts
@@ -1,7 +1,7 @@
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
-import { db, messages, chats } from '@buster/database';
-import { eq } from 'drizzle-orm';
+import { chats, db, messages } from '@buster/database';
 import { createTestChat, createTestMessage } from '@buster/test-utils';
+import { eq } from 'drizzle-orm';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { markMessageComplete } from './mark-message-complete-step';
 
 describe('markMessageComplete integration test', () => {
@@ -20,7 +20,7 @@ describe('markMessageComplete integration test', () => {
     // Create a test message - returns just the ID string
     testMessageId = await createTestMessage(testChatId, testUserId, {
       isCompleted: false, // Start with uncompleted message for testing
-      finalReasoningMessage: null,
+      finalReasoningMessage: '',
     });
   });
 
@@ -40,11 +40,8 @@ describe('markMessageComplete integration test', () => {
     expect(result.completedAt).toBeDefined();
 
     // Verify the message was updated in the database
-    const [updatedMessage] = await db
-      .select()
-      .from(messages)
-      .where(eq(messages.id, testMessageId));
-    
+    const [updatedMessage] = await db.select().from(messages).where(eq(messages.id, testMessageId));
+
     expect(updatedMessage?.isCompleted).toBe(true);
     expect(updatedMessage?.finalReasoningMessage).toBe('Task completed successfully');
   });
@@ -67,14 +64,11 @@ describe('markMessageComplete integration test', () => {
     expect(result.messageId).toBe(testMessageId);
 
     // Wait a moment for database update to complete
-    await new Promise(resolve => setTimeout(resolve, 100));
+    await new Promise((resolve) => setTimeout(resolve, 100));
 
     // Verify the chat was updated with file information
-    const [updatedChat] = await db
-      .select()
-      .from(chats)
-      .where(eq(chats.id, testChatId));
-    
+    const [updatedChat] = await db.select().from(chats).where(eq(chats.id, testChatId));
+
     expect(updatedChat).toBeDefined();
     expect(updatedChat?.mostRecentFileId).toBe(selectedFile.fileId);
     expect(updatedChat?.mostRecentFileType).toBe(selectedFile.fileType);
@@ -106,7 +100,7 @@ describe('markMessageComplete integration test', () => {
         .select()
         .from(messages)
         .where(eq(messages.id, testMessageId));
-      
+
       expect(updatedMessage?.finalReasoningMessage).toBe('complete');
     }
   });
@@ -128,11 +122,8 @@ describe('markMessageComplete integration test', () => {
     expect(result.success).toBe(true);
 
     // Verify the chat was not updated
-    const [chat] = await db
-      .select()
-      .from(chats)
-      .where(eq(chats.id, testChatId));
-    
+    const [chat] = await db.select().from(chats).where(eq(chats.id, testChatId));
+
     expect(chat?.mostRecentFileId).toBeNull();
   });
-});
\ No newline at end of file
+});
diff --git a/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts b/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts
index 2d5a4adba..564440993 100644
--- a/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts
+++ b/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts
@@ -39,6 +39,11 @@ describe('runThinkAndPrepAgentStep', () => {
       options: {
         messageId: 'test-message-id',
         sql_dialect_guidance: 'postgres',
+        userId: 'test-user-id',
+        chatId: 'test-chat-id',
+        organizationId: 'test-organization-id',
+        dataSourceId: 'test-data-source-id',
+        dataSourceSyntax: 'test-data-source-syntax',
       },
       streamOptions: {
         messages: [{ role: 'user', content: 'Test prompt' }],
@@ -71,6 +76,11 @@ describe('runThinkAndPrepAgentStep', () => {
         options: {
           messageId: 'test-message-id',
           sql_dialect_guidance: 'postgres',
+          userId: 'test-user-id',
+          chatId: 'test-chat-id',
+          organizationId: 'test-organization-id',
+          dataSourceId: 'test-data-source-id',
+          dataSourceSyntax: 'test-data-source-syntax',
         },
         streamOptions: {
           messages: [{ role: 'user', content: 'Test prompt' }],
@@ -100,6 +110,11 @@ describe('runThinkAndPrepAgentStep', () => {
         options: {
           messageId: 'test-message-id',
           sql_dialect_guidance: 'postgres',
+          userId: 'test-user-id',
+          chatId: 'test-chat-id',
+          organizationId: 'test-organization-id',
+          dataSourceId: 'test-data-source-id',
+          dataSourceSyntax: 'test-data-source-syntax',
         },
         streamOptions: {
           messages: [{ role: 'user', content: 'Test prompt' }],
@@ -127,6 +142,11 @@ describe('runThinkAndPrepAgentStep', () => {
         options: {
           messageId: 'test-message-id',
           sql_dialect_guidance: 'postgres',
+          userId: 'test-user-id',
+          chatId: 'test-chat-id',
+          organizationId: 'test-organization-id',
+          dataSourceId: 'test-data-source-id',
+          dataSourceSyntax: 'test-data-source-syntax',
         },
         streamOptions: {
           messages: [{ role: 'user', content: 'Test prompt' }],