Merge branch 'staging' into big-nate/bus-1483-quick-win-for-filter-dashboard-drill-downexplore-metric

apps/trigger/package.json

@@ -27,6 +27,7 @@
     "@buster/test-utils": "workspace:*",
     "@buster/typescript-config": "workspace:*",
     "@buster/vitest-config": "workspace:*",
+    "@buster/web-tools": "workspace:*",
     "@mastra/core": "catalog:",
     "@trigger.dev/sdk": "catalog:",
     "ai": "catalog:",

apps/trigger/src/tasks/message-post-processing/helpers/message-fetchers.ts

@@ -1,5 +1,16 @@
 import { getPermissionedDatasets } from '@buster/access-controls';
-import { and, chats, eq, getDb, isNotNull, isNull, lte, messages, users } from '@buster/database';
+import {
+  and,
+  chats,
+  eq,
+  getChatConversationHistory,
+  getDb,
+  isNotNull,
+  isNull,
+  lte,
+  messages,
+  users,
+} from '@buster/database';
 import type { CoreMessage } from 'ai';
 import {
   DataFetchError,
@@ -37,12 +48,25 @@ export async function fetchMessageWithContext(messageId: string): Promise<Messag
       throw new MessageNotFoundError(messageId);
     }
 
+    // Get the complete conversation history using the fixed helper
+    // This ensures we get the most recent completed message with valid rawLlmMessages
+    let conversationHistory: CoreMessage[] = [];
+    try {
+      conversationHistory = await getChatConversationHistory({ messageId });
+    } catch (_error) {
+      // If we can't get conversation history, fall back to the current message's rawLlmMessages
+      // This handles edge cases where there are no completed messages yet
+      if (messageData.rawLlmMessages && Array.isArray(messageData.rawLlmMessages)) {
+        conversationHistory = messageData.rawLlmMessages as CoreMessage[];
+      }
+    }
+
     return {
       id: messageData.id,
       chatId: messageData.chatId,
       createdBy: messageData.createdBy,
       createdAt: new Date(messageData.createdAt),
-      rawLlmMessages: messageData.rawLlmMessages as CoreMessage[],
+      rawLlmMessages: conversationHistory,
       userName: messageData.userName ?? messageData.userEmail ?? 'Unknown',
       organizationId: messageData.organizationId,
     };

apps/web/src/components/ui/charts/MetricChart/BusterMetricChart.tsx

@@ -163,7 +163,7 @@ export const BusterMetricChart: React.FC<BusterMetricChartProps> = React.memo(
           {...memoizedAnimation}>
           <AnimatedTitleWrapper title={formattedHeader} type="header" />
           <div className="w-full overflow-hidden p-2 text-center">
-            <div className="truncate">{formattedValue}</div>
+            <h2 className="text-foreground truncate text-4xl font-normal!">{formattedValue}</h2>
           </div>
           <AnimatedTitleWrapper title={formattedSubHeader} type="subHeader" />
         </motion.div>

packages/ai/package.json

@@ -45,7 +45,7 @@
     "@buster/test-utils": "workspace:*",
     "@buster/typescript-config": "workspace:*",
     "@buster/vitest-config": "workspace:*",
-    "@buster-tools/web-tools": "workspace:*",
+    "@buster/web-tools": "workspace:*",
     "@mastra/core": "catalog:",
     "@mastra/loggers": "^0.10.3",
     "ai": "catalog:",

packages/ai/src/steps/extract-values-search-step.test.ts

@@ -19,17 +19,16 @@ vi.mock('braintrust', () => ({
   wrapAISDKModel: vi.fn((model) => model),
 }));
 
-// Create a ref object to hold the mock generate function
-const mockGenerateRef = { current: vi.fn() };
+// Mock the AI SDK
+vi.mock('ai', () => ({
+  generateObject: vi.fn(),
+}));
 
-// Mock the Agent class from Mastra with the generate function
+// Mock Mastra
 vi.mock('@mastra/core', async () => {
   const actual = await vi.importActual('@mastra/core');
   return {
     ...actual,
-    Agent: vi.fn().mockImplementation(() => ({
-      generate: (...args: any[]) => mockGenerateRef.current(...args),
-    })),
     createStep: actual.createStep,
   };
 });
@@ -41,18 +40,17 @@ import { extractValuesSearchStep } from './extract-values-search-step';
 
 // Import the mocked functions
 import { generateEmbedding, searchValuesByEmbedding } from '@buster/stored-values/search';
+import { generateObject } from 'ai';
 
 const mockGenerateEmbedding = generateEmbedding as ReturnType<typeof vi.fn>;
 const mockSearchValuesByEmbedding = searchValuesByEmbedding as ReturnType<typeof vi.fn>;
-
-// Access the mock generate function through the ref
-const mockGenerate = mockGenerateRef.current;
+const mockGenerateObject = generateObject as ReturnType<typeof vi.fn>;
 
 describe('extractValuesSearchStep', () => {
   beforeEach(() => {
     vi.clearAllMocks();
     // Set default mock behavior
-    mockGenerate.mockResolvedValue({
+    mockGenerateObject.mockResolvedValue({
       object: { values: [] },
     });
   });
@@ -72,7 +70,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock the LLM response for keyword extraction
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['Red Bull', 'California'] },
       });
 
@@ -141,7 +139,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock empty keyword extraction
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: [] },
       });
 
@@ -195,7 +193,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock successful keyword extraction
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['Red Bull'] },
       });
 
@@ -226,7 +224,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock LLM extraction success but embedding failure
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['test keyword'] },
       });
 
@@ -254,7 +252,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock successful keyword extraction
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['test keyword'] },
       });
 
@@ -284,7 +282,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock two keywords: one succeeds, one fails
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['keyword1', 'keyword2'] },
       });
 
@@ -327,7 +325,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock everything to fail
-      mockGenerate.mockRejectedValue(new Error('LLM failure'));
+      mockGenerateObject.mockRejectedValue(new Error('LLM failure'));
       mockGenerateEmbedding.mockRejectedValue(new Error('Embedding failure'));
       mockSearchValuesByEmbedding.mockRejectedValue(new Error('Database failure'));
 
@@ -378,7 +376,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock successful keyword extraction
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['Red Bull'] },
       });
 
@@ -437,7 +435,7 @@ describe('extractValuesSearchStep', () => {
       runtimeContext.set('dataSourceId', 'test-datasource-id');
 
       // Mock successful keyword extraction
-      mockGenerate.mockResolvedValue({
+      mockGenerateObject.mockResolvedValue({
         object: { values: ['test'] },
       });
 

packages/ai/src/steps/extract-values-search-step.ts

@@ -1,7 +1,8 @@
 import type { StoredValueResult } from '@buster/stored-values';
 import { generateEmbedding, searchValuesByEmbedding } from '@buster/stored-values/search';
-import { Agent, createStep } from '@mastra/core';
+import { createStep } from '@mastra/core';
 import type { RuntimeContext } from '@mastra/core/runtime-context';
+import { generateObject } from 'ai';
 import type { CoreMessage } from 'ai';
 import { wrapTraced } from 'braintrust';
 import { z } from 'zod';
@@ -12,6 +13,11 @@ import type { AnalystRuntimeContext } from '../workflows/analyst-workflow';
 
 const inputSchema = thinkAndPrepWorkflowInputSchema;
 
+// Schema for what the LLM returns
+const llmOutputSchema = z.object({
+  values: z.array(z.string()).describe('The values that the agent will search for.'),
+});
+
 // Step output schema - what the step returns after performing the search
 export const extractValuesSearchOutputSchema = z.object({
   values: z.array(z.string()).describe('The values that the agent will search for.'),
@@ -231,12 +237,6 @@ async function searchStoredValues(
   }
 }
 
-const valuesAgent = new Agent({
-  name: 'Extract Values',
-  instructions: extractValuesInstructions,
-  model: Haiku35,
-});
-
 const extractValuesSearchStepExecution = async ({
   inputData,
   runtimeContext,
@@ -264,12 +264,19 @@ const extractValuesSearchStepExecution = async ({
     try {
       const tracedValuesExtraction = wrapTraced(
         async () => {
-          const response = await valuesAgent.generate(messages, {
-            maxSteps: 0,
-            output: extractValuesSearchOutputSchema,
+          const { object } = await generateObject({
+            model: Haiku35,
+            schema: llmOutputSchema,
+            messages: [
+              {
+                role: 'system',
+                content: extractValuesInstructions,
+              },
+              ...messages,
+            ],
           });
 
-          return response.object;
+          return object;
         },
         {
           name: 'Extract Values',

packages/ai/src/steps/generate-chat-title-step.ts

@@ -1,6 +1,7 @@
 import { updateChat, updateMessage } from '@buster/database';
-import { Agent, createStep } from '@mastra/core';
+import { createStep } from '@mastra/core';
 import type { RuntimeContext } from '@mastra/core/runtime-context';
+import { generateObject } from 'ai';
 import type { CoreMessage } from 'ai';
 import { wrapTraced } from 'braintrust';
 import { z } from 'zod';
@@ -11,6 +12,12 @@ import type { AnalystRuntimeContext } from '../workflows/analyst-workflow';
 
 const inputSchema = thinkAndPrepWorkflowInputSchema;
 
+// Schema for what the LLM returns
+const llmOutputSchema = z.object({
+  title: z.string().describe('The title for the chat.'),
+});
+
+// Schema for what the step returns (includes pass-through data)
 export const generateChatTitleOutputSchema = z.object({
   title: z.string().describe('The title for the chat.'),
   // Pass through dashboard context
@@ -28,13 +35,9 @@ export const generateChatTitleOutputSchema = z.object({
 
 const generateChatTitleInstructions = `
 I am a chat title generator that is responsible for generating a title for the chat.
-`;
 
-const todosAgent = new Agent({
-  name: 'Extract Values',
-  instructions: generateChatTitleInstructions,
-  model: Haiku35,
-});
+The title should be 3-8 words, capturing the main topic or intent of the conversation.
+`;
 
 const generateChatTitleExecution = async ({
   inputData,
@@ -63,12 +66,19 @@ const generateChatTitleExecution = async ({
     try {
       const tracedChatTitle = wrapTraced(
         async () => {
-          const response = await todosAgent.generate(messages, {
-            maxSteps: 0,
-            output: generateChatTitleOutputSchema,
+          const { object } = await generateObject({
+            model: Haiku35,
+            schema: llmOutputSchema,
+            messages: [
+              {
+                role: 'system',
+                content: generateChatTitleInstructions,
+              },
+              ...messages,
+            ],
           });
 
-          return response.object;
+          return object;
         },
         {
           name: 'Generate Chat Title',

packages/ai/src/tools/web-tools/web-search-tool.test.ts

@@ -1,7 +1,7 @@
 import { beforeEach, describe, expect, it, vi } from 'vitest';
 import { webSearch } from './web-search-tool';
 
-vi.mock('@buster-tools/web-tools', () => {
+vi.mock('@buster/web-tools', () => {
   const mockFirecrawlService = {
     webSearch: vi.fn(),
   };
@@ -17,9 +17,7 @@ describe('webSearch tool', () => {
 
   beforeEach(async () => {
     vi.clearAllMocks();
-    const { mockFirecrawlService: mock } = vi.mocked(
-      await import('@buster-tools/web-tools')
-    ) as any;
+    const { mockFirecrawlService: mock } = vi.mocked(await import('@buster/web-tools')) as any;
     mockFirecrawlService = mock;
   });
 

packages/ai/src/tools/web-tools/web-search-tool.ts

@@ -1,8 +1,4 @@
-import {
-  FirecrawlService,
-  type WebSearchOptions,
-  type WebSearchResult,
-} from '@buster-tools/web-tools';
+import { FirecrawlService, type WebSearchOptions, type WebSearchResult } from '@buster/web-tools';
 import type { RuntimeContext } from '@mastra/core/runtime-context';
 import { createTool } from '@mastra/core/tools';
 import { wrapTraced } from 'braintrust';

packages/ai/src/utils/retry/retry-agent-stream.ts

@@ -14,7 +14,7 @@ import type { RetryableError, WorkflowContext } from './types';
  * Creates a workflow-aware healing message for NoSuchToolError
  */
 function createWorkflowAwareHealingMessage(toolName: string, context?: WorkflowContext): string {
-  const baseMessage = `Tool "${toolName}" is not available in the current mode.`;
+  const baseMessage = `error: Tool "${toolName}" is not available.`;
 
   if (!context) {
     return `${baseMessage} Please use one of the available tools instead.`;
@@ -22,28 +22,44 @@ function createWorkflowAwareHealingMessage(toolName: string, context?: WorkflowC
 
   const { currentStep, availableTools } = context;
 
+  const nextMode = currentStep === 'think-and-prep' ? 'analyst' : 'think-and-prep';
+  const transitionDescription =
+    currentStep === 'think-and-prep'
+      ? 'after thinking, understanding the data, and submitting your thoughts'
+      : 'after completing your analysis';
+
   // Use actual available tools if provided
   if (availableTools && availableTools.size > 0) {
-    const toolList = Array.from(availableTools).sort().join(', ');
-    return `${baseMessage}
+    const currentToolList = Array.from(availableTools).sort().join(', ');
 
-You are currently in ${currentStep} mode. The available tools are: ${toolList}.
+    const nextModeTools =
+      nextMode === 'analyst'
+        ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool'
+        : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion';
 
-Please use one of these tools to continue with your task. Make sure to use the exact tool name as listed above.`;
+    return `${baseMessage} For reference you are currently in ${currentStep} mode which has access to the following tools:
+${currentToolList}
+
+The next mode that you'll transition to ${transitionDescription} will be the ${nextMode} mode which has access to the following tools:
+${nextModeTools}`;
   }
 
   // Fallback to static message if tools not provided
-  const pipelineContext = `
+  const currentModeTools =
+    currentStep === 'think-and-prep'
+      ? 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion'
+      : 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool';
 
-This workflow has two steps:
-1. think-and-prep mode: Available tools are sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion
-2. analyst mode: Available tools are createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool
+  const nextModeTools =
+    nextMode === 'analyst'
+      ? 'createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool'
+      : 'sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion';
 
-You are currently in ${currentStep} mode. Please use one of the tools available in your current mode.
+  return `${baseMessage} For reference you are currently in ${currentStep} mode which has access to the following tools:
+${currentModeTools}
 
-You should proceed with the proper tool calls in the context of the current step. There is a chance you might be a little confused about where you are in the workflow. or the tools available to you.`;
-
-  return baseMessage + pipelineContext;
+The next mode that you'll transition to ${transitionDescription} will be the ${nextMode} mode which has access to the following tools:
+${nextModeTools}`;
 }
 
 /**

packages/ai/src/utils/sql-permissions/sql-parser-helpers.test.ts

@@ -426,7 +426,7 @@ models:
       const sql = 'SELECT * FROM users';
       const result = validateWildcardUsage(sql);
       expect(result.isValid).toBe(false);
-      expect(result.error).toContain('Wildcard usage on physical tables is not allowed');
+      expect(result.error).toContain("You're not allowed to use a wildcard on physical tables");
       expect(result.blockedTables).toContain('users');
     });
 
@@ -434,7 +434,7 @@ models:
       const sql = 'SELECT u.* FROM users u';
       const result = validateWildcardUsage(sql);
       expect(result.isValid).toBe(false);
-      expect(result.error).toContain('Wildcard usage on physical tables is not allowed');
+      expect(result.error).toContain("You're not allowed to use a wildcard on physical tables");
       expect(result.blockedTables).toContain('u');
     });
 
@@ -470,7 +470,7 @@ models:
       `;
       const result = validateWildcardUsage(sql);
       expect(result.isValid).toBe(false);
-      expect(result.error).toContain('Wildcard usage on physical tables is not allowed');
+      expect(result.error).toContain("You're not allowed to use a wildcard on physical tables");
       expect(result.blockedTables).toContain('users');
     });
 
@@ -514,7 +514,7 @@ models:
       const sql = 'SELECT * FROM public.users';
       const result = validateWildcardUsage(sql);
       expect(result.isValid).toBe(false);
-      expect(result.error).toContain('Wildcard usage on physical tables is not allowed');
+      expect(result.error).toContain("You're not allowed to use a wildcard on physical tables");
     });
 
     it('should handle invalid SQL gracefully', () => {

packages/ai/src/utils/sql-permissions/sql-parser-helpers.ts

@@ -410,10 +410,9 @@ export function validateWildcardUsage(
     }
 
     if (blockedTables.length > 0) {
-      const tableList = blockedTables.join(', ');
       return {
         isValid: false,
-        error: `Wildcard usage on physical tables is not allowed: ${tableList}. Please specify explicit column names.`,
+        error: `You're not allowed to use a wildcard on physical tables, please be specific about which columns you'd like to work with`,
         blockedTables,
       };
     }

packages/database/src/queries/messages/chatConversationHistory.ts

@@ -35,12 +35,14 @@ async function getAllMessagesForChat(chatId: string): Promise<
     id: string;
     rawLlmMessages: unknown;
     createdAt: string;
+    isCompleted: boolean;
   }>
 > {
   let chatMessages: Array<{
     id: string;
     rawLlmMessages: unknown;
     createdAt: string;
+    isCompleted: boolean;
   }>;
   try {
     chatMessages = await db
@@ -48,6 +50,7 @@ async function getAllMessagesForChat(chatId: string): Promise<
         id: messages.id,
         rawLlmMessages: messages.rawLlmMessages,
         createdAt: messages.createdAt,
+        isCompleted: messages.isCompleted,
       })
       .from(messages)
       .where(and(eq(messages.chatId, chatId), isNull(messages.deletedAt)))
@@ -61,25 +64,44 @@ async function getAllMessagesForChat(chatId: string): Promise<
   return chatMessages;
 }
 
-// Helper function to combine raw LLM messages
-function combineRawLlmMessages(chatMessages: Array<{ rawLlmMessages: unknown }>): CoreMessage[] {
-  const conversationHistory: CoreMessage[] = [];
-
+// Helper function to get the most recent raw LLM messages
+function getMostRecentRawLlmMessages(
+  chatMessages: Array<{ rawLlmMessages: unknown; isCompleted: boolean }>
+): CoreMessage[] {
   try {
-    for (const message of chatMessages) {
-      if (message.rawLlmMessages && Array.isArray(message.rawLlmMessages)) {
-        // Preserve the exact message structure from the database
-        // Each rawLlmMessages array should contain properly unbundled messages
-        conversationHistory.push(...(message.rawLlmMessages as CoreMessage[]));
+    // Find the most recent completed message with valid rawLlmMessages
+    // We iterate backwards to find the most recent valid message
+    for (let i = chatMessages.length - 1; i >= 0; i--) {
+      const message = chatMessages[i];
+
+      // Skip if message is not completed
+      if (!message?.isCompleted) {
+        continue;
+      }
+
+      // Skip if rawLlmMessages is empty, null, or default empty object
+      const rawMessages = message.rawLlmMessages;
+      if (
+        !rawMessages ||
+        (typeof rawMessages === 'object' && Object.keys(rawMessages).length === 0) ||
+        rawMessages === '{}'
+      ) {
+        continue;
+      }
+
+      // Check if it's a valid array
+      if (Array.isArray(rawMessages) && rawMessages.length > 0) {
+        return rawMessages as CoreMessage[];
       }
     }
+
+    // No valid messages found
+    return [];
   } catch (processingError) {
     throw new Error(
       `Failed to process conversation history: ${processingError instanceof Error ? processingError.message : 'Unknown processing error'}`
     );
   }
-
-  return conversationHistory;
 }
 
 // Zod schemas for validation
@@ -94,7 +116,8 @@ export type ChatConversationHistoryOutput = z.infer<typeof ChatConversationHisto
 
 /**
  * Get complete conversation history for a chat from any message in that chat
- * Finds the chat from the given messageId, then loads ALL rawLlmMessages from ALL messages in that chat
+ * Finds the chat from the given messageId, then returns the most recent message's rawLlmMessages
+ * which contains the complete conversation history up to that point
  */
 export async function getChatConversationHistory(
   input: ChatConversationHistoryInput
@@ -109,8 +132,8 @@ export async function getChatConversationHistory(
     // Get all messages for this chat
     const chatMessages = await getAllMessagesForChat(chatId);
 
-    // Combine all rawLlmMessages into conversation history
-    const conversationHistory = combineRawLlmMessages(chatMessages);
+    // Get the most recent rawLlmMessages which contains the complete conversation history
+    const conversationHistory = getMostRecentRawLlmMessages(chatMessages);
 
     // Validate output
     try {

packages/web-tools/package.json

@@ -1,5 +1,5 @@
 {
-  "name": "@buster-tools/web-tools",
+  "name": "@buster/web-tools",
   "version": "0.1.0",
   "description": "Web scraping and research tools using Firecrawl and other services",
   "type": "module",

pnpm-lock.yaml

@@ -208,6 +208,9 @@ importers:
       '@buster/vitest-config':
         specifier: workspace:*
         version: link:../../packages/vitest-config
+      '@buster/web-tools':
+        specifier: workspace:*
+        version: link:../../packages/web-tools
       '@mastra/core':
         specifier: 'catalog:'
         version: 0.10.8(openapi-types@12.1.3)(react@18.3.1)(zod@3.25.1)
@@ -681,9 +684,6 @@ importers:
       '@ai-sdk/provider':
         specifier: ^1.1.3
         version: 1.1.3
-      '@buster-tools/web-tools':
-        specifier: workspace:*
-        version: link:../web-tools
       '@buster/access-controls':
         specifier: workspace:*
         version: link:../access-controls
@@ -714,6 +714,9 @@ importers:
       '@buster/vitest-config':
         specifier: workspace:*
         version: link:../vitest-config
+      '@buster/web-tools':
+        specifier: workspace:*
+        version: link:../web-tools
       '@mastra/core':
         specifier: 'catalog:'
         version: 0.10.8(openapi-types@12.1.3)(react@18.3.1)(zod@3.25.1)