mppx 0.0.1 → 0.1.0

package/src/tempo/server/Sse.test.ts

@@ -0,0 +1,289 @@
+import type { Address, Hex } from 'viem'
+import { describe, expect, test } from 'vitest'
+import type * as ChannelStore from '../stream/ChannelStore.js'
+import { serve, toResponse } from '../stream/Sse.js'
+
+const channelId = '0x0000000000000000000000000000000000000000000000000000000000000001' as Hex
+const challengeId = 'test-challenge-id'
+const tickCost = 1000000n
+
+function memoryStore(): ChannelStore.ChannelStore {
+  const channels = new Map()
+  return {
+    async getChannel(id) {
+      return channels.get(id) ?? null
+    },
+    async updateChannel(id, fn) {
+      const result = fn(channels.get(id) ?? null)
+      if (result) channels.set(id, result)
+      else channels.delete(id)
+      return result
+    },
+  }
+}
+
+function seedChannel(
+  store: ChannelStore.ChannelStore,
+  balance: bigint,
+): Promise<ChannelStore.State | null> {
+  return store.updateChannel(channelId, () => ({
+    channelId,
+    payer: '0x0000000000000000000000000000000000000001' as Address,
+    payee: '0x0000000000000000000000000000000000000002' as Address,
+    token: '0x0000000000000000000000000000000000000003' as Address,
+    authorizedSigner: '0x0000000000000000000000000000000000000004' as Address,
+    chainId: 42431,
+    escrowContract: '0x542831e3E4Ace07559b7C8787395f4Fb99F70787' as Address,
+    deposit: balance,
+    settledOnChain: 0n,
+    highestVoucherAmount: balance,
+    highestVoucher: null,
+    spent: 0n,
+    units: 0,
+    finalized: false,
+    createdAt: new Date().toISOString(),
+  }))
+}
+
+async function readStream(stream: ReadableStream<Uint8Array>): Promise<string> {
+  const reader = stream.getReader()
+  const decoder = new TextDecoder()
+  let result = ''
+  while (true) {
+    const { done, value } = await reader.read()
+    if (done) break
+    result += decoder.decode(value, { stream: true })
+  }
+  return result
+}
+
+describe('Sse.serve', () => {
+  test('emits message events for each yielded value (StreamController)', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 3000000n)
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost,
+        generate: async function* (stream) {
+          await stream.charge()
+          yield 'hello'
+          await stream.charge()
+          yield 'world'
+          await stream.charge()
+          yield 'done'
+        },
+      }),
+    )
+
+    expect(response).toBeInstanceOf(Response)
+    expect(response.headers.get('Content-Type')).toBe('text/event-stream; charset=utf-8')
+
+    const output = await readStream(response.body!)
+
+    expect(output).toContain('event: message\ndata: hello\n\n')
+    expect(output).toContain('event: message\ndata: world\n\n')
+    expect(output).toContain('event: message\ndata: done\n\n')
+    expect(output).toContain('event: payment-receipt\n')
+
+    const channel = await store.getChannel(channelId)
+    expect(channel!.spent).toBe(3000000n)
+    expect(channel!.units).toBe(3)
+  })
+
+  test('emits payment-need-voucher when balance exhausted and resumes after top-up', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 1000000n)
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost,
+        pollIntervalMs: 10,
+        generate: async function* (stream) {
+          await stream.charge()
+          yield 'first'
+          await stream.charge()
+          yield 'second'
+        },
+      }),
+    )
+
+    const reader = response.body!.getReader()
+    const decoder = new TextDecoder()
+    const chunks: string[] = []
+
+    const { value: chunk1 } = await reader.read()
+    chunks.push(decoder.decode(chunk1, { stream: true }))
+    expect(chunks[0]).toContain('event: message\ndata: first\n\n')
+
+    const readNext = reader.read().then(({ value }) => {
+      const text = decoder.decode(value, { stream: true })
+      chunks.push(text)
+      return text
+    })
+
+    await new Promise((r) => setTimeout(r, 30))
+
+    await store.updateChannel(channelId, (current) => {
+      if (!current) return null
+      return { ...current, highestVoucherAmount: current.highestVoucherAmount + 2000000n }
+    })
+
+    const secondChunk = await readNext
+    expect(secondChunk).toContain('event: payment-need-voucher\n')
+
+    const remaining: string[] = []
+    while (true) {
+      const { done, value } = await reader.read()
+      if (done) break
+      remaining.push(decoder.decode(value, { stream: true }))
+    }
+    const all = remaining.join('')
+    expect(all).toContain('event: message\ndata: second\n\n')
+    expect(all).toContain('event: payment-receipt\n')
+  })
+
+  test('respects abort signal', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 10000000n)
+
+    const controller = new AbortController()
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost,
+        signal: controller.signal,
+        generate: async function* (stream) {
+          let i = 0
+          while (true) {
+            await stream.charge()
+            yield `chunk-${i++}`
+            await new Promise((r) => setTimeout(r, 5))
+          }
+        },
+      }),
+    )
+
+    const reader = response.body!.getReader()
+    const decoder = new TextDecoder()
+
+    const { value: first } = await reader.read()
+    expect(decoder.decode(first)).toContain('event: message\ndata: chunk-0\n\n')
+
+    controller.abort()
+
+    while (true) {
+      const { done } = await reader.read()
+      if (done) break
+    }
+  })
+
+  test('emits receipt with correct spent and units', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 2000000n)
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost,
+        generate: async function* (stream) {
+          await stream.charge()
+          yield 'a'
+          await stream.charge()
+          yield 'b'
+        },
+      }),
+    )
+
+    const output = await readStream(response.body!)
+    const receiptRaw = output.split('event: payment-receipt\ndata: ')[1]?.split('\n\n')[0]
+    const receipt = JSON.parse(receiptRaw!)
+
+    expect(receipt.challengeId).toBe('test-challenge-id')
+    expect(receipt.channelId).toBe(channelId)
+    expect(receipt.spent).toBe('2000000')
+    expect(receipt.units).toBe(2)
+  })
+
+  test('handles empty generator', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 1000000n)
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost,
+        generate: async function* () {},
+      }),
+    )
+
+    const output = await readStream(response.body!)
+    expect(output).toContain('event: payment-receipt\n')
+    expect(output).not.toContain('event: message\n')
+
+    const channel = await store.getChannel(channelId)
+    expect(channel!.spent).toBe(0n)
+    expect(channel!.units).toBe(0)
+  })
+
+  test('allows tickCost override', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 500n)
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost: 100n,
+        generate: async function* (stream) {
+          for (let i = 0; i < 5; i++) {
+            await stream.charge()
+            yield `tok-${i}`
+          }
+        },
+      }),
+    )
+
+    const output = await readStream(response.body!)
+    for (let i = 0; i < 5; i++) {
+      expect(output).toContain(`event: message\ndata: tok-${i}\n\n`)
+    }
+
+    const channel = await store.getChannel(channelId)
+    expect(channel!.spent).toBe(500n)
+    expect(channel!.units).toBe(5)
+  })
+
+  test('sets correct SSE response headers', async () => {
+    const store = memoryStore()
+    await seedChannel(store, 1000000n)
+
+    const response = toResponse(
+      serve({
+        store,
+        channelId,
+        challengeId,
+        tickCost,
+        generate: async function* () {},
+      }),
+    )
+
+    expect(response.headers.get('Cache-Control')).toBe('no-cache, no-transform')
+    expect(response.headers.get('Connection')).toBe('keep-alive')
+    expect(response.headers.get('Content-Type')).toBe('text/event-stream; charset=utf-8')
+  })
+})

package/src/tempo/server/index.ts

@@ -0,0 +1,5 @@
+export * as ChannelStore from '../stream/ChannelStore.js'
+export * as Sse from '../stream/Sse.js'
+export { charge } from './Charge.js'
+export { tempo } from './MethodIntents.js'
+export { session, settle } from './Session.js'

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

@@ -0,0 +1,153 @@
+/**
+ * Tempo-specific SSE transport that wraps the base HTTP transport
+ * with metering logic (context capture from credentials, per-token
+ * charging via Sse.serve).
+ *
+ * @internal
+ */
+import * as Transport from '../../../server/Transport.js'
+import type * as ChannelStore from '../../stream/ChannelStore.js'
+import * as Sse_core from '../../stream/Sse.js'
+
+/** SSE transport with Tempo stream controller. */
+export type Sse = Transport.Sse<Sse_core.StreamController>
+
+/**
+ * Creates a Tempo-metered SSE transport.
+ *
+ * Wraps an HTTP transport with:
+ * - Context capture from credentials (channelId, tickCost)
+ * - Per-token charging via Sse.serve for generator/iterable responses
+ * - Auto-detection of upstream SSE responses
+ * - Fallback to standard HTTP receipt handling for plain Response
+ */
+export function sse(options: sse.Options & { store: ChannelStore.ChannelStore }): Sse {
+  const { pollingInterval, poll } = options
+
+  // When `poll` is true, strip `waitForUpdate` so the SSE charge loop
+  // falls back to polling. This is needed for runtimes like Cloudflare Workers
+  // where resolving promises across request contexts is not supported.
+  const store = (() => {
+    if (!poll) return options.store
+    const { waitForUpdate: _, ...store } = options.store
+    return store
+  })()
+
+  const contextMap = new Map<string, Sse_core.fromRequest.Context>()
+
+  const base = Transport.http()
+  return Transport.from<Request, Response, Transport.ReceiptResponseOf<Sse>, Response>({
+    name: 'sse',
+
+    getCredential(request) {
+      const credential = base.getCredential(request)
+      if (credential) {
+        try {
+          const ctx = Sse_core.fromRequest(request)
+          contextMap.set(ctx.challengeId, ctx)
+        } catch {
+          // ignore — non-SSE credentials won't have stream context
+        }
+      }
+      return credential
+    },
+
+    respondChallenge(options) {
+      return base.respondChallenge(options) as Response
+    },
+
+    respondReceipt({ receipt, response, challengeId }) {
+      // Auto-detect upstream SSE responses and parse them into an
+      // AsyncIterable so they flow through the metered pipeline.
+      // This lets proxy consumers simply pass `result.withReceipt(upstreamRes)`
+      // and get per-event charging automatically.
+      const resolved =
+        response instanceof Response && Sse_core.isEventStream(response) && response.body
+          ? Sse_core.iterateData(response, { skip: (d) => d === '[DONE]' })
+          : response
+
+      if (isAsyncGeneratorFunction(resolved) || isAsyncIterable(resolved)) {
+        const ctx = contextMap.get(challengeId)
+        if (!ctx) throw new Error('No SSE context available — credential was not parsed')
+        contextMap.delete(challengeId)
+
+        // Pass async generator functions directly so Sse.serve gives them
+        // a StreamController for manual charge(). Pass raw AsyncIterables
+        // as-is so Sse.serve auto-charges per yielded value.
+        const generate: Sse_core.serve.Options['generate'] = isAsyncGeneratorFunction(resolved)
+          ? (resolved as Sse_core.serve.Options['generate'])
+          : (resolved as AsyncIterable<string>)
+        const stream = Sse_core.serve({
+          store,
+          channelId: ctx.channelId,
+          challengeId,
+          tickCost: ctx.tickCost,
+          pollIntervalMs: pollingInterval,
+          generate,
+        })
+        return Sse_core.toResponse(stream)
+      }
+
+      return base.respondReceipt({ receipt, response: response as Response, challengeId })
+    },
+  })
+}
+
+export declare namespace sse {
+  type Options = {
+    /**
+     * When true, the charge loop uses polling instead of `waitForUpdate()`.
+     *
+     * Required for runtimes like Cloudflare Workers where resolving promises
+     * across request contexts is not supported. Without this flag, a mid-stream
+     * voucher POST (Request B) would resolve a waiter created in the streaming
+     * request context (Request A), causing a Workers error.
+     *
+     * @default false
+     */
+    poll?: boolean | undefined
+    /** Polling interval (in milliseconds). @default 10 */
+    pollingInterval?: number | undefined
+  }
+}
+
+/** Default SSE serve: iterates values and emits `event: message` per value. */
+export function defaultServe(options: {
+  generate: AsyncIterable<string> | ((...args: any[]) => AsyncIterable<string>)
+  challengeId: string
+}): Response {
+  const iterable =
+    typeof options.generate === 'function' ? options.generate(undefined as any) : options.generate
+  const encoder = new TextEncoder()
+  const stream = new ReadableStream<Uint8Array>({
+    async start(controller) {
+      try {
+        for await (const value of iterable) {
+          controller.enqueue(encoder.encode(`event: message\ndata: ${value}\n\n`))
+        }
+      } catch (e) {
+        controller.error(e)
+      } finally {
+        controller.close()
+      }
+    },
+  })
+  return new Response(stream, {
+    headers: {
+      'Content-Type': 'text/event-stream; charset=utf-8',
+      'Cache-Control': 'no-cache, no-transform',
+      Connection: 'keep-alive',
+    },
+  })
+}
+
+function isAsyncGeneratorFunction(
+  value: unknown,
+): value is (...args: unknown[]) => AsyncIterable<string> {
+  if (typeof value !== 'function') return false
+  return value.constructor?.name === 'AsyncGeneratorFunction'
+}
+
+function isAsyncIterable(value: unknown): value is AsyncIterable<string> {
+  return value !== null && typeof value === 'object' && Symbol.asyncIterator in (value as object)
+}