buster/packages/ai/src/tools/communication-tools/done-tool/helpers/file-selection.test.ts

import type { ChatMessageReasoningMessage } from '@buster/server-shared/chats';
import { describe, expect, it } from 'vitest';
import {
  type ExtractedFile,
  createFileResponseMessages,
  extractFilesFromReasoning,
  selectFilesForResponse,
} from './file-selection';

describe('file-selection', () => {
  describe('extractFilesFromReasoning', () => {
    it('should extract created metrics with version numbers', () => {
      const reasoningHistory: ChatMessageReasoningMessage[] = [
        {
          id: 'test-1',
          type: 'files',
          title: 'Created 2 metrics',
          status: 'completed',
          file_ids: ['metric-1', 'metric-2'],
          files: {
            'metric-1': {
              id: 'metric-1',
              file_type: 'metric',
              file_name: 'Revenue Metric',
              version_number: 1,
              status: 'completed',
              file: {
                text: '{"name": "Revenue Metric", "sql": "SELECT revenue FROM sales"}',
              },
            },
            'metric-2': {
              id: 'metric-2',
              file_type: 'metric',
              file_name: 'Growth Metric',
              version_number: 1,
              status: 'completed',
              file: {
                text: '{"name": "Growth Metric", "sql": "SELECT growth FROM analytics"}',
              },
            },
          },
        },
      ];

      const files = extractFilesFromReasoning(reasoningHistory);

      expect(files).toHaveLength(2);
      expect(files[0]).toMatchObject({
        id: 'metric-1',
        fileType: 'metric',
        fileName: 'Revenue Metric',
        status: 'completed',
        operation: 'created',
        versionNumber: 1,
      });
      expect(files[1]).toMatchObject({
        id: 'metric-2',
        fileType: 'metric',
        fileName: 'Growth Metric',
        status: 'completed',
        operation: 'created',
        versionNumber: 1,
      });
    });

    it('should extract modified dashboards with version numbers', () => {
      const reasoningHistory: ChatMessageReasoningMessage[] = [
        {
          id: 'test-1',
          type: 'files',
          title: 'Modified 1 dashboard',
          status: 'completed',
          file_ids: ['dashboard-1'],
          files: {
            'dashboard-1': {
              id: 'dashboard-1',
              file_type: 'dashboard',
              file_name: 'Sales Dashboard',
              version_number: 3,
              status: 'completed',
              file: {
                text: '{"name": "Sales Dashboard", "rows": [{"id": 1, "items": [{"id": "metric-1"}], "columnSizes": [12]}]}',
              },
            },
          },
        },
      ];

      const files = extractFilesFromReasoning(reasoningHistory);

      expect(files).toHaveLength(1);
      expect(files[0]).toMatchObject({
        id: 'dashboard-1',
        fileType: 'dashboard',
        fileName: 'Sales Dashboard',
        status: 'completed',
        operation: 'modified',
        versionNumber: 3,
      });
    });

    it('should build metric-to-dashboard relationships', () => {
      const reasoningHistory: ChatMessageReasoningMessage[] = [
        {
          id: 'test-1',
          type: 'files',
          title: 'Created 1 metric',
          status: 'completed',
          file_ids: ['metric-1'],
          files: {
            'metric-1': {
              id: 'metric-1',
              file_type: 'metric',
              file_name: 'Revenue Metric',
              version_number: 1,
              status: 'completed',
              file: {
                text: '{"name": "Revenue Metric"}',
              },
            },
          },
        },
        {
          id: 'test-2',
          type: 'files',
          title: 'Created 1 dashboard',
          status: 'completed',
          file_ids: ['dashboard-1'],
          files: {
            'dashboard-1': {
              id: 'dashboard-1',
              file_type: 'dashboard',
              file_name: 'Sales Dashboard',
              version_number: 1,
              status: 'completed',
              file: {
                text: '{"name": "Sales Dashboard", "rows": [{"id": 1, "items": [{"id": "metric-1"}], "columnSizes": [12]}]}',
              },
            },
          },
        },
      ];

      const files = extractFilesFromReasoning(reasoningHistory);
      const metric = files.find((f) => f.id === 'metric-1');

      expect(metric?.containedInDashboards).toEqual(['dashboard-1']);
    });
  });

  describe('selectFilesForResponse', () => {
    it('should return dashboard when metric inside it was modified', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 2,
          containedInDashboards: ['dashboard-1'],
        },
        {
          id: 'dashboard-1',
          fileType: 'dashboard',
          fileName: 'Sales Dashboard',
          status: 'completed',
          operation: 'created',
          versionNumber: 1,
          ymlContent:
            '{"name": "Sales Dashboard", "rows": [{"id": 1, "items": [{"id": "metric-1"}], "columnSizes": [12]}]}',
        },
      ];

      const selected = selectFilesForResponse(files);

      expect(selected).toHaveLength(1);
      expect(selected[0]?.id).toBe('dashboard-1');
    });

    it('should return standalone modified metrics', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 2,
          containedInDashboards: [],
        },
        {
          id: 'metric-2',
          fileType: 'metric',
          fileName: 'Growth Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 3,
          containedInDashboards: [],
        },
      ];

      const selected = selectFilesForResponse(files);

      expect(selected).toHaveLength(2);
      expect(selected.map((f) => f.id)).toEqual(['metric-1', 'metric-2']);
    });

    it('should return dashboard and standalone metric when both exist', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 2,
          containedInDashboards: ['dashboard-1'],
        },
        {
          id: 'metric-2',
          fileType: 'metric',
          fileName: 'Standalone Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 1,
          containedInDashboards: [],
        },
        {
          id: 'dashboard-1',
          fileType: 'dashboard',
          fileName: 'Sales Dashboard',
          status: 'completed',
          operation: 'created',
          versionNumber: 1,
          ymlContent:
            '{"name": "Sales Dashboard", "rows": [{"id": 1, "items": [{"id": "metric-1"}], "columnSizes": [12]}]}',
        },
      ];

      const selected = selectFilesForResponse(files);

      expect(selected).toHaveLength(2);
      expect(selected.map((f) => f.id).sort()).toEqual(['dashboard-1', 'metric-2'].sort());
    });

    it('should prioritize dashboards over metrics in standard cases', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'created',
          versionNumber: 1,
        },
        {
          id: 'dashboard-1',
          fileType: 'dashboard',
          fileName: 'Sales Dashboard',
          status: 'completed',
          operation: 'created',
          versionNumber: 1,
        },
      ];

      const selected = selectFilesForResponse(files);

      expect(selected).toHaveLength(1);
      expect(selected[0]?.id).toBe('dashboard-1');
    });

    it('should deduplicate files by ID and keep highest version', () => {
      const files: ExtractedFile[] = [
        {
          id: 'dashboard-1',
          fileType: 'dashboard',
          fileName: 'Sales Dashboard',
          status: 'completed',
          operation: 'created',
          versionNumber: 1,
        },
        {
          id: 'dashboard-1',
          fileType: 'dashboard',
          fileName: 'Sales Dashboard',
          status: 'completed',
          operation: 'modified',
          versionNumber: 2,
        },
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'created',
          versionNumber: 3,
        },
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 1,
        },
      ];

      const selected = selectFilesForResponse(files);

      expect(selected).toHaveLength(2);

      const dashboard = selected.find((f) => f.fileType === 'dashboard');
      const metric = selected.find((f) => f.fileType === 'metric');

      expect(dashboard?.versionNumber).toBe(2);
      expect(dashboard?.operation).toBe('modified');

      expect(metric?.versionNumber).toBe(3);
      expect(metric?.operation).toBe('created');
    });

    it('should default missing version numbers to 1 during deduplication', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'created',
          // No versionNumber provided (should default to 1)
        },
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 2,
        },
      ];

      const selected = selectFilesForResponse(files);

      expect(selected).toHaveLength(1);
      expect(selected[0]?.versionNumber).toBe(2);
      expect(selected[0]?.operation).toBe('modified');
    });
  });

  describe('createFileResponseMessages', () => {
    it('should create response messages with correct version numbers', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          operation: 'modified',
          versionNumber: 3,
        },
        {
          id: 'dashboard-1',
          fileType: 'dashboard',
          fileName: 'Sales Dashboard',
          status: 'completed',
          operation: 'created',
          versionNumber: 1,
        },
      ];

      const messages = createFileResponseMessages(files);

      expect(messages).toHaveLength(2);
      expect(messages[0]).toMatchObject({
        id: 'metric-1',
        type: 'file',
        file_type: 'metric',
        file_name: 'Revenue Metric',
        version_number: 3,
      });
      expect((messages[0] as any).metadata?.[0]?.message).toBe('Metric modified successfully');

      expect(messages[1]).toMatchObject({
        id: 'dashboard-1',
        type: 'file',
        file_type: 'dashboard',
        file_name: 'Sales Dashboard',
        version_number: 1,
      });
      expect((messages[1] as any).metadata?.[0]?.message).toBe('Dashboard created successfully');
    });

    it('should default to version 1 if version number is missing', () => {
      const files: ExtractedFile[] = [
        {
          id: 'metric-1',
          fileType: 'metric',
          fileName: 'Revenue Metric',
          status: 'completed',
          // No versionNumber provided
        },
      ];

      const messages = createFileResponseMessages(files);

      expect((messages[0] as any).version_number).toBe(1);
    });
  });
});