fix(llm-proxy): reserve quota before the provider call and stop cutting off long streams
This commit is contained in:
Yun Chan
parent
2f94d24c99
commit
b306034bfc
7 changed files with 1489 additions and 307 deletions
112
server/supabase/functions/_shared/quota-reservation.test.ts
Normal file
112
server/supabase/functions/_shared/quota-reservation.test.ts
Normal file
|
|
@ -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')
|
||||
})
|
||||
|
|
@ -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<string, unknown>
|
||||
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<string, unknown>).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<string, unknown>
|
||||
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<typeof createClient>,
|
||||
): Promise<'completed' | 'released'> {
|
||||
): Promise<QuotaReservationFinalStatus> {
|
||||
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<string, unknown>).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<typeof createClient>,
|
||||
): Promise<QuotaReservation> {
|
||||
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<typeof createClient>,
|
||||
): Promise<QuotaReservationFinalStatus> {
|
||||
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<typeof createClient>,
|
||||
): Promise<Tier> {
|
||||
return (await readSubscriptionQuotaState(userId, serviceRoleClient)).tier
|
||||
}
|
||||
|
||||
/** 기간 창(daily=오늘, weekly=오늘-6..오늘) 안의 daily_usage 합계. */
|
||||
async function readWindowUsage(
|
||||
userId: string,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue