feat: support ws.onUpgrade for handling connection upgrades (#2732)

Author: kettanaito

src/core/experimental/frames/http-frame.test.ts

@@ -4,7 +4,8 @@ import { ws } from '../../ws'
 import { bypass } from '../../bypass'
 import type { HttpNetworkFrameEventMap } from './http-frame'
 import { HttpNetworkFrame } from './http-frame'
-import { InMemoryHandlersController } from '#core/experimental/handlers-controller'
+import { InMemoryHandlersController } from '../../experimental/handlers-controller'
+import { getSiblingHandlers } from '../../utils/internal/attachSiblingHandlers'
 
 beforeAll(() => {
   vi.spyOn(console, 'error').mockImplementation(() => {})
@@ -46,6 +47,9 @@ it('filters only request type handlers', async () => {
   const webSocketHandlers = [
     ws.link('ws://localhost').addEventListener('connection', () => {}),
   ]
+  const webSocketSiblingHandlers = webSocketHandlers.flatMap((handler) =>
+    getSiblingHandlers(handler),
+  )
 
   const controller = new InMemoryHandlersController([
     ...httpHandlers,
@@ -55,6 +59,7 @@ it('filters only request type handlers', async () => {
 
   expect(frame.getHandlers(controller)).toEqual([
     ...httpHandlers,
+    ...webSocketSiblingHandlers,
     ...graphqlHandlers,
   ])
   expect(frame.getHandlers(new InMemoryHandlersController([]))).toEqual([])

src/core/experimental/handlers-controller.test.ts

@@ -1,8 +1,68 @@
 import { http } from '../http'
 import { graphql } from '../graphql'
 import { ws } from '../ws'
+import { getSiblingHandlers } from '../utils/internal/attachSiblingHandlers'
 import { InMemoryHandlersController } from './handlers-controller'
 
+describe('constructor', () => {
+  it('places the sibling in its own kind bucket', () => {
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    const controller = new InMemoryHandlersController([wsHandler])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('interleaves the sibling at the owner position when grouping by kind', () => {
+    const httpOne = http.get('/', () => {})
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+    const httpTwo = http.get('/', () => {})
+
+    const controller = new InMemoryHandlersController([
+      httpOne,
+      wsHandler,
+      httpTwo,
+    ])
+
+    expect(controller.getHandlersByKind('request')).toEqual([
+      httpOne,
+      upgradeHandler,
+      httpTwo,
+    ])
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+  })
+
+  it('extracts siblings from every owner in the input list', () => {
+    const wsOne = ws.link('*').addEventListener('connection', () => {})
+    const wsTwo = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeOne] = getSiblingHandlers(wsOne)
+    const [upgradeTwo] = getSiblingHandlers(wsTwo)
+
+    const controller = new InMemoryHandlersController([wsOne, wsTwo])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsOne, wsTwo])
+    expect(controller.getHandlersByKind('request')).toEqual([
+      upgradeOne,
+      upgradeTwo,
+    ])
+  })
+
+  it('dedupes the shared upgrade sibling across multiple handlers from the same link', () => {
+    const chat = ws.link('*')
+    const wsOne = chat.addEventListener('connection', () => {})
+    const wsTwo = chat.addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsOne)
+
+    const controller = new InMemoryHandlersController([wsOne, wsTwo])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsOne, wsTwo])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+})
+
 describe(InMemoryHandlersController.prototype.use, () => {
   it('prepends a handler to an empty controller', () => {
     const controller = new InMemoryHandlersController([])
@@ -51,6 +111,44 @@ describe(InMemoryHandlersController.prototype.use, () => {
 
     expect(controller.currentHandlers()).toEqual([graphqlOne, httpTwo, httpOne])
   })
+
+  it('propagates siblings to their kind buckets at runtime', () => {
+    const controller = new InMemoryHandlersController([])
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    controller.use([wsHandler])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('prepends incoming siblings before existing handlers of the same kind', () => {
+    const existingHttp = http.get('/existing', () => {})
+    const controller = new InMemoryHandlersController([existingHttp])
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    controller.use([wsHandler])
+
+    expect(controller.getHandlersByKind('request')).toEqual([
+      upgradeHandler,
+      existingHttp,
+    ])
+  })
+
+  it('dedupes the shared upgrade sibling when called with multiple handlers from the same link', () => {
+    const chat = ws.link('*')
+    const wsOne = chat.addEventListener('connection', () => {})
+    const wsTwo = chat.addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsOne)
+
+    const controller = new InMemoryHandlersController([])
+    controller.use([wsOne, wsTwo])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsOne, wsTwo])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
 })
 
 describe(InMemoryHandlersController.prototype.reset, () => {
@@ -96,6 +194,42 @@ describe(InMemoryHandlersController.prototype.reset, () => {
      */
     expect(controller.currentHandlers()).toEqual([httpTwo])
   })
+
+  it('places siblings into their kind buckets when resetting to next handlers', () => {
+    const controller = new InMemoryHandlersController([])
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+
+    controller.reset([wsHandler])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('restores siblings when resetting to the initial handlers', () => {
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsHandler)
+    const controller = new InMemoryHandlersController([wsHandler])
+
+    controller.use([http.get('/runtime', () => {})])
+    controller.reset([])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsHandler])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
+
+  it('dedupes the shared upgrade sibling when reset with multiple handlers from the same link', () => {
+    const chat = ws.link('*')
+    const wsOne = chat.addEventListener('connection', () => {})
+    const wsTwo = chat.addEventListener('connection', () => {})
+    const [upgradeHandler] = getSiblingHandlers(wsOne)
+
+    const controller = new InMemoryHandlersController([])
+    controller.reset([wsOne, wsTwo])
+
+    expect(controller.getHandlersByKind('websocket')).toEqual([wsOne, wsTwo])
+    expect(controller.getHandlersByKind('request')).toEqual([upgradeHandler])
+  })
 })
 
 describe(InMemoryHandlersController.prototype.getHandlersByKind, () => {
@@ -112,11 +246,10 @@ describe(InMemoryHandlersController.prototype.getHandlersByKind, () => {
       ]).getHandlersByKind('websocket'),
     ).toEqual([])
 
+    const wsHandler = ws.link('*').addEventListener('connection', () => {})
     expect(
-      new InMemoryHandlersController([
-        ws.link('*').addEventListener('connection', () => {}),
-      ]).getHandlersByKind('request'),
-    ).toEqual([])
+      new InMemoryHandlersController([wsHandler]).getHandlersByKind('request'),
+    ).toEqual(getSiblingHandlers(wsHandler))
   })
 
   it('returns all handlers if they all match', () => {
@@ -142,14 +275,15 @@ describe(InMemoryHandlersController.prototype.getHandlersByKind, () => {
     const httpHandler = http.get('/', () => {})
     const graphqlHandler = graphql.query('', () => {})
     const wsHandler = ws.link('*').addEventListener('connection', () => {})
+    const wsHandlerSiblings = getSiblingHandlers(wsHandler)
 
     expect(
       new InMemoryHandlersController([
         httpHandler,
         graphqlHandler,
         wsHandler,
       ]).getHandlersByKind('request'),
-    ).toEqual([httpHandler, graphqlHandler])
+    ).toEqual([httpHandler, graphqlHandler, ...wsHandlerSiblings])
 
     expect(
       new InMemoryHandlersController([

src/core/experimental/handlers-controller.ts

@@ -2,18 +2,31 @@ import { invariant } from 'outvariant'
 import { type RequestHandler } from '../handlers/RequestHandler'
 import { type WebSocketHandler } from '../handlers/WebSocketHandler'
 import { devUtils } from '../utils/internal/devUtils'
+import { getSiblingHandlers } from '../utils/internal/attachSiblingHandlers'
 
 export type AnyHandler = RequestHandler | WebSocketHandler
 export type HandlersMap = Partial<Record<AnyHandler['kind'], Array<AnyHandler>>>
 
 export function groupHandlersByKind(handlers: Array<AnyHandler>): HandlersMap {
   const groups: HandlersMap = {}
 
+  const pushUnique = (kind: AnyHandler['kind'], handler: AnyHandler) => {
+    const bucket = (groups[kind] ||= [])
+
+    if (!bucket.includes(handler)) {
+      bucket.push(handler)
+    }
+  }
+
   /**
    * @note `Object.groupBy` is not implemented in Node.js v20.
    */
   for (const handler of handlers) {
-    ;(groups[handler.kind] ||= []).push(handler)
+    pushUnique(handler.kind, handler)
+
+    for (const sibling of getSiblingHandlers(handler)) {
+      pushUnique(sibling.kind, sibling)
+    }
   }
 
   return groups
@@ -69,14 +82,16 @@ export abstract class HandlersController {
     }
 
     const { handlers } = this.getState()
-
-    // Iterate over next handlers and prepend them to their respective lists.
-    // Iterate in a reverse order to the keep the order of the runtime handlers as provided.
-    for (let i = nextHandlers.length - 1; i >= 0; i--) {
-      const handler = nextHandlers[i]
-      handlers[handler.kind] = handlers[handler.kind]
-        ? [handler, ...handlers[handler.kind]!]
-        : [handler]
+    const overrides = groupHandlersByKind(nextHandlers)
+
+    // Prepend overrides to their respective kind buckets so they take
+    // priority over existing handlers while preserving input order.
+    for (const kind in overrides) {
+      const overridesForKind = overrides[kind as AnyHandler['kind']]!
+      const existingForKind = handlers[kind as AnyHandler['kind']]
+      handlers[kind as AnyHandler['kind']] = existingForKind
+        ? [...overridesForKind, ...existingForKind]
+        : overridesForKind
     }
 
     this.setState({ handlers })

src/core/handlers/WebSocketHandler.ts

@@ -118,7 +118,7 @@ export class WebSocketHandler {
       params: parsedResult.match.params || {},
     }
 
-    if (resolutionContext?.[kAutoConnect]) {
+    if (resolutionContext?.[kAutoConnect] ?? true) {
       if (this[kConnect](resolvedConnection)) {
         return resolvedConnection
       }

src/core/utils/internal/attachSiblingHandlers.ts

@@ -0,0 +1,28 @@
+import { invariant } from 'outvariant'
+import type { AnyHandler } from '../../experimental/handlers-controller'
+
+const kSiblingHandlers = Symbol('kSiblingHandlers')
+
+export function attachSiblingHandlers<T extends AnyHandler>(
+  owner: T,
+  siblings: Array<AnyHandler>,
+): T {
+  invariant(
+    getSiblingHandlers(owner).length === 0,
+    'Failed to merge handlers: the owner "%s" handler is already merged',
+    owner.kind,
+  )
+
+  Object.defineProperty(owner, kSiblingHandlers, {
+    value: siblings,
+    enumerable: false,
+    writable: false,
+    configurable: false,
+  })
+
+  return owner
+}
+
+export function getSiblingHandlers(owner: AnyHandler): Array<AnyHandler> {
+  return Reflect.get(owner, kSiblingHandlers) || []
+}