[Security Assistant] Fix system prompts (#189130)

x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_settings/conversation_settings_editor.tsx

@@ -13,6 +13,7 @@ import { FormattedMessage } from '@kbn/i18n-react';
 import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/public/common';
 import { noop } from 'lodash/fp';
 import { PromptResponse } from '@kbn/elastic-assistant-common';
+import { QueryObserverResult } from '@tanstack/react-query';
 import { Conversation } from '../../../..';
 import * as i18n from './translations';
 import * as i18nModel from '../../../connectorland/models/model_selector/translations';
@@ -37,6 +38,7 @@ export interface ConversationSettingsEditorProps {
     React.SetStateAction<ConversationsBulkActions>
   >;
   onSelectedConversationChange: (conversation?: Conversation) => void;
+  refetchConversations?: () => Promise<QueryObserverResult<Record<string, Conversation>, unknown>>;
 }
 
 /**
@@ -53,6 +55,7 @@ export const ConversationSettingsEditor: React.FC<ConversationSettingsEditorProp
     conversationsSettingsBulkActions,
     setConversationsSettingsBulkActions,
     onSelectedConversationChange,
+    refetchConversations,
   }) => {
     const { data: connectors, isSuccess: areConnectorsFetched } = useLoadConnectors({
       http,
@@ -276,6 +279,7 @@ export const ConversationSettingsEditor: React.FC<ConversationSettingsEditorProp
             conversation={selectedConversation}
             isDisabled={isDisabled}
             onSystemPromptSelectionChange={handleOnSystemPromptSelectionChange}
+            refetchConversations={refetchConversations}
             selectedPrompt={selectedSystemPrompt}
             isSettingsModalVisible={true}
             setIsSettingsModalVisible={noop} // noop, already in settings

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

@@ -321,6 +321,7 @@ const ConversationSettingsManagementComponent: React.FC<Props> = ({
             conversationsSettingsBulkActions={conversationsSettingsBulkActions}
             http={http}
             isDisabled={isDisabled}
+            refetchConversations={refetchConversations}
             selectedConversation={selectedConversation}
             setConversationSettings={setConversationSettings}
             setConversationsSettingsBulkActions={setConversationsSettingsBulkActions}

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

@@ -256,6 +256,13 @@ const AssistantComponent: React.FC<Props> = ({
           conversations[WELCOME_CONVERSATION_TITLE] ??
           getDefaultConversation({ cTitle: WELCOME_CONVERSATION_TITLE });
 
+        // updated selected system prompt
+        setEditingSystemPromptId(
+          getDefaultSystemPrompt({
+            allSystemPrompts,
+            conversation: conversationToReturn,
+          })?.id
+        );
         if (
           prev &&
           prev.id === conversationToReturn.id &&
@@ -273,6 +280,7 @@ const AssistantComponent: React.FC<Props> = ({
       });
     }
   }, [
+    allSystemPrompts,
     areConnectorsFetched,
     conversationTitle,
     conversations,
@@ -647,6 +655,7 @@ const AssistantComponent: React.FC<Props> = ({
             actionTypeId: (defaultConnector?.actionTypeId as string) ?? '.gen-ai',
             provider: apiConfig?.apiProvider,
             model: apiConfig?.defaultModel,
+            defaultSystemPromptId: allSystemPrompts.find((sp) => sp.isNewConversationDefault)?.id,
           },
         });
       },
@@ -665,14 +674,14 @@ const AssistantComponent: React.FC<Props> = ({
 
   useEffect(() => {
     (async () => {
-      if (areConnectorsFetched && currentConversation?.id === '') {
+      if (areConnectorsFetched && currentConversation?.id === '' && !isLoadingPrompts) {
         const conversation = await mutateAsync(currentConversation);
         if (currentConversation.id === '' && conversation) {
           setCurrentConversationId(conversation.id);
         }
       }
     })();
-  }, [areConnectorsFetched, currentConversation, mutateAsync]);
+  }, [areConnectorsFetched, currentConversation, isLoadingPrompts, mutateAsync]);
 
   const handleCreateConversation = useCallback(async () => {
     const newChatExists = find(conversations, ['title', NEW_CHAT]);
@@ -791,6 +800,7 @@ const AssistantComponent: React.FC<Props> = ({
                     isSettingsModalVisible={isSettingsModalVisible}
                     setIsSettingsModalVisible={setIsSettingsModalVisible}
                     allSystemPrompts={allSystemPrompts}
+                    refetchConversations={refetchResults}
                   />
                 </EuiFlexItem>
                 <EuiFlexItem grow={false}>
@@ -823,6 +833,7 @@ const AssistantComponent: React.FC<Props> = ({
     handleOnSystemPromptSelectionChange,
     isSettingsModalVisible,
     isWelcomeSetup,
+    refetchResults,
   ]);
 
   return (

x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt_editor/system_prompt/helpers.test.tsx

@@ -16,13 +16,13 @@ import { getOptions, getOptionFromPrompt } from './helpers';
 describe('helpers', () => {
   describe('getOptionFromPrompt', () => {
     it('returns an EuiSuperSelectOption with the correct value', () => {
-      const option = getOptionFromPrompt({ ...mockSystemPrompt });
+      const option = getOptionFromPrompt({ ...mockSystemPrompt, isCleared: false });
 
       expect(option.value).toBe(mockSystemPrompt.id);
     });
 
     it('returns an EuiSuperSelectOption with the correct inputDisplay', () => {
-      const option = getOptionFromPrompt({ ...mockSystemPrompt });
+      const option = getOptionFromPrompt({ ...mockSystemPrompt, isCleared: false });
 
       render(<>{option.inputDisplay}</>);
 
@@ -30,7 +30,7 @@ describe('helpers', () => {
     });
 
     it('shows the expected name in the dropdownDisplay', () => {
-      const option = getOptionFromPrompt({ ...mockSystemPrompt });
+      const option = getOptionFromPrompt({ ...mockSystemPrompt, isCleared: false });
 
       render(<TestProviders>{option.dropdownDisplay}</TestProviders>);
 
@@ -38,7 +38,7 @@ describe('helpers', () => {
     });
 
     it('shows the expected prompt content in the dropdownDisplay', () => {
-      const option = getOptionFromPrompt({ ...mockSystemPrompt });
+      const option = getOptionFromPrompt({ ...mockSystemPrompt, isCleared: false });
 
       render(<TestProviders>{option.dropdownDisplay}</TestProviders>);
 
@@ -51,7 +51,7 @@ describe('helpers', () => {
       const prompts = [mockSystemPrompt, mockSuperheroSystemPrompt];
       const promptIds = prompts.map(({ id }) => id);
 
-      const options = getOptions({ prompts });
+      const options = getOptions({ prompts, isCleared: false });
       const optionValues = options.map(({ value }) => value);
 
       expect(optionValues).toEqual(promptIds);

x-pack/packages/kbn-elastic-assistant/impl/assistant/prompt_editor/system_prompt/helpers.tsx

@@ -12,19 +12,38 @@ import styled from '@emotion/styled';
 import { isEmpty } from 'lodash/fp';
 import { euiThemeVars } from '@kbn/ui-theme';
 import { PromptResponse } from '@kbn/elastic-assistant-common';
+import { css } from '@emotion/react';
 import { EMPTY_PROMPT } from './translations';
 
 const Strong = styled.strong`
   margin-right: ${euiThemeVars.euiSizeS};
 `;
 
+interface GetOptionFromPromptProps extends PromptResponse {
+  content: string;
+  id: string;
+  name: string;
+  isCleared: boolean;
+}
+
 export const getOptionFromPrompt = ({
   content,
   id,
+  isCleared,
   name,
-}: PromptResponse): EuiSuperSelectOption<string> => ({
+}: GetOptionFromPromptProps): EuiSuperSelectOption<string> => ({
   value: id,
-  inputDisplay: <span data-test-subj="systemPromptText">{name}</span>,
+  inputDisplay: (
+    <span
+      data-test-subj="systemPromptText"
+      // @ts-ignore
+      css={css`
+        color: ${isCleared ? euiThemeVars.euiColorLightShade : euiThemeVars.euiColorDarkestShade};
+      `}
+    >
+      {name}
+    </span>
+  ),
   dropdownDisplay: (
     <>
       <Strong data-test-subj="name">{name}</Strong>
@@ -41,6 +60,10 @@ export const getOptionFromPrompt = ({
 
 interface GetOptionsProps {
   prompts: PromptResponse[] | undefined;
+  isCleared: boolean;
 }
-export const getOptions = ({ prompts }: GetOptionsProps): Array<EuiSuperSelectOption<string>> =>
-  prompts?.map(getOptionFromPrompt) ?? [];
+export const getOptions = ({
+  prompts,
+  isCleared,
+}: GetOptionsProps): Array<EuiSuperSelectOption<string>> =>
+  prompts?.map((p) => getOptionFromPrompt({ ...p, isCleared })) ?? [];

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

@@ -5,8 +5,9 @@
  * 2.0.
  */
 
-import React, { useCallback, useMemo } from 'react';
+import React, { useCallback, useMemo, useState } from 'react';
 import { PromptResponse } from '@kbn/elastic-assistant-common';
+import { QueryObserverResult } from '@tanstack/react-query';
 import { Conversation } from '../../../..';
 import { SelectSystemPrompt } from './select_system_prompt';
 
@@ -17,6 +18,7 @@ interface Props {
   onSystemPromptSelectionChange: (systemPromptId: string | undefined) => void;
   setIsSettingsModalVisible: React.Dispatch<React.SetStateAction<boolean>>;
   allSystemPrompts: PromptResponse[];
+  refetchConversations?: () => Promise<QueryObserverResult<Record<string, Conversation>, unknown>>;
 }
 
 const SystemPromptComponent: React.FC<Props> = ({
@@ -26,9 +28,12 @@ const SystemPromptComponent: React.FC<Props> = ({
   onSystemPromptSelectionChange,
   setIsSettingsModalVisible,
   allSystemPrompts,
+  refetchConversations,
 }) => {
+  const [isCleared, setIsCleared] = useState(false);
   const selectedPrompt = useMemo(() => {
     if (editingSystemPromptId !== undefined) {
+      setIsCleared(false);
       return allSystemPrompts.find((p) => p.id === editingSystemPromptId);
     } else {
       return allSystemPrompts.find((p) => p.id === conversation?.apiConfig?.defaultSystemPromptId);
@@ -36,10 +41,21 @@ const SystemPromptComponent: React.FC<Props> = ({
   }, [allSystemPrompts, conversation?.apiConfig?.defaultSystemPromptId, editingSystemPromptId]);
 
   const handleClearSystemPrompt = useCallback(() => {
-    if (conversation) {
+    if (editingSystemPromptId === undefined) {
+      setIsCleared(false);
+      onSystemPromptSelectionChange(
+        allSystemPrompts.find((p) => p.id === conversation?.apiConfig?.defaultSystemPromptId)?.id
+      );
+    } else {
+      setIsCleared(true);
       onSystemPromptSelectionChange(undefined);
     }
-  }, [conversation, onSystemPromptSelectionChange]);
+  }, [
+    allSystemPrompts,
+    conversation?.apiConfig?.defaultSystemPromptId,
+    editingSystemPromptId,
+    onSystemPromptSelectionChange,
+  ]);
 
   return (
     <SelectSystemPrompt
@@ -48,6 +64,8 @@ const SystemPromptComponent: React.FC<Props> = ({
       conversation={conversation}
       data-test-subj="systemPrompt"
       isClearable={true}
+      isCleared={isCleared}
+      refetchConversations={refetchConversations}
       isSettingsModalVisible={isSettingsModalVisible}
       onSystemPromptSelectionChange={onSystemPromptSelectionChange}
       selectedPrompt={selectedPrompt}

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

@@ -22,6 +22,7 @@ import {
   PromptResponse,
   PromptTypeEnum,
 } from '@kbn/elastic-assistant-common/impl/schemas/prompts/bulk_crud_prompts_route.gen';
+import { QueryObserverResult } from '@tanstack/react-query';
 import { Conversation } from '../../../../..';
 import { getOptions } from '../helpers';
 import * as i18n from '../translations';
@@ -38,6 +39,7 @@ export interface Props {
   selectedPrompt: PromptResponse | undefined;
   clearSelectedSystemPrompt?: () => void;
   isClearable?: boolean;
+  isCleared?: boolean;
   isDisabled?: boolean;
   isOpen?: boolean;
   isSettingsModalVisible: boolean;
@@ -46,6 +48,7 @@ export interface Props {
   onSelectedConversationChange?: (result: Conversation) => void;
   setConversationSettings?: React.Dispatch<React.SetStateAction<Record<string, Conversation>>>;
   setConversationsSettingsBulkActions?: React.Dispatch<Record<string, Conversation>>;
+  refetchConversations?: () => Promise<QueryObserverResult<Record<string, Conversation>, unknown>>;
 }
 
 const ADD_NEW_SYSTEM_PROMPT = 'ADD_NEW_SYSTEM_PROMPT';
@@ -57,8 +60,10 @@ const SelectSystemPromptComponent: React.FC<Props> = ({
   selectedPrompt,
   clearSelectedSystemPrompt,
   isClearable = false,
+  isCleared = false,
   isDisabled = false,
   isOpen = false,
+  refetchConversations,
   isSettingsModalVisible,
   onSystemPromptSelectionChange,
   setIsSettingsModalVisible,
@@ -89,10 +94,11 @@ const SelectSystemPromptComponent: React.FC<Props> = ({
             defaultSystemPromptId: promptId,
           },
         });
+        await refetchConversations?.();
         return result;
       }
     },
-    [conversation, setApiConfig]
+    [conversation, refetchConversations, setApiConfig]
   );
 
   const addNewSystemPrompt = useMemo(() => {
@@ -116,7 +122,10 @@ const SelectSystemPromptComponent: React.FC<Props> = ({
   }, []);
 
   // SuperSelect State/Actions
-  const options = useMemo(() => getOptions({ prompts: allSystemPrompts }), [allSystemPrompts]);
+  const options = useMemo(
+    () => getOptions({ prompts: allSystemPrompts, isCleared }),
+    [allSystemPrompts, isCleared]
+  );
 
   const onChange = useCallback(
     async (selectedSystemPromptId) => {
@@ -160,9 +169,8 @@ const SelectSystemPromptComponent: React.FC<Props> = ({
   );
 
   const clearSystemPrompt = useCallback(() => {
-    setSelectedSystemPrompt(undefined);
     clearSelectedSystemPrompt?.();
-  }, [clearSelectedSystemPrompt, setSelectedSystemPrompt]);
+  }, [clearSelectedSystemPrompt]);
 
   return (
     <EuiFlexGroup
@@ -226,10 +234,14 @@ const SelectSystemPromptComponent: React.FC<Props> = ({
                 inline-size: 16px;
                 block-size: 16px;
                 border-radius: 16px;
-                background: ${euiThemeVars.euiColorMediumShade};
+                background: ${isCleared
+                  ? euiThemeVars.euiColorLightShade
+                  : euiThemeVars.euiColorMediumShade};
 
                 :hover:not(:disabled) {
-                  background: ${euiThemeVars.euiColorMediumShade};
+                  background: ${isCleared
+                    ? euiThemeVars.euiColorLightShade
+                    : euiThemeVars.euiColorMediumShade};
                   transform: none;
                 }
 

x-pack/plugins/security_solution/public/assistant/provider.test.tsx

@@ -5,13 +5,38 @@
  * 2.0.
  */
 
+import React from 'react';
 import { act, renderHook } from '@testing-library/react-hooks';
 import { httpServiceMock, type HttpSetupMock } from '@kbn/core-http-browser-mocks';
 import type { Storage } from '@kbn/kibana-utils-plugin/public';
-import { createConversations } from './provider';
+import { AssistantProvider, createConversations } from './provider';
 import { coreMock } from '@kbn/core/public/mocks';
+import { useKibana as mockUseKibana } from '../common/lib/kibana/__mocks__';
 import { loadAllActions as loadConnectors } from '@kbn/triggers-actions-ui-plugin/public/common/constants';
+import { useKibana } from '../common/lib/kibana';
+import { render, waitFor } from '@testing-library/react';
+import { TestProviders } from '../common/mock';
+import { useAssistantAvailability } from './use_assistant_availability';
+import {
+  bulkUpdatePrompts,
+  getPrompts,
+  getUserConversations,
+} from '@kbn/elastic-assistant/impl/assistant/api';
+import { BASE_SECURITY_SYSTEM_PROMPTS } from './content/prompts/system';
 
+const mockedUseKibana = mockUseKibana();
+jest.mock('./use_assistant_availability');
+jest.mock('../common/lib/kibana');
+
+jest.mock('@kbn/elastic-assistant/impl/assistant/api');
+jest.mock('../common/hooks/use_license', () => ({
+  useLicense: () => ({
+    isEnterprise: () => true,
+  }),
+  licenseService: {
+    isEnterprise: () => true,
+  },
+}));
 jest.mock('@kbn/triggers-actions-ui-plugin/public/common/constants');
 let http: HttpSetupMock = coreMock.createSetup().http;
 export const mockConnectors = [
@@ -199,3 +224,85 @@ describe('createConversations', () => {
     });
   });
 });
+describe('AssistantProvider', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    (useKibana as jest.Mock).mockReturnValue({
+      ...mockedUseKibana,
+      services: {
+        ...mockedUseKibana.services,
+      },
+    });
+    jest.mocked(useAssistantAvailability).mockReturnValue({
+      hasAssistantPrivilege: true,
+      hasConnectorsAllPrivilege: true,
+      hasConnectorsReadPrivilege: true,
+      hasUpdateAIAssistantAnonymization: true,
+      isAssistantEnabled: true,
+    });
+
+    (getUserConversations as jest.Mock).mockResolvedValue({
+      page: 1,
+      perPage: 5,
+      total: 5,
+      data: [],
+    });
+    (getPrompts as jest.Mock).mockResolvedValue({
+      page: 1,
+      perPage: 5,
+      total: 0,
+      data: [],
+    });
+  });
+  it('should not render the assistant when no prompts have been returned', async () => {
+    const { queryByTestId } = render(
+      <AssistantProvider>
+        <span data-test-subj="ourAssistant" />
+      </AssistantProvider>,
+      {
+        wrapper: TestProviders,
+      }
+    );
+    expect(queryByTestId('ourAssistant')).toBeNull();
+  });
+  it('should render the assistant when prompts are returned', async () => {
+    (getPrompts as jest.Mock).mockResolvedValue({
+      page: 1,
+      perPage: 5,
+      total: 2,
+      data: BASE_SECURITY_SYSTEM_PROMPTS,
+    });
+    const { getByTestId } = render(
+      <AssistantProvider>
+        <span data-test-subj="ourAssistant" />
+      </AssistantProvider>,
+      {
+        wrapper: TestProviders,
+      }
+    );
+    await waitFor(() => {
+      expect(getByTestId('ourAssistant')).not.toBeNull();
+    });
+  });
+  it('should render the assistant once prompts have been created', async () => {
+    (bulkUpdatePrompts as jest.Mock).mockResolvedValue({
+      success: true,
+      attributes: {
+        results: {
+          created: BASE_SECURITY_SYSTEM_PROMPTS,
+        },
+      },
+    });
+    const { getByTestId } = render(
+      <AssistantProvider>
+        <span data-test-subj="ourAssistant" />
+      </AssistantProvider>,
+      {
+        wrapper: TestProviders,
+      }
+    );
+    await waitFor(() => {
+      expect(getByTestId('ourAssistant')).not.toBeNull();
+    });
+  });
+});

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

@@ -5,7 +5,7 @@
  * 2.0.
  */
 import type { FC, PropsWithChildren } from 'react';
-import React, { useEffect } from 'react';
+import React, { useEffect, useState } from 'react';
 import { parse } from '@kbn/datemath';
 import type { Storage } from '@kbn/kibana-utils-plugin/public';
 import { i18n } from '@kbn/i18n';
@@ -128,7 +128,7 @@ export const createBasePrompts = async (notifications: NotificationsStart, http:
     notifications.toasts
   );
   if (bulkResult && bulkResult.success) {
-    return true;
+    return bulkResult.attributes.results.created;
   }
 };
 
@@ -176,6 +176,8 @@ export const AssistantProvider: FC<PropsWithChildren<unknown>> = ({ children })
     storage,
   ]);
 
+  const [basePromptsLoaded, setBasePromptsLoaded] = useState(false);
+
   useEffect(() => {
     const createSecurityPrompts = once(async () => {
       if (
@@ -183,15 +185,20 @@ export const AssistantProvider: FC<PropsWithChildren<unknown>> = ({ children })
         assistantAvailability.isAssistantEnabled &&
         assistantAvailability.hasAssistantPrivilege
       ) {
-        const res = await getPrompts({
-          http,
-          toasts: notifications.toasts,
-        });
+        try {
+          const res = await getPrompts({
+            http,
+            toasts: notifications.toasts,
+          });
 
-        if (res.total === 0) {
-          await createBasePrompts(notifications, http);
-        }
+          if (res.total === 0) {
+            await createBasePrompts(notifications, http);
+          }
+          // eslint-disable-next-line no-empty
+        } catch (e) {}
       }
+
+      setBasePromptsLoaded(true);
     });
     createSecurityPrompts();
   }, [
@@ -205,6 +212,9 @@ export const AssistantProvider: FC<PropsWithChildren<unknown>> = ({ children })
   const { signalIndexName } = useSignalIndex();
   const alertsIndexPattern = signalIndexName ?? undefined;
   const toasts = useAppToasts() as unknown as IToasts; // useAppToasts is the current, non-deprecated method of getting the toasts service in the Security Solution, but it doesn't return the IToasts interface (defined by core)
+  // Because our conversations need an assigned system prompt at create time,
+  // we want to make sure the prompts are there before creating the first conversation
+  // however if there is an error fetching the prompts, we don't want to block the app
 
   return (
     <ElasticAssistantProvider
@@ -224,7 +234,7 @@ export const AssistantProvider: FC<PropsWithChildren<unknown>> = ({ children })
       toasts={toasts}
       currentAppId={currentAppId ?? 'securitySolutionUI'}
     >
-      {children}
+      {basePromptsLoaded ? children : null}
     </ElasticAssistantProvider>
   );
 };