From 31c6e4edef242ffa115aac3816bc7ff62d0459df Mon Sep 17 00:00:00 2001 From: razhayat Date: Thu, 11 Dec 2025 19:03:04 +0200 Subject: [PATCH 1/7] chore: move consume aware signal logic to utils --- .../query-core/src/infiniteQueryBehavior.ts | 25 ++++++++----------- packages/query-core/src/utils.ts | 20 +++++++++++++++ 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/packages/query-core/src/infiniteQueryBehavior.ts b/packages/query-core/src/infiniteQueryBehavior.ts index 476d90ce15..9fc0f8954f 100644 --- a/packages/query-core/src/infiniteQueryBehavior.ts +++ b/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( 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) diff --git a/packages/query-core/src/utils.ts b/packages/query-core/src/utils.ts index dadc6d7bc3..fb813f430d 100644 --- a/packages/query-core/src/utils.ts +++ b/packages/query-core/src/utils.ts @@ -465,3 +465,23 @@ export function shouldThrowError) => boolean>( return !!throwOnError } + +export function addConsumeAwareSignal( + object: unknown, + signal: AbortSignal, + onCancelled: VoidFunction, +) { + Object.defineProperty(object, 'signal', { + enumerable: true, + get: () => { + if (signal.aborted) { + onCancelled() + } else { + signal.addEventListener('abort', () => { + onCancelled() + }) + } + return signal + }, + }) +} From 2211c1acafc9d1002432dd3333c300f98c86411d Mon Sep 17 00:00:00 2001 From: razhayat Date: Thu, 11 Dec 2025 19:03:29 +0200 Subject: [PATCH 2/7] fix: add consume aware signal to streamedQuery --- packages/query-core/src/streamedQuery.ts | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/packages/query-core/src/streamedQuery.ts b/packages/query-core/src/streamedQuery.ts index 2eb944d699..81163dee95 100644 --- a/packages/query-core/src/streamedQuery.ts +++ b/packages/query-core/src/streamedQuery.ts @@ -1,4 +1,4 @@ -import { addToEnd } from './utils' +import { addConsumeAwareSignal, addToEnd } from './utils' import type { QueryFunction, QueryFunctionContext, QueryKey } from './types' type BaseStreamedQueryParams = { @@ -73,10 +73,18 @@ export function streamedQuery< let result = initialValue - const stream = await streamFn(context) + let cancelled = false + const streamFnContext = { ...context } + addConsumeAwareSignal( + streamFnContext, + context.signal, + () => (cancelled = true), + ) + + const stream = await streamFn(streamFnContext) for await (const chunk of stream) { - if (context.signal.aborted) { + if (cancelled) { break } @@ -90,7 +98,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(context.queryKey, result) } From aec8321912ce9c5f497abac8ce93f4f8fa99a67c Mon Sep 17 00:00:00 2001 From: razhayat Date: Thu, 11 Dec 2025 20:09:37 +0200 Subject: [PATCH 3/7] test: add a test for not a signal that is not consumed by streamFn --- .../src/__tests__/streamedQuery.test.tsx | 43 ++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/packages/query-core/src/__tests__/streamedQuery.test.tsx b/packages/query-core/src/__tests__/streamedQuery.test.tsx index 084ea7d9db..8768f442e0 100644 --- a/packages/query-core/src/__tests__/streamedQuery.test.tsx +++ b/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: 'idle', + data: [0, 1], + }) + }) + test('should support custom reducer', async () => { const key = queryKey() From b23c4f4da734f60b24f7e4982faabd81aff72353 Mon Sep 17 00:00:00 2001 From: razhayat Date: Thu, 11 Dec 2025 20:10:00 +0200 Subject: [PATCH 4/7] chore: add changeset --- .changeset/dry-streets-exist.md | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 .changeset/dry-streets-exist.md diff --git a/.changeset/dry-streets-exist.md b/.changeset/dry-streets-exist.md new file mode 100644 index 0000000000..0d60838f42 --- /dev/null +++ b/.changeset/dry-streets-exist.md @@ -0,0 +1,5 @@ +--- +'@tanstack/query-core': patch +--- + +Made context.signal consume aware with streamedQuery From 118b914f371206ebd636be0df5c3a51b8c9dff36 Mon Sep 17 00:00:00 2001 From: razhayat Date: Thu, 11 Dec 2025 23:50:43 +0200 Subject: [PATCH 5/7] fix: made queryFn of streamedQuery not consume signal directly (so that only streamFn can consume it) --- .../src/__tests__/streamedQuery.test.tsx | 2 +- .../query-core/src/infiniteQueryBehavior.ts | 2 +- packages/query-core/src/streamedQuery.ts | 22 ++++++++++++++----- packages/query-core/src/utils.ts | 11 ++++++---- 4 files changed, 26 insertions(+), 11 deletions(-) diff --git a/packages/query-core/src/__tests__/streamedQuery.test.tsx b/packages/query-core/src/__tests__/streamedQuery.test.tsx index 8768f442e0..6adced18ea 100644 --- a/packages/query-core/src/__tests__/streamedQuery.test.tsx +++ b/packages/query-core/src/__tests__/streamedQuery.test.tsx @@ -456,7 +456,7 @@ describe('streamedQuery', () => { expect(queryClient.getQueryState(key)).toMatchObject({ status: 'success', - fetchStatus: 'idle', + fetchStatus: 'fetching', data: [0, 1], }) }) diff --git a/packages/query-core/src/infiniteQueryBehavior.ts b/packages/query-core/src/infiniteQueryBehavior.ts index 9fc0f8954f..af9c50e550 100644 --- a/packages/query-core/src/infiniteQueryBehavior.ts +++ b/packages/query-core/src/infiniteQueryBehavior.ts @@ -30,7 +30,7 @@ export function infiniteQueryBehavior( const addSignalProperty = (object: unknown) => { addConsumeAwareSignal( object, - context.signal, + () => context.signal, () => (cancelled = true), ) } diff --git a/packages/query-core/src/streamedQuery.ts b/packages/query-core/src/streamedQuery.ts index 81163dee95..788a3c2c74 100644 --- a/packages/query-core/src/streamedQuery.ts +++ b/packages/query-core/src/streamedQuery.ts @@ -1,5 +1,10 @@ import { addConsumeAwareSignal, addToEnd } from './utils' -import type { QueryFunction, QueryFunctionContext, QueryKey } from './types' +import type { + OmitKeyof, + QueryFunction, + QueryFunctionContext, + QueryKey, +} from './types' type BaseStreamedQueryParams = { streamFn: ( @@ -74,10 +79,17 @@ export function streamedQuery< let result = initialValue let cancelled = false - const streamFnContext = { ...context } - addConsumeAwareSignal( - streamFnContext, - context.signal, + const streamFnContext = addConsumeAwareSignal< + OmitKeyof + >( + { + client: context.client, + meta: context.meta, + queryKey: context.queryKey, + pageParam: context.pageParam, + direction: context.direction, + }, + () => context.signal, () => (cancelled = true), ) diff --git a/packages/query-core/src/utils.ts b/packages/query-core/src/utils.ts index fb813f430d..d4ad8073c7 100644 --- a/packages/query-core/src/utils.ts +++ b/packages/query-core/src/utils.ts @@ -466,14 +466,15 @@ export function shouldThrowError) => boolean>( return !!throwOnError } -export function addConsumeAwareSignal( - object: unknown, - signal: AbortSignal, +export function addConsumeAwareSignal( + object: T, + getSignal: () => AbortSignal, onCancelled: VoidFunction, -) { +): T & { signal: AbortSignal } { Object.defineProperty(object, 'signal', { enumerable: true, get: () => { + const signal = getSignal() if (signal.aborted) { onCancelled() } else { @@ -484,4 +485,6 @@ export function addConsumeAwareSignal( return signal }, }) + + return object as T & { signal: AbortSignal } } From 7adee5455b8ef86ea5bc46f3fa35ba4e2e6dc94e Mon Sep 17 00:00:00 2001 From: razhayat Date: Fri, 12 Dec 2025 10:45:41 +0200 Subject: [PATCH 6/7] fix: made onCancelled get called only once --- packages/query-core/src/utils.ts | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/packages/query-core/src/utils.ts b/packages/query-core/src/utils.ts index d4ad8073c7..7fc0be85b2 100644 --- a/packages/query-core/src/utils.ts +++ b/packages/query-core/src/utils.ts @@ -471,17 +471,23 @@ export function addConsumeAwareSignal( getSignal: () => AbortSignal, onCancelled: VoidFunction, ): T & { signal: AbortSignal } { + let consumed = false + Object.defineProperty(object, 'signal', { enumerable: true, get: () => { const signal = getSignal() + if (consumed) { + return signal + } + + consumed = true if (signal.aborted) { onCancelled() } else { - signal.addEventListener('abort', () => { - onCancelled() - }) + signal.addEventListener('abort', onCancelled, { once: true }) } + return signal }, }) From 0346c4ec109ef026909dd6e2b1cac24ca13e6aab Mon Sep 17 00:00:00 2001 From: razhayat Date: Fri, 12 Dec 2025 12:10:12 +0200 Subject: [PATCH 7/7] fix: memoized getSignal and made sure it is called only once --- packages/query-core/src/utils.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/query-core/src/utils.ts b/packages/query-core/src/utils.ts index 7fc0be85b2..feee9c36c6 100644 --- a/packages/query-core/src/utils.ts +++ b/packages/query-core/src/utils.ts @@ -472,11 +472,12 @@ export function addConsumeAwareSignal( onCancelled: VoidFunction, ): T & { signal: AbortSignal } { let consumed = false + let signal: AbortSignal | undefined Object.defineProperty(object, 'signal', { enumerable: true, get: () => { - const signal = getSignal() + signal ??= getSignal() if (consumed) { return signal }