From 13a77e9278af767397a538403dd79c2dd16eec0d Mon Sep 17 00:00:00 2001
From: Walter Rafelsberger <walter.rafelsberger@elastic.co>
Date: Tue, 7 May 2024 12:57:31 +0200
Subject: [PATCH] tweak dispatch

---
 x-pack/packages/ml/response_stream/client/use_fetch_stream.ts | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/x-pack/packages/ml/response_stream/client/use_fetch_stream.ts b/x-pack/packages/ml/response_stream/client/use_fetch_stream.ts
index ff51af87f30986..55950c6ee8f774 100644
--- a/x-pack/packages/ml/response_stream/client/use_fetch_stream.ts
+++ b/x-pack/packages/ml/response_stream/client/use_fetch_stream.ts
@@ -151,10 +151,14 @@ export function useFetchStream<B extends object, R extends Reducer<any, any>>(
     setIsRunning(false);
   };
 
+  // This custom dispatch function allows us to update the `dataRef` value and will
+  // then trigger an update of `data` right away as we don't want to have the
+  // throttling in place for these types of updates.
   const dispatch = (action: ReducerAction<FetchStreamCustomReducer<R>['reducer']>) => {
     dataRef.current = reducerWithFallback.reducer(dataRef.current, action) as ReducerState<
       CustomReducer<R>
     >;
+    setData(dataRef.current);
   };
 
   const cancel = () => {