[v10.2.x] SaveDashboardPrompt: Reduce time to open drawer when many changes applied (#78308)

* SaveDashboard: Reduce time to open drawer when many changes applied (#78283)

(cherry picked from commit f32f8a160e)

public/app/features/dashboard/components/GenAI/GenAIButton.test.tsx

@@ -13,6 +13,7 @@ import { EventTrackingSrc } from './tracking';
 import { Role } from './utils';
 
 const mockedUseOpenAiStreamState = {
+  messages: [],
   setMessages: jest.fn(),
   reply: 'I am a robot',
   streamStatus: StreamStatus.IDLE,
@@ -43,6 +44,7 @@ describe('GenAIButton', () => {
   describe('when LLM plugin is not configured', () => {
     beforeAll(() => {
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: undefined,
         streamStatus: StreamStatus.IDLE,
         reply: 'Some completed genereated text',
@@ -64,7 +66,10 @@ describe('GenAIButton', () => {
   describe('when LLM plugin is properly configured, so it is enabled', () => {
     const setMessagesMock = jest.fn();
     beforeEach(() => {
+      setMessagesMock.mockClear();
+
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: undefined,
         streamStatus: StreamStatus.IDLE,
         reply: 'Some completed genereated text',
@@ -100,6 +105,20 @@ describe('GenAIButton', () => {
       expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]);
     });
 
+    it('should call the messages when they are provided as callback', async () => {
+      const onGenerate = jest.fn();
+      const messages = jest.fn().mockReturnValue([{ content: 'Generate X', role: 'system' as Role }]);
+      const onClick = jest.fn();
+      setup({ onGenerate, messages, temperature: 3, onClick, eventTrackingSrc });
+
+      const generateButton = await screen.findByRole('button');
+      await fireEvent.click(generateButton);
+
+      expect(messages).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]);
+    });
+
     it('should call the onClick callback', async () => {
       const onGenerate = jest.fn();
       const onClick = jest.fn();
@@ -116,6 +135,7 @@ describe('GenAIButton', () => {
   describe('when it is generating data', () => {
     beforeEach(() => {
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: undefined,
         streamStatus: StreamStatus.GENERATING,
         reply: 'Some incomplete generated text',
@@ -160,7 +180,10 @@ describe('GenAIButton', () => {
   describe('when there is an error generating data', () => {
     const setMessagesMock = jest.fn();
     beforeEach(() => {
+      setMessagesMock.mockClear();
+
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: new Error('Something went wrong'),
         streamStatus: StreamStatus.IDLE,
         reply: '',
@@ -224,5 +247,19 @@ describe('GenAIButton', () => {
 
       await waitFor(() => expect(onClick).toHaveBeenCalledTimes(1));
     });
+
+    it('should call the messages when they are provided as callback', async () => {
+      const onGenerate = jest.fn();
+      const messages = jest.fn().mockReturnValue([{ content: 'Generate X', role: 'system' as Role }]);
+      const onClick = jest.fn();
+      setup({ onGenerate, messages, temperature: 3, onClick, eventTrackingSrc });
+
+      const generateButton = await screen.findByRole('button');
+      await fireEvent.click(generateButton);
+
+      expect(messages).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]);
+    });
   });
 });

public/app/features/dashboard/components/GenAI/GenAIButton.tsx

@@ -18,7 +18,7 @@ export interface GenAIButtonProps {
   // Button click handler
   onClick?: (e: React.MouseEvent<HTMLButtonElement>) => void;
   // Messages to send to the LLM plugin
-  messages: Message[];
+  messages: Message[] | (() => Message[]);
   // Callback function that the LLM plugin streams responses to
   onGenerate: (response: string) => void;
   // Temperature for the LLM plugin. Default is 1.
@@ -43,8 +43,14 @@ export const GenAIButton = ({
 }: GenAIButtonProps) => {
   const styles = useStyles2(getStyles);
 
-  const { setMessages, reply, value, error, streamStatus } = useOpenAIStream(OPEN_AI_MODEL, temperature);
-
+  const {
+    messages: streamMessages,
+    setMessages,
+    reply,
+    value,
+    error,
+    streamStatus,
+  } = useOpenAIStream(OPEN_AI_MODEL, temperature);
   const [history, setHistory] = useState<string[]>([]);
   const [showHistory, setShowHistory] = useState(true);
 
@@ -56,7 +62,7 @@ export const GenAIButton = ({
   const onClick = (e: React.MouseEvent<HTMLButtonElement>) => {
     if (!hasHistory) {
       onClickProp?.(e);
-      setMessages(messages);
+      setMessages(typeof messages === 'function' ? messages() : messages);
     } else {
       if (setShowHistory) {
         setShowHistory(true);
@@ -154,7 +160,7 @@ export const GenAIButton = ({
           content={
             <GenAIHistory
               history={history}
-              messages={messages}
+              messages={streamMessages}
               onApplySuggestion={onApplySuggestion}
               updateHistory={pushHistoryEntry}
               eventTrackingSrc={eventTrackingSrc}

public/app/features/dashboard/components/GenAI/GenAIDashboardChangesButton.tsx

@@ -1,4 +1,4 @@
-import React, { useMemo } from 'react';
+import React, { useCallback } from 'react';
 
 import { DashboardModel } from '../../state';
 
@@ -27,7 +27,7 @@ const CHANGES_GENERATION_STANDARD_PROMPT = [
 ].join('.\n');
 
 export const GenAIDashboardChangesButton = ({ dashboard, onGenerate, disabled }: GenAIDashboardChangesButtonProps) => {
-  const messages = useMemo(() => getMessages(dashboard), [dashboard]);
+  const messages = useCallback(() => getMessages(dashboard), [dashboard]);
 
   return (
     <GenAIButton

public/app/features/dashboard/components/GenAI/hooks.ts

@@ -26,6 +26,7 @@ export function useOpenAIStream(
   temperature = 1
 ): {
   setMessages: React.Dispatch<React.SetStateAction<Message[]>>;
+  messages: Message[];
   reply: string;
   streamStatus: StreamStatus;
   error: Error | undefined;
@@ -138,6 +139,7 @@ export function useOpenAIStream(
 
   return {
     setMessages,
+    messages,
     reply,
     streamStatus,
     error,