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