diff --git a/server/supabase/functions/_shared/quota-reservation.test.ts b/server/supabase/functions/_shared/quota-reservation.test.ts new file mode 100644 index 0000000..5b9e666 --- /dev/null +++ b/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 + +class FakeRpcClient { + calls: Array<{ name: string; args: Record }> = [] + constructor(private readonly result: unknown, private readonly error: unknown = null) {} + rpc(name: string, args: Record): 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') +}) diff --git a/server/supabase/functions/_shared/quota.ts b/server/supabase/functions/_shared/quota.ts index 54ed7db..53d2177 100644 --- a/server/supabase/functions/_shared/quota.ts +++ b/server/supabase/functions/_shared/quota.ts @@ -60,7 +60,8 @@ export interface QuotaConsumeResult { consumedFrom: QuotaConsumedFrom } -export interface SttQuotaReservation { +/** reserve_*_quota RPC 결과 — 공급자 호출 전에 한 단위를 잡아 두는 리스(lease). */ +export interface QuotaReservation { allowed: boolean reservationId: string | null status: 'reserved' | 'completed' | 'released' | 'denied' @@ -72,6 +73,59 @@ export interface SttQuotaReservation { consumedFrom: 'base' | 'overage' | 'unlimited' | 'none' } +export type SttQuotaReservation = QuotaReservation + +export type QuotaReservationFinalStatus = 'completed' | 'released' + +const RESERVATION_STATUSES: readonly unknown[] = ['reserved', 'completed', 'released', 'denied'] +const QUOTA_PERIODS: readonly unknown[] = ['daily', 'weekly'] +const QUOTA_TIERS: readonly unknown[] = ['free', 'pro', 'pro_plus', 'team', 'enterprise'] +const RESERVATION_SOURCES: readonly unknown[] = ['base', 'overage', 'unlimited', 'none'] + +/** reserve_stt_quota / reserve_llm_quota 공통 응답 검증. 형식이 어긋나면 throw. */ +export function parseQuotaReservationResponse(data: unknown, label: string): QuotaReservation { + if (!data || typeof data !== 'object' || Array.isArray(data)) { + throw new Error(`Invalid ${label} quota reservation response.`) + } + const result = data as Record + if ( + typeof result.allowed !== 'boolean' + || (result.reservation_id !== null && typeof result.reservation_id !== 'string') + || !RESERVATION_STATUSES.includes(result.status) + || typeof result.current !== 'number' + || typeof result.limit !== 'number' + || !QUOTA_PERIODS.includes(result.period) + || !QUOTA_TIERS.includes(result.tier) + || typeof result.overage_credits !== 'number' + || !RESERVATION_SOURCES.includes(result.consumed_from) + ) { + throw new Error(`Invalid ${label} quota reservation response.`) + } + return { + allowed: result.allowed, + reservationId: result.reservation_id as string | null, + status: result.status as QuotaReservation['status'], + current: result.current, + limit: result.limit, + period: result.period as QuotaPeriod, + tier: result.tier as Tier, + overageCredits: result.overage_credits, + consumedFrom: result.consumed_from as QuotaReservation['consumedFrom'], + } +} + +/** finalize_*_quota RPC 공통 응답 검증. */ +export function parseQuotaFinalizeResponse(data: unknown, label: string): QuotaReservationFinalStatus { + if (!data || typeof data !== 'object' || Array.isArray(data)) { + throw new Error(`Invalid ${label} quota finalization response.`) + } + const status = (data as Record).status + if (status !== 'completed' && status !== 'released') { + throw new Error(`Invalid ${label} quota finalization response.`) + } + return status +} + export async function reserveSttQuota( userId: string, reservationId: string, @@ -84,38 +138,14 @@ export async function reserveSttQuota( if (error || !data || typeof data !== 'object' || Array.isArray(data)) { throw new Error('Failed to reserve STT quota.') } - const result = data as Record - if ( - typeof result.allowed !== 'boolean' - || (result.reservation_id !== null && typeof result.reservation_id !== 'string') - || !['reserved', 'completed', 'released', 'denied'].includes(String(result.status)) - || typeof result.current !== 'number' - || typeof result.limit !== 'number' - || !['daily', 'weekly'].includes(String(result.period)) - || !['free', 'pro', 'pro_plus', 'team', 'enterprise'].includes(String(result.tier)) - || typeof result.overage_credits !== 'number' - || !['base', 'overage', 'unlimited', 'none'].includes(String(result.consumed_from)) - ) { - throw new Error('Invalid STT quota reservation response.') - } - return { - allowed: result.allowed, - reservationId: result.reservation_id as string | null, - status: result.status as SttQuotaReservation['status'], - current: result.current, - limit: result.limit, - period: result.period as QuotaPeriod, - tier: result.tier as Tier, - overageCredits: result.overage_credits, - consumedFrom: result.consumed_from as SttQuotaReservation['consumedFrom'], - } + return parseQuotaReservationResponse(data, 'STT') } export async function finalizeSttQuota( reservationId: string, succeeded: boolean, serviceRoleClient: ReturnType, -): Promise<'completed' | 'released'> { +): Promise { const { data, error } = await serviceRoleClient.rpc('finalize_stt_quota', { p_reservation_id: reservationId, p_succeeded: succeeded, @@ -123,11 +153,50 @@ export async function finalizeSttQuota( if (error || !data || typeof data !== 'object' || Array.isArray(data)) { throw new Error('Failed to finalize STT quota.') } - const status = (data as Record).status - if (status !== 'completed' && status !== 'released') { - throw new Error('Invalid STT quota finalization response.') + return parseQuotaFinalizeResponse(data, 'STT') +} + +/** + * LLM 쿼터 한 단위를 공급자 호출 **전에** 잡는다 (`reserve_llm_quota`). + * consume_quota 와 같은 per-user/feature advisory lock 아래에서 진행 중인 요청까지 + * 사용량에 포함하므로, 동시 요청이 한도를 넘어 공급자 비용을 쓰지 못한다. + * 성공/실패는 finalizeLlmQuota 로 확정하고, 확정되지 않은 리스는 10분 뒤 환불된다. + */ +export async function reserveLlmQuota( + userId: string, + reservationId: string, + feature: QuotaFeature, + baseLimit: number, + period: QuotaPeriod, + serviceRoleClient: ReturnType, +): Promise { + const { data, error } = await serviceRoleClient.rpc('reserve_llm_quota', { + p_user_id: userId, + p_reservation_id: reservationId, + p_feature: feature, + p_base_limit: baseLimit, + p_period: period, + }) + if (error || !data || typeof data !== 'object' || Array.isArray(data)) { + throw new Error('Failed to reserve LLM quota.') } - return status + return parseQuotaReservationResponse(data, 'LLM') +} + +/** 예약을 확정한다 — succeeded=true 면 소비 유지, false 면 사용량과 overage 크레딧 환불. */ +export async function finalizeLlmQuota( + reservationId: string, + succeeded: boolean, + serviceRoleClient: ReturnType, +): Promise { + const { data, error } = await serviceRoleClient.rpc('finalize_llm_quota', { + p_reservation_id: reservationId, + p_succeeded: succeeded, + }) + if (error || !data || typeof data !== 'object' || Array.isArray(data)) { + throw new Error('Failed to finalize LLM quota.') + } + return parseQuotaFinalizeResponse(data, 'LLM') } interface SubscriptionQuotaState { @@ -151,6 +220,14 @@ async function readSubscriptionQuotaState( } } +/** 구독 티어만 조회 (행이 없으면 free). 사용량은 읽지 않는다. */ +export async function readSubscriptionTier( + userId: string, + serviceRoleClient: ReturnType, +): Promise { + return (await readSubscriptionQuotaState(userId, serviceRoleClient)).tier +} + /** 기간 창(daily=오늘, weekly=오늘-6..오늘) 안의 daily_usage 합계. */ async function readWindowUsage( userId: string, diff --git a/server/supabase/functions/llm-proxy/handler.test.ts b/server/supabase/functions/llm-proxy/handler.test.ts new file mode 100644 index 0000000..2470194 --- /dev/null +++ b/server/supabase/functions/llm-proxy/handler.test.ts @@ -0,0 +1,378 @@ +// Regression tests for llm-proxy: +// * quota is reserved before the provider call, so parallel requests past the +// allowance never reach Anthropic; +// * the provider's first-byte deadline no longer cuts off a long stream, while +// idle/total stream deadlines and client disconnects still stop it; +// * provider-side failures refund the reservation. +import type { QuotaFeature, QuotaPeriod, QuotaReservation } from '../_shared/quota.ts' +import { + createLlmProxyHandler, + type LlmProxyDeps, + type LlmQuotaStore, + type ProviderFetch, +} from './handler.ts' +import { createSseCompletionTracker } from './provider-deadline.ts' + +function assert(condition: boolean, message: string): asserts condition { + if (!condition) throw new Error(message) +} + +function assertEquals(actual: T, expected: T, message: string): void { + const a = JSON.stringify(actual) + const e = JSON.stringify(expected) + if (a !== e) throw new Error(`${message}: expected ${e}, got ${a}`) +} + +const delay = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)) +const encoder = new TextEncoder() + +/** In-memory reserve/finalize with the same in-flight semantics as reserve_llm_quota. */ +class FakeLlmQuotaStore implements LlmQuotaStore { + used = 0 + reserveCalls = 0 + finalized: Array<{ id: string; succeeded: boolean }> = [] + private reservations = new Map() + constructor(private readonly limit: number, private readonly tier: 'free' | 'pro' = 'pro') {} + + readTier(): Promise<'free' | 'pro'> { + return Promise.resolve(this.tier) + } + + reserve( + _userId: string, + reservationId: string, + _feature: QuotaFeature, + _baseLimit: number, + period: QuotaPeriod, + ): Promise { + this.reserveCalls++ + if (this.used >= this.limit) { + return Promise.resolve({ + allowed: false, reservationId: null, status: 'denied', current: this.used, limit: this.limit, + period, tier: this.tier, overageCredits: 0, consumedFrom: 'none', + }) + } + this.used++ + this.reservations.set(reservationId, 'reserved') + return Promise.resolve({ + allowed: true, reservationId, status: 'reserved', current: this.used, limit: this.limit, + period, tier: this.tier, overageCredits: 0, consumedFrom: 'base', + }) + } + + finalize(reservationId: string, succeeded: boolean): Promise<'completed' | 'released'> { + this.finalized.push({ id: reservationId, succeeded }) + if (this.reservations.get(reservationId) === 'reserved' && !succeeded) this.used-- + const status = succeeded ? 'completed' : 'released' + this.reservations.set(reservationId, status) + return Promise.resolve(status) + } +} + +interface UpstreamPlan { + chunks: string[] + intervalMs: number + /** Stop producing after this many chunks and hang (no more data, no EOF). */ + stallAfter?: number +} + +/** + * Mock of Deno's fetch body semantics: aborting the request signal after the + * headers arrived errors the response body stream. + */ +function slowSseResponse(plan: UpstreamPlan, signal: AbortSignal): { response: Response; state: { aborted: boolean } } { + const state = { aborted: false } + let index = 0 + let timer: ReturnType | null = null + let wake: (() => void) | null = null + let streamController: ReadableStreamDefaultController | null = null + const onAbort = () => { + state.aborted = true + if (timer !== null) clearTimeout(timer) + timer = null + wake?.() + try { + streamController?.error(signal.reason) + } catch { + // closed + } + } + signal.addEventListener('abort', onAbort, { once: true }) + const body = new ReadableStream({ + start(controller) { + streamController = controller + }, + async pull(controller) { + if (state.aborted) return + if (plan.stallAfter !== undefined && index >= plan.stallAfter) { + await new Promise((resolve) => { + wake = resolve + }) + return + } + await new Promise((resolve) => { + wake = resolve + timer = setTimeout(resolve, plan.intervalMs) + }) + timer = null + if (state.aborted) return + if (index >= plan.chunks.length) { + signal.removeEventListener('abort', onAbort) + controller.close() + return + } + controller.enqueue(encoder.encode(plan.chunks[index++])) + }, + cancel() { + if (timer !== null) clearTimeout(timer) + signal.removeEventListener('abort', onAbort) + }, + }) + return { + response: new Response(body, { status: 200, headers: { 'Content-Type': 'text/event-stream' } }), + state, + } +} + +const SSE_CHUNKS = [ + 'event: message_start\ndata: {"type":"message_start"}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"text_delta","text":"1. "}}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"text_delta","text":"item"}}\n\n', + 'event: content_block_delta\ndata: {"type":"content_block_delta","delta":{"type":"text_delta","text":" two"}}\n\n', + 'event: message_delta\ndata: {"type":"message_delta"}\n\n', + 'event: message_st', + 'op\ndata: {"type":"message_stop"}\n\n', +] + +function buildDeps( + store: FakeLlmQuotaStore, + fetchProvider: ProviderFetch, + timeouts: LlmProxyDeps['timeouts'] = {}, +): LlmProxyDeps { + let seq = 0 + return { + authenticate: () => Promise.resolve({ id: 'user-1' }), + providerKey: () => 'test-key', + quotaStore: () => store, + issueGenerationReceipt: () => Promise.resolve('00000000-0000-4000-8000-000000000001'), + fetchProvider, + newReservationId: () => `res-${++seq}`, + timeouts, + } +} + +function llmRequest(stream: boolean, signal?: AbortSignal): Request { + return new Request('http://localhost/llm-proxy', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + messages: [{ role: 'user', content: 'Write a detailed itemised summary.' }], + max_tokens: 2048, + stream, + }), + signal, + }) +} + +function okJson(): Response { + return new Response(JSON.stringify({ content: [{ type: 'text', text: 'hello' }] }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }) +} + +Deno.test('stream longer than the first-byte deadline is relayed in full and charged once', async () => { + const store = new FakeLlmQuotaStore(10) + const upstreams: Array<{ aborted: boolean }> = [] + const fetchProvider: ProviderFetch = (_input, init) => { + const { response, state } = slowSseResponse({ chunks: SSE_CHUNKS, intervalMs: 25 }, init.signal) + upstreams.push(state) + return Promise.resolve(response) + } + // 7 chunks x 25 ms ≈ 200 ms of streaming against a 40 ms first-byte deadline. + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider, { + ttfbMs: 40, + streamIdleMs: 500, + streamTotalMs: 5_000, + })) + + const resp = await handler(llmRequest(true)) + assertEquals(resp.status, 200, 'status') + const text = await resp.text() + assertEquals(text, SSE_CHUNKS.join(''), 'relayed body') + assert(upstreams.length === 1 && !upstreams[0].aborted, 'upstream must not be aborted') + assertEquals(store.finalized, [{ id: 'res-1', succeeded: true }], 'finalized') + assertEquals(store.used, 1, 'one unit spent') +}) + +Deno.test('stream that stalls past the idle deadline is aborted and refunded', async () => { + const store = new FakeLlmQuotaStore(10) + const upstreams: Array<{ aborted: boolean }> = [] + const fetchProvider: ProviderFetch = (_input, init) => { + const { response, state } = slowSseResponse( + { chunks: SSE_CHUNKS, intervalMs: 5, stallAfter: 2 }, + init.signal, + ) + upstreams.push(state) + return Promise.resolve(response) + } + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider, { + ttfbMs: 1_000, + streamIdleMs: 60, + streamTotalMs: 5_000, + })) + + const resp = await handler(llmRequest(true)) + assertEquals(resp.status, 200, 'status') + let errored = false + try { + await resp.text() + } catch (err) { + errored = err instanceof DOMException && err.name === 'TimeoutError' + } + assert(errored, 'client stream must error with a timeout') + assert(upstreams.length === 1 && upstreams[0].aborted, 'upstream must be aborted') + assertEquals(store.finalized, [{ id: 'res-1', succeeded: false }], 'refunded') + assertEquals(store.used, 0, 'unit returned') +}) + +Deno.test('stream that exceeds the total deadline is aborted and refunded', async () => { + const store = new FakeLlmQuotaStore(10) + const many = Array.from({ length: 200 }, () => 'event: ping\ndata: {"type":"ping"}\n\n') + const fetchProvider: ProviderFetch = (_input, init) => + Promise.resolve(slowSseResponse({ chunks: many, intervalMs: 10 }, init.signal).response) + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider, { + ttfbMs: 1_000, + streamIdleMs: 500, + streamTotalMs: 120, + })) + + const resp = await handler(llmRequest(true)) + let errored = false + try { + await resp.text() + } catch { + errored = true + } + assert(errored, 'client stream must error') + assertEquals(store.finalized, [{ id: 'res-1', succeeded: false }], 'refunded') +}) + +Deno.test('stream that ends without message_stop is refunded', async () => { + const store = new FakeLlmQuotaStore(10) + const fetchProvider: ProviderFetch = (_input, init) => + Promise.resolve(slowSseResponse({ chunks: SSE_CHUNKS.slice(0, 3), intervalMs: 5 }, init.signal).response) + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider, { ttfbMs: 1_000 })) + + const resp = await handler(llmRequest(true)) + await resp.text() + assertEquals(store.finalized, [{ id: 'res-1', succeeded: false }], 'refunded') +}) + +Deno.test('client disconnect cancels the upstream stream and keeps the charge', async () => { + const store = new FakeLlmQuotaStore(10) + const upstreams: Array<{ aborted: boolean }> = [] + const fetchProvider: ProviderFetch = (_input, init) => { + const { response, state } = slowSseResponse({ chunks: SSE_CHUNKS, intervalMs: 20 }, init.signal) + upstreams.push(state) + return Promise.resolve(response) + } + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider, { ttfbMs: 1_000 })) + + const resp = await handler(llmRequest(true)) + const reader = resp.body!.getReader() + await reader.read() + await reader.cancel('client left') + await delay(10) + assert(upstreams.length === 1 && upstreams[0].aborted, 'upstream must be aborted') + assertEquals(store.finalized, [{ id: 'res-1', succeeded: true }], 'charged') +}) + +Deno.test('parallel requests past the allowance never reach the provider', async () => { + const store = new FakeLlmQuotaStore(2) + let providerCalls = 0 + let release: () => void = () => {} + const gate = new Promise((resolve) => { + release = resolve + }) + const fetchProvider: ProviderFetch = async () => { + providerCalls++ + await gate + return okJson() + } + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider)) + + const pending = Array.from({ length: 6 }, () => handler(llmRequest(false))) + await delay(10) + // All reservations are decided while the first two provider calls are still in flight. + assertEquals(providerCalls, 2, 'provider calls while in flight') + release() + const statuses = (await Promise.all(pending)).map((r) => r.status).sort() + assertEquals(statuses, [200, 200, 429, 429, 429, 429], 'statuses') + assertEquals(providerCalls, 2, 'total provider calls') + assertEquals(store.used, 2, 'units spent') +}) + +Deno.test('provider error refunds the reservation', async () => { + const store = new FakeLlmQuotaStore(1) + const fetchProvider: ProviderFetch = () => + Promise.resolve(new Response('overloaded', { status: 529 })) + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider)) + + const resp = await handler(llmRequest(false)) + assertEquals(resp.status, 502, 'status') + await resp.body?.cancel() + assertEquals(store.finalized, [{ id: 'res-1', succeeded: false }], 'refunded') + assertEquals(store.used, 0, 'unit returned') +}) + +Deno.test('first-byte timeout returns 504 and refunds the reservation', async () => { + const store = new FakeLlmQuotaStore(1) + const fetchProvider: ProviderFetch = (_input, init) => + new Promise((_resolve, reject) => { + const signal = init.signal + signal.addEventListener('abort', () => reject(signal.reason), { once: true }) + }) + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider, { ttfbMs: 30 })) + + const resp = await handler(llmRequest(false)) + assertEquals(resp.status, 504, 'status') + assertEquals(await resp.json(), { error: 'provider_timeout' }, 'body') + assertEquals(store.finalized, [{ id: 'res-1', succeeded: false }], 'refunded') +}) + +Deno.test('denied reservation returns 429 without calling the provider', async () => { + const store = new FakeLlmQuotaStore(0, 'free') + let providerCalls = 0 + const fetchProvider: ProviderFetch = () => { + providerCalls++ + return Promise.resolve(okJson()) + } + const handler = createLlmProxyHandler(buildDeps(store, fetchProvider)) + + const resp = await handler(llmRequest(false)) + assertEquals(resp.status, 429, 'status') + const body = await resp.json() + assertEquals(body.error, 'quota_exceeded', 'error') + assertEquals(body.model, 'claude-haiku-4-5-20251001', 'model') + assertEquals(body.tier, 'free', 'tier') + assertEquals(providerCalls, 0, 'provider calls') +}) + +Deno.test('non-stream success settles the reservation as completed', async () => { + const store = new FakeLlmQuotaStore(5) + const handler = createLlmProxyHandler(buildDeps(store, () => Promise.resolve(okJson()))) + + const resp = await handler(llmRequest(false)) + assertEquals(resp.status, 200, 'status') + assertEquals(await resp.json(), { content: [{ type: 'text', text: 'hello' }] }, 'body') + assertEquals(store.finalized, [{ id: 'res-1', succeeded: true }], 'completed') +}) + +Deno.test('SSE completion tracker finds message_stop split across chunks', () => { + const tracker = createSseCompletionTracker() + tracker.push(encoder.encode('event: message_st')) + assert(!tracker.completed, 'not yet') + tracker.push(encoder.encode('op\ndata: {}\n\n')) + assert(tracker.completed, 'completed') +}) diff --git a/server/supabase/functions/llm-proxy/handler.ts b/server/supabase/functions/llm-proxy/handler.ts new file mode 100644 index 0000000..c5f3454 --- /dev/null +++ b/server/supabase/functions/llm-proxy/handler.ts @@ -0,0 +1,336 @@ +// server/supabase/functions/llm-proxy/handler.ts +// Use case: proxy one Anthropic Messages API call for an authenticated user. +// +// Order of work: +// 1. authenticate, validate the request, resolve the tier and model +// 2. reserve one quota unit (atomic, counts in-flight requests) +// 3. call the provider only while holding the reservation +// 4. settle the reservation: completed when the user received the answer, +// released (refunded) when the provider failed, timed out, or the stream +// ended before message_stop +// +// IO goes through ports (auth, quota store, receipt issuer, provider fetch) so +// the quota ordering and the stream deadlines are testable without Supabase or +// Anthropic. index.ts wires the real adapters. + +import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts' +import { authErrorResponse, type AuthError } from '../_shared/auth.ts' +import { + getQuotaPolicy, + modelToQuotaKey, + type QuotaFeature, + type QuotaPeriod, + type QuotaReservation, + type QuotaReservationFinalStatus, + type Tier, +} from '../_shared/quota.ts' +import { + hasAssistantText, + LlmRequestError, + parseLlmRequest, +} from '../_shared/llm-contract.ts' +import { + GENERATION_ID_HEADER, + GENERATION_PURPOSE_HEADER, + GenerationReceiptError, + type GenerationPurpose, + parseGenerationPurpose, +} from '../_shared/generation-receipt.ts' +import { buildAnthropicSystemBlocks } from '../_shared/generative-ai-safety.ts' +import { + createProviderCallDeadline, + DEFAULT_LLM_PROVIDER_TIMEOUTS, + type LlmProviderTimeouts, + relayProviderStream, + type StreamRelayOutcome, +} from './provider-deadline.ts' + +export const ANTHROPIC_MESSAGES_URL = 'https://api.anthropic.com/v1/messages' + +/** 티어별 허용 모델 — free는 Haiku만, pro는 +Sonnet, pro_plus는 +Opus, team/enterprise는 전 모델 */ +export const TIER_MODELS: Readonly> = { + free: ['claude-haiku-4-5-20251001'], + pro: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], + pro_plus: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], + team: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], + enterprise: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], +} + +export const DEFAULT_MODEL: Readonly> = { + free: 'claude-haiku-4-5-20251001', + pro: 'claude-sonnet-4-6', + pro_plus: 'claude-sonnet-4-6', + team: 'claude-sonnet-4-6', + enterprise: 'claude-sonnet-4-6', +} + +/** Port over the LLM quota tables (subscriptions, daily_usage, llm_quota_reservations). */ +export interface LlmQuotaStore { + readTier(userId: string): Promise + reserve( + userId: string, + reservationId: string, + feature: QuotaFeature, + baseLimit: number, + period: QuotaPeriod, + ): Promise + finalize(reservationId: string, succeeded: boolean): Promise +} + +/** Request passed to the provider port. */ +export interface ProviderRequest { + method: 'POST' + headers: Record + body: string + signal: AbortSignal +} + +/** Port over the provider HTTP call (global fetch in production). */ +export type ProviderFetch = (url: string, init: ProviderRequest) => Promise + +export interface LlmProxyDeps { + authenticate(req: Request): Promise<{ id: string }> + /** null when the provider key is not configured. */ + providerKey(): string | null + quotaStore(): LlmQuotaStore + /** Issue a content-generation receipt; throws GenerationReceiptError when unavailable. */ + issueGenerationReceipt(userId: string, purpose: GenerationPurpose, model: string): Promise + fetchProvider?: ProviderFetch + newReservationId?(): string + timeouts?: Partial +} + +const NO_STORE_JSON = { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' } + +function json(status: number, body: Record, noStore = true): Response { + return new Response(JSON.stringify(body), { + status, + headers: noStore ? NO_STORE_JSON : { ...corsHeaders, 'Content-Type': 'application/json' }, + }) +} + +/** 쓰지 않을 공급자 응답 본문을 닫아 연결과 생성을 정리한다. */ +async function discardBody(resp: Response): Promise { + if (!resp.body || resp.bodyUsed) return + try { + await resp.body.cancel() + } catch { + // 이미 닫힌 스트림 — 무시 + } +} + +function generationHeaders(generationId: string | null): Record { + return generationId === null ? {} : { [GENERATION_ID_HEADER]: generationId } +} + +function isAuthLikeError(err: unknown): err is AuthError { + return Boolean(err && typeof err === 'object' && 'status' in err && 'message' in err) +} + +/** + * Settle a reservation exactly once. A failed finalize is logged; an unsettled + * reservation is released when its lease expires. + */ +function reservationSettler(store: LlmQuotaStore, reservationId: string): (succeeded: boolean) => Promise { + let done = false + return async (succeeded: boolean) => { + if (done) return + done = true + try { + await store.finalize(reservationId, succeeded) + } catch { + console.error('LLM quota finalization failed', { succeeded }) + } + } +} + +export function createLlmProxyHandler(deps: LlmProxyDeps): (req: Request) => Promise { + const fetchProvider: ProviderFetch = deps.fetchProvider ?? ((url, init) => fetch(url, init)) + const newReservationId = deps.newReservationId ?? (() => crypto.randomUUID()) + const timeouts: LlmProviderTimeouts = { ...DEFAULT_LLM_PROVIDER_TIMEOUTS, ...deps.timeouts } + + return async (req: Request): Promise => { + const preflight = handleCorsPreflightRequest(req) + if (preflight) return preflight + + if (req.method !== 'POST') { + return json(405, { error: 'Method not allowed' }, false) + } + + let settleReservation: ((succeeded: boolean) => Promise) | null = null + try { + const user = await deps.authenticate(req) + + let rawBody: unknown + try { + rawBody = await req.json() + } catch { + throw new LlmRequestError('Invalid JSON body') + } + const body = parseLlmRequest(rawBody) + const generationPurpose = parseGenerationPurpose(req.headers.get(GENERATION_PURPOSE_HEADER)) + + // A deployment without a provider must not consume quota or fabricate an answer. + const anthropicKey = deps.providerKey() + if (!anthropicKey) return json(503, { error: 'provider_unavailable' }) + + const quota = deps.quotaStore() + + // 1단계: 티어 조회 + 모델 검증 + const tier = await quota.readTier(user.id) + const requestedModel = body.model ?? DEFAULT_MODEL[tier] + if (!TIER_MODELS[tier].includes(requestedModel)) { + return json(403, { + error: 'model_not_allowed', + tier, + requested: requestedModel, + allowed: TIER_MODELS[tier], + }, false) + } + + // 2단계: 공급자 호출 전에 쿼터 한 단위를 원자적으로 예약한다. + // 진행 중인 요청도 사용량에 포함되므로 병렬 요청이 한도를 넘어 공급자 비용을 쓰지 못한다. + const quotaKey = modelToQuotaKey(requestedModel) + const policy = getQuotaPolicy(tier, quotaKey) + const reservation = await quota.reserve( + user.id, + newReservationId(), + quotaKey, + policy.limit, + policy.period, + ) + if (!reservation.allowed || reservation.reservationId === null) { + return json(429, { + error: 'quota_exceeded', + model: requestedModel, + current: reservation.current, + limit: reservation.limit, + period: reservation.period, + tier: reservation.tier, + overage_credits: reservation.overageCredits, + }, false) + } + const settle = reservationSettler(quota, reservation.reservationId) + settleReservation = settle + + // 3단계: 공급자 호출 (Prompt Caching 2024-07-31 활성화). + // 공급자 5xx/과부하/타임아웃으로 답을 주지 못한 요청은 예약을 해제(환불)한다. + const deadline = createProviderCallDeadline(timeouts.ttfbMs, req.signal) + let anthropicResp: Response + try { + anthropicResp = await fetchProvider(ANTHROPIC_MESSAGES_URL, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'x-api-key': anthropicKey, + 'anthropic-version': '2023-06-01', + 'anthropic-beta': 'prompt-caching-2024-07-31', + }, + body: JSON.stringify({ + model: requestedModel, + max_tokens: body.max_tokens, + system: buildAnthropicSystemBlocks(body.system), + messages: body.messages, + stream: body.stream, + }), + signal: deadline.signal, + }) + } catch (err) { + deadline.dispose() + throw err + } + + if (!anthropicResp.ok) { + deadline.dispose() + console.error('Anthropic request failed', { status: anthropicResp.status }) + await discardBody(anthropicResp) + await settle(false) + return json(502, { error: 'provider_request_failed' }) + } + + if (body.stream && anthropicResp.body) { + // The first-byte deadline covers the provider call only; the streamed + // body gets its own idle/total deadlines. + deadline.headersReceived() + let generationId: string | null + try { + generationId = generationPurpose === null + ? null + : await deps.issueGenerationReceipt(user.id, generationPurpose, requestedModel) + } catch (err) { + deadline.abort(err) + deadline.dispose() + await discardBody(anthropicResp) + throw err + } + const relayed = relayProviderStream(anthropicResp.body, { + idleMs: timeouts.streamIdleMs, + totalMs: timeouts.streamTotalMs, + abortUpstream: (reason) => deadline.abort(reason), + onSettled: async (outcome: StreamRelayOutcome) => { + deadline.dispose() + if (outcome === 'failed') console.error('Anthropic stream did not complete') + // A client that stops reading keeps the charge: the provider already + // generated what was sent. Only provider-side failures are refunded. + await settle(outcome !== 'failed') + }, + }) + return new Response(relayed, { + status: 200, + headers: { + ...corsHeaders, + ...generationHeaders(generationId), + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-store', + Connection: 'keep-alive', + }, + }) + } + + // 비스트리밍: 본문 전체가 답이므로 첫 바이트 기한을 본문 읽기까지 유지한다. + let nonStreamData: unknown + try { + nonStreamData = await anthropicResp.json() + } finally { + deadline.dispose() + } + if (!hasAssistantText(nonStreamData)) { + console.error('Anthropic returned an invalid response shape') + await settle(false) + return json(502, { error: 'provider_invalid_response' }) + } + + const generationId = generationPurpose === null + ? null + : await deps.issueGenerationReceipt(user.id, generationPurpose, requestedModel) + await settle(true) + return new Response(JSON.stringify(nonStreamData), { + status: 200, + headers: { + ...corsHeaders, + ...generationHeaders(generationId), + 'Content-Type': 'application/json', + 'Cache-Control': 'no-store', + }, + }) + } catch (err) { + // Anything that failed while a reservation was held gives the unit back. + await settleReservation?.(false) + if (err instanceof GenerationReceiptError) { + const invalidPurpose = err.code === 'invalid_generation_purpose' + return json(invalidPurpose ? 400 : 503, { + error: invalidPurpose ? 'invalid_request' : 'generation_receipt_unavailable', + }) + } + if (err instanceof LlmRequestError) { + return json(err.status, { error: 'invalid_request', message: err.message }) + } + if (isAuthLikeError(err)) { + return authErrorResponse(err, corsHeaders) + } + const timedOut = err instanceof DOMException && err.name === 'TimeoutError' + console.error('llm-proxy failed', { kind: timedOut ? 'provider_timeout' : 'internal_error' }) + return json(timedOut ? 504 : 500, { error: timedOut ? 'provider_timeout' : 'internal_error' }) + } + } +} diff --git a/server/supabase/functions/llm-proxy/index.ts b/server/supabase/functions/llm-proxy/index.ts index e5870c0..57034b4 100644 --- a/server/supabase/functions/llm-proxy/index.ts +++ b/server/supabase/functions/llm-proxy/index.ts @@ -1,293 +1,50 @@ // server/supabase/functions/llm-proxy/index.ts -// Anthropic Claude Messages API 프록시. +// Anthropic Claude Messages API 프록시 — composition root. // Phase 3.2: 모델별 쿼터 (Haiku/Sonnet/Opus × Free/Pro/Pro+) // 요청: application/json { messages, system?, max_tokens?, model? } // 응답: JSON (non-stream) 또는 SSE (stream=true) +// 유스케이스와 쿼터/기한 규칙은 handler.ts, 공급자 기한은 provider-deadline.ts 에 있다. -import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts' -import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts' +import { requireUser } from '../_shared/auth.ts' import { - checkQuota, - consumeQuota, createServiceRoleClient, - modelToQuotaKey, - getQuotaPolicy, - type Tier, + finalizeLlmQuota, + readSubscriptionTier, + reserveLlmQuota, } from '../_shared/quota.ts' -import { - hasAssistantText, - LlmRequestError, - parseLlmRequest, -} from '../_shared/llm-contract.ts' -import { - GENERATION_ID_HEADER, - GENERATION_PURPOSE_HEADER, - GenerationReceiptError, - parseGenerationPurpose, - parseGenerationReceiptId, - type GenerationPurpose, -} from '../_shared/generation-receipt.ts' -import { buildAnthropicSystemBlocks } from '../_shared/generative-ai-safety.ts' +import { GenerationReceiptError, parseGenerationReceiptId } from '../_shared/generation-receipt.ts' import { readProviderKey } from '../_shared/provider-key.ts' +import { createLlmProxyHandler, type LlmQuotaStore } from './handler.ts' -/** 티어별 허용 모델 — free는 Haiku만, pro는 +Sonnet, pro_plus는 +Opus, team/enterprise는 전 모델 */ -const TIER_MODELS: Record = { - free: ['claude-haiku-4-5-20251001'], - pro: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], - pro_plus: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], - team: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], - enterprise: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'], +let serviceClient: ReturnType | null = null +function service(): ReturnType { + serviceClient ??= createServiceRoleClient() + return serviceClient } -const DEFAULT_MODEL: Record = { - free: 'claude-haiku-4-5-20251001', - pro: 'claude-sonnet-4-6', - pro_plus: 'claude-sonnet-4-6', - team: 'claude-sonnet-4-6', - enterprise: 'claude-sonnet-4-6', -} - -async function issueGenerationReceipt( - serviceClient: ReturnType, - userId: string, - purpose: GenerationPurpose | null, - model: string, -): Promise { - if (purpose === null) return null - const { data, error } = await serviceClient.rpc('issue_content_generation_receipt_v1', { - p_actor_id: userId, - p_purpose: purpose, - p_model: model, - }) - if (error) { - console.error('Generation receipt issuance failed', { code: error.code ?? 'unknown' }) - throw new GenerationReceiptError('generation_receipt_unavailable') - } - return parseGenerationReceiptId(data) -} - -/** 쓰지 않을 공급자 응답 본문을 닫아 연결과 생성을 정리한다. */ -async function discardBody(resp: Response): Promise { - if (!resp.body || resp.bodyUsed) return - try { - await resp.body.cancel() - } catch { - // 이미 닫힌 스트림 — 무시 +function supabaseLlmQuotaStore(client: ReturnType): LlmQuotaStore { + return { + readTier: (userId) => readSubscriptionTier(userId, client), + reserve: (userId, reservationId, feature, baseLimit, period) => + reserveLlmQuota(userId, reservationId, feature, baseLimit, period, client), + finalize: (reservationId, succeeded) => finalizeLlmQuota(reservationId, succeeded, client), } } -function generationHeaders(generationId: string | null): Record { - return generationId === null ? {} : { [GENERATION_ID_HEADER]: generationId } -} - -Deno.serve(async (req: Request) => { - const preflight = handleCorsPreflightRequest(req) - if (preflight) return preflight - - if (req.method !== 'POST') { - return new Response(JSON.stringify({ error: 'Method not allowed' }), { - status: 405, - headers: { ...corsHeaders, 'Content-Type': 'application/json' }, +Deno.serve(createLlmProxyHandler({ + authenticate: requireUser, + providerKey: () => readProviderKey('ANTHROPIC_API_KEY'), + quotaStore: () => supabaseLlmQuotaStore(service()), + issueGenerationReceipt: async (userId, purpose, model) => { + const { data, error } = await service().rpc('issue_content_generation_receipt_v1', { + p_actor_id: userId, + p_purpose: purpose, + p_model: model, }) - } - - try { - const user = await requireUser(req) - const serviceClient = createServiceRoleClient() - - let rawBody: unknown - try { - rawBody = await req.json() - } catch { - throw new LlmRequestError('Invalid JSON body') + if (error) { + console.error('Generation receipt issuance failed', { code: error.code ?? 'unknown' }) + throw new GenerationReceiptError('generation_receipt_unavailable') } - const body = parseLlmRequest(rawBody) - const generationPurpose = parseGenerationPurpose( - req.headers.get(GENERATION_PURPOSE_HEADER), - ) - - // A deployment without a provider must not consume quota or fabricate an answer. - const anthropicKey = readProviderKey('ANTHROPIC_API_KEY') - if (!anthropicKey) { - return new Response(JSON.stringify({ error: 'provider_unavailable' }), { - status: 503, - headers: { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }, - }) - } - - // 1단계: 티어 조회 (어떤 모델이든 한 번만 읽으면 됨 — haiku로 대리 조회) - const tierCheck = await checkQuota(user.id, 'llm_haiku', serviceClient) - const tier = tierCheck.tier - - // 모델 선택 + 티어 검증 - const requestedModel = body.model ?? DEFAULT_MODEL[tier] - if (!TIER_MODELS[tier].includes(requestedModel)) { - return new Response( - JSON.stringify({ - error: 'model_not_allowed', - tier, - requested: requestedModel, - allowed: TIER_MODELS[tier], - }), - { - status: 403, - headers: { ...corsHeaders, 'Content-Type': 'application/json' }, - }, - ) - } - - // 2단계: 해당 모델의 쿼터 확인 (모델별 일간/주간) - const quotaKey = modelToQuotaKey(requestedModel) - const modelQuota = await checkQuota(user.id, quotaKey, serviceClient) - if (!modelQuota.allowed) { - return new Response( - JSON.stringify({ - error: 'quota_exceeded', - model: requestedModel, - current: modelQuota.current, - limit: modelQuota.limit, - period: modelQuota.period, - tier, - overage_credits: modelQuota.overageCredits, - }), - { - status: 429, - headers: { ...corsHeaders, 'Content-Type': 'application/json' }, - }, - ) - } - - // 3단계: 공급자 호출 (Prompt Caching 2024-07-31 활성화) - // 쿼터는 공급자가 요청을 받아들인 뒤에만 소비한다. 공급자 5xx/과부하/타임아웃으로 - // 답을 주지 못한 요청이 쿼터를 깎고, 클라이언트 재시도가 한 번 더 깎는 일을 막는다. - const systemPayload = buildAnthropicSystemBlocks(body.system) - - const anthropicResp = await fetch('https://api.anthropic.com/v1/messages', { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - 'x-api-key': anthropicKey, - 'anthropic-version': '2023-06-01', - 'anthropic-beta': 'prompt-caching-2024-07-31', - }, - body: JSON.stringify({ - model: requestedModel, - max_tokens: body.max_tokens, - system: systemPayload, - messages: body.messages, - stream: body.stream, - }), - signal: AbortSignal.timeout(45_000), - }) - - if (!anthropicResp.ok) { - console.error('Anthropic request failed', { status: anthropicResp.status }) - await discardBody(anthropicResp) - return new Response(JSON.stringify({ error: 'provider_request_failed' }), { - status: 502, - headers: { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }, - }) - } - - // 비스트리밍 응답은 형식 검증까지 통과한 뒤에 소비한다. - let nonStreamData: unknown = null - if (!body.stream || !anthropicResp.body) { - nonStreamData = await anthropicResp.json() - if (!hasAssistantText(nonStreamData)) { - console.error('Anthropic returned an invalid response shape') - return new Response(JSON.stringify({ error: 'provider_invalid_response' }), { - status: 502, - headers: { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }, - }) - } - } - - // 4단계: 쿼터 소비 (원자적 base → overage fallback). - // 동시 요청이 checkQuota 를 함께 통과했다면 여기서 초과분이 거절된다. - const policy = getQuotaPolicy(tier, quotaKey) - let consume: Awaited> - try { - consume = await consumeQuota(user.id, quotaKey, serviceClient, policy.limit) - } catch (err) { - await discardBody(anthropicResp) - throw err - } - if (!consume.allowed) { - await discardBody(anthropicResp) - return new Response( - JSON.stringify({ - error: 'quota_exceeded', - model: requestedModel, - current: consume.current, - limit: consume.limit, - tier, - overage_credits: consume.overageCredits, - }), - { - status: 429, - headers: { ...corsHeaders, 'Content-Type': 'application/json' }, - }, - ) - } - - if (body.stream && anthropicResp.body) { - const generationId = await issueGenerationReceipt( - serviceClient, - user.id, - generationPurpose, - requestedModel, - ) - return new Response(anthropicResp.body, { - status: 200, - headers: { - ...corsHeaders, - ...generationHeaders(generationId), - 'Content-Type': 'text/event-stream', - 'Cache-Control': 'no-store', - Connection: 'keep-alive', - }, - }) - } - - const generationId = await issueGenerationReceipt( - serviceClient, - user.id, - generationPurpose, - requestedModel, - ) - return new Response(JSON.stringify(nonStreamData), { - status: 200, - headers: { - ...corsHeaders, - ...generationHeaders(generationId), - 'Content-Type': 'application/json', - 'Cache-Control': 'no-store', - }, - }) - } catch (err) { - if (err instanceof GenerationReceiptError) { - const invalidPurpose = err.code === 'invalid_generation_purpose' - return new Response(JSON.stringify({ - error: invalidPurpose ? 'invalid_request' : 'generation_receipt_unavailable', - }), { - status: invalidPurpose ? 400 : 503, - headers: { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }, - }) - } - if (err instanceof LlmRequestError) { - return new Response(JSON.stringify({ error: 'invalid_request', message: err.message }), { - status: err.status, - headers: { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }, - }) - } - if (err && typeof err === 'object' && 'status' in err && 'message' in err) { - return authErrorResponse(err as AuthError, corsHeaders) - } - const timedOut = err instanceof DOMException && err.name === 'TimeoutError' - console.error('llm-proxy failed', { kind: timedOut ? 'provider_timeout' : 'internal_error' }) - return new Response(JSON.stringify({ error: timedOut ? 'provider_timeout' : 'internal_error' }), { - status: timedOut ? 504 : 500, - headers: { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }, - }) - } -}) + return parseGenerationReceiptId(data) + }, +})) diff --git a/server/supabase/functions/llm-proxy/provider-deadline.ts b/server/supabase/functions/llm-proxy/provider-deadline.ts new file mode 100644 index 0000000..c9c5fde --- /dev/null +++ b/server/supabase/functions/llm-proxy/provider-deadline.ts @@ -0,0 +1,226 @@ +// server/supabase/functions/llm-proxy/provider-deadline.ts +// Deadlines for the Anthropic call, split by phase. +// +// A single AbortSignal.timeout(45 s) on fetch also errors the response body, +// so it used to cut off streamed answers that took longer than 45 s in total. +// The phases now have their own limits: +// * time to first byte (response headers): ttfbMs — covers the provider call. +// For non-stream requests the whole body is the answer, so the deadline +// stays armed until the JSON has been read. +// * streamed body: an idle limit between chunks and a total limit for the +// stream, sized for a max_tokens=4096 answer. +// The client's request signal is tied in, so a client that disconnects cancels +// the upstream generation instead of letting it run to completion. + +export interface LlmProviderTimeouts { + /** Response headers must arrive within this time. */ + ttfbMs: number + /** Longest gap allowed between two streamed chunks. */ + streamIdleMs: number + /** Longest a streamed body may run after the headers arrived. */ + streamTotalMs: number +} + +export const DEFAULT_LLM_PROVIDER_TIMEOUTS: Readonly = Object.freeze({ + ttfbMs: 45_000, + streamIdleMs: 30_000, + streamTotalMs: 180_000, +}) + +export function providerTimeoutError(message: string): DOMException { + return new DOMException(message, 'TimeoutError') +} + +export interface ProviderCallDeadline { + /** Pass to fetch; aborts on the first-byte deadline, a stream deadline or a client disconnect. */ + readonly signal: AbortSignal + /** Stop the first-byte timer (the response headers arrived). */ + headersReceived(): void + abort(reason: unknown): void + /** Clear timers and listeners. Safe to call more than once. */ + dispose(): void +} + +export function createProviderCallDeadline( + ttfbMs: number, + clientSignal: AbortSignal | null, +): ProviderCallDeadline { + const controller = new AbortController() + let ttfbTimer: ReturnType | null = setTimeout(() => { + ttfbTimer = null + controller.abort(providerTimeoutError('provider response timeout')) + }, ttfbMs) + + const onClientAbort = (): void => { + controller.abort(clientSignal?.reason ?? new DOMException('client disconnected', 'AbortError')) + } + if (clientSignal?.aborted) onClientAbort() + else clientSignal?.addEventListener('abort', onClientAbort, { once: true }) + + const clearTtfb = (): void => { + if (ttfbTimer !== null) clearTimeout(ttfbTimer) + ttfbTimer = null + } + + return { + signal: controller.signal, + headersReceived: clearTtfb, + abort(reason: unknown): void { + controller.abort(reason) + }, + dispose(): void { + clearTtfb() + clientSignal?.removeEventListener('abort', onClientAbort) + }, + } +} + +/** + * Tracks whether an Anthropic SSE stream reached `message_stop`. + * Chunks can split the marker, so a short tail of the previous text is kept. + */ +export function createSseCompletionTracker(): { + push(chunk: Uint8Array): void + readonly completed: boolean +} { + const MARKER = 'message_stop' + const decoder = new TextDecoder() + let tail = '' + let completed = false + return { + push(chunk: Uint8Array): void { + if (completed) return + const text = tail + decoder.decode(chunk, { stream: true }) + if (text.includes(MARKER)) { + completed = true + tail = '' + return + } + tail = text.slice(-(MARKER.length - 1)) + }, + get completed(): boolean { + return completed + }, + } +} + +/** How a relayed stream ended. */ +export type StreamRelayOutcome = + /** Upstream ended after message_stop: the client received the full answer. */ + | 'completed' + /** Upstream errored, timed out, or ended without message_stop. */ + | 'failed' + /** The client stopped reading (disconnect/cancel). */ + | 'cancelled' + +export interface StreamRelayOptions { + idleMs: number + totalMs: number + /** Abort the upstream request (fetch signal). */ + abortUpstream(reason: unknown): void + /** Called exactly once when the relay ends. Awaited before the client stream closes or errors. */ + onSettled(outcome: StreamRelayOutcome): Promise +} + +/** + * Relay a provider SSE body to the client with an idle deadline between + * chunks and a total deadline for the whole stream. + */ +export function relayProviderStream( + upstream: ReadableStream, + options: StreamRelayOptions, +): ReadableStream { + const reader = upstream.getReader() + const tracker = createSseCompletionTracker() + let settled = false + let totalTimer: ReturnType | null = null + let idleTimer: ReturnType | null = null + let rejectIdle: ((reason: unknown) => void) | null = null + let downstream: ReadableStreamDefaultController | null = null + + const clearTimers = (): void => { + if (totalTimer !== null) clearTimeout(totalTimer) + if (idleTimer !== null) clearTimeout(idleTimer) + totalTimer = null + idleTimer = null + } + + const settle = async (outcome: StreamRelayOutcome): Promise => { + if (settled) return false + settled = true + clearTimers() + try { + await options.onSettled(outcome) + } catch { + // Settlement is best effort; the caller logs its own failures. + } + return true + } + + const terminate = async (reason: unknown): Promise => { + if (settled) return + options.abortUpstream(reason) + reader.cancel(reason).catch(() => undefined) + rejectIdle?.(reason) + if (await settle('failed')) { + try { + downstream?.error(reason) + } catch { + // Already closed. + } + } + } + + const readWithIdleDeadline = (): Promise> => + new Promise((resolve, reject) => { + rejectIdle = reject + idleTimer = setTimeout(() => { + idleTimer = null + void terminate(providerTimeoutError('provider stream idle timeout')) + }, options.idleMs) + reader.read().then(resolve, reject).finally(() => { + if (idleTimer !== null) clearTimeout(idleTimer) + idleTimer = null + rejectIdle = null + }) + }) + + return new ReadableStream({ + start(controller) { + downstream = controller + // Also fires while the client applies backpressure and no read is pending. + totalTimer = setTimeout(() => { + totalTimer = null + void terminate(providerTimeoutError('provider stream total timeout')) + }, options.totalMs) + }, + async pull(controller) { + if (settled) return + let result: ReadableStreamReadResult + try { + result = await readWithIdleDeadline() + } catch (err) { + if (settled) return + options.abortUpstream(err) + if (await settle('failed')) controller.error(err) + return + } + if (settled) return + if (result.done) { + const outcome: StreamRelayOutcome = tracker.completed ? 'completed' : 'failed' + // A stream that ended without message_stop is still closed cleanly: the + // client already treats a missing message_stop as incomplete, and any + // in-band `event: error` the provider sent stays readable. + if (await settle(outcome)) controller.close() + return + } + tracker.push(result.value) + controller.enqueue(result.value) + }, + async cancel(reason) { + options.abortUpstream(reason) + reader.cancel(reason).catch(() => undefined) + await settle('cancelled') + }, + }) +} diff --git a/server/supabase/migrations/20260928020800_llm_quota_reservations.sql b/server/supabase/migrations/20260928020800_llm_quota_reservations.sql new file mode 100644 index 0000000..aeb45f2 --- /dev/null +++ b/server/supabase/migrations/20260928020800_llm_quota_reservations.sql @@ -0,0 +1,296 @@ +-- llm-proxy: reserve LLM quota before calling the provider. +-- +-- llm-proxy used to run a read-only checkQuota (a SELECT on daily_usage), +-- call Anthropic, and only then charge with consume_quota. Nothing counted +-- requests that were still in flight, so N parallel requests from one user all +-- passed the pre-check and each paid for a full Anthropic generation before +-- consume_quota rejected every request past the allowance and threw the paid +-- reply away (the same flaw 20260928000037 fixed for meeting documents). +-- +-- reserve_llm_quota now takes one unit up front, under the same +-- per-user/per-feature advisory lock that consume_quota uses, and records it +-- in daily_usage immediately, so a concurrent request (or a consume_quota +-- call for the same feature) sees it. The edge function calls the provider +-- only while it holds a reservation, then settles it: +-- * finalize_llm_quota(id, true) -> 'completed' (the unit stays spent) +-- * finalize_llm_quota(id, false) -> 'released' (daily_usage -1, and the +-- overage credit is returned when the unit came from overage) +-- A crashed worker cannot hold a unit forever: reservations whose 10 minute +-- lease has expired are released by the next reserve call for that +-- user/feature, well beyond the proxy's 45 s time-to-first-byte and 180 s +-- stream deadlines. +-- +-- p_base_limit / p_period come from PLAN_QUOTA (packages/core plan catalog), +-- exactly like consume_quota. Only service_role may call these functions. +-- Apply this migration before deploying the llm-proxy that calls it. + +CREATE TABLE IF NOT EXISTS public.llm_quota_reservations ( + id uuid PRIMARY KEY, + user_id uuid NOT NULL REFERENCES auth.users(id) ON DELETE CASCADE, + feature text NOT NULL CHECK (feature IN ('llm_haiku', 'llm_sonnet', 'llm_opus')), + usage_date date NOT NULL DEFAULT CURRENT_DATE, + consumed_from text NOT NULL CHECK (consumed_from IN ('base', 'overage', 'unlimited')), + status text NOT NULL DEFAULT 'reserved' CHECK (status IN ('reserved', 'completed', 'released')), + tier text NOT NULL, + quota_period text NOT NULL CHECK (quota_period IN ('daily', 'weekly')), + quota_limit integer NOT NULL, + current_count integer NOT NULL, + overage_after integer NOT NULL, + lease_expires_at timestamptz NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + finalized_at timestamptz, + release_reason text +); + +ALTER TABLE public.llm_quota_reservations ENABLE ROW LEVEL SECURITY; +REVOKE ALL ON TABLE public.llm_quota_reservations FROM PUBLIC, anon, authenticated; +GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE public.llm_quota_reservations TO service_role; + +CREATE INDEX IF NOT EXISTS idx_llm_quota_reservations_reclaim + ON public.llm_quota_reservations(user_id, feature, lease_expires_at) + WHERE status = 'reserved'; + +CREATE OR REPLACE FUNCTION public.reserve_llm_quota( + p_user_id uuid, + p_reservation_id uuid, + p_feature text, + p_base_limit integer, + p_period text +) RETURNS jsonb +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ +DECLARE + existing public.llm_quota_reservations%ROWTYPE; + expired public.llm_quota_reservations%ROWTYPE; + subscription_tier text := 'free'; + overage integer := 0; + new_overage integer; + current_count integer := 0; + consumed_from text; +BEGIN + IF p_user_id IS NULL + OR p_reservation_id IS NULL + OR p_feature IS NULL + OR p_feature NOT IN ('llm_haiku', 'llm_sonnet', 'llm_opus') + OR p_base_limit IS NULL + OR p_base_limit < -1 + OR p_period IS NULL + OR p_period NOT IN ('daily', 'weekly') THEN + RAISE EXCEPTION 'invalid_llm_quota_reservation' USING ERRCODE = '22023'; + END IF; + + -- Same lock key as consume_quota so both paths serialise per user/feature. + PERFORM pg_advisory_xact_lock(hashtextextended(p_user_id::text || ':' || p_feature, 20260928)); + + SELECT * INTO existing + FROM public.llm_quota_reservations + WHERE id = p_reservation_id + FOR UPDATE; + + IF FOUND THEN + IF existing.user_id <> p_user_id OR existing.feature <> p_feature THEN + RAISE EXCEPTION 'llm_quota_reservation_conflict' USING ERRCODE = 'PT409'; + END IF; + RETURN jsonb_build_object( + 'allowed', existing.status IN ('reserved', 'completed'), + 'reservation_id', existing.id, + 'status', existing.status, + 'current', existing.current_count, + 'limit', existing.quota_limit, + 'period', existing.quota_period, + 'tier', existing.tier, + 'overage_credits', existing.overage_after, + 'consumed_from', existing.consumed_from + ); + END IF; + + -- Reclaim crashed requests before calculating the next allowance. + FOR expired IN + SELECT * + FROM public.llm_quota_reservations + WHERE user_id = p_user_id + AND feature = p_feature + AND status = 'reserved' + AND lease_expires_at <= now() + FOR UPDATE + LOOP + UPDATE public.daily_usage + SET count = greatest(count - 1, 0) + WHERE user_id = expired.user_id + AND date = expired.usage_date + AND feature = expired.feature; + + IF expired.consumed_from = 'overage' THEN + UPDATE public.subscriptions + SET overage_credits = overage_credits + 1, + updated_at = now() + WHERE user_id = expired.user_id; + END IF; + + UPDATE public.llm_quota_reservations + SET status = 'released', finalized_at = now(), release_reason = 'lease_expired' + WHERE id = expired.id; + END LOOP; + + SELECT coalesce(tier, 'free'), coalesce(overage_credits, 0) + INTO subscription_tier, overage + FROM public.subscriptions + WHERE user_id = p_user_id + FOR UPDATE; + + IF NOT FOUND THEN + subscription_tier := 'free'; + overage := 0; + END IF; + + -- Not available: never spend credits on a model the tier does not include. + IF p_base_limit = 0 THEN + RETURN jsonb_build_object( + 'allowed', false, + 'reservation_id', NULL, + 'status', 'denied', + 'current', 0, + 'limit', 0, + 'period', p_period, + 'tier', subscription_tier, + 'overage_credits', overage, + 'consumed_from', 'none' + ); + END IF; + + -- daily_usage already includes units held by in-flight reservations. + SELECT coalesce(sum(count), 0)::integer INTO current_count + FROM public.daily_usage + WHERE user_id = p_user_id + AND feature = p_feature + AND date >= CASE WHEN p_period = 'weekly' THEN CURRENT_DATE - 6 ELSE CURRENT_DATE END + AND date <= CURRENT_DATE; + + IF p_base_limit = -1 THEN + consumed_from := 'unlimited'; + ELSIF current_count < p_base_limit THEN + consumed_from := 'base'; + ELSE + UPDATE public.subscriptions + SET overage_credits = overage_credits - 1, + updated_at = now() + WHERE user_id = p_user_id + AND overage_credits > 0 + RETURNING overage_credits INTO new_overage; + + IF NOT FOUND THEN + RETURN jsonb_build_object( + 'allowed', false, + 'reservation_id', NULL, + 'status', 'denied', + 'current', current_count, + 'limit', p_base_limit, + 'period', p_period, + 'tier', subscription_tier, + 'overage_credits', 0, + 'consumed_from', 'none' + ); + END IF; + + overage := new_overage; + consumed_from := 'overage'; + END IF; + + INSERT INTO public.daily_usage(user_id, date, feature, count) + VALUES (p_user_id, CURRENT_DATE, p_feature, 1) + ON CONFLICT (user_id, date, feature) + DO UPDATE SET count = public.daily_usage.count + 1; + + current_count := current_count + 1; + INSERT INTO public.llm_quota_reservations( + id, user_id, feature, consumed_from, tier, quota_period, quota_limit, + current_count, overage_after, lease_expires_at + ) VALUES ( + p_reservation_id, p_user_id, p_feature, consumed_from, subscription_tier, p_period, p_base_limit, + current_count, overage, now() + interval '10 minutes' + ); + + RETURN jsonb_build_object( + 'allowed', true, + 'reservation_id', p_reservation_id, + 'status', 'reserved', + 'current', current_count, + 'limit', p_base_limit, + 'period', p_period, + 'tier', subscription_tier, + 'overage_credits', overage, + 'consumed_from', consumed_from + ); +END; +$$; + +CREATE OR REPLACE FUNCTION public.finalize_llm_quota( + p_reservation_id uuid, + p_succeeded boolean +) RETURNS jsonb +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = public, pg_temp +AS $$ +DECLARE + reservation public.llm_quota_reservations%ROWTYPE; + final_status text; +BEGIN + IF p_reservation_id IS NULL OR p_succeeded IS NULL THEN + RAISE EXCEPTION 'invalid_llm_quota_finalize' USING ERRCODE = '22023'; + END IF; + + SELECT * INTO reservation + FROM public.llm_quota_reservations + WHERE id = p_reservation_id; + IF NOT FOUND THEN + RAISE EXCEPTION 'llm_quota_reservation_not_found' USING ERRCODE = 'P0002'; + END IF; + + PERFORM pg_advisory_xact_lock(hashtextextended(reservation.user_id::text || ':' || reservation.feature, 20260928)); + SELECT * INTO reservation + FROM public.llm_quota_reservations + WHERE id = p_reservation_id + FOR UPDATE; + + IF reservation.status <> 'reserved' THEN + RETURN jsonb_build_object('reservation_id', reservation.id, 'status', reservation.status); + END IF; + + IF p_succeeded THEN + final_status := 'completed'; + ELSE + UPDATE public.daily_usage + SET count = greatest(count - 1, 0) + WHERE user_id = reservation.user_id + AND date = reservation.usage_date + AND feature = reservation.feature; + + IF reservation.consumed_from = 'overage' THEN + UPDATE public.subscriptions + SET overage_credits = overage_credits + 1, + updated_at = now() + WHERE user_id = reservation.user_id; + END IF; + final_status := 'released'; + END IF; + + UPDATE public.llm_quota_reservations + SET status = final_status, + finalized_at = now(), + release_reason = CASE WHEN p_succeeded THEN NULL ELSE 'provider_failed' END + WHERE id = reservation.id; + + RETURN jsonb_build_object('reservation_id', reservation.id, 'status', final_status); +END; +$$; + +REVOKE ALL ON FUNCTION public.reserve_llm_quota(uuid, uuid, text, integer, text) FROM PUBLIC, anon, authenticated; +REVOKE ALL ON FUNCTION public.finalize_llm_quota(uuid, boolean) FROM PUBLIC, anon, authenticated; +GRANT EXECUTE ON FUNCTION public.reserve_llm_quota(uuid, uuid, text, integer, text) TO service_role; +GRANT EXECUTE ON FUNCTION public.finalize_llm_quota(uuid, boolean) TO service_role; + +COMMENT ON TABLE public.llm_quota_reservations IS + 'Service-only leases that reserve LLM quota before the Anthropic call in llm-proxy and refund failed or expired generations.';