RSC setState equivalent. Streaming server subscription re-renders via AsyncIterable ReactElements over the wire via RSC protocol and React 19

FragmentSubscription.tsx

import "server-only"

import { StreamingFragment } from "@frens/reframe/Micro.ClientComponents"
import { ReactElement } from "react"
import { AsyncIterable$forEach } from "../etc/AsyncIterable$forEach"
import { MaybePromise, RenderQ } from "../types"

interface FragmentSubscriptionProps<Event> {
  signal: AbortSignal
  fallback?: ReactElement
  subscriptions: Array<AsyncIterable<Event>>
  render: (props?: Event) => MaybePromise<ReactElement | void>
}

// TODO: document this
export default function FragmentSubscription<Event>(props: FragmentSubscriptionProps<Event>) {
  const { signal, subscriptions, fallback, render } = props
  if (signal.aborted) return null
  const { readable, writable } = new TransformStream<RenderQ>()

  const dispose = () => (writer = null)
  let writer: null | WritableStreamDefaultWriter<RenderQ> = writable.getWriter()
  signal.addEventListener("abort", dispose)
  writer.closed.then(dispose, dispose)

  if (fallback) writer.write({ replace: fallback })

  async function refresh(props?: Event) {
    if (!writer?.write) return
    const children = await render(props)
    if (!children) return
    writer?.write({ replace: children })
  }

  void refresh() // initial load

  subscriptions.forEach((sub) => {
    AsyncIterable$forEach(sub, signal, refresh).catch((cause) => {
      signal.throwIfAborted()
      if (signal.aborted) return
      console.error(cause)
    })
  })

  return <StreamingFragment>{readable}</StreamingFragment>
}

StreamingFragment.tsx

"use client"

import { isConnectionClosedException } from "@frens/reframe/client/isConnectionClosedException"
import { AsyncIterable$forEach } from "@frens/reframe/etc/AsyncIterable$forEach"
import { isAbortError } from "@frens/reframe/etc/isAbortError"
import { ReactElement, ReactNode, useDeferredValue, useEffect, useState } from "react"
import { useStableValue } from "./useStableValue"

type StreamingNodeOrCommand =
  | ReactElement // not ReactNode so that type is always an object
  | { append?: undefined; replace: ReactNode }
  | { replace?: undefined; append: ReactNode }

interface StreamingFragmentProps {
  /**
   * render before the stream sends the first child
   * NOTE: this is only read on first render; changes to this prop will be silently ignored
   */
  initialChildren?: ReactElement
  children: AsyncIterable<StreamingNodeOrCommand>
}

/**
 * @private internal implementation
 *
 * # Streaming Fragment
 * The practical upshot of rendering sequentially in time is that you can effectively achieve
 * subscription based re-rendering from the server.
 * This is useful for streaming in AI LLM responses word by word.
 * {@link StreamingFragment} is the client-side counterpart to `FragmentSubscription`
 *
 * {@link React.Fragment} lets you synchronously group {@link ReactNode}s in space without a wrapper node.
 * {@link StreamingFragment} lets you Asynchronously group {@link ReactNode}s in space AND time!
 *
 * Whereas a normal {@link React.Fragment} can render a synchronous {@link Iterable} list of children,
 * a {@link StreamingFragment} can render an {@link AsyncIterable} list of children.
 *
 * ## append / replace === space / time
 * Normally, each child in the props.children is rendered sequentially in space.
 * However, if a child is an object with a `replace` property, the previous children are replaced with the new child instead of appended.
 * This effectively renders each child sequentially in time instead of space.
 *
 * TL;DR -- use `append` for stuff like AI LLM streaming word by word,
 *          use `replace` for stuff like subscription based re-rendering.
 *
 * ## abort vs error handling
 * EASY: When rendering synchronously in space, it's easy to know when to stop rendering.
 * HARD: When rendering asynchronously in time, it's hard to know when to stop rendering.
 *
 * When the iterator is aborted, this is fine, just stop changing anything. UI stays valid.
 *
 * When the iterator throws an error, this is bad, the UI is now invalid. The error bubbles up to the nearest error boundary.
 *
 * When the cause was because the connection was closed, this is fine, just stop changing anything. UI stays valid.
 */
export default function StreamingFragment(props: StreamingFragmentProps) {
  return useDeferredValue(useRemoteControlledChildren(props))
}

/** @private internal implementation */
function useRemoteControlledChildren(props: StreamingFragmentProps): ReactNode {
  // TODO: DX; maybe someday; warn if initialChildren changes after first render
  const originalView = props.initialChildren ?? null
  const updateStream = useStableValue(`${useRemoteControlledChildren.name} props.children`, props.children)

  const [currentView, updateChildren] = useState<ReactNode[]>([originalView])

  useEffect(() => {
    const mounted = new AbortController()

    /**
     * CRITICAL: DO NOT EVER re-render {@link originalView}
     * (even if the {@link updateStream} changes)
     *
     * Why not?
     * Conceptually, the {@link originalView} may be an earlier state of the same view.
     * resetting to an earlier state would be totally bizarre and confusing to the user.
     * Imagine if this was rendering live price data.
     * Resetting to an earlier state could make it appear as if the price
     */
    // setChildren([props.initialChildren]) // <-- DO NOT DO THIS!

    // render each child as it arrives
    AsyncIterable$forEach(updateStream, mounted.signal, onEach).then(onDone, onException)

    return function onUnmount() {
      mounted.abort("unmount")
    }

    function onEach(child: StreamingNodeOrCommand) {
      if (mounted.signal.aborted) return
      if ("replace" in child) return updateChildren([child.replace])
      if ("append" in child) return updateChildren((children) => [...children, child.append])
      updateChildren((children) => [...children, child])
    }

    function onDone() {
      console.info(`${useRemoteControlledChildren.name} stream ended normally`)
    }

    // NOTE: expected reasons will not trigger a re-render
    // unexpected reasons bubble up to the nearest error boundary
    function onException(cause: unknown) {
      // component was unmounted; nothing to do
      if (mounted.signal.aborted) return

      // the stream was aborted somehow; the next child will never arrive
      // this is expected behavior; keep the last sent child rendered; UI stays valid
      if (isAbortError(cause)) return console.info(`${useRemoteControlledChildren.name} aborted`)

      // RSC stream was closed; the next child will never arrive
      // this is expected behavior; keep the last sent child rendered; UI stays valid
      if (isConnectionClosedException(cause))
        return console.info(`${useRemoteControlledChildren.name} connection closed`)

      // unexpected error; the UI is now invalid; pass rendering responsibility up to the nearest error boundary
      console.error(`${useRemoteControlledChildren.name} unexpected error`, cause)
      updateChildren(() => {
        // bubble up to the nearest React error boundary
        throw cause
      })
    }
  }, [updateStream])

  return currentView
}