fix(query-core): have revertState capture manual updates (#9558)

Author: TkDodo

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

@@ -350,6 +350,47 @@ describe('streamedQuery', () => {
     unsubscribe()
   })
 
+  test('should abort when unsubscribed', async () => {
+    const key = queryKey()
+    const observer = new QueryObserver(queryClient, {
+      queryKey: key,
+      queryFn: streamedQuery({
+        queryFn: (context) => {
+          // just consume the signal
+          // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+          const numbers = context.signal ? 3 : 0
+          return createAsyncNumberGenerator(numbers)
+        },
+      }),
+    })
+
+    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(10)
+
+    expect(queryClient.getQueryState(key)).toMatchObject({
+      status: 'success',
+      fetchStatus: 'idle',
+      data: [0],
+    })
+  })
+
   test('should support maxChunks', async () => {
     const key = queryKey()
     const observer = new QueryObserver(queryClient, {

packages/query-core/src/query.ts

@@ -604,22 +604,26 @@ export class Query<
             fetchMeta: action.meta ?? null,
           }
         case 'success':
-          // If fetching ends successfully, we don't need revertState as a fallback anymore.
-          this.#revertState = undefined
-          return {
+          const newState = {
             ...state,
             data: action.data,
             dataUpdateCount: state.dataUpdateCount + 1,
             dataUpdatedAt: action.dataUpdatedAt ?? Date.now(),
             error: null,
             isInvalidated: false,
-            status: 'success',
+            status: 'success' as const,
             ...(!action.manual && {
-              fetchStatus: 'idle',
+              fetchStatus: 'idle' as const,
               fetchFailureCount: 0,
               fetchFailureReason: null,
             }),
           }
+
+          // If fetching ends successfully, we don't need revertState as a fallback anymore.
+          // For manual updates, capture the state to revert to it in case of a cancellation.
+          this.#revertState = action.manual ? newState : undefined
+
+          return newState
         case 'error':
           const error = action.error