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

package/src/do-rpc/server.ts

@@ -1,6 +1,7 @@
 import {
   Effect,
   Exit,
+  Fiber,
   Headers,
   type Layer,
   Option,
@@ -19,17 +20,29 @@ import {
 
 import type * as CfTypes from '../cf-types.ts'
 
-/**
- * The erased dynamic RPC schemas (`Rpc.exitSchema(...) as Schema.Top`, `rpc.payloadSchema`) are
- * plain data — they need NO decoding/encoding services. `Schema.Top` types those service channels
- * as `unknown`, which surfaces as `unknown` in the requirements channel of `encode/decodeUnknownEffect`
- * (anyUnknownInErrorContext / TS377030). Assert the true `never` services in ONE place.
- *
- * TODO(effect): revisit if a cleaner/idiomatic API lands — https://github.com/Effect-TS/effect/issues/6489
- */
-const erasedJsonCodec = (schema: Schema.Top) =>
+const isEncodedRpcMessage = Schema.is(RpcMessage.EncodedSchema)
+
+const isRequestEncoded = (message: unknown): message is RpcMessage.RequestEncoded =>
+  isEncodedRpcMessage(message) === true && message._tag === 'Request'
+
+const erasedCodec = (codecFor: RpcSerialization.CodecFor, schema: Schema.Top) =>
   // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- erased RPC data schemas require no codec services
-  Schema.toCodecJson(schema as Schema.Codec<unknown, unknown, never, never>)
+  codecFor(schema as Schema.Codec<unknown, unknown, never, never>)
+
+const schemaBinarySerialization = RpcSerialization.RpcSerialization.pipe(
+  Effect.provide(RpcSerialization.layerSchemaBinary()),
+)
+
+const encodeBytes = (parser: RpcSerialization.Parser, value: unknown): Uint8Array<ArrayBuffer> => {
+  const encoded = parser.encode(value)
+  if (encoded instanceof Uint8Array === false) {
+    throw new TypeError('SchemaBinary RPC serialization did not produce bytes')
+  }
+  if (encoded.buffer instanceof ArrayBuffer === false) {
+    throw new TypeError('SchemaBinary RPC serialization produced shared bytes')
+  }
+  return new Uint8Array(encoded.buffer, encoded.byteOffset, encoded.byteLength)
+}
 
 export interface ClientDoWithRpcCallback {
   __DURABLE_OBJECT_BRAND: never
@@ -58,22 +71,10 @@ export const toDurableObjectHandler =
   ) => Effect.Effect<Uint8Array<ArrayBuffer> | CfTypes.ReadableStream>) =>
   (serializedPayload) =>
     Effect.gen(function* () {
-      const parser = RpcSerialization.msgPack.makeUnsafe()
-
-      // Decode incoming requests - client sends array of requests
-      const decoded = parser.decode(serializedPayload)
-
-      // Handle potential nested array from client serialization
-      let requests: RpcMessage.FromClientEncoded[]
-      if (Array.isArray(decoded) === true && decoded.length === 1 && Array.isArray(decoded[0]) === true) {
-        // Double-wrapped array [[{...}]] -> [{...}]
-        requests = decoded[0]
-      } else if (Array.isArray(decoded) === true) {
-        // Single array [{...}]
-        requests = decoded
-      } else {
-        requests = []
-      }
+      const serialization = yield* schemaBinarySerialization
+      const parser = serialization.makeUnsafe()
+
+      const requests = parser.decode(serializedPayload).filter(isRequestEncoded)
 
       // Get the context with handlers
       const context = yield* Effect.context<Rpc.ToHandler<Rpcs> | Rpc.Middleware<Rpcs>>()
@@ -82,9 +83,6 @@ export const toDurableObjectHandler =
       const responses: any[] = []
 
       for (const request of requests) {
-        if (request._tag !== 'Request') {
-          continue
-        }
         const requestId = RpcMessage.RequestId(request.id)
 
         // Find the RPC handler
@@ -102,19 +100,17 @@ export const toDurableObjectHandler =
           continue
         }
 
-        const payloadResult = yield* Schema.decodeUnknownEffect(erasedJsonCodec(rpc.payloadSchema))(
+        const payloadResult = yield* Schema.decodeUnknownEffect(erasedCodec(serialization.codecFor, rpc.payloadSchema))(
           request.payload,
         ).pipe(Effect.provideContext(entry.context), Effect.result)
 
         if (Result.isFailure(payloadResult) === true) {
-          // Request payloads are encoded with the JSON codec by Effect's RPC client. Decode them
-          // before dispatch so JSON-only representations such as `null` become their schema values.
           // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- Rpc.exitSchema requires AnyWithProps; type narrowing already done above
           const exitSchema = Rpc.exitSchema(rpc as any) as Schema.Top
           const rawExit = Exit.die(SchemaIssue.makeFormatterDefault()(payloadResult.failure.issue))
-          const encodedExit = yield* Schema.encodeUnknownEffect(erasedJsonCodec(exitSchema))(rawExit).pipe(
-            Effect.provideContext(entry.context),
-          )
+          const encodedExit = yield* Schema.encodeUnknownEffect(erasedCodec(serialization.codecFor, exitSchema))(
+            rawExit,
+          ).pipe(Effect.provideContext(entry.context))
           responses.push({
             _tag: 'Exit',
             requestId,
@@ -130,7 +126,15 @@ export const toDurableObjectHandler =
 
         // For streaming RPCs with only one request, return ReadableStream directly
         if (isStream === true && requests.length === 1) {
-          return yield* createStreamingResponse(rpc, entry, requestId, payload, parser, options.layer)
+          return yield* createStreamingResponse(
+            rpc,
+            entry,
+            requestId,
+            payload,
+            parser,
+            serialization.codecFor,
+            options.layer,
+          )
         }
 
         // Execute the handler
@@ -143,7 +147,8 @@ export const toDurableObjectHandler =
             }),
             rpc,
           })
-          const effectOrStream = Rpc.isWrapper(handlerResult) === true ? handlerResult.value : handlerResult
+          const effectOrStream: Effect.Effect<unknown, unknown> | Stream.Stream<unknown, unknown> =
+            Rpc.isWrapper(handlerResult) === true ? handlerResult.value : handlerResult
 
           let value: any
           if (Effect.isEffect(effectOrStream) === true) {
@@ -161,7 +166,7 @@ export const toDurableObjectHandler =
           if (exitSchema !== undefined) {
             // Use schema encoding for proper serialization
             const rawExit = Exit.succeed(value)
-            encodedExit = yield* Schema.encodeUnknownEffect(erasedJsonCodec(exitSchema))(rawExit)
+            encodedExit = yield* Schema.encodeUnknownEffect(erasedCodec(serialization.codecFor, exitSchema))(rawExit)
           } else {
             // Fallback to direct exit
             encodedExit = Exit.succeed(value)
@@ -183,7 +188,9 @@ export const toDurableObjectHandler =
               if (exitSchema !== undefined) {
                 // Use schema encoding for proper serialization
                 const rawExit = Exit.failCause(cause)
-                encodedExit = yield* Schema.encodeUnknownEffect(erasedJsonCodec(exitSchema))(rawExit)
+                encodedExit = yield* Schema.encodeUnknownEffect(erasedCodec(serialization.codecFor, exitSchema))(
+                  rawExit,
+                )
               } else {
                 // Fallback to direct exit
                 encodedExit = Exit.failCause(cause)
@@ -201,9 +208,7 @@ export const toDurableObjectHandler =
         responses.push(result)
       }
 
-      // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- msgPack parser.encode returns unknown; cast to expected wire format
-      const encoded = parser.encode(responses) as Uint8Array<ArrayBuffer>
-      return encoded
+      return encodeBytes(parser, responses)
     }).pipe(Effect.provide(options.layer), Effect.scoped, Effect.orDie) as Effect.Effect<
       Uint8Array<ArrayBuffer> | CfTypes.ReadableStream
     >
@@ -213,6 +218,7 @@ export const emitStreamResponse = Effect.fn('do-rpc/emitStreamResponse')(functio
   callerContext,
   env,
   requestId,
+  schema,
   storeId,
   values,
 }: {
@@ -220,7 +226,8 @@ export const emitStreamResponse = Effect.fn('do-rpc/emitStreamResponse')(functio
   callerContext: { bindingName: string; durableObjectId: string }
   requestId: string
   storeId: string
-  values: ReadonlyArray.NonEmptyReadonlyArray<any>
+  schema: Schema.Top
+  values: ReadonlyArray.NonEmptyReadonlyArray<unknown>
 }) {
   // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- CF worker env bindings are typed as Record<string, any>; narrowing to known DO namespace
   const clientDoNamespace = env[callerContext.bindingName] as
@@ -233,12 +240,15 @@ export const emitStreamResponse = Effect.fn('do-rpc/emitStreamResponse')(functio
 
   const clientDo = clientDoNamespace.get(clientDoNamespace.idFromString(callerContext.durableObjectId))
 
-  const res: RpcMessage.ResponseChunkEncoded = { _tag: 'Chunk', requestId, values }
-  const parser = RpcSerialization.msgPack.makeUnsafe()
+  const serialization = yield* schemaBinarySerialization
+  const parser = serialization.makeUnsafe()
+  const encodedValues = yield* Schema.encodeUnknownEffect(
+    erasedCodec(serialization.codecFor, Schema.NonEmptyArray(schema)),
+  )(values)
+  const res = { _tag: 'Chunk', requestId, values: encodedValues }
   // Native Cloudflare RPC rejects schema values with custom prototypes. Keep the callback
   // boundary clone-safe by sending the already-encoded Effect RPC message as bytes.
-  // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- msgPack parser.encode returns unknown; the encoded result is a byte payload
-  const serializedRes = parser.encode(res) as Uint8Array<ArrayBuffer>
+  const serializedRes = encodeBytes(parser, res)
 
   yield* Effect.tryPromise(() => clientDo.syncUpdateRpc(serializedRes, storeId))
 })
@@ -252,7 +262,8 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
   entry: Rpc.Handler<Rpcs['_tag']>,
   requestId: RpcMessage.RequestId,
   payload: unknown,
-  parser: ReturnType<typeof RpcSerialization.msgPack.makeUnsafe>,
+  parser: RpcSerialization.Parser,
+  codecFor: RpcSerialization.CodecFor,
   layer: Layer.Layer<Rpc.ToHandler<Rpcs> | Rpc.Middleware<Rpcs>, LE>,
 ): Effect.Effect<CfTypes.ReadableStream, never, Scope.Scope> =>
   Effect.gen(function* () {
@@ -265,11 +276,15 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
       }),
       rpc,
     })
-    const effectOrStream = Rpc.isWrapper(handlerResult) === true ? handlerResult.value : handlerResult
+    const effectOrStream: Effect.Effect<unknown, unknown> | Stream.Stream<unknown, unknown> =
+      Rpc.isWrapper(handlerResult) === true ? handlerResult.value : handlerResult
 
-    const stream: Stream.Stream<any, any> =
-      // @effect-diagnostics-next-line anyUnknownInErrorContext:off -- `Rpc.Handler.handler` returns `Effect<any, any>` due to dynamic dispatch; orDie converts the error to a defect handled by the downstream catchCause
-      Effect.isEffect(effectOrStream) === true ? yield* Effect.orDie(effectOrStream) : effectOrStream
+    // @effect-diagnostics-next-line anyUnknownInErrorContext:off -- dynamically dispatched handlers have an unknown error; orDie converts it to a defect handled by the downstream catchCause
+    const result = Effect.isEffect(effectOrStream) === true ? yield* Effect.orDie(effectOrStream) : effectOrStream
+    if (Stream.isStream(result) === false) {
+      return yield* Effect.die(`Streaming RPC ${rpc._tag} did not produce a Stream`)
+    }
+    const stream = result
 
     // Get the stream schemas for proper chunk-level encoding
     // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- Rpc.Handler doesn't expose successSchema publicly; see https://github.com/Effect-TS/effect/issues/6064
@@ -282,10 +297,12 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
         : Option.none()
     const arrayEncoder =
       Option.isSome(streamSchemas) === true
-        ? Schema.encodeUnknownEffect(Schema.toCodecJson(Schema.Array(streamSchemas.value.success)))
-        : Schema.encodeUnknownEffect(Schema.toCodecJson(Schema.Array(Schema.Any)))
+        ? Schema.encodeUnknownEffect(erasedCodec(codecFor, Schema.NonEmptyArray(streamSchemas.value.success)))
+        : Schema.encodeUnknownEffect(erasedCodec(codecFor, Schema.NonEmptyArray(Schema.Any)))
 
     // Convert stream to ReadableStream
+    let producer: Fiber.Fiber<void> | undefined
+    let cancelled = false
     const readableStream = new ReadableStream({
       start(controller) {
         // Run the stream and send chunks + final exit
@@ -304,9 +321,7 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
                 values: encodedValues,
               }
 
-              // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- msgPack parser.encode returns unknown; cast to expected wire format
-              const serialized = parser.encode([chunkMessage]) as Uint8Array<ArrayBuffer>
-              controller.enqueue(serialized)
+              controller.enqueue(encodeBytes(parser, [chunkMessage]))
             }),
           )
 
@@ -314,7 +329,7 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
           const rawExit = Exit.void
           // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- Rpc.exitSchema requires AnyWithProps; type narrowing already done above
           const exitSchema = Rpc.exitSchema(rpc as any) as Schema.Top
-          const encodedExit = yield* Schema.encodeUnknownEffect(erasedJsonCodec(exitSchema))(rawExit)
+          const encodedExit = yield* Schema.encodeUnknownEffect(erasedCodec(codecFor, exitSchema))(rawExit)
 
           const exitMessage = {
             _tag: 'Exit' as const,
@@ -322,18 +337,17 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
             exit: encodedExit,
           }
 
-          // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- msgPack parser.encode returns unknown; cast to expected wire format
-          const exitSerialized = parser.encode([exitMessage]) as Uint8Array<ArrayBuffer>
-          controller.enqueue(exitSerialized)
+          controller.enqueue(encodeBytes(parser, [exitMessage]))
           controller.close()
         }).pipe(
           Effect.catchCause((cause) =>
             Effect.gen(function* () {
+              if (cancelled === true) return
               // Send error exit with proper schema encoding
               const rawExit = Exit.failCause(cause)
               // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- Rpc.exitSchema requires AnyWithProps; type narrowing already done above
               const exitSchema = Rpc.exitSchema(rpc as any) as Schema.Top
-              const encodedExit = yield* Schema.encodeUnknownEffect(erasedJsonCodec(exitSchema))(rawExit)
+              const encodedExit = yield* Schema.encodeUnknownEffect(erasedCodec(codecFor, exitSchema))(rawExit)
 
               const exitMessage = {
                 _tag: 'Exit' as const,
@@ -341,23 +355,26 @@ const createStreamingResponse = <Rpcs extends Rpc.Any, LE>(
                 exit: encodedExit,
               }
 
-              // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- msgPack parser.encode returns unknown; cast to expected wire format
-              const exitSerialized = parser.encode([exitMessage]) as Uint8Array<ArrayBuffer>
-              controller.enqueue(exitSerialized)
+              controller.enqueue(encodeBytes(parser, [exitMessage]))
               controller.close()
             }),
           ),
         )
 
         // Run the stream processing
-        runStream.pipe(
+        producer = runStream.pipe(
           Effect.provide(layer),
           Effect.scoped,
           Effect.tapCauseLogPretty,
           (_) => _ as Effect.Effect<void>,
-          Effect.runPromise,
+          Effect.runFork,
         )
       },
+      cancel() {
+        cancelled = true
+        // @effect-diagnostics-next-line runEffectInsideEffect:off -- native stream callback runs outside Effect; interruption needs no services
+        return producer === undefined ? undefined : Effect.runPromise(Fiber.interrupt(producer))
+      },
       // oxlint-disable-next-line typescript-eslint(no-unsafe-type-assertion) -- bridging standard Web API ReadableStream to Cloudflare Worker ReadableStream type
     }) as any as CfTypes.ReadableStream
 

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

@@ -28,6 +28,7 @@ import {
   RpcMessage,
   RpcSerialization,
   RpcServer,
+  Schema,
   Scope,
   Stream,
 } from '@livestore/utils/effect'
@@ -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') {
@@ -171,7 +178,7 @@ export const setupDurableObjectWebSocketRpc = ({
         ws,
         scope,
         incomingQueue,
-        ...omitUndefineds({ onMessage }),
+        ...omitUndefineds({ onMessage, onRequestExit, onProtocolDefect }),
       }).pipe(Layer.provide(RpcSerialization.layerJson))
 
       const ServerLive = rpcLayer.pipe(Layer.provide(ProtocolLive))
@@ -239,6 +246,8 @@ export interface WsRpcServerArgs {
   ws: CfTypes.WebSocket
   scope: Scope.Scope
   onMessage?: (message: RpcMessage.FromClientEncoded, ws: CfTypes.WebSocket) => void
+  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>
 }
@@ -274,30 +283,79 @@ export const layerRpcServerWebsocket = (args: WsRpcServerArgs) =>
  *
  * @internal Used internally by `layerRpcServerWebsocket`
  */
-const makeSocketProtocol = ({ incomingQueue, scope, ws, onMessage }: WsRpcServerArgs) =>
+const makeSocketProtocol = ({
+  incomingQueue,
+  scope,
+  ws,
+  onMessage,
+  onRequestExit,
+  onProtocolDefect,
+}: WsRpcServerArgs) =>
   Effect.gen(function* () {
     const serialization = yield* RpcSerialization.RpcSerialization
     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.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_
 
@@ -312,15 +370,22 @@ const makeSocketProtocol = ({ incomingQueue, scope, ws, onMessage }: WsRpcServer
               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: Function.constVoid,
             })
           } catch (cause) {
-            return Effect.orDie(writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!))
+            return writeDefect(cause)
           }
         }),
         Stream.runDrain,
@@ -342,6 +407,7 @@ const makeSocketProtocol = ({ incomingQueue, scope, ws, onMessage }: WsRpcServer
         supportsTransferables: false,
         supportsSpanPropagation: true,
         supportsNotifications: true,
+        codecFor: serialization.codecFor,
       }))
     })