@livestore/sync-cf 0.4.0-dev.21 → 0.4.0-dev.23

package/src/cf-worker/do/durable-object.ts

@@ -1,6 +1,7 @@
 /// <reference types="@cloudflare/workers-types" />
 
 import { DurableObject } from 'cloudflare:workers'
+
 import { type CfTypes, setupDurableObjectWebSocketRpc } from '@livestore/common-cf'
 import { CfDeclare } from '@livestore/common-cf/declare'
 import {
@@ -14,8 +15,10 @@ import {
   Schema,
   type Scope,
 } from '@livestore/utils/effect'
+
 import {
   type Env,
+  extractForwardedHeaders,
   type MakeDurableObjectClassOptions,
   matchSyncRequest,
   type SyncBackendRpcInterface,
@@ -33,7 +36,7 @@ declare class Response extends CfDeclare.Response {}
 declare class WebSocketPair extends CfDeclare.WebSocketPair {}
 declare class WebSocketRequestResponsePair extends CfDeclare.WebSocketRequestResponsePair {}
 
-const DurableObjectBase = DurableObject<Env> as any as new (
+const DurableObjectBase = DurableObject as any as new (
   state: CfTypes.DurableObjectState,
   env: Env,
 ) => CfTypes.DurableObject & { ctx: CfTypes.DurableObjectState; env: Env }
@@ -87,14 +90,14 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
 
   const Logging = Logger.consoleWithThread('SyncDo')
 
-  const Observability = options?.otel?.baseUrl
-    ? Otlp.layer({
+  const Observability: Layer.Layer<never> = options?.otel?.baseUrl !== undefined
+    ? (Otlp.layer({
         baseUrl: options.otel.baseUrl,
         tracerExportInterval: 50,
         resource: {
           serviceName: options.otel.serviceName ?? 'sync-cf-do',
         },
-      }).pipe(Layer.provide(FetchHttpClient.layer))
+      }).pipe(Layer.provide(FetchHttpClient.layer)) as Layer.Layer<never>)
     : Layer.empty
 
   return class SyncBackendDOBase extends DurableObjectBase implements SyncBackendRpcInterface {
@@ -106,7 +109,7 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
       const WebSocketRpcServerLive = makeRpcServer({ doSelf: this, doOptions: options })
 
       // This registers the `webSocketMessage` and `webSocketClose` handlers
-      if (enabledTransports.has('ws')) {
+      if (enabledTransports.has('ws') === true) {
         setupDurableObjectWebSocketRpc({
           doSelf: this,
           rpcLayer: WebSocketRpcServerLive,
@@ -142,7 +145,7 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
       }
     }
 
-    fetch = async (request: Request): Promise<Response> =>
+    override fetch = async (request: Request): Promise<Response> =>
       Effect.gen(this, function* () {
         const searchParams = matchSyncRequest(request)
         if (searchParams === undefined) {
@@ -155,16 +158,20 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
           throw new Error(`Transport ${transport} is not enabled (based on \`options.enabledTransports\`)`)
         }
 
+        // Extract headers to forward based on configuration (available for all transports)
+        const headers = extractForwardedHeaders(request, options?.forwardHeaders)
+
         if (transport === 'http') {
-          return yield* this.handleHttp(request)
+          return yield* this.handleHttp(request, headers)
         }
 
         if (transport === 'ws') {
           const { 0: client, 1: server } = new WebSocketPair()
 
           // Since we're using websocket hibernation, we need to remember the storeId for subsequent `webSocketMessage` calls
+          // Also store forwarded headers so they're available after hibernation resume
           server.serializeAttachment(
-            Schema.encodeSync(WebSocketAttachmentSchema)({ storeId, payload, pullRequestIds: [] }),
+            Schema.encodeSync(WebSocketAttachmentSchema)({ storeId, payload, pullRequestIds: [], headers }),
           )
 
           // See https://developers.cloudflare.com/durable-objects/examples/websocket-hibernation-server
@@ -220,10 +227,11 @@ export const makeDurableObject: MakeDurableObjectClass = (options) => {
      *
      * Requires the `enable_request_signal` compatibility flag to properly support `pull` streaming responses
      */
-    private handleHttp = (request: CfTypes.Request) =>
+    private handleHttp = (request: CfTypes.Request, forwardedHeaders: Record<string, string> | undefined) =>
       createHttpRpcHandler({
         request,
-        responseHeaders: options?.http?.responseHeaders,
+        ...(options?.http?.responseHeaders !== undefined ? { responseHeaders: options.http.responseHeaders } : {}),
+        ...(forwardedHeaders !== undefined ? { forwardedHeaders } : {}),
       }).pipe(Effect.withSpan('@livestore/sync-cf:durable-object:handleHttp'))
 
     private runEffectAsPromise = <T, E = never>(effect: Effect.Effect<T, E, Scope.Scope>): Promise<T> =>

package/src/cf-worker/do/layer.ts

@@ -1,9 +1,10 @@
 import { UnknownError } from '@livestore/common'
-import { EventSequenceNumber, State } from '@livestore/common/schema'
 import type { CfTypes } from '@livestore/common-cf'
+import { EventSequenceNumber, State } from '@livestore/common/schema'
 import { shouldNeverHappen } from '@livestore/utils'
 import { Effect, Predicate } from '@livestore/utils/effect'
 import { nanoid } from '@livestore/utils/nanoid'
+
 import type { Env, MakeDurableObjectClassOptions, RpcSubscription } from '../shared.ts'
 import { contextTable, eventlogTable } from './sqlite.ts'
 import { makeStorage } from './sync-storage.ts'
@@ -27,7 +28,7 @@ export class DoCtx extends Effect.Service<DoCtx>()('DoCtx', {
       }
 
       const getStoreId = (from: CfTypes.Request | { storeId: string }) => {
-        if (Predicate.hasProperty(from, 'url')) {
+        if (Predicate.hasProperty(from, 'url') === true) {
           const url = new URL(from.url)
           return (
             url.searchParams.get('storeId') ?? shouldNeverHappen(`No storeId provided in request URL search params`)
@@ -42,7 +43,7 @@ export class DoCtx extends Effect.Service<DoCtx>()('DoCtx', {
         const opt = doOptions?.storage
         if (opt?._tag === 'd1') {
           const db = (doSelf.env as any)[opt.binding]
-          if (!db) {
+          if (db == null) {
             return yield* UnknownError.make({ cause: new Error(`D1 binding '${opt.binding}' not found on env`) })
           }
           return { _tag: 'd1' as const, db }

package/src/cf-worker/do/pull.ts

@@ -1,8 +1,10 @@
-import { BackendIdMismatchError, InvalidPullError, SyncBackend, UnknownError } from '@livestore/common'
+import { BackendIdMismatchError, SyncBackend, UnknownError } from '@livestore/common'
 import { splitChunkBySize } from '@livestore/common/sync'
 import { Chunk, Effect, Option, Schema, Stream } from '@livestore/utils/effect'
+
 import { MAX_PULL_EVENTS_PER_MESSAGE, MAX_WS_MESSAGE_BYTES } from '../../common/constants.ts'
 import { SyncMessage } from '../../common/mod.ts'
+import type { ForwardedHeaders } from '../shared.ts'
 import { DoCtx } from './layer.ts'
 
 const encodePullResponse = Schema.encodeSync(SyncMessage.PullResponse)
@@ -16,15 +18,28 @@ const encodePullResponse = Schema.encodeSync(SyncMessage.PullResponse)
 // DO RPC:
 // - Further chunks will be emitted manually in `push.ts`
 // - If the client sends a `Interrupt` RPC message, TODO
-export const makeEndingPullStream = (
-  req: SyncMessage.PullRequest,
-  payload: Schema.JsonValue | undefined,
-): Stream.Stream<SyncMessage.PullResponse, InvalidPullError, DoCtx> =>
+export const makeEndingPullStream = ({
+  req,
+  payload,
+  headers,
+}: {
+  req: SyncMessage.PullRequest
+  payload: Schema.JsonValue | undefined
+  headers: ForwardedHeaders | undefined
+}): Stream.Stream<SyncMessage.PullResponse, UnknownError | BackendIdMismatchError, DoCtx> =>
   Effect.gen(function* () {
     const { doOptions, backendId, storeId, storage } = yield* DoCtx
 
-    if (doOptions?.onPull) {
-      yield* Effect.tryAll(() => doOptions!.onPull!(req, { storeId, payload })).pipe(UnknownError.mapToUnknownError)
+    if (doOptions?.onPull !== undefined) {
+      yield* Effect.tryAll(() =>
+        doOptions.onPull!(req, {
+          storeId,
+          ...(payload !== undefined ? { payload } : {}),
+          ...(headers !== undefined ? { headers } : {}),
+        }),
+      ).pipe(
+        UnknownError.mapToUnknownError,
+      )
     }
 
     if (req.cursor._tag === 'Some' && req.cursor.value.backendId !== backendId) {
@@ -61,7 +76,7 @@ export const makeEndingPullStream = (
       }),
       Stream.tap(
         Effect.fn(function* (res) {
-          if (doOptions?.onPullRes) {
+          if (doOptions?.onPullRes !== undefined) {
             yield* Effect.tryAll(() => doOptions.onPullRes!(res)).pipe(UnknownError.mapToUnknownError)
           }
         }),
@@ -70,6 +85,10 @@ export const makeEndingPullStream = (
     )
   }).pipe(
     Stream.unwrap,
-    Stream.mapError((cause) => InvalidPullError.make({ cause })),
+    Stream.mapError((cause) =>
+      cause._tag === 'BackendIdMismatchError' || cause._tag === 'UnknownError'
+        ? cause
+        : new UnknownError({ cause }),
+    ),
     Stream.withSpan('cloudflare-provider:pull'),
   )

package/src/cf-worker/do/push.ts

@@ -1,30 +1,39 @@
 import {
   BackendIdMismatchError,
-  InvalidPushError,
   ServerAheadError,
   SyncBackend,
   UnknownError,
 } from '@livestore/common'
-import { splitChunkBySize } from '@livestore/common/sync'
 import { type CfTypes, emitStreamResponse } from '@livestore/common-cf'
+import { splitChunkBySize } from '@livestore/common/sync'
 import { Chunk, Effect, Option, type RpcMessage, Schema } from '@livestore/utils/effect'
+
 import { MAX_PUSH_EVENTS_PER_REQUEST, MAX_WS_MESSAGE_BYTES } from '../../common/constants.ts'
 import { SyncMessage } from '../../common/mod.ts'
-import { type Env, type MakeDurableObjectClassOptions, type StoreId, WebSocketAttachmentSchema } from '../shared.ts'
+import {
+  type Env,
+  type ForwardedHeaders,
+  type MakeDurableObjectClassOptions,
+  type StoreId,
+  WebSocketAttachmentSchema,
+} from '../shared.ts'
 import { DoCtx } from './layer.ts'
 
 const encodePullResponse = Schema.encodeSync(SyncMessage.PullResponse)
+const jsonStringify = Schema.encodeSync(Schema.parseJson())
 type PullBatchItem = SyncMessage.PullResponse['batch'][number]
 
 export const makePush =
   ({
     payload,
+    headers,
     options,
     storeId,
     ctx,
     env,
   }: {
     payload: Schema.JsonValue | undefined
+    headers: ForwardedHeaders | undefined
     options: MakeDurableObjectClassOptions | undefined
     storeId: StoreId
     ctx: CfTypes.DurableObjectState
@@ -39,8 +48,14 @@ export const makePush =
         return SyncMessage.PushAck.make({})
       }
 
-      if (options?.onPush) {
-        yield* Effect.tryAll(() => options.onPush!(pushRequest, { storeId, payload })).pipe(
+      if (options?.onPush !== undefined) {
+        yield* Effect.tryAll(() =>
+          options.onPush!(pushRequest, {
+            storeId,
+            ...(payload !== undefined ? { payload } : {}),
+            ...(headers !== undefined ? { headers } : {}),
+          }),
+        ).pipe(
           UnknownError.mapToUnknownError,
         )
       }
@@ -126,13 +141,13 @@ export const makePush =
         if (connectedClients.length > 0) {
           for (const { response, encoded } of responses) {
             // Only calling once for now.
-            if (options?.onPullRes) {
+            if (options?.onPullRes !== undefined) {
               yield* Effect.tryAll(() => options.onPullRes!(response)).pipe(UnknownError.mapToUnknownError)
             }
 
             // NOTE we're also sending the pullRes chunk to the pushing ws client as confirmation
             for (const conn of connectedClients) {
-              const attachment = Schema.decodeSync(WebSocketAttachmentSchema)(conn.deserializeAttachment())
+              const attachment = yield* Schema.decode(WebSocketAttachmentSchema)(conn.deserializeAttachment())
 
               // We're doing something a bit "advanced" here as we're directly emitting Effect RPC-compatible
               // response messsages on the Effect RPC-managed websocket connection to the WS client.
@@ -143,7 +158,7 @@ export const makePush =
                   requestId,
                   values: [encoded],
                 }
-                conn.send(JSON.stringify(res))
+                conn.send(jsonStringify(res))
               }
             }
           }
@@ -180,12 +195,16 @@ export const makePush =
     }).pipe(
       Effect.tap(
         Effect.fn(function* (message) {
-          if (options?.onPushRes) {
+          if (options?.onPushRes !== undefined) {
             yield* Effect.tryAll(() => options.onPushRes!(message)).pipe(UnknownError.mapToUnknownError)
           }
         }),
       ),
-      Effect.mapError((cause) => InvalidPushError.make({ cause })),
+      Effect.mapError((cause) =>
+        cause._tag === 'BackendIdMismatchError' || cause._tag === 'ServerAheadError' || cause._tag === 'UnknownError'
+          ? cause
+          : new UnknownError({ cause }),
+      ),
       Effect.withSpan('sync-cf:do:push', { attributes: { storeId, batchSize: pushRequest.batch.length } }),
     )
 

package/src/cf-worker/do/sqlite.ts

@@ -1,5 +1,6 @@
 import { EventSequenceNumber, State } from '@livestore/common/schema'
 import { Schema } from '@livestore/utils/effect'
+
 import { PERSISTENCE_FORMAT_VERSION } from '../shared.ts'
 
 /**

package/src/cf-worker/do/sync-storage.ts

@@ -1,7 +1,8 @@
 import { UnknownError } from '@livestore/common'
-import type { LiveStoreEvent } from '@livestore/common/schema'
 import type { CfTypes } from '@livestore/common-cf'
+import type { LiveStoreEvent } from '@livestore/common/schema'
 import { Chunk, Effect, Option, Schema, Stream } from '@livestore/utils/effect'
+
 import { SyncMetadata } from '../../common/sync-message-types.ts'
 import { PERSISTENCE_FORMAT_VERSION, type StoreId } from '../shared.ts'
 import { eventlogTable } from './sqlite.ts'
@@ -51,6 +52,7 @@ export const makeStorage = (
   const D1_MIN_PAGE_SIZE = 1
 
   const decodeEventlogRows = Schema.decodeUnknownSync(Schema.Array(eventlogTable.rowSchema))
+  const jsonStringify = Schema.encodeSync(Schema.parseJson())
   const textEncoder = new TextEncoder()
 
   const decreaseLimit = (limit: number) => Math.max(D1_MIN_PAGE_SIZE, Math.floor(limit / 2))
@@ -120,7 +122,7 @@ export const makeStorage = (
             return Option.none()
           }
 
-          const encodedSize = textEncoder.encode(JSON.stringify(rawEvents)).byteLength
+          const encodedSize = textEncoder.encode(jsonStringify(rawEvents)).byteLength
 
           if (encodedSize > D1_TARGET_RESPONSE_BYTES && state.limit > D1_MIN_PAGE_SIZE) {
             const nextLimit = decreaseLimit(state.limit)

package/src/cf-worker/do/transport/do-rpc-server.ts

@@ -1,4 +1,4 @@
-import { InvalidPullError, InvalidPushError } from '@livestore/common'
+import { UnknownError } from '@livestore/common'
 import { type CfTypes, toDurableObjectHandler } from '@livestore/common-cf'
 import {
   Effect,
@@ -11,6 +11,7 @@ import {
   RpcSerialization,
   Stream,
 } from '@livestore/utils/effect'
+
 import { SyncDoRpc } from '../../../common/do-rpc-schema.ts'
 import { SyncMessage } from '../../../common/mod.ts'
 import { DoCtx, type DoCtxInput } from '../layer.ts'
@@ -39,17 +40,18 @@ export const createDoRpcHandler = (
           const { rpcSubscriptions } = yield* DoCtx
 
           // TODO rename `req.rpcContext` to something more appropriate
-          if (req.rpcContext) {
+          if (req.rpcContext !== undefined) {
             rpcSubscriptions.set(req.storeId, {
               storeId: req.storeId,
-              payload: req.payload,
               subscribedAt: Date.now(),
               requestId: Headers.get(headers, 'x-rpc-request-id').pipe(Option.getOrThrow),
               callerContext: req.rpcContext.callerContext,
+              ...(req.payload !== undefined ? { payload: req.payload } : {}),
             })
           }
 
-          return makeEndingPullStream(req, req.payload)
+          // DO-RPC doesn't have HTTP headers context - headers are undefined
+          return makeEndingPullStream({ req, payload: req.payload, headers: undefined })
         }).pipe(
           Stream.unwrap,
           Stream.map((res) => ({
@@ -57,18 +59,27 @@ export const createDoRpcHandler = (
             rpcRequestId: Headers.get(headers, 'x-rpc-request-id').pipe(Option.getOrThrow),
           })),
           Stream.provideLayer(DoCtx.Default({ ...input, from: { storeId: req.storeId } })),
-          Stream.mapError((cause) => (cause._tag === 'InvalidPullError' ? cause : InvalidPullError.make({ cause }))),
+          Stream.mapError((cause) =>
+            cause._tag === 'UnknownError' || cause._tag === 'BackendIdMismatchError'
+              ? cause
+              : new UnknownError({ cause }),
+          ),
           Stream.tapErrorCause(Effect.log),
         ),
       'SyncDoRpc.Push': (req) =>
         Effect.gen(this, function* () {
           const { doOptions, ctx, env, storeId } = yield* DoCtx
-          const push = makePush({ storeId, payload: req.payload, options: doOptions, ctx, env })
+          // DO-RPC doesn't have HTTP headers context - headers are undefined
+          const push = makePush({ storeId, payload: req.payload, headers: undefined, options: doOptions, ctx, env })
 
           return yield* push(req)
         }).pipe(
           Effect.provide(DoCtx.Default({ ...input, from: { storeId: req.storeId } })),
-          Effect.mapError((cause) => (cause._tag === 'InvalidPushError' ? cause : InvalidPushError.make({ cause }))),
+          Effect.mapError((cause) =>
+            cause._tag === 'UnknownError' || cause._tag === 'ServerAheadError' || cause._tag === 'BackendIdMismatchError'
+              ? cause
+              : new UnknownError({ cause }),
+          ),
           Effect.tapCauseLogPretty,
         ),
     })

package/src/cf-worker/do/transport/http-rpc-server.ts

@@ -1,45 +1,50 @@
 import type { CfTypes } from '@livestore/common-cf'
 import { Effect, HttpApp, Layer, RpcSerialization, RpcServer } from '@livestore/utils/effect'
+
 import { SyncHttpRpc } from '../../../common/http-rpc-schema.ts'
 import * as SyncMessage from '../../../common/sync-message-types.ts'
+import { headersRecordToMap } from '../../shared.ts'
 import { DoCtx } from '../layer.ts'
 import { makeEndingPullStream } from '../pull.ts'
 import { makePush } from '../push.ts'
 
-export const createHttpRpcHandler = ({
+export const createHttpRpcHandler = Effect.fn('createHttpRpcHandler')(function* ({
   request,
   responseHeaders,
+  forwardedHeaders,
 }: {
   request: CfTypes.Request
   responseHeaders?: Record<string, string>
-}) =>
-  Effect.gen(function* () {
-    const handlerLayer = createHttpRpcLayer
-    const httpApp = RpcServer.toHttpApp(SyncHttpRpc).pipe(Effect.provide(handlerLayer))
-    const webHandler = yield* httpApp.pipe(Effect.map(HttpApp.toWebHandler))
-
-    const response = yield* Effect.promise(
-      () => webHandler(request as TODO as Request) as TODO as Promise<CfTypes.Response>,
-    ).pipe(Effect.timeout(10000))
-
-    if (responseHeaders !== undefined) {
-      for (const [key, value] of Object.entries(responseHeaders)) {
-        response.headers.set(key, value)
-      }
+  forwardedHeaders?: Record<string, string>
+}) {
+  const handlerLayer = createHttpRpcLayer(forwardedHeaders)
+  const httpApp = RpcServer.toHttpApp(SyncHttpRpc).pipe(Effect.provide(handlerLayer))
+  const webHandler = yield* httpApp.pipe(Effect.map(HttpApp.toWebHandler))
+
+  const response = yield* Effect.promise(
+    () => webHandler(request as TODO as Request) as TODO as Promise<CfTypes.Response>,
+  ).pipe(Effect.timeout(10000))
+
+  if (responseHeaders !== undefined) {
+    for (const [key, value] of Object.entries(responseHeaders)) {
+      response.headers.set(key, value)
     }
+  }
+
+  return response
+})
 
-    return response
-  }).pipe(Effect.withSpan('createHttpRpcHandler'))
+const createHttpRpcLayer = (forwardedHeaders: Record<string, string> | undefined) => {
+  const headers = headersRecordToMap(forwardedHeaders)
 
-const createHttpRpcLayer =
   // TODO implement admin requests
-  SyncHttpRpc.toLayer({
-    'SyncHttpRpc.Pull': (req) => makeEndingPullStream(req, req.payload),
+  return SyncHttpRpc.toLayer({
+    'SyncHttpRpc.Pull': (req) => makeEndingPullStream({ req, payload: req.payload, headers }),
 
     'SyncHttpRpc.Push': (req) =>
       Effect.gen(function* () {
         const { ctx, env, doOptions, storeId } = yield* DoCtx
-        const push = makePush({ payload: undefined, options: doOptions, storeId, ctx, env })
+        const push = makePush({ payload: undefined, headers, options: doOptions, storeId, ctx, env })
 
         return yield* push(req)
       }),
@@ -49,3 +54,4 @@ const createHttpRpcLayer =
     Layer.provideMerge(RpcServer.layerProtocolHttp({ path: '/http-rpc' })),
     Layer.provideMerge(RpcSerialization.layerJson),
   )
+}

package/src/cf-worker/do/transport/ws-rpc-server.ts

@@ -1,34 +1,62 @@
-import { InvalidPullError, InvalidPushError } from '@livestore/common'
-import { Effect, identity, Layer, RpcServer, Stream } from '@livestore/utils/effect'
+import { UnknownError } from '@livestore/common'
+import { WsContext } from '@livestore/common-cf'
+import { Effect, identity, Layer, RpcServer, Schema, Stream } from '@livestore/utils/effect'
+
 import { SyncWsRpc } from '../../../common/ws-rpc-schema.ts'
+import { headersRecordToMap, WebSocketAttachmentSchema } from '../../shared.ts'
 import { DoCtx, type DoCtxInput } from '../layer.ts'
 import { makeEndingPullStream } from '../pull.ts'
 import { makePush } from '../push.ts'
 
 export const makeRpcServer = ({ doSelf, doOptions }: Omit<DoCtxInput, 'from'>) => {
-  // TODO implement admin requests
   const handlersLayer = SyncWsRpc.toLayer({
     'SyncWsRpc.Pull': (req) =>
-      makeEndingPullStream(req, req.payload).pipe(
-        // Needed to keep the stream alive on the client side for phase 2 (i.e. not send the `Exit` stream RPC message)
-        req.live ? Stream.concat(Stream.never) : identity,
-        Stream.provideLayer(DoCtx.Default({ doSelf, doOptions, from: { storeId: req.storeId } })),
-        Stream.mapError((cause) => (cause._tag === 'InvalidPullError' ? cause : InvalidPullError.make({ cause }))),
-        // Stream.tapErrorCause(Effect.log),
-      ),
+      Effect.gen(function* () {
+        const headers = yield* getForwardedHeaders
+        return makeEndingPullStream({ req, payload: req.payload, headers }).pipe(
+          // Needed to keep the stream alive on the client side for phase 2 (i.e. not send the `Exit` stream RPC message)
+          req.live === true ? Stream.concat(Stream.never) : identity,
+          Stream.provideLayer(DoCtx.Default({ doSelf, doOptions, from: { storeId: req.storeId } })),
+          Stream.mapError((cause) =>
+            cause._tag === 'UnknownError' || cause._tag === 'BackendIdMismatchError'
+              ? cause
+              : new UnknownError({ cause }),
+          ),
+        )
+      }).pipe(Stream.unwrap),
     'SyncWsRpc.Push': (req) =>
       Effect.gen(function* () {
         const { doOptions, storeId, ctx, env } = yield* DoCtx
+        const headers = yield* getForwardedHeaders
 
-        const push = makePush({ options: doOptions, storeId, payload: req.payload, ctx, env })
+        const push = makePush({ options: doOptions, storeId, payload: req.payload, headers, ctx, env })
 
         return yield* push(req)
       }).pipe(
         Effect.provide(DoCtx.Default({ doSelf, doOptions, from: { storeId: req.storeId } })),
-        Effect.mapError((cause) => (cause._tag === 'InvalidPushError' ? cause : InvalidPushError.make({ cause }))),
+        Effect.mapError((cause) =>
+          cause._tag === 'UnknownError' || cause._tag === 'ServerAheadError' || cause._tag === 'BackendIdMismatchError'
+            ? cause
+            : new UnknownError({ cause }),
+        ),
         Effect.tapCauseLogPretty,
       ),
   })
 
   return RpcServer.layer(SyncWsRpc).pipe(Layer.provide(handlersLayer))
 }
+
+/** Extracts forwarded headers from the WebSocket attachment */
+const getForwardedHeaders = Effect.gen(function* () {
+  const { ws } = yield* WsContext
+  const attachment = ws.deserializeAttachment()
+  const decoded = Schema.decodeUnknownEither(WebSocketAttachmentSchema)(attachment)
+  if (decoded._tag === 'Left') {
+    yield* Effect.logError('Failed to decode WebSocket attachment for forwarded headers', { error: decoded.left })
+    ws.close(1011, 'invalid-attachment')
+    return yield* Effect.die('Invalid WebSocket attachment (headers decode failed)')
+  }
+
+  const headers = headersRecordToMap(decoded.right.headers)
+  return headers
+})