fix(session): track delivered chunks to prevent overpayment on WS fallback close

Author: o-az

pnpm-workspace.yaml

@@ -28,3 +28,5 @@ overrides:
   yaml@<2.8.3: '>=2.8.3'
   brace-expansion@<5.0.5: '>=5.0.5'
   lodash@<=4.17.23: '>=4.18.0'
+
+nodeOptions: '--disable-warning=ExperimentalWarning --disable-warning=DeprecationWarning'

src/tempo/client/SessionManager.ts

@@ -111,6 +111,8 @@ export function sessionManager(parameters: sessionManager.Parameters): SessionMa
   let closeReadyWaiter: CloseReadyWaiter | null = null
   let expectedSocketCloseAmount: string | null = null
   let receiptWaiter: ReceiptWaiter | null = null
+  let wsDeliveredChunks = 0n
+  let wsTickCost = 0n
 
   const method = sessionPlugin({
     account: parameters.account,
@@ -195,6 +197,19 @@ export function sessionManager(parameters: sessionManager.Parameters): SessionMa
     }
 
     const cumulative = channel?.channelId === channelId ? channel.cumulativeAmount : 0n
+
+    // For WS sessions, use delivered chunk count × tick cost as a tight spend
+    // estimate.  Without this, a socket death before close-ready would cause
+    // the client to sign for the full cumulative voucher authorization —
+    // potentially orders of magnitude more than what was actually consumed.
+    // The estimate may undercount by at most 1 chunk (if the server committed
+    // a charge but the socket died before delivering the message).
+    if (wsTickCost > 0n) {
+      const deliveryEstimate = wsDeliveredChunks * wsTickCost
+      const bestSpent = spent > deliveryEstimate ? spent : deliveryEstimate
+      return (bestSpent > cumulative ? cumulative : bestSpent).toString()
+    }
+
     return (cumulative > spent ? cumulative : spent).toString()
   }
 
@@ -495,6 +510,8 @@ export function sessionManager(parameters: sessionManager.Parameters): SessionMa
 
       closeReadyReceipt = null
       activeSocketChallenge = challenge
+      wsDeliveredChunks = 0n
+      wsTickCost = BigInt(challenge.request.amount as string)
       const openCredential = PaymentCredential.deserialize<SessionCredentialPayload>(credential)
       activeSocketChannelId = openCredential.payload.channelId
       const rawSocket = new WebSocketImpl(wsUrl, protocols)
@@ -562,6 +579,7 @@ export function sessionManager(parameters: sessionManager.Parameters): SessionMa
           case 'authorization':
             break
           case 'message':
+            wsDeliveredChunks += 1n
             managedSocket.emit('message', { data: message.data, type: 'message' })
             break
           case 'payment-close-ready':

src/tempo/server/Session.test.ts

@@ -3744,6 +3744,108 @@ describe.runIf(isLocalnet)('session', () => {
       }
     })
 
+    test('fallback close after socket death signs for delivered amount, not full voucher', async () => {
+      const backingStore = Store.memory()
+      const routeHandler = Mppx_server.create({
+        methods: [
+          tempo_server.session({
+            store: backingStore,
+            getClient: () => client,
+            account: recipientAccount,
+            currency,
+            escrowContract,
+            chainId: chain.id,
+          }),
+        ],
+        realm: 'api.example.com',
+        secretKey: 'secret',
+      }).session({ amount: '1', decimals: 6, unitType: 'token' })
+
+      const route = (request: Request) => routeHandler(request)
+
+      const httpHandler = NodeRequest.toNodeListener(async (request) => {
+        const result = await route(request)
+        if (result.status === 402) return result.challenge
+        return result.withReceipt(new Response('ok'))
+      })
+
+      const nodeServer = node_http.createServer(httpHandler)
+      const wsServer = new WebSocketServer({ noServer: true })
+      let serverSocket: import('ws').WebSocket | null = null
+
+      await new Promise<void>((resolve) => nodeServer.listen(0, resolve))
+      const { port } = nodeServer.address() as { port: number }
+      const server = Http.wrapServer(nodeServer, { port, url: `http://localhost:${port}` })
+
+      wsServer.on('connection', (socket: import('ws').WebSocket) => {
+        serverSocket = socket
+        void TempoWs.serve({
+          socket,
+          store: backingStore,
+          url: `${server.url}/ws`,
+          route,
+          generate: async function* (stream: TempoWs.SessionController) {
+            await stream.charge()
+            yield 'chunk-1'
+            await stream.charge()
+            yield 'chunk-2'
+            await stream.charge()
+            yield 'chunk-3'
+            await new Promise((resolve) => setTimeout(resolve, 60_000))
+          },
+        })
+      })
+
+      nodeServer.on('upgrade', (req, socket, head) => {
+        if (req.url !== '/ws') {
+          socket.destroy()
+          return
+        }
+
+        wsServer.handleUpgrade(req, socket, head, (websocket: import('ws').WebSocket) => {
+          wsServer.emit('connection', websocket, req)
+        })
+      })
+
+      try {
+        const manager = sessionManager({
+          account: payer,
+          client,
+          escrowContract,
+          fetch: globalThis.fetch,
+          maxDeposit: '3',
+        })
+
+        const ws = await manager.ws(`ws://localhost:${port}/ws`)
+        const chunks: string[] = []
+
+        await new Promise<void>((resolve) => {
+          ws.addEventListener('message', (event) => {
+            if (typeof event.data !== 'string') return
+            chunks.push(event.data)
+            if (chunks.length === 3) serverSocket?.terminate()
+          })
+          ws.addEventListener('close', () => resolve(), { once: true })
+        })
+
+        expect(chunks).toEqual(['chunk-1', 'chunk-2', 'chunk-3'])
+
+        const closeReceipt = await manager.close()
+        expect(closeReceipt).toBeDefined()
+
+        const settledAmount = BigInt(closeReceipt!.spent)
+        const expectedCost = 3n * 1000000n
+        expect(settledAmount).toBeLessThanOrEqual(expectedCost)
+        expect(settledAmount).toBeGreaterThan(0n)
+
+        const fullDeposit = 3000000n
+        expect(settledAmount).toBeLessThan(fullDeposit)
+      } finally {
+        wsServer.close()
+        server.close()
+      }
+    })
+
     test('rejects tx-bearing open receipts replayed during websocket close', async () => {
       const routeHandler = Mppx_server.create({
         methods: [