mppx 0.5.7 → 0.5.9

package/src/tempo/server/internal/transport.ts

@@ -38,6 +38,16 @@ export function sse(options: sse.Options & { store: ChannelStore.ChannelStore })
   return Transport.from<Request, Response, Transport.ReceiptResponseOf<Sse>, Response>({
     name: 'sse',
 
+    captureRequest(request) {
+      return (
+        base.captureRequest?.(request) ?? {
+          headers: new Headers(request.headers),
+          method: request.method,
+          url: new URL(request.url),
+        }
+      )
+    },
+
     getCredential(request) {
       return base.getCredential(request)
     },
@@ -46,11 +56,13 @@ export function sse(options: sse.Options & { store: ChannelStore.ChannelStore })
       return base.respondChallenge(options) as Response
     },
 
-    respondReceipt({ credential, receipt, response, challengeId, input }) {
-      const payload = credential.payload as Partial<SessionCredentialPayload>
+    respondReceipt({ credential, envelope, receipt, response, challengeId, input }) {
+      const verifiedCredential = envelope?.credential ?? credential
+      const verifiedChallengeId = envelope?.challenge.id ?? challengeId
+      const payload = verifiedCredential.payload as Partial<SessionCredentialPayload>
       if (!payload.channelId) throw new Error('No SSE context available')
       const channelId = payload.channelId
-      const tickCost = BigInt(credential.challenge.request.amount as string)
+      const tickCost = BigInt(verifiedCredential.challenge.request.amount as string)
 
       // Auto-detect upstream SSE responses and parse them into an
       // AsyncIterable so they flow through the metered pipeline.
@@ -71,7 +83,7 @@ export function sse(options: sse.Options & { store: ChannelStore.ChannelStore })
         const stream = Sse_core.serve({
           store,
           channelId,
-          challengeId,
+          challengeId: verifiedChallengeId,
           tickCost,
           pollIntervalMs: pollingInterval,
           generate,
@@ -81,11 +93,12 @@ export function sse(options: sse.Options & { store: ChannelStore.ChannelStore })
       }
 
       const baseResponse = base.respondReceipt({
-        credential,
+        credential: verifiedCredential,
+        envelope,
         input,
         receipt,
         response: response as Response,
-        challengeId,
+        challengeId: verifiedChallengeId,
       })
 
       // Non-SSE response (e.g. upstream returned JSON instead of event-stream).

package/src/tempo/session/ChannelStore.test.ts

@@ -37,7 +37,7 @@ function seedChannel(
   return store.updateChannel(channelId, () => makeChannel(overrides))
 }
 
-function stripUpdateMethod(store: Store.Store): Store.Store {
+function stripUpdateMethod(store: Store.Store | Store.AtomicStore): Store.Store {
   return {
     get: store.get.bind(store),
     put: store.put.bind(store),
@@ -209,6 +209,25 @@ describe('channelStore', () => {
       await sleep(10)
       expect(ch1Resolved).toBe(false)
     })
+
+    test('resolves on successful deductFromChannel with atomic store.update', async () => {
+      const cs = ChannelStore.fromStore(Store.memory())
+      await seedChannel(cs, { highestVoucherAmount: 5_000_000n, spent: 0n })
+
+      let resolved = false
+      const waiter = cs.waitForUpdate!(channelId).then(() => {
+        resolved = true
+      })
+
+      await sleep(10)
+      expect(resolved).toBe(false)
+
+      const result = await ChannelStore.deductFromChannel(cs, channelId, 1_000_000n)
+      expect(result.ok).toBe(true)
+
+      await waiter
+      expect(resolved).toBe(true)
+    })
   })
 })
 

package/src/tempo/session/ChannelStore.ts

@@ -65,6 +65,9 @@ export interface State {
  * guarantee that no concurrent mutation occurs between reading `current`
  * and writing the return value.
  *
+ * Callbacks should be synchronous and deterministic. When a `ChannelStore`
+ * is backed by `Store.update()`, adapters may retry them internally.
+ *
  * Backends implement this via their native mechanisms:
  * - **In-memory / JS single-thread**: Synchronous callback execution
  * - **Durable Objects**: Single-threaded execution model
@@ -90,6 +93,18 @@ export type ChannelStore = {
    * When not implemented, callers fall back to polling.
    */
   waitForUpdate?(channelId: Hex): Promise<void>
+
+  /**
+   * Atomic read-modify-write that returns the callback's `result` directly.
+   *
+   * Used by {@link deductFromChannel} to atomically compute the deduction
+   * outcome. When backed by `Store.update()`, this delegates to the store's
+   * native atomic primitive.
+   */
+  updateChannelResult?<result>(
+    channelId: Hex,
+    fn: (current: State | null) => Store.Change<State, result>,
+  ): Promise<result>
 }
 
 export type DeductResult = { ok: true; channel: State } | { ok: false; channel: State }
@@ -106,19 +121,41 @@ export async function deductFromChannel(
   channelId: Hex,
   amount: bigint,
 ): Promise<DeductResult> {
-  let deducted = false
+  if (store.updateChannelResult) {
+    const result = await store.updateChannelResult<DeductResult | null>(
+      channelId,
+      (current): Store.Change<State, DeductResult | null> => {
+        if (!current) return { op: 'noop', result: null }
+        if (current.finalized)
+          return { op: 'noop', result: { ok: false, channel: current } as const }
+        if (current.highestVoucherAmount - current.spent >= amount) {
+          const next = { ...current, spent: current.spent + amount, units: current.units + 1 }
+          return { op: 'set', value: next, result: { ok: true, channel: next } as const }
+        }
+        return { op: 'noop', result: { ok: false, channel: current } as const }
+      },
+    )
+    if (!result) throw new Error('channel not found')
+    return result
+  }
+
+  let result: DeductResult | null = null
   const channel = await store.updateChannel(channelId, (current) => {
-    deducted = false
     if (!current) return null
-    if (current.finalized) return current
+    if (current.finalized) {
+      result = { ok: false, channel: current }
+      return current
+    }
     if (current.highestVoucherAmount - current.spent >= amount) {
-      deducted = true
-      return { ...current, spent: current.spent + amount, units: current.units + 1 }
+      const next = { ...current, spent: current.spent + amount, units: current.units + 1 }
+      result = { ok: true, channel: next }
+      return next
     }
+    result = { ok: false, channel: current }
     return current
   })
   if (!channel) throw new Error('channel not found')
-  return { ok: deducted, channel }
+  return result ?? { ok: false, channel }
 }
 
 /**
@@ -140,10 +177,12 @@ export async function deductFromChannel(
  */
 const storeCache = new WeakMap<Store.Store, ChannelStore>()
 
-export function fromStore(store: Store.Store): ChannelStore {
+export function fromStore(store: Store.Store | Store.AtomicStore): ChannelStore {
   const cached = storeCache.get(store)
   if (cached) return cached
 
+  const atomicUpdate = 'update' in store ? (store as Store.AtomicStore).update : undefined
+
   const waiters = new Map<string, Set<() => void>>()
   const locks = new Map<string, Promise<void>>()
 
@@ -158,6 +197,29 @@ export function fromStore(store: Store.Store): ChannelStore {
     channelId: Hex,
     fn: (current: State | null) => State | null,
   ): Promise<State | null> {
+    return updateResult(channelId, (current) => {
+      const next = fn(current)
+      if (next) return { op: 'set', value: next, result: next }
+      return { op: 'delete', result: null }
+    })
+  }
+
+  async function updateResult<result>(
+    channelId: Hex,
+    fn: (current: State | null) => Store.Change<State, result>,
+  ): Promise<result> {
+    let change: Store.Change<State, result> | undefined
+
+    if (atomicUpdate) {
+      const result = await atomicUpdate(channelId, (current) => {
+        change = fn((current as State | null) ?? null)
+        if (change.op !== 'set') return change
+        return { ...change, value: change.value as never }
+      })
+      if (change?.op !== 'noop') notify(channelId)
+      return result
+    }
+
     while (locks.has(channelId)) await locks.get(channelId)
 
     let release!: () => void
@@ -170,10 +232,11 @@ export function fromStore(store: Store.Store): ChannelStore {
 
     try {
       const current = (await store.get(channelId)) as State | null
-      const next = fn(current)
-      if (next) await store.put(channelId, next as never)
-      else await store.delete(channelId)
-      return next
+      change = fn(current)
+      if (change.op === 'set') await store.put(channelId, change.value as never)
+      if (change.op === 'delete') await store.delete(channelId)
+      if (change.op !== 'noop') notify(channelId)
+      return change.result
     } finally {
       locks.delete(channelId)
       release()
@@ -185,9 +248,7 @@ export function fromStore(store: Store.Store): ChannelStore {
       return (await store.get(channelId)) as State | null
     },
     async updateChannel(channelId, fn) {
-      const result = await update(channelId, fn)
-      notify(channelId)
-      return result
+      return update(channelId, fn)
     },
     waitForUpdate(channelId) {
       return new Promise<void>((resolve) => {
@@ -201,6 +262,8 @@ export function fromStore(store: Store.Store): ChannelStore {
     },
   }
 
+  cs.updateChannelResult = updateResult
+
   storeCache.set(store, cs)
   return cs
 }

package/src/tempo/session/Sse.ts

@@ -80,6 +80,13 @@ export function parseEvent(raw: string): SseEvent | null {
 }
 
 export type SessionController = {
+  /**
+   * Reserve voucher coverage for the next emitted chunk.
+   *
+   * The reservation blocks until sufficient voucher headroom exists, but the
+   * charge is only committed once a chunk is actually emitted. If the stream
+   * ends or aborts before that emission, the reservation is dropped.
+   */
   charge(): Promise<void>
 }
 
@@ -93,11 +100,13 @@ export type SessionController = {
  *   generator controls when charges happen by calling `stream.charge()`.
  *
  * For each emitted value the stream:
- * 1. Deducts `tickCost` from the channel balance atomically (auto or manual).
+ * 1. Reserves `tickCost` from the channel's available voucher headroom
+ *    (auto or manual).
  * 2. If balance is sufficient, emits `event: message` with the value.
  * 3. If balance is exhausted, emits `event: payment-need-voucher`
  *    and polls store until the client tops up the channel.
- * 4. On generator completion, emits a final `event: payment-receipt`.
+ * 4. Commits the reserved charge immediately before the chunk is emitted.
+ * 5. On generator completion, emits a final `event: payment-receipt`.
  *
  * Returns a `ReadableStream<Uint8Array>` suitable for use as an HTTP response body.
  */
@@ -118,15 +127,21 @@ export function serve(options: serve.Options): ReadableStream<Uint8Array> {
     async start(controller) {
       const aborted = () => signal?.aborted ?? false
       const emit = (event: string) => controller.enqueue(encoder.encode(event))
+      let reservedAmount = 0n
+      let reservedUnits = 0
 
       const charge = () =>
-        chargeOrWait({
+        reserveChargeOrWait({
           store,
           channelId,
           amount: tickCost,
+          reservedAmount,
           emit,
           pollIntervalMs,
           signal,
+        }).then(() => {
+          reservedAmount += tickCost
+          reservedUnits += 1
         })
 
       const iterable: AsyncIterable<string> =
@@ -137,7 +152,14 @@ export function serve(options: serve.Options): ReadableStream<Uint8Array> {
           if (aborted()) break
 
           if (typeof generate !== 'function') await charge()
-
+          await commitReservedCharges({
+            store,
+            channelId,
+            amount: reservedAmount,
+            units: reservedUnits,
+          })
+          reservedAmount = 0n
+          reservedUnits = 0
           controller.enqueue(encoder.encode(`event: message\ndata: ${value}\n\n`))
         }
 
@@ -221,45 +243,76 @@ export declare namespace fromRequest {
 }
 
 /**
- * Atomically deduct `amount` from a channel, retrying when balance is
- * insufficient. Uses `store.waitForUpdate()` when available for
+ * Reserve `amount` of voucher headroom for a future emission, retrying when
+ * balance is insufficient. Uses `store.waitForUpdate()` when available for
  * event-driven wakeups, falling back to polling otherwise. Emits
  * `payment-need-voucher` events via `emit` while waiting.
  */
-async function chargeOrWait(options: {
+async function reserveChargeOrWait(options: {
   store: ChannelStore.ChannelStore
   channelId: Hex
   amount: bigint
-  emit: (event: string) => void
+  reservedAmount: bigint
+  emit: (event: string) => void | Promise<void>
   pollIntervalMs: number
   signal?: AbortSignal | undefined
 }): Promise<void> {
-  const { store, channelId, amount, emit, pollIntervalMs, signal } = options
+  const { store, channelId, amount, emit, pollIntervalMs, reservedAmount, signal } = options
+
+  let channel = await store.getChannel(channelId)
+  if (!channel) throw new Error('channel not found')
 
-  let result = await ChannelStore.deductFromChannel(store, channelId, amount)
+  const hasHeadroom = (state: ChannelStore.State) =>
+    state.highestVoucherAmount - state.spent - reservedAmount >= amount
 
-  if (!result.ok) {
-    // Emit a single need-voucher event, then poll/wait until the client
-    // sends an updated voucher. The requiredCumulative is constant here:
-    // `spent` only changes on successful deduction (which exits the loop),
-    // so re-emitting on every poll cycle would just cause redundant
-    // voucher POSTs from the client.
+  if (hasHeadroom(channel)) return
+
+  // Emit a single need-voucher event, then wait until the accepted voucher
+  // headroom covers both already-reserved units and the next requested unit.
+  await Promise.resolve(
     emit(
       formatNeedVoucherEvent({
         channelId,
-        requiredCumulative: (result.channel.spent + amount).toString(),
-        acceptedCumulative: result.channel.highestVoucherAmount.toString(),
-        deposit: result.channel.deposit.toString(),
+        requiredCumulative: (channel.spent + reservedAmount + amount).toString(),
+        acceptedCumulative: channel.highestVoucherAmount.toString(),
+        deposit: channel.deposit.toString(),
       }),
-    )
+    ),
+  )
 
-    while (!result.ok) {
-      await waitForUpdate(store, channelId, pollIntervalMs, signal)
-      result = await ChannelStore.deductFromChannel(store, channelId, amount)
-    }
+  while (!hasHeadroom(channel)) {
+    await waitForUpdate(store, channelId, pollIntervalMs, signal)
+    channel = await store.getChannel(channelId)
+    if (!channel) throw new Error('channel not found')
   }
 }
 
+async function commitReservedCharges(options: {
+  store: ChannelStore.ChannelStore
+  channelId: Hex
+  amount: bigint
+  units: number
+}): Promise<void> {
+  const { store, channelId, amount, units } = options
+  if (amount === 0n || units === 0) return
+
+  let committed = false
+  const channel = await store.updateChannel(channelId, (current) => {
+    if (!current) return null
+    if (current.finalized) return current
+    if (current.highestVoucherAmount - current.spent < amount) return current
+    committed = true
+    return {
+      ...current,
+      spent: current.spent + amount,
+      units: current.units + units,
+    }
+  })
+
+  if (!channel) throw new Error('channel not found')
+  if (!committed) throw new Error('reserved voucher coverage is no longer available')
+}
+
 async function waitForUpdate(
   store: ChannelStore.ChannelStore,
   channelId: Hex,