diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx
index 81720b31ea94b..0a7cfea7768bd 100644
--- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx
+++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx
@@ -275,75 +275,74 @@ export function ChatBody({
     navigator.clipboard?.writeText(content || '');
   };
 
-  const handleActionClick = ({
-    message,
-    payload,
-  }: {
-    message: Message;
-    payload: ChatActionClickPayload;
-  }) => {
-    setStickToBottom(true);
-    switch (payload.type) {
-      case ChatActionClickType.executeEsqlQuery:
-        next(
-          messages.concat({
-            '@timestamp': new Date().toISOString(),
-            message: {
-              role: MessageRole.Assistant,
-              content: '',
-              function_call: {
-                name: 'execute_query',
-                arguments: JSON.stringify({
-                  query: payload.query,
-                }),
-                trigger: MessageRole.User,
+  const handleActionClick = useCallback(
+    ({ message, payload }: { message: Message; payload: ChatActionClickPayload }) => {
+      setStickToBottom(true);
+      switch (payload.type) {
+        case ChatActionClickType.executeEsqlQuery:
+          next(
+            messages.concat({
+              '@timestamp': new Date().toISOString(),
+              message: {
+                role: MessageRole.Assistant,
+                content: '',
+                function_call: {
+                  name: 'execute_query',
+                  arguments: JSON.stringify({
+                    query: payload.query,
+                  }),
+                  trigger: MessageRole.User,
+                },
               },
-            },
-          })
-        );
-        break;
-
-      case ChatActionClickType.updateVisualization:
-        const visualizeQueryResponse = message;
-
-        const visualizeQueryResponseData = JSON.parse(visualizeQueryResponse.message.data ?? '{}');
-
-        next(
-          messages.slice(0, messages.indexOf(visualizeQueryResponse)).concat({
-            '@timestamp': new Date().toISOString(),
-            message: {
-              name: 'visualize_query',
-              content: visualizeQueryResponse.message.content,
-              data: JSON.stringify({
-                ...visualizeQueryResponseData,
-                userOverrides: payload.userOverrides,
-              }),
-              role: MessageRole.User,
-            },
-          })
-        );
-        break;
-      case ChatActionClickType.visualizeEsqlQuery:
-        next(
-          messages.concat({
-            '@timestamp': new Date().toISOString(),
-            message: {
-              role: MessageRole.Assistant,
-              content: '',
-              function_call: {
+            })
+          );
+          break;
+
+        case ChatActionClickType.updateVisualization:
+          const visualizeQueryResponse = message;
+
+          const visualizeQueryResponseData = JSON.parse(
+            visualizeQueryResponse.message.data ?? '{}'
+          );
+
+          next(
+            messages.slice(0, messages.indexOf(visualizeQueryResponse)).concat({
+              '@timestamp': new Date().toISOString(),
+              message: {
                 name: 'visualize_query',
-                arguments: JSON.stringify({
-                  query: payload.query,
-                  intention: VisualizeESQLUserIntention.visualizeAuto,
+                content: visualizeQueryResponse.message.content,
+                data: JSON.stringify({
+                  ...visualizeQueryResponseData,
+                  userOverrides: payload.userOverrides,
                 }),
-                trigger: MessageRole.User,
+                role: MessageRole.User,
               },
-            },
-          })
-        );
-        break;
-    }
-  };
+            })
+          );
+          break;
+        case ChatActionClickType.visualizeEsqlQuery:
+          next(
+            messages.concat({
+              '@timestamp': new Date().toISOString(),
+              message: {
+                role: MessageRole.Assistant,
+                content: '',
+                function_call: {
+                  name: 'visualize_query',
+                  arguments: JSON.stringify({
+                    query: payload.query,
+                    intention: VisualizeESQLUserIntention.visualizeAuto,
+                  }),
+                  trigger: MessageRole.User,
+                },
+              },
+            })
+          );
+          break;
+      }
+    },
+    [messages, next]
+  );
 
   if (!hasCorrectLicense && !initialConversationId) {
     footer = (
diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx b/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx
index a473e34ba2a0c..2d2bc79550dcf 100644
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx
@@ -18,7 +18,7 @@ import {
 import { css } from '@emotion/css';
 import classNames from 'classnames';
 import type { Code, InlineCode, Parent, Text } from 'mdast';
-import React, { useMemo, useRef } from 'react';
+import React, { useMemo } from 'react';
 import type { Node } from 'unist';
 import { ChatActionClickHandler } from '../chat/types';
 import { CodeBlock, EsqlCodeBlock } from './esql_code_block';
@@ -120,10 +120,6 @@ export function MessageText({ loading, content, onActionClick }: Props) {
     overflow-wrap: anywhere;
   `;
 
-  const onActionClickRef = useRef(onActionClick);
-
-  onActionClickRef.current = onActionClick;
-
   const { parsingPluginList, processingPluginList } = useMemo(() => {
     const parsingPlugins = getDefaultEuiMarkdownParsingPlugins();
 
@@ -148,7 +144,7 @@ export function MessageText({ loading, content, onActionClick }: Props) {
             <EsqlCodeBlock
               value={props.value}
               actionsDisabled={loading}
-              onActionClick={onActionClickRef.current}
+              onActionClick={onActionClick}
             />
             <EuiSpacer size="m" />
           </>
@@ -186,7 +182,7 @@ export function MessageText({ loading, content, onActionClick }: Props) {
       parsingPluginList: [loadingCursorPlugin, esqlLanguagePlugin, ...parsingPlugins],
       processingPluginList: processingPlugins,
     };
-  }, [loading]);
+  }, [loading, onActionClick]);
 
   return (
     <EuiText size="s" className={containerClassName}>