[Backport 2.x] [query assist] use badge to show agent errors

src/plugins/data/public/query/query_string/query_string_manager.mock.ts

@@ -28,6 +28,7 @@
  * under the License.
  */
 
+import { of } from 'rxjs';
 import { QueryStringContract } from '.';
 import { Query, Dataset } from '../../../common';
 import { datasetServiceMock } from './dataset_service/dataset_service.mock';
@@ -44,10 +45,10 @@ const createSetupContractMock = (isEnhancementsEnabled: boolean = false) => {
   };
 
   const queryStringManagerMock: jest.Mocked<QueryStringContract> = {
-    getQuery: jest.fn(),
+    getQuery: jest.fn().mockReturnValue(defaultQuery),
     setQuery: jest.fn(),
-    getUpdates$: jest.fn(),
-    getDefaultQuery: jest.fn(),
+    getUpdates$: jest.fn().mockReturnValue(of(defaultQuery)),
+    getDefaultQuery: jest.fn().mockReturnValue(defaultQuery),
     formatQuery: jest.fn(),
     clearQuery: jest.fn(),
     addToQueryHistory: jest.fn(),

src/plugins/query_enhancements/public/plugin.tsx

@@ -4,7 +4,6 @@
  */
 
 import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../core/public';
-import { IStorageWrapper, Storage } from '../../opensearch_dashboards_utils/public';
 import { ConfigSchema } from '../common/config';
 import { setData, setStorage } from './services';
 import { createQueryAssistExtension } from './query_assist';
@@ -19,6 +18,7 @@ import { LanguageConfig, Query } from '../../data/public';
 import { s3TypeConfig } from './datasets';
 import { createEditor, DefaultInput, SingleLineInput } from '../../data/public';
 import { QueryLanguageReference } from './query_editor/query_language_reference';
+import { DataStorage } from '../../data/common';
 
 export class QueryEnhancementsPlugin
   implements
@@ -28,12 +28,12 @@ export class QueryEnhancementsPlugin
       QueryEnhancementsPluginSetupDependencies,
       QueryEnhancementsPluginStartDependencies
     > {
-  private readonly storage: IStorageWrapper;
+  private readonly storage: DataStorage;
   private readonly config: ConfigSchema;
 
   constructor(initializerContext: PluginInitializerContext) {
     this.config = initializerContext.config.get<ConfigSchema>();
-    this.storage = new Storage(window.localStorage);
+    this.storage = new DataStorage(window.localStorage, 'discover.queryAssist.');
   }
 
   public setup(

src/plugins/query_enhancements/public/query_assist/_index.scss

@@ -7,6 +7,31 @@
   &.queryAssist__callout {
     margin-top: $euiSizeXS;
   }
+
+  .queryAssist__popoverWrapper {
+    display: flex;
+    align-items: center;
+    padding-right: $euiSizeS;
+  }
+
+  // Hide the error badge when the input is focused
+  .queryAssist__inputWrapper:has(.queryAssist__input:focus) .queryAssist__popoverWrapper {
+    display: none;
+  }
+
+  .queryAssist__input {
+    text-overflow: ellipsis;
+  }
+}
+
+.queryAssist__popoverText {
+  padding: 0 $euiSizeM;
+}
+
+#queryAssistErrorTitle {
+  // there is a medium size icon with two extra-small gutter between the icon
+  // and the title pushing to the left so title is aligned with other content
+  margin-left: -($euiSizeM + $euiSizeS);
 }
 
 .osdQueryEditorExtensionComponent__query-assist {

src/plugins/query_enhancements/public/query_assist/components/call_outs.test.tsx

@@ -4,13 +4,13 @@
  */
 
 import { render } from '@testing-library/react';
-import React, { ComponentProps } from 'react';
+import React, { ComponentProps, PropsWithChildren } from 'react';
 import { IntlProvider } from 'react-intl';
 import { QueryAssistCallOut } from './call_outs';
 
 type Props = ComponentProps<typeof QueryAssistCallOut>;
 
-const IntlWrapper = ({ children }: { children: unknown }) => (
+const IntlWrapper = ({ children }: PropsWithChildren<unknown>) => (
   <IntlProvider locale="en">{children}</IntlProvider>
 );
 

src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.test.tsx

@@ -0,0 +1,218 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { fireEvent, render, screen, waitFor } from '@testing-library/react';
+import React, { ComponentProps, PropsWithChildren } from 'react';
+import { IntlProvider } from 'react-intl';
+import { of } from 'rxjs';
+import { QueryAssistBar } from '.';
+import { notificationServiceMock, uiSettingsServiceMock } from '../../../../../core/public/mocks';
+import { DataStorage } from '../../../../data/common';
+import { QueryEditorExtensionDependencies, QueryStringContract } from '../../../../data/public';
+import { dataPluginMock } from '../../../../data/public/mocks';
+import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
+import { setData, setStorage } from '../../services';
+import { useGenerateQuery } from '../hooks';
+import { AgentError, ProhibitedQueryError } from '../utils';
+import { QueryAssistInput } from './query_assist_input';
+
+jest.mock('../../../../opensearch_dashboards_react/public', () => ({
+  useOpenSearchDashboards: jest.fn(),
+  withOpenSearchDashboards: jest.fn((component: React.Component) => component),
+}));
+
+jest.mock('../hooks', () => ({
+  useGenerateQuery: jest.fn().mockReturnValue({ generateQuery: jest.fn(), loading: false }),
+}));
+
+jest.mock('./query_assist_input', () => ({
+  QueryAssistInput: ({ inputRef, error }: ComponentProps<typeof QueryAssistInput>) => (
+    <>
+      <input ref={inputRef} />
+      <div>{JSON.stringify(error)}</div>
+    </>
+  ),
+}));
+
+const dataMock = dataPluginMock.createStartContract(true);
+const queryStringMock = dataMock.query.queryString as jest.Mocked<QueryStringContract>;
+const uiSettingsMock = uiSettingsServiceMock.createStartContract();
+const notificationsMock = notificationServiceMock.createStartContract();
+
+setData(dataMock);
+setStorage(new DataStorage(window.localStorage, 'mock-prefix'));
+
+const dependencies: QueryEditorExtensionDependencies = {
+  language: 'PPL',
+  onSelectLanguage: jest.fn(),
+  isCollapsed: false,
+  setIsCollapsed: jest.fn(),
+};
+
+type Props = ComponentProps<typeof QueryAssistBar>;
+
+const IntlWrapper = ({ children }: PropsWithChildren<unknown>) => (
+  <IntlProvider locale="en">{children}</IntlProvider>
+);
+
+const renderQueryAssistBar = (overrideProps: Partial<Props> = {}) => {
+  const props: Props = Object.assign<Props, Partial<Props>>({ dependencies }, overrideProps);
+  const component = render(<QueryAssistBar {...props} />, {
+    wrapper: IntlWrapper,
+  });
+  return { component, props: props as jest.MockedObjectDeep<Props> };
+};
+
+describe('QueryAssistBar', () => {
+  beforeEach(() => {
+    (useOpenSearchDashboards as jest.Mock).mockReturnValue({
+      services: {
+        data: dataMock,
+        uiSettings: uiSettingsMock,
+        notifications: notificationsMock,
+      },
+    });
+  });
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('renders null if collapsed', () => {
+    const { component } = renderQueryAssistBar({
+      dependencies: { ...dependencies, isCollapsed: true },
+    });
+    expect(component.container).toBeEmptyDOMElement();
+  });
+
+  it('matches snapshot', () => {
+    const { component } = renderQueryAssistBar();
+    expect(component.container).toMatchSnapshot();
+  });
+
+  it('displays callout when query input is empty on submit', async () => {
+    renderQueryAssistBar();
+
+    fireEvent.click(screen.getByRole('button'));
+
+    await waitFor(() => {
+      expect(screen.getByTestId('query-assist-empty-query-callout')).toBeInTheDocument();
+    });
+  });
+
+  it('displays callout when dataset is not selected on submit', async () => {
+    queryStringMock.getQuery.mockReturnValueOnce({ query: '', language: 'kuery' });
+    queryStringMock.getUpdates$.mockReturnValueOnce(of({ query: '', language: 'kuery' }));
+    renderQueryAssistBar();
+
+    fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
+    fireEvent.click(screen.getByRole('button'));
+
+    await waitFor(() => {
+      expect(screen.getByTestId('query-assist-empty-index-callout')).toBeInTheDocument();
+    });
+  });
+
+  it('displays callout for guardrail errors', async () => {
+    const generateQueryMock = jest.fn().mockResolvedValue({ error: new ProhibitedQueryError() });
+    (useGenerateQuery as jest.Mock).mockReturnValue({
+      generateQuery: generateQueryMock,
+      loading: false,
+    });
+
+    renderQueryAssistBar();
+
+    fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
+    fireEvent.click(screen.getByRole('button'));
+
+    await waitFor(() => {
+      expect(screen.getByTestId('query-assist-guard-callout')).toBeInTheDocument();
+    });
+  });
+
+  it('passes agent errors to input', async () => {
+    const generateQueryMock = jest.fn().mockResolvedValue({
+      error: new AgentError({
+        error: { type: 'mock-type', reason: 'mock-reason', details: 'mock-details' },
+        status: 303,
+      }),
+    });
+    (useGenerateQuery as jest.Mock).mockReturnValue({
+      generateQuery: generateQueryMock,
+      loading: false,
+    });
+
+    renderQueryAssistBar();
+
+    fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
+    fireEvent.click(screen.getByRole('button'));
+
+    await waitFor(() => {
+      expect(screen.getByText(/mock-reason/)).toBeInTheDocument();
+    });
+  });
+
+  it('displays toast for other unknown errors', async () => {
+    const mockError = new Error('mock-error');
+    const generateQueryMock = jest.fn().mockResolvedValue({
+      error: mockError,
+    });
+    (useGenerateQuery as jest.Mock).mockReturnValue({
+      generateQuery: generateQueryMock,
+      loading: false,
+    });
+
+    renderQueryAssistBar();
+
+    fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
+    fireEvent.click(screen.getByRole('button'));
+
+    await waitFor(() => {
+      expect(notificationsMock.toasts.addError).toHaveBeenCalledWith(mockError, {
+        title: 'Failed to generate results',
+      });
+    });
+  });
+
+  it('submits a valid query and updates services', async () => {
+    const generateQueryMock = jest
+      .fn()
+      .mockResolvedValue({ response: { query: 'generated query' } });
+    (useGenerateQuery as jest.Mock).mockReturnValue({
+      generateQuery: generateQueryMock,
+      loading: false,
+    });
+
+    renderQueryAssistBar();
+
+    fireEvent.change(screen.getByRole('textbox'), { target: { value: 'test query' } });
+    fireEvent.click(screen.getByRole('button'));
+
+    await waitFor(() => {
+      expect(generateQueryMock).toHaveBeenCalledWith({
+        question: 'test query',
+        index: 'Default Index Pattern',
+        language: 'PPL',
+        dataSourceId: 'mock-data-source-id',
+      });
+    });
+
+    expect(queryStringMock.setQuery).toHaveBeenCalledWith({
+      dataset: {
+        dataSource: {
+          id: 'mock-data-source-id',
+          title: 'Default Data Source',
+          type: 'OpenSearch',
+        },
+        id: 'default-index-pattern',
+        timeFieldName: '@timestamp',
+        title: 'Default Index Pattern',
+        type: 'INDEX_PATTERN',
+      },
+      language: 'PPL',
+      query: 'generated query',
+    });
+    expect(screen.getByTestId('query-assist-query-generated-callout')).toBeInTheDocument();
+  });
+});