fix: fixed non cancellation on unmount

Author: mbret

src/lib/queries/QueryClientProvider$.tsx

@@ -1,4 +1,9 @@
-import { hashKey, type QueryKey, useQueryClient } from "@tanstack/react-query"
+import {
+  hashKey,
+  type QueryClient,
+  type QueryKey,
+  useQueryClient,
+} from "@tanstack/react-query"
 import { createContext, memo, useContext, useEffect, useState } from "react"
 import {
   fromEvent,
@@ -9,6 +14,7 @@ import {
 } from "rxjs"
 
 type CacheEntry = {
+  queryKey: QueryKey
   query$: Observable<unknown>
   signal: AbortSignal
   sub: Subscription | undefined
@@ -19,6 +25,8 @@ type CacheEntry = {
 export class QueryClient$ {
   public readonly queryMap: Map<string, CacheEntry> = new Map()
 
+  constructor(public readonly queryClient: QueryClient) {}
+
   getQuery(queryHash: string) {
     return this.queryMap.get(queryHash)
   }
@@ -31,17 +39,12 @@ export class QueryClient$ {
     const queryHash = hashKey(queryKey)
 
     const sharedQuery$ = query$.pipe(
-      /**
-       * abort signal is triggered on:
-       * - manual cancellation from user
-       * - unmounting the component
-       * @see https://tanstack.com/query/latest/docs/framework/react/guides/query-cancellation
-       */
       takeUntil(fromEvent(signal, "abort")),
       share(),
     )
 
     const cacheEntry: CacheEntry = {
+      queryKey,
       query$: sharedQuery$,
       signal,
       sub: undefined,
@@ -60,7 +63,9 @@ export class QueryClient$ {
         }
       },
       complete: () => {
-        this.deleteQuery(queryHash)
+        if (this.queryMap.get(queryHash) === cacheEntry) {
+          this.deleteQuery(queryHash)
+        }
       },
     })
 
@@ -74,14 +79,29 @@ export class QueryClient$ {
 
     if (!entry) return
 
+    entry.isCompleted = true
+
+    this.queryMap.delete(queryHash)
+
     if (entry.sub) {
       entry.sub.unsubscribe()
       entry.sub = undefined
     }
 
-    entry.isCompleted = true
-
-    this.queryMap.delete(queryHash)
+    /**
+     * Only cancel queries whose stream already emitted at least once.
+     * Those are "live" streams stuck in the take(1) refetch loop and
+     * need an explicit cancel so the pending queryFnAsync promise
+     * settles. One-shot observables (used like promises) that haven't
+     * resolved yet should keep react-query's default behavior: stay
+     * pending and resolve naturally even after unmount.
+     */
+    if (!entry.signal.aborted && entry.lastData !== undefined) {
+      this.queryClient?.cancelQueries({
+        queryKey: entry.queryKey,
+        exact: true,
+      })
+    }
   }
 
   destroy() {
@@ -99,44 +119,42 @@ export const Context = createContext<QueryClient$ | undefined>(undefined)
  * (including via `queryClient.clear()`), matching `useQuery$` observable
  * entries are dropped.
  */
-export const QueryClientProvider$ = memo(
-  ({
-    children,
-    client: _client,
-  }: {
-    children: React.ReactNode
-    client?: QueryClient$
-  }) => {
-    const [client] = useState(() => _client ?? new QueryClient$())
-    const queryClient = useQueryClient()
-
-    useEffect(
-      function subscribeToQueryCache() {
-        return queryClient.getQueryCache().subscribe((event) => {
-          const activeObservers = event.query.getObserversCount()
-
-          /**
-           * When observers unmount, we need to delete the query from the cache
-           * because if the query contains a stream, it will continue. Nothing stops it.
-           * However this is only valid when there is no more observers.
-           */
-          if (event.type === "observerRemoved" && activeObservers === 0) {
-            client.deleteQuery(event.query.queryHash)
-          }
-        })
-      },
-      [queryClient, client],
-    )
-
-    useEffect(() => {
-      return () => {
-        client.destroy()
-      }
-    }, [client])
+export const QueryClientProvider$ = memo(function QueryClientProvider$({
+  children,
+  client: _client,
+}: {
+  children: React.ReactNode
+  client?: QueryClient$
+}) {
+  const queryClient = useQueryClient()
+  const [client] = useState(() => _client ?? new QueryClient$(queryClient))
+
+  useEffect(
+    function subscribeToQueryCache() {
+      return queryClient.getQueryCache().subscribe((event) => {
+        const activeObservers = event.query.getObserversCount()
+
+        /**
+         * When all observers unmount we need to delete the query from the cache
+         * because if the query contains a stream, it will continue. Nothing stops it.
+         * However this is only valid when there is no more observers.
+         */
+        if (event.type === "observerRemoved" && activeObservers === 0) {
+          client.deleteQuery(event.query.queryHash)
+        }
+      })
+    },
+    [queryClient, client],
+  )
+
+  useEffect(() => {
+    return () => {
+      client.destroy()
+    }
+  }, [client])
 
-    return <Context.Provider value={client}>{children}</Context.Provider>
-  },
-)
+  return <Context.Provider value={client}>{children}</Context.Provider>
+})
 
 export const useQueryClient$ = () => {
   const client = useContext(Context)