fix(vrowzer): restore the project after the Service Worker restarts

packages/fs/README.md

@@ -142,7 +142,7 @@ publisher.mkdir('/new-dir')
 publisher.initFiles({ '/main.js': 'code' }, { '/logo.png': arrayBuffer }) // bulk init (each target gets its own copy of the binary files)
 ```
 
-`writeFile()` and `unlink()` accept an operation ID, which is set on the message. A Worker that applies such a message can acknowledge it with a `V_FS_ACK` message (`{ type: 'V_FS_ACK', id, error? }`), so that the main thread knows when the change is applied. `@vrowzer/fs` does not send acknowledgements by itself: the Worker that handles the messages decides when a change counts as applied.
+`writeFile()`, `unlink()` and `initFiles()` accept an operation ID, which is set on the message. A Worker that applies such a message can acknowledge it with a `V_FS_ACK` message (`{ type: 'V_FS_ACK', id, error? }`), so that the main thread knows when the change is applied. `@vrowzer/fs` does not send acknowledgements by itself: the Worker that handles the messages decides when a change counts as applied.
 
 ```ts
 publisher.writeFile('/main.js', 'export const x = 2', { id: crypto.randomUUID() })

packages/fs/docs/watcher/interfaces/FSInitMessage.md

@@ -19,4 +19,5 @@ export interface FSInitMessage
 | --- | --- | --- |
 | `binaryFiles` _(optional)_ | `Record<string, ArrayBuffer>` | Binary files: path -> ArrayBuffer content (a copy transferred via postMessage's transfer list) |
 | `files` _(optional)_ | `Record<string, string>` | Text files: path -> UTF-8 string content |
+| `id` _(optional)_ | `string` | Operation ID. A Worker that applies the message can acknowledge it with [FSAckMessage](/packages/fs/docs/watcher/interfaces/FSAckMessage.md). |
 | `type` | `"V_FS_INIT"` |  |

packages/fs/docs/watcher/interfaces/FileSystemPublisher.md

@@ -38,7 +38,7 @@ Add a postMessage target.
 ### initFiles()
 
 ```ts
-initFiles(files?: Record<string, string>, binaryFiles?: Record<string, ArrayBuffer>): void;
+initFiles(files?: Record<string, string>, binaryFiles?: Record<string, ArrayBuffer>, options?: FileSystemPublishOptions): void;
 ```
 
 Initialize files in bulk.
@@ -49,6 +49,7 @@ Initialize files in bulk.
 | --- | --- | --- |
 | `files` | `Record<string, string>` | Text files: path -> UTF-8 string content _(optional)_ |
 | `binaryFiles` | `Record<string, ArrayBuffer>` | Binary files: path -> ArrayBuffer content. Each target receives its own copy via postMessage's transfer list, so the caller's ArrayBuffers stay usable. _(optional)_ |
+| `options` | [`FileSystemPublishOptions`](/packages/fs/docs/watcher/interfaces/FileSystemPublishOptions.md) | Options for this operation, such as its ID _(optional)_ |
 
 #### Returns
 

packages/fs/src/watcher/protocol.ts

@@ -100,6 +100,10 @@ export interface FSInitMessage {
    * Binary files: path -> ArrayBuffer content (a copy transferred via postMessage's transfer list)
    */
   binaryFiles?: Record<string, ArrayBuffer>
+  /**
+   * Operation ID. A Worker that applies the message can acknowledge it with {@link FSAckMessage}.
+   */
+  id?: string
 }
 
 export type FileSystemSyncMessage =

packages/fs/src/watcher/publisher.test.ts

@@ -320,5 +320,29 @@ describe('FileSystemPublisher', () => {
       expect([...new Uint8Array(message.binaryFiles['/a.bin'])]).toEqual([1, 2, 3])
       expect([...new Uint8Array(message.binaryFiles['/b.bin'])]).toEqual([1, 2, 3])
     })
+
+    test('sets the operation id on the message for every target', () => {
+      const target1 = createMockTarget()
+      const target2 = createMockTarget()
+      const publisher = createFileSystemPublisher([target1, target2])
+
+      publisher.initFiles(
+        { '/main.js': 'code' },
+        { '/app.wasm': new ArrayBuffer(4) },
+        { id: 'op-4' }
+      )
+
+      expect(target1.calls[0]!.message.id).toBe('op-4')
+      expect(target2.calls[0]!.message.id).toBe('op-4')
+    })
+
+    test('omits the operation id without the option', () => {
+      const target = createMockTarget()
+      const publisher = createFileSystemPublisher([target])
+
+      publisher.initFiles({ '/main.js': 'code' })
+
+      expect(target.calls[0]!.message).not.toHaveProperty('id')
+    })
   })
 })

packages/fs/src/watcher/publisher.ts

@@ -76,8 +76,13 @@ export interface FileSystemPublisher {
    * @param files - Text files: path -> UTF-8 string content
    * @param binaryFiles - Binary files: path -> ArrayBuffer content. Each target receives its own copy
    *   via postMessage's transfer list, so the caller's ArrayBuffers stay usable.
+   * @param options - Options for this operation, such as its ID
    */
-  initFiles(files?: Record<string, string>, binaryFiles?: Record<string, ArrayBuffer>): void
+  initFiles(
+    files?: Record<string, string>,
+    binaryFiles?: Record<string, ArrayBuffer>,
+    options?: FileSystemPublishOptions
+  ): void
   /**
    * Add a postMessage target.
    *
@@ -143,11 +148,11 @@ export function createFileSystemPublisher(
       broadcast({ type: 'V_FS_MKDIR', path })
     },
 
-    initFiles(files, binaryFiles) {
+    initFiles(files, binaryFiles, options) {
       // Each target gets its own copies of the binary files in the transfer list.
       // The caller's buffers are never transferred, so they stay usable.
       for (const target of _targets) {
-        const msg: FSInitMessage = { type: 'V_FS_INIT' }
+        const msg: FSInitMessage = { type: 'V_FS_INIT', ...idOf(options) }
         if (files) {
           msg.files = files
         }

packages/vite-dev-server/src/node/server/middlewares/beforeRequest.test.ts

@@ -0,0 +1,77 @@
+import { describe, expect, test } from 'vite-plus/test'
+import { Hono } from 'hono'
+import { beforeRequestMiddleware } from './beforeRequest'
+import { crossOriginMiddleware } from './crossOrigin'
+
+function createApp(beforeRequest: (request: Request) => Promise<Response | undefined>) {
+  const app = new Hono()
+  app.use(crossOriginMiddleware())
+  app.use(beforeRequestMiddleware(beforeRequest))
+  app.get('*', c => c.text('served'))
+  return app
+}
+
+describe('beforeRequestMiddleware', () => {
+  test('waits for the hook before handling the request', async () => {
+    let release!: () => void
+    const app = createApp(() => new Promise(resolve => {
+      release = () => resolve(undefined)
+    }))
+
+    let settled = false
+    const response = app.request('/main.js').then((res) => {
+      settled = true
+      return res
+    })
+    for (let index = 0; index < 10; index++) {
+      await Promise.resolve()
+    }
+    expect(settled).toBe(false)
+
+    release()
+    const res = await response
+    expect(res.status).toBe(200)
+    expect(await res.text()).toBe('served')
+  })
+
+  test('handles the request as usual when the hook resolves to undefined', async () => {
+    const app = createApp(async () => undefined)
+
+    const res = await app.request('/main.js')
+
+    expect(res.status).toBe(200)
+    expect(await res.text()).toBe('served')
+  })
+
+  test('answers with the response of the hook, with the headers of earlier middlewares', async () => {
+    const app = createApp(async () => new Response('not ready', {
+      status: 503,
+      headers: { 'Content-Type': 'text/plain' },
+    }))
+
+    const res = await app.request('/main.js')
+
+    expect(res.status).toBe(503)
+    expect(await res.text()).toBe('not ready')
+    expect(res.headers.get('Content-Type')).toBe('text/plain')
+    expect(res.headers.get('Cross-Origin-Resource-Policy')).toBe('same-origin')
+    expect(res.headers.get('Cross-Origin-Embedder-Policy')).toBe('require-corp')
+    expect(res.headers.get('Cross-Origin-Opener-Policy')).toBe('same-origin')
+  })
+
+  test('passes the request to the hook', async () => {
+    const requests: Request[] = []
+    const app = createApp(async (request) => {
+      requests.push(request)
+      return undefined
+    })
+
+    await app.request('http://localhost/__preview__/main.js?import', {
+      headers: { Accept: 'text/javascript' },
+    })
+
+    expect(requests).toHaveLength(1)
+    expect(requests[0]!.url).toBe('http://localhost/__preview__/main.js?import')
+    expect(requests[0]!.headers.get('Accept')).toBe('text/javascript')
+  })
+})

packages/vite-dev-server/src/node/server/middlewares/beforeRequest.ts

@@ -0,0 +1,32 @@
+/**
+ * Before request Hono middleware
+ *
+ * Lets the owner of the Service Worker dev server hold a request until it can be served,
+ * or answer it instead.
+ */
+
+/**
+ * @author kazuya kawaguchi (a.k.a. kazupon)
+ * @license MIT
+ */
+
+import type { MiddlewareHandler } from 'hono'
+import type { ViteEnv } from '../index'
+
+/**
+ * Middleware that waits for `beforeRequest` before the request is handled.
+ * When the hook resolves to a `Response`, that response is returned instead.
+ */
+export function beforeRequestMiddleware(
+  beforeRequest: (request: Request) => Promise<Response | undefined>,
+): MiddlewareHandler<ViteEnv> {
+  return async function viteBeforeRequestMiddleware(c, next) {
+    const response = await beforeRequest(c.req.raw)
+    if (response) {
+      // A response returned as it is would lose the headers that earlier middlewares set,
+      // such as the cross-origin isolation headers, so rebuild it through the context
+      return c.newResponse(response.body, response)
+    }
+    await next()
+  }
+}

packages/vite-dev-server/src/node/service-worker.ts

@@ -28,6 +28,7 @@ import { isResolvedConfig, resolveConfig } from './config'
 import { syncPublicFiles } from './public-files-sync'
 import { initPublicFiles } from './publicDir'
 import { baseMiddleware } from './server/middlewares/base'
+import { beforeRequestMiddleware } from './server/middlewares/beforeRequest'
 import { crossOriginMiddleware } from './server/middlewares/crossOrigin'
 import { errorMiddleware } from './server/middlewares/error'
 import { htmlFallbackMiddleware } from './server/middlewares/htmlFallback'
@@ -306,6 +307,18 @@ export interface CreateServerOptions {
    * FSWawtcher factory function to create a custom FSWatcher instance.
    */
   watcherFactory?: (targets: string[], options: WatchOptions) => FSWatcher
+  /**
+   * Called before a request within the base path is handled. The request waits until the returned
+   * promise resolves. When it resolves to a `Response`, that response is returned instead, with the
+   * headers that the earlier middlewares set.
+   */
+  beforeRequest?: (request: Request) => Promise<Response | undefined>
+  /**
+   * Called each time the channel handshake with a Web Worker (`V_WW_CONNECT_PORT`) completes,
+   * before the client is told that the connection is established. A later connection replaces
+   * the previous channel.
+   */
+  onWorkerChannelReady?: () => void
   /**
    * @internal
    */
@@ -547,6 +560,13 @@ export function createServer(
     // Cross-origin isolation headers (CORP/COEP/COOP) for credentialless iframe + SW
     middlewares.use(crossOriginMiddleware())
 
+    // Let the owner hold requests until they can be served, e.g. until a restarted Service Worker
+    // has the project files and the Web Worker channel again. This is registered before the
+    // configureServer hooks, so that the middlewares of plugins wait as well.
+    if (options.beforeRequest) {
+      middlewares.use(beforeRequestMiddleware(options.beforeRequest))
+    }
+
     // TODO(kazupon): disable middlewares, after implementing them
     // middlewares.use(rejectInvalidRequestMiddleware())
     // middlewares.use(rejectNoCorsRequestMiddleware())
@@ -722,6 +742,10 @@ export function createServer(
               }
             )
 
+            // Requests can be transformed again. Tell the owner before the client hears that the
+            // connection is established, so that requests held by `beforeRequest` can go on by then.
+            options.onWorkerChannelReady?.()
+
             // Notify the originating client that the connection is established
             if (clientId) {
               const client = await serviceWorkerScope.clients.get(clientId)