Show partial results after search has been canceled

src/platform/packages/shared/kbn-unified-histogram/components/chart/hooks/use_total_hits.ts

@@ -61,8 +61,8 @@ export const useTotalHits = ({
     return () => subscription.unsubscribe();
   }, [fetch, fetch$]);
 
-  const onAbort = useCallback(() => {
-    abortController.current?.abort();
+  const onAbort = useCallback((e: Event) => {
+    abortController.current?.abort((e.target as AbortSignal)?.reason);
   }, []);
 
   useEffect(() => {

src/platform/plugins/shared/dashboard/public/dashboard_app/no_data/dashboard_app_no_data.tsx

@@ -19,6 +19,7 @@ import {
 import { withSuspense } from '@kbn/shared-ux-utility';
 import type { LensSerializedState } from '@kbn/lens-plugin/public';
 import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 import {
   coreServices,
   dataService,
@@ -59,12 +60,12 @@ export const DashboardAppNoDataPage = ({
 
   useEffect(() => {
     return () => {
-      abortController?.abort();
+      abortController?.abort(AbortReason.CLEANUP);
     };
   }, [abortController]);
 
   const onTryESQL = useCallback(async () => {
-    abortController?.abort();
+    abortController?.abort(AbortReason.REPLACED);
     if (lensHelpersAsync.value) {
       const abc = new AbortController();
       const { dataViews } = dataService;

src/platform/plugins/shared/data/common/search/poll_search.test.ts

@@ -8,7 +8,7 @@
  */
 
 import { pollSearch } from './poll_search';
-import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import { AbortError, AbortReason } from '@kbn/kibana-utils-plugin/common';
 
 describe('pollSearch', () => {
   function getMockedSearch$(resolveOnI = 1) {
@@ -86,6 +86,29 @@ describe('pollSearch', () => {
     expect(cancelFn).toBeCalledTimes(1);
   });
 
+  test('Does not throw or cancel if abort reason is CANCELED', async () => {
+    const searchFn = jest.fn().mockResolvedValue({
+      isRunning: false,
+      isPartial: false,
+      rawResponse: {},
+    });
+    const cancelFn = jest.fn();
+
+    const abortController = new AbortController();
+    setTimeout(() => abortController.abort(AbortReason.CANCELED), 100);
+
+    // Poll should complete without throwing AbortError and without calling cancel
+    await expect(
+      pollSearch(searchFn, cancelFn, { abortSignal: abortController.signal }).toPromise()
+    ).resolves.toEqual({
+      isRunning: false,
+      isPartial: false,
+      rawResponse: {},
+    });
+
+    expect(cancelFn).not.toHaveBeenCalled();
+  });
+
   test('Does not leak unresolved promises on cancel', async () => {
     const searchFn = getMockedSearch$(20);
     const cancelFn = jest.fn().mockRejectedValueOnce({ error: 'Oh no!' });

src/platform/plugins/shared/data/common/search/poll_search.ts

@@ -8,8 +8,19 @@
  */
 
 import type { Observable } from 'rxjs';
-import { from, timer, defer, fromEvent, EMPTY } from 'rxjs';
-import { expand, map, switchMap, takeUntil, takeWhile, tap } from 'rxjs';
+import {
+  defer,
+  EMPTY,
+  expand,
+  from,
+  fromEvent,
+  switchMap,
+  takeUntil,
+  takeWhile,
+  tap,
+  throwError,
+  timer,
+} from 'rxjs';
 import { AbortError } from '@kbn/kibana-utils-plugin/common';
 import type { IKibanaSearchResponse } from '@kbn/search-types';
 import type { IAsyncSearchOptions } from '..';
@@ -54,9 +65,7 @@ export const pollSearch = <Response extends IKibanaSearchResponse>(
     }
 
     const aborted$ = (abortSignal ? fromEvent(abortSignal, 'abort') : EMPTY).pipe(
-      map(() => {
-        throw new AbortError();
-      })
+      switchMap((e) => throwError(() => new AbortError((e.target as AbortSignal)?.reason)))
     );
 
     return from(search()).pipe(

src/platform/plugins/shared/data/public/search/search_interceptor/search_abort_controller.test.ts

@@ -7,6 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 import { SearchAbortController } from './search_abort_controller';
 
 const timeTravel = (msToRun = 0) => {
@@ -45,10 +46,26 @@ describe('search abort controller', () => {
     const controller2 = new AbortController();
     sac.addAbortSignal(controller2.signal);
     expect(sac.getSignal().aborted).toBe(false);
-    controller.abort();
+    controller.abort(AbortReason.CANCELED);
     expect(sac.getSignal().aborted).toBe(false);
-    controller2.abort();
-    expect(sac.getSignal().aborted).toBe(true);
+    controller2.abort(AbortReason.CANCELED);
+    const signal = sac.getSignal();
+    expect(signal.aborted).toBe(true);
+    expect(signal.reason).toBe(AbortReason.CANCELED);
+  });
+
+  test('when the abort reason is CANCELED', () => {
+    const sac = new SearchAbortController();
+    sac.abort(AbortReason.CANCELED);
+    expect(sac.isCanceled()).toBe(true);
+    expect(sac.isTimeout()).toBe(false);
+  });
+
+  test('when the abort reason is TIMEOUT', () => {
+    const sac = new SearchAbortController();
+    sac.abort(AbortReason.TIMEOUT);
+    expect(sac.isTimeout()).toBe(true);
+    expect(sac.isCanceled()).toBe(false);
   });
 
   test('aborts explicitly even if all inputs are not aborted', () => {

src/platform/plugins/shared/data/public/search/search_interceptor/search_abort_controller.ts

@@ -9,23 +9,18 @@
 
 import type { Subscription } from 'rxjs';
 import { timer } from 'rxjs';
-
-export enum AbortReason {
-  Timeout = 'timeout',
-}
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 
 export class SearchAbortController {
   private inputAbortSignals: AbortSignal[] = new Array();
   private abortController: AbortController = new AbortController();
   private timeoutSub?: Subscription;
   private destroyed = false;
-  private reason?: AbortReason;
 
   constructor(timeout?: number) {
     if (timeout) {
       this.timeoutSub = timer(timeout).subscribe(() => {
-        this.reason = AbortReason.Timeout;
-        this.abortController.abort();
+        this.abortController.abort(AbortReason.TIMEOUT);
         this.timeoutSub!.unsubscribe();
       });
     }
@@ -34,7 +29,7 @@ export class SearchAbortController {
   private abortHandler = () => {
     const allAborted = this.inputAbortSignals.every((signal) => signal.aborted);
     if (allAborted) {
-      this.abortController.abort();
+      this.abortController.abort(this.inputAbortSignals[0].reason);
       this.cleanup();
     }
   };
@@ -67,12 +62,16 @@ export class SearchAbortController {
     return this.abortController.signal;
   }
 
-  public abort() {
+  public abort(reason?: AbortReason) {
     this.cleanup();
-    this.abortController.abort();
+    this.abortController.abort(reason);
   }
 
   public isTimeout() {
-    return this.reason === AbortReason.Timeout;
+    return this.abortController.signal.reason === AbortReason.TIMEOUT;
+  }
+
+  public isCanceled() {
+    return this.abortController.signal.reason === AbortReason.CANCELED;
   }
 }

src/platform/plugins/shared/data/public/search/search_interceptor/search_interceptor.test.ts

@@ -17,16 +17,15 @@ import { AbortError } from '@kbn/kibana-utils-plugin/public';
 import { EsError, type IEsError } from '@kbn/search-errors';
 import type { ISessionService } from '..';
 import { SearchSessionState } from '..';
-
 import * as searchPhaseException from '../../../common/search/test_data/search_phase_execution_exception.json';
 import * as resourceNotFoundException from '../../../common/search/test_data/resource_not_found_exception.json';
 import { BehaviorSubject } from 'rxjs';
 import { dataPluginMock } from '../../mocks';
+import { AbortReason } from '@kbn/kibana-utils-plugin/common';
 import { ESQL_ASYNC_SEARCH_STRATEGY, UI_SETTINGS } from '../../../common';
 import type { SearchServiceStartDependencies } from '../search_service';
 import type { Start as InspectorStart } from '@kbn/inspector-plugin/public';
 import { SearchTimeoutError, TimeoutErrorMode } from './timeout_error';
-
 import { SearchSessionIncompleteWarning } from './search_session_incomplete_warning';
 import { getMockSearchConfig } from '../../../config.mock';
 
@@ -145,9 +144,9 @@ describe('SearchInterceptor', () => {
       }
     });
 
-    next.mockClear();
-    error.mockClear();
-    complete.mockClear();
+    next.mockReset();
+    error.mockReset();
+    complete.mockReset();
     jest.clearAllTimers();
     jest.clearAllMocks();
 
@@ -1050,7 +1049,6 @@ describe('SearchInterceptor', () => {
       afterEach(() => {
         const sessionServiceMock = sessionService as jest.Mocked<ISessionService>;
         sessionServiceMock.getSearchOptions.mockReset();
-        mockCoreSetup.http.post.mockReset();
       });
 
       test('gets session search options from session service', async () => {
@@ -1359,7 +1357,7 @@ describe('SearchInterceptor', () => {
         const abort = sessionService.trackSearch.mock.calls[0][0].abort;
         expect(abort).toBeInstanceOf(Function);
 
-        abort();
+        abort(AbortReason.REPLACED);
 
         await timeTravel(10);
 
@@ -2129,5 +2127,111 @@ describe('SearchInterceptor', () => {
         response.subscribe({ error });
       });
     });
+
+    describe('partial results', () => {
+      beforeEach(() => {
+        mockCoreSetup.http.post.mockResolvedValue(
+          getMockSearchResponse({
+            id: '1',
+            isPartial: true,
+            isRunning: true,
+            rawResponse: {},
+          })
+        );
+      });
+
+      test('should request partial results and throw error if timed out', async () => {
+        const abortController = new AbortController();
+        setTimeout(() => {
+          abortController.abort(AbortReason.TIMEOUT);
+        }, 50);
+
+        const response = searchInterceptor.search(
+          {},
+          { abortSignal: abortController.signal, pollInterval: 100 }
+        );
+        response.subscribe({ next, error });
+
+        await timeTravel(); // Run first request/response
+
+        expect(next).toHaveBeenCalled();
+        expect(error).not.toHaveBeenCalled();
+
+        await timeTravel(50); // Run until abort
+
+        expect(mockCoreSetup.http.post).toHaveBeenCalledTimes(2);
+        expect(mockCoreSetup.http.post.mock.calls[1]).toMatchInlineSnapshot(`
+          Array [
+            "/internal/search/ese/1",
+            Object {
+              "asResponse": true,
+              "body": "{\\"id\\":\\"1\\",\\"params\\":{},\\"retrieveResults\\":true,\\"stream\\":true}",
+              "context": undefined,
+              "signal": AbortSignal {},
+              "version": "1",
+            },
+          ]
+        `);
+        expect(error).toHaveBeenCalled();
+      });
+
+      test('should request partial results and not throw error if canceled', async () => {
+        const abortController = new AbortController();
+        setTimeout(() => {
+          abortController.abort(AbortReason.CANCELED);
+        }, 50);
+
+        const response = searchInterceptor.search(
+          {},
+          { abortSignal: abortController.signal, pollInterval: 100 }
+        );
+        response.subscribe({ next, error });
+
+        await timeTravel(); // Run first request/response
+
+        expect(next).toHaveBeenCalled();
+        expect(error).not.toHaveBeenCalled();
+
+        await timeTravel(50); // Run until abort
+
+        expect(mockCoreSetup.http.post).toHaveBeenCalledTimes(2);
+        expect(mockCoreSetup.http.post.mock.calls[1]).toMatchInlineSnapshot(`
+          Array [
+            "/internal/search/ese/1",
+            Object {
+              "asResponse": true,
+              "body": "{\\"id\\":\\"1\\",\\"params\\":{},\\"retrieveResults\\":true,\\"stream\\":true}",
+              "context": undefined,
+              "signal": AbortSignal {},
+              "version": "1",
+            },
+          ]
+        `);
+        expect(error).not.toHaveBeenCalled();
+      });
+
+      test('should not request partial results and throw error if canceled for a reason other than CANCELED/TIMEOUT', async () => {
+        const abortController = new AbortController();
+        setTimeout(() => {
+          abortController.abort(AbortReason.CLEANUP);
+        }, 50);
+
+        const response = searchInterceptor.search(
+          {},
+          { abortSignal: abortController.signal, pollInterval: 100 }
+        );
+        response.subscribe({ next, error });
+
+        await timeTravel(); // Run first request/response
+
+        expect(next).toHaveBeenCalled();
+        expect(error).not.toHaveBeenCalled();
+
+        await timeTravel(50); // Run until abort
+
+        expect(mockCoreSetup.http.post).toHaveBeenCalledTimes(1);
+        expect(error).toHaveBeenCalled();
+      });
+    });
   });
 });