From 8852912717de99a08b62fb8984cbe84804f902be Mon Sep 17 00:00:00 2001
From: Marco Vettorello <marco.vettorello@elastic.co>
Date: Tue, 10 Dec 2024 09:05:29 +0100
Subject: [PATCH] bring back updateSessionId usage on useLoadingState

---
 .../asset_details/hooks/use_loading_state.ts  | 26 +++++++++++++++----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/x-pack/plugins/observability_solution/infra/public/components/asset_details/hooks/use_loading_state.ts b/x-pack/plugins/observability_solution/infra/public/components/asset_details/hooks/use_loading_state.ts
index 662723cccb420..f3098b6dca8f8 100644
--- a/x-pack/plugins/observability_solution/infra/public/components/asset_details/hooks/use_loading_state.ts
+++ b/x-pack/plugins/observability_solution/infra/public/components/asset_details/hooks/use_loading_state.ts
@@ -26,6 +26,7 @@ import { useCallback, useEffect, useMemo } from 'react';
 import { SearchSessionState, waitUntilNextSessionCompletes$ } from '@kbn/data-plugin/public';
 import { useKibanaContextForPlugin } from '../../../hooks/use_kibana';
 import { useDatePickerContext } from './use_date_picker';
+import { useSearchSessionContext } from '../../../hooks/use_search_session';
 
 export type RequestState = 'running' | 'done' | 'error';
 const WAIT_MS = 1000;
@@ -36,11 +37,16 @@ export const useLoadingState = () => {
   const {
     data: { search },
   } = services;
+  const { updateSearchSessionId } = useSearchSessionContext();
 
   const isAutoRefreshRequestPending$ = useMemo(() => new BehaviorSubject<boolean>(false), []);
   const requestsCount$ = useMemo(() => new BehaviorSubject(0), []);
   const requestState$ = useMemo(() => new BehaviorSubject<RequestState | null>(null), []);
 
+  useEffect(() => {
+    updateSearchSessionId();
+  }, [updateSearchSessionId]);
+
   const waitUntilRequestsCompletes$ = useCallback(
     () =>
       requestsCount$.pipe(
@@ -126,17 +132,27 @@ export const useLoadingState = () => {
 
       autoRefreshTick$.pipe(
         skipUntil(isAutoRefreshEnabled$()),
-        withLatestFrom(search.session.state$),
-        switchMap(([_, state]) => {
+        withLatestFrom(search.session.state$, isAutoRefreshRequestPending$),
+        switchMap(([_, state, isAutoRefreshRequestPending]) => {
           // if the current state$ value is not Completed
           if (state === SearchSessionState.Loading) {
             // Wait until queries using data.search complete before processing the next tick
             // This will only be called when Lens is used in the Asset Details page
-            return waitUntilNextSessionCompletes$(search.session);
+            return waitUntilNextSessionCompletes$(search.session).pipe(
+              tap(() => {
+                updateSearchSessionId();
+              })
+            );
           }
 
           // Else immediately emit true if session state is already completed
-          return of(null);
+          return of(null).pipe(
+            tap(() => {
+              if (!isAutoRefreshRequestPending) {
+                updateSearchSessionId();
+              }
+            })
+          );
         })
       )
     ).subscribe();
@@ -153,7 +169,7 @@ export const useLoadingState = () => {
     requestState$,
     requestsCount$,
     search.session,
-    // updateSearchSessionId,
+    updateSearchSessionId,
     waitUntilRequestsCompletes$,
   ]);