feat: add WebSocket support to dev proxy with cookie rewriting (#39370)

packages/dev-proxy/README.md

@@ -212,6 +212,7 @@ const cookieRewrite: CookieRewriteOptions = {
 ## Behavior
 
 - The proxy preserves the matched path prefix when forwarding requests.
+- WebSocket Upgrade requests use the same routes, targets, origin policy, and cookie rewriting.
 - Request bodies are forwarded as streams.
 - Hop-by-hop headers are removed before forwarding.
 - Local credentialed CORS and preflight requests are handled by the proxy.

packages/dev-proxy/src/cli.ts

@@ -1,4 +1,5 @@
 import type { ServerType } from '@hono/node-server'
+import type { Duplex } from 'node:stream'
 import type { DevProxyCliOptions, DevProxyConfig } from './types'
 import process from 'node:process'
 import { serve } from '@hono/node-server'
@@ -10,7 +11,7 @@ import {
   resolveDevProxyServerOptions,
   watchDevProxyConfig,
 } from './config'
-import { createDevProxyApp } from './server'
+import { createDevProxyApp, createWebSocketUpgradeHandler } from './server'
 
 function printUsage() {
   console.log(`Usage:
@@ -43,19 +44,32 @@ const closeServer = (server: ServerType) =>
 
 const startDevProxyServer = (config: DevProxyConfig, cliOptions: DevProxyCliOptions) => {
   let app = createDevProxyApp(config)
+  let handleWebSocketUpgrade = createWebSocketUpgradeHandler(config)
+  const upgradedSockets = new Set<Duplex>()
   const { host, port } = resolveDevProxyServerOptions(config.server, cliOptions)
   const server = serve({
     fetch: (request, env) => app.fetch(request, env),
     hostname: host,
     port,
   })
+  server.on('upgrade', (request, socket, head) => {
+    upgradedSockets.add(socket)
+    socket.once('close', () => upgradedSockets.delete(socket))
+    handleWebSocketUpgrade(request, socket, head)
+  })
 
   return {
     host,
     port,
     server,
+    close() {
+      upgradedSockets.forEach((socket) => socket.destroy())
+      upgradedSockets.clear()
+      return closeServer(server)
+    },
     updateConfig(nextConfig: DevProxyConfig) {
       app = createDevProxyApp(nextConfig)
+      handleWebSocketUpgrade = createWebSocketUpgradeHandler(nextConfig)
     },
   }
 }
@@ -76,7 +90,7 @@ const createDevProxyRuntime = (initialConfig: DevProxyConfig, cliOptions: DevPro
       return
     }
 
-    await closeServer(runtime.server)
+    await runtime.close()
     runtime = startDevProxyServer(nextConfig, cliOptions)
     console.log(`[dev-proxy] restarted on http://${runtime.host}:${runtime.port} after ${reason}`)
   }
@@ -98,7 +112,7 @@ const createDevProxyRuntime = (initialConfig: DevProxyConfig, cliOptions: DevPro
     enqueueReload,
     close: async () => {
       await reloadTask
-      await closeServer(runtime.server)
+      await runtime.close()
     },
   }
 }

packages/dev-proxy/src/server.spec.ts

@@ -1,9 +1,39 @@
 /**
  * @vitest-environment node
  */
+import { Buffer } from 'node:buffer'
+import http from 'node:http'
+import net from 'node:net'
+import { PassThrough } from 'node:stream'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 import { resolveCookieRewriteLocalScopeKey, toScopedLocalCookieName } from './cookies'
-import { buildUpstreamUrl, createDevProxyApp, isAllowedDevOrigin } from './server'
+import {
+  buildUpstreamUrl,
+  createDevProxyApp,
+  createWebSocketUpgradeHandler,
+  isAllowedDevOrigin,
+} from './server'
+
+const listen = (server: http.Server) =>
+  new Promise<number>((resolve, reject) => {
+    server.once('error', reject)
+    server.listen(0, '127.0.0.1', () => {
+      const address = server.address()
+      if (!address || typeof address === 'string') {
+        reject(new Error('Failed to start test server.'))
+        return
+      }
+      resolve(address.port)
+    })
+  })
+
+const close = (server: http.Server) =>
+  new Promise<void>((resolve, reject) => {
+    server.close((error) => {
+      if (error) reject(error)
+      else resolve()
+    })
+  })
 
 describe('dev proxy server', () => {
   beforeEach(() => {
@@ -259,6 +289,252 @@ describe('dev proxy server', () => {
     ])
   })
 
+  // Scenario: Socket.IO collaboration must reuse the auth cookies stored by the local dev proxy.
+  it('should proxy WebSocket upgrades with the configured cookie rewriting', async () => {
+    // Arrange
+    const sockets = new Set<net.Socket>()
+    let upstreamRequest:
+      | {
+          connection: string | undefined
+          cookie: string | undefined
+          keepAlive: string | string[] | undefined
+          origin: string | undefined
+          proxyAuthorization: string | undefined
+          te: string | string[] | undefined
+          upgrade: string | undefined
+          url: string | undefined
+          xHop: string | string[] | undefined
+        }
+      | undefined
+    const upstreamServer = http.createServer()
+    upstreamServer.on('connection', (socket) => sockets.add(socket))
+    upstreamServer.on('upgrade', (request, socket) => {
+      upstreamRequest = {
+        connection: request.headers.connection,
+        cookie: request.headers.cookie,
+        keepAlive: request.headers['keep-alive'],
+        origin: request.headers.origin,
+        proxyAuthorization: request.headers['proxy-authorization'],
+        te: request.headers.te,
+        upgrade: request.headers.upgrade,
+        url: request.url,
+        xHop: request.headers['x-hop'],
+      }
+      socket.write(
+        [
+          'HTTP/1.1 101 Switching Protocols',
+          'Connection: Upgrade, X-Upstream-Hop',
+          'Upgrade: websocket',
+          'X-Upstream-Hop: upstream-only',
+          'Keep-Alive: timeout=5',
+          'Proxy-Authenticate: Basic realm="proxy"',
+          'Set-Cookie: __Host-access_token=next; Path=/socket.io; Domain=cloud.example.com; Secure; SameSite=None',
+          'Set-Cookie: __Secure-refresh_token=renewed; Path=/socket.io; Domain=cloud.example.com; Secure; HttpOnly',
+          '',
+          'proxied websocket',
+        ].join('\r\n'),
+      )
+    })
+    const upstreamPort = await listen(upstreamServer)
+    const upstreamOrigin = `http://127.0.0.1:${upstreamPort}`
+    const cookieRewrite = {
+      hostPrefixCookies: ['access_token', 'refresh_token'],
+      localCookieScope: 'target-origin' as const,
+    }
+    const localScopeKey = resolveCookieRewriteLocalScopeKey(cookieRewrite, new URL(upstreamOrigin))!
+    const accessTokenCookieName = toScopedLocalCookieName('access_token', localScopeKey)
+    const refreshTokenCookieName = toScopedLocalCookieName('refresh_token', localScopeKey)
+
+    const proxyServer = http.createServer()
+    proxyServer.on(
+      'upgrade',
+      createWebSocketUpgradeHandler({
+        routes: [{ paths: '/socket.io', target: upstreamOrigin, cookieRewrite }],
+      }),
+    )
+    const proxyPort = await listen(proxyServer)
+    const client = net.connect(proxyPort, '127.0.0.1')
+    sockets.add(client)
+
+    try {
+      // Act
+      const response = await new Promise<string>((resolve, reject) => {
+        let value = ''
+        const timeout = setTimeout(
+          () => reject(new Error('Timed out waiting for WebSocket upgrade response.')),
+          3000,
+        )
+
+        client.setEncoding('utf8')
+        client.on('connect', () => {
+          client.write(
+            [
+              'GET /socket.io/?EIO=4&transport=websocket HTTP/1.1',
+              `Host: 127.0.0.1:${proxyPort}`,
+              'Connection: Upgrade, X-Hop',
+              'Upgrade: websocket',
+              'X-Hop: client-only',
+              'Keep-Alive: timeout=5',
+              'Proxy-Authorization: Basic secret',
+              'TE: trailers',
+              'Origin: http://localhost:3000',
+              `Cookie: ${accessTokenCookieName}=secret`,
+              '',
+              '',
+            ].join('\r\n'),
+          )
+        })
+        client.on('data', (chunk) => {
+          value += chunk
+          if (!value.includes('proxied websocket')) return
+
+          clearTimeout(timeout)
+          resolve(value)
+        })
+        client.on('error', (error) => {
+          clearTimeout(timeout)
+          reject(error)
+        })
+      })
+
+      // Assert
+      expect(response).toContain('101 Switching Protocols')
+      expect(response).toContain('proxied websocket')
+      expect(response).toContain('Connection: Upgrade\r\n')
+      expect(response).toContain('Upgrade: websocket\r\n')
+      expect(response).not.toMatch(/X-Upstream-Hop:/i)
+      expect(response).not.toMatch(/Keep-Alive:/i)
+      expect(response).not.toMatch(/Proxy-Authenticate:/i)
+      expect(response).toContain(`Set-Cookie: ${accessTokenCookieName}=next; Path=/; SameSite=Lax`)
+      expect(response).toContain(`Set-Cookie: ${refreshTokenCookieName}=renewed; Path=/; HttpOnly`)
+      expect(response).not.toContain('__Host-access_token')
+      expect(response).not.toContain('__Secure-refresh_token')
+      expect(upstreamRequest).toEqual({
+        connection: 'Upgrade',
+        cookie: 'access_token=secret',
+        keepAlive: undefined,
+        origin: upstreamOrigin,
+        proxyAuthorization: undefined,
+        te: undefined,
+        upgrade: 'websocket',
+        url: '/socket.io/?EIO=4&transport=websocket',
+        xHop: undefined,
+      })
+    } finally {
+      sockets.forEach((socket) => socket.destroy())
+      await Promise.all([close(proxyServer), close(upstreamServer)])
+    }
+  })
+
+  // Scenario: invalid route targets should fail the individual Upgrade request, not the proxy process.
+  it('should return a bad gateway response when a WebSocket target is invalid', async () => {
+    // Arrange
+    const logger = { error: vi.fn() }
+    const clientSocket = new PassThrough()
+    const request = {
+      headers: {},
+      method: 'GET',
+      url: '/socket.io/?EIO=4&transport=websocket',
+    } as http.IncomingMessage
+    let response = ''
+    clientSocket.setEncoding('utf8')
+    clientSocket.on('data', (chunk) => {
+      response += chunk
+    })
+    const handleUpgrade = createWebSocketUpgradeHandler({
+      routes: [{ paths: '/socket.io', target: 'not a URL' }],
+      logger,
+    })
+
+    // Act
+    expect(() => handleUpgrade(request, clientSocket, Buffer.alloc(0))).not.toThrow()
+    await new Promise<void>((resolve) => clientSocket.once('finish', resolve))
+
+    // Assert
+    expect(response).toContain('502 Bad Gateway')
+    expect(logger.error).toHaveBeenCalledOnce()
+  })
+
+  // Scenario: parsed chunked bodies need close-delimited framing when an Upgrade is rejected upstream.
+  it('should safely forward chunked non-upgrade responses', async () => {
+    // Arrange
+    const sockets = new Set<net.Socket>()
+    const upstreamServer = http.createServer()
+    upstreamServer.on('connection', (socket) => sockets.add(socket))
+    upstreamServer.on('upgrade', (_request, socket) => {
+      socket.end(
+        [
+          'HTTP/1.1 401 Unauthorized',
+          'Connection: keep-alive',
+          'Transfer-Encoding: chunked',
+          'Content-Type: text/plain',
+          '',
+          '6',
+          'denied',
+          '0',
+          '',
+          '',
+        ].join('\r\n'),
+      )
+    })
+    const upstreamPort = await listen(upstreamServer)
+    const proxyServer = http.createServer()
+    proxyServer.on(
+      'upgrade',
+      createWebSocketUpgradeHandler({
+        routes: [{ paths: '/socket.io', target: `http://127.0.0.1:${upstreamPort}` }],
+      }),
+    )
+    const proxyPort = await listen(proxyServer)
+    const client = net.connect(proxyPort, '127.0.0.1')
+    sockets.add(client)
+
+    try {
+      // Act
+      const response = await new Promise<string>((resolve, reject) => {
+        let value = ''
+        const timeout = setTimeout(
+          () => reject(new Error('Timed out waiting for rejected Upgrade response.')),
+          3000,
+        )
+
+        client.setEncoding('utf8')
+        client.on('connect', () => {
+          client.write(
+            [
+              'GET /socket.io/?EIO=4&transport=websocket HTTP/1.1',
+              `Host: 127.0.0.1:${proxyPort}`,
+              'Connection: Upgrade',
+              'Upgrade: websocket',
+              '',
+              '',
+            ].join('\r\n'),
+          )
+        })
+        client.on('data', (chunk) => {
+          value += chunk
+        })
+        client.on('end', () => {
+          clearTimeout(timeout)
+          resolve(value)
+        })
+        client.on('error', (error) => {
+          clearTimeout(timeout)
+          reject(error)
+        })
+      })
+
+      // Assert
+      expect(response).toContain('401 Unauthorized')
+      expect(response).toContain('Connection: close')
+      expect(response).not.toMatch(/Transfer-Encoding:/i)
+      expect(response.endsWith('\r\n\r\ndenied')).toBe(true)
+    } finally {
+      sockets.forEach((socket) => socket.destroy())
+      await Promise.all([close(proxyServer), close(upstreamServer)])
+    }
+  })
+
   // Scenario: routes are matched in config order so callers can put specific routes first.
   it('should prefer earlier route entries', async () => {
     // Arrange

packages/dev-proxy/src/server.ts

@@ -1,10 +1,15 @@
 import type { Context, Hono } from 'hono'
+import type { Buffer } from 'node:buffer'
+import type { IncomingMessage } from 'node:http'
+import type { Duplex } from 'node:stream'
 import type {
   CookieRewriteOptions,
   CreateDevProxyAppOptions,
   DevProxyCorsAllowedOrigins,
   DevProxyRoute,
 } from './types'
+import { request as httpRequest } from 'node:http'
+import { request as httpsRequest } from 'node:https'
 import { Hono as HonoApp } from 'hono'
 import {
   getCookieHeaderValue,
@@ -18,10 +23,8 @@ const LOCAL_DEV_HOSTS = new Set(['localhost', '127.0.0.1', '[::1]', '::1'])
 const ALLOW_METHODS = 'GET,HEAD,POST,PUT,PATCH,DELETE,OPTIONS'
 const DEFAULT_ALLOW_HEADERS = 'Authorization, Content-Type, X-CSRF-Token'
 const UPSTREAM_ACCEPT_ENCODING = 'identity'
-const RESPONSE_HEADERS_TO_DROP = [
+const HOP_BY_HOP_HEADERS = [
   'connection',
-  'content-encoding',
-  'content-length',
   'keep-alive',
   'proxy-authenticate',
   'proxy-authorization',
@@ -30,6 +33,20 @@ const RESPONSE_HEADERS_TO_DROP = [
   'transfer-encoding',
   'upgrade',
 ] as const
+const DECODED_RESPONSE_HEADERS_TO_DROP = ['content-encoding', 'content-length'] as const
+
+const createHopByHopHeaderNames = (connectionHeader?: string | null) =>
+  new Set([
+    ...HOP_BY_HOP_HEADERS,
+    ...(connectionHeader
+      ?.split(',')
+      .map((header) => header.trim().toLowerCase())
+      .filter(Boolean) || []),
+  ])
+
+const removeHopByHopHeaders = (headers: Headers) => {
+  createHopByHopHeaderNames(headers.get('connection')).forEach((header) => headers.delete(header))
+}
 
 const appendHeaderValue = (headers: Headers, name: string, value: string) => {
   const currentValue = headers.get(name)
@@ -102,29 +119,30 @@ export const buildUpstreamUrl = (target: string, requestPath: string, search = '
 }
 
 const createProxyRequestHeaders = (
-  request: Request,
+  headers: Headers,
   targetUrl: URL,
   cookieRewrite: CookieRewriteOptions | false | undefined,
 ) => {
-  const headers = new Headers(request.headers)
-  headers.delete('host')
-  headers.set('accept-encoding', UPSTREAM_ACCEPT_ENCODING)
+  const upstreamHeaders = new Headers(headers)
+  removeHopByHopHeaders(upstreamHeaders)
+  upstreamHeaders.delete('host')
+  upstreamHeaders.set('accept-encoding', UPSTREAM_ACCEPT_ENCODING)
 
-  if (headers.has('origin')) headers.set('origin', targetUrl.origin)
+  if (upstreamHeaders.has('origin')) upstreamHeaders.set('origin', targetUrl.origin)
 
   if (cookieRewrite) {
-    const originalCookieHeader = headers.get('cookie') || undefined
+    const originalCookieHeader = upstreamHeaders.get('cookie') || undefined
     const localScopeKey = resolveCookieRewriteLocalScopeKey(cookieRewrite, targetUrl)
     const rewrittenCookieHeader = rewriteCookieHeaderForUpstream(
-      headers.get('cookie') || undefined,
+      upstreamHeaders.get('cookie') || undefined,
       {
         ...cookieRewrite,
         localScopeKey,
         useHostPrefix: targetUrl.protocol === 'https:',
       },
     )
-    if (rewrittenCookieHeader) headers.set('cookie', rewrittenCookieHeader)
-    else headers.delete('cookie')
+    if (rewrittenCookieHeader) upstreamHeaders.set('cookie', rewrittenCookieHeader)
+    else upstreamHeaders.delete('cookie')
 
     if (localScopeKey && cookieRewrite.csrfHeader) {
       const scopedCsrfCookieName = toScopedLocalCookieName(
@@ -132,12 +150,12 @@ const createProxyRequestHeaders = (
         localScopeKey,
       )
       const scopedCsrfToken = getCookieHeaderValue(originalCookieHeader, scopedCsrfCookieName)
-      if (scopedCsrfToken) headers.set(cookieRewrite.csrfHeader.headerName, scopedCsrfToken)
-      else headers.delete(cookieRewrite.csrfHeader.headerName)
+      if (scopedCsrfToken) upstreamHeaders.set(cookieRewrite.csrfHeader.headerName, scopedCsrfToken)
+      else upstreamHeaders.delete(cookieRewrite.csrfHeader.headerName)
     }
   }
 
-  return headers
+  return upstreamHeaders
 }
 
 const getSetCookieHeaders = (headers: Headers) => {
@@ -157,7 +175,8 @@ const createUpstreamResponseHeaders = (
   cookieRewrite: CookieRewriteOptions | false | undefined,
 ) => {
   const headers = new Headers(response.headers)
-  RESPONSE_HEADERS_TO_DROP.forEach((header) => headers.delete(header))
+  removeHopByHopHeaders(headers)
+  DECODED_RESPONSE_HEADERS_TO_DROP.forEach((header) => headers.delete(header))
   headers.delete('set-cookie')
 
   const localScopeKey = cookieRewrite
@@ -187,7 +206,11 @@ const proxyRequest = async (
 ) => {
   const requestUrl = new URL(context.req.url)
   const targetUrl = buildUpstreamUrl(route.target, requestUrl.pathname, requestUrl.search)
-  const requestHeaders = createProxyRequestHeaders(context.req.raw, targetUrl, route.cookieRewrite)
+  const requestHeaders = createProxyRequestHeaders(
+    context.req.raw.headers,
+    targetUrl,
+    route.cookieRewrite,
+  )
   const requestInit: RequestInit & { duplex?: 'half' } = {
     method: context.req.method,
     headers: requestHeaders,
@@ -218,6 +241,192 @@ const proxyRequest = async (
 const normalizeRoutePaths = (paths: DevProxyRoute['paths']) =>
   Array.isArray(paths) ? paths : [paths]
 
+const findProxyRoute = (routes: readonly DevProxyRoute[], requestPath: string) =>
+  routes.find((route) =>
+    normalizeRoutePaths(route.paths).some(
+      (routePath) => requestPath === routePath || requestPath.startsWith(`${routePath}/`),
+    ),
+  )
+
+const createHeadersFromIncomingMessage = (request: IncomingMessage) => {
+  const headers = new Headers()
+  Object.entries(request.headers).forEach(([name, value]) => {
+    if (Array.isArray(value)) value.forEach((item) => headers.append(name, item))
+    else if (value !== undefined) headers.set(name, value)
+  })
+  return headers
+}
+
+type ResponseHeader = readonly [name: string, value: string]
+
+const createIncomingResponseHeaders = (
+  response: IncomingMessage,
+  targetUrl: URL,
+  cookieRewrite: CookieRewriteOptions | false | undefined,
+  connectionMode: 'close' | 'upgrade',
+) => {
+  const headers: ResponseHeader[] = []
+  const setCookieHeaders: string[] = []
+  const headersToDrop = createHopByHopHeaderNames(response.headers.connection)
+
+  for (let index = 0; index < response.rawHeaders.length; index += 2) {
+    const name = response.rawHeaders[index]
+    const value = response.rawHeaders[index + 1]
+    if (!name || value === undefined) continue
+
+    const normalizedName = name.toLowerCase()
+    if (normalizedName === 'set-cookie') setCookieHeaders.push(value)
+    else if (!headersToDrop.has(normalizedName)) headers.push([name, value])
+  }
+
+  const localScopeKey = cookieRewrite
+    ? resolveCookieRewriteLocalScopeKey(cookieRewrite, targetUrl)
+    : undefined
+  const responseSetCookieHeaders = cookieRewrite
+    ? rewriteSetCookieHeadersForLocal(setCookieHeaders, {
+        ...cookieRewrite,
+        localScopeKey,
+      })
+    : setCookieHeaders
+  responseSetCookieHeaders.forEach((cookie) => headers.push(['Set-Cookie', cookie]))
+
+  if (connectionMode === 'upgrade') {
+    headers.push(['Connection', 'Upgrade'])
+    headers.push(['Upgrade', response.headers.upgrade || 'websocket'])
+  } else {
+    headers.push(['Connection', 'close'])
+  }
+  return headers
+}
+
+const writeIncomingResponseHead = (
+  socket: Duplex,
+  response: IncomingMessage,
+  headers: readonly ResponseHeader[],
+) => {
+  const statusCode = response.statusCode || 502
+  const statusMessage = response.statusMessage || 'Bad Gateway'
+  socket.write(`HTTP/1.1 ${statusCode} ${statusMessage}\r\n`)
+  headers.forEach(([name, value]) => socket.write(`${name}: ${value}\r\n`))
+  socket.write('\r\n')
+}
+
+const closeUpgradeRequest = (socket: Duplex, statusCode: number, statusMessage: string) => {
+  socket.end(
+    `HTTP/1.1 ${statusCode} ${statusMessage}\r\nConnection: close\r\nContent-Length: 0\r\n\r\n`,
+  )
+}
+
+export const createWebSocketUpgradeHandler = (
+  options: Pick<CreateDevProxyAppOptions, 'routes' | 'cors' | 'logger'>,
+) => {
+  const logger = options.logger || console
+  const allowedOrigins = options.cors?.allowedOrigins || 'local'
+
+  const handleUpgrade = (request: IncomingMessage, clientSocket: Duplex, head: Buffer) => {
+    const requestOrigin = request.headers.origin
+    if (requestOrigin && !isAllowedDevOrigin(requestOrigin, allowedOrigins)) {
+      closeUpgradeRequest(clientSocket, 403, 'Forbidden')
+      return
+    }
+
+    const requestUrl = new URL(request.url || '/', 'http://localhost')
+    const route = findProxyRoute(options.routes, requestUrl.pathname)
+    if (!route) {
+      closeUpgradeRequest(clientSocket, 404, 'Not Found')
+      return
+    }
+
+    const targetUrl = buildUpstreamUrl(route.target, requestUrl.pathname, requestUrl.search)
+    const upgradeProtocol = request.headers.upgrade || 'websocket'
+    const requestHeaders = createProxyRequestHeaders(
+      createHeadersFromIncomingMessage(request),
+      targetUrl,
+      route.cookieRewrite,
+    )
+    requestHeaders.set('connection', 'Upgrade')
+    requestHeaders.set('upgrade', upgradeProtocol)
+    const requestImpl =
+      targetUrl.protocol === 'https:'
+        ? httpsRequest
+        : targetUrl.protocol === 'http:'
+          ? httpRequest
+          : undefined
+    if (!requestImpl) {
+      logger.error(
+        '[dev-proxy]',
+        new Error(`Unsupported proxy target protocol: ${targetUrl.protocol}`),
+      )
+      closeUpgradeRequest(clientSocket, 502, 'Bad Gateway')
+      return
+    }
+
+    const upstreamRequest = requestImpl(targetUrl, {
+      headers: Object.fromEntries(requestHeaders.entries()),
+      method: request.method,
+    })
+    let upstreamSocket: Duplex | undefined
+
+    const closeUpstream = () => {
+      if (upstreamSocket) upstreamSocket.destroy()
+      else upstreamRequest.destroy()
+    }
+    clientSocket.once('close', closeUpstream)
+    clientSocket.once('error', closeUpstream)
+
+    upstreamRequest.on('upgrade', (response, socket, upstreamHead) => {
+      upstreamSocket = socket
+      if (clientSocket.destroyed) {
+        socket.destroy()
+        return
+      }
+
+      const responseHeaders = createIncomingResponseHeaders(
+        response,
+        targetUrl,
+        route.cookieRewrite,
+        'upgrade',
+      )
+      writeIncomingResponseHead(clientSocket, response, responseHeaders)
+      if (upstreamHead.length) clientSocket.write(upstreamHead)
+      if (head.length) socket.write(head)
+
+      socket.once('error', () => clientSocket.destroy())
+      socket.once('close', () => clientSocket.destroy())
+      clientSocket.pipe(socket)
+      socket.pipe(clientSocket)
+    })
+
+    upstreamRequest.on('response', (response) => {
+      const responseHeaders = createIncomingResponseHeaders(
+        response,
+        targetUrl,
+        route.cookieRewrite,
+        'close',
+      )
+      writeIncomingResponseHead(clientSocket, response, responseHeaders)
+      response.once('error', () => clientSocket.destroy())
+      response.pipe(clientSocket)
+    })
+
+    upstreamRequest.on('error', (error) => {
+      logger.error('[dev-proxy]', error)
+      if (!clientSocket.destroyed) closeUpgradeRequest(clientSocket, 502, 'Bad Gateway')
+    })
+
+    upstreamRequest.end()
+  }
+
+  return (request: IncomingMessage, clientSocket: Duplex, head: Buffer) => {
+    try {
+      handleUpgrade(request, clientSocket, head)
+    } catch (error) {
+      logger.error('[dev-proxy]', error)
+      if (!clientSocket.destroyed) closeUpgradeRequest(clientSocket, 502, 'Bad Gateway')
+    }
+  }
+}
+
 const registerProxyRoute = (
   app: Hono,
   route: DevProxyRoute,

web/.env.example

@@ -17,7 +17,7 @@ NEXT_PUBLIC_API_PREFIX=http://localhost:5001/console/api
 NEXT_PUBLIC_PUBLIC_API_PREFIX=http://localhost:5001/api
 # When the frontend and backend run on different subdomains, set NEXT_PUBLIC_COOKIE_DOMAIN=1.
 NEXT_PUBLIC_COOKIE_DOMAIN=
-# WebSocket server URL.
+# WebSocket server URL. Keep this on the local proxy when using DEV_PROXY_TARGET.
 NEXT_PUBLIC_SOCKET_URL=ws://localhost:5001
 
 # Dev proxy routes are configured in web/dev-proxy.config.ts.

web/app/components/workflow/__tests__/workflow-edge-events.spec.tsx

@@ -25,7 +25,7 @@ const reactFlowBridge = vi.hoisted(() => ({
 }))
 
 const collaborationBridge = vi.hoisted(() => ({
-  canPersistLocalGraph: vi.fn(),
+  canFlushGraphOnPageClose: vi.fn(),
   graphImportHandler: null as null | ((payload: { nodes: Node[]; edges: Edge[] }) => void),
   historyActionHandler: null as null | ((payload: unknown) => void),
   restoreIntentHandler: null as
@@ -197,7 +197,7 @@ vi.mock('@langgenius/dify-ui/toast', () => ({
 
 vi.mock('../collaboration/core/collaboration-manager', () => ({
   collaborationManager: {
-    canPersistLocalGraph: collaborationBridge.canPersistLocalGraph,
+    canFlushGraphOnPageClose: collaborationBridge.canFlushGraphOnPageClose,
     onGraphImport: (handler: (payload: { nodes: Node[]; edges: Edge[] }) => void) => {
       collaborationBridge.graphImportHandler = handler
       return vi.fn()
@@ -522,7 +522,7 @@ vi.mock('@/context/permission-state', async () => {
 describe('Workflow edge event wiring', () => {
   beforeEach(() => {
     vi.clearAllMocks()
-    collaborationBridge.canPersistLocalGraph.mockReturnValue(true)
+    collaborationBridge.canFlushGraphOnPageClose.mockReturnValue(true)
     eventEmitterState.subscription = null
     reactFlowBridge.store = null
     collaborationBridge.graphImportHandler = null
@@ -633,8 +633,8 @@ describe('Workflow edge event wiring', () => {
     })
   })
 
-  it('should skip the unmount save while the collaborative graph is not ready', () => {
-    collaborationBridge.canPersistLocalGraph.mockReturnValue(false)
+  it('should skip the unmount save when the current collaborator is not the draft leader', () => {
+    collaborationBridge.canFlushGraphOnPageClose.mockReturnValue(false)
 
     const { unmount } = renderSubject({ isCollaborationEnabled: true })
 

web/app/components/workflow/index.tsx

@@ -366,7 +366,7 @@ export const Workflow: FC<WorkflowProps> = memo(
 
     useEffect(() => {
       return () => {
-        if (isCollaborationEnabled && !collaborationManager.canPersistLocalGraph()) return
+        if (isCollaborationEnabled && !collaborationManager.canFlushGraphOnPageClose()) return
 
         handleSyncWorkflowDraft(true, true, {
           onError: () => {

web/dev-proxy.config.ts

@@ -45,7 +45,7 @@ export default {
       cookieRewrite: difyCookieRewrite,
     },
     {
-      paths: ['/console/api'],
+      paths: ['/console/api', '/socket.io'],
       target: DEV_PROXY_TARGET,
       cookieRewrite: difyCookieRewrite,
     },