@livestore/common-cf 0.4.0-dev.9 → 0.5.0-dev.0

package/src/ws-rpc/ws-rpc-server.ts

@@ -17,21 +17,30 @@
 
 import { notYetImplemented, omitUndefineds } from '@livestore/utils'
 import {
-  constVoid,
+  Context,
   Effect,
+  Function,
   Exit,
   Layer,
   Logger,
-  LogLevel,
-  Mailbox,
+  Queue,
+  References,
   RpcMessage,
   RpcSerialization,
   RpcServer,
   Scope,
   Stream,
 } from '@livestore/utils/effect'
+
 import type * as CfTypes from '../cf-types.ts'
 
+/**
+ * Context service providing access to the current WebSocket.
+ * This is useful for reading WebSocket attachment data (e.g., forwarded headers)
+ * inside RPC handlers.
+ */
+export class WsContext extends Context.Service<WsContext, { readonly ws: CfTypes.WebSocket }>()('WsContext') {}
+
 /**
  * Configuration options for setting up WebSocket RPC on a Durable Object.
  */
@@ -44,11 +53,25 @@ export interface DurableObjectWebSocketRpcConfig {
    * - 'accept': Use traditional WebSocket handling (not yet implemented)
    */
   webSocketMode: 'hibernate' | 'accept'
-  /** Effect RPC layer that defines the available RPC methods and handlers */
-  rpcLayer: Layer.Layer<never, never, RpcServer.Protocol>
+  /**
+   * Effect RPC layer that requires `RpcServer.Protocol` (and `WsContext` if used)
+   * and provides the RPC server services.
+   *
+   * This is typically created by:
+   * ```typescript
+   * RpcServer.layer(MyRpcs).pipe(Layer.provide(handlersLayer))
+   * ```
+   *
+   * `WsContext` is provided by the WebSocket protocol layer, so handlers can access
+   * WebSocket attachment data (e.g., forwarded headers stored after WebSocket upgrade).
+   *
+   * The layer requirements (`RIn`) must be a subset of `RpcServer.Protocol | WsContext`,
+   * which are both provided by the WebSocket protocol layer.
+   */
+  rpcLayer: Layer.Layer<never, never, RpcServer.Protocol | WsContext>
   /** Function to get access to incoming requests */
   onMessage?: (msg: RpcMessage.FromClientEncoded, ws: CfTypes.WebSocket) => void
-  mainLayer?: Layer.Layer<never, never, never>
+  mainLayer?: Layer.Layer<never>
 }
 
 /**
@@ -125,14 +148,14 @@ export const setupDurableObjectWebSocketRpc = ({
   const serverCtxMap = new Map<
     CfTypes.WebSocket,
     {
-      scope: Scope.CloseableScope
+      scope: Scope.Closeable
       onMessage: (message: string | ArrayBuffer) => Promise<void>
     }
   >()
 
   const launchServer = (ws: CfTypes.WebSocket) =>
     Effect.gen(function* () {
-      if (serverCtxMap.has(ws)) {
+      if (serverCtxMap.has(ws) === true) {
         return serverCtxMap.get(ws)!
       }
 
@@ -140,12 +163,13 @@ export const setupDurableObjectWebSocketRpc = ({
 
       const scope = yield* Scope.make()
 
-      const incomingQueue = yield* Mailbox.make<Uint8Array<ArrayBufferLike> | string>()
+      const incomingQueue = yield* Queue.unbounded<Uint8Array | string>()
 
-      yield* Scope.addFinalizer(scope, incomingQueue.shutdown)
+      yield* Scope.addFinalizer(scope, Queue.shutdown(incomingQueue).pipe(Effect.asVoid))
 
       const ProtocolLive = layerRpcServerWebsocket({
         ws,
+        scope,
         incomingQueue,
         ...omitUndefineds({ onMessage }),
       }).pipe(Layer.provide(RpcSerialization.layerJson))
@@ -154,19 +178,16 @@ export const setupDurableObjectWebSocketRpc = ({
 
       yield* Layer.launch(ServerLive).pipe(Effect.tapCauseLogPretty, Effect.forkIn(scope))
 
-      const runtime = yield* Effect.runtime()
+      const services = yield* Effect.context()
 
       const ctx = {
         scope,
         onMessage: (message: string | ArrayBuffer) =>
-          incomingQueue
-            .offer(message as Uint8Array<ArrayBufferLike> | string)
-            .pipe(
-              Effect.asVoid,
-              Effect.withSpan('ws-rpc-server/onMessage', { root: true }),
-              Effect.provide(runtime),
-              Effect.runPromise,
-            ),
+          Queue.offer(incomingQueue, message as Uint8Array | string).pipe(
+            Effect.asVoid,
+            Effect.withSpan('ws-rpc-server/onMessage', { root: true }),
+            Effect.runPromiseWith(services),
+          ),
       }
 
       serverCtxMap.set(ws, ctx)
@@ -174,8 +195,14 @@ export const setupDurableObjectWebSocketRpc = ({
       return ctx
     }).pipe(
       Effect.tapCauseLogPretty,
-      Logger.withMinimumLogLevel(LogLevel.Debug), // Useful for debugging
-      Effect.provide(Layer.mergeAll(Logger.consoleWithThread('ws-rpc-server'), mainLayer ?? Layer.empty)),
+      Effect.annotateLogs({ thread: 'ws-rpc-server' }),
+      Effect.provide(
+        Layer.mergeAll(
+          Logger.layer([Logger.consoleStructured]),
+          Layer.succeed(References.MinimumLogLevel, 'Debug'), // Useful for debugging
+          mainLayer ?? Layer.empty,
+        ),
+      ),
       Effect.withSpan('effect-ws-rpc-server'),
       Effect.runPromise,
     )
@@ -190,7 +217,7 @@ export const setupDurableObjectWebSocketRpc = ({
   const webSocketClose: CfTypes.DurableObject['webSocketClose'] = async (ws, _code, _reason, _wasClean) => {
     const ctx = serverCtxMap.get(ws)
     // console.log('webSocketClose', ctx, ws)
-    if (ctx) {
+    if (ctx !== undefined) {
       await Scope.close(ctx.scope, Exit.void).pipe(Effect.runPromise)
       serverCtxMap.delete(ws)
     }
@@ -210,9 +237,10 @@ export const setupDurableObjectWebSocketRpc = ({
  */
 export interface WsRpcServerArgs {
   ws: CfTypes.WebSocket
+  scope: Scope.Scope
   onMessage?: (message: RpcMessage.FromClientEncoded, ws: CfTypes.WebSocket) => void
-  /** Mailbox queue for receiving incoming messages from the WebSocket */
-  incomingQueue: Mailbox.Mailbox<Uint8Array<ArrayBufferLike> | string>
+  /** Queue for receiving incoming messages from the WebSocket */
+  incomingQueue: Queue.Queue<Uint8Array | string>
 }
 
 /**
@@ -221,13 +249,19 @@ export interface WsRpcServerArgs {
  * This layer handles the low-level WebSocket protocol details for RPC communication,
  * including message serialization, routing, and error handling.
  *
+ * Also provides `WsContext` with the current WebSocket so handlers can access
+ * WebSocket attachment data (e.g., forwarded headers).
+ *
  * @param args Configuration for WebSocket RPC protocol
- * @returns Effect layer that provides RPC server protocol functionality
+ * @returns Effect layer that provides RPC server protocol functionality and WsContext
  *
  * @internal This is typically used internally by `setupDurableObjectWebSocketRpc`
  */
 export const layerRpcServerWebsocket = (args: WsRpcServerArgs) =>
-  Layer.scoped(RpcServer.Protocol, makeSocketProtocol(args))
+  Layer.mergeAll(
+    Layer.effect(RpcServer.Protocol, makeSocketProtocol(args)),
+    Layer.succeed(WsContext, WsContext.of({ ws: args.ws })),
+  )
 
 /**
  * Creates the low-level RPC protocol implementation for WebSocket communication.
@@ -240,16 +274,16 @@ export const layerRpcServerWebsocket = (args: WsRpcServerArgs) =>
  *
  * @internal Used internally by `layerRpcServerWebsocket`
  */
-const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =>
+const makeSocketProtocol = ({ incomingQueue, scope, ws, onMessage }: WsRpcServerArgs) =>
   Effect.gen(function* () {
     const serialization = yield* RpcSerialization.RpcSerialization
-    const disconnects = yield* Mailbox.make<number>()
+    const disconnects = yield* Queue.unbounded<number>()
 
-    const writeRaw = (msg: Uint8Array<ArrayBufferLike> | string) => Effect.succeed(ws.send(msg))
+    const writeRaw = (msg: Uint8Array | string) => Effect.succeed(ws.send(msg))
 
     let writeRequest!: (clientId: number, message: RpcMessage.FromClientEncoded) => Effect.Effect<void>
 
-    const parser = serialization.unsafeMake()
+    const parser = serialization.makeUnsafe()
     const id = 0
 
     const write = (response: RpcMessage.FromServerEncoded) => {
@@ -268,7 +302,7 @@ const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =
       writeRequest = writeRequest_
 
       // Start processing messages now that writeRequest is available
-      const startProcessing = Mailbox.toStream(incomingQueue).pipe(
+      const startProcessing = Stream.fromQueue(incomingQueue).pipe(
         Stream.tap((data) => {
           try {
             const decoded = parser.decode(data) as ReadonlyArray<RpcMessage.FromClientEncoded>
@@ -278,12 +312,12 @@ const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =
               while: () => i < decoded.length,
               body: () => {
                 const request = decoded[i++]!
-                if (onMessage) {
+                if (onMessage !== undefined) {
                   onMessage(request, ws)
                 }
                 return writeRequest(id, request)
               },
-              step: constVoid,
+              step: Function.constVoid,
             })
           } catch (cause) {
             return Effect.orDie(writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!))
@@ -291,7 +325,7 @@ const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =
         }),
         Stream.runDrain,
         Effect.tapCauseLogPretty,
-        Effect.fork,
+        Effect.forkIn(scope),
       )
 
       // Start the message processing
@@ -307,6 +341,7 @@ const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =
         supportsAck: true,
         supportsTransferables: false,
         supportsSpanPropagation: true,
+        supportsNotifications: true,
       }))
     })
 

package/src/ws-rpc/ws-rpc.test.ts

@@ -1,20 +1,20 @@
+import { expect } from 'vitest'
+
+import { Vitest } from '@livestore/utils-dev/node-vitest'
+import { WranglerDevServer } from '@livestore/utils-dev/wrangler'
 import {
-  Chunk,
   Effect,
   FetchHttpClient,
   Layer,
-  Logger,
-  LogLevel,
   Option,
+  References,
   RpcClient,
   RpcSerialization,
   Socket,
   Stream,
 } from '@livestore/utils/effect'
 import { PlatformNode } from '@livestore/utils/node'
-import { Vitest } from '@livestore/utils-dev/node-vitest'
-import { WranglerDevServerService } from '@livestore/utils-dev/wrangler'
-import { expect } from 'vitest'
+
 import { TestRpcs } from './test-fixtures/rpc-schema.ts'
 
 const testTimeout = 60_000
@@ -22,29 +22,33 @@ const testTimeout = 60_000
 const withWranglerTest = Vitest.makeWithTestCtx({
   timeout: testTimeout,
   makeLayer: () =>
-    WranglerDevServerService.Default({
+    WranglerDevServer.layer({
       cwd: `${import.meta.dirname}/test-fixtures`,
     }).pipe(
       Layer.provide(
-        Layer.mergeAll(PlatformNode.NodeContext.layer, FetchHttpClient.layer, Logger.minimumLogLevel(LogLevel.Debug)),
+        Layer.mergeAll(
+          PlatformNode.NodeServices.layer,
+          FetchHttpClient.layer,
+          Layer.succeed(References.MinimumLogLevel, 'Debug'),
+        ),
       ),
     ),
 })
 
 const ProtocolLive = Layer.suspend(() =>
   Effect.gen(function* () {
-    const server = yield* WranglerDevServerService
+    const server = yield* WranglerDevServer.WranglerDevServer
     return RpcClient.layerProtocolSocket().pipe(
       Layer.provide(Socket.layerWebSocket(`ws://localhost:${server.port}`)),
       Layer.provide(Socket.layerWebSocketConstructorGlobal),
       Layer.provide(RpcSerialization.layerJson),
     )
-  }).pipe(Layer.unwrapEffect),
+  }).pipe(Layer.unwrap),
 )
 
 Vitest.describe('Durable Object WebSocket RPC', { timeout: testTimeout }, () => {
   // Direct HTTP RPC client tests
-  Vitest.scopedLive('should call ping method', (test) =>
+  Vitest.live('should call ping method', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const result = yield* client.Ping({ message: 'Hello HTTP RPC' })
@@ -52,7 +56,7 @@ Vitest.describe('Durable Object WebSocket RPC', { timeout: testTimeout }, () =>
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should call echo method', (test) =>
+  Vitest.live('should call echo method', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const result = yield* client.Echo({ text: 'Echo' })
@@ -60,7 +64,7 @@ Vitest.describe('Durable Object WebSocket RPC', { timeout: testTimeout }, () =>
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should call add method', (test) =>
+  Vitest.live('should call add method', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const result = yield* client.Add({ a: 15, b: 25 })
@@ -68,104 +72,72 @@ Vitest.describe('Durable Object WebSocket RPC', { timeout: testTimeout }, () =>
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle RPC fail method', (test) =>
+  Vitest.live('should handle RPC fail method', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const error = yield* client.Fail({ message: 'test http failure' }).pipe(Effect.exit)
       expect(error.toString()).toMatchInlineSnapshot(`
-        "{
-          "_id": "Exit",
-          "_tag": "Failure",
-          "cause": {
-            "_id": "Cause",
-            "_tag": "Fail",
-            "failure": "RPC failure: test http failure"
-          }
-        }"
+        "Failure(Cause([Fail("RPC failure: test http failure")]))"
       `)
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle defect method', (test) =>
+  Vitest.live('should handle defect method', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const error = yield* client.Defect({ message: 'test http defect' }).pipe(Effect.exit)
       expect(error.toString()).toMatchInlineSnapshot(`
-        "{
-          "_id": "Exit",
-          "_tag": "Failure",
-          "cause": {
-            "_id": "Cause",
-            "_tag": "Die",
-            "defect": "some defect: test http defect"
-          }
-        }"
+        "Failure(Cause([Die("some defect: test http defect")]))"
       `)
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle streaming RPC via HTTP', (test) =>
+  Vitest.live('should handle streaming RPC via HTTP', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const stream = client.Stream({}).pipe(
         Stream.take(4),
         Stream.map((c) => c.maybeNumber.pipe(Option.getOrUndefined)),
       )
-      const chunks = yield* Stream.runCollect(stream)
-      expect(Chunk.toReadonlyArray(chunks)).toEqual([1, 4, 9, 16]) // squares of 1,2,3,4
+      const result = yield* Stream.runCollect(stream)
+      expect(result).toEqual([1, 4, 9, 16]) // squares of 1,2,3,4
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle streaming RPC with error via HTTP', (test) =>
+  Vitest.live('should handle streaming RPC with error via HTTP', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const stream = client.StreamError({ count: 5, errorAfter: 4 })
       const error = yield* Stream.runCollect(stream).pipe(Effect.exit)
       expect(error.toString()).toMatchInlineSnapshot(`
-        "{
-          "_id": "Exit",
-          "_tag": "Failure",
-          "cause": {
-            "_id": "Cause",
-            "_tag": "Fail",
-            "failure": "Stream error after 4: got 9"
-          }
-        }"
+        "Failure(Cause([Fail("Stream error after 4: got 9")]))"
       `)
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle streaming RPC with defect via HTTP', (test) =>
+  Vitest.live('should handle streaming RPC with defect via HTTP', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const stream = client.StreamDefect({ count: 4, defectAfter: 1 })
       const error = yield* Stream.runCollect(stream).pipe(Effect.exit)
       expect(error.toString()).toMatchInlineSnapshot(`
-        "{
-          "_id": "Exit",
-          "_tag": "Failure",
-          "cause": {
-            "_id": "Cause",
-            "_tag": "Die",
-            "defect": "Stream defect after 1: got 4"
-          }
-        }"
+        "Failure(Cause([Die("Stream defect after 1: got 4")]))"
       `)
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle stream interruption via HTTP', (test) =>
+  Vitest.live('should handle stream interruption via HTTP', (test) =>
     Effect.gen(function* () {
       const client = yield* RpcClient.make(TestRpcs)
       const stream = client.StreamInterruptible({ delay: 50, interruptAfterCount: 3 }).pipe(Stream.take(3))
-      const chunks = yield* Stream.runCollect(stream)
-      expect(Chunk.toReadonlyArray(chunks)).toEqual([1, 2, 3])
+      const result = yield* Stream.runCollect(stream)
+      expect(result).toEqual([1, 2, 3])
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 })
 
 Vitest.describe('Hibernation Tests', { timeout: 25000 }, () => {
-  Vitest.scopedLive('should maintain RPC functionality after hibernation', (test) =>
+  Vitest.live('should maintain RPC functionality after hibernation', (test) =>
     Effect.gen(function* () {
       console.log('🧪 Testing RPC server persistence across hibernation...')
 
@@ -220,15 +192,15 @@ Vitest.describe('Hibernation Tests', { timeout: 25000 }, () => {
         Stream.take(3),
         Stream.map((c) => c.maybeNumber.pipe(Option.getOrUndefined)),
       )
-      const chunks = yield* Stream.runCollect(stream)
-      expect(Chunk.toReadonlyArray(chunks)).toEqual([1, 4, 9]) // squares of 1,2,3
+      const result = yield* Stream.runCollect(stream)
+      expect(result).toEqual([1, 4, 9]) // squares of 1,2,3
       console.log('✅ Streaming after hibernation successful')
 
       console.log('🎉 All RPC operations successful after hibernation!')
     }).pipe(Effect.provide(ProtocolLive), withWranglerTest(test)),
   )
 
-  Vitest.scopedLive('should handle rapid operations after hibernation', (test) =>
+  Vitest.live('should handle rapid operations after hibernation', (test) =>
     Effect.gen(function* () {
       console.log('🧪 Testing rapid operations after hibernation...')