keep subscription on data while query is running

Author: phryneas

packages/toolkit/src/query/core/buildMiddleware/batchActions.ts

@@ -45,13 +45,23 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
       const {
         meta: { arg, requestId },
       } = action
+      const substate = (mutableState[arg.queryCacheKey] ??= {})
+      substate[`${requestId}_running`] = {}
       if (arg.subscribe) {
-        const substate = (mutableState[arg.queryCacheKey] ??= {})
         substate[requestId] =
           arg.subscriptionOptions ?? substate[requestId] ?? {}
-
-        return true
       }
+      return true
+    }
+    let mutated = false
+    if (
+      queryThunk.fulfilled.match(action) ||
+      queryThunk.rejected.match(action)
+    ) {
+      const state = mutableState[action.meta.arg.queryCacheKey] || {}
+      const key = `${action.meta.requestId}_running`
+      mutated ||= !!state[key]
+      delete state[key]
     }
     if (queryThunk.rejected.match(action)) {
       const {
@@ -62,11 +72,11 @@ export const buildBatchedActionsHandler: InternalHandlerBuilder<
         substate[requestId] =
           arg.subscriptionOptions ?? substate[requestId] ?? {}
 
-        return true
+        mutated = true
       }
     }
 
-    return false
+    return mutated
   }
 
   return (

packages/toolkit/src/query/core/buildSlice.ts

@@ -147,12 +147,9 @@ export function buildSlice({
       builder
         .addCase(queryThunk.pending, (draft, { meta, meta: { arg } }) => {
           const upserting = isUpsertQuery(arg)
-          if (arg.subscribe || upserting) {
-            // only initialize substate if we want to subscribe to it
-            draft[arg.queryCacheKey] ??= {
-              status: QueryStatus.uninitialized,
-              endpointName: arg.endpointName,
-            }
+          draft[arg.queryCacheKey] ??= {
+            status: QueryStatus.uninitialized,
+            endpointName: arg.endpointName,
           }
 
           updateQuerySubstateIfExists(draft, arg.queryCacheKey, (substate) => {

packages/toolkit/src/query/tests/buildInitiate.test.tsx

@@ -13,6 +13,12 @@ const api = createApi({
         return { data }
       },
     }),
+    failing: build.query<void, void>({
+      async queryFn() {
+        await Promise.resolve()
+        return { error: { status: 500, data: 'error' } }
+      },
+    }),
   }),
 })
 
@@ -52,3 +58,78 @@ test('multiple synchonrous initiate calls with pre-existing cache entry', async
     requestId: thirdValue.requestId,
   })
 })
+
+describe('calling initiate without a cache entry, with subscribe: false still returns correct values', () => {
+  test('successful query', async () => {
+    const { store, api } = storeRef
+    calls = 0
+    const promise = store.dispatch(
+      api.endpoints.increment.initiate(undefined, { subscribe: false })
+    )
+    expect(
+      store.dispatch(
+        api.internalActions.internal_probeSubscription({
+          queryCacheKey: 'increment(undefined)',
+          requestId: promise.requestId,
+        })
+      )
+    ).toBe(false)
+    expect(
+      store.dispatch(
+        api.internalActions.internal_probeSubscription({
+          queryCacheKey: 'increment(undefined)',
+          requestId: `${promise.requestId}_running`,
+        })
+      )
+    ).toBe(true)
+
+    await expect(promise).resolves.toMatchObject({
+      data: 0,
+      status: 'fulfilled',
+    })
+    expect(
+      store.dispatch(
+        api.internalActions.internal_probeSubscription({
+          queryCacheKey: 'increment(undefined)',
+          requestId: `${promise.requestId}_running`,
+        })
+      )
+    ).toBe(false)
+  })
+
+  test('rejected query', async () => {
+    const { store, api } = storeRef
+    calls = 0
+    const promise = store.dispatch(
+      api.endpoints.failing.initiate(undefined, { subscribe: false })
+    )
+    expect(
+      store.dispatch(
+        api.internalActions.internal_probeSubscription({
+          queryCacheKey: 'failing(undefined)',
+          requestId: promise.requestId,
+        })
+      )
+    ).toBe(false)
+    expect(
+      store.dispatch(
+        api.internalActions.internal_probeSubscription({
+          queryCacheKey: 'failing(undefined)',
+          requestId: `${promise.requestId}_running`,
+        })
+      )
+    ).toBe(true)
+
+    await expect(promise).resolves.toMatchObject({
+      status: 'rejected',
+    })
+    expect(
+      store.dispatch(
+        api.internalActions.internal_probeSubscription({
+          queryCacheKey: 'failing(undefined)',
+          requestId: `${promise.requestId}_running`,
+        })
+      )
+    ).toBe(false)
+  })
+})