diff --git a/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/query_assist_bar.test.tsx.snap b/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/query_assist_bar.test.tsx.snap
index d9adc26b4d10..c9ffcf90c721 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/query_assist_bar.test.tsx.snap
+++ b/src/plugins/query_enhancements/public/query_assist/components/__snapshots__/query_assist_bar.test.tsx.snap
@@ -21,9 +21,6 @@ exports[`QueryAssistBar matches snapshot 1`] = `
           >
             <input />
           </div>
-          <div
-            class="euiFlexItem euiFlexItem--flexGrowZero queryAssist__badgeContainer"
-          />
           <div
             class="euiFlexItem euiFlexItem--flexGrowZero"
           >
diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.test.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.test.tsx
index 3a170cfe1ec0..44036a11c859 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.test.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.test.tsx
@@ -28,8 +28,11 @@ jest.mock('../hooks', () => ({
 }));
 
 jest.mock('./query_assist_input', () => ({
-  QueryAssistInput: ({ inputRef }: ComponentProps<typeof QueryAssistInput>) => (
-    <input ref={inputRef} />
+  QueryAssistInput: ({ inputRef, error }: ComponentProps<typeof QueryAssistInput>) => (
+    <>
+      <input ref={inputRef} />
+      <div>{JSON.stringify(error)}</div>
+    </>
   ),
 }));
 
@@ -128,7 +131,7 @@ describe('QueryAssistBar', () => {
     });
   });
 
-  it('displays badge for agent errors', async () => {
+  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' },
@@ -146,7 +149,7 @@ describe('QueryAssistBar', () => {
     fireEvent.click(screen.getByRole('button'));
 
     await waitFor(() => {
-      expect(screen.getByTestId('queryAssistErrorBadge')).toBeInTheDocument();
+      expect(screen.getByText(/mock-reason/)).toBeInTheDocument();
     });
   });
 
diff --git a/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx b/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx
index 1e6eb89da1f8..64bfcae6f99f 100644
--- a/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx
+++ b/src/plugins/query_enhancements/public/query_assist/components/query_assist_input.test.tsx
@@ -4,9 +4,10 @@
  */
 
 import { I18nProvider } from '@osd/i18n/react';
-import { fireEvent, render } from '@testing-library/react';
+import { fireEvent, render, screen, waitFor } from '@testing-library/react';
 import React, { ComponentProps } from 'react';
 import { SuggestionsComponentProps } from '../../../../data/public/ui/typeahead/suggestions_component';
+import { AgentError } from '../utils';
 import { QueryAssistInput } from './query_assist_input';
 
 jest.mock('../../services', () => ({
@@ -73,4 +74,14 @@ describe('<QueryAssistInput /> spec', () => {
     fireEvent.click(suggestionButton);
     expect(inputElement.value).toBe('mock suggestion 1');
   });
+
+  it('should show error badge if there is an error', async () => {
+    renderQueryAssistInput({
+      error: new AgentError({
+        error: { type: 'mock-type', reason: 'mock-reason', details: 'mock-details' },
+        status: 303,
+      }),
+    });
+    expect(screen.getByTestId('queryAssistErrorBadge')).toBeInTheDocument();
+  });
 });