diff --git a/packages/ai/src/steps/docs-agent-steps/get-repository-tree-step/get-repository-tree-step.test.ts b/packages/ai/src/steps/docs-agent-steps/get-repository-tree-step/get-repository-tree-step.test.ts index 5523fb39e..09f1fc5a5 100644 --- a/packages/ai/src/steps/docs-agent-steps/get-repository-tree-step/get-repository-tree-step.test.ts +++ b/packages/ai/src/steps/docs-agent-steps/get-repository-tree-step/get-repository-tree-step.test.ts @@ -70,12 +70,18 @@ describe('runGetRepositoryTreeStep', () => { }); it('should handle missing sandbox gracefully', async () => { + // Create a minimal sandbox object that passes validation but doesn't have full functionality + const minimalSandbox = { + id: 'mock-sandbox-minimal', + fs: {}, + } as unknown as Sandbox; + const input = { message: 'Test message', organizationId: 'org-123', contextInitialized: true, context: { - sandbox: null as any, + sandbox: minimalSandbox, dataSourceId: validDataSourceId, todoList: '', clarificationQuestions: [], diff --git a/packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-streaming.test.ts b/packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-streaming.test.ts index 1bb67397a..ac786d7c7 100644 --- a/packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-streaming.test.ts +++ b/packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-streaming.test.ts @@ -11,6 +11,7 @@ import type { vi.mock('@buster/database', () => ({ updateMessageEntries: vi.fn().mockResolvedValue({ success: true }), + updateMessage: vi.fn().mockResolvedValue({ success: true }), })); describe('Respond Without Asset Creation Tool Streaming Tests', () => { diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.test.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.test.ts index 541b58128..919bee089 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.test.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.test.ts @@ -235,7 +235,7 @@ describe('create-reports-execute', () => { const firstCall = mockUpdateMessageEntries.mock.calls[0]?.[0]; expect(firstCall.messageId).toBe('msg-001'); expect(firstCall.reasoningMessages).toBeDefined(); - expect(firstCall.rawLlmMessages).toBeDefined(); + // rawLlmMessages are intentionally not created in initial entries to avoid duplicates // State should be updated expect(state.initialEntriesCreated).toBe(true); diff --git a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts index c6b5947ac..0853c4e61 100644 --- a/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts +++ b/packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts @@ -95,7 +95,7 @@ export function createCreateReportsExecute( try { const toolCallId = state.toolCallId || `tool-${Date.now()}`; const reasoningEntry = createCreateReportsReasoningEntry(state, toolCallId); - const rawLlmMessage = createCreateReportsRawLlmMessageEntry(state, toolCallId); + // Skip creating rawLlmMessage here to avoid duplicates - it will be created with the result later const updates: Parameters[0] = { messageId: context.messageId, @@ -105,11 +105,7 @@ export function createCreateReportsExecute( updates.reasoningMessages = [reasoningEntry]; } - if (rawLlmMessage) { - updates.rawLlmMessages = [rawLlmMessage]; - } - - if (reasoningEntry || rawLlmMessage) { + if (reasoningEntry) { await updateMessageEntries(updates); state.initialEntriesCreated = true; } diff --git a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts index e1aa2bfca..f28c1b59a 100644 --- a/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts +++ b/packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts @@ -1,9 +1,6 @@ import { updateMessageEntries } from '@buster/database'; import type { ToolCallOptions } from 'ai'; -import { - createModifyReportsRawLlmMessageEntry, - createModifyReportsReasoningEntry, -} from './helpers/modify-reports-transform-helper'; +import { createModifyReportsReasoningEntry } from './helpers/modify-reports-transform-helper'; import type { ModifyReportsContext, ModifyReportsState } from './modify-reports-tool'; export function modifyReportsStart(context: ModifyReportsContext, state: ModifyReportsState) { @@ -24,12 +21,12 @@ export function modifyReportsStart(context: ModifyReportsContext, state: ModifyR if (context.messageId) { try { if (context.messageId && state.toolCallId) { - // Update database with both reasoning and raw LLM entries + // Update database with reasoning entry only - raw LLM message will be created with result later try { const reasoningEntry = createModifyReportsReasoningEntry(state, options.toolCallId); - const rawLlmMessage = createModifyReportsRawLlmMessageEntry(state, options.toolCallId); + // Skip creating rawLlmMessage here to avoid duplicates - it will be created with the result later - // Update both entries together if they exist + // Update reasoning entry if it exists const updates: Parameters[0] = { messageId: context.messageId, }; @@ -38,11 +35,7 @@ export function modifyReportsStart(context: ModifyReportsContext, state: ModifyR updates.reasoningMessages = [reasoningEntry]; } - if (rawLlmMessage) { - updates.rawLlmMessages = [rawLlmMessage]; - } - - if (reasoningEntry || rawLlmMessage) { + if (reasoningEntry) { await updateMessageEntries(updates); } } catch (error) { diff --git a/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts b/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts index ad24b9037..088e08d4a 100644 --- a/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts +++ b/packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts @@ -1,18 +1,38 @@ -import { NoSuchToolError, generateText } from 'ai'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { z } from 'zod'; + +// Mock the dependencies - must be before imports +vi.mock('ai', () => { + class MockNoSuchToolError extends Error { + toolName: string; + availableTools: string[]; + + constructor(options: { toolName: string; availableTools: string[] }) { + super(`Tool ${options.toolName} not found`); + this.toolName = options.toolName; + this.availableTools = options.availableTools; + } + + static isInstance(error: any): boolean { + return error instanceof MockNoSuchToolError; + } + } + + return { + NoSuchToolError: MockNoSuchToolError, + generateText: vi.fn(), + streamText: vi.fn(), + tool: vi.fn((config: any) => config), + stepCountIs: vi.fn((count: number) => ({ type: 'stepCount', count })), + hasToolCall: vi.fn((toolName: string) => ({ type: 'toolCall', toolName })), + }; +}); + +import { NoSuchToolError, generateText, streamText } from 'ai'; import { ANALYST_AGENT_NAME, THINK_AND_PREP_AGENT_NAME } from '../../../agents'; import type { RepairContext } from '../types'; import { canHandleNoSuchTool, repairWrongToolName } from './re-ask-strategy'; -// Mock the dependencies -vi.mock('ai', async () => { - const actual = await vi.importActual('ai'); - return { - ...actual, - generateText: vi.fn(), - }; -}); - vi.mock('braintrust', () => ({ wrapTraced: (fn: any) => fn, })); @@ -47,17 +67,18 @@ describe('re-ask-strategy', () => { describe('repairWrongToolName', () => { it('should re-ask and get corrected tool call', async () => { - const mockGenerateText = vi.mocked(generateText); + const mockStreamText = vi.mocked(streamText); const correctedToolCall = { toolName: 'correctTool', input: { param: 'value' }, }; - mockGenerateText.mockResolvedValueOnce({ - toolCalls: [correctedToolCall], - text: '', - usage: {}, + mockStreamText.mockReturnValueOnce({ + textStream: (async function* () {})(), + toolCalls: Promise.resolve([correctedToolCall]), + text: Promise.resolve(''), + usage: Promise.resolve({}), } as any); const context: RepairContext = { @@ -68,8 +89,8 @@ describe('re-ask-strategy', () => { input: { param: 'value' }, } as any, tools: { - correctTool: { inputSchema: {} }, - anotherTool: { inputSchema: {} }, + correctTool: { inputSchema: z.object({}) }, + anotherTool: { inputSchema: z.object({}) }, } as any, error: new NoSuchToolError({ toolName: 'wrongTool', @@ -90,7 +111,7 @@ describe('re-ask-strategy', () => { }); // Verify the tool input is properly formatted as an object in the messages - const calls = mockGenerateText.mock.calls[0]; + const calls = mockStreamText.mock.calls[0]; const messages = calls?.[0]?.messages; const assistantMessage = messages?.find((m: any) => m.role === 'assistant'); const content = assistantMessage?.content?.[0]; @@ -98,7 +119,7 @@ describe('re-ask-strategy', () => { expect(content.input).toEqual({ param: 'value' }); } - expect(mockGenerateText).toHaveBeenCalledWith( + expect(mockStreamText).toHaveBeenCalledWith( expect.objectContaining({ model: 'mock-model', messages: expect.arrayContaining([ @@ -114,11 +135,12 @@ describe('re-ask-strategy', () => { }); it('should use analyst agent context for error message', async () => { - const mockGenerateText = vi.mocked(generateText); - mockGenerateText.mockResolvedValueOnce({ - toolCalls: [], - text: '', - usage: {}, + const mockStreamText = vi.mocked(streamText); + mockStreamText.mockReturnValueOnce({ + textStream: (async function* () {})(), + toolCalls: Promise.resolve([]), + text: Promise.resolve(''), + usage: Promise.resolve({}), } as any); const context: RepairContext = { @@ -129,8 +151,8 @@ describe('re-ask-strategy', () => { input: '{}', } as any, tools: { - createMetrics: {}, - modifyMetrics: {}, + createMetrics: { inputSchema: z.object({}) }, + modifyMetrics: { inputSchema: z.object({}) }, } as any, error: new NoSuchToolError({ toolName: 'executeSql', @@ -146,8 +168,8 @@ describe('re-ask-strategy', () => { await repairWrongToolName(context); - const calls = mockGenerateText.mock.calls[0]; - if (!calls) throw new Error('generateText not called'); + const calls = mockStreamText.mock.calls[0]; + if (!calls) throw new Error('streamText not called'); const messages = calls[0]?.messages; if (!messages) throw new Error('No messages found'); const toolResultMessage = messages.find((m: any) => m.role === 'tool'); @@ -164,11 +186,12 @@ describe('re-ask-strategy', () => { }); it('should use think-and-prep agent context for error message', async () => { - const mockGenerateText = vi.mocked(generateText); - mockGenerateText.mockResolvedValueOnce({ - toolCalls: [], - text: '', - usage: {}, + const mockStreamText = vi.mocked(streamText); + mockStreamText.mockReturnValueOnce({ + textStream: (async function* () {})(), + toolCalls: Promise.resolve([]), + text: Promise.resolve(''), + usage: Promise.resolve({}), } as any); const context: RepairContext = { @@ -179,8 +202,8 @@ describe('re-ask-strategy', () => { input: '{}', } as any, tools: { - executeSql: {}, - sequentialThinking: {}, + executeSql: { inputSchema: z.object({}) }, + sequentialThinking: { inputSchema: z.object({}) }, } as any, error: new NoSuchToolError({ toolName: 'createMetrics', @@ -197,8 +220,8 @@ describe('re-ask-strategy', () => { await repairWrongToolName(context); - const calls = mockGenerateText.mock.calls[0]; - if (!calls) throw new Error('generateText not called'); + const calls = mockStreamText.mock.calls[0]; + if (!calls) throw new Error('streamText not called'); const messages = calls[0]?.messages; if (!messages) throw new Error('No messages found'); const toolResultMessage = messages.find((m: any) => m.role === 'tool'); @@ -213,11 +236,12 @@ describe('re-ask-strategy', () => { }); it('should return null if no valid tool call is returned', async () => { - const mockGenerateText = vi.mocked(generateText); - mockGenerateText.mockResolvedValueOnce({ - toolCalls: [], - text: '', - usage: {}, + const mockStreamText = vi.mocked(streamText); + mockStreamText.mockReturnValueOnce({ + textStream: (async function* () {})(), + toolCalls: Promise.resolve([]), + text: Promise.resolve(''), + usage: Promise.resolve({}), } as any); const context: RepairContext = { @@ -228,7 +252,7 @@ describe('re-ask-strategy', () => { input: '{}', } as any, tools: { - correctTool: {}, + correctTool: { inputSchema: z.object({}) }, } as any, error: new NoSuchToolError({ toolName: 'wrongTool', @@ -243,8 +267,10 @@ describe('re-ask-strategy', () => { }); it('should handle errors during re-ask', async () => { - const mockGenerateText = vi.mocked(generateText); - mockGenerateText.mockRejectedValueOnce(new Error('Generation failed')); + const mockStreamText = vi.mocked(streamText); + mockStreamText.mockImplementationOnce(() => { + throw new Error('Generation failed'); + }); const context: RepairContext = { toolCall: { @@ -267,11 +293,12 @@ describe('re-ask-strategy', () => { }); it('should wrap non-JSON string inputs in an object', async () => { - const mockGenerateText = vi.mocked(generateText); - mockGenerateText.mockResolvedValueOnce({ - toolCalls: [{ toolName: 'correctTool', input: { wrapped: true } }], - text: '', - usage: {}, + const mockStreamText = vi.mocked(streamText); + mockStreamText.mockReturnValueOnce({ + textStream: (async function* () {})(), + toolCalls: Promise.resolve([{ toolName: 'correctTool', input: { wrapped: true } }]), + text: Promise.resolve(''), + usage: Promise.resolve({}), } as any); const context: RepairContext = { @@ -282,7 +309,7 @@ describe('re-ask-strategy', () => { input: 'plain text input', } as any, tools: { - correctTool: { inputSchema: {} }, + correctTool: { inputSchema: z.object({}) }, } as any, error: new NoSuchToolError({ toolName: 'wrongTool', @@ -295,7 +322,7 @@ describe('re-ask-strategy', () => { await repairWrongToolName(context); // Verify the non-JSON string was wrapped in an object - const calls = mockGenerateText.mock.calls[0]; + const calls = mockStreamText.mock.calls[0]; const messages = calls?.[0]?.messages; const assistantMessage = messages?.find((m: any) => m.role === 'assistant'); const content = assistantMessage?.content?.[0]; @@ -305,11 +332,12 @@ describe('re-ask-strategy', () => { }); it('should handle already valid JSON string inputs', async () => { - const mockGenerateText = vi.mocked(generateText); - mockGenerateText.mockResolvedValueOnce({ - toolCalls: [{ toolName: 'correctTool', input: { handled: true } }], - text: '', - usage: {}, + const mockStreamText = vi.mocked(streamText); + mockStreamText.mockReturnValueOnce({ + textStream: (async function* () {})(), + toolCalls: Promise.resolve([{ toolName: 'correctTool', input: { handled: true } }]), + text: Promise.resolve(''), + usage: Promise.resolve({}), } as any); const context: RepairContext = { @@ -320,7 +348,7 @@ describe('re-ask-strategy', () => { input: '{"already":"valid"}', } as any, tools: { - correctTool: { inputSchema: {} }, + correctTool: { inputSchema: z.object({}) }, } as any, error: new NoSuchToolError({ toolName: 'wrongTool', @@ -333,7 +361,7 @@ describe('re-ask-strategy', () => { await repairWrongToolName(context); // Verify the valid JSON string was parsed to an object - const calls = mockGenerateText.mock.calls[0]; + const calls = mockStreamText.mock.calls[0]; const messages = calls?.[0]?.messages; const assistantMessage = messages?.find((m: any) => m.role === 'assistant'); const content = assistantMessage?.content?.[0];