[Security Solution][Detection Engine] update query automatically in rule create form through AI assistant

x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx

@@ -263,7 +263,7 @@ const AssistantComponent: React.FC<Props> = ({
   // Add min-height to all codeblocks so timeline icon doesn't overflow
   const codeBlockContainers = [...document.getElementsByClassName('euiCodeBlock')];
   // @ts-ignore-expect-error
-  codeBlockContainers.forEach((e) => (e.style.minHeight = '75px'));
+  codeBlockContainers.forEach((e) => (e.style.minHeight = '85px'));
   ////
 
   const onToggleShowAnonymizedValues = useCallback(() => {

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx

@@ -8,7 +8,7 @@
 import { EuiCommentProps } from '@elastic/eui';
 import type { HttpSetup } from '@kbn/core-http-browser';
 import { omit } from 'lodash/fp';
-import React, { useCallback, useMemo, useState } from 'react';
+import React, { useCallback, useMemo, useState, useRef } from 'react';
 import type { IToasts } from '@kbn/core-notifications-browser';
 import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
 import { useLocalStorage, useSessionStorage } from 'react-use';
@@ -137,6 +137,7 @@ export interface UseAssistantContext {
   basePromptContexts: PromptContextTemplate[];
   unRegisterPromptContext: UnRegisterPromptContext;
   currentAppId: string;
+  codeBlockRef: React.MutableRefObject<(codeBlock: string) => void>;
 }
 
 const AssistantContext = React.createContext<UseAssistantContext | undefined>(undefined);
@@ -237,6 +238,11 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
    */
   const [selectedSettingsTab, setSelectedSettingsTab] = useState<SettingsTabs | null>(null);
 
+  /**
+   * Setting code block ref that can be used to store callback from parent components
+   */
+  const codeBlockRef = useRef(() => {});
+
   const getLastConversationId = useCallback(
     // if a conversationId has been provided, use that
     // if not, check local storage
@@ -284,6 +290,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       setLastConversationId: setLocalStorageLastConversationId,
       baseConversations,
       currentAppId,
+      codeBlockRef,
     }),
     [
       actionTypeRegistry,
@@ -316,6 +323,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
       setLocalStorageLastConversationId,
       baseConversations,
       currentAppId,
+      codeBlockRef,
     ]
   );
 

x-pack/packages/kbn-elastic-assistant/impl/new_chat/index.test.tsx

@@ -18,6 +18,11 @@ jest.mock('../assistant/use_assistant_overlay', () => ({
   useAssistantOverlay: () => mockUseAssistantOverlay,
 }));
 
+let mockUseAssistantContext = { codeBlockRef: { current: null } };
+jest.mock('../..', () => ({
+  useAssistantContext: () => mockUseAssistantContext,
+}));
+
 const defaultProps: Props = {
   category: 'alert',
   description: 'Test description',
@@ -27,6 +32,9 @@ const defaultProps: Props = {
 };
 
 describe('NewChat', () => {
+  beforeEach(() => {
+    mockUseAssistantContext = { codeBlockRef: { current: null } };
+  });
   afterEach(() => {
     jest.clearAllMocks();
   });
@@ -118,4 +126,17 @@ describe('NewChat', () => {
 
     expect(onShowOverlaySpy).toHaveBeenCalled();
   });
+
+  it('assigns onExportCodeBlock callback to context codeBlock reference', () => {
+    const onExportCodeBlock = jest.fn();
+    render(<NewChat {...defaultProps} onExportCodeBlock={onExportCodeBlock} />);
+
+    expect(mockUseAssistantContext.codeBlockRef.current).toBe(onExportCodeBlock);
+  });
+
+  it('does not change assigns context codeBlock reference if onExportCodeBlock not defined', () => {
+    render(<NewChat {...defaultProps} />);
+
+    expect(mockUseAssistantContext.codeBlockRef.current).toBe(null);
+  });
 });

x-pack/packages/kbn-elastic-assistant/impl/new_chat/index.tsx

@@ -6,7 +6,8 @@
  */
 
 import { EuiButtonEmpty, EuiLink } from '@elastic/eui';
-import React, { useCallback, useMemo } from 'react';
+import React, { useCallback, useMemo, useEffect } from 'react';
+import { useAssistantContext } from '../..';
 
 import { PromptContext } from '../assistant/prompt_context/types';
 import { useAssistantOverlay } from '../assistant/use_assistant_overlay';
@@ -29,6 +30,8 @@ export type Props = Omit<PromptContext, 'id'> & {
   asLink?: boolean;
   /** Optional callback when overlay shows */
   onShowOverlay?: () => void;
+  /** Optional callback that returns copied code block */
+  onExportCodeBlock?: (codeBlock: string) => void;
 };
 
 const NewChatComponent: React.FC<Props> = ({
@@ -45,6 +48,7 @@ const NewChatComponent: React.FC<Props> = ({
   isAssistantEnabled,
   asLink = false,
   onShowOverlay,
+  onExportCodeBlock,
 }) => {
   const { showAssistantOverlay } = useAssistantOverlay(
     category,
@@ -56,12 +60,25 @@ const NewChatComponent: React.FC<Props> = ({
     tooltip,
     isAssistantEnabled
   );
+  const { codeBlockRef } = useAssistantContext();
 
   const showOverlay = useCallback(() => {
     showAssistantOverlay(true);
     onShowOverlay?.();
   }, [showAssistantOverlay, onShowOverlay]);
 
+  useEffect(() => {
+    if (onExportCodeBlock) {
+      codeBlockRef.current = onExportCodeBlock;
+    }
+
+    return () => {
+      if (onExportCodeBlock) {
+        codeBlockRef.current = () => {};
+      }
+    };
+  }, [codeBlockRef, onExportCodeBlock]);
+
   const icon = useMemo(() => {
     if (iconType === null) {
       return undefined;

x-pack/plugins/security_solution/public/assistant/content/conversations/index.tsx

@@ -10,6 +10,7 @@ import type { Conversation } from '@kbn/elastic-assistant';
 import { DATA_QUALITY_DASHBOARD_CONVERSATION_ID } from '@kbn/ecs-data-quality-dashboard';
 
 import { DETECTION_RULES_CONVERSATION_ID } from '../../../detections/pages/detection_engine/rules/translations';
+import { DETECTION_RULES_CREATE_FORM_CONVERSATION_ID } from '../../../detections/pages/detection_engine/translations';
 import {
   ALERT_SUMMARY_CONVERSATION_ID,
   EVENT_SUMMARY_CONVERSATION_ID,
@@ -41,6 +42,14 @@ export const BASE_SECURITY_CONVERSATIONS: Record<string, Conversation> = {
     messages: [],
     replacements: {},
   },
+  [DETECTION_RULES_CREATE_FORM_CONVERSATION_ID]: {
+    id: '',
+    title: DETECTION_RULES_CREATE_FORM_CONVERSATION_ID,
+    category: 'assistant',
+    isDefault: true,
+    messages: [],
+    replacements: {},
+  },
   [EVENT_SUMMARY_CONVERSATION_ID]: {
     id: '',
     title: EVENT_SUMMARY_CONVERSATION_ID,

x-pack/plugins/security_solution/public/assistant/helpers.tsx

@@ -13,9 +13,9 @@ import { replaceAnonymizedValuesWithOriginalValues } from '@kbn/elastic-assistan
 import type { TimelineEventsDetailsItem } from '../../common/search_strategy';
 import type { Rule } from '../detection_engine/rule_management/logic';
 import { SendToTimelineButton } from './send_to_timeline';
-
+import { DETECTION_RULES_CREATE_FORM_CONVERSATION_ID } from '../detections/pages/detection_engine/translations';
 export const LOCAL_STORAGE_KEY = `securityAssistant`;
-
+import { UpdateQueryInFormButton } from './update_query_in_form';
 export interface QueryField {
   field: string;
   values: string;
@@ -84,30 +84,37 @@ export const augmentMessageCodeBlocks = (
           document.querySelectorAll(`.message-${messageIndex} .euiCodeBlock__controls`)[
             codeBlockIndex
           ],
-        button: sendToTimelineEligibleQueryTypes.includes(codeBlock.type) ? (
-          <SendToTimelineButton
-            asEmptyButton={true}
-            dataProviders={[
-              {
-                id: 'assistant-data-provider',
-                name: `Assistant Query from conversation ${currentConversation.id}`,
-                enabled: true,
-                excluded: false,
-                queryType: codeBlock.type,
-                kqlQuery: codeBlock.content ?? '',
-                queryMatch: {
-                  field: 'host.name',
-                  operator: ':',
-                  value: 'test',
-                },
-                and: [],
-              },
-            ]}
-            keepDataView={true}
-          >
-            <EuiIcon type="timeline" />
-          </SendToTimelineButton>
-        ) : null,
+        button: (
+          <>
+            {sendToTimelineEligibleQueryTypes.includes(codeBlock.type) ? (
+              <SendToTimelineButton
+                asEmptyButton={true}
+                dataProviders={[
+                  {
+                    id: 'assistant-data-provider',
+                    name: `Assistant Query from conversation ${currentConversation.id}`,
+                    enabled: true,
+                    excluded: false,
+                    queryType: codeBlock.type,
+                    kqlQuery: codeBlock.content ?? '',
+                    queryMatch: {
+                      field: 'host.name',
+                      operator: ':',
+                      value: 'test',
+                    },
+                    and: [],
+                  },
+                ]}
+                keepDataView={true}
+              >
+                <EuiIcon type="timeline" />
+              </SendToTimelineButton>
+            ) : null}
+            {DETECTION_RULES_CREATE_FORM_CONVERSATION_ID === currentConversation.title ? (
+              <UpdateQueryInFormButton query={codeBlock.content ?? ''} />
+            ) : null}
+          </>
+        ),
       };
     })
   );

x-pack/plugins/security_solution/public/assistant/update_query_in_form/index.test.tsx

@@ -0,0 +1,32 @@
+/*
+ * 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 React from 'react';
+import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+
+import { UpdateQueryInFormButton } from '.';
+
+const mockUseAssistantContext = { codeBlockRef: { current: jest.fn() } };
+jest.mock('@kbn/elastic-assistant', () => ({
+  useAssistantContext: () => mockUseAssistantContext,
+}));
+
+describe('UpdateQueryInFormButton', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('calls codeBlockRef callback on click', () => {
+    const testQuery = 'from auditbeat* | limit 10';
+    render(<UpdateQueryInFormButton query={testQuery} />);
+
+    userEvent.click(screen.getByTestId('update-query-in-form-button'));
+
+    expect(mockUseAssistantContext.codeBlockRef.current).toHaveBeenCalledWith(testQuery);
+  });
+});

x-pack/plugins/security_solution/public/assistant/update_query_in_form/index.tsx

@@ -0,0 +1,44 @@
+/*
+ * 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 { FC, PropsWithChildren } from 'react';
+import React from 'react';
+import { EuiButtonEmpty, EuiToolTip, EuiIcon } from '@elastic/eui';
+import { useAssistantContext } from '@kbn/elastic-assistant';
+
+import { UPDATE_QUERY_IN_FORM_TOOLTIP } from './translations';
+
+export interface UpdateQueryInFormButtonProps {
+  query: string;
+}
+
+export const UpdateQueryInFormButton: FC<PropsWithChildren<UpdateQueryInFormButtonProps>> = ({
+  query,
+}) => {
+  const { codeBlockRef } = useAssistantContext();
+
+  const handleClick = () => {
+    codeBlockRef?.current?.(query);
+  };
+
+  return (
+    <EuiButtonEmpty
+      data-test-subj="update-query-in-form-button"
+      aria-label={UPDATE_QUERY_IN_FORM_TOOLTIP}
+      onClick={handleClick}
+      color="text"
+      flush="both"
+      size="xs"
+    >
+      <EuiToolTip position="right" content={UPDATE_QUERY_IN_FORM_TOOLTIP}>
+        <EuiIcon type="documentEdit" />
+      </EuiToolTip>
+    </EuiButtonEmpty>
+  );
+};
+
+UpdateQueryInFormButton.displayName = 'UpdateQueryInFormButton';

x-pack/plugins/security_solution/public/assistant/update_query_in_form/translations.ts

@@ -0,0 +1,15 @@
+/*
+ * 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 { i18n } from '@kbn/i18n';
+
+export const UPDATE_QUERY_IN_FORM_TOOLTIP = i18n.translate(
+  'xpack.securitySolution.assistant.updateQueryInFormTooltip',
+  {
+    defaultMessage: 'Update query in form',
+  }
+);

x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/ai_assistant/index.test.tsx

@@ -26,14 +26,14 @@ describe('AiAssistant', () => {
   it('does not render chat component when does not have hasAssistantPrivilege', () => {
     useAssistantAvailabilityMock.mockReturnValue({ hasAssistantPrivilege: false });
 
-    const { container } = render(<AiAssistant getFields={jest.fn()} />, {
+    const { container } = render(<AiAssistant getFields={jest.fn()} setFieldValue={jest.fn()} />, {
       wrapper: TestProviders,
     });
 
     expect(container).toBeEmptyDOMElement();
   });
   it('renders chat component when has hasAssistantPrivilege', () => {
-    render(<AiAssistant getFields={jest.fn()} />, {
+    render(<AiAssistant getFields={jest.fn()} setFieldValue={jest.fn()} />, {
       wrapper: TestProviders,
     });