server-context.ts

import React, { cache } from 'react'
import 'server-only'

type ContextType<Value> = React.FC<{
  /** The children of the context provider. */
  children: React.ReactNode

  /** Sets the context value for its descendants. */
  value: Value
}> & {
  /** Fallback when nothing has been set */
  initialValue: Value

  /** Server‐side cache getter */
  getCache: () => { value?: Value }
}

/** Creates a context value provider with an initial value. */
export function createContext<Value>(initialValue?: Value): ContextType<Value> {
  const getCache = cache<() => { value?: Value }>(() => ({ value: undefined }))

  /** Sets the context value for its descendants. */
  function Context({
    children,
    value,
  }: {
    children: React.ReactNode
    value: Value
  }) {
    getCache().value = value
    return children
  }

  Object.assign(Context, {
    initialValue,
    getCache,
  })

  return Context as ContextType<Value>
}

/** Returns the current context value. */
export function getContext<Value>(Context: ContextType<Value>): Value {
  const cacheValue = Context.getCache().value
  return (cacheValue === undefined ? Context.initialValue : cacheValue) as Value
}