diff --git a/.changeset/dynamic-websocket-session-charges.md b/.changeset/dynamic-websocket-session-charges.md new file mode 100644 index 00000000..a7a9b3c2 --- /dev/null +++ b/.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. diff --git a/src/tempo/session/server/MeteredStream.ts b/src/tempo/session/server/MeteredStream.ts index 232c1f01..bbd094c1 100644 --- a/src/tempo/session/server/MeteredStream.ts +++ b/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 + 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. */ @@ -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 diff --git a/src/tempo/session/server/Ws.test.ts b/src/tempo/session/server/Ws.test.ts index e98857b0..b2c295f9 100644 --- a/src/tempo/session/server/Ws.test.ts +++ b/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((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() 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