feat: support dynamic WebSocket session charges

.changeset/dynamic-websocket-session-charges.md

@@ -0,0 +1,5 @@
+---
+'mppx': patch
+---
+
+Added explicit per-message amounts and cancellation signals to manual WebSocket session metering.

src/tempo/session/server/MeteredStream.ts

@@ -12,8 +12,14 @@ export type SessionController = {
    * 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.
+   *
+   * Pass an explicit raw-unit `amount` for request-aware or otherwise dynamic
+   * pricing. When omitted, the session challenge's configured tick cost is
+   * used.
    */
-  charge(): Promise<void>
+  charge(amount?: bigint): Promise<void>
+  /** Aborted when the client closes or requests the final session receipt. */
+  signal: AbortSignal
 }
 
 /** Async stream source accepted by paid session transports. */
@@ -49,7 +55,7 @@ export async function* meterIterable(options: MeteredStreamOptions): AsyncGenera
   let reservedAmount = 0n
   let reservedUnits = 0
 
-  const charge = async () => {
+  const charge = async (amount = options.tickCost) => {
     if (prepaidUnits > 0) {
       prepaidUnits -= 1
       return
@@ -58,19 +64,20 @@ export async function* meterIterable(options: MeteredStreamOptions): AsyncGenera
     await reserveChargeOrWait({
       store: options.store,
       channelId: options.channelId,
-      amount: options.tickCost,
+      amount,
       reservedAmount,
       emit: options.emitNeedVoucher,
       formatNeedVoucher: options.formatNeedVoucher,
       pollIntervalMs: options.pollIntervalMs,
       signal: options.signal,
     })
-    reservedAmount += options.tickCost
+    reservedAmount += amount
     reservedUnits += 1
   }
 
+  const signal = options.signal ?? new AbortController().signal
   const iterable =
-    typeof options.generate === 'function' ? options.generate({ charge }) : options.generate
+    typeof options.generate === 'function' ? options.generate({ charge, signal }) : options.generate
 
   for await (const value of iterable) {
     if (options.signal?.aborted) break

src/tempo/session/server/Ws.test.ts

@@ -408,6 +408,119 @@ describe('isows', () => {
     expect(channel?.units).toBe(0)
   })
 
+  test('aborts a blocked application generator before sending close-ready', async () => {
+    const socket = new MockSocket()
+    const store = memoryChannelStore()
+    await seedChannel(store, 1n)
+    let generatorAborted = false
+
+    await Ws.serve({
+      socket,
+      store,
+      url: 'ws://example.test/stream',
+      route: async () => ({
+        status: 200,
+        withReceipt(response = new Response(null, { status: 204 })) {
+          response.headers.set(
+            'Payment-Receipt',
+            serializeSessionReceipt(
+              createSessionReceipt({
+                challengeId: challenge.id,
+                channelId,
+                acceptedCumulative: 1n,
+                spent: 0n,
+                units: 0,
+              }),
+            ),
+          )
+          return response
+        },
+      }),
+      generate: async function* (stream) {
+        await new Promise<void>((resolve) => {
+          stream.signal.addEventListener('abort', () => resolve(), { once: true })
+        })
+        generatorAborted = true
+        yield* []
+      },
+    })
+
+    socket.receive(
+      Ws.formatAuthorizationMessage(
+        makeCredential({
+          action: 'open',
+          channelId,
+          cumulativeAmount: '1',
+          signature: `0x${'77'.repeat(65)}`,
+          transaction: '0x01',
+          type: 'transaction',
+        }),
+      ),
+    )
+    await sleep(10)
+    socket.receive(Ws.formatCloseRequestMessage())
+    await sleep(10)
+
+    expect(generatorAborted).toBe(true)
+    expect(
+      socket.sent
+        .map((message) => Ws.parseMessage(message))
+        .some((message) => message?.mpp === 'payment-close-ready'),
+    ).toBe(true)
+  })
+
+  test('meters an explicit per-message amount independently of the challenge tick cost', async () => {
+    const socket = new MockSocket()
+    const store = memoryChannelStore()
+    await seedChannel(store, 25n)
+
+    await Ws.serve({
+      socket,
+      store,
+      url: 'ws://example.test/stream',
+      route: async () => ({
+        status: 200,
+        withReceipt(response = new Response(null, { status: 204 })) {
+          response.headers.set(
+            'Payment-Receipt',
+            serializeSessionReceipt(
+              createSessionReceipt({
+                challengeId: challenge.id,
+                channelId,
+                acceptedCumulative: 25n,
+                spent: 0n,
+                units: 0,
+              }),
+            ),
+          )
+          return response
+        },
+      }),
+      generate: async function* (stream) {
+        await stream.charge(7n)
+        yield 'priced-response'
+      },
+    })
+
+    socket.receive(
+      Ws.formatAuthorizationMessage(
+        makeCredential({
+          action: 'open',
+          channelId,
+          cumulativeAmount: '25',
+          signature: `0x${'77'.repeat(65)}`,
+          transaction: '0x01',
+          type: 'transaction',
+        }),
+      ),
+    )
+
+    await sleep(10)
+
+    expect(await store.getChannel(channelId)).toMatchObject({ spent: 7n, units: 1 })
+    expect(socket.sent.some((message) => message.includes('priced-response'))).toBe(true)
+  })
+
   test('does not meter or emit application messages after close is requested on-chain', async () => {
     const socket = new MockSocket()
     const store = memoryChannelStore()

src/tempo/session/server/Ws.ts

@@ -322,6 +322,8 @@ export declare namespace serve {
      *  multiple offers — otherwise a client can select the cheapest offer
      *  and still receive the same stream. */
     amount?: string | undefined
+    /** Application stream. A manual stream can call `charge(amount)` with a
+     *  per-message raw-unit amount; omitting it uses the challenge tick cost. */
     generate: AsyncIterable<string> | ((stream: SessionController) => AsyncIterable<string>)
     pollIntervalMs?: number | undefined
     /** Payment route handler. Receives synthetic `POST` requests with only