experimental-a2 0.9.0 → 0.10.0

package/src/client.ts

@@ -217,16 +217,21 @@ export type A2Client<
 }
 
 /**
- * The wire the client rides. The string form is one route serving GET
+ * The wire the client rides. The first form is one route serving GET
  * (SSE stream) + POST (push). `http` splits the verbs across two
  * routes, for when platform duration limits differ per verb. `ws`
- * rides everything — every session of this client — over one
- * multiplexed WebSocket, served by `sessionsSocket`.
+ * rides everything over WebSocket, one socket per session by default.
+ * Multiplexing every session of the client is an explicit opt-in.
  */
+export type SessionUrl = string | ((sessionId: string) => string)
+
 export type ClientApi =
-  | string // one route: GET SSE stream + POST push
-  | { type: 'http'; push: string; stream: string } // split routes
-  | { type: 'ws'; url: string } // one socket, all sessions, both directions
+  | SessionUrl // one route: GET SSE stream + POST push
+  | { type: 'http'; push: SessionUrl; stream: SessionUrl } // split routes
+  /** One socket per session. */
+  | { type: 'ws'; url: SessionUrl; multiplex?: false }
+  /** One socket for every session of this client. */
+  | { type: 'ws'; url: string; multiplex: true }
 
 /**
  * The socket surface the `ws` transport drives — satisfied structurally
@@ -344,17 +349,26 @@ type Transport = {
   awaitReady?(sessionId: string, maxWaitMs: number): Promise<void>
 }
 
+const sessionUrl = (url: SessionUrl, sessionId: string): string => {
+  const resolved = typeof url === 'function' ? url(sessionId) : url
+  if (typeof resolved !== 'string') {
+    throw new TypeError('a session URL function must return a string')
+  }
+  return resolved
+}
+
 const httpTransport = (
-  routes: { push: string; stream: string },
+  routes: { push: SessionUrl; stream: SessionUrl },
   fetchImpl: typeof globalThis.fetch,
   // Contracts that declare no presence skip presence frames unparsed,
   // exactly like unknown named frames.
   presence: boolean,
 ): Transport => ({
   async *connect({ sessionId, startAfter, signal }) {
-    const sep = routes.stream.includes('?') ? '&' : '?'
+    const streamUrl = sessionUrl(routes.stream, sessionId)
+    const sep = streamUrl.includes('?') ? '&' : '?'
     const res = await fetchImpl(
-      `${routes.stream}${sep}sessionId=${encodeURIComponent(sessionId)}&index=${startAfter}`,
+      `${streamUrl}${sep}sessionId=${encodeURIComponent(sessionId)}&index=${startAfter}`,
       { headers: { accept: 'text/event-stream' }, signal },
     )
     if (!res.ok || !res.body) {
@@ -413,7 +427,7 @@ const httpTransport = (
   },
 
   async push(body) {
-    const res = await fetchImpl(routes.push, {
+    const res = await fetchImpl(sessionUrl(routes.push, body.sessionId), {
       method: 'POST',
       headers: { 'content-type': 'application/json' },
       body: JSON.stringify(body),
@@ -435,9 +449,10 @@ const httpTransport = (
   },
 
   async history({ sessionId, gte, lte }) {
-    const sep = routes.stream.includes('?') ? '&' : '?'
+    const streamUrl = sessionUrl(routes.stream, sessionId)
+    const sep = streamUrl.includes('?') ? '&' : '?'
     const res = await fetchImpl(
-      `${routes.stream}${sep}sessionId=${encodeURIComponent(sessionId)}&gte=${gte}&lte=${lte}`,
+      `${streamUrl}${sep}sessionId=${encodeURIComponent(sessionId)}&gte=${gte}&lte=${lte}`,
       { headers: { accept: 'application/json' } },
     )
     if (!res.ok) {
@@ -459,11 +474,14 @@ const httpTransport = (
   sendPresence(envelope) {
     void (async () => {
       try {
-        const res = await fetchImpl(routes.push, {
-          method: 'POST',
-          headers: { 'content-type': 'application/json' },
-          body: JSON.stringify(envelope),
-        })
+        const res = await fetchImpl(
+          sessionUrl(routes.push, envelope.sessionId),
+          {
+            method: 'POST',
+            headers: { 'content-type': 'application/json' },
+            body: JSON.stringify(envelope),
+          },
+        )
         await res.arrayBuffer()
       } catch {
         // Dropped by design — nothing was ever true.
@@ -475,6 +493,289 @@ const httpTransport = (
 /** `WebSocket.OPEN` — fixed by the spec, identical in every runtime. */
 const WS_OPEN = 1
 
+type PendingAck = {
+  resolve: (events: WireEvent[]) => void
+  reject: (error: A2Error) => void
+}
+
+const webSocketUrl = (value: string): URL => {
+  let resolved: URL
+  if (/^(https?|wss?):\/\//i.test(value)) {
+    resolved = new URL(value)
+  } else {
+    const base = (globalThis as { location?: { href: string } }).location?.href
+    if (base === undefined) {
+      throw new TypeError(
+        `cannot resolve the relative ws url '${value}' without a browser location — pass an absolute ws:// or wss:// url`,
+      )
+    }
+    resolved = new URL(value, base)
+  }
+  if (resolved.protocol === 'http:') resolved.protocol = 'ws:'
+  else if (resolved.protocol === 'https:') resolved.protocol = 'wss:'
+  return resolved
+}
+
+const socketReadyWaiters = () => {
+  const sessions = new Map<string, Set<() => void>>()
+
+  const notify = (sessionId: string): void => {
+    const waiters = sessions.get(sessionId)
+    if (!waiters) return
+    sessions.delete(sessionId)
+    for (const waiter of waiters) waiter()
+  }
+
+  const wait = (sessionId: string, maxWaitMs: number): Promise<void> =>
+    new Promise((resolve) => {
+      let waiters = sessions.get(sessionId)
+      if (!waiters) {
+        waiters = new Set()
+        sessions.set(sessionId, waiters)
+      }
+      let settled = false
+      const done = (): void => {
+        if (settled) return
+        settled = true
+        clearTimeout(timer)
+        waiters.delete(done)
+        if (waiters.size === 0) sessions.delete(sessionId)
+        // oxlint-disable-next-line promise/no-multiple-resolved -- `settled` guards the timer and notification callers
+        resolve()
+      }
+      const timer = setTimeout(done, maxWaitMs)
+      ;(timer as { unref?: () => void }).unref?.()
+      waiters.add(done)
+    })
+
+  return { notify, wait }
+}
+
+const rejectPendingAcks = (pending: Map<number, PendingAck>): void => {
+  if (pending.size === 0) return
+  const error = new A2Error(
+    'STORE_UNAVAILABLE',
+    'socket closed with the ack outstanding',
+  )
+  for (const waiter of pending.values()) waiter.reject(error)
+  pending.clear()
+}
+
+const sessionWsTransport = (
+  url: SessionUrl,
+  webSocketImpl: ClientWebSocketConstructor | undefined,
+  presence: boolean,
+): Transport => {
+  type SocketConnection = {
+    socket: ClientWebSocket
+    pending: Map<number, PendingAck>
+  }
+
+  const connections = new Map<string, SocketConnection>()
+  const ready = socketReadyWaiters()
+  let nextReq = 1
+
+  const socketUrl = (sessionId: string, startAfter: number): string => {
+    const resolved = webSocketUrl(sessionUrl(url, sessionId))
+    resolved.searchParams.set('sessionId', sessionId)
+    resolved.searchParams.set('index', String(startAfter))
+    return resolved.href
+  }
+
+  return {
+    async *connect({ sessionId, startAfter, signal }) {
+      if (webSocketImpl === undefined) {
+        throw new TypeError(
+          'no WebSocket implementation available — pass one via createClient({ webSocket }) where the global is missing',
+        )
+      }
+      const socket = new webSocketImpl(socketUrl(sessionId, startAfter))
+      const connection: SocketConnection = { socket, pending: new Map() }
+      connections.set(sessionId, connection)
+
+      const queue: TransportFrame[] = []
+      let ended = false
+      let failure: Error | null = null
+      let wake: (() => void) | null = null
+      const notify = (): void => {
+        wake?.()
+        wake = null
+      }
+      const finish = (error: Error | null): void => {
+        if (ended) return
+        ended = true
+        failure = error
+        if (connections.get(sessionId) === connection) {
+          connections.delete(sessionId)
+        }
+        rejectPendingAcks(connection.pending)
+        notify()
+      }
+
+      socket.addEventListener('open', () => {
+        if (connections.get(sessionId) !== connection || ended) return
+        queue.push({ kind: 'ping' })
+        notify()
+        ready.notify(sessionId)
+      })
+      socket.addEventListener('message', (event) => {
+        if (connections.get(sessionId) !== connection || ended) return
+        if (typeof event.data !== 'string') return
+        const frame = parseSocketFrame(event.data)
+        if (frame === null) return
+        switch (frame.kind) {
+          case 'ack': {
+            const waiter = connection.pending.get(frame.req)
+            if (waiter === undefined) return
+            connection.pending.delete(frame.req)
+            if (
+              frame.sessionId !== undefined ||
+              ('events' in frame &&
+                frame.events.some((row) => row.sessionId !== sessionId))
+            ) {
+              waiter.reject(
+                new A2Error(
+                  'STORE_UNAVAILABLE',
+                  'ack does not belong to the socket session',
+                ),
+              )
+              return
+            }
+            if ('events' in frame) waiter.resolve(frame.events)
+            else waiter.reject(frame.error)
+            return
+          }
+          case 'event':
+            if (frame.event.sessionId !== sessionId) return
+            queue.push({ kind: 'event', event: frame.event })
+            break
+          case 'presence':
+            if (presence && frame.sessionId === undefined) {
+              queue.push({ kind: 'presence', patch: frame.patch })
+            }
+            break
+          case 'presence-snapshot':
+            if (presence && frame.sessionId === undefined) {
+              queue.push({
+                kind: 'presence-snapshot',
+                snapshot: frame.snapshot,
+              })
+            }
+            break
+          case 'ping':
+            queue.push(frame)
+            break
+          case 'subscribed':
+          case 'unsubscribed':
+            return
+        }
+        notify()
+      })
+      socket.addEventListener('close', (event) => {
+        finish(
+          event.code === 1000
+            ? null
+            : new Error(
+                `socket closed (${event.code}${event.reason ? `: ${event.reason}` : ''})`,
+              ),
+        )
+      })
+      socket.addEventListener('error', () => {
+        finish(new Error('socket error'))
+      })
+
+      const onAbort = (): void => {
+        finish(
+          signal.reason instanceof Error
+            ? signal.reason
+            : new Error('stream aborted'),
+        )
+        try {
+          socket.close(1000)
+        } catch {
+          // Already closed.
+        }
+      }
+      if (signal.aborted) onAbort()
+      else signal.addEventListener('abort', onAbort, { once: true })
+
+      try {
+        for (;;) {
+          while (queue.length > 0) yield queue.shift()!
+          if (ended) {
+            if (failure) throw failure
+            return
+          }
+          await new Promise<void>((resolve) => {
+            wake = resolve
+          })
+        }
+      } finally {
+        signal.removeEventListener('abort', onAbort)
+        finish(null)
+        try {
+          socket.close(1000)
+        } catch {
+          // Already closed.
+        }
+      }
+    },
+
+    async push(body) {
+      const connection = connections.get(body.sessionId)
+      if (
+        connection === undefined ||
+        connection.socket.readyState !== WS_OPEN
+      ) {
+        throw new A2Error('STORE_UNAVAILABLE', 'no open socket for the session')
+      }
+      const req = nextReq
+      nextReq += 1
+      return await new Promise<WireEvent[]>((resolve, reject) => {
+        connection.pending.set(req, { resolve, reject })
+        const frame: SocketPushFrame = {
+          kind: 'push',
+          req,
+          events: body.events,
+        }
+        try {
+          connection.socket.send(JSON.stringify(frame))
+        } catch (cause) {
+          connection.pending.delete(req)
+          reject(
+            new A2Error('STORE_UNAVAILABLE', 'socket send failed', { cause }),
+          )
+        }
+      })
+    },
+
+    awaitReady(sessionId, maxWaitMs) {
+      const connection = connections.get(sessionId)
+      if (connection?.socket.readyState === WS_OPEN) return Promise.resolve()
+      return ready.wait(sessionId, maxWaitMs)
+    },
+
+    sendPresence(envelope) {
+      const connection = connections.get(envelope.sessionId)
+      if (
+        connection === undefined ||
+        connection.socket.readyState !== WS_OPEN
+      ) {
+        return
+      }
+      const frame: SocketPresenceFrame = {
+        kind: 'presence',
+        ...envelope.presence,
+      }
+      try {
+        connection.socket.send(JSON.stringify(frame))
+      } catch {
+        return
+      }
+    },
+  }
+}
+
 /** One session's lane on the shared ws socket — created per `connect`
  * attempt, so `startAfter` is that attempt's resume frontier. */
 type WsChannel = {
@@ -502,15 +803,11 @@ const subscribeFrame = (
   return JSON.stringify(frame)
 }
 
-const wsTransport = (
+const multiplexedWsTransport = (
   url: string,
   webSocketImpl: ClientWebSocketConstructor | undefined,
   presence: boolean,
 ): Transport => {
-  type PendingAck = {
-    resolve: (events: WireEvent[]) => void
-    reject: (error: A2Error) => void
-  }
   // ONE socket serves every session of this client (specs §13's
   // multiplexed superset): channels subscribe on it at their own
   // frontiers, down frames route back by sessionId, acks by req.
@@ -518,41 +815,12 @@ const wsTransport = (
   const pending = new Map<number, PendingAck>()
   /** Push retries parked in `awaitReady`, flushed per session on its
    * subscribe confirmation. */
-  const readyWaiters = new Map<string, Set<() => void>>()
+  const ready = socketReadyWaiters()
   let socket: ClientWebSocket | null = null
   let socketOpen = false
   let watchdog: ReturnType<typeof setTimeout> | undefined
   let nextReq = 1
 
-  const flushReadyWaiters = (sessionId: string): void => {
-    const waiters = readyWaiters.get(sessionId)
-    if (!waiters) return
-    readyWaiters.delete(sessionId)
-    for (const waiter of waiters) waiter()
-  }
-
-  const socketUrl = (): string => {
-    let resolved: URL
-    if (/^(https?|wss?):\/\//i.test(url)) {
-      resolved = new URL(url)
-    } else {
-      // A relative url resolves against the page, exactly like fetch.
-      // Outside a browser (SSR, tests) there is nothing to resolve
-      // against — demand an absolute url there.
-      const base = (globalThis as { location?: { href: string } }).location
-        ?.href
-      if (base === undefined) {
-        throw new TypeError(
-          `cannot resolve the relative ws url '${url}' without a browser location — pass an absolute ws:// or wss:// url`,
-        )
-      }
-      resolved = new URL(url, base)
-    }
-    if (resolved.protocol === 'http:') resolved.protocol = 'ws:'
-    else if (resolved.protocol === 'https:') resolved.protocol = 'wss:'
-    return resolved.href
-  }
-
   const trySend = (data: string): void => {
     if (socket === null || !socketOpen) return
     try {
@@ -584,14 +852,7 @@ const wsTransport = (
     socketOpen = false
     clearTimeout(watchdog)
     watchdog = undefined
-    if (pending.size > 0) {
-      const rejection = new A2Error(
-        'STORE_UNAVAILABLE',
-        'socket closed with the ack outstanding',
-      )
-      for (const waiter of pending.values()) waiter.reject(rejection)
-      pending.clear()
-    }
+    rejectPendingAcks(pending)
     return current
   }
 
@@ -662,7 +923,7 @@ const wsTransport = (
         if (channel === undefined || channel.ended || channel.live) return
         channel.live = true
         deliver(channel, { kind: 'ping' }) // the connected signal
-        flushReadyWaiters(frame.sessionId)
+        ready.notify(frame.sessionId)
         return
       }
       case 'unsubscribed': {
@@ -723,7 +984,7 @@ const wsTransport = (
         'no WebSocket implementation available — pass one via createClient({ webSocket }) where the global is missing',
       )
     }
-    const created = new webSocketImpl(socketUrl())
+    const created = new webSocketImpl(webSocketUrl(url).href)
     socket = created
     socketOpen = false
     armWatchdog(created)
@@ -881,29 +1142,7 @@ const wsTransport = (
       ) {
         return Promise.resolve()
       }
-      return new Promise((resolve) => {
-        let waiters = readyWaiters.get(sessionId)
-        if (!waiters) {
-          waiters = new Set()
-          readyWaiters.set(sessionId, waiters)
-        }
-        let settled = false
-        const done = (): void => {
-          if (settled) return
-          settled = true
-          clearTimeout(timer)
-          waiters.delete(done)
-          if (waiters.size === 0) readyWaiters.delete(sessionId)
-          // oxlint-disable-next-line promise/no-multiple-resolved -- `settled` guards the two callers (timer, open flush); the rule cannot see through the flag
-          resolve()
-        }
-        // The cap keeps a session nobody reconnects (write-only usage,
-        // an unreachable server) from parking a push forever — the
-        // attempt then fails fast and the loop moves on.
-        const timer = setTimeout(done, maxWaitMs)
-        ;(timer as { unref?: () => void }).unref?.()
-        waiters.add(done)
-      })
+      return ready.wait(sessionId, maxWaitMs)
     },
     sendPresence(envelope) {
       const channel = channels.get(envelope.sessionId)
@@ -926,7 +1165,7 @@ const transportFor = (
   webSocketImpl: ClientWebSocketConstructor | undefined,
   presence: boolean,
 ): Transport => {
-  if (typeof api === 'string') {
+  if (typeof api === 'string' || typeof api === 'function') {
     return httpTransport({ push: api, stream: api }, fetchImpl, presence)
   }
   if (api.type === 'http') {
@@ -936,7 +1175,13 @@ const transportFor = (
       presence,
     )
   }
-  return wsTransport(api.url, webSocketImpl, presence)
+  if (api.multiplex === true) {
+    if (typeof api.url !== 'string') {
+      throw new TypeError('a multiplexed WebSocket URL must be a string')
+    }
+    return multiplexedWsTransport(api.url, webSocketImpl, presence)
+  }
+  return sessionWsTransport(api.url, webSocketImpl, presence)
 }
 
 type Pending<D extends EventDefs> = {

package/src/server-fetch.ts

@@ -14,6 +14,7 @@ import {
 } from './push-envelope.ts'
 import type { A2Server } from './server.ts'
 import {
+  sessionSocket,
   sessionsSocket,
   type A2Socket,
   type SocketSession,
@@ -76,6 +77,8 @@ export type ServerFetchOptions<
 > = {
   authorize?: (operation: A2Operation<D, P>) => boolean | Promise<boolean>
   upgradeWebSocket?: UpgradeWebSocket
+  /** Allow one upgraded connection to select more than one session. */
+  multiplexWebSocket?: boolean
 }
 
 export type ServerIngressContext = {
@@ -212,7 +215,27 @@ export function createServerFetch<
       await serverFetchHooks.get(fetch)?.validateIngress?.(context)
     }
 
-    const attach = (socket: A2Socket): void => {
+    const gateSocketPush = async (
+      sessionId: string,
+      events: ParsedPushEvent[],
+      pushedPresence?: ParsedPushPresence,
+    ): Promise<A2Error | null> => {
+      try {
+        await authorize(
+          freezePush<D, P>(sessionId, events, pushedPresence, 'websocket'),
+        )
+        await validateIngress({
+          sessionId,
+          events,
+          ...(pushedPresence === undefined ? {} : { presence: pushedPresence }),
+        })
+        return null
+      } catch (error) {
+        return asA2Error(error)
+      }
+    }
+
+    const attachMultiplexed = (socket: A2Socket): void => {
       const deadline = invocationDeadlineMs()
       sessionsSocket<D>(
         async (sessionId, startAfter) => {
@@ -230,28 +253,7 @@ export function createServerFetch<
         {
           presence,
           ...(deadline === null ? {} : { deadline }),
-          gatePush: async (sessionId, events, pushedPresence) => {
-            try {
-              await authorize(
-                freezePush<D, P>(
-                  sessionId,
-                  events,
-                  pushedPresence,
-                  'websocket',
-                ),
-              )
-              await validateIngress({
-                sessionId,
-                events,
-                ...(pushedPresence === undefined
-                  ? {}
-                  : { presence: pushedPresence }),
-              })
-              return null
-            } catch (error) {
-              return asA2Error(error)
-            }
-          },
+          gatePush: gateSocketPush,
         },
       )
     }
@@ -265,6 +267,32 @@ export function createServerFetch<
               { status: 426 },
             )
           }
+          if (options.multiplexWebSocket === true) {
+            return await options.upgradeWebSocket(attachMultiplexed)
+          }
+
+          const { searchParams } = new URL(request.url)
+          const sessionId = searchParams.get('sessionId')
+          if (sessionId === null || sessionId.length === 0) {
+            throw new A2Error('INVALID_PAYLOAD', 'missing sessionId')
+          }
+          const startAfter = parseResumeIndex(searchParams.get('index'))
+          await authorize({
+            type: 'stream',
+            sessionId,
+            startAfter,
+            transport: 'websocket',
+          })
+          const target = server.session(sessionId)
+          const attach = (socket: A2Socket): void => {
+            const deadline = invocationDeadlineMs()
+            sessionSocket(sessionId, target, socket, {
+              startAfter,
+              presence,
+              ...(deadline === null ? {} : { deadline }),
+              gatePush: gateSocketPush,
+            })
+          }
           return await options.upgradeWebSocket(attach)
         }