Merge pull request #620 from buster-so/staging

Release

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:",

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
+// Mock the AI SDK
-const mockGenerateRef = { current: vi.fn() };
+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>;
-
+const mockGenerateObject = generateObject as ReturnType<typeof vi.fn>;
-// Access the mock generate function through the ref
-const mockGenerate = mockGenerateRef.current;
 
 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, {
+          const { object } = await generateObject({
-            maxSteps: 0,
+            model: Haiku35,
-            output: extractValuesSearchOutputSchema,
+            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({
+The title should be 3-8 words, capturing the main topic or intent of the conversation.
-  name: 'Extract Values',
+`;
-  instructions: generateChatTitleInstructions,
-  model: Haiku35,
-});
 
 const generateChatTitleExecution = async ({
   inputData,
@@ -63,12 +66,19 @@ const generateChatTitleExecution = async ({
     try {
       const tracedChatTitle = wrapTraced(
         async () => {
-          const response = await todosAgent.generate(messages, {
+          const { object } = await generateObject({
-            maxSteps: 0,
+            model: Haiku35,
-            output: generateChatTitleOutputSchema,
+            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(
+    const { mockFirecrawlService: mock } = vi.mocked(await import('@buster/web-tools')) as any;
-      await import('@buster-tools/web-tools')
-    ) as any;
     mockFirecrawlService = mock;
   });
 

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

@@ -1,8 +1,4 @@
-import {
+import { FirecrawlService, type WebSearchOptions, type WebSearchResult } from '@buster/web-tools';
-  FirecrawlService,
-  type WebSearchOptions,
-  type WebSearchResult,
-} from '@buster-tools/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(', ');
+    const currentToolList = Array.from(availableTools).sort().join(', ');
-    return `${baseMessage}
 
-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:
+  const nextModeTools =
-1. think-and-prep mode: Available tools are sequentialThinking, executeSql, respondWithoutAssetCreation, submitThoughts, messageUserClarifyingQuestion
+    nextMode === 'analyst'
-2. analyst mode: Available tools are createMetrics, modifyMetrics, createDashboards, modifyDashboards, doneTool
+      ? '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.`;
+The next mode that you'll transition to ${transitionDescription} will be the ${nextMode} mode which has access to the following tools:
-
+${nextModeTools}`;
-  return baseMessage + pipelineContext;
 }
 
 /**

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/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)