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

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

@@ -1,36 +1,39 @@
 import {
   BackendIdMismatchError,
-  InvalidPushError,
   ServerAheadError,
   SyncBackend,
-  UnexpectedError,
+  UnknownError,
 } from '@livestore/common'
 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,
-  MAX_PULL_EVENTS_PER_MESSAGE,
-  MAX_WS_MESSAGE_BYTES,
+  type ForwardedHeaders,
   type MakeDurableObjectClassOptions,
   type StoreId,
   WebSocketAttachmentSchema,
 } from '../shared.ts'
 import { DoCtx } from './layer.ts'
-import { splitChunkBySize } from './ws-chunking.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
@@ -45,9 +48,15 @@ export const makePush =
         return SyncMessage.PushAck.make({})
       }
 
-      if (options?.onPush) {
-        yield* Effect.tryAll(() => options.onPush!(pushRequest, { storeId, payload })).pipe(
-          UnexpectedError.mapToUnexpectedError,
+      if (options?.onPush !== undefined) {
+        yield* Effect.tryAll(() =>
+          options.onPush!(pushRequest, {
+            storeId,
+            ...(payload !== undefined ? { payload } : {}),
+            ...(headers !== undefined ? { headers } : {}),
+          }),
+        ).pipe(
+          UnknownError.mapToUnknownError,
         )
       }
 
@@ -87,9 +96,9 @@ export const makePush =
         const connectedClients = ctx.getWebSockets()
 
         // Preparing chunks of responses to make sure we don't exceed the WS message size limit.
-        const responses = Chunk.fromIterable(pushRequest.batch).pipe(
+        const responses = yield* Chunk.fromIterable(pushRequest.batch).pipe(
           splitChunkBySize({
-            maxItems: MAX_PULL_EVENTS_PER_MESSAGE,
+            maxItems: MAX_PUSH_EVENTS_PER_REQUEST,
             maxBytes: MAX_WS_MESSAGE_BYTES,
             encode: (items) =>
               encodePullResponse(
@@ -105,23 +114,25 @@ export const makePush =
                 }),
               ),
           }),
-          Chunk.map((eventsChunk) => {
-            const batchWithMetadata = Chunk.toReadonlyArray(eventsChunk).map((eventEncoded) => ({
-              eventEncoded,
-              metadata: Option.some(SyncMessage.SyncMetadata.make({ createdAt })),
-            }))
-
-            const response = SyncMessage.PullResponse.make({
-              batch: batchWithMetadata,
-              pageInfo: SyncBackend.pageInfoNoMore,
-              backendId,
-            })
-
-            return {
-              response,
-              encoded: Schema.encodeSync(SyncMessage.PullResponse)(response),
-            }
-          }),
+          Effect.map(
+            Chunk.map((eventsChunk) => {
+              const batchWithMetadata = Chunk.toReadonlyArray(eventsChunk).map((eventEncoded) => ({
+                eventEncoded,
+                metadata: Option.some(SyncMessage.SyncMetadata.make({ createdAt })),
+              }))
+
+              const response = SyncMessage.PullResponse.make({
+                batch: batchWithMetadata,
+                pageInfo: SyncBackend.pageInfoNoMore,
+                backendId,
+              })
+
+              return {
+                response,
+                encoded: Schema.encodeSync(SyncMessage.PullResponse)(response),
+              }
+            }),
+          ),
         )
 
         // Dual broadcasting: WebSocket + RPC clients
@@ -130,13 +141,13 @@ export const makePush =
         if (connectedClients.length > 0) {
           for (const { response, encoded } of responses) {
             // Only calling once for now.
-            if (options?.onPullRes) {
-              yield* Effect.tryAll(() => options.onPullRes!(response)).pipe(UnexpectedError.mapToUnexpectedError)
+            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.
@@ -147,7 +158,7 @@ export const makePush =
                   requestId,
                   values: [encoded],
                 }
-                conn.send(JSON.stringify(res))
+                conn.send(jsonStringify(res))
               }
             }
           }
@@ -184,12 +195,16 @@ export const makePush =
     }).pipe(
       Effect.tap(
         Effect.fn(function* (message) {
-          if (options?.onPushRes) {
-            yield* Effect.tryAll(() => options.onPushRes!(message)).pipe(UnexpectedError.mapToUnexpectedError)
+          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,13 +1,19 @@
 import { EventSequenceNumber, State } from '@livestore/common/schema'
 import { Schema } from '@livestore/utils/effect'
+
 import { PERSISTENCE_FORMAT_VERSION } from '../shared.ts'
 
+/**
+ * Main event log table storing all LiveStore events.
+ *
+ * ⚠️  IMPORTANT: Any changes to this schema require bumping PERSISTENCE_FORMAT_VERSION in shared.ts
+ */
 export const eventlogTable = State.SQLite.table({
   // NOTE actual table name is determined at runtime to use proper storeId
   name: `eventlog_${PERSISTENCE_FORMAT_VERSION}_$storeId`,
   columns: {
-    seqNum: State.SQLite.integer({ primaryKey: true, schema: EventSequenceNumber.GlobalEventSequenceNumber }),
-    parentSeqNum: State.SQLite.integer({ schema: EventSequenceNumber.GlobalEventSequenceNumber }),
+    seqNum: State.SQLite.integer({ primaryKey: true, schema: EventSequenceNumber.Global.Schema }),
+    parentSeqNum: State.SQLite.integer({ schema: EventSequenceNumber.Global.Schema }),
     name: State.SQLite.text({}),
     args: State.SQLite.text({ schema: Schema.parseJson(Schema.Any), nullable: true }),
     /** ISO date format. Currently only used for debugging purposes. */
@@ -17,12 +23,16 @@ export const eventlogTable = State.SQLite.table({
   },
 })
 
-/** Will only ever have one row per durable object. */
+/**
+ * Context metadata table - one row per Durable Object.
+ *
+ * ⚠️  IMPORTANT: Any changes to this schema require bumping PERSISTENCE_FORMAT_VERSION in shared.ts
+ */
 export const contextTable = State.SQLite.table({
   name: `context_${PERSISTENCE_FORMAT_VERSION}`,
   columns: {
     storeId: State.SQLite.text({ primaryKey: true }),
-    currentHead: State.SQLite.integer({ schema: EventSequenceNumber.GlobalEventSequenceNumber }),
+    currentHead: State.SQLite.integer({ schema: EventSequenceNumber.Global.Schema }),
     backendId: State.SQLite.text({}),
   },
 })

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

@@ -1,9 +1,10 @@
-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 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 { 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 SyncStorage = {
@@ -12,26 +13,30 @@ export type SyncStorage = {
     {
       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: { _tag: 'd1'; db: CfTypes.D1Database } | { _tag: 'do-sqlite' },
+): 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 +52,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.parseJson())
   const textEncoder = new TextEncoder()
 
   const decreaseLimit = (limit: number) => Math.max(D1_MIN_PAGE_SIZE, Math.floor(limit / 2))
@@ -66,17 +72,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 +98,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<Option.Option<readonly [Chunk.Chunk<EmittedEvent>, State]>, UnknownError> =>
         Effect.gen(function* () {
           const statement =
             state.cursor === undefined
@@ -116,7 +122,7 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
             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)
@@ -143,11 +149,13 @@ export const makeStorage = (ctx: CfTypes.DurableObjectState, env: Env, storeId:
 
       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 +190,136 @@ 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<Option.Option<readonly [Chunk.Chunk<EmittedEvent>, 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) {
+              return Option.none()
+            }
+
+            const decodedRows = Chunk.fromIterable(decodeEventlogRows(rows))
+            const eventsChunk = Chunk.map(decodedRows, ({ createdAt, ...eventEncoded }) => ({
+              eventEncoded,
+              metadata: Option.some(SyncMetadata.make({ createdAt })),
+            }))
+
+            const lastSeqNum = Chunk.unsafeLast(decodedRows).seqNum
+            const nextState: State = { cursor: lastSeqNum }
+
+            return Option.some([eventsChunk, nextState] as const)
+          },
+          catch: (error) => new UnknownError({ cause: error, payload: { dbName, stage: 'select' } }),
+        })
+
+      const stream = Stream.unfoldChunkEffect(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,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,31 +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 = ({ 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 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 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 push = makePush({ payload: undefined, headers, options: doOptions, storeId, ctx, env })
 
         return yield* push(req)
       }),
@@ -35,3 +54,4 @@ const createHttpRpcLayer =
     Layer.provideMerge(RpcServer.layerProtocolHttp({ path: '/http-rpc' })),
     Layer.provideMerge(RpcSerialization.layerJson),
   )
+}