@livestore/common-cf 0.4.0 → 0.5.0-dev.1

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

@@ -18,16 +18,17 @@
 import { notYetImplemented, omitUndefineds } from '@livestore/utils'
 import {
   Context,
-  constVoid,
   Effect,
+  Function,
   Exit,
   Layer,
   Logger,
-  LogLevel,
-  Mailbox,
+  Queue,
+  References,
   RpcMessage,
   RpcSerialization,
   RpcServer,
+  Schema,
   Scope,
   Stream,
 } from '@livestore/utils/effect'
@@ -39,7 +40,7 @@ import type * as CfTypes from '../cf-types.ts'
  * This is useful for reading WebSocket attachment data (e.g., forwarded headers)
  * inside RPC handlers.
  */
-export class WsContext extends Context.Tag('WsContext')<WsContext, { readonly ws: CfTypes.WebSocket }>() {}
+export class WsContext extends Context.Service<WsContext, { readonly ws: CfTypes.WebSocket }>()('WsContext') {}
 
 /**
  * Configuration options for setting up WebSocket RPC on a Durable Object.
@@ -55,7 +56,7 @@ export interface DurableObjectWebSocketRpcConfig {
   webSocketMode: 'hibernate' | 'accept'
   /**
    * Effect RPC layer that requires `RpcServer.Protocol` (and `WsContext` if used)
-   * and provides the RPC server runtime.
+   * and provides the RPC server services.
    *
    * This is typically created by:
    * ```typescript
@@ -71,6 +72,10 @@ export interface DurableObjectWebSocketRpcConfig {
   rpcLayer: Layer.Layer<never, never, RpcServer.Protocol | WsContext>
   /** Function to get access to incoming requests */
   onMessage?: (msg: RpcMessage.FromClientEncoded, ws: CfTypes.WebSocket) => void
+  /** Called after a terminal RPC response has been accepted by the serializer. */
+  onRequestExit?: (requestId: string | number, ws: CfTypes.WebSocket) => void
+  /** Called after a connection-level RPC defect has been accepted by the serializer. */
+  onProtocolDefect?: (ws: CfTypes.WebSocket) => void
   mainLayer?: Layer.Layer<never>
 }
 
@@ -139,6 +144,8 @@ export const setupDurableObjectWebSocketRpc = ({
   rpcLayer,
   webSocketMode,
   onMessage,
+  onRequestExit,
+  onProtocolDefect,
   mainLayer,
 }: DurableObjectWebSocketRpcConfig) => {
   if (webSocketMode === 'accept') {
@@ -148,7 +155,7 @@ export const setupDurableObjectWebSocketRpc = ({
   const serverCtxMap = new Map<
     CfTypes.WebSocket,
     {
-      scope: Scope.CloseableScope
+      scope: Scope.Closeable
       onMessage: (message: string | ArrayBuffer) => Promise<void>
     }
   >()
@@ -163,33 +170,31 @@ export const setupDurableObjectWebSocketRpc = ({
 
       const scope = yield* Scope.make()
 
-      const incomingQueue = yield* Mailbox.make<Uint8Array | 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 }),
+        ...omitUndefineds({ onMessage, onRequestExit, onProtocolDefect }),
       }).pipe(Layer.provide(RpcSerialization.layerJson))
 
       const ServerLive = rpcLayer.pipe(Layer.provide(ProtocolLive))
 
       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 | 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)
@@ -197,8 +202,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,
     )
@@ -233,9 +244,12 @@ 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 | string>
+  onRequestExit?: (requestId: string | number, ws: CfTypes.WebSocket) => void
+  onProtocolDefect?: (ws: CfTypes.WebSocket) => void
+  /** Queue for receiving incoming messages from the WebSocket */
+  incomingQueue: Queue.Queue<Uint8Array | string>
 }
 
 /**
@@ -253,7 +267,10 @@ export interface WsRpcServerArgs {
  * @internal This is typically used internally by `setupDurableObjectWebSocketRpc`
  */
 export const layerRpcServerWebsocket = (args: WsRpcServerArgs) =>
-  Layer.mergeAll(Layer.scoped(RpcServer.Protocol, makeSocketProtocol(args)), Layer.succeed(WsContext, { ws: args.ws }))
+  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.
@@ -266,35 +283,84 @@ export const layerRpcServerWebsocket = (args: WsRpcServerArgs) =>
  *
  * @internal Used internally by `layerRpcServerWebsocket`
  */
-const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =>
+const makeSocketProtocol = ({
+  incomingQueue,
+  scope,
+  ws,
+  onMessage,
+  onRequestExit,
+  onProtocolDefect,
+}: WsRpcServerArgs) =>
   Effect.gen(function* () {
     const serialization = yield* RpcSerialization.RpcSerialization
-    const disconnects = yield* Mailbox.make<number>()
+    const disconnects = yield* Queue.unbounded<number>()
+    const requestIdsWithSchemas = new Set<string | number>()
+    const interruptedExitEncoded = yield* Schema.encodeUnknownEffect(
+      serialization.codecFor(Schema.Exit(Schema.Void, Schema.Never, Schema.Never)),
+    )(Exit.interrupt()).pipe(Effect.orDie)
 
-    const writeRaw = (msg: Uint8Array | string) => Effect.succeed(ws.send(msg))
+    const writeRaw = (msg: Uint8Array | string) => Effect.sync(() => 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) => {
+    const runLifecycleHook = (name: string, hook: () => void) =>
+      Effect.sync(hook).pipe(
+        Effect.catchCause((cause) => Effect.logError(`WebSocket RPC ${name} hook failed`, { cause })),
+      )
+
+    const writeDefect = (defect: unknown) => {
+      try {
+        const encoded = parser.encode(RpcMessage.ResponseDefectEncoded(defect))
+        const responseEffect = encoded === undefined ? Effect.void : Effect.orDie(writeRaw(encoded))
+        const onAccepted = Effect.sync(() => requestIdsWithSchemas.clear()).pipe(
+          Effect.andThen(
+            onProtocolDefect !== undefined
+              ? runLifecycleHook('onProtocolDefect', () => onProtocolDefect(ws))
+              : Effect.void,
+          ),
+        )
+        return Effect.ensuring(responseEffect, onAccepted)
+      } catch (cause) {
+        return Effect.die(cause)
+      }
+    }
+
+    const writeResponse = (response: unknown, onAccepted?: Effect.Effect<void>) => {
       try {
         const encoded = parser.encode(response)
-        if (encoded === undefined) {
-          return Effect.void
-        }
-        return Effect.orDie(writeRaw(encoded))
+        const responseEffect = encoded === undefined ? Effect.void : Effect.orDie(writeRaw(encoded))
+        return onAccepted === undefined ? responseEffect : Effect.ensuring(responseEffect, onAccepted)
       } catch (cause) {
-        return Effect.orDie(writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!))
+        return writeDefect(cause)
       }
     }
 
+    const writeExit = (requestId: string | number, exit: unknown) =>
+      writeResponse(
+        { _tag: 'Exit', requestId, exit },
+        Effect.sync(() => requestIdsWithSchemas.delete(requestId)).pipe(
+          Effect.andThen(
+            onRequestExit !== undefined
+              ? runLifecycleHook('onRequestExit', () => onRequestExit(requestId, ws))
+              : Effect.void,
+          ),
+        ),
+      )
+
+    const write = (response: RpcMessage.FromServerEncoded) => {
+      if (response._tag === 'Exit') return writeExit(response.requestId, response.exit)
+      if (response._tag === 'Defect') return writeDefect(response.defect)
+      return writeResponse(response)
+    }
+
     const protocol = yield* RpcServer.Protocol.make((writeRequest_) => {
       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>
@@ -304,20 +370,27 @@ const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =
               while: () => i < decoded.length,
               body: () => {
                 const request = decoded[i++]!
-                if (onMessage !== undefined) {
-                  onMessage(request, ws)
+                if (onMessage !== undefined) onMessage(request, ws)
+
+                if (request._tag === 'Request') {
+                  requestIdsWithSchemas.add(request.id)
+                } else if (request._tag === 'Interrupt' && requestIdsWithSchemas.has(request.requestId) === false) {
+                  // Effect creates this Exit for an unknown request, but its encoder drops it when hibernation has
+                  // discarded the request schema.
+                  return writeExit(request.requestId, interruptedExitEncoded)
                 }
+
                 return writeRequest(id, request)
               },
-              step: constVoid,
+              step: Function.constVoid,
             })
           } catch (cause) {
-            return Effect.orDie(writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!))
+            return writeDefect(cause)
           }
         }),
         Stream.runDrain,
         Effect.tapCauseLogPretty,
-        Effect.fork,
+        Effect.forkIn(scope),
       )
 
       // Start the message processing
@@ -333,6 +406,8 @@ const makeSocketProtocol = ({ incomingQueue, ws, onMessage }: WsRpcServerArgs) =
         supportsAck: true,
         supportsTransferables: false,
         supportsSpanPropagation: true,
+        supportsNotifications: true,
+        codecFor: serialization.codecFor,
       }))
     })
 

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

@@ -1,15 +1,13 @@
 import { expect } from 'vitest'
 
 import { Vitest } from '@livestore/utils-dev/node-vitest'
-import { WranglerDevServerService } from '@livestore/utils-dev/wrangler'
+import { WranglerDevServer } from '@livestore/utils-dev/wrangler'
 import {
-  Chunk,
   Effect,
   FetchHttpClient,
   Layer,
-  Logger,
-  LogLevel,
   Option,
+  References,
   RpcClient,
   RpcSerialization,
   Socket,
@@ -24,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' })
@@ -54,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' })
@@ -62,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 })
@@ -70,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...')
 
@@ -222,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...')