From 0f0e8618a39caf0d11ad2871c8883841423ee006 Mon Sep 17 00:00:00 2001 From: Georgios Konstantopoulos Date: Sun, 19 Jul 2026 12:47:01 -0700 Subject: [PATCH 1/2] feat: support dynamic websocket session charges --- .../dynamic-websocket-session-charges.md | 5 ++ src/tempo/session/server/MeteredStream.ts | 12 +++-- src/tempo/session/server/Ws.test.ts | 52 +++++++++++++++++++ src/tempo/session/server/Ws.ts | 2 + 4 files changed, 67 insertions(+), 4 deletions(-) create mode 100644 .changeset/dynamic-websocket-session-charges.md diff --git a/.changeset/dynamic-websocket-session-charges.md b/.changeset/dynamic-websocket-session-charges.md new file mode 100644 index 00000000..0509f7b2 --- /dev/null +++ b/.changeset/dynamic-websocket-session-charges.md @@ -0,0 +1,5 @@ +--- +'mppx': patch +--- + +Added explicit per-message amounts to manual WebSocket session metering. diff --git a/src/tempo/session/server/MeteredStream.ts b/src/tempo/session/server/MeteredStream.ts index 232c1f01..3be395e0 100644 --- a/src/tempo/session/server/MeteredStream.ts +++ b/src/tempo/session/server/MeteredStream.ts @@ -12,8 +12,12 @@ 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 + charge(amount?: bigint): Promise } /** Async stream source accepted by paid session transports. */ @@ -49,7 +53,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,14 +62,14 @@ 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 } diff --git a/src/tempo/session/server/Ws.test.ts b/src/tempo/session/server/Ws.test.ts index e98857b0..72d0d97b 100644 --- a/src/tempo/session/server/Ws.test.ts +++ b/src/tempo/session/server/Ws.test.ts @@ -408,6 +408,58 @@ describe('isows', () => { expect(channel?.units).toBe(0) }) + 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() diff --git a/src/tempo/session/server/Ws.ts b/src/tempo/session/server/Ws.ts index 7c0b5022..c33b5f29 100644 --- a/src/tempo/session/server/Ws.ts +++ b/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 | ((stream: SessionController) => AsyncIterable) pollIntervalMs?: number | undefined /** Payment route handler. Receives synthetic `POST` requests with only From 0abb0fa96415206aa5b0a5bcb21008c4d21f9b2a Mon Sep 17 00:00:00 2001 From: Georgios Konstantopoulos Date: Sun, 19 Jul 2026 13:41:01 -0700 Subject: [PATCH 2/2] fix: cancel websocket session generators on close --- .../dynamic-websocket-session-charges.md | 2 +- src/tempo/session/server/MeteredStream.ts | 5 +- src/tempo/session/server/Ws.test.ts | 61 +++++++++++++++++++ 3 files changed, 66 insertions(+), 2 deletions(-) diff --git a/.changeset/dynamic-websocket-session-charges.md b/.changeset/dynamic-websocket-session-charges.md index 0509f7b2..a7a9b3c2 100644 --- a/.changeset/dynamic-websocket-session-charges.md +++ b/.changeset/dynamic-websocket-session-charges.md @@ -2,4 +2,4 @@ 'mppx': patch --- -Added explicit per-message amounts to manual WebSocket session metering. +Added explicit per-message amounts and cancellation signals to manual WebSocket session metering. diff --git a/src/tempo/session/server/MeteredStream.ts b/src/tempo/session/server/MeteredStream.ts index 3be395e0..bbd094c1 100644 --- a/src/tempo/session/server/MeteredStream.ts +++ b/src/tempo/session/server/MeteredStream.ts @@ -18,6 +18,8 @@ export type SessionController = { * used. */ charge(amount?: bigint): Promise + /** Aborted when the client closes or requests the final session receipt. */ + signal: AbortSignal } /** Async stream source accepted by paid session transports. */ @@ -73,8 +75,9 @@ export async function* meterIterable(options: MeteredStreamOptions): AsyncGenera 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 diff --git a/src/tempo/session/server/Ws.test.ts b/src/tempo/session/server/Ws.test.ts index 72d0d97b..b2c295f9 100644 --- a/src/tempo/session/server/Ws.test.ts +++ b/src/tempo/session/server/Ws.test.ts @@ -408,6 +408,67 @@ 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((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()