fix(server-functions): prevent HMR race condition by regenerating virtual manifest module from directiveFnsById state

[Crsk]

This error occurred when triggering server functions during development

The idea basically is to avoid manual module invalidation (which is async and prone to race conditions), relying on Vite's built-in mechanisms

Fixes #4486

didn’t add e2e test because the race condition does not happen in the production build environment

to see if working, i setup a cookies loaded theme, the expected result is to hit the theme toggle, see the style update and no Error: Server function module export not resolved for serverFn ID: xxx error.

---

Here's the theme setup i used

// theme-provider.tsx

import { useRouter } from '@tanstack/react-router'
import { createContext, PropsWithChildren, useContext, useState } from 'react'
import { setThemeServerFn } from './theme'

export type Theme = 'light' | 'dark' | 'system'

interface ThemeContextValue {
  currentTheme: Theme
  setTheme: (theme: Theme) => Promise<void>
}

const ThemeContext = createContext<ThemeContextValue | null>(null)

export const ThemeProvider = ({ children, initialTheme }: PropsWithChildren<{ initialTheme: Theme }>) => {
  const [currentTheme, setCurrentTheme] = useState<Theme>(initialTheme)
  const router = useRouter()

  const setTheme = async (theme: Theme) => {
    setThemeServerFn({ data: theme })
    setCurrentTheme(theme)
    router.invalidate()
  }

  return <ThemeContext.Provider value={{ currentTheme, setTheme }}>{children}</ThemeContext.Provider>
}

export const useTheme = () => {
  const context = useContext(ThemeContext)
  if (!context) throw new Error('useTheme must be used within a ThemeProvider')

  return context
}

// theme.ts

import { type Theme } from './theme-provider'
import { createServerFn } from '@tanstack/react-start'
import { getCookie, setCookie } from '@tanstack/react-start/server'

export const THEME_COOKIE_KEY = 'ui-theme'
export const THEME_VALUES = ['light', 'dark', 'system'] as const

export const getThemeServerFn = createServerFn().handler(async () => {
  return (getCookie(THEME_COOKIE_KEY) || 'light') as Theme
})

export const setThemeServerFn = createServerFn({ method: 'POST' })
  .validator((data: Theme) => {
    if (!THEME_VALUES.includes(data)) throw new Error('Invalid theme provided')

    return data
  })
  .handler(async ({ data }) => {
    setCookie(THEME_COOKIE_KEY, data)
  })

// theme-toggle.tsx

import { useTheme } from '~/theme-provider'


export const ThemeToggle = () => {
  const { currentTheme, setTheme } = useTheme()
  const toggle = () => setTheme(currentTheme === 'dark' ? 'light' : 'dark')

  return (
    <div className="p-2">
      <button onClick={toggle} className="bg-emerald-500 text-white px-2 py-1 rounded uppercase font-black text-sm">
        {currentTheme === 'light' ? 'Dark' : 'Light'}
      </button>
    </div>
  )
}

// __root.tsx

/// <reference types="vite/client" />
import {
  HeadContent,
  Link,
  Outlet,
  Scripts,
  createRootRoute,
} from '@tanstack/react-router'
import { TanStackRouterDevtools } from '@tanstack/react-router-devtools'
import * as React from 'react'
import { DefaultCatchBoundary } from '~/components/DefaultCatchBoundary'
import { NotFound } from '~/components/NotFound'
import appCss from '~/styles/app.css?url'
import { getThemeServerFn } from '~/theme'
import { seo } from '~/utils/seo'
import { ThemeProvider, useTheme } from '~/theme-provider'
import { ThemeToggle } from '~/theme-toggle'
import { useMemo } from 'react'

export const Route = createRootRoute({
  head: () => ({
    meta: [
      {
        charSet: 'utf-8',
      },
      {
        name: 'viewport',
        content: 'width=device-width, initial-scale=1',
      },
      ...seo({
        title:
          'TanStack Start | Type-Safe, Client-First, Full-Stack React Framework',
        description: `TanStack Start is a type-safe, client-first, full-stack React framework. `,
      }),
    ],
    links: [
      { rel: 'stylesheet', href: appCss },
      {
        rel: 'apple-touch-icon',
        sizes: '180x180',
        href: '/apple-touch-icon.png',
      },
      {
        rel: 'icon',
        type: 'image/png',
        sizes: '32x32',
        href: '/favicon-32x32.png',
      },
      {
        rel: 'icon',
        type: 'image/png',
        sizes: '16x16',
        href: '/favicon-16x16.png',
      },
      { rel: 'manifest', href: '/site.webmanifest', color: '#fffff' },
      { rel: 'icon', href: '/favicon.ico' },
    ],
    scripts: [
      {
        src: '/customScript.js',
        type: 'text/javascript',
      },
    ],
  }),
  errorComponent: DefaultCatchBoundary,
  notFoundComponent: () => <NotFound />,
  shellComponent: RootComponent,
  loader: () => getThemeServerFn()
})

function RootComponent() {
  const initialTheme = Route.useLoaderData()

  return (
    <ThemeProvider initialTheme={initialTheme}>
      <RootDocument>
        <Outlet />
      </RootDocument>
    </ThemeProvider>
  )
}

function RootDocument({ children }: { children: React.ReactNode }) {
  const { currentTheme } = useTheme()
  const colorMap = useMemo(() => ({
    bg: currentTheme === 'dark' ? 'black' : 'white',
    color: currentTheme === 'dark' ? 'white' : 'black',
    hr: currentTheme === 'dark' ? '#161616' : '#ededed',
  }), [currentTheme])

  return (
    <html lang="en" style={{ backgroundColor: colorMap.bg }}>
      <head>
        <HeadContent />
      </head>
      <body style={{ backgroundColor: colorMap.bg, color: colorMap.color }}>
        <div className="p-2 flex gap-2 text-lg">
          <Link
            to="/"
            activeProps={{
              className: 'font-bold',
            }}
            activeOptions={{ exact: true }}
          >
            Home
          </Link>
        </div>
        <hr style={{ borderColor: colorMap.hr }} />
        <ThemeToggle />
        {children}
        <TanStackRouterDevtools position="bottom-right" />
        <Scripts />
      </body>
    </html>
  )
}

[schiller-manuel]

Thanks for taking a stab at this, but i don't understand how this would solve the problem (maybe i am not seeing it?). this still invalidates the virtual module (but only in the plugin we don't use in start), but does not do any invalidation the env-aware plugin we use in start.

server functions are discovered incrementally in dev, so we need to at least invalidate the virtual module.
i had this (only invalidating the virtual module) in a previous version but this also had the same issue as described in #4486

[Crsk]

@schiller-manuel

Got it, i have no idea what i’m doing tbh, because reverting #4551 solves the issue for me but that does not explain why people is getting this same error before #4551 was introduced.

i’ll just share what i have learned about the issue and pray for a solution to come:

• from commit f0146e1 it seems like the simplifying (ironic) logic moving from static in-disk to incremental in-memory state is causing the race condition.
• i still believe the buildOnDirectiveFnsByIdCallback micromanager should be removed, it’s trying to do Vite’s job by manually invalidating the modules, that’s prone to race condition since Vite's reprocessing is async (after invalidation it takes some time to compile and load), meaning the manifest could reload before the function was ready, so should be a better deal to tell Vite hey bro this manifest is stale viteDevServer.moduleGraph.invalidateModule(...), figure out a new one from the plugin’s loader callback and the directiveFnsById state, but before do your internal thing by analyzing the dependency chain solving them properly so that we have no manifest pointing to stale or non-existing references.
• ok that one was long, in essence probably not the best idea ever to tell Vite to do both things asynchronously: reprocess the file/module and update its index/manifest.
• invalidation ≠ readiness. Vite's moduleGraph.invalidateModule() is synchronous, but the actual module reprocessing is asynchronous.
• if i git check out before fix: invalidate modules after server function manifest update #4551 then i cannot reproduce the error. Probably we’re seeing two sources of the same issue, one of them solved by reverting 4551 and no idea about the other one reported in Server function module export not resolved for serverFn (probably some race condition) #4486.