address document-context comments

Author: arturbien

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"
   },

packages/react/document-context/src/document-context.test.tsx

@@ -8,16 +8,6 @@ function TestComponent() {
 }
 
 describe('DocumentContext', () => {
-  it('provides default document when no custom document is provided', () => {
-    render(
-      <DocumentProvider>
-        <TestComponent />
-      </DocumentProvider>
-    );
-
-    expect(screen.getByText('Has document: true')).toBeInTheDocument();
-  });
-
   it('provides custom document when specified', () => {
     const mockDocument = {} as Document;
 
@@ -59,18 +49,20 @@ describe('DocumentContext', () => {
     expect(screen.getByText('Document ID: 2')).toBeInTheDocument();
   });
 
-  it('provides global document when no value is passed to provider', () => {
+  it('useDocument returns global document when no value is passed to provider', () => {
     function TestDocumentConsumer() {
       const doc = useDocument();
       return <div>{doc === globalThis.document ? 'global document' : 'other document'}</div>;
     }
 
-    render(
-      <DocumentProvider>
-        <TestDocumentConsumer />
-      </DocumentProvider>
-    );
+    render(<TestDocumentConsumer />);
 
     expect(screen.getByText('global document')).toBeInTheDocument();
   });
+
+  it('useDocument returns default document when custom document is not provided', () => {
+    render(<TestComponent />);
+
+    expect(screen.getByText('Has document: true')).toBeInTheDocument();
+  });
 });

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);
 }

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"