refactor: enhance create dashboards tool with improved state management and type definitions

- Removed unused properties from the dashboard state and output schemas for clarity.
- Introduced a new state file schema to better represent dashboard files.
- Updated helper functions to streamline file transformation and reasoning message creation.
- Improved type safety by utilizing the newly defined `StatusSchema`.

These changes optimize the dashboard creation process and enhance code maintainability.

packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-start.ts

@@ -12,8 +12,6 @@ export function createDashboardsStart(
 ) {
   return async (options: ToolCallOptions) => {
     state.toolCallId = options.toolCallId;
-    state.argsText = '';
-    state.files = [];
 
     if (context.messageId) {
       try {

packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool.ts

@@ -1,3 +1,4 @@
+import { StatusSchema } from '@buster/server-shared/chats';
 import { tool } from 'ai';
 import { z } from 'zod';
 import { getDashboardToolDescription } from '../helpers/get-dashboard-tool-description';
@@ -19,11 +20,6 @@ const CreateDashboardsInputSchema = z.object({
 const CreateDashboardsOutputFileSchema = z.object({
   id: z.string(),
   name: z.string(),
-  file_type: z.string(),
-  result_message: z.string().optional(),
-  results: z.array(z.record(z.any())).optional(),
-  created_at: z.string(),
-  updated_at: z.string(),
   version_number: z.number(),
 });
 
@@ -35,7 +31,6 @@ const CreateDashboardsOutputFailedFileSchema = z.object({
 // Output schema for the create dashboards tool
 const CreateDashboardsOutputSchema = z.object({
   message: z.string(),
-  duration: z.number(),
   files: z.array(CreateDashboardsOutputFileSchema),
   failed_files: z.array(CreateDashboardsOutputFailedFileSchema),
 });
@@ -50,20 +45,23 @@ const CreateDashboardsContextSchema = z.object({
   messageId: z.string().optional().describe('The message ID'),
 });
 
-const CreateDashboardsReasoningFileSchema = z.object({
-  name: z.string(),
-  yml_content: z.string(),
-  status: z.enum(['processing', 'completed', 'failed']).optional(),
-  id: z.string().optional(),
-  version: z.number().optional(),
-  error: z.string().optional(),
+const CreateDashboardStateFileSchema = z.object({
+  id: z.string().uuid(),
+  file_name: z.string().optional(),
+  file_type: z.string(),
+  version_number: z.number(),
+  file: z
+    .object({
+      text: z.string(),
+    })
+    .optional(),
+  status: StatusSchema,
 });
 
 const CreateDashboardsStateSchema = z.object({
   toolCallId: z.string().optional(),
   argsText: z.string().optional(),
-  parsedArgs: CreateDashboardsInputSchema.optional(),
-  files: z.array(CreateDashboardsReasoningFileSchema).optional(),
+  files: z.array(CreateDashboardStateFileSchema).optional(),
 });
 
 // Export types
@@ -74,8 +72,8 @@ export type CreateDashboardsOutputFile = z.infer<typeof CreateDashboardsOutputFi
 export type CreateDashboardsOutputFailedFile = z.infer<
   typeof CreateDashboardsOutputFailedFileSchema
 >;
-export type CreateDashboardsReasoningFile = z.infer<typeof CreateDashboardsReasoningFileSchema>;
 export type CreateDashboardsState = z.infer<typeof CreateDashboardsStateSchema>;
+export type CreateDashboardStateFile = z.infer<typeof CreateDashboardStateFileSchema>;
 
 // Factory function that accepts agent context and maps to tool context
 export function createCreateDashboardsTool(context: CreateDashboardsContext) {
@@ -83,7 +81,6 @@ export function createCreateDashboardsTool(context: CreateDashboardsContext) {
   const state: CreateDashboardsState = {
     argsText: undefined,
     files: undefined,
-    parsedArgs: undefined,
     toolCallId: undefined,
   };
 

packages/ai/src/tools/visualization-tools/dashboards/create-dashboards-tool/helpers/create-dashboards-tool-transform-helper.ts

@@ -1,4 +1,7 @@
-import type { ChatMessageReasoningMessage } from '@buster/server-shared/chats';
+import type {
+  ChatMessageReasoningMessage,
+  ChatMessageReasoningMessage_File,
+} from '@buster/server-shared/chats';
 import type { ModelMessage } from 'ai';
 import type { CreateDashboardsState } from '../create-dashboards-tool';
 
@@ -9,72 +12,42 @@ export function createCreateDashboardsReasoningEntry(
   state: CreateDashboardsState,
   toolCallId: string
 ): ChatMessageReasoningMessage | undefined {
+  state.toolCallId = toolCallId;
+
   if (!state.files || state.files.length === 0) {
     return undefined;
   }
 
-  // Transform files to the format expected by reasoning messages
-  interface FileEntry {
-    id: string;
-    file_type: string;
-    file_name: string;
-    version_number?: number;
-    status: string;
-    file: {
-      text: string;
-    };
-    error?: string;
-  }
-  const fileEntries: Record<string, FileEntry> = {};
+  // Build Record<string, ReasoningFile> as required by schema
+  const filesRecord: Record<string, ChatMessageReasoningMessage_File> = {};
   const fileIds: string[] = [];
-
-  state.files.forEach((file) => {
-    const fileId = file.id || `dashboard-${toolCallId}-${file.name}`;
-    fileIds.push(fileId);
-    fileEntries[fileId] = {
-      id: fileId,
+  for (const f of state.files) {
+    // Skip entries that do not yet have a file_name
+    if (!f.file_name) continue;
+    const id = f.id;
+    fileIds.push(id);
+    filesRecord[id] = {
+      id,
       file_type: 'dashboard',
-      file_name: file.name,
-      status: file.status || 'processing',
+      file_name: f.file_name,
+      version_number: f.version_number,
+      status: f.status,
       file: {
-        text: file.yml_content,
+        text: f.file?.text,
       },
-      ...(file.version !== undefined ? { version_number: file.version } : {}),
-      ...(file.error && { error: file.error }),
     };
-  });
-
-  // Determine status and title based on file states
-  let title = 'Creating dashboards';
-  let status: 'loading' | 'completed' | 'failed' = 'loading';
-
-  const processingCount = state.files.filter((f) => f.status === 'processing').length;
-  const completedCount = state.files.filter((f) => f.status === 'completed').length;
-  const failedCount = state.files.filter((f) => f.status === 'failed').length;
-
-  if (processingCount > 0) {
-    title = `Creating ${state.files.length} ${state.files.length === 1 ? 'dashboard' : 'dashboards'}`;
-    status = 'loading';
-  } else if (failedCount === state.files.length) {
-    title = 'Failed to create dashboards';
-    status = 'failed';
-  } else if (completedCount > 0 || failedCount > 0) {
-    if (failedCount > 0) {
-      title = `Created ${completedCount} ${completedCount === 1 ? 'dashboard' : 'dashboards'}, ${failedCount} failed`;
-      status = 'failed';
-    } else {
-      title = `Created ${completedCount} ${completedCount === 1 ? 'dashboard' : 'dashboards'}`;
-      status = 'completed';
-    }
   }
 
+  // If nothing valid to show yet, skip emitting a files reasoning message
+  if (fileIds.length === 0) return undefined;
+
   return {
     id: toolCallId,
     type: 'files',
-    title,
-    status,
+    title: 'Creating dashboards...',
+    status: 'loading',
     file_ids: fileIds,
-    files: fileEntries,
+    files: filesRecord,
   } as ChatMessageReasoningMessage;
 }
 
@@ -85,9 +58,8 @@ export function createCreateDashboardsRawLlmMessageEntry(
   state: CreateDashboardsState,
   toolCallId: string
 ): ModelMessage | undefined {
-  if (!state.parsedArgs || !state.parsedArgs.files || state.parsedArgs.files.length === 0) {
-    return undefined;
-  }
+  // If we don't have files yet, skip emitting raw LLM entry
+  if (!state.files || state.files.length === 0) return undefined;
 
   return {
     role: 'assistant',
@@ -97,7 +69,13 @@ export function createCreateDashboardsRawLlmMessageEntry(
         toolCallId,
         toolName: 'createDashboards',
         input: {
-          files: state.parsedArgs.files,
+          files: state.files
+            .map((file) => ({
+              name: file.file_name ?? 'dashboard',
+              yml_content: file.file?.text ?? '',
+            }))
+            // Filter out clearly invalid entries
+            .filter((f) => f.name && f.yml_content),
         },
       },
     ],

packages/server-shared/src/chats/chat-message.types.ts

@@ -14,7 +14,7 @@ const ChatUserMessageSchema = z
   })
   .or(z.null());
 
-const StatusSchema = z.enum(['loading', 'completed', 'failed']);
+export const StatusSchema = z.enum(['loading', 'completed', 'failed']);
 
 const ResponseMessage_TextSchema = z.object({
   id: z.string(),