test: update integration tests to include additional properties for agent options

- Added userId, chatId, organizationId, dataSourceId, and dataSourceSyntax to the integration tests for the Think and Prep Agent and Mark Message Complete Step.
- Cleaned up import statements for better organization and readability.
- Ensured that the tests reflect the latest schema changes for improved tracking and functionality.

These updates enhance the clarity and robustness of the integration tests, aligning them with recent changes in agent options.

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({

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();
   });
-});
+});

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' }],