diff --git a/apps/server/src/api/v2/chats/handler.test.ts b/apps/server/src/api/v2/chats/handler.test.ts index 98f348ce7..3b614e860 100644 --- a/apps/server/src/api/v2/chats/handler.test.ts +++ b/apps/server/src/api/v2/chats/handler.test.ts @@ -15,6 +15,7 @@ vi.mock('./services/chat-service', () => ({ vi.mock('./services/chat-helpers', () => ({ handleAssetChat: vi.fn(), + handleAssetChatWithPrompt: vi.fn(), })); vi.mock('@buster/database', () => ({ @@ -41,7 +42,7 @@ vi.mock('@buster/database', () => ({ import { getUserOrganizationId, updateMessage } from '@buster/database'; import { tasks } from '@trigger.dev/sdk/v3'; import { createChatHandler } from './handler'; -import { handleAssetChat } from './services/chat-helpers'; +import { handleAssetChat, handleAssetChatWithPrompt } from './services/chat-helpers'; import { initializeChat } from './services/chat-service'; describe('createChatHandler', () => { @@ -120,8 +121,51 @@ describe('createChatHandler', () => { expect(result).toEqual(mockChat); }); - it('should handle asset-based chat creation', async () => { - const assetChat = { ...mockChat, title: 'Asset Chat' }; + it('should handle asset-based chat creation and NOT trigger analyst task', async () => { + // Asset chat should match Rust implementation exactly + const assetChat = { + ...mockChat, + title: 'Test Metric', // Should be the asset name + message_ids: ['asset-msg-123'], + messages: { + 'asset-msg-123': { + id: 'asset-msg-123', + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + request_message: null, // No request message per Rust implementation + response_messages: { + 'text-msg-id': { + type: 'text' as const, + id: 'text-msg-id', + message: + 'Test Metric has been pulled into a new chat.\n\nContinue chatting to modify or make changes to it.', + is_final_message: true, + }, + 'asset-123': { + type: 'file' as const, + id: 'asset-123', + file_type: 'metric' as const, + file_name: 'Test Metric', + version_number: 1, + filter_version_id: null, + metadata: [ + { + status: 'completed' as const, + message: 'Pulled into new chat', + timestamp: expect.any(Number), + }, + ], + }, + }, + response_message_ids: ['text-msg-id', 'asset-123'], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: '', + feedback: null, + is_completed: true, + }, + }, + }; vi.mocked(handleAssetChat).mockResolvedValue(assetChat); const result = await createChatHandler( @@ -137,15 +181,8 @@ describe('createChatHandler', () => { mockUser, mockChat ); - expect(tasks.trigger).toHaveBeenCalledWith( - 'analyst-agent-task', - { - message_id: 'msg-123', - }, - { - concurrencyKey: 'chat-123', - } - ); + // IMPORTANT: Should NOT trigger analyst task for asset-only requests + expect(tasks.trigger).not.toHaveBeenCalled(); expect(result).toEqual(assetChat); }); @@ -158,23 +195,177 @@ describe('createChatHandler', () => { expect(tasks.trigger).not.toHaveBeenCalled(); }); - it('should not call handleAssetChat when prompt is provided with asset', async () => { + it('should call handleAssetChatWithPrompt when prompt is provided with asset', async () => { + // Chat should start empty when we have asset+prompt + const emptyChat = { + ...mockChat, + message_ids: [], + messages: {}, + }; + + // After handleAssetChatWithPrompt, we should have import message then user message + const chatWithPrompt = { + ...mockChat, + message_ids: ['import-msg-123', 'user-msg-123'], + messages: { + 'import-msg-123': { + id: 'import-msg-123', + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + request_message: null, + response_messages: {}, + response_message_ids: [], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: true, + }, + 'user-msg-123': { + id: 'user-msg-123', + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + request_message: { + request: 'Hello', + sender_id: 'user-123', + sender_name: 'Test User', + }, + response_messages: {}, + response_message_ids: [], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: false, + }, + }, + }; + + // Mock initializeChat to return empty chat (no initial message created) + vi.mocked(initializeChat).mockResolvedValue({ + chatId: 'chat-123', + messageId: 'msg-123', + chat: emptyChat, + }); + + vi.mocked(handleAssetChatWithPrompt).mockResolvedValueOnce(chatWithPrompt); + const result = await createChatHandler( { prompt: 'Hello', asset_id: 'asset-123', asset_type: 'metric' }, mockUser ); + // Verify initializeChat was called without prompt (to avoid duplicate message) + expect(initializeChat).toHaveBeenCalledWith( + { prompt: undefined, asset_id: 'asset-123', asset_type: 'metric' }, + mockUser, + '550e8400-e29b-41d4-a716-446655440000' + ); + expect(handleAssetChat).not.toHaveBeenCalled(); + expect(handleAssetChatWithPrompt).toHaveBeenCalledWith( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + 'Hello', + mockUser, + emptyChat + ); expect(tasks.trigger).toHaveBeenCalledWith( 'analyst-agent-task', { - message_id: 'msg-123', + message_id: 'user-msg-123', // Should use the last message ID (user's prompt) }, { concurrencyKey: 'chat-123', } ); - expect(result).toEqual(mockChat); + expect(result).toEqual(chatWithPrompt); + }); + + it('should ensure correct message order: import first, then user prompt', async () => { + const chatWithMessages = { + ...mockChat, + message_ids: ['import-msg-123', 'user-msg-123'], + messages: { + 'import-msg-123': { + id: 'import-msg-123', + created_at: '2025-07-25T12:00:00.000Z', + updated_at: '2025-07-25T12:00:00.000Z', + request_message: null, // Import messages have no request + response_messages: { + 'asset-123': { + type: 'file' as const, + id: 'asset-123', + file_type: 'metric' as const, + file_name: 'Test Metric', + version_number: 1, + }, + }, + response_message_ids: ['asset-123'], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: true, + }, + 'user-msg-123': { + id: 'user-msg-123', + created_at: '2025-07-25T12:00:01.000Z', // After import + updated_at: '2025-07-25T12:00:01.000Z', + request_message: { + request: 'What is this metric?', + sender_id: 'user-123', + sender_name: 'Test User', + }, + response_messages: {}, + response_message_ids: [], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: false, + }, + }, + }; + + vi.mocked(initializeChat).mockResolvedValue({ + chatId: 'chat-123', + messageId: 'msg-123', + chat: { ...mockChat, message_ids: [], messages: {} }, // Empty chat + }); + + vi.mocked(handleAssetChatWithPrompt).mockResolvedValueOnce(chatWithMessages); + + const result = await createChatHandler( + { prompt: 'What is this metric?', asset_id: 'asset-123', asset_type: 'metric' }, + mockUser + ); + + // Verify message order is correct + expect(result.message_ids).toHaveLength(2); + expect(result.message_ids[0]).toBe('import-msg-123'); + expect(result.message_ids[1]).toBe('user-msg-123'); + + // Verify import message has no request + const importMsg = result.messages['import-msg-123']; + expect(importMsg).toBeDefined(); + expect(importMsg?.request_message).toBeNull(); + expect(importMsg?.is_completed).toBe(true); + + // Verify user message has request + const userMsg = result.messages['user-msg-123']; + expect(userMsg).toBeDefined(); + expect(userMsg?.request_message?.request).toBe('What is this metric?'); + expect(userMsg?.is_completed).toBe(false); + + // Verify analyst task is triggered with user message ID + expect(tasks.trigger).toHaveBeenCalledWith( + 'analyst-agent-task', + { message_id: 'user-msg-123' }, + { concurrencyKey: 'chat-123' } + ); }); it('should handle trigger errors gracefully', async () => { diff --git a/apps/server/src/api/v2/chats/handler.ts b/apps/server/src/api/v2/chats/handler.ts index 3c89be23b..b9d42af44 100644 --- a/apps/server/src/api/v2/chats/handler.ts +++ b/apps/server/src/api/v2/chats/handler.ts @@ -7,7 +7,7 @@ import { type ChatWithMessages, } from '@buster/server-shared/chats'; import { tasks } from '@trigger.dev/sdk/v3'; -import { handleAssetChat } from './services/chat-helpers'; +import { handleAssetChat, handleAssetChatWithPrompt } from './services/chat-helpers'; import { initializeChat } from './services/chat-service'; /** @@ -57,30 +57,60 @@ export async function createChatHandler( } // Initialize chat (new or existing) - const { chatId, messageId, chat } = await initializeChat(request, user, organizationId); + // When we have both asset and prompt, we'll skip creating the initial message + // since handleAssetChatWithPrompt will create both the import and prompt messages + const shouldCreateInitialMessage = !(request.asset_id && request.asset_type && request.prompt); + const modifiedRequest = shouldCreateInitialMessage + ? request + : { ...request, prompt: undefined }; + + const { chatId, messageId, chat } = await initializeChat(modifiedRequest, user, organizationId); // Handle asset-based chat if needed let finalChat: ChatWithMessages = chat; - if (request.asset_id && request.asset_type && !request.prompt) { - finalChat = await handleAssetChat( - chatId, - messageId, - request.asset_id, - request.asset_type, - user, - chat - ); + let actualMessageId = messageId; // Track the actual message ID to use for triggering + let shouldTriggerAnalyst = true; // Flag to control whether to trigger analyst task + + if (request.asset_id && request.asset_type) { + if (!request.prompt) { + // Original flow: just import the asset without a prompt + finalChat = await handleAssetChat( + chatId, + messageId, + request.asset_id, + request.asset_type, + user, + chat + ); + // For asset-only chats, don't trigger analyst task - just return the chat with asset + shouldTriggerAnalyst = false; + } else { + // New flow: import asset then process the prompt + finalChat = await handleAssetChatWithPrompt( + chatId, + messageId, + request.asset_id, + request.asset_type, + request.prompt, + user, + chat + ); + // For asset+prompt chats, use the last message ID (the user's prompt message) + const lastMessageId = finalChat.message_ids[finalChat.message_ids.length - 1]; + if (lastMessageId) { + actualMessageId = lastMessageId; + } + } } - // Trigger background analysis if we have content - // This should be very fast (just queuing the job, not waiting for completion) - if (request.prompt || request.asset_id) { + // Trigger background analysis only if we have a prompt or it's not an asset-only request + if (shouldTriggerAnalyst && (request.prompt || !request.asset_id)) { try { // Just queue the background job - should be <100ms const taskHandle = await tasks.trigger( 'analyst-agent-task', { - message_id: messageId, + message_id: actualMessageId, }, { concurrencyKey: chatId, // Ensure sequential processing per chat @@ -96,7 +126,7 @@ export async function createChatHandler( // Update the message with the trigger run ID const { updateMessage } = await import('@buster/database'); - await updateMessage(messageId, { + await updateMessage(actualMessageId, { triggerRunId: taskHandle.id, }); diff --git a/apps/server/src/api/v2/chats/services/asset-import-service.ts b/apps/server/src/api/v2/chats/services/asset-import-service.ts new file mode 100644 index 000000000..fe154d92d --- /dev/null +++ b/apps/server/src/api/v2/chats/services/asset-import-service.ts @@ -0,0 +1,133 @@ +import type { User } from '@buster/database'; +import { + type AssetDetailsResult, + type Message, + chats, + createMessage, + createMessageFileAssociation, + db, + getAssetDetailsById, +} from '@buster/database'; +import type { + ChatAssetType, + ChatMessage, + ChatMessageResponseMessage, +} from '@buster/server-shared/chats'; +import { eq } from 'drizzle-orm'; +import { convertChatAssetTypeToDatabaseAssetType } from './server-asset-conversion'; + +/** + * Creates an import message for an asset + * This message represents the initial import of the asset into the chat + */ +export async function createAssetImportMessage( + chatId: string, + messageId: string, + assetId: string, + assetType: ChatAssetType, + assetDetails: AssetDetailsResult, + user: User +): Promise { + // Create the import message content + const importContent = `Imported ${assetType === 'metric' ? 'metric' : 'dashboard'} "${ + assetDetails.name + }"`; + + // Create the message in the database + const message = await createMessage({ + messageId, + chatId, + content: importContent, + userId: user.id, + }); + + // Update the message to include response and mark as completed + const { updateMessage } = await import('@buster/database'); + await updateMessage(messageId, { + isCompleted: true, + responseMessages: [ + { + id: assetId, + type: 'file', + file_type: assetType === 'metric' ? 'metric' : 'dashboard', + file_name: assetDetails.name, + version_number: assetDetails.versionNumber, + }, + ], + }); + + // Create the file association + const dbAssetType = convertChatAssetTypeToDatabaseAssetType(assetType); + await createMessageFileAssociation({ + messageId, + fileId: assetId, + fileType: dbAssetType, + version: assetDetails.versionNumber, + }); + + // Update the chat with most recent file information and title (matching Rust behavior) + const fileType = assetType === 'metric' ? 'metric' : 'dashboard'; + await db + .update(chats) + .set({ + title: assetDetails.name, // Set chat title to asset name + mostRecentFileId: assetId, + mostRecentFileType: fileType, + mostRecentVersionNumber: assetDetails.versionNumber, + updatedAt: new Date().toISOString(), + }) + .where(eq(chats.id, chatId)); + + // Return the message with the updated fields + return { + ...message, + isCompleted: true, + responseMessages: [ + { + id: assetId, + type: 'file', + file_type: assetType === 'metric' ? 'metric' : 'dashboard', + file_name: assetDetails.name, + version_number: assetDetails.versionNumber, + }, + ], + }; +} + +/** + * Builds a ChatMessage from a database Message for an asset import + */ +export function buildAssetImportChatMessage(message: Message, user: User): ChatMessage { + const responseMessages: Record = {}; + + // Parse response messages if they exist + if (Array.isArray(message.responseMessages)) { + for (const resp of message.responseMessages as ChatMessageResponseMessage[]) { + if ('id' in resp && resp.id) { + responseMessages[resp.id] = resp; + } + } + } + + return { + id: message.id, + created_at: message.createdAt, + updated_at: message.updatedAt, + request_message: message.requestMessage + ? { + request: message.requestMessage, + sender_id: user.id, + sender_name: user.name || user.email || 'Unknown User', + sender_avatar: user.avatarUrl || undefined, + } + : null, + response_messages: responseMessages, + response_message_ids: Object.keys(responseMessages), + reasoning_messages: {}, + reasoning_message_ids: [], + final_reasoning_message: null, + feedback: null, + is_completed: true, + post_processing_message: undefined, + }; +} diff --git a/apps/server/src/api/v2/chats/services/chat-helpers.test.ts b/apps/server/src/api/v2/chats/services/chat-helpers.test.ts index fcb54dd31..cee10e097 100644 --- a/apps/server/src/api/v2/chats/services/chat-helpers.test.ts +++ b/apps/server/src/api/v2/chats/services/chat-helpers.test.ts @@ -1,540 +1,531 @@ +import type { Chat, Message, User } from '@buster/database'; +import type { ChatAssetType, ChatWithMessages } from '@buster/server-shared/chats'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -// Mock the database connection and database functions BEFORE any other imports -vi.mock('@buster/database/connection', () => ({ - initializePool: vi.fn(), - getPool: vi.fn(), -})); +// Mock dependencies vi.mock('@buster/database', () => ({ - createMessage: vi.fn(), db: { - transaction: vi.fn((callback: any) => callback({ insert: vi.fn() })), + update: vi.fn().mockReturnThis(), + insert: vi.fn().mockReturnThis(), + select: vi.fn().mockReturnThis(), + from: vi.fn().mockReturnThis(), + where: vi.fn().mockReturnThis(), + set: vi.fn().mockReturnThis(), + returning: vi.fn().mockReturnThis(), + values: vi.fn().mockReturnThis(), + limit: vi.fn().mockReturnThis(), + transaction: vi.fn(), }, - getChatWithDetails: vi.fn(), - getMessagesForChat: vi.fn(), chats: {}, messages: {}, + messagesToFiles: {}, + dashboardFiles: {}, + metricFiles: {}, + eq: vi.fn(), + and: vi.fn(), + isNull: vi.fn(), + inArray: vi.fn(), + generateAssetMessages: vi.fn(), + createMessage: vi.fn(), + getChatWithDetails: vi.fn(), + getMessagesForChat: vi.fn(), + createMessageFileAssociation: vi.fn(), })); -// Mock the access-controls package vi.mock('@buster/access-controls', () => ({ canUserAccessChatCached: vi.fn(), })); -import { canUserAccessChatCached } from '@buster/access-controls'; -import * as database from '@buster/database'; -import type { Chat, Message } from '@buster/database'; -import { ChatError, ChatErrorCode } from '@buster/server-shared/chats'; -import { buildChatWithMessages, handleExistingChat, handleNewChat } from './chat-helpers'; +vi.mock('./server-asset-conversion', () => ({ + convertChatAssetTypeToDatabaseAssetType: vi.fn((type: ChatAssetType) => + type === 'metric' ? 'metric_file' : 'dashboard_file' + ), +})); -const mockUser = { - id: '550e8400-e29b-41d4-a716-446655440001', - name: 'Test User', - email: 'test@example.com', - avatarUrl: null, -}; +import { createMessage, db, generateAssetMessages } from '@buster/database'; +import { eq } from 'drizzle-orm'; +import { handleAssetChat, handleAssetChatWithPrompt } from './chat-helpers'; -const mockChat = { - id: 'chat-123', - title: 'Test Chat', - organizationId: '550e8400-e29b-41d4-a716-446655440000', - createdBy: '550e8400-e29b-41d4-a716-446655440001', - updatedBy: '550e8400-e29b-41d4-a716-446655440001', - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - publiclyAccessible: false, - deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - mostRecentVersionNumber: null, - slackChatAuthorization: null, - slackThreadTs: null, - slackChannelId: null, - workspaceSharingEnabledBy: null, - workspaceSharingEnabledAt: null, -} as Chat; +describe('chat-helpers', () => { + const mockUser: User = { + id: 'user-123', + name: 'Test User', + email: 'test@example.com', + avatarUrl: null, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + phoneNumber: null, + isBusterAdmin: false, + isEmailVerified: true, + authUserId: 'auth-123', + deletedAt: null, + } as User; -const mockMessage: Message = { - id: 'msg-123', - chatId: 'chat-123', - createdBy: 'user-123', - requestMessage: 'Test message', - responseMessages: {}, - reasoning: {}, - title: 'Test message', - rawLlmMessages: {}, - finalReasoningMessage: null, - isCompleted: false, - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - deletedAt: null, - feedback: null, - postProcessingMessage: null, - triggerRunId: null, -}; + const createMockChat = (): ChatWithMessages => ({ + id: 'chat-123', + title: 'Test Chat', + is_favorited: false, + message_ids: [], + messages: {}, + created_at: new Date().toISOString(), + updated_at: new Date().toISOString(), + created_by: 'user-123', + created_by_id: 'user-123', + created_by_name: 'Test User', + created_by_avatar: null, + publicly_accessible: false, + }); -describe('buildChatWithMessages', () => { - it('should build a ChatWithMessages object from database entities', () => { - const result = buildChatWithMessages(mockChat, [mockMessage], mockUser, true); - expect(result).toMatchObject({ - id: 'chat-123', - title: 'Test Chat', - is_favorited: true, - message_ids: ['msg-123'], - messages: { - 'msg-123': { - id: 'msg-123', - created_at: expect.any(String), - updated_at: expect.any(String), - request_message: { - request: 'Test message', - sender_id: 'user-123', - sender_name: 'Test User', - }, - response_messages: {}, - response_message_ids: [], - reasoning_message_ids: [], - reasoning_messages: {}, - final_reasoning_message: null, - feedback: null, - is_completed: false, - }, + const mockMetricAssetMessage: Message = { + id: 'import-msg-123', + chatId: 'chat-123', + createdBy: 'user-123', + requestMessage: null, // No request message for asset imports + responseMessages: [ + { + type: 'text', + id: 'text-msg-123', + message: + 'Test Metric has been pulled into a new chat.\n\nContinue chatting to modify or make changes to it.', + is_final_message: true, }, - created_at: expect.any(String), - updated_at: expect.any(String), - created_by: '550e8400-e29b-41d4-a716-446655440001', - created_by_id: '550e8400-e29b-41d4-a716-446655440001', - created_by_name: 'Test User', - created_by_avatar: null, - publicly_accessible: false, - permission: 'owner', - }); - }); + { + type: 'file', + id: 'asset-123', + file_type: 'metric', + file_name: 'Test Metric', + version_number: 1, + filter_version_id: null, + metadata: [ + { + status: 'completed', + message: 'Pulled into new chat', + timestamp: 1234567890, + }, + ], + }, + ], + reasoning: [], + finalReasoningMessage: '', + title: 'Test Metric', + rawLlmMessages: [ + { + role: 'user', + content: `I've imported the following metric:\n\nSuccessfully imported 1 metric file.\n\nFile details:\n[...]`, + }, + ], + isCompleted: true, + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + deletedAt: null, + feedback: null, + triggerRunId: null, + postProcessingMessage: null, + } as Message; - it('should handle missing creator details', () => { - const result = buildChatWithMessages(mockChat, [], null, false); - - expect(result.created_by_name).toBe('Unknown User'); - expect(result.created_by_avatar).toBeNull(); - }); -}); - -describe('handleExistingChat', () => { - const mockUser = { - id: 'user-1', - email: 'test@example.com', - name: 'Test User', - avatarUrl: null, + const mockDashboardAssetMessage: Message = { + ...mockMetricAssetMessage, + id: 'import-dashboard-msg-123', + title: 'Test Dashboard', + responseMessages: [ + { + type: 'text', + id: 'text-dashboard-msg-123', + message: + 'Test Dashboard has been pulled into a new chat.\n\nContinue chatting to modify or make changes to it.', + is_final_message: true, + }, + { + type: 'file', + id: 'dashboard-123', + file_type: 'dashboard', + file_name: 'Test Dashboard', + version_number: 1, + filter_version_id: null, + metadata: [ + { + status: 'completed', + message: 'Pulled into new chat', + timestamp: 1234567890, + }, + ], + }, + ], + rawLlmMessages: [ + { + role: 'user', + content: `I've imported the following dashboard:\n\nSuccessfully imported 1 dashboard file with 2 additional context files.\n\nFile details:\n[...]`, + }, + ], }; beforeEach(() => { vi.clearAllMocks(); }); - it('should handle existing chat with new message', async () => { - const mockChat = { - id: 'chat-1', - title: 'Test Chat', + describe('handleAssetChat', () => { + it('should handle metric import correctly', async () => { + vi.mocked(generateAssetMessages).mockReset().mockResolvedValue([mockMetricAssetMessage]); + vi.mocked(db.update).mockReturnValue({ + set: vi.fn().mockReturnValue({ + where: vi.fn().mockResolvedValue(undefined), + }), + } as any); + + const result = await handleAssetChat( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + mockUser, + createMockChat() + ); + + // Verify generateAssetMessages was called correctly + expect(generateAssetMessages).toHaveBeenCalledWith({ + assetId: 'asset-123', + assetType: 'metric_file', + userId: 'user-123', + chatId: 'chat-123', + }); + + // Verify chat was updated with asset info + expect(db.update).toHaveBeenCalledWith(expect.anything()); + const updateCall = vi.mocked(db.update).mock.calls[0]; + expect(updateCall).toBeDefined(); + + // Verify message was added to chat + expect(result.message_ids).toContain('import-msg-123'); + expect(result.messages['import-msg-123']).toBeDefined(); + + const importMessage = result.messages['import-msg-123']; + expect(importMessage?.request_message).toBeNull(); + expect(importMessage?.is_completed).toBe(true); + expect(importMessage?.response_messages).toHaveProperty('text-msg-123'); + expect(importMessage?.response_messages).toHaveProperty('asset-123'); + + // Verify response messages are in correct format + const textMessage = importMessage?.response_messages['text-msg-123']; + expect(textMessage).toMatchObject({ + type: 'text', + id: 'text-msg-123', + message: expect.stringContaining('Test Metric has been pulled into a new chat'), + is_final_message: true, + }); + + const fileMessage = importMessage?.response_messages['asset-123']; + expect(fileMessage).toMatchObject({ + type: 'file', + id: 'asset-123', + file_type: 'metric', + file_name: 'Test Metric', + version_number: 1, + }); + + // Verify chat title was updated + expect(result.title).toBe('Test Metric'); + }); + + it('should handle dashboard import with metrics correctly', async () => { + vi.mocked(generateAssetMessages).mockReset().mockResolvedValue([mockDashboardAssetMessage]); + vi.mocked(db.update).mockReturnValue({ + set: vi.fn().mockReturnValue({ + where: vi.fn().mockResolvedValue(undefined), + }), + } as any); + + const result = await handleAssetChat( + 'chat-123', + 'msg-123', + 'dashboard-123', + 'dashboard', + mockUser, + createMockChat() + ); + + // Verify generateAssetMessages was called correctly + expect(generateAssetMessages).toHaveBeenCalledWith({ + assetId: 'dashboard-123', + assetType: 'dashboard_file', + userId: 'user-123', + chatId: 'chat-123', + }); + + // Verify message was added to chat + expect(result.message_ids).toContain('import-dashboard-msg-123'); + expect(result.messages['import-dashboard-msg-123']).toBeDefined(); + + const importMessage = result.messages['import-dashboard-msg-123']; + expect(importMessage?.response_messages).toHaveProperty('dashboard-123'); + + // Verify dashboard file message + const fileMessage = importMessage?.response_messages['dashboard-123']; + expect(fileMessage).toMatchObject({ + type: 'file', + id: 'dashboard-123', + file_type: 'dashboard', + file_name: 'Test Dashboard', + version_number: 1, + }); + + // Verify chat title was updated + expect(result.title).toBe('Test Dashboard'); + }); + + it('should handle errors gracefully', async () => { + vi.mocked(generateAssetMessages).mockReset().mockRejectedValue(new Error('Database error')); + const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {}); + + const result = await handleAssetChat( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + mockUser, + createMockChat() + ); + + expect(consoleSpy).toHaveBeenCalledWith( + 'Failed to handle asset chat:', + expect.objectContaining({ + chatId: 'chat-123', + assetId: 'asset-123', + chatAssetType: 'metric', + userId: 'user-123', + }) + ); + + // Should return original chat without modifications + // Note: We can't do exact comparison due to timestamp differences + expect(result.id).toBe('chat-123'); + expect(result.title).toBe('Test Chat'); + expect(result.message_ids).toEqual([]); + expect(result.messages).toEqual({}); + + consoleSpy.mockRestore(); + }); + }); + + describe('handleAssetChatWithPrompt', () => { + const mockUserMessage: Message = { + id: 'user-msg-123', + chatId: 'chat-123', + createdBy: 'user-123', + requestMessage: 'Tell me about this metric', + responseMessages: {}, + reasoning: {}, + finalReasoningMessage: null, + title: 'Tell me about this metric', + rawLlmMessages: {}, + isCompleted: false, createdAt: new Date().toISOString(), updatedAt: new Date().toISOString(), - createdBy: 'user-1', - updatedBy: 'user-1', - organizationId: 'org-1', - publiclyAccessible: false, deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - } as Chat; - - const mockMessage = { - id: 'message-1', - chatId: 'chat-1', - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - createdBy: 'user-1', - requestMessage: 'Test message', - responseMessages: {}, - reasoning: {}, - title: 'Test message', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, feedback: null, + triggerRunId: null, + postProcessingMessage: null, } as Message; - vi.mocked(database.getChatWithDetails).mockResolvedValue({ - chat: mockChat, - user: mockUser as unknown as any, - isFavorited: false, + it('should create import message then user prompt message', async () => { + // Only return the metric message for this test + const metricOnlyMessage = { ...mockMetricAssetMessage }; + vi.mocked(generateAssetMessages).mockReset().mockResolvedValueOnce([metricOnlyMessage]); + vi.mocked(createMessage).mockResolvedValue(mockUserMessage); + vi.mocked(db.update).mockReturnValue({ + set: vi.fn().mockReturnValue({ + where: vi.fn().mockResolvedValue(undefined), + }), + } as any); + + const result = await handleAssetChatWithPrompt( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + 'Tell me about this metric', + mockUser, + createMockChat() + ); + + // Verify generateAssetMessages was called (same as handleAssetChat) + expect(generateAssetMessages).toHaveBeenCalledWith({ + assetId: 'asset-123', + assetType: 'metric_file', + userId: 'user-123', + chatId: 'chat-123', + }); + + // Verify createMessage was called for user prompt + expect(createMessage).toHaveBeenCalledWith({ + messageId: expect.any(String), + chatId: 'chat-123', + content: 'Tell me about this metric', + userId: 'user-123', + }); + + // Verify both messages were added to chat in correct order + expect(result.message_ids).toHaveLength(2); + expect(result.message_ids[0]).toBe('import-msg-123'); + expect(result.message_ids[1]).toBe('user-msg-123'); + + // Verify import message structure + const importMessage = result.messages['import-msg-123']; + expect(importMessage?.request_message).toBeNull(); + expect(importMessage?.is_completed).toBe(true); + expect(importMessage?.response_messages).toHaveProperty('asset-123'); + + // Verify user message structure + const userMessage = result.messages['user-msg-123']; + expect(userMessage?.request_message).toMatchObject({ + request: 'Tell me about this metric', + sender_id: 'user-123', + sender_name: 'Test User', + }); + expect(userMessage?.is_completed).toBe(false); + expect(userMessage?.response_messages).toEqual({}); + + // Verify chat title was updated to asset name + expect(result.title).toBe('Test Metric'); }); - vi.mocked(canUserAccessChatCached).mockResolvedValue(true); - vi.mocked(database.createMessage).mockResolvedValue(mockMessage); - vi.mocked(database.getMessagesForChat).mockResolvedValue([mockMessage]); + it('should handle dashboard with prompt correctly', async () => { + // Only return the dashboard message for this test + vi.mocked(generateAssetMessages) + .mockReset() + .mockResolvedValueOnce([mockDashboardAssetMessage]); + vi.mocked(createMessage).mockResolvedValue({ + ...mockUserMessage, + requestMessage: 'Explain this dashboard', + } as Message); + vi.mocked(db.update).mockReturnValue({ + set: vi.fn().mockReturnValue({ + where: vi.fn().mockResolvedValue(undefined), + }), + } as any); - const result = await handleExistingChat('chat-1', 'message-1', 'Test message', mockUser); + const result = await handleAssetChatWithPrompt( + 'chat-123', + 'msg-123', + 'dashboard-123', + 'dashboard', + 'Explain this dashboard', + mockUser, + createMockChat() + ); - expect(result.chatId).toBe('chat-1'); - expect(result.messageId).toBe('message-1'); - expect(result.chat.messages['message-1']).toBeDefined(); - }); + // Verify generateAssetMessages was called for dashboard + expect(generateAssetMessages).toHaveBeenCalledWith({ + assetId: 'dashboard-123', + assetType: 'dashboard_file', + userId: 'user-123', + chatId: 'chat-123', + }); - it('should throw error if chat not found', async () => { - vi.mocked(database.getChatWithDetails).mockResolvedValue(null); + // Verify both messages exist + expect(result.message_ids).toHaveLength(2); + expect(result.messages['import-dashboard-msg-123']).toBeDefined(); + expect(result.messages['user-msg-123']).toBeDefined(); - await expect( - handleExistingChat('chat-1', 'message-1', 'Test message', mockUser) - ).rejects.toThrow(new ChatError(ChatErrorCode.CHAT_NOT_FOUND, 'Chat not found', 404)); - }); - - it('should throw error if permission denied', async () => { - vi.mocked(database.getChatWithDetails).mockResolvedValue({ - chat: { id: 'chat-1' } as Chat, - user: null, - isFavorited: false, - }); - vi.mocked(canUserAccessChatCached).mockResolvedValue(false); - - await expect( - handleExistingChat('chat-1', 'message-1', 'Test message', mockUser) - ).rejects.toThrow( - new ChatError( - ChatErrorCode.PERMISSION_DENIED, - 'You do not have permission to access this chat', - 403 - ) - ); - }); - - it('should prepend new message to maintain descending order (newest first)', async () => { - const baseTime = new Date('2024-01-01T10:00:00Z'); - - const mockChat = { - id: 'chat-1', - title: 'Test Chat', - createdAt: baseTime.toISOString(), - updatedAt: baseTime.toISOString(), - createdBy: 'user-1', - updatedBy: 'user-1', - organizationId: 'org-1', - publiclyAccessible: false, - deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - } as Chat; - - const existingMessage1 = { - id: 'message-1', - chatId: 'chat-1', - createdAt: new Date(baseTime.getTime() + 2000).toISOString(), // 2 seconds later - updatedAt: new Date(baseTime.getTime() + 2000).toISOString(), - createdBy: 'user-1', - requestMessage: 'Second message', - responseMessages: {}, - reasoning: {}, - title: 'Second message', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, - feedback: null, - } as Message; - - const existingMessage2 = { - id: 'message-2', - chatId: 'chat-1', - createdAt: baseTime.toISOString(), // Original time - updatedAt: baseTime.toISOString(), - createdBy: 'user-1', - requestMessage: 'First message', - responseMessages: {}, - reasoning: {}, - title: 'First message', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, - feedback: null, - } as Message; - - const newMessage = { - id: 'message-3', - chatId: 'chat-1', - createdAt: new Date(baseTime.getTime() + 4000).toISOString(), // 4 seconds later (newest) - updatedAt: new Date(baseTime.getTime() + 4000).toISOString(), - createdBy: 'user-1', - requestMessage: 'Third message (newest)', - responseMessages: {}, - reasoning: {}, - title: 'Third message (newest)', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, - feedback: null, - } as Message; - - vi.mocked(database.getChatWithDetails).mockResolvedValue({ - chat: mockChat, - user: mockUser as unknown as any, - isFavorited: false, + // Verify chat title is dashboard name + expect(result.title).toBe('Test Dashboard'); }); - vi.mocked(canUserAccessChatCached).mockResolvedValue(true); - vi.mocked(database.createMessage).mockResolvedValue(newMessage); - // getMessagesForChat returns existing messages in descending order (newest first) - vi.mocked(database.getMessagesForChat).mockResolvedValue([existingMessage1, existingMessage2]); + it('should handle missing asset gracefully', async () => { + vi.mocked(generateAssetMessages).mockReset().mockResolvedValueOnce([]); + vi.mocked(createMessage).mockResolvedValue(mockUserMessage); + const consoleSpy = vi.spyOn(console, 'warn').mockImplementation(() => {}); - const result = await handleExistingChat( - 'chat-1', - 'message-3', - 'Third message (newest)', - mockUser - ); + const result = await handleAssetChatWithPrompt( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + 'Tell me about this metric', + mockUser, + createMockChat() + ); - expect(result.chat.messages['message-3']).toBeDefined(); - expect(result.chat.messages['message-1']).toBeDefined(); - expect(result.chat.messages['message-2']).toBeDefined(); + expect(consoleSpy).toHaveBeenCalledWith( + 'No asset messages generated', + expect.objectContaining({ + assetId: 'asset-123', + assetType: 'metric_file', + userId: 'user-123', + chatId: 'chat-123', + }) + ); - expect(result.chat.message_ids).toEqual(['message-3', 'message-1', 'message-2']); + // Should still create user message + expect(createMessage).toHaveBeenCalled(); + expect(result.message_ids).toContain('user-msg-123'); - expect(result.chat.message_ids[0]).toBe('message-3'); - - const message3 = result.chat.messages['message-3']; - const message1 = result.chat.messages['message-1']; - const message2 = result.chat.messages['message-2']; - - expect(message3).toBeDefined(); - expect(message1).toBeDefined(); - expect(message2).toBeDefined(); - - const message3Time = new Date(message3!.created_at).getTime(); - const message1Time = new Date(message1!.created_at).getTime(); - const message2Time = new Date(message2!.created_at).getTime(); - - expect(message3Time).toBeGreaterThan(message1Time); - expect(message1Time).toBeGreaterThan(message2Time); - }); - - it('should handle single existing message with new message correctly', async () => { - const baseTime = new Date('2024-01-01T10:00:00Z'); - - const mockChat = { - id: 'chat-1', - title: 'Test Chat', - createdAt: baseTime.toISOString(), - updatedAt: baseTime.toISOString(), - createdBy: 'user-1', - updatedBy: 'user-1', - organizationId: 'org-1', - publiclyAccessible: false, - deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - } as Chat; - - const existingMessage = { - id: 'message-1', - chatId: 'chat-1', - createdAt: baseTime.toISOString(), - updatedAt: baseTime.toISOString(), - createdBy: 'user-1', - requestMessage: 'First message', - responseMessages: {}, - reasoning: {}, - title: 'First message', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, - feedback: null, - } as Message; - - const newMessage = { - id: 'message-2', - chatId: 'chat-1', - createdAt: new Date(baseTime.getTime() + 1000).toISOString(), // 1 second later - updatedAt: new Date(baseTime.getTime() + 1000).toISOString(), - createdBy: 'user-1', - requestMessage: 'Second message', - responseMessages: {}, - reasoning: {}, - title: 'Second message', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, - feedback: null, - } as Message; - - vi.mocked(database.getChatWithDetails).mockResolvedValue({ - chat: mockChat, - user: mockUser as unknown as any, - isFavorited: false, + consoleSpy.mockRestore(); }); - vi.mocked(canUserAccessChatCached).mockResolvedValue(true); - vi.mocked(database.createMessage).mockResolvedValue(newMessage); - vi.mocked(database.getMessagesForChat).mockResolvedValue([existingMessage]); + it('should handle errors and still create user message', async () => { + vi.mocked(generateAssetMessages) + .mockReset() + .mockRejectedValueOnce(new Error('Database error')); + vi.mocked(createMessage).mockResolvedValue(mockUserMessage); + const consoleSpy = vi.spyOn(console, 'error').mockImplementation(() => {}); - const result = await handleExistingChat('chat-1', 'message-2', 'Second message', mockUser); + const result = await handleAssetChatWithPrompt( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + 'Tell me about this metric', + mockUser, + createMockChat() + ); - expect(result.chat.messages['message-1']).toBeDefined(); - expect(result.chat.messages['message-2']).toBeDefined(); + expect(consoleSpy).toHaveBeenCalled(); - expect(result.chat.message_ids).toEqual(['message-2', 'message-1']); + // Should still create user message despite error + expect(createMessage).toHaveBeenCalled(); + expect(result.message_ids).toContain('user-msg-123'); + expect(result.messages['user-msg-123']?.request_message?.request).toBe( + 'Tell me about this metric' + ); - expect(result.chat.message_ids[0]).toBe('message-2'); - }); - - it('should handle empty existing messages with new message', async () => { - const baseTime = new Date('2024-01-01T10:00:00Z'); - - const mockChat = { - id: 'chat-1', - title: 'Test Chat', - createdAt: baseTime.toISOString(), - updatedAt: baseTime.toISOString(), - createdBy: 'user-1', - updatedBy: 'user-1', - organizationId: 'org-1', - publiclyAccessible: false, - deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - } as Chat; - - const newMessage = { - id: 'message-1', - chatId: 'chat-1', - createdAt: baseTime.toISOString(), - updatedAt: baseTime.toISOString(), - createdBy: 'user-1', - requestMessage: 'First message', - responseMessages: {}, - reasoning: {}, - title: 'First message', - rawLlmMessages: {}, - isCompleted: false, - deletedAt: null, - finalReasoningMessage: null, - feedback: null, - } as Message; - - vi.mocked(database.getChatWithDetails).mockResolvedValue({ - chat: mockChat, - user: mockUser as unknown as any, - isFavorited: false, + consoleSpy.mockRestore(); }); - vi.mocked(canUserAccessChatCached).mockResolvedValue(true); - vi.mocked(database.createMessage).mockResolvedValue(newMessage); - vi.mocked(database.getMessagesForChat).mockResolvedValue([]); // No existing messages + it('should maintain message order when multiple messages exist', async () => { + const chatWithExistingMessages = { + ...createMockChat(), + message_ids: ['existing-msg-1', 'existing-msg-2'], + messages: { + 'existing-msg-1': {} as any, + 'existing-msg-2': {} as any, + }, + }; - const result = await handleExistingChat('chat-1', 'message-1', 'First message', mockUser); + vi.mocked(generateAssetMessages).mockReset().mockResolvedValueOnce([mockMetricAssetMessage]); + vi.mocked(createMessage).mockResolvedValue(mockUserMessage); + vi.mocked(db.update).mockReturnValue({ + set: vi.fn().mockReturnValue({ + where: vi.fn().mockResolvedValue(undefined), + }), + } as any); - expect(result.chat.messages['message-1']).toBeDefined(); - expect(Object.keys(result.chat.messages)).toHaveLength(1); + const result = await handleAssetChatWithPrompt( + 'chat-123', + 'msg-123', + 'asset-123', + 'metric', + 'Tell me about this metric', + mockUser, + chatWithExistingMessages + ); - expect(result.chat.message_ids).toEqual(['message-1']); - }); -}); - -describe('handleNewChat', () => { - const mockUser = { - id: 'user-1', - email: 'test@example.com', - name: 'Test User', - avatarUrl: null, - }; - - beforeEach(() => { - vi.clearAllMocks(); - }); - - it('should create new chat with message', async () => { - const mockChat = { - id: 'chat-1', - title: 'Test Chat', - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - createdBy: 'user-1', - updatedBy: 'user-1', - organizationId: 'org-1', - publiclyAccessible: false, - deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - } as Chat; - - const mockTx = { - insert: vi.fn().mockReturnThis(), - values: vi.fn().mockReturnThis(), - returning: vi.fn().mockResolvedValue([mockChat]), - }; - - vi.mocked(database.db.transaction).mockImplementation((callback: any) => callback(mockTx)); - - const result = await handleNewChat({ - title: 'Test Chat', - messageId: 'message-1', - prompt: 'Test message', - user: mockUser, - organizationId: 'org-1', - }); - - expect(result.chatId).toBe('chat-1'); - expect(result.messageId).toBe('message-1'); - }); - - it('should create new chat without message', async () => { - const mockChat = { - id: 'chat-1', - title: 'Test Chat', - createdAt: new Date().toISOString(), - updatedAt: new Date().toISOString(), - createdBy: 'user-1', - updatedBy: 'user-1', - organizationId: 'org-1', - publiclyAccessible: false, - deletedAt: null, - publiclyEnabledBy: null, - publicExpiryDate: null, - mostRecentFileId: null, - mostRecentFileType: null, - } as Chat; - - const mockTx = { - insert: vi.fn().mockReturnThis(), - values: vi.fn().mockReturnThis(), - returning: vi.fn().mockResolvedValue([mockChat]), - }; - - vi.mocked(database.db.transaction).mockImplementation((callback: any) => callback(mockTx)); - - const result = await handleNewChat({ - title: 'Test Chat', - messageId: 'message-1', - prompt: undefined, - user: mockUser, - organizationId: 'org-1', - }); - - expect(result.chatId).toBe('chat-1'); - expect(result.messageId).toBe('message-1'); - expect(Object.keys(result.chat.messages)).toHaveLength(0); + // Should preserve existing messages and add new ones in order + expect(result.message_ids).toEqual([ + 'existing-msg-1', + 'existing-msg-2', + 'import-msg-123', + 'user-msg-123', + ]); + }); }); }); diff --git a/apps/server/src/api/v2/chats/services/chat-helpers.ts b/apps/server/src/api/v2/chats/services/chat-helpers.ts index 4a8313e9d..f629648bf 100644 --- a/apps/server/src/api/v2/chats/services/chat-helpers.ts +++ b/apps/server/src/api/v2/chats/services/chat-helpers.ts @@ -63,12 +63,7 @@ function buildMessages( } } - // Early return for already-correct format - if (parsedMessages && typeof parsedMessages === 'object' && !Array.isArray(parsedMessages)) { - return parsedMessages as Record; - } - - // Optimized array processing with pre-allocation and validation + // Handle array format (new format from generateAssetMessages) if (Array.isArray(parsedMessages)) { const result: Record = {}; for (let i = 0; i < parsedMessages.length; i++) { @@ -80,6 +75,11 @@ function buildMessages( return result; } + // Handle object format (legacy format with IDs as keys) + if (parsedMessages && typeof parsedMessages === 'object' && !Array.isArray(parsedMessages)) { + return parsedMessages as Record; + } + return {}; } @@ -388,6 +388,10 @@ export async function handleAssetChat( // Convert and add to chat for (const msg of assetMessages) { + // Build response messages from the database message + const responseMessages = buildResponseMessages(msg.responseMessages); + const responseMessageIds = Object.keys(responseMessages); + const chatMessage: ChatMessage = { id: msg.id, created_at: msg.createdAt, @@ -400,13 +404,13 @@ export async function handleAssetChat( sender_avatar: chat.created_by_avatar || undefined, } : null, - response_messages: {}, - response_message_ids: [], + response_messages: responseMessages, + response_message_ids: responseMessageIds, reasoning_message_ids: [], reasoning_messages: {}, - final_reasoning_message: null, + final_reasoning_message: msg.finalReasoningMessage || null, feedback: null, - is_completed: false, + is_completed: msg.isCompleted || false, post_processing_message: validateNullableJsonb( msg.postProcessingMessage, PostProcessingMessageSchema @@ -420,6 +424,26 @@ export async function handleAssetChat( chat.messages[msg.id] = chatMessage; } + // Update the chat with most recent file information and title (matching Rust behavior) + const fileType = chatAssetType === 'metric' ? 'metric' : 'dashboard'; + + // Get the asset name from the first message + const assetName = assetMessages[0]?.title || ''; + + await db + .update(chats) + .set({ + title: assetName, // Set chat title to asset name + mostRecentFileId: assetId, + mostRecentFileType: fileType, + mostRecentVersionNumber: 1, // Asset imports always start at version 1 + updatedAt: new Date().toISOString(), + }) + .where(eq(chats.id, chatId)); + + // Update the chat object with the new title + chat.title = assetName; + return chat; } catch (error) { console.error('Failed to handle asset chat:', { @@ -442,6 +466,222 @@ export async function handleAssetChat( } } +/** + * Handle asset-based chat initialization with a prompt + * This creates an import message for the asset, then adds the user's prompt as a follow-up + */ +export async function handleAssetChatWithPrompt( + chatId: string, + _messageId: string, // Initial message ID (not used since we create two messages) + assetId: string, + chatAssetType: ChatAssetType, + prompt: string, + user: User, + chat: ChatWithMessages +): Promise { + const userId = user.id; + try { + // First, use the exact same logic as handleAssetChat to import the asset + // This ensures we get dashboard metrics and proper formatting + const assetType = convertChatAssetTypeToDatabaseAssetType(chatAssetType); + const assetMessages = await generateAssetMessages({ + assetId, + assetType, + userId, + chatId, + }); + + if (!assetMessages || assetMessages.length === 0) { + console.warn('No asset messages generated', { + assetId, + assetType, + userId, + chatId, + }); + // Still create the user message with prompt + const userMessageId = crypto.randomUUID(); + const userMessage = await createMessage({ + messageId: userMessageId, + chatId, + content: prompt, + userId: user.id, + }); + + // Add to chat + const chatMessage: ChatMessage = { + id: userMessage.id, + created_at: userMessage.createdAt, + updated_at: userMessage.updatedAt, + request_message: { + request: prompt, + sender_id: user.id, + sender_name: chat.created_by_name, + sender_avatar: chat.created_by_avatar || undefined, + }, + response_messages: {}, + response_message_ids: [], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: false, + post_processing_message: undefined, + }; + + if (!chat.message_ids.includes(userMessage.id)) { + chat.message_ids.push(userMessage.id); + } + chat.messages[userMessage.id] = chatMessage; + + return chat; + } + + // Add the import message to chat (exact same logic as handleAssetChat) + for (const msg of assetMessages) { + // Build response messages from the database message + const responseMessages = buildResponseMessages(msg.responseMessages); + const responseMessageIds = Object.keys(responseMessages); + + const chatMessage: ChatMessage = { + id: msg.id, + created_at: msg.createdAt, + updated_at: msg.updatedAt, + request_message: msg.requestMessage + ? { + request: msg.requestMessage, + sender_id: msg.createdBy, + sender_name: chat.created_by_name, + sender_avatar: chat.created_by_avatar || undefined, + } + : null, + response_messages: responseMessages, + response_message_ids: responseMessageIds, + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: msg.finalReasoningMessage || null, + feedback: null, + is_completed: msg.isCompleted || false, + post_processing_message: validateNullableJsonb( + msg.postProcessingMessage, + PostProcessingMessageSchema + ), + }; + + // Only add message ID if it doesn't already exist + if (!chat.message_ids.includes(msg.id)) { + chat.message_ids.push(msg.id); + } + chat.messages[msg.id] = chatMessage; + } + + // Update the chat with most recent file information and title (matching handleAssetChat) + const fileType = chatAssetType === 'metric' ? 'metric' : 'dashboard'; + const assetName = assetMessages[0]?.title || ''; + + await db + .update(chats) + .set({ + title: assetName, // Set chat title to asset name + mostRecentFileId: assetId, + mostRecentFileType: fileType, + mostRecentVersionNumber: 1, // Asset imports always start at version 1 + updatedAt: new Date().toISOString(), + }) + .where(eq(chats.id, chatId)); + + // Update the chat object with the new title + chat.title = assetName; + + // Then, create the user's prompt message as a follow-up + const userMessageId = crypto.randomUUID(); + const userMessage = await createMessage({ + messageId: userMessageId, + chatId, + content: prompt, + userId: user.id, + }); + + // Add user message to chat + const userChatMessage: ChatMessage = { + id: userMessage.id, + created_at: userMessage.createdAt, + updated_at: userMessage.updatedAt, + request_message: { + request: prompt, + sender_id: user.id, + sender_name: chat.created_by_name, + sender_avatar: chat.created_by_avatar || undefined, + }, + response_messages: {}, + response_message_ids: [], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: false, + post_processing_message: undefined, + }; + + if (!chat.message_ids.includes(userMessage.id)) { + chat.message_ids.push(userMessage.id); + } + chat.messages[userMessage.id] = userChatMessage; + + return chat; + } catch (error) { + console.error('Failed to handle asset chat with prompt:', { + chatId, + assetId, + chatAssetType, + userId, + error: + error instanceof Error + ? { + message: error.message, + stack: error.stack, + name: error.name, + } + : String(error), + }); + + // Don't fail the entire request, create the user message anyway + const userMessageId = crypto.randomUUID(); + const userMessage = await createMessage({ + messageId: userMessageId, + chatId, + content: prompt, + userId: user.id, + }); + + const chatMessage: ChatMessage = { + id: userMessage.id, + created_at: userMessage.createdAt, + updated_at: userMessage.updatedAt, + request_message: { + request: prompt, + sender_id: user.id, + sender_name: chat.created_by_name, + sender_avatar: chat.created_by_avatar || undefined, + }, + response_messages: {}, + response_message_ids: [], + reasoning_message_ids: [], + reasoning_messages: {}, + final_reasoning_message: null, + feedback: null, + is_completed: false, + post_processing_message: undefined, + }; + + if (!chat.message_ids.includes(userMessage.id)) { + chat.message_ids.push(userMessage.id); + } + chat.messages[userMessage.id] = chatMessage; + + return chat; + } +} + /** * Soft delete a message and all subsequent messages in the same chat * Used for "redo from this point" functionality diff --git a/apps/web/src/components/features/metrics/hooks/useMetricDrilldownItem.tsx b/apps/web/src/components/features/metrics/hooks/useMetricDrilldownItem.tsx new file mode 100644 index 000000000..8a302216e --- /dev/null +++ b/apps/web/src/components/features/metrics/hooks/useMetricDrilldownItem.tsx @@ -0,0 +1,25 @@ +import React, { useMemo } from 'react'; +import type { DropdownItem } from '@/components/ui/dropdown'; +import { WandSparkle } from '@/components/ui/icons'; +import { FollowUpWithAssetContent } from '@/components/features/popups/FollowUpWithAsset'; + +export const useMetricDrilldownItem = ({ metricId }: { metricId: string }): DropdownItem => { + return useMemo( + () => ({ + value: 'drilldown', + label: 'Drill down & filter', + items: [ + + ], + icon: + }), + [metricId] + ); +}; diff --git a/apps/web/src/components/features/popups/FollowUpWithAsset.tsx b/apps/web/src/components/features/popups/FollowUpWithAsset.tsx index 45f558a43..a070d26ff 100644 --- a/apps/web/src/components/features/popups/FollowUpWithAsset.tsx +++ b/apps/web/src/components/features/popups/FollowUpWithAsset.tsx @@ -8,6 +8,8 @@ import { AppTooltip } from '../../ui/tooltip'; import { useAppLayoutContextSelector } from '@/context/BusterAppLayout'; import { assetParamsToRoute } from '../../../lib/assets'; +type FollowUpMode = 'filter' | 'drilldown'; + type FollowUpWithAssetProps = { assetType: Exclude; assetId: string; @@ -16,6 +18,7 @@ type FollowUpWithAssetProps = { align?: PopoverProps['align']; placeholder?: string; buttonText?: string; + mode?: FollowUpMode; }; export const FollowUpWithAssetContent: React.FC<{ @@ -23,22 +26,41 @@ export const FollowUpWithAssetContent: React.FC<{ assetId: string; placeholder?: string; buttonText?: string; + mode?: FollowUpMode; }> = React.memo( ({ assetType, assetId, placeholder = 'Describe the filter you want to apply', - buttonText = 'Apply custom filter' + buttonText = 'Apply custom filter', + mode }) => { const { mutateAsync: startChatFromAsset, isPending } = useStartChatFromAsset(); const onChangePage = useAppLayoutContextSelector((x) => x.onChangePage); + const transformPrompt = useMemoizedFn((userPrompt: string): string => { + if (!mode) return userPrompt; + + if (mode === 'filter') { + return `Hey Buster. Please recreate this dashboard applying this filter to the metrics on the dashboard: ${userPrompt}`; + } + + if (mode === 'drilldown') { + return `Hey Buster. Can you filter or drill down into this metric based on the following request: ${userPrompt}`; + } + + return userPrompt; + }); + const onSubmit = useMemoizedFn(async (prompt: string) => { if (!prompt || !assetId || !assetType || isPending) return; + + const transformedPrompt = transformPrompt(prompt); + const res = await startChatFromAsset({ asset_id: assetId, asset_type: assetType, - prompt + prompt: transformedPrompt }); const link = assetParamsToRoute({ assetId, @@ -63,7 +85,16 @@ export const FollowUpWithAssetContent: React.FC<{ FollowUpWithAssetContent.displayName = 'FollowUpWithAssetContent'; export const FollowUpWithAssetPopup: React.FC = React.memo( - ({ assetType, assetId, side = 'bottom', align = 'end', children, placeholder, buttonText }) => { + ({ + assetType, + assetId, + side = 'bottom', + align = 'end', + children, + placeholder, + buttonText, + mode + }) => { return ( = React.me assetId={assetId} placeholder={placeholder} buttonText={buttonText} + mode={mode} /> }> {children} diff --git a/apps/web/src/components/ui/card/InputCard.tsx b/apps/web/src/components/ui/card/InputCard.tsx index c68b62798..eb4d0cca9 100644 --- a/apps/web/src/components/ui/card/InputCard.tsx +++ b/apps/web/src/components/ui/card/InputCard.tsx @@ -32,6 +32,13 @@ export const InputCard: React.FC = ({ const disableSubmit = !inputHasText(inputValue) || loading; + const handlePressEnter = (e: React.KeyboardEvent) => { + if (!disableSubmit) { + e.preventDefault(); + onSubmit?.(inputValue); + } + }; + const spacingClass = 'py-2.5 px-3'; return ( @@ -42,6 +49,7 @@ export const InputCard: React.FC = ({ value={inputValue} readOnly={loading} onChange={handleChange} + onPressEnter={handlePressEnter} autoResize={{ minRows: 5, maxRows: 10 diff --git a/apps/web/src/controllers/DashboardController/DashboardViewDashboardController/DashboardContentController/DashboardMetricItem/MetricItemCardThreeDotMenu.tsx b/apps/web/src/controllers/DashboardController/DashboardViewDashboardController/DashboardContentController/DashboardMetricItem/MetricItemCardThreeDotMenu.tsx index 87ceb09c3..ededc1162 100644 --- a/apps/web/src/controllers/DashboardController/DashboardViewDashboardController/DashboardContentController/DashboardMetricItem/MetricItemCardThreeDotMenu.tsx +++ b/apps/web/src/controllers/DashboardController/DashboardViewDashboardController/DashboardContentController/DashboardMetricItem/MetricItemCardThreeDotMenu.tsx @@ -7,7 +7,6 @@ import { Dropdown, type DropdownItems, type DropdownItem } from '@/components/ui import { DotsVertical, Trash, - WandSparkle, ShareRight, PenSparkle, SquareChartPen, @@ -17,7 +16,6 @@ import { cn } from '@/lib/utils'; import { ASSET_ICONS } from '@/components/features/config/assetIcons'; import { assetParamsToRoute } from '@/lib/assets/assetParamsToRoute'; import { useChatLayoutContextSelector } from '@/layouts/ChatLayout'; -import { FollowUpWithAssetContent } from '@/components/features/popups/FollowUpWithAsset'; import { useGetMetric } from '@/api/buster_rest/metrics'; import { getShareAssetConfig } from '@/components/features/ShareMenu/helpers'; import { getIsEffectiveOwner } from '@/lib/share'; @@ -28,6 +26,7 @@ import { useFavoriteMetricSelectMenu, useVersionHistorySelectMenu } from '@/components/features/metrics/ThreeDotMenu'; +import { useMetricDrilldownItem } from '@/components/features/metrics/hooks/useMetricDrilldownItem'; export const MetricItemCardThreeDotMenu: React.FC<{ dashboardId: string; @@ -50,7 +49,7 @@ const MetricItemCardThreeDotMenuPopover: React.FC<{ const chatId = useChatLayoutContextSelector((x) => x.chatId); const removeFromDashboardItem = useRemoveFromDashboardItem({ dashboardId, metricId }); const openChartItem = useOpenChartItem({ dashboardId, metricId, chatId }); - const drilldownItem = useDrilldownItem({ metricId }); + const drilldownItem = useMetricDrilldownItem({ metricId }); const shareMenu = useShareMenuSelectMenu({ metricId }); const editWithAI = useEditWithAI({ metricId, dashboardId, chatId }); const editChartButton = useEditChartButton({ metricId, dashboardId, chatId }); @@ -166,26 +165,6 @@ const useOpenChartItem = ({ }; }; -const useDrilldownItem = ({ metricId }: { metricId: string }): DropdownItem => { - return useMemo( - () => ({ - value: 'drilldown', - label: 'Drill down & filter', - items: [ - - ], - icon: - }), - [metricId] - ); -}; - const useShareMenuSelectMenu = ({ metricId }: { metricId: string }): DropdownItem | undefined => { const { data: shareAssetConfig } = useGetMetric( { id: metricId }, diff --git a/apps/web/src/layouts/ChatLayout/FileContainer/FileContainerHeader/DashboardContainerHeaderButtons/DashboardContainerHeaderButtons.tsx b/apps/web/src/layouts/ChatLayout/FileContainer/FileContainerHeader/DashboardContainerHeaderButtons/DashboardContainerHeaderButtons.tsx index b1bbe7a85..9a2dabe51 100644 --- a/apps/web/src/layouts/ChatLayout/FileContainer/FileContainerHeader/DashboardContainerHeaderButtons/DashboardContainerHeaderButtons.tsx +++ b/apps/web/src/layouts/ChatLayout/FileContainer/FileContainerHeader/DashboardContainerHeaderButtons/DashboardContainerHeaderButtons.tsx @@ -84,7 +84,7 @@ AddContentToDashboardButton.displayName = 'AddContentToDashboardButton'; const FollowUpWithAssetButton = React.memo(({ dashboardId }: { dashboardId: string }) => { return ( - +