Patching Vite HMR To Work With Tailwind JIT

During my work at Learnlife, we use React + Vite + Tailwind JIT in our frontend, because we love them! At some point in time, we lost the joy of the fast HMR which was our favorite feature of Vite. We thought it’s a problem with Vite and we tried to revert back to Create React App, but that was never the case, the problem was trickier than we thought. In this article, I will walk you through how we fixed the issue.

UPDATE 19/12/2021: This solution is no longer needed, it seems that we’ve been using the @vitejs/plugin-react-refresh which is replaced by @vitejs/plugin-react and actually solves the problem.

How Did It Look Like?

Each time we change a React component the change is reflected directly to the browser, which is what we expect from HMR, however, immediately after that, Vite performs a page reload, that takes several seconds, given our setup inside docker-compose with mounted volumes hence the I/O is very on macOS. I think you might imagine our frustration waiting for several seconds to see a little change on screen, sometimes I even forgot what I am experimenting with. So, I started looking deeper into the problem.

The First Attempt: Looking Into The Source Code

On my first attempt, I started googling if anyone has the same problem, I found a Github issue states that the HMR was broken after JIT then Vite started doing a full reload after detecting a change from Tailwind. It also says that someone contributed to fixing the problem but I still see it in our setup after the latest versions, but I never gave up!

I started looking for answers in the source code of both Tailwind and Vite.

I found out—in v2.1—that Tailwind touches some empty tmp files and added them to be tracked by Vite, then Vite will do something with that file. What I understood is that Vite has no idea what to do with that file after the dead-end since it’s just an empty file with a random name, and given that Tailwind is the only known (at the time of writing the article) that uses JIT, so Vite doesn’t know what should be hot-ly reloaded.

The code in Tailwind JIT (v2.1 not v3.0) looks like this:

// Source: https://github.com/tailwindlabs/tailwindcss/blob/v2.1/jit/lib/setupContext.js#L36

// Earmarks a directory for our touch files.
// If the directory already exists we delete any existing touch files,
// invalidating any caches associated with them.
const touchDir =
  env.TAILWIND_TOUCH_DIR || path.join(os.homedir() || os.tmpdir(), '.tailwindcss', 'touch')

if (!sharedState.env.TAILWIND_DISABLE_TOUCH) {
  if (fs.existsSync(touchDir)) {
    for (let file of fs.readdirSync(touchDir)) {
      try {
        fs.unlinkSync(path.join(touchDir, file))
      } catch (_err) {}
    }
  } else {
    fs.mkdirSync(touchDir, { recursive: true })
  }
}

In Vite, the code looks like this:

function propagateUpdate(
  node: ModuleNode,
  boundaries: Set<{
    boundary: ModuleNode
    acceptedVia: ModuleNode
  }>,
  currentChain: ModuleNode[] = [node]
): boolean /* hasDeadEnd */ {
  if (node.isSelfAccepting) {
    boundaries.add({
      boundary: node,
      acceptedVia: node
    })

    // additionally check for CSS importers, since a PostCSS plugin like
    // Tailwind JIT may register any file as a dependency to a CSS file.
    for (const importer of node.importers) {
      if (isCSSRequest(importer.url) && !currentChain.includes(importer)) {
        propagateUpdate(importer, boundaries, currentChain.concat(importer))
      }
    }

    return false
  }
  
  // ... more
}

I tried to fork Vite and Tailwind JIT for ourselves—that comes at cost of maintenance for sure—but was never easy to have it fully ready for our use in the timeframe, so I slept on it.

The Second Attempt: A Working Solution

A few months later, while working on the front end, I found myself stuck at removing a border from input that I don’t know how it came in the first place, so I was experimenting with Tailwind classes, you might have already imagined how irritating it is to wait for seconds till you can see the change in the browser a change that led to nothing. I sat there till the end of the day. The very next day, I was angrily looking to solve the problem, because it wasn’t productive to work like that.

If forks from Vite and Tailwind weren’t feasibly possible, then I had to look for another solution. A monkey patch is very likely, and I hope that this solution helps us to solve the problem in the projects themselves.

I started looking into Vite plugins in the hope to find a way to prevent the page reload at my will, notably the vite-plugin-full-reload, it has some examples that give me a few ideas, since I read the source of Vite before, I know what messages to be sent to the client when a change happens. I though I about injecting my own functions into the objects of Vite via a Plugin.

At first, I thought about preventing the tmp file from being tracked by Vite. The first trial looked like this:

export function tailwindJitPatchPlugin() {
  return {
    name: 'tailwind-jit-patch-plugin',
    configureServer({ watcher }: ViteDevServer) {
   
      const originalAdd = watcher.add.bind(watcher)
      watcher.add = file => {
        if (
          file.includes(".tailwindcss") // it was working in v2.1 since all touches happen in a folder
        )
          return
        return originalAdd(file)
      }
    },
  }
}

This prevented the full reload, but it never triggers Tailwind built to get the new classes added. However, I noticed that if I save the entry CSS file (i.e. src/style/index.scss in our case) it will do a Hot Reload.

I thought about touching the file when that happens.

import { utimes } from 'fs/promises'
import path from 'path'

const touchFile = async path => {
  const time = new Date()
  return utimes(path, time, time)
}

export function tailwindJitPatchPlugin() {
  return {
    name: 'tailwind-jit-patch-plugin',
    configureServer({ watcher }: ViteDevServer) {
   
      const originalAdd = watcher.add.bind(watcher)
      watcher.add = file => {
        if (
          file.includes(".tailwindcss") // in v2.1
        ) {
          touchFile(path.resolve(process.cwd(), entryFile))
          return
        }
        return originalAdd(file)
      }
    },
  }
}

It worked for a moment, however, the adding to the watcher is only needed if the file is added for the first type, so the touch will not be triggered unless it’s a different tmp file. Probably it’s expected that if I save the CSS file, it will be a recursive touch and will never stop.

Then I thought about the send() method in the ws server. Because it’s mainly worked for letting the client know there’s file change, it might not cause the recursive touching problem.

So, the plugin now looks like this:

import { utimes } from 'fs/promises'
import path from 'path'
import { ViteDevServer } from 'vite'

const touchFile = async path => {
  const time = new Date()
  return utimes(path, time, time)
}

export function tailwindJitPatchPlugin({
  entryFile = './src/index.css',
} = {}) {
  return {
    name: 'tailwind-jit-patch-plugin',
    configureServer({ ws, watcher, config: { logger } }: ViteDevServer) {
      const originalSend = ws.send.bind(ws)
      ws.send = update => {
        originalSend(update)
        logger.info(`touching ${entryFile}`)
        touchFile(path.resolve(process.cwd(), entryFile))
      }
    },
  }
}

It worked as expected, but Tailwind still touching tmp files then it causes the reload, the good thing is that the touch can be disabled (in v2.1) when passing TAILWIND_DISABLE_TOUCH=true argument in the command line.

For Tailwind 2.1, that was the first working solution but needed some performance improvements.

import { utimes } from 'fs/promises'
import { debounce } from 'lodash'
import path from 'path'
import { ViteDevServer } from 'vite'

const touchFile = async path => {
  const time = new Date()
  return utimes(path, time, time)
}

const debouncedTouch = debounce(touchFile, 100)

export function tailwindJitPatchPlugin({
  entryFile = './src/index.css',
  enabled = true,
} = {}) {
  return {
    name: 'tailwind-jit-patch-plugin',
    apply: 'serve', // Only during serve
    configureServer({ ws, watcher, config: { logger } }: ViteDevServer) {
      if (!enabled) return

      // Logging
      logger.info(
        'EXPERIMENTAL! Tailwind Jit Patch Plugin - it might be buggy, but this is the only solution at the moment'
      )

      const originalSend = ws.send.bind(ws)
      ws.send = update => {
        originalSend(update)

        const anyJsUpdate = update.updates?.some(u => u.type === 'js-update')
        // Only if js update, there's more?
        if (!anyJsUpdate) return 

        const theEntryFile = update.updates?.find(
          u => `.${u.path}` === entryFile
        )
        
        // Don't touch the css if it's changed just now.
        if (theEntryFile) return
        logger.info(`touching ${entryFile}`)
        
        // Noticed that JS HMR was more than instant especially if no CSS change needed.
        // It delays the touching of the file so the first update is not blocked.
        // Now the JS or old CSS classes updates run in an instant,
        // but the new classes need about 100ms to take effect,
        // Which is the same update time before the debounce.
        debouncedTouch(path.resolve(process.cwd(), entryFile)) 
      }
    },
  }
}

The solution above is the final solution for Tailwind v2.1, however, in version 3.0, the TAILWIND_DISABLE_TOUCH variable seems to have no effect. So, I needed to prevent the watching of the file, that was the final solution.

The Final Solution

The final solution looks like this, it handles the case of Tailwind 3.0 as well.

import { utimes } from 'fs/promises'
import { debounce } from 'lodash'
import path from 'path'
import { ViteDevServer } from 'vite'

const touchFile = async path => {
  const time = new Date()
  return utimes(path, time, time)
}

const debouncedTouch = debounce(touchFile, 100)

export function tailwindJitPatchPlugin({
  entryFile = './src/index.css',
  enabled = true,
} = {}) {
  return {
    name: 'tailwind-jit-patch-plugin',
    apply: 'serve',
    configureServer({ ws, watcher, config: { logger } }: ViteDevServer) {
      if (!enabled) return

      logger.info(
        'EXPERIMENTAL! Tailwind Jit Patch Plugin - it might be buggy, but this is the only solution at the moment'
      )

      const originalSend = ws.send.bind(ws)
      ws.send = update => {
        originalSend(update)

        const anyJsUpdate = update.updates?.some(u => u.type === 'js-update')
        if (!anyJsUpdate) return

        const theEntryFile = update.updates?.find(
          u => `.${u.path}` === entryFile
        )
        if (theEntryFile) return
        logger.info(`touching ${entryFile}`)
        debouncedTouch(path.resolve(process.cwd(), entryFile))
      }

      // For Tailwind 3.0
      const originalAdd = watcher.add.bind(watcher)
      watcher.add = file => {
        if (
          ['true', true, 1, '1'].includes(process.env.TAILWIND_DISABLE_TOUCH) &&
          file.match(/tmp-\d+-\w{12}/)
        )
          return
        return originalAdd(file)
      }
    },
  }
}

The touched file name in Tailwind 3.0 looks random with some pattern. Unlike in 2.1, the files are not inside a single directory that the name implies is related to Tailwind (i.e. .tailwindcss), and the environment variable TAILWIND_TOUCH_DIR has no effect.

A Caveat

In Tailwind 3.0, since it’s not possible to disable the touch from environment variables, the solution to add files to a watcher can be quirky and might cause undesireable behaviour if there’s another framework that files with the same file name pattern. For now, I used the same old environment variable to control the touch.

Conclusion

I hope that this plugin solves your problem and a step towards solving the problem in Tailwind and Vite.

I might check if it’s possible to return those environment variables back or find a better solution, till then, see you in a future article.