fix: streamed query consume aware signal (#9963)

* chore: move consume aware signal logic to utils

* fix: add consume aware signal to streamedQuery

* test: add a test for not a signal that is not consumed by streamFn

* chore: add changeset

* fix: made queryFn of streamedQuery not consume signal directly (so that only streamFn can consume it)

* fix: made onCancelled get called only once

* fix: memoized getSignal and made sure it is called only once

* fix: no-unnecessary-condition

* chore: size-limit

---------

Co-authored-by: Dominik Dorfmeister <[email protected]>

Author: razhayat

.changeset/dry-streets-exist.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/query-core': patch
+---
+
+Made context.signal consume aware with streamedQuery

package.json

@@ -94,7 +94,7 @@
     {
       "name": "react full",
       "path": "packages/react-query/build/modern/index.js",
-      "limit": "12.02 kB",
+      "limit": "12.10 kB",
       "ignore": [
         "react",
         "react-dom"
@@ -103,7 +103,7 @@
     {
       "name": "react minimal",
       "path": "packages/react-query/build/modern/index.js",
-      "limit": "9.08 kB",
+      "limit": "9.11 kB",
       "import": "{ useQuery, QueryClient, QueryClientProvider }",
       "ignore": [
         "react",

packages/query-core/src/__tests__/streamedQuery.test.tsx

@@ -329,7 +329,12 @@ describe('streamedQuery', () => {
     const observer = new QueryObserver(queryClient, {
       queryKey: key,
       queryFn: streamedQuery({
-        streamFn: () => createAsyncNumberGenerator(3),
+        streamFn: (context) => {
+          // just consume the signal
+          // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+          const numbers = context.signal ? 3 : 0
+          return createAsyncNumberGenerator(numbers)
+        },
         refetchMode: 'append',
       }),
     })
@@ -420,6 +425,42 @@ describe('streamedQuery', () => {
     })
   })
 
+  test('should not abort when signal not consumed', async () => {
+    const key = queryKey()
+    const observer = new QueryObserver(queryClient, {
+      queryKey: key,
+      queryFn: streamedQuery({
+        streamFn: () => createAsyncNumberGenerator(3),
+      }),
+    })
+
+    const unsubscribe = observer.subscribe(vi.fn())
+
+    expect(queryClient.getQueryState(key)).toMatchObject({
+      status: 'pending',
+      fetchStatus: 'fetching',
+      data: undefined,
+    })
+
+    await vi.advanceTimersByTimeAsync(60)
+
+    expect(queryClient.getQueryState(key)).toMatchObject({
+      status: 'success',
+      fetchStatus: 'fetching',
+      data: [0],
+    })
+
+    unsubscribe()
+
+    await vi.advanceTimersByTimeAsync(50)
+
+    expect(queryClient.getQueryState(key)).toMatchObject({
+      status: 'success',
+      fetchStatus: 'fetching',
+      data: [0, 1],
+    })
+  })
+
   test('should support custom reducer', async () => {
     const key = queryKey()
 

packages/query-core/src/infiniteQueryBehavior.ts

@@ -1,4 +1,9 @@
-import { addToEnd, addToStart, ensureQueryFn } from './utils'
+import {
+  addConsumeAwareSignal,
+  addToEnd,
+  addToStart,
+  ensureQueryFn,
+} from './utils'
 import type { QueryBehavior } from './query'
 import type {
   InfiniteData,
@@ -23,19 +28,11 @@ export function infiniteQueryBehavior<TQueryFnData, TError, TData, TPageParam>(
       const fetchFn = async () => {
         let cancelled = false
         const addSignalProperty = (object: unknown) => {
-          Object.defineProperty(object, 'signal', {
-            enumerable: true,
-            get: () => {
-              if (context.signal.aborted) {
-                cancelled = true
-              } else {
-                context.signal.addEventListener('abort', () => {
-                  cancelled = true
-                })
-              }
-              return context.signal
-            },
-          })
+          addConsumeAwareSignal(
+            object,
+            () => context.signal,
+            () => (cancelled = true),
+          )
         }
 
         const queryFn = ensureQueryFn(context.options, context.fetchOptions)

packages/query-core/src/streamedQuery.ts

@@ -1,5 +1,10 @@
-import { addToEnd } from './utils'
-import type { QueryFunction, QueryFunctionContext, QueryKey } from './types'
+import { addConsumeAwareSignal, addToEnd } from './utils'
+import type {
+  OmitKeyof,
+  QueryFunction,
+  QueryFunctionContext,
+  QueryKey,
+} from './types'
 
 type BaseStreamedQueryParams<TQueryFnData, TQueryKey extends QueryKey> = {
   streamFn: (
@@ -73,10 +78,25 @@ export function streamedQuery<
 
     let result = initialValue
 
-    const stream = await streamFn(context)
+    let cancelled: boolean = false as boolean
+    const streamFnContext = addConsumeAwareSignal<
+      OmitKeyof<typeof context, 'signal'>
+    >(
+      {
+        client: context.client,
+        meta: context.meta,
+        queryKey: context.queryKey,
+        pageParam: context.pageParam,
+        direction: context.direction,
+      },
+      () => context.signal,
+      () => (cancelled = true),
+    )
+
+    const stream = await streamFn(streamFnContext)
 
     for await (const chunk of stream) {
-      if (context.signal.aborted) {
+      if (cancelled) {
         break
       }
 
@@ -90,7 +110,7 @@ export function streamedQuery<
     }
 
     // finalize result: replace-refetching needs to write to the cache
-    if (isRefetch && refetchMode === 'replace' && !context.signal.aborted) {
+    if (isRefetch && refetchMode === 'replace' && !cancelled) {
       context.client.setQueryData<TData>(context.queryKey, result)
     }
 

packages/query-core/src/utils.ts

@@ -465,3 +465,33 @@ export function shouldThrowError<T extends (...args: Array<any>) => boolean>(
 
   return !!throwOnError
 }
+
+export function addConsumeAwareSignal<T>(
+  object: T,
+  getSignal: () => AbortSignal,
+  onCancelled: VoidFunction,
+): T & { signal: AbortSignal } {
+  let consumed = false
+  let signal: AbortSignal | undefined
+
+  Object.defineProperty(object, 'signal', {
+    enumerable: true,
+    get: () => {
+      signal ??= getSignal()
+      if (consumed) {
+        return signal
+      }
+
+      consumed = true
+      if (signal.aborted) {
+        onCancelled()
+      } else {
+        signal.addEventListener('abort', onCancelled, { once: true })
+      }
+
+      return signal
+    },
+  })
+
+  return object as T & { signal: AbortSignal }
+}