fix(llm-proxy): reserve quota before the provider call and stop cutting off long streams

server/supabase/functions/_shared/quota-reservation.test.ts

@@ -0,0 +1,112 @@
+import type { createClient } from '@supabase/supabase-js'
+import {
+  finalizeLlmQuota,
+  parseQuotaFinalizeResponse,
+  parseQuotaReservationResponse,
+  reserveLlmQuota,
+} from './quota.ts'
+
+function assert(condition: boolean, message: string): asserts condition {
+  if (!condition) throw new Error(message)
+}
+
+function assertThrows(fn: () => unknown, message: string): void {
+  try {
+    fn()
+  } catch {
+    return
+  }
+  throw new Error(message)
+}
+
+type ServiceClient = ReturnType<typeof createClient>
+
+class FakeRpcClient {
+  calls: Array<{ name: string; args: Record<string, unknown> }> = []
+  constructor(private readonly result: unknown, private readonly error: unknown = null) {}
+  rpc(name: string, args: Record<string, unknown>): Promise<{ data: unknown; error: unknown }> {
+    this.calls.push({ name, args })
+    return Promise.resolve({ data: this.result, error: this.error })
+  }
+}
+
+const RESERVED = {
+  allowed: true,
+  reservation_id: '6f1d8a0e-0000-4000-8000-000000000001',
+  status: 'reserved',
+  current: 3,
+  limit: 50,
+  period: 'daily',
+  tier: 'pro',
+  overage_credits: 0,
+  consumed_from: 'base',
+}
+
+Deno.test('parseQuotaReservationResponse accepts reserve RPC payloads', () => {
+  const parsed = parseQuotaReservationResponse(RESERVED, 'LLM')
+  assert(parsed.allowed && parsed.reservationId === RESERVED.reservation_id, 'reservation id')
+  assert(parsed.period === 'daily' && parsed.tier === 'pro' && parsed.consumedFrom === 'base', 'fields')
+
+  const denied = parseQuotaReservationResponse(
+    { ...RESERVED, allowed: false, reservation_id: null, status: 'denied', consumed_from: 'none' },
+    'LLM',
+  )
+  assert(!denied.allowed && denied.reservationId === null && denied.status === 'denied', 'denied')
+})
+
+Deno.test('parseQuotaReservationResponse rejects malformed payloads', () => {
+  assertThrows(() => parseQuotaReservationResponse(null, 'LLM'), 'null')
+  assertThrows(() => parseQuotaReservationResponse([], 'LLM'), 'array')
+  assertThrows(() => parseQuotaReservationResponse({ ...RESERVED, status: 'held' }, 'LLM'), 'status')
+  assertThrows(() => parseQuotaReservationResponse({ ...RESERVED, tier: 'gold' }, 'LLM'), 'tier')
+  assertThrows(() => parseQuotaReservationResponse({ ...RESERVED, period: 'monthly' }, 'LLM'), 'period')
+  assertThrows(() => parseQuotaReservationResponse({ ...RESERVED, current: '3' }, 'LLM'), 'current')
+})
+
+Deno.test('parseQuotaFinalizeResponse accepts completed/released only', () => {
+  assert(parseQuotaFinalizeResponse({ status: 'completed' }, 'LLM') === 'completed', 'completed')
+  assert(parseQuotaFinalizeResponse({ status: 'released' }, 'LLM') === 'released', 'released')
+  assertThrows(() => parseQuotaFinalizeResponse({ status: 'reserved' }, 'LLM'), 'reserved')
+})
+
+Deno.test('reserveLlmQuota calls reserve_llm_quota with feature, limit and period', async () => {
+  const client = new FakeRpcClient(RESERVED)
+  const result = await reserveLlmQuota(
+    'user-1',
+    RESERVED.reservation_id,
+    'llm_opus',
+    50,
+    'daily',
+    client as unknown as ServiceClient,
+  )
+  assert(result.reservationId === RESERVED.reservation_id, 'result')
+  assert(client.calls.length === 1 && client.calls[0].name === 'reserve_llm_quota', 'rpc name')
+  const args = client.calls[0].args
+  assert(
+    args.p_user_id === 'user-1'
+      && args.p_reservation_id === RESERVED.reservation_id
+      && args.p_feature === 'llm_opus'
+      && args.p_base_limit === 50
+      && args.p_period === 'daily',
+    'rpc args',
+  )
+})
+
+Deno.test('reserveLlmQuota throws on RPC error', async () => {
+  const client = new FakeRpcClient(null, { message: 'boom' })
+  let threw = false
+  try {
+    await reserveLlmQuota('user-1', 'r', 'llm_haiku', 250, 'weekly', client as unknown as ServiceClient)
+  } catch {
+    threw = true
+  }
+  assert(threw, 'must throw')
+})
+
+Deno.test('finalizeLlmQuota calls finalize_llm_quota', async () => {
+  const client = new FakeRpcClient({ reservation_id: 'r', status: 'released' })
+  const status = await finalizeLlmQuota('r', false, client as unknown as ServiceClient)
+  assert(status === 'released', 'status')
+  const call = client.calls[0]
+  assert(call.name === 'finalize_llm_quota' && call.args.p_reservation_id === 'r' && call.args.p_succeeded === false, 'rpc')
+})