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

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',
@@ -66,7 +68,11 @@ describe('GenAIButton', () => {
     const setMessagesMock = jest.fn();
     const setShouldStopMock = jest.fn();
     beforeEach(() => {
+      setMessagesMock.mockClear();
+      setShouldStopMock.mockClear();
+
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: undefined,
         streamStatus: StreamStatus.IDLE,
         reply: 'Some completed genereated text',
@@ -103,6 +109,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();
@@ -121,6 +141,7 @@ describe('GenAIButton', () => {
 
     beforeEach(() => {
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: undefined,
         streamStatus: StreamStatus.GENERATING,
         reply: 'Some incomplete generated text',
@@ -177,7 +198,11 @@ describe('GenAIButton', () => {
     const setMessagesMock = jest.fn();
     const setShouldStopMock = jest.fn();
     beforeEach(() => {
+      setMessagesMock.mockClear();
+      setShouldStopMock.mockClear();
+
       jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
         error: new Error('Something went wrong'),
         streamStatus: StreamStatus.IDLE,
         reply: '',
@@ -242,5 +267,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

@@ -16,7 +16,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,7 +43,15 @@ export const GenAIButton = ({
 }: GenAIButtonProps) => {
   const styles = useStyles2(getStyles);
 
-  const { setMessages, setStopGeneration, reply, value, error, streamStatus } = useOpenAIStream(model, temperature);
+  const {
+    messages: streamMessages,
+    setMessages,
+    setStopGeneration,
+    reply,
+    value,
+    error,
+    streamStatus,
+  } = useOpenAIStream(model, temperature);
 
   const [history, setHistory] = useState<string[]>([]);
   const [showHistory, setShowHistory] = useState(true);
@@ -59,7 +67,7 @@ export const GenAIButton = ({
     } else {
       if (!hasHistory) {
         onClickProp?.(e);
-        setMessages(messages);
+        setMessages(typeof messages === 'function' ? messages() : messages);
       } else {
         if (setShowHistory) {
           setShowHistory(true);
@@ -161,7 +169,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';
 
@@ -35,7 +35,7 @@ const CHANGES_GENERATION_POSTFIX_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

@@ -27,6 +27,7 @@ export function useOpenAIStream(
 ): {
   setMessages: React.Dispatch<React.SetStateAction<Message[]>>;
   setStopGeneration: React.Dispatch<React.SetStateAction<boolean>>;
+  messages: Message[];
   reply: string;
   streamStatus: StreamStatus;
   error: Error | undefined;
@@ -154,6 +155,7 @@ export function useOpenAIStream(
   return {
     setMessages,
     setStopGeneration,
+    messages,
     reply,
     streamStatus,
     error,