From 1d10b6efc566a3285f3be90e2d7d16eac7b157b6 Mon Sep 17 00:00:00 2001
From: Jerel Miller <jerelmiller@gmail.com>
Date: Wed, 18 Sep 2024 13:55:43 -0600
Subject: [PATCH] Handle different cache ids in useSuspenseFragment

---
 src/react/hooks/useSuspenseFragment.ts    | 15 +++++++++++----
 src/react/internal/cache/SuspenseCache.ts |  4 ++--
 src/react/internal/cache/types.ts         |  6 ++++++
 3 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/react/hooks/useSuspenseFragment.ts b/src/react/hooks/useSuspenseFragment.ts
index 47244ff89f4..45c708fb2dc 100644
--- a/src/react/hooks/useSuspenseFragment.ts
+++ b/src/react/hooks/useSuspenseFragment.ts
@@ -8,9 +8,8 @@ import { canonicalStringify } from "../../cache/index.js";
 import type { Cache } from "../../cache/index.js";
 import { useApolloClient } from "./useApolloClient.js";
 import { getSuspenseCache } from "../internal/index.js";
-import type { CacheKey } from "../internal/index.js";
-import React from "rehackt";
-import type { FragmentKey } from "../internal/cache/types.js";
+import React, { useMemo } from "rehackt";
+import type { FragmentCacheKey, FragmentKey } from "../internal/cache/types.js";
 import { __use } from "./internal/__use.js";
 import { wrapHook } from "./internal/index.js";
 
@@ -59,8 +58,16 @@ function _useSuspenseFragment<
   options: UseSuspenseFragmentOptions<TData, TVariables>
 ): UseSuspenseFragmentResult<TData> {
   const client = useApolloClient(options.client);
+  const { from } = options;
+  const { cache } = client;
 
-  const cacheKey: CacheKey = [
+  const id = useMemo(
+    () => (typeof from === "string" ? from : cache.identify(from)),
+    [cache, from]
+  )!;
+
+  const cacheKey: FragmentCacheKey = [
+    id,
     options.fragment,
     canonicalStringify(options.variables),
   ];
diff --git a/src/react/internal/cache/SuspenseCache.ts b/src/react/internal/cache/SuspenseCache.ts
index b96de20431d..7ae77b49473 100644
--- a/src/react/internal/cache/SuspenseCache.ts
+++ b/src/react/internal/cache/SuspenseCache.ts
@@ -6,7 +6,7 @@ import type {
 import type { Observable } from "../../../utilities/index.js";
 import { canUseWeakMap } from "../../../utilities/index.js";
 import { InternalQueryReference } from "./QueryReference.js";
-import type { CacheKey } from "./types.js";
+import type { CacheKey, FragmentCacheKey } from "./types.js";
 import { FragmentReference } from "./FragmentReference.js";
 
 export interface SuspenseCacheOptions {
@@ -58,7 +58,7 @@ export class SuspenseCache {
   }
 
   getFragmentRef<TData>(
-    cacheKey: CacheKey,
+    cacheKey: FragmentCacheKey,
     createObservable: () => Observable<WatchFragmentResult<TData>>
   ) {
     const ref = this.fragmentRefs.lookupArray(cacheKey) as {
diff --git a/src/react/internal/cache/types.ts b/src/react/internal/cache/types.ts
index f45a55eda8f..a163431ad9d 100644
--- a/src/react/internal/cache/types.ts
+++ b/src/react/internal/cache/types.ts
@@ -6,6 +6,12 @@ export type CacheKey = [
   ...queryKey: any[],
 ];
 
+export type FragmentCacheKey = [
+  cacheId: string,
+  fragment: DocumentNode,
+  stringifiedVariables: string,
+];
+
 export interface QueryKey {
   __queryKey?: string;
 }