diff --git a/apps/server/src/api/v2/deploy/POST.test.ts b/apps/server/src/api/v2/deploy/POST.test.ts index 387d7df19..20af1031c 100644 --- a/apps/server/src/api/v2/deploy/POST.test.ts +++ b/apps/server/src/api/v2/deploy/POST.test.ts @@ -179,20 +179,17 @@ describe('deployHandler', () => { const result = await deployHandler(mockRequest, mockUser); - expect(mockUpsertDataset).toHaveBeenCalledWith( - {}, - { - name: 'test_model', - dataSourceId: 'ds-123', - organizationId: 'org-123', - database: 'test_db', - schema: 'public', - description: 'Test model', - sql_definition: 'SELECT * FROM test', - yml_file: 'model: test', - userId: 'user-123', - } - ); + expect(mockUpsertDataset).toHaveBeenCalledWith({ + name: 'test_model', + dataSourceId: 'ds-123', + organizationId: 'org-123', + database: 'test_db', + schema: 'public', + description: 'Test model', + sql_definition: 'SELECT * FROM test', + yml_file: 'model: test', + userId: 'user-123', + }); expect(result.models.success).toEqual([ { diff --git a/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts b/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts index d56ed0864..63a2e473c 100644 --- a/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts +++ b/packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts @@ -15,7 +15,7 @@ import type { export const ANALYST_ENGINEER_AGENT_NAME = 'analyticsEngineerAgent'; -const STOP_CONDITIONS = [stepCountIs(100), hasToolCall(IDLE_TOOL_NAME)]; +const STOP_CONDITIONS = [stepCountIs(250)]; export function createAnalyticsEngineerAgent( analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions diff --git a/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts b/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts index 5f33e39af..323eb84e7 100644 --- a/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts +++ b/packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts @@ -3,7 +3,6 @@ import { EDIT_FILE_TOOL_NAME, GLOB_TOOL_NAME, GREP_TOOL_NAME, - IDLE_TOOL_NAME, LS_TOOL_NAME, MULTI_EDIT_FILE_TOOL_NAME, READ_FILE_TOOL_NAME, @@ -15,7 +14,6 @@ import { createEditFileTool, createGlobTool, createGrepTool, - createIdleTool, createLsTool, createMultiEditFileTool, createReadFileTool, @@ -32,7 +30,6 @@ import type { AnalyticsEngineerAgentOptions } from './types'; export async function createAnalyticsEngineerToolset( analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions ) { - const idleTool = createIdleTool({}); const writeFileTool = createWriteFileTool({ messageId: analyticsEngineerAgentOptions.messageId, projectDirectory: analyticsEngineerAgentOptions.folder_structure, @@ -71,7 +68,6 @@ export async function createAnalyticsEngineerToolset( todosList: analyticsEngineerAgentOptions.todosList, }); const runSqlTool = createRunSqlTool({ - messageId: analyticsEngineerAgentOptions.messageId, apiKey: analyticsEngineerAgentOptions.apiKey || process.env.BUSTER_API_KEY || '', apiUrl: analyticsEngineerAgentOptions.apiUrl || process.env.BUSTER_API_URL || 'http://localhost:3000', @@ -99,7 +95,6 @@ export async function createAnalyticsEngineerToolset( // : null; return { - [IDLE_TOOL_NAME]: idleTool, [WRITE_FILE_TOOL_NAME]: writeFileTool, [GREP_TOOL_NAME]: grepTool, [GLOB_TOOL_NAME]: globTool, diff --git a/packages/ai/src/llm/providers/gateway.ts b/packages/ai/src/llm/providers/gateway.ts index 17c76adc9..4be7855f5 100644 --- a/packages/ai/src/llm/providers/gateway.ts +++ b/packages/ai/src/llm/providers/gateway.ts @@ -57,8 +57,8 @@ export const DEFAULT_ANTHROPIC_OPTIONS: AnthropicProviderOptions = { cacheControl: { type: 'ephemeral' }, thinking: { type: 'enabled', - budgetTokens: 10000 // Set desired tokens for reasoning - } + budgetTokens: 10000, // Set desired tokens for reasoning + }, }, bedrock: { cachePoint: { type: 'default' }, @@ -66,8 +66,8 @@ export const DEFAULT_ANTHROPIC_OPTIONS: AnthropicProviderOptions = { anthropic_beta: ['fine-grained-tool-streaming-2025-05-14'], reasoning_config: { type: 'enabled', - budget_tokens: 10000 // Adjust as needed - } + budget_tokens: 10000, // Adjust as needed + }, }, }, }; @@ -88,8 +88,8 @@ export const DEFAULT_ANALYTICS_ENGINEER_OPTIONS = { cacheControl: { type: 'ephemeral' }, thinking: { type: 'enabled', - budgetTokens: 10000 // Set desired tokens for reasoning - } + budgetTokens: 10000, // Set desired tokens for reasoning + }, }, bedrock: { cachePoint: { type: 'default' }, @@ -97,8 +97,8 @@ export const DEFAULT_ANALYTICS_ENGINEER_OPTIONS = { anthropic_beta: ['fine-grained-tool-streaming-2025-05-14'], reasoning_config: { type: 'enabled', - budget_tokens: 10000 // Adjust as needed - } + budget_tokens: 10000, // Adjust as needed + }, }, }, }; diff --git a/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.test.ts b/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.test.ts new file mode 100644 index 000000000..c1fe42c12 --- /dev/null +++ b/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.test.ts @@ -0,0 +1,161 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { RetrieveMetadataContext, RetrieveMetadataOutput } from './retrieve-metadata'; +import { createRetrieveMetadataExecute } from './retrieve-metadata-execute'; + +// Mock global fetch +global.fetch = vi.fn() as unknown as typeof fetch; + +describe('retrieve-metadata-execute error handling', () => { + const mockContext: RetrieveMetadataContext = { + apiKey: 'test-api-key', + apiUrl: 'http://localhost:3000', + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('createRetrieveMetadataExecute', () => { + it('should retrieve metadata successfully', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + const mockResponse: RetrieveMetadataOutput = { + metadata: { + columns: [ + { name: 'id', type: 'integer' }, + { name: 'name', type: 'string' }, + ], + row_count: 1000, + }, + }; + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: true, + json: async () => mockResponse, + } as Response); + + const result = await executeHandler({ + database: 'test_db', + schema: 'public', + name: 'users', + }); + + expect(result).toEqual(mockResponse); + expect(fetch).toHaveBeenCalledWith( + 'http://localhost:3000/api/v2/tools/metadata?database=test_db&schema=public&name=users', + expect.objectContaining({ + method: 'GET', + headers: { + Authorization: 'Bearer test-api-key', + }, + }) + ); + }); + + it('should throw error with clear message on API errors', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: false, + status: 404, + statusText: 'Not Found', + json: async () => ({ error: 'Table not found' }), + } as Response); + + await expect( + executeHandler({ + database: 'test_db', + schema: 'public', + name: 'nonexistent_table', + }) + ).rejects.toThrow('Table not found'); + }); + + it('should handle network errors gracefully', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + vi.mocked(fetch).mockRejectedValueOnce(new Error('Network connection failed')); + + await expect( + executeHandler({ + database: 'test_db', + schema: 'public', + name: 'users', + }) + ).rejects.toThrow('Network connection failed'); + }); + + it('should handle permission errors', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: async () => ({ error: 'Insufficient permissions' }), + } as Response); + + await expect( + executeHandler({ + database: 'test_db', + schema: 'restricted', + name: 'sensitive_table', + }) + ).rejects.toThrow('Insufficient permissions'); + }); + + it('should handle server errors', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: false, + status: 500, + statusText: 'Internal Server Error', + json: async () => ({ error: 'Database connection failed' }), + } as Response); + + await expect( + executeHandler({ + database: 'test_db', + schema: 'public', + name: 'users', + }) + ).rejects.toThrow('Database connection failed'); + }); + + it('should handle malformed error responses', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: false, + status: 400, + statusText: 'Bad Request', + json: async () => { + throw new Error('Invalid JSON'); + }, + } as unknown as Response); + + await expect( + executeHandler({ + database: 'test_db', + schema: 'public', + name: 'users', + }) + ).rejects.toThrow('HTTP 400: Bad Request'); + }); + + it('should handle generic metadata retrieval failures', async () => { + const executeHandler = createRetrieveMetadataExecute(mockContext); + + vi.mocked(fetch).mockRejectedValueOnce(new Error()); + + await expect( + executeHandler({ + database: 'test_db', + schema: 'public', + name: 'users', + }) + ).rejects.toThrow('Metadata retrieval failed'); + }); + }); +}); diff --git a/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.ts b/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.ts index 0dd7d8bd4..277ea277e 100644 --- a/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.ts +++ b/packages/ai/src/tools/database-tools/retrieve-metadata/retrieve-metadata-execute.ts @@ -75,7 +75,7 @@ export function createRetrieveMetadataExecute(context: RetrieveMetadataContext) return result.data; } - // Throw error if retrieval failed + // Throw error with clear message - API server handles logging throw new Error(result.error || 'Metadata retrieval failed'); }, { name: RETRIEVE_METADATA_TOOL_NAME } diff --git a/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.test.ts b/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.test.ts new file mode 100644 index 000000000..ba5963d6e --- /dev/null +++ b/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.test.ts @@ -0,0 +1,203 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { RunSqlContext, RunSqlOutput } from './run-sql'; +import { createRunSqlExecute } from './run-sql-execute'; + +// Mock global fetch +global.fetch = vi.fn() as unknown as typeof fetch; + +describe('run-sql-execute error handling', () => { + const mockContext: RunSqlContext = { + apiKey: 'test-api-key', + apiUrl: 'http://localhost:3000', + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe('createRunSqlExecute', () => { + it('should execute SQL query successfully', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + const mockResponse: RunSqlOutput = { + data: [ + { id: 1, name: 'Test 1' }, + { id: 2, name: 'Test 2' }, + ], + data_metadata: { + id: { type: 'integer' }, + name: { type: 'string' }, + }, + has_more_records: false, + }; + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: true, + json: async () => mockResponse, + } as Response); + + const result = await executeHandler({ + data_source_id: 'test-ds-id', + sql: 'SELECT * FROM users LIMIT 2', + }); + + expect(result).toEqual(mockResponse); + expect(fetch).toHaveBeenCalledWith( + 'http://localhost:3000/api/v2/tools/sql', + expect.objectContaining({ + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: 'Bearer test-api-key', + }, + body: JSON.stringify({ + data_source_id: 'test-ds-id', + sql: 'SELECT * FROM users LIMIT 2', + }), + }) + ); + }); + + it('should throw error with clear message on API errors', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: false, + status: 400, + statusText: 'Bad Request', + json: async () => ({ error: 'Invalid SQL query' }), + } as Response); + + await expect( + executeHandler({ + data_source_id: 'test-ds-id', + sql: 'INVALID SQL', + }) + ).rejects.toThrow('Invalid SQL query'); + }); + + it('should retry on timeout errors', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + const mockResponse: RunSqlOutput = { + data: [{ result: 'success' }], + data_metadata: {}, + has_more_records: false, + }; + + // First call times out, second succeeds + vi.mocked(fetch) + .mockResolvedValueOnce({ + ok: false, + status: 500, + statusText: 'Internal Server Error', + json: async () => ({ error: 'Query timed out' }), + } as Response) + .mockResolvedValueOnce({ + ok: true, + json: async () => mockResponse, + } as Response); + + const result = await executeHandler({ + data_source_id: 'test-ds-id', + sql: 'SELECT * FROM large_table', + }); + + expect(result).toEqual(mockResponse); + expect(fetch).toHaveBeenCalledTimes(2); + }); + + it('should fail after max retries', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + // Mock all retries to timeout + vi.mocked(fetch).mockResolvedValue({ + ok: false, + status: 500, + statusText: 'Internal Server Error', + json: async () => ({ error: 'Query timeout' }), + } as Response); + + await expect( + executeHandler({ + data_source_id: 'test-ds-id', + sql: 'SELECT * FROM large_table', + }) + ).rejects.toThrow(); + + // Should have tried 4 times (1 initial + 3 retries) + expect(fetch).toHaveBeenCalledTimes(4); + }); + + it('should handle network errors gracefully', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + vi.mocked(fetch).mockRejectedValueOnce(new Error('Network error')); + + await expect( + executeHandler({ + data_source_id: 'test-ds-id', + sql: 'SELECT 1', + }) + ).rejects.toThrow('Network error'); + }); + + it('should throw error for empty SQL query', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + await expect( + executeHandler({ + data_source_id: 'test-ds-id', + sql: '', + }) + ).rejects.toThrow('SQL query cannot be empty'); + }); + + it('should handle non-timeout errors without retrying', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + vi.mocked(fetch).mockResolvedValueOnce({ + ok: false, + status: 403, + statusText: 'Forbidden', + json: async () => ({ error: 'Permission denied' }), + } as Response); + + await expect( + executeHandler({ + data_source_id: 'test-ds-id', + sql: 'SELECT * FROM secure_table', + }) + ).rejects.toThrow('Permission denied'); + + // Should not retry for non-timeout errors + expect(fetch).toHaveBeenCalledTimes(1); + }); + + it('should handle fetch failures with timeout detection', async () => { + const executeHandler = createRunSqlExecute(mockContext); + + const mockResponse: RunSqlOutput = { + data: [{ id: 1 }], + data_metadata: {}, + has_more_records: false, + }; + + // First call throws fetch failed, second succeeds + vi.mocked(fetch) + .mockRejectedValueOnce(new Error('fetch failed')) + .mockResolvedValueOnce({ + ok: true, + json: async () => mockResponse, + } as Response); + + const result = await executeHandler({ + data_source_id: 'test-ds-id', + sql: 'SELECT 1', + }); + + expect(result).toEqual(mockResponse); + expect(fetch).toHaveBeenCalledTimes(2); + }); + }); +}); diff --git a/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.ts b/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.ts index b81fab6c4..460fd7afd 100644 --- a/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.ts +++ b/packages/ai/src/tools/database-tools/run-sql/run-sql-execute.ts @@ -130,7 +130,7 @@ export function createRunSqlExecute(context: RunSqlContext) { return result.data; } - // Throw error if execution failed + // Throw error with clear message - API server handles logging throw new Error(result.error || 'Query execution failed'); }, { name: RUN_SQL_TOOL_NAME } diff --git a/packages/ai/src/tools/database-tools/run-sql/run-sql.ts b/packages/ai/src/tools/database-tools/run-sql/run-sql.ts index c99c9b83e..b45ecadfd 100644 --- a/packages/ai/src/tools/database-tools/run-sql/run-sql.ts +++ b/packages/ai/src/tools/database-tools/run-sql/run-sql.ts @@ -21,7 +21,6 @@ export const RunSqlInputSchema = z.object({ const RunSqlContextSchema = z.object({ apiKey: z.string().describe('API key for authentication'), apiUrl: z.string().describe('Base URL of the API server'), - messageId: z.string().describe('Message ID for database updates'), }); export type RunSqlInput = z.infer; diff --git a/packages/ai/src/tools/file-tools/bash-tool/bash-tool-execute.ts b/packages/ai/src/tools/file-tools/bash-tool/bash-tool-execute.ts index edae4ad18..3ef79d6b1 100644 --- a/packages/ai/src/tools/file-tools/bash-tool/bash-tool-execute.ts +++ b/packages/ai/src/tools/file-tools/bash-tool/bash-tool-execute.ts @@ -46,7 +46,7 @@ function validateDbtCommand(command: string): { isValid: boolean; error?: string // Handle cases like: "dbt run", "dbt run --select model", "cd path && dbt run" const dbtMatch = command.match(/\bdbt\s+([a-z-]+)/); - if (!dbtMatch) { + if (!dbtMatch || !dbtMatch[1]) { // Not a dbt command, allow it return { isValid: true }; }