build and tests

This commit is contained in:
dal 2025-10-02 09:39:55 -06:00
parent 4338cfd5d3
commit c8dfd4ccb3
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
6 changed files with 146 additions and 86 deletions

View File

@ -180,11 +180,15 @@ export function createDoneToolDelta(context: DoneToolContext, doneToolState: Don
}
}
// Store ALL assets (including reports) for chat update later
if (newAssets.length > 0) {
// Store assets for chat update later (excluding reasoning which is not an asset type)
const assetsForChatUpdate = newAssets.filter(
(a): a is typeof a & { assetType: Exclude<typeof a.assetType, 'reasoning'> } =>
a.assetType !== 'reasoning'
);
if (assetsForChatUpdate.length > 0) {
doneToolState.addedAssets = [
...(doneToolState.addedAssets || []),
...newAssets.map((a) => ({
...assetsForChatUpdate.map((a) => ({
assetId: a.assetId,
assetType: a.assetType,
versionNumber: a.versionNumber,

View File

@ -63,8 +63,10 @@ async function processDone(
return {
output,
sequenceNumber: updateResult.sequenceNumber,
skipped: updateResult.skipped,
...(updateResult.sequenceNumber !== undefined && {
sequenceNumber: updateResult.sequenceNumber,
}),
...(updateResult.skipped !== undefined && { skipped: updateResult.skipped }),
};
} catch (error) {
console.error('[done-tool] Error updating message entries:', error);

View File

@ -1,12 +1,4 @@
import { randomUUID } from 'node:crypto';
import {
getAssetLatestVersion,
isMessageUpdateQueueClosed,
updateChat,
updateMessage,
updateMessageEntries,
waitForPendingUpdates,
} from '@buster/database/queries';
import type { ModelMessage, ToolCallOptions } from 'ai';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { CREATE_DASHBOARDS_TOOL_NAME } from '../../visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool';
@ -17,6 +9,7 @@ import type { DoneToolContext, DoneToolState } from './done-tool';
import { createDoneToolDelta } from './done-tool-delta';
import { createDoneToolStart } from './done-tool-start';
// Mock database queries
vi.mock('@buster/database/queries', () => ({
updateChat: vi.fn(),
updateMessage: vi.fn(),
@ -30,6 +23,21 @@ vi.mock('@buster/database/queries', () => ({
getAssetLatestVersion: vi.fn().mockResolvedValue(1),
}));
// Import mocked functions after the mock definition
import {
getAssetLatestVersion,
isMessageUpdateQueueClosed,
updateChat,
updateMessage,
updateMessageEntries,
waitForPendingUpdates,
} from '@buster/database/queries';
// Type assertion for mocked functions
const mockedIsMessageUpdateQueueClosed = vi.mocked(isMessageUpdateQueueClosed);
const mockedWaitForPendingUpdates = vi.mocked(waitForPendingUpdates);
const mockedGetAssetLatestVersion = vi.mocked(getAssetLatestVersion);
describe('done-tool-start', () => {
const mockContext: DoneToolContext = {
chatId: 'chat-123',
@ -43,11 +51,20 @@ describe('done-tool-start', () => {
finalResponse: undefined,
};
// Helper to create mock ToolCallOptions
const createMockToolCallOptions = (
overrides: Partial<ToolCallOptions> = {}
): ToolCallOptions => ({
messages: [],
toolCallId: 'done-call',
...overrides,
});
beforeEach(() => {
vi.clearAllMocks();
isMessageUpdateQueueClosed.mockReturnValue(false);
waitForPendingUpdates.mockResolvedValue(undefined);
getAssetLatestVersion.mockResolvedValue(1);
mockedIsMessageUpdateQueueClosed.mockReturnValue(false);
mockedWaitForPendingUpdates.mockResolvedValue(undefined);
mockedGetAssetLatestVersion.mockResolvedValue(1);
});
describe('mostRecentFile selection', () => {
@ -142,10 +159,12 @@ describe('done-tool-start', () => {
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
// Start phase - initializes state
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - streams in the assets and final response
const deltaInput = JSON.stringify({
@ -161,8 +180,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
expect(updateChat).toHaveBeenCalledWith('chat-123', {
mostRecentFileId: reportId,
@ -248,10 +267,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the first metric as the asset to return
const deltaInput = JSON.stringify({
@ -267,8 +288,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Should select the first metric (first in extractedFiles)
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -359,10 +380,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the report and standalone metric as assets to return
const deltaInput = JSON.stringify({
@ -383,8 +406,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Report should be selected as mostRecentFile
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -459,10 +482,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the first metric
const deltaInput = JSON.stringify({
@ -478,8 +503,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Should select the first metric
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -526,10 +551,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the first dashboard
const deltaInput = JSON.stringify({
@ -545,8 +572,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Should select the first dashboard
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -612,10 +639,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the dashboard
const deltaInput = JSON.stringify({
@ -631,8 +660,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Should select the dashboard (first in extractedFiles)
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -644,10 +673,12 @@ describe('done-tool-start', () => {
it('should handle empty file lists gracefully', async () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: [],
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: [],
})
);
// Should not call updateChat when no files exist
expect(updateChat).not.toHaveBeenCalled();
@ -697,10 +728,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the report
const deltaInput = JSON.stringify({
@ -716,8 +749,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Report should still be selected as mostRecentFile
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -788,10 +821,12 @@ describe('done-tool-start', () => {
const doneToolStart = createDoneToolStart(mockContext, mockDoneToolState);
const doneToolDelta = createDoneToolDelta(mockContext, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Delta phase - stream in the dashboard (metrics are embedded)
const deltaInput = JSON.stringify({
@ -807,8 +842,8 @@ describe('done-tool-start', () => {
await doneToolDelta({
inputTextDelta: deltaInput,
toolCallId: 'done-call',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'done-call' }),
});
// Should select the dashboard since that's what we're returning
expect(updateChat).toHaveBeenCalledWith('chat-123', {
@ -852,10 +887,12 @@ describe('done-tool-start', () => {
];
const doneToolStart = createDoneToolStart(contextWithEmptyChatId, mockDoneToolState);
await doneToolStart({
toolCallId: 'done-call',
messages: mockMessages,
} as ToolCallOptions);
await doneToolStart(
createMockToolCallOptions({
toolCallId: 'done-call',
messages: mockMessages,
})
);
// Should not call updateChat when chatId is missing
expect(updateChat).not.toHaveBeenCalled();

View File

@ -71,6 +71,15 @@ describe('Done Tool Streaming Tests', () => {
workflowStartTime: Date.now(),
};
// Helper to create mock ToolCallOptions
const createMockToolCallOptions = (
overrides: Partial<ToolCallOptions> = {}
): ToolCallOptions => ({
messages: [],
toolCallId: 'test-call-id',
...overrides,
});
describe('createDoneToolStart', () => {
test('should initialize state with entry_id on start', async () => {
const state: DoneToolState = {
@ -252,8 +261,8 @@ describe('Done Tool Streaming Tests', () => {
});
await deltaHandler({
inputTextDelta: deltaInput,
toolCallId: 'call-1',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'call-1' }),
});
const queries = await import('@buster/database/queries');
@ -375,8 +384,8 @@ describe('Done Tool Streaming Tests', () => {
});
await deltaHandler({
inputTextDelta: deltaInput,
toolCallId: 'call-2',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'call-2' }),
});
const queries = await import('@buster/database/queries');
@ -492,8 +501,8 @@ describe('Done Tool Streaming Tests', () => {
});
await deltaHandler({
inputTextDelta: deltaInput,
toolCallId: 'call-3',
} as ToolCallOptions);
...createMockToolCallOptions({ toolCallId: 'call-3' }),
});
const queries = await import('@buster/database/queries');
const updateArgs = ((queries.updateChat as unknown as { mock: { calls: unknown[][] } }).mock
@ -688,6 +697,7 @@ describe('Done Tool Streaming Tests', () => {
const finishHandler = createDoneToolFinish(mockContext, state);
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: 'This is the final response message',
};
@ -710,6 +720,7 @@ describe('Done Tool Streaming Tests', () => {
const finishHandler = createDoneToolFinish(mockContext, state);
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: 'Response without prior start',
};
@ -746,6 +757,7 @@ The following items were processed:
`;
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: markdownResponse,
};
@ -810,6 +822,7 @@ The following items were processed:
expect(state.finalResponse).toBeTypeOf('string');
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: 'Final test',
};
await finishHandler({ input, toolCallId: 'test-123', messages: [] });
@ -859,6 +872,7 @@ The following items were processed:
);
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: 'This is a streaming response that comes in multiple chunks',
};
await finishHandler({ input, toolCallId, messages: [] });

View File

@ -108,6 +108,7 @@ describe('Done Tool Integration Tests', () => {
await startHandler({ toolCallId, messages: [] });
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: 'This is the complete final response',
};
@ -168,6 +169,7 @@ All operations completed successfully.`;
expect(state.finalResponse).toBe(expectedResponse);
const input: DoneToolInput = {
assetsToReturn: [],
finalResponse: expectedResponse,
};
@ -207,14 +209,14 @@ All operations completed successfully.`;
await startHandler1({ toolCallId: toolCallId1, messages: [] });
await finishHandler1({
input: { finalResponse: 'First response' },
input: { assetsToReturn: [], finalResponse: 'First response' },
toolCallId: toolCallId1,
messages: [],
});
await startHandler2({ toolCallId: toolCallId2, messages: [] });
await finishHandler2({
input: { finalResponse: 'Second response' },
input: { assetsToReturn: [], finalResponse: 'Second response' },
toolCallId: toolCallId2,
messages: [],
});

View File

@ -1,3 +1,4 @@
import { randomUUID } from 'node:crypto';
import type { ModelMessage } from 'ai';
import { beforeEach, describe, expect, it, vi } from 'vitest';
import { getChatConversationHistory } from './chatConversationHistory';
@ -79,7 +80,7 @@ describe('getChatConversationHistory - Orphaned Tool Call Cleanup', () => {
]);
const result = await getChatConversationHistory({
messageId: 'test-message-id',
messageId: randomUUID(),
});
// Should have removed the orphaned tool call but kept the valid one
@ -144,7 +145,7 @@ describe('getChatConversationHistory - Orphaned Tool Call Cleanup', () => {
]);
const result = await getChatConversationHistory({
messageId: 'test-message-id',
messageId: randomUUID(),
});
// Should keep both messages (tool call and result)
@ -206,7 +207,7 @@ describe('getChatConversationHistory - Orphaned Tool Call Cleanup', () => {
]);
const result = await getChatConversationHistory({
messageId: 'test-message-id',
messageId: randomUUID(),
});
// Should keep the assistant message because it has at least one valid tool call
@ -265,7 +266,7 @@ describe('getChatConversationHistory - Orphaned Tool Call Cleanup', () => {
]);
const result = await getChatConversationHistory({
messageId: 'test-message-id',
messageId: randomUUID(),
});
// Should have removed the assistant message with only orphaned tool call
@ -293,7 +294,7 @@ describe('getChatConversationHistory - Orphaned Tool Call Cleanup', () => {
]);
const result = await getChatConversationHistory({
messageId: 'test-message-id',
messageId: randomUUID(),
});
expect(result).toEqual([]);
@ -328,7 +329,7 @@ describe('getChatConversationHistory - Orphaned Tool Call Cleanup', () => {
]);
const result = await getChatConversationHistory({
messageId: 'test-message-id',
messageId: randomUUID(),
});
expect(result).toHaveLength(2);