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

## Summary Fixes `ES|QL` codeblocks from not being able to be sent to Timeline. <p align="center"> <img width="500" src="fdc3b2a0-5b4b-4584-b304-c4d24de1917c" /> </p> ## Test instructions Either request the assistant to generate an ESQL query or just paste this codeblock into the conversation to test the action directly from the user message. Be sure to declare the codeblock language as `esql` or include one of the [string match patterns](https://github.com/elastic/kibana/pull/169478/files#diff-f70f0b96568e024e53bfbb62adcca72051f0a2e824d4ab22664eed0e149be248R38) above the code block so the action can be recognized. ```` Below is an `Elasticsearch Query Language` query: ```esql FROM logs-endpoint* | WHERE event.category == \"process\" | STATS proc_count = COUNT(process.name) BY host.name | KEEP host.name, proc_count ``` ```` Note: The `send to timeline` actions appear to only reliably show up when using the Assistant instance within Timeline. Now that we have a more reliable way of attaching actions via markdown plugins/parsers, we should refactor this code to use that method as opposed to the code block/dom inspection route that is used currently. In the meantime I will see if there is a low-impact fix that can be made here.
2025-04-23 01:13:23 -04:00 · 2023-10-23 16:01:16 -06:00 · 2023-10-23 16:01:16 -06:00 · 93636d9fc0
commit 93636d9fc0
parent f5d7c86cc4
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 ? (