RFC: Server Context for React Server Components

fixtures/flight/src/App.js

@@ -38,6 +38,17 @@ async function ServerComponent() {
   await new Promise(resolve => setTimeout(() => resolve('deferred text'), 50));
 }
 
+const ServerContext = React.createServerContext(41);
+
+function ServerContextProvider(props) {
+  return <ServerContext value={props.value}>{props.children}</ServerContext>;
+}
+
+async function ServerContextConsumer(props) {
+  const value = React.useServerContext(ServerContext);
+  return <div>value: {value}<br/>{props.children}</div>;
+}
+
 export default async function App({prerender}) {
   const res = await fetch('http://localhost:3001/todos');
   const todos = await res.json();
@@ -89,6 +100,13 @@ export default async function App({prerender}) {
           <Note />
           <Foo>{dedupedChild}</Foo>
           <Bar>{Promise.resolve([dedupedChild])}</Bar>
+          <ServerContextProvider value={42}>
+            <ServerContextConsumer >
+              <ServerContextProvider value={43}>
+                <ServerContextConsumer />
+              </ServerContextProvider>
+            </ServerContextConsumer>
+          </ServerContextProvider>
         </Container>
       </body>
     </html>

packages/react-server/src/ReactFlightServer.js

@@ -120,6 +120,7 @@ import {
   REACT_LAZY_TYPE,
   REACT_MEMO_TYPE,
   REACT_POSTPONE_TYPE,
+  REACT_SERVER_CONTEXT_TYPE,
   ASYNC_ITERATOR,
 } from 'shared/ReactSymbols';
 
@@ -2476,6 +2477,29 @@ function renderModelDestructive(
           // Attribution on the client is still correct since it has a pop.
         }
 
+        // Handle Server Context Provider elements
+        if (element.type !== null && element.type.$$typeof === REACT_SERVER_CONTEXT_TYPE) {
+          const context = element.type;
+
+          const prevKeyPath = task.keyPath;
+          task.keyPath = task.keyPath !== null ? task.keyPath + ',' + element.key : element.key;
+
+          const result = ReactSharedInternals.setServerContextValue(
+            context,
+            props.value,
+            () => renderModelDestructive(
+              request,
+              task,
+              emptyRoot,
+              '',
+              props.children,
+            )
+          );
+
+          task.keyPath = prevKeyPath;
+          return result;
+        }
+
         const newChild = renderElement(
           request,
           task,

packages/react/src/ReactContextRegistryServer.js

@@ -0,0 +1,62 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import {AsyncLocalStorage} from 'async_hooks';
+import type {ReactContext} from 'shared/ReactTypes';
+
+// Create a dedicated AsyncLocalStorage for server contexts
+export const SERVER_CONTEXT_STORAGE = new AsyncLocalStorage();
+
+/**
+ * Retrieves the current value of a server context from the registry.
+ * Falls back to the default value if not found.
+ */
+export function getServerContextValue<T>(context: ReactContext<T>): T {
+  // Get the current store from AsyncLocalStorage
+  const store = SERVER_CONTEXT_STORAGE.getStore();
+  if (!store || !store.has(context._serverContextId)) {
+    // No value in the store, use the default value
+    return context._currentValue;
+  }
+
+  // Return the value from the store
+  return store.get(context._serverContextId);
+}
+
+/**
+ * Sets a server context value and runs a callback with that value.
+ * Uses AsyncLocalStorage to ensure the context is available throughout
+ * the asynchronous server component rendering process.
+ */
+export function setServerContextValue<T, R>(
+  context: ReactContext<T>,
+  value: T,
+  callback: () => R,
+): R {
+  if (!context._isServerContext) {
+    // Not a server context, use the normal context behavior
+    const prevValue = context._currentValue;
+    context._currentValue = value;
+    try {
+      return callback();
+    } finally {
+      context._currentValue = prevValue;
+    }
+  }
+
+  // Get the existing store or create a new one
+  const prevStore = SERVER_CONTEXT_STORAGE.getStore() || new Map();
+  const nextStore = new Map(prevStore);
+
+  // Set the new value in the store
+  nextStore.set(context._serverContextId, value);
+
+  // Run the callback with the new context value in the store
+  return SERVER_CONTEXT_STORAGE.run(nextStore, callback);
+}

packages/react/src/ReactContextServer.js

@@ -0,0 +1,81 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import {REACT_SERVER_CONTEXT_TYPE, REACT_PROVIDER_TYPE} from 'shared/ReactSymbols';
+import type {ReactContext, ReactProviderType} from 'shared/ReactTypes';
+import ReactSharedInternals from './ReactSharedInternalsServer';
+
+/**
+ * Creates a context object that works with React Server Components.
+ * Uses AsyncLocalStorage to maintain context across async operations.
+ */
+export function createServerContext<T>(
+  defaultValue: T,
+  displayName?: string,
+): ReactContext<T> {
+  const contextId = Symbol('ServerContext');
+
+  const context: ReactContext<T> = {
+    $$typeof: REACT_SERVER_CONTEXT_TYPE,
+    // Maintain compatibility with existing context API
+    _currentValue: defaultValue,
+    _currentValue2: defaultValue,
+    _threadCount: 0,
+
+    // Server context specific properties
+    _serverContextId: contextId,
+    _isServerContext: true,
+
+    // Will be populated below
+    Provider: (null: any),
+    Consumer: (null: any),
+  };
+
+  // Create the Provider component
+  const Provider: ReactProviderType<T> = {
+    $$typeof: REACT_PROVIDER_TYPE,
+    _context: context,
+  };
+
+  context.Provider = Provider;
+
+  if (__DEV__) {
+    context._debugName = displayName || 'ServerContext';
+  }
+
+  return context;
+}
+
+/**
+ * Returns the current value of a server context within the AsyncLocalStorage.
+ * Falls back to the default value if not found.
+ */
+export function useServerContext<T>(context: ReactContext<T>): T {
+  if (!context._isServerContext) {
+    throw new Error(
+      'useServerContext: Expected a context created with createServerContext. ' +
+      'Did you use React.createContext instead?'
+    );
+  }
+
+  return ReactSharedInternals.getServerContextValue(context);
+}
+
+/**
+ * Sets a server context value and runs a callback with that value.
+ * Uses AsyncLocalStorage to ensure the context is available throughout
+ * the asynchronous server component rendering process.
+ */
+export function setServerContext<T, R>(
+  context: ReactContext<T>,
+  value: T,
+  callback: () => R,
+): R {
+  return ReactSharedInternals.setServerContextValue(context, value, callback);
+}

packages/react/src/ReactServer.experimental.development.js

@@ -40,6 +40,10 @@ import {startTransition} from './ReactStartTransition';
 import {postpone} from './ReactPostpone';
 import {captureOwnerStack} from './ReactOwnerStack';
 import version from 'shared/ReactVersion';
+import {
+  createServerContext,
+  useServerContext,
+} from './ReactContextServer';
 
 const Children = {
   map,
@@ -82,4 +86,6 @@ export {
   REACT_SUSPENSE_LIST_TYPE as unstable_SuspenseList,
   REACT_VIEW_TRANSITION_TYPE as unstable_ViewTransition,
   captureOwnerStack, // DEV-only
+  createServerContext,
+  useServerContext,
 };

packages/react/src/ReactServer.experimental.js

@@ -40,6 +40,12 @@ import {cache} from './ReactCacheServer';
 import {startTransition} from './ReactStartTransition';
 import {postpone} from './ReactPostpone';
 import version from 'shared/ReactVersion';
+// Explicitly import and re-export the server context functionality
+import {
+  createServerContext,
+  useServerContext,
+  setServerContext,
+} from './ReactContextServer';
 
 const Children = {
   map,
@@ -82,4 +88,8 @@ export {
   REACT_SUSPENSE_LIST_TYPE as unstable_SuspenseList,
   REACT_VIEW_TRANSITION_TYPE as unstable_ViewTransition,
   REACT_ACTIVITY_TYPE as unstable_Activity,
+  createServerContext,
+  useServerContext,
+  // Also export the internal setServerContext function to ensure it's available
+  setServerContext,
 };

packages/react/src/ReactSharedInternalsServer.js

@@ -23,6 +23,12 @@ import {
   TaintRegistryPendingRequests,
 } from './ReactTaintRegistry';
 
+import {
+  SERVER_CONTEXT_STORAGE,
+  getServerContextValue,
+  setServerContextValue,
+} from './ReactContextRegistryServer';
+
 import {enableTaint} from 'shared/ReactFeatureFlags';
 
 export type SharedStateServer = {
@@ -35,6 +41,11 @@ export type SharedStateServer = {
   TaintRegistryByteLengths: Set<number>,
   TaintRegistryPendingRequests: Set<RequestCleanupQueue>,
 
+  // Server Context
+  SERVER_CONTEXT_STORAGE: AsyncLocalStorage<any>,
+  getServerContextValue: typeof getServerContextValue,
+  setServerContextValue: typeof setServerContextValue,
+
   // DEV-only
 
   // ReactDebugCurrentFrame
@@ -59,6 +70,11 @@ if (enableTaint) {
     TaintRegistryPendingRequests;
 }
 
+// Server Context Registry
+ReactSharedInternals.SERVER_CONTEXT_STORAGE = SERVER_CONTEXT_STORAGE;
+ReactSharedInternals.getServerContextValue = getServerContextValue;
+ReactSharedInternals.setServerContextValue = setServerContextValue;
+
 if (__DEV__) {
   // Stack implementation injected by the current renderer.
   ReactSharedInternals.getCurrentStack = (null: null | (() => string));

packages/shared/ReactSymbols.js

@@ -25,6 +25,7 @@ export const REACT_PROFILER_TYPE: symbol = Symbol.for('react.profiler');
 export const REACT_PROVIDER_TYPE: symbol = Symbol.for('react.provider'); // TODO: Delete with enableRenderableContext
 export const REACT_CONSUMER_TYPE: symbol = Symbol.for('react.consumer');
 export const REACT_CONTEXT_TYPE: symbol = Symbol.for('react.context');
+export const REACT_SERVER_CONTEXT_TYPE: symbol = Symbol.for('react.server_context');
 export const REACT_FORWARD_REF_TYPE: symbol = Symbol.for('react.forward_ref');
 export const REACT_SUSPENSE_TYPE: symbol = Symbol.for('react.suspense');
 export const REACT_SUSPENSE_LIST_TYPE: symbol = Symbol.for(