Merge remote-tracking branch 'origin/staging' into add-report-to-collection-functionality

apps/web/src/controllers/ReportPageControllers/GeneratingContent.tsx

@@ -1,7 +1,5 @@
-import { queryKeys } from '@/api/query_keys';
 import { ShimmerText } from '@/components/ui/typography/ShimmerText';
 import { cn } from '@/lib/classMerge';
-import { useQuery } from '@tanstack/react-query';
 
 export const GeneratingContent = ({
   messageId,
@@ -10,18 +8,9 @@ export const GeneratingContent = ({
   messageId: string;
   className?: string;
 }) => {
-  const { data: text } = useQuery({
-    ...queryKeys.chatsBlackBoxMessages(messageId),
-    notifyOnChangeProps: ['data'],
-    select: (data) => data
-  });
-
   return (
-    <div
-      className={cn('right-0 bottom-0 left-0 -mt-68 flex items-center justify-center', className)}>
-      <div className="border-border item-center flex w-full justify-center rounded border px-8 py-1.5 shadow">
-        <ShimmerText text={text || 'Generating content...'} className="text-lg" />
-      </div>
+    <div className={cn('right-0 bottom-0 left-0 -mt-68', className)}>
+      <ShimmerText text="Generating..." className="text-lg" />
     </div>
   );
 };

apps/web/src/controllers/ReportPageControllers/ReportPageController.tsx

@@ -11,6 +11,9 @@ import { ReportEditorSkeleton } from '@/components/ui/report/ReportEditorSkeleto
 import { useChatIndividualContextSelector } from '@/layouts/ChatLayout/ChatContext';
 import { useTrackAndUpdateReportChanges } from '@/api/buster-electric/reports/hooks';
 import { GeneratingContent } from './GeneratingContent';
+import { useQuery } from '@tanstack/react-query';
+import { queryKeys } from '@/api/query_keys';
+import type { BusterChatMessage } from '@/api/asset_interfaces/chat';
 
 export const ReportPageController: React.FC<{
   reportId: string;
@@ -24,8 +27,25 @@ export const ReportPageController: React.FC<{
     const isStreamingMessage = useChatIndividualContextSelector((x) => x.isStreamingMessage);
     const messageId = useChatIndividualContextSelector((x) => x.currentMessageId);
 
+    // Fetch the current message to check which files are being generated
+    const { data: currentMessage } = useQuery<BusterChatMessage>({
+      ...queryKeys.chatsMessages(messageId || ''),
+      enabled: !!messageId && isStreamingMessage
+    });
+
+    // Check if this specific report is being generated in the current message
+    const isThisReportBeingGenerated = React.useMemo(() => {
+      if (!currentMessage || !isStreamingMessage || !messageId) return false;
+
+      // Check if the current report ID matches any file being generated
+      const responseMessages = Object.values(currentMessage.response_messages || {});
+      return responseMessages.some(
+        (msg) => msg.type === 'file' && msg.file_type === 'report' && msg.id === reportId
+      );
+    }, [currentMessage, isStreamingMessage, messageId, reportId]);
+
     const content = report?.content || '';
-    const showGeneratingContent = messageId && isStreamingMessage;
+    const showGeneratingContent = isThisReportBeingGenerated;
     const commonClassName = 'sm:px-[max(64px,calc(50%-350px))]';
 
     const { mutate: updateReport } = useUpdateReport();

packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-delta.ts

@@ -219,9 +219,6 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
 
         // Update report content for all reports that have content
         if (contentUpdates.length > 0) {
-          // Track response messages to create in batch
-          const responseMessagesToCreate: ChatMessageResponseMessage[] = [];
-
           for (const update of contentUpdates) {
             try {
               await updateReportContent({
@@ -229,37 +226,16 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
                 content: update.content,
               });
 
-              // Mark the file as completed in state
+              // Keep the file status as 'loading' during streaming
+              // Status will be updated to 'completed' in the execute phase
               const stateFile = state.files?.find((f) => f.id === update.reportId);
               if (stateFile) {
-                stateFile.status = 'completed';
-
-                // Create response message for this report if not already created
-                if (!state.responseMessagesCreated?.has(update.reportId)) {
-                  // Create response message for this report
-                  responseMessagesToCreate.push({
-                    id: update.reportId,
-                    type: 'file' as const,
-                    file_type: 'report' as const,
-                    file_name: stateFile.file_name || '',
-                    version_number: stateFile.version_number || 1,
-                    filter_version_id: null,
-                    metadata: [
-                      {
-                        status: 'completed' as const,
-                        message: 'Report created successfully',
-                        timestamp: Date.now(),
-                      },
-                    ],
-                  });
-
-                  // Track that we've created a response message for this report
-                  if (!state.responseMessagesCreated) {
-                    state.responseMessagesCreated = new Set<string>();
-                  }
-                  state.responseMessagesCreated.add(update.reportId);
-                }
+                // Ensure status remains 'loading' during delta phase
+                stateFile.status = 'loading';
               }
+
+              // Note: Response messages should only be created in execute phase
+              // after all processing is complete
             } catch (error) {
               const errorMessage = error instanceof Error ? error.message : 'Content update failed';
               console.error('[create-reports] Error updating report content:', {
@@ -268,35 +244,15 @@ export function createCreateReportsDelta(context: CreateReportsContext, state: C
                 stack: error instanceof Error ? error.stack : undefined,
               });
 
-              // Mark the file as failed in state with error message
+              // Keep file as loading during delta phase even on error
+              // The execute phase will handle final status
               const stateFile = state.files?.find((f) => f.id === update.reportId);
               if (stateFile) {
-                stateFile.status = 'failed';
+                stateFile.status = 'loading';
                 stateFile.error = `Failed to update report content: ${errorMessage}`;
               }
             }
           }
-
-          // Update database with response messages if we have any
-          if (responseMessagesToCreate.length > 0 && context.messageId) {
-            try {
-              await updateMessageEntries({
-                messageId: context.messageId,
-                responseMessages: responseMessagesToCreate,
-              });
-
-              console.info('[create-reports] Created response messages during delta', {
-                count: responseMessagesToCreate.length,
-                reportIds: responseMessagesToCreate.map((m) => m.id),
-              });
-            } catch (error) {
-              console.error(
-                '[create-reports] Error creating response messages during delta:',
-                error
-              );
-              // Don't throw - continue processing
-            }
-          }
         }
       }
     }

packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/helpers/modify-reports-transform-helper.ts

@@ -46,15 +46,28 @@ export function createModifyReportsReasoningEntry(
   let status: 'loading' | 'completed' | 'failed' = 'loading';
   let secondaryTitle: string | undefined;
 
-  // Check if modification is complete based on state
-  if (state.finalContent !== undefined) {
+  // Check if all edits have a final status (completed or failed), not just 'loading'
+  const allEditsComplete =
+    state.edits && state.edits.length > 0
+      ? state.edits.every((edit) => edit.status === 'completed' || edit.status === 'failed')
+      : false;
+
+  // Only mark as complete when all edits are actually done, not during streaming
+  if (allEditsComplete) {
     // Check if any edits failed
     const hasFailedEdits = state.edits?.some((edit) => edit.status === 'failed') ?? false;
 
     if (hasFailedEdits) {
       title = 'Failed to modify report';
       status = 'failed';
-    } else if (state.finalContent) {
+      // Update the file status in filesRecord
+      if (state.reportId) {
+        const file = filesRecord[state.reportId];
+        if (file) {
+          file.status = 'failed';
+        }
+      }
+    } else {
       title = 'Modified 1 report';
       status = 'completed';
       // Update the file status in filesRecord
@@ -66,14 +79,15 @@ export function createModifyReportsReasoningEntry(
       }
     }
 
-    // Only show elapsed time when all edits are complete (not during streaming)
-    // Check if all edits have a final status (completed or failed), not just 'loading'
-    const allEditsComplete =
-      state.edits?.every((edit) => edit.status === 'completed' || edit.status === 'failed') ??
-      false;
-
-    if (allEditsComplete) {
-      secondaryTitle = formatElapsedTime(state.startTime);
+    // Show elapsed time when complete
+    secondaryTitle = formatElapsedTime(state.startTime);
+  } else {
+    // Keep file status as loading during streaming
+    if (state.reportId) {
+      const file = filesRecord[state.reportId];
+      if (file) {
+        file.status = 'loading';
+      }
     }
   }
 

packages/database/src/queries/messages/update-message-entries.ts

@@ -127,6 +127,7 @@ export async function updateMessageEntries({
             WITH new_messages AS (
               SELECT 
                 value,
+                ordinality as input_order,
                 value->>'role' AS role,
                 COALESCE(
                   CASE 
@@ -138,7 +139,7 @@ export async function updateMessageEntries({
                   END,
                   ''
                 ) AS tool_calls
-              FROM jsonb_array_elements(${newData}::jsonb) AS value
+              FROM jsonb_array_elements(${newData}::jsonb) WITH ORDINALITY AS t(value, ordinality)
             ),
             existing_messages AS (
               SELECT 
@@ -170,8 +171,8 @@ export async function updateMessageEntries({
                 WHERE n.role = e.role AND n.tool_calls = e.tool_calls
               )
               UNION ALL
-              -- Add all new messages
-              SELECT n.value, 1000000 + row_number() OVER () AS ord
+              -- Add all new messages, preserving their input order
+              SELECT n.value, 1000000 + n.input_order AS ord
               FROM new_messages n
             ) combined
           )