[Security Solution] Security Assistant: Data Quality dashboard _New C…

…hat_ button & API updates (elastic#16)

### Summary

- adds a _New chat_ button to the Data Quality dashboard
- updates the `useAssistantOverlay` API
- `NewChat` takes an optional `children` prop
- added `NewChatById` component
- improves test coverage

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

@@ -384,6 +384,7 @@ const AssistantComponent: React.FC<Props> = ({
                       clearConversation(selectedConversationId);
                       setSelectedSystemPromptId(getDefaultSystemPrompt().id);
                       setSelectedPromptContextIds([]);
+                      setSuggestedUserPrompt('');
                     }}
                   />
                 </EuiToolTip>

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

@@ -0,0 +1,102 @@
+/*
+ * 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 { act, renderHook } from '@testing-library/react-hooks';
+
+import { useAssistantOverlay } from '.';
+
+const mockUseAssistantContext = {
+  registerPromptContext: jest.fn(),
+  showAssistantOverlay: jest.fn(),
+  unRegisterPromptContext: jest.fn(),
+};
+jest.mock('../../assistant_context', () => {
+  const original = jest.requireActual('../../assistant_context');
+
+  return {
+    ...original,
+    useAssistantContext: () => mockUseAssistantContext,
+  };
+});
+
+describe('useAssistantOverlay', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('calls registerPromptContext with the expected context', async () => {
+    const category = 'event';
+    const description = 'test description';
+    const getPromptContext = jest.fn(() => Promise.resolve('test data'));
+    const id = 'test-id';
+    const suggestedUserPrompt = 'test user prompt';
+    const tooltip = 'test tooltip';
+
+    renderHook(() =>
+      useAssistantOverlay(
+        category,
+        null,
+        description,
+        getPromptContext,
+        id,
+        suggestedUserPrompt,
+        tooltip
+      )
+    );
+
+    expect(mockUseAssistantContext.registerPromptContext).toHaveBeenCalledWith({
+      category,
+      description,
+      getPromptContext,
+      id,
+      suggestedUserPrompt,
+      tooltip,
+    });
+  });
+
+  it('calls unRegisterPromptContext on unmount', () => {
+    const { unmount } = renderHook(() =>
+      useAssistantOverlay(
+        'event',
+        null,
+        'description',
+        () => Promise.resolve('data'),
+        'id',
+        null,
+        'tooltip'
+      )
+    );
+
+    unmount();
+
+    expect(mockUseAssistantContext.unRegisterPromptContext).toHaveBeenCalledWith('id');
+  });
+
+  it('calls `showAssistantOverlay` from the assistant context', () => {
+    const { result } = renderHook(() =>
+      useAssistantOverlay(
+        'event',
+        'conversation-id',
+        'description',
+        () => Promise.resolve('data'),
+        'id',
+        null,
+        'tooltip'
+      )
+    );
+
+    act(() => {
+      result.current.showAssistantOverlay(true);
+    });
+
+    expect(mockUseAssistantContext.showAssistantOverlay).toHaveBeenCalledWith({
+      showOverlay: true,
+      promptContextId: 'id',
+      conversationId: 'conversation-id',
+    });
+  });
+});

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

@@ -5,41 +5,84 @@
  * 2.0.
  */
 
-import { useCallback, useEffect, useMemo, useRef } from 'react';
+import { useCallback, useEffect, useMemo } from 'react';
 
 import { useAssistantContext } from '../../assistant_context';
 import { getUniquePromptContextId } from '../../assistant_context/helpers';
 import type { PromptContext } from '../prompt_context/types';
 
-interface Props {
-  promptContext?: Omit<PromptContext, 'id'>;
-  promptContextId?: string;
-  conversationId?: string;
-}
 interface UseAssistantOverlay {
   showAssistantOverlay: (show: boolean) => void;
   promptContextId: string;
 }
 
-export const useAssistantOverlay = ({
-  conversationId,
-  promptContext,
-  promptContextId,
-}: Props): UseAssistantOverlay => {
-  // create a unique prompt context id if one is not provided:
-  const _promptContextId = useMemo(
-    () => promptContextId ?? getUniquePromptContextId(),
-    [promptContextId]
-  );
+/**
+ * `useAssistantOverlay` is a hook that registers context with the assistant overlay, and
+ * returns an optional `showAssistantOverlay` function to display the assistant overlay.
+ * As an alterative to using the `showAssistantOverlay` returned from this hook, you may
+ * use the `NewChatById` component and pass it the `promptContextId` returned by this hook.
+ *
+ * USE THIS WHEN: You want to register context in one part of the tree, and then show
+ * a _New chat_ button in another part of the tree without passing around the data, or when
+ * you want to build a custom `New chat` button with features not not provided by the
+ * `NewChat` component.
+ */
+export const useAssistantOverlay = (
+  /**
+   * The category of data, e.g. `alert | alerts | event | events | string`
+   *
+   * `category` helps the assistant display the most relevant user prompts
+   */
+  category: PromptContext['category'],
+
+  /**
+   * optionally automatically add this context to a specific conversation when the assistant is displayed
+   */
+  conversationId: string | null,
+
+  /**
+   * The assistant will display this **short**, static description
+   * in the context pill
+   */
+  description: PromptContext['description'],
+
+  /**
+   * The assistant will invoke this function to retrieve the context data,
+   * which will be included in a prompt (e.g. the contents of an alert or an event)
+   */
+  getPromptContext: PromptContext['getPromptContext'],
+
+  /**
+   * Optionally provide a unique identifier for this prompt context, or accept the uuid default.
+   */
+  id: PromptContext['id'] | null,
+
+  /**
+   * An optional user prompt that's filled in, but not sent, when the Elastic Assistant opens
+   */
+  suggestedUserPrompt: PromptContext['suggestedUserPrompt'] | null,
 
-  const _promptContextRef = useRef<PromptContext | undefined>(
-    promptContext != null
-      ? {
-          ...promptContext,
-          id: _promptContextId,
-        }
-      : undefined
+  /**
+   * The assistant will display this tooltip when the user hovers over the context pill
+   */
+  tooltip: PromptContext['tooltip']
+): UseAssistantOverlay => {
+  // memoize the props so that we can use them in the effect below:
+  const _category: PromptContext['category'] = useMemo(() => category, [category]);
+  const _description: PromptContext['description'] = useMemo(() => description, [description]);
+  const _getPromptContext: PromptContext['getPromptContext'] = useMemo(
+    () => getPromptContext,
+    [getPromptContext]
   );
+  const promptContextId: PromptContext['id'] = useMemo(
+    () => id ?? getUniquePromptContextId(),
+    [id]
+  );
+  const _suggestedUserPrompt: PromptContext['suggestedUserPrompt'] = useMemo(
+    () => suggestedUserPrompt ?? undefined,
+    [suggestedUserPrompt]
+  );
+  const _tooltip = useMemo(() => tooltip, [tooltip]);
 
   // the assistant context is used to show/hide the assistant overlay:
   const {
@@ -51,31 +94,42 @@ export const useAssistantOverlay = ({
   // proxy show / hide calls to assistant context, using our internal prompt context id:
   const showAssistantOverlay = useCallback(
     (showOverlay: boolean) => {
-      assistantContextShowOverlay({
-        showOverlay,
-        promptContextId: _promptContextId,
-        conversationId,
-      });
+      if (promptContextId != null) {
+        assistantContextShowOverlay({
+          showOverlay,
+          promptContextId,
+          conversationId: conversationId ?? undefined,
+        });
+      }
     },
-    [assistantContextShowOverlay, _promptContextId, conversationId]
+    [assistantContextShowOverlay, conversationId, promptContextId]
   );
 
-  useEffect(
-    () => () => unRegisterPromptContext(_promptContextId),
-    [_promptContextId, unRegisterPromptContext]
-  );
+  useEffect(() => {
+    unRegisterPromptContext(promptContextId); // a noop if the current prompt context id is not registered
+
+    const newContext: PromptContext = {
+      category: _category,
+      description: _description,
+      getPromptContext: _getPromptContext,
+      id: promptContextId,
+      suggestedUserPrompt: _suggestedUserPrompt,
+      tooltip: _tooltip,
+    };
+
+    registerPromptContext(newContext);
+
+    return () => unRegisterPromptContext(promptContextId);
+  }, [
+    _category,
+    _description,
+    _getPromptContext,
+    _suggestedUserPrompt,
+    _tooltip,
+    promptContextId,
+    registerPromptContext,
+    unRegisterPromptContext,
+  ]);
 
-  if (
-    promptContext != null &&
-    (_promptContextRef.current?.category !== promptContext?.category ||
-      _promptContextRef.current?.description !== promptContext?.description ||
-      _promptContextRef.current?.getPromptContext !== promptContext?.getPromptContext ||
-      _promptContextRef.current?.suggestedUserPrompt !== promptContext?.suggestedUserPrompt ||
-      _promptContextRef.current?.tooltip !== promptContext?.tooltip)
-  ) {
-    _promptContextRef.current = { ...promptContext, id: _promptContextId };
-    registerPromptContext(_promptContextRef.current);
-  }
-
-  return { promptContextId: _promptContextId, showAssistantOverlay };
+  return { promptContextId, showAssistantOverlay };
 };

x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts

@@ -19,6 +19,9 @@ export interface CodeBlockDetails {
 export type QueryType = 'eql' | 'kql' | 'dsl' | 'json' | 'no-type';
 
 /**
+ * `analyzeMarkdown` is a helper that enriches content returned from a query
+ * with action buttons
+ *
  * Returns a list of code block details for each code block in the markdown,
  * including the type of code block and the content of the code block.
  *

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

@@ -91,7 +91,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
   const registerPromptContext: RegisterPromptContext = useCallback(
     (promptContext: PromptContext) => {
       setPromptContexts((prevPromptContexts) => {
-        if (prevPromptContexts[promptContext.id] == null) {
+        if (promptContext != null && prevPromptContexts[promptContext.id] == null) {
           return updatePromptContexts({
             prevPromptContexts,
             promptContext,