[8.11] [Security Solution] [Elastic AI Assistant] Fixes ES|QL codeblocks from not being able to be sent to Timeline (#169478) (#169594)

# Backport This will backport the following commits from `main` to `8.11`: - [[Security Solution] [Elastic AI Assistant] Fixes `ES|QL` codeblocks from not being able to be sent to Timeline (#169478)](https://github.com/elastic/kibana/pull/169478)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Garrett Spong <spong@users.noreply.github.com>
2025-04-24 09:48:58 -04:00 · 2023-10-23 19:33:17 -04:00 · 2023-10-23 19:33:17 -04:00 · 03866cb70b
commit 03866cb70b
parent 68e8bc6328
7 changed files with 184 additions and 31 deletions
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts
@ -18,7 +18,7 @@ export interface CodeBlockDetails {
  button?: React.ReactNode;
 }

-export type QueryType = 'eql' | 'kql' | 'dsl' | 'json' | 'no-type';
+export type QueryType = 'eql' | 'esql' | 'kql' | 'dsl' | 'json' | 'no-type' | 'sql';

 /**
 * `analyzeMarkdown` is a helper that enriches content returned from a query
@ -35,6 +35,7 @@ export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => {
  // If your codeblocks aren't getting tagged with the right language, add keywords to the array.
  const types = {
    eql: ['Event Query Language', 'EQL sequence query', 'EQL'],
+    esql: ['Elasticsearch Query Language', 'ESQL', 'ES|QL', 'SQL'],
    kql: ['Kibana Query Language', 'KQL Query', 'KQL'],
    dsl: [
      'Elasticsearch QueryDSL',
--- a/x-pack/plugins/security_solution/public/app/app.tsx
+++ b/x-pack/plugins/security_solution/public/app/app.tsx
@ -66,29 +66,29 @@ const StartAppComponent: FC<StartAppComponent> = ({
          <ReduxStoreProvider store={store}>
            <KibanaThemeProvider theme$={theme$}>
              <EuiThemeProvider darkMode={darkMode}>
-                <AssistantProvider>
-                  <MlCapabilitiesProvider>
-                    <UserPrivilegesProvider kibanaCapabilities={capabilities}>
-                      <ManageUserInfo>
-                        <NavigationProvider core={services}>
-                          <ReactQueryClientProvider>
-                            <CellActionsProvider
-                              getTriggerCompatibleActions={uiActions.getTriggerCompatibleActions}
-                            >
-                              <UpsellingProvider upsellingService={upselling}>
-                                <DiscoverInTimelineContextProvider>
+                <MlCapabilitiesProvider>
+                  <UserPrivilegesProvider kibanaCapabilities={capabilities}>
+                    <ManageUserInfo>
+                      <NavigationProvider core={services}>
+                        <ReactQueryClientProvider>
+                          <CellActionsProvider
+                            getTriggerCompatibleActions={uiActions.getTriggerCompatibleActions}
+                          >
+                            <UpsellingProvider upsellingService={upselling}>
+                              <DiscoverInTimelineContextProvider>
+                                <AssistantProvider>
                                  <PageRouter history={history} onAppLeave={onAppLeave}>
                                    {children}
                                  </PageRouter>
-                                </DiscoverInTimelineContextProvider>
-                              </UpsellingProvider>
-                            </CellActionsProvider>
-                          </ReactQueryClientProvider>
-                        </NavigationProvider>
-                      </ManageUserInfo>
-                    </UserPrivilegesProvider>
-                  </MlCapabilitiesProvider>
-                </AssistantProvider>
+                                </AssistantProvider>
+                              </DiscoverInTimelineContextProvider>
+                            </UpsellingProvider>
+                          </CellActionsProvider>
+                        </ReactQueryClientProvider>
+                      </NavigationProvider>
+                    </ManageUserInfo>
+                  </UserPrivilegesProvider>
+                </MlCapabilitiesProvider>
              </EuiThemeProvider>
            </KibanaThemeProvider>
            <ErrorToastDispatcher />
--- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_code_block.tsx
@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { EuiCodeBlock, EuiFlexGroup, EuiFlexItem, EuiPanel, useEuiTheme } from '@elastic/eui';
+import { css } from '@emotion/css';
+import React from 'react';
+
+export const CustomCodeBlock = ({ value }: { value: string }) => {
+  const theme = useEuiTheme();
+
+  return (
+    <EuiPanel
+      hasShadow={false}
+      hasBorder={false}
+      paddingSize="s"
+      className={css`
+        background-color: ${theme.euiTheme.colors.lightestShade};
+        .euiCodeBlock__pre {
+          margin-bottom: 0;
+          padding: ${theme.euiTheme.size.m};
+          min-block-size: 48px;
+        }
+        .euiCodeBlock__controls {
+          inset-block-start: ${theme.euiTheme.size.m};
+          inset-inline-end: ${theme.euiTheme.size.m};
+        }
+      `}
+    >
+      <EuiFlexGroup direction="column" gutterSize="xs">
+        <EuiFlexItem grow={false}>
+          <EuiCodeBlock isCopyable fontSize="m">
+            {value}
+          </EuiCodeBlock>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </EuiPanel>
+  );
+};
--- a/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/get_comments/custom_codeblock/custom_codeblock_markdown_plugin.tsx
@ -0,0 +1,35 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Node } from 'unist';
+import type { Parent } from 'mdast';
+
+export const customCodeBlockLanguagePlugin = () => {
+  const visitor = (node: Node, parent?: Parent) => {
+    if ('children' in node) {
+      const nodeAsParent = node as Parent;
+      nodeAsParent.children.forEach((child) => {
+        visitor(child, nodeAsParent);
+      });
+    }
+
+    if (
+      node.type === 'code' &&
+      (node.lang === 'eql' ||
+        node.lang === 'esql' ||
+        node.lang === 'kql' ||
+        node.lang === 'dsl' ||
+        node.lang === 'json')
+    ) {
+      node.type = 'customCodeBlock';
+    }
+  };
+
+  return (tree: Node) => {
+    visitor(tree);
+  };
+};
--- a/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/get_comments/index.tsx
@ -7,7 +7,15 @@

 import type { EuiCommentProps } from '@elastic/eui';
 import type { Conversation } from '@kbn/elastic-assistant';
-import { EuiAvatar, EuiMarkdownFormat, EuiText, tint } from '@elastic/eui';
+import {
+  EuiAvatar,
+  EuiMarkdownFormat,
+  EuiSpacer,
+  EuiText,
+  getDefaultEuiMarkdownParsingPlugins,
+  getDefaultEuiMarkdownProcessingPlugins,
+  tint,
+} from '@elastic/eui';
 import React from 'react';

 import { AssistantAvatar } from '@kbn/elastic-assistant';
@ -15,6 +23,8 @@ import { css } from '@emotion/react';
 import { euiThemeVars } from '@kbn/ui-theme';
 import { CommentActions } from '../comment_actions';
 import * as i18n from './translations';
+import { customCodeBlockLanguagePlugin } from './custom_codeblock/custom_codeblock_markdown_plugin';
+import { CustomCodeBlock } from './custom_codeblock/custom_code_block';

 export const getComments = ({
  currentConversation,
@ -24,8 +34,29 @@ export const getComments = ({
  currentConversation: Conversation;
  lastCommentRef: React.MutableRefObject<HTMLDivElement | null>;
  showAnonymizedValues: boolean;
-}): EuiCommentProps[] =>
-  currentConversation.messages.map((message, index) => {
+}): EuiCommentProps[] => {
+  const parsingPlugins = getDefaultEuiMarkdownParsingPlugins();
+  const processingPlugins = getDefaultEuiMarkdownProcessingPlugins();
+
+  const { components } = processingPlugins[1][1];
+
+  processingPlugins[1][1].components = {
+    ...components,
+    customCodeBlock: (props) => {
+      return (
+        <>
+          <CustomCodeBlock value={props.value} />
+          <EuiSpacer size="m" />
+        </>
+      );
+    },
+  };
+
+  // Fun fact: must spread existing parsingPlugins last
+  const parsingPluginList = [customCodeBlockLanguagePlugin, ...parsingPlugins];
+  const processingPluginList = processingPlugins;
+
+  return currentConversation.messages.map((message, index) => {
    const isUser = message.role === 'user';
    const replacements = currentConversation.replacements;
    const messageContentWithReplacements =
@ -45,13 +76,21 @@ export const getComments = ({
      children:
        index !== currentConversation.messages.length - 1 ? (
          <EuiText>
-            <EuiMarkdownFormat className={`message-${index}`}>
+            <EuiMarkdownFormat
+              className={`message-${index}`}
+              parsingPluginList={parsingPluginList}
+              processingPluginList={processingPluginList}
+            >
              {showAnonymizedValues ? message.content : transformedMessage.content}
            </EuiMarkdownFormat>
          </EuiText>
        ) : (
          <EuiText>
-            <EuiMarkdownFormat className={`message-${index}`}>
+            <EuiMarkdownFormat
+              className={`message-${index}`}
+              parsingPluginList={parsingPluginList}
+              processingPluginList={processingPluginList}
+            >
              {showAnonymizedValues ? message.content : transformedMessage.content}
            </EuiMarkdownFormat>
            <span ref={lastCommentRef} />
@ -82,3 +121,4 @@ export const getComments = ({
        : {}),
    };
  });
+};
--- a/x-pack/plugins/security_solution/public/assistant/helpers.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/helpers.tsx
@ -49,7 +49,13 @@ export const getPromptContextFromEventDetailsItem = (data: TimelineEventsDetails
  return getFieldsAsCsv(allFields);
 };

-const sendToTimelineEligibleQueryTypes: Array<CodeBlockDetails['type']> = ['kql', 'dsl', 'eql'];
+const sendToTimelineEligibleQueryTypes: Array<CodeBlockDetails['type']> = [
+  'kql',
+  'dsl',
+  'eql',
+  'esql',
+  'sql', // Models often put the code block language as sql, for esql, so adding this as a fallback
+];

 /**
 * Returns message contents with replacements applied.
--- a/x-pack/plugins/security_solution/public/assistant/send_to_timeline/index.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/send_to_timeline/index.tsx
@ -26,9 +26,11 @@ import {
  applyKqlFilterQuery,
  setActiveTabTimeline,
  setFilters,
+  showTimeline,
  updateDataView,
  updateEqlOptions,
 } from '../../timelines/store/timeline/actions';
+import { useDiscoverInTimelineContext } from '../../common/components/discover_in_timeline/use_discover_in_timeline_context';

 export interface SendToTimelineButtonProps {
  asEmptyButton: boolean;
@ -50,6 +52,8 @@ export const SendToTimelineButton: React.FunctionComponent<SendToTimelineButtonP
 }) => {
  const dispatch = useDispatch();

+  const { discoverStateContainer } = useDiscoverInTimelineContext();
+
  const getDataViewsSelector = useMemo(
    () => sourcererSelectors.getSourcererDataViewsSelector(),
    []
@ -68,6 +72,30 @@ export const SendToTimelineButton: React.FunctionComponent<SendToTimelineButtonP

  const configureAndOpenTimeline = useCallback(() => {
    if (dataProviders || filters) {
+      // If esql, don't reset filters or mess with dataview & time range
+      if (dataProviders?.[0]?.queryType === 'esql' || dataProviders?.[0]?.queryType === 'sql') {
+        discoverStateContainer.current?.appState.update({
+          query: {
+            query: dataProviders[0].kqlQuery,
+            language: 'esql',
+          },
+        });
+
+        dispatch(
+          setActiveTabTimeline({
+            id: TimelineId.active,
+            activeTab: TimelineTabs.esql,
+          })
+        );
+        dispatch(
+          showTimeline({
+            id: TimelineId.active,
+            show: true,
+          })
+        );
+        return;
+      }
+
      // Reset the current timeline
      if (timeRange) {
        clearTimeline({
@ -147,6 +175,7 @@ export const SendToTimelineButton: React.FunctionComponent<SendToTimelineButtonP
            break;
        }
      }
+
      // Use filters if more than a certain amount of ids for dom performance.
      if (filters) {
        dispatch(
@ -172,13 +201,14 @@ export const SendToTimelineButton: React.FunctionComponent<SendToTimelineButtonP
    }
  }, [
    dataProviders,
-    clearTimeline,
-    dispatch,
-    defaultDataView.id,
-    signalIndexName,
    filters,
    timeRange,
    keepDataView,
+    dispatch,
+    clearTimeline,
+    discoverStateContainer,
+    defaultDataView.id,
+    signalIndexName,
  ]);

  return asEmptyButton ? (