From 048c30fbba14bd84ef513b29e820b98d8bbd0a19 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Artur=20Bie=C5=84?= <artur.bien@gmail.com>
Date: Wed, 12 Feb 2025 23:26:48 +0100
Subject: [PATCH] address document-context comments

---
 packages/react/document-context/package.json  |  4 +++-
 .../document-context/src/document-context.tsx | 22 ++++++++++---------
 yarn.lock                                     | 18 +++++++++++++++
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/packages/react/document-context/package.json b/packages/react/document-context/package.json
index aabeaad58..baf8735d5 100644
--- a/packages/react/document-context/package.json
+++ b/packages/react/document-context/package.json
@@ -28,11 +28,13 @@
     "version": "yarn version"
   },
   "dependencies": {
-    "@radix-ui/react-primitive": "workspace:*"
+    "@radix-ui/react-primitive": "workspace:*",
+    "use-sync-external-store": "^1.4.0"
   },
   "devDependencies": {
     "@repo/typescript-config": "workspace:*",
     "@types/react": "^19.0.7",
+    "@types/use-sync-external-store": "^0.0.6",
     "react": "^19.0.0",
     "typescript": "^5.7.3"
   },
diff --git a/packages/react/document-context/src/document-context.tsx b/packages/react/document-context/src/document-context.tsx
index 00c730583..9b440b66c 100644
--- a/packages/react/document-context/src/document-context.tsx
+++ b/packages/react/document-context/src/document-context.tsx
@@ -1,23 +1,25 @@
 import * as React from 'react';
-
+import { useSyncExternalStore } from 'use-sync-external-store/shim';
 // Use null as initial value to handle SSR safely
 const DocumentContext = React.createContext<Document | null>(null);
 
 interface DocumentProviderProps {
-  document?: Document;
+  document: Document;
   children: React.ReactNode;
 }
 
-export function DocumentProvider({ document: doc, children }: DocumentProviderProps) {
-  const value = React.useMemo(
-    () => doc ?? (typeof document !== 'undefined' ? globalThis?.document : null),
-    [doc]
-  );
-  return <DocumentContext.Provider value={value}>{children}</DocumentContext.Provider>;
+export function DocumentProvider({ document, children }: DocumentProviderProps) {
+  return <DocumentContext.Provider value={document}>{children}</DocumentContext.Provider>;
 }
 
+const subscribe = () => () => {};
+
 export function useDocument() {
   const doc = React.useContext(DocumentContext);
-  // Return default document if available and no context value
-  return doc ?? (typeof document !== 'undefined' ? globalThis?.document : null);
+  const isHydrated = useSyncExternalStore(
+    subscribe,
+    () => true,
+    () => false
+  );
+  return doc ?? (isHydrated ? document : null);
 }
diff --git a/yarn.lock b/yarn.lock
index d9bd0f815..3eb2697e9 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2073,8 +2073,10 @@ __metadata:
     "@radix-ui/react-primitive": "workspace:*"
     "@repo/typescript-config": "workspace:*"
     "@types/react": "npm:^19.0.7"
+    "@types/use-sync-external-store": "npm:^0.0.6"
     react: "npm:^19.0.0"
     typescript: "npm:^5.7.3"
+    use-sync-external-store: "npm:^1.4.0"
   peerDependencies:
     "@types/react": "*"
     react: ^16.8 || ^17.0 || ^18.0 || ^19.0
@@ -4534,6 +4536,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@types/use-sync-external-store@npm:^0.0.6":
+  version: 0.0.6
+  resolution: "@types/use-sync-external-store@npm:0.0.6"
+  checksum: 10/a95ce330668501ad9b1c5b7f2b14872ad201e552a0e567787b8f1588b22c7040c7c3d80f142cbb9f92d13c4ea41c46af57a20f2af4edf27f224d352abcfe4049
+  languageName: node
+  linkType: hard
+
 "@types/uuid@npm:^9.0.1":
   version: 9.0.8
   resolution: "@types/uuid@npm:9.0.8"
@@ -13437,6 +13446,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"use-sync-external-store@npm:^1.4.0":
+  version: 1.4.0
+  resolution: "use-sync-external-store@npm:1.4.0"
+  peerDependencies:
+    react: ^16.8.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
+  checksum: 10/08bf581a8a2effaefc355e9d18ed025d436230f4cc973db2f593166df357cf63e47b9097b6e5089b594758bde322e1737754ad64905e030d70f8ff7ee671fd01
+  languageName: node
+  linkType: hard
+
 "util-deprecate@npm:^1.0.2":
   version: 1.0.2
   resolution: "util-deprecate@npm:1.0.2"