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

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

@@ -1,37 +1,40 @@
-import { UnexpectedError } from '@livestore/common'
-import type { LiveStoreEvent } from '@livestore/common/schema'
+import { UnknownError } from '@livestore/common'
 import type { CfTypes } from '@livestore/common-cf'
-import { Chunk, Effect, Option, Schema, Stream } from '@livestore/utils/effect'
+import type { LiveStoreEvent } from '@livestore/common/schema'
+import { Effect, Option, ReadonlyArray as EffectArray, Schema, Stream } from '@livestore/utils/effect'
+
 import { SyncMetadata } from '../../common/sync-message-types.ts'
-import { type Env, PERSISTENCE_FORMAT_VERSION, type StoreId } from '../shared.ts'
+import { PERSISTENCE_FORMAT_VERSION, type StoreId } from '../shared.ts'
 import { eventlogTable } from './sqlite.ts'
 
+export type StorageEngine = { _tag: 'd1'; db: CfTypes.D1Database } | { _tag: 'do-sqlite' }
+
 export type SyncStorage = {
   dbName: string
   getEvents: (cursor: number | undefined) => Effect.Effect<
     {
       total: number
       stream: Stream.Stream<
-        { eventEncoded: LiveStoreEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> },
-        UnexpectedError
+        { eventEncoded: LiveStoreEvent.Global.Encoded; metadata: Option.Option<SyncMetadata> },
+        UnknownError
       >
     },
-    UnexpectedError
+    UnknownError
   >
   appendEvents: (
-    batch: ReadonlyArray<LiveStoreEvent.AnyEncodedGlobal>,
+    batch: ReadonlyArray<LiveStoreEvent.Global.Encoded>,
     createdAt: string,
-  ) => Effect.Effect<void, UnexpectedError>
-  resetStore: Effect.Effect<void, UnexpectedError>
+  ) => Effect.Effect<void, UnknownError>
+  resetStore: Effect.Effect<void, UnknownError>
 }
 
-export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId: StoreId): SyncStorage => {
+export const makeStorage = (ctx: CfTypes.DurableObjectState, storeId: StoreId, engine: StorageEngine): SyncStorage => {
   const dbName = `eventlog_${PERSISTENCE_FORMAT_VERSION}_${toValidTableName(storeId)}`
 
   const execDb = <T>(cb: (db: CfTypes.D1Database) => Promise<CfTypes.D1Result<T>>) =>
     Effect.tryPromise({
-      try: () => cb(env.DB),
-      catch: (error) => new UnexpectedError({ cause: error, payload: { dbName } }),
+      try: () => cb(engine._tag === 'd1' ? engine.db : (undefined as never)),
+      catch: (error) => new UnknownError({ cause: error, payload: { dbName } }),
     }).pipe(
       Effect.map((_) => _.results),
       Effect.withSpan('@livestore/sync-cf:durable-object:execDb'),
@@ -47,6 +50,7 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
   const D1_MIN_PAGE_SIZE = 1
 
   const decodeEventlogRows = Schema.decodeUnknownSync(Schema.Array(eventlogTable.rowSchema))
+  const jsonStringify = Schema.encodeSync(Schema.fromJsonString(Schema.Unknown))
   const textEncoder = new TextEncoder()
 
   const decreaseLimit = (limit: number) => Math.max(D1_MIN_PAGE_SIZE, Math.floor(limit / 2))
@@ -66,17 +70,17 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
     return limit
   }
 
-  const getEvents = (
+  const getEventsD1 = (
     cursor: number | undefined,
   ): Effect.Effect<
     {
       total: number
       stream: Stream.Stream<
-        { eventEncoded: LiveStoreEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> },
-        UnexpectedError
+        { eventEncoded: LiveStoreEvent.Global.Encoded; metadata: Option.Option<SyncMetadata> },
+        UnknownError
       >
     },
-    UnexpectedError
+    UnknownError
   > =>
     Effect.gen(function* () {
       const countStatement =
@@ -92,13 +96,13 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
       const total = Number(countRows[0]?.total ?? 0)
 
       type State = { cursor: number | undefined; limit: number }
-      type EmittedEvent = { eventEncoded: LiveStoreEvent.AnyEncodedGlobal; metadata: Option.Option<SyncMetadata> }
+      type EmittedEvent = { eventEncoded: LiveStoreEvent.Global.Encoded; metadata: Option.Option<SyncMetadata> }
 
       const initialState: State = { cursor, limit: D1_INITIAL_PAGE_SIZE }
 
       const fetchPage = (
         state: State,
-      ): Effect.Effect<Option.Option<readonly [Chunk.Chunk<EmittedEvent>, State]>, UnexpectedError> =>
+      ): Effect.Effect<readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>], UnknownError> =>
         Effect.gen(function* () {
           const statement =
             state.cursor === undefined
@@ -113,10 +117,11 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
           })
 
           if (rawEvents.length === 0) {
-            return Option.none()
+            const done: readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>] = [[], Option.none()]
+            return done
           }
 
-          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)
@@ -126,28 +131,36 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
             }
           }
 
-          const decodedRows = Chunk.fromIterable(decodeEventlogRows(rawEvents))
+          const decodedRows = decodeEventlogRows(rawEvents)
+
+          if (EffectArray.isReadonlyArrayNonEmpty(decodedRows) === false) {
+            const done: readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>] = [[], Option.none()]
+            return done
+          }
 
-          const eventsChunk = Chunk.map(decodedRows, ({ createdAt, ...eventEncoded }) => ({
+          const events = decodedRows.map(({ createdAt, ...eventEncoded }) => ({
             eventEncoded,
             metadata: Option.some(SyncMetadata.make({ createdAt })),
           }))
 
-          const lastSeqNum = Chunk.unsafeLast(decodedRows).seqNum
+          const lastSeqNum = EffectArray.lastNonEmpty(decodedRows).seqNum
           const nextState: State = { cursor: lastSeqNum, limit: computeNextLimit(state.limit, encodedSize) }
 
-          return Option.some([eventsChunk, nextState] as const)
+          const page: readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>] = [events, Option.some(nextState)]
+          return page
         })
 
-      const stream = Stream.unfoldChunkEffect(initialState, fetchPage)
+      const stream = Stream.paginate(initialState, fetchPage)
 
       return { total, stream }
     }).pipe(
-      UnexpectedError.mapToUnexpectedError,
-      Effect.withSpan('@livestore/sync-cf:durable-object:getEvents', { attributes: { dbName, cursor } }),
+      UnknownError.mapToUnknownError,
+      Effect.withSpan('@livestore/sync-cf:durable-object:getEvents', {
+        attributes: { dbName, cursor, engine: engine._tag },
+      }),
     )
 
-  const appendEvents: SyncStorage['appendEvents'] = (batch, createdAt) =>
+  const appendEventsD1: SyncStorage['appendEvents'] = (batch, createdAt) =>
     Effect.gen(function* () {
       // If there are no events, do nothing.
       if (batch.length === 0) return
@@ -182,24 +195,144 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
         )
       }
     }).pipe(
-      UnexpectedError.mapToUnexpectedError,
+      UnknownError.mapToUnknownError,
       Effect.withSpan('@livestore/sync-cf:durable-object:appendEvents', {
-        attributes: { dbName, batchLength: batch.length },
+        attributes: { dbName, batchLength: batch.length, engine: engine._tag },
       }),
     )
 
   const resetStore = Effect.promise(() => ctx.storage.deleteAll()).pipe(
-    UnexpectedError.mapToUnexpectedError,
+    UnknownError.mapToUnknownError,
     Effect.withSpan('@livestore/sync-cf:durable-object:resetStore'),
   )
 
-  return {
-    dbName,
-    // getHead,
-    getEvents,
-    appendEvents,
-    resetStore,
+  // DO SQLite engine implementation
+  const getEventsDoSqlite = (
+    cursor: number | undefined,
+  ): Effect.Effect<
+    {
+      total: number
+      stream: Stream.Stream<
+        { eventEncoded: LiveStoreEvent.Global.Encoded; metadata: Option.Option<SyncMetadata> },
+        UnknownError
+      >
+    },
+    UnknownError
+  > =>
+    Effect.gen(function* () {
+      const selectCountSql =
+        cursor === undefined
+          ? `SELECT COUNT(*) as total FROM "${dbName}"`
+          : `SELECT COUNT(*) as total FROM "${dbName}" WHERE seqNum > ?`
+
+      const total = yield* Effect.try({
+        try: () => {
+          const cursorIter =
+            cursor === undefined ? ctx.storage.sql.exec(selectCountSql) : ctx.storage.sql.exec(selectCountSql, cursor)
+          let computed = 0
+          for (const row of cursorIter) {
+            computed = Number((row as any).total ?? 0)
+          }
+          return computed
+        },
+        catch: (error) => new UnknownError({ cause: error, payload: { dbName, stage: 'count' } }),
+      })
+
+      type State = { cursor: number | undefined }
+      type EmittedEvent = { eventEncoded: LiveStoreEvent.Global.Encoded; metadata: Option.Option<SyncMetadata> }
+
+      const DO_PAGE_SIZE = 256
+      const initialState: State = { cursor }
+
+      const fetchPage = (
+        state: State,
+      ): Effect.Effect<readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>], UnknownError> =>
+        Effect.try({
+          try: () => {
+            const sql =
+              state.cursor === undefined
+                ? `SELECT * FROM "${dbName}" ORDER BY seqNum ASC LIMIT ?`
+                : `SELECT * FROM "${dbName}" WHERE seqNum > ? ORDER BY seqNum ASC LIMIT ?`
+
+            const iter =
+              state.cursor === undefined
+                ? ctx.storage.sql.exec(sql, DO_PAGE_SIZE)
+                : ctx.storage.sql.exec(sql, state.cursor, DO_PAGE_SIZE)
+
+            const rows: any[] = []
+            for (const row of iter) rows.push(row)
+
+            if (rows.length === 0) {
+              const done: readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>] = [[], Option.none()]
+              return done
+            }
+
+            const decodedRows = decodeEventlogRows(rows)
+
+            if (EffectArray.isReadonlyArrayNonEmpty(decodedRows) === false) {
+              const done: readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>] = [[], Option.none()]
+              return done
+            }
+
+            const events = decodedRows.map(({ createdAt, ...eventEncoded }) => ({
+              eventEncoded,
+              metadata: Option.some(SyncMetadata.make({ createdAt })),
+            }))
+
+            const lastSeqNum = EffectArray.lastNonEmpty(decodedRows).seqNum
+            const nextState: State = { cursor: lastSeqNum }
+
+            const page: readonly [ReadonlyArray<EmittedEvent>, Option.Option<State>] = [events, Option.some(nextState)]
+            return page
+          },
+          catch: (error) => new UnknownError({ cause: error, payload: { dbName, stage: 'select' } }),
+        })
+
+      const stream = Stream.paginate(initialState, fetchPage)
+
+      return { total, stream }
+    }).pipe(
+      UnknownError.mapToUnknownError,
+      Effect.withSpan('@livestore/sync-cf:durable-object:getEvents', {
+        attributes: { dbName, cursor, engine: engine._tag },
+      }),
+    )
+
+  const appendEventsDoSqlite: SyncStorage['appendEvents'] = (batch, createdAt) =>
+    Effect.try({
+      try: () => {
+        if (batch.length === 0) return
+        // Keep params per statement within conservative limits (align with D1 bound params ~100)
+        const CHUNK_SIZE = 14
+        for (let i = 0; i < batch.length; i += CHUNK_SIZE) {
+          const chunk = batch.slice(i, i + CHUNK_SIZE)
+          const placeholders = chunk.map(() => '(?, ?, ?, ?, ?, ?, ?)').join(', ')
+          const sql = `INSERT INTO "${dbName}" (seqNum, parentSeqNum, args, name, createdAt, clientId, sessionId) VALUES ${placeholders}`
+          const params = chunk.flatMap((event) => [
+            event.seqNum,
+            event.parentSeqNum,
+            event.args === undefined ? null : JSON.stringify(event.args),
+            event.name,
+            createdAt,
+            event.clientId,
+            event.sessionId,
+          ])
+          ctx.storage.sql.exec(sql, ...params)
+        }
+      },
+      catch: (error) => new UnknownError({ cause: error, payload: { dbName, stage: 'insert' } }),
+    }).pipe(
+      Effect.withSpan('@livestore/sync-cf:durable-object:appendEvents', {
+        attributes: { dbName, batchLength: batch.length, engine: engine._tag },
+      }),
+      UnknownError.mapToUnknownError,
+    )
+
+  if (engine._tag === 'd1') {
+    return { dbName, getEvents: getEventsD1, appendEvents: appendEventsD1, resetStore }
   }
+
+  return { dbName, getEvents: getEventsDoSqlite, appendEvents: appendEventsDoSqlite, resetStore }
 }
 
 const toValidTableName = (str: string) => str.replaceAll(/[^a-zA-Z0-9]/g, '_')

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

@@ -1,84 +1,84 @@
-import { InvalidPullError, InvalidPushError } from '@livestore/common'
+import { UnknownError } from '@livestore/common'
 import { type CfTypes, toDurableObjectHandler } from '@livestore/common-cf'
-import {
-  Effect,
-  Headers,
-  HttpServer,
-  Layer,
-  Logger,
-  LogLevel,
-  Option,
-  RpcSerialization,
-  Stream,
-} from '@livestore/utils/effect'
+import { Effect, Headers, Option, 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'
+import { rpcSubscriptionKeyPrefix, type RpcSubscription } from '../../shared.ts'
+import * as DoCtx from '../layer.ts'
 import { makeEndingPullStream } from '../pull.ts'
 import { makePush } from '../push.ts'
 
 export interface DoRpcHandlerOptions {
   payload: Uint8Array<ArrayBuffer>
-  input: Omit<DoCtxInput, 'from'>
+  input: Omit<DoCtx.DoCtxInput, 'from'>
 }
 
 export const createDoRpcHandler = (
   options: DoRpcHandlerOptions,
 ): Effect.Effect<Uint8Array<ArrayBuffer> | CfTypes.ReadableStream> =>
-  Effect.gen(this, function* () {
+  Effect.gen({ self: this }, function* () {
     const { payload, input } = options
-    // const { rpcSubscriptions, backendId, doOptions, ctx, env } = yield* DoCtx
 
     // TODO add admin RPCs
     const RpcLive = SyncDoRpc.toLayer({
-      'SyncDoRpc.Ping': (_req) => {
-        return Effect.succeed(SyncMessage.Pong.make({}))
-      },
+      'SyncDoRpc.Ping': () => Effect.void,
       'SyncDoRpc.Pull': (req, { headers }) =>
-        Effect.gen(this, function* () {
-          const { rpcSubscriptions } = yield* DoCtx
+        Effect.gen({ self: this }, function* () {
+          const { ctx } = yield* DoCtx.DoCtx
 
           // TODO rename `req.rpcContext` to something more appropriate
-          if (req.rpcContext) {
-            rpcSubscriptions.set(req.storeId, {
+          if (req.rpcContext !== undefined) {
+            const subscription: RpcSubscription = {
               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 } : {}),
+            }
+            // Keyed by client DO id (not storeId — one backend serves many client DOs). The DO's synchronous KV
+            // storage is the single source of truth, so the entry outlives backend reconstruction.
+            ctx.storage.kv.put(`${rpcSubscriptionKeyPrefix}${subscription.callerContext.durableObjectId}`, subscription)
           }
 
-          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) => ({
             ...res,
             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.tapErrorCause(Effect.log),
+          Stream.provide(DoCtx.layer({ ...input, from: { storeId: req.storeId } })),
+          Stream.mapError((cause) =>
+            cause._tag === 'UnknownError' || cause._tag === 'BackendIdMismatchError'
+              ? cause
+              : new UnknownError({ cause }),
+          ),
+          Stream.tapCause(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 })
+        Effect.gen({ self: this }, function* () {
+          const { doOptions, ctx, env, storeId } = yield* DoCtx.DoCtx
+          // 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.provide(DoCtx.layer({ ...input, from: { storeId: req.storeId } })),
+          Effect.mapError((cause) =>
+            cause._tag === 'UnknownError' ||
+            cause._tag === 'ServerAheadError' ||
+            cause._tag === 'BackendIdMismatchError'
+              ? cause
+              : new UnknownError({ cause }),
+          ),
           Effect.tapCauseLogPretty,
         ),
     })
 
     const handler = toDurableObjectHandler(SyncDoRpc, {
-      layer: Layer.mergeAll(RpcLive, RpcSerialization.layerJson, HttpServer.layerContext).pipe(
-        Layer.provide(Logger.consoleWithThread('SyncDo')),
-        Layer.provide(Logger.minimumLogLevel(LogLevel.Debug)),
-      ),
+      layer: RpcLive,
     })
 
-    return yield* handler(payload)
+    return yield* handler(payload).pipe(Effect.annotateLogs({ thread: 'SyncDo' }))
   }).pipe(Effect.withSpan('createDoRpcHandler'))

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

@@ -1,37 +1,54 @@
 import type { CfTypes } from '@livestore/common-cf'
-import { Effect, HttpApp, Layer, RpcSerialization, RpcServer } from '@livestore/utils/effect'
+import { Effect, HttpEffect, 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 { DoCtx } from '../layer.ts'
+import { headersRecordToMap } from '../../shared.ts'
+import * as DoCtx from '../layer.ts'
 import { makeEndingPullStream } from '../pull.ts'
 import { makePush } from '../push.ts'
 
-export const createHttpRpcHandler = ({ request }: { request: CfTypes.Request }) =>
-  Effect.gen(function* () {
-    const handlerLayer = createHttpRpcLayer
-    const httpApp = RpcServer.toHttpApp(SyncHttpRpc).pipe(Effect.provide(handlerLayer))
-    const webHandler = yield* httpApp.pipe(Effect.map(HttpApp.toWebHandler))
+export const createHttpRpcHandler = Effect.fn('createHttpRpcHandler')(function* ({
+  request,
+  responseHeaders,
+  forwardedHeaders,
+}: {
+  request: CfTypes.Request
+  responseHeaders?: Record<string, string>
+  forwardedHeaders?: Record<string, string>
+}) {
+  const handlerLayer = createHttpRpcLayer(forwardedHeaders)
+  const httpEffect = yield* RpcServer.toHttpEffect(SyncHttpRpc).pipe(Effect.provide(handlerLayer))
+  const webHandler = HttpEffect.toWebHandler(httpEffect)
+
+  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 yield* Effect.promise(
-      () => webHandler(request as TODO as Request) as TODO as Promise<CfTypes.Response>,
-    ).pipe(Effect.timeout(10000))
-  }).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 { ctx, env, doOptions, storeId } = yield* DoCtx.DoCtx
+        const push = makePush({ payload: req.payload, headers, options: doOptions, storeId, ctx, env })
 
         return yield* push(req)
       }),
 
     'SyncHttpRpc.Ping': () => Effect.succeed(SyncMessage.Pong.make({})),
-  }).pipe(
-    Layer.provideMerge(RpcServer.layerProtocolHttp({ path: '/http-rpc' })),
-    Layer.provideMerge(RpcSerialization.layerJson),
-  )
+  }).pipe(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, Result, RpcServer, Schema, Stream } from '@livestore/utils/effect'
+
 import { SyncWsRpc } from '../../../common/ws-rpc-schema.ts'
-import { DoCtx, type DoCtxInput } from '../layer.ts'
+import { headersRecordToMap, WebSocketAttachmentSchema } from '../../shared.ts'
+import * as DoCtx 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
+export const makeRpcServer = ({ doSelf, doOptions }: Omit<DoCtx.DoCtxInput, 'from'>) => {
   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.provide(DoCtx.layer({ 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 { doOptions, storeId, ctx, env } = yield* DoCtx.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.provide(DoCtx.layer({ doSelf, doOptions, from: { storeId: req.storeId } })),
+        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.decodeUnknownResult(WebSocketAttachmentSchema)(attachment)
+  if (Result.isFailure(decoded) === true) {
+    yield* Effect.logError('Failed to decode WebSocket attachment for forwarded headers', { error: decoded.failure })
+    ws.close(1011, 'invalid-attachment')
+    return yield* Effect.die('Invalid WebSocket attachment (headers decode failed)')
+  }
+
+  const headers = headersRecordToMap(decoded.success.headers)
+  return headers
+})