fix(llm-proxy): reserve quota before the provider call and stop cutting off long streams
This commit is contained in:
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
|
consumedFrom: QuotaConsumedFrom
|
||||||
}
|
}
|
||||||
|
|
||||||
export interface SttQuotaReservation {
|
/** reserve_*_quota RPC 결과 — 공급자 호출 전에 한 단위를 잡아 두는 리스(lease). */
|
||||||
|
export interface QuotaReservation {
|
||||||
allowed: boolean
|
allowed: boolean
|
||||||
reservationId: string | null
|
reservationId: string | null
|
||||||
status: 'reserved' | 'completed' | 'released' | 'denied'
|
status: 'reserved' | 'completed' | 'released' | 'denied'
|
||||||
|
|
@ -72,6 +73,59 @@ export interface SttQuotaReservation {
|
||||||
consumedFrom: 'base' | 'overage' | 'unlimited' | 'none'
|
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(
|
export async function reserveSttQuota(
|
||||||
userId: string,
|
userId: string,
|
||||||
reservationId: string,
|
reservationId: string,
|
||||||
|
|
@ -84,38 +138,14 @@ export async function reserveSttQuota(
|
||||||
if (error || !data || typeof data !== 'object' || Array.isArray(data)) {
|
if (error || !data || typeof data !== 'object' || Array.isArray(data)) {
|
||||||
throw new Error('Failed to reserve STT quota.')
|
throw new Error('Failed to reserve STT quota.')
|
||||||
}
|
}
|
||||||
const result = data as Record<string, unknown>
|
return parseQuotaReservationResponse(data, 'STT')
|
||||||
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'],
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
export async function finalizeSttQuota(
|
export async function finalizeSttQuota(
|
||||||
reservationId: string,
|
reservationId: string,
|
||||||
succeeded: boolean,
|
succeeded: boolean,
|
||||||
serviceRoleClient: ReturnType<typeof createClient>,
|
serviceRoleClient: ReturnType<typeof createClient>,
|
||||||
): Promise<'completed' | 'released'> {
|
): Promise<QuotaReservationFinalStatus> {
|
||||||
const { data, error } = await serviceRoleClient.rpc('finalize_stt_quota', {
|
const { data, error } = await serviceRoleClient.rpc('finalize_stt_quota', {
|
||||||
p_reservation_id: reservationId,
|
p_reservation_id: reservationId,
|
||||||
p_succeeded: succeeded,
|
p_succeeded: succeeded,
|
||||||
|
|
@ -123,11 +153,50 @@ export async function finalizeSttQuota(
|
||||||
if (error || !data || typeof data !== 'object' || Array.isArray(data)) {
|
if (error || !data || typeof data !== 'object' || Array.isArray(data)) {
|
||||||
throw new Error('Failed to finalize STT quota.')
|
throw new Error('Failed to finalize STT quota.')
|
||||||
}
|
}
|
||||||
const status = (data as Record<string, unknown>).status
|
return parseQuotaFinalizeResponse(data, 'STT')
|
||||||
if (status !== 'completed' && status !== 'released') {
|
}
|
||||||
throw new Error('Invalid STT quota finalization response.')
|
|
||||||
|
/**
|
||||||
|
* 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 {
|
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 합계. */
|
/** 기간 창(daily=오늘, weekly=오늘-6..오늘) 안의 daily_usage 합계. */
|
||||||
async function readWindowUsage(
|
async function readWindowUsage(
|
||||||
userId: string,
|
userId: string,
|
||||||
|
|
|
||||||
378
server/supabase/functions/llm-proxy/handler.test.ts
Normal file
378
server/supabase/functions/llm-proxy/handler.test.ts
Normal file
|
|
@ -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<T>(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<void>((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<string, 'reserved' | 'completed' | 'released'>()
|
||||||
|
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<QuotaReservation> {
|
||||||
|
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<typeof setTimeout> | null = null
|
||||||
|
let wake: (() => void) | null = null
|
||||||
|
let streamController: ReadableStreamDefaultController<Uint8Array> | 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<Uint8Array>({
|
||||||
|
start(controller) {
|
||||||
|
streamController = controller
|
||||||
|
},
|
||||||
|
async pull(controller) {
|
||||||
|
if (state.aborted) return
|
||||||
|
if (plan.stallAfter !== undefined && index >= plan.stallAfter) {
|
||||||
|
await new Promise<void>((resolve) => {
|
||||||
|
wake = resolve
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
await new Promise<void>((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<void>((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<Response>((_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')
|
||||||
|
})
|
||||||
336
server/supabase/functions/llm-proxy/handler.ts
Normal file
336
server/supabase/functions/llm-proxy/handler.ts
Normal file
|
|
@ -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<Record<Tier, readonly string[]>> = {
|
||||||
|
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<Record<Tier, string>> = {
|
||||||
|
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<Tier>
|
||||||
|
reserve(
|
||||||
|
userId: string,
|
||||||
|
reservationId: string,
|
||||||
|
feature: QuotaFeature,
|
||||||
|
baseLimit: number,
|
||||||
|
period: QuotaPeriod,
|
||||||
|
): Promise<QuotaReservation>
|
||||||
|
finalize(reservationId: string, succeeded: boolean): Promise<QuotaReservationFinalStatus>
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Request passed to the provider port. */
|
||||||
|
export interface ProviderRequest {
|
||||||
|
method: 'POST'
|
||||||
|
headers: Record<string, string>
|
||||||
|
body: string
|
||||||
|
signal: AbortSignal
|
||||||
|
}
|
||||||
|
|
||||||
|
/** Port over the provider HTTP call (global fetch in production). */
|
||||||
|
export type ProviderFetch = (url: string, init: ProviderRequest) => Promise<Response>
|
||||||
|
|
||||||
|
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<string>
|
||||||
|
fetchProvider?: ProviderFetch
|
||||||
|
newReservationId?(): string
|
||||||
|
timeouts?: Partial<LlmProviderTimeouts>
|
||||||
|
}
|
||||||
|
|
||||||
|
const NO_STORE_JSON = { ...corsHeaders, 'Content-Type': 'application/json', 'Cache-Control': 'no-store' }
|
||||||
|
|
||||||
|
function json(status: number, body: Record<string, unknown>, 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<void> {
|
||||||
|
if (!resp.body || resp.bodyUsed) return
|
||||||
|
try {
|
||||||
|
await resp.body.cancel()
|
||||||
|
} catch {
|
||||||
|
// 이미 닫힌 스트림 — 무시
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function generationHeaders(generationId: string | null): Record<string, string> {
|
||||||
|
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<void> {
|
||||||
|
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<Response> {
|
||||||
|
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<Response> => {
|
||||||
|
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<void>) | 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' })
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -1,293 +1,50 @@
|
||||||
// server/supabase/functions/llm-proxy/index.ts
|
// 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+)
|
// Phase 3.2: 모델별 쿼터 (Haiku/Sonnet/Opus × Free/Pro/Pro+)
|
||||||
// 요청: application/json { messages, system?, max_tokens?, model? }
|
// 요청: application/json { messages, system?, max_tokens?, model? }
|
||||||
// 응답: JSON (non-stream) 또는 SSE (stream=true)
|
// 응답: JSON (non-stream) 또는 SSE (stream=true)
|
||||||
|
// 유스케이스와 쿼터/기한 규칙은 handler.ts, 공급자 기한은 provider-deadline.ts 에 있다.
|
||||||
|
|
||||||
import { corsHeaders, handleCorsPreflightRequest } from '../_shared/cors.ts'
|
import { requireUser } from '../_shared/auth.ts'
|
||||||
import { requireUser, authErrorResponse, type AuthError } from '../_shared/auth.ts'
|
|
||||||
import {
|
import {
|
||||||
checkQuota,
|
|
||||||
consumeQuota,
|
|
||||||
createServiceRoleClient,
|
createServiceRoleClient,
|
||||||
modelToQuotaKey,
|
finalizeLlmQuota,
|
||||||
getQuotaPolicy,
|
readSubscriptionTier,
|
||||||
type Tier,
|
reserveLlmQuota,
|
||||||
} from '../_shared/quota.ts'
|
} from '../_shared/quota.ts'
|
||||||
import {
|
import { GenerationReceiptError, parseGenerationReceiptId } from '../_shared/generation-receipt.ts'
|
||||||
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 { readProviderKey } from '../_shared/provider-key.ts'
|
import { readProviderKey } from '../_shared/provider-key.ts'
|
||||||
|
import { createLlmProxyHandler, type LlmQuotaStore } from './handler.ts'
|
||||||
|
|
||||||
/** 티어별 허용 모델 — free는 Haiku만, pro는 +Sonnet, pro_plus는 +Opus, team/enterprise는 전 모델 */
|
let serviceClient: ReturnType<typeof createServiceRoleClient> | null = null
|
||||||
const TIER_MODELS: Record<Tier, string[]> = {
|
function service(): ReturnType<typeof createServiceRoleClient> {
|
||||||
free: ['claude-haiku-4-5-20251001'],
|
serviceClient ??= createServiceRoleClient()
|
||||||
pro: ['claude-haiku-4-5-20251001', 'claude-sonnet-4-6', 'claude-opus-4-6'],
|
return serviceClient
|
||||||
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'],
|
|
||||||
}
|
}
|
||||||
|
|
||||||
const DEFAULT_MODEL: Record<Tier, string> = {
|
function supabaseLlmQuotaStore(client: ReturnType<typeof createServiceRoleClient>): LlmQuotaStore {
|
||||||
free: 'claude-haiku-4-5-20251001',
|
return {
|
||||||
pro: 'claude-sonnet-4-6',
|
readTier: (userId) => readSubscriptionTier(userId, client),
|
||||||
pro_plus: 'claude-sonnet-4-6',
|
reserve: (userId, reservationId, feature, baseLimit, period) =>
|
||||||
team: 'claude-sonnet-4-6',
|
reserveLlmQuota(userId, reservationId, feature, baseLimit, period, client),
|
||||||
enterprise: 'claude-sonnet-4-6',
|
finalize: (reservationId, succeeded) => finalizeLlmQuota(reservationId, succeeded, client),
|
||||||
}
|
|
||||||
|
|
||||||
async function issueGenerationReceipt(
|
|
||||||
serviceClient: ReturnType<typeof createServiceRoleClient>,
|
|
||||||
userId: string,
|
|
||||||
purpose: GenerationPurpose | null,
|
|
||||||
model: string,
|
|
||||||
): Promise<string | null> {
|
|
||||||
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<void> {
|
|
||||||
if (!resp.body || resp.bodyUsed) return
|
|
||||||
try {
|
|
||||||
await resp.body.cancel()
|
|
||||||
} catch {
|
|
||||||
// 이미 닫힌 스트림 — 무시
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
function generationHeaders(generationId: string | null): Record<string, string> {
|
Deno.serve(createLlmProxyHandler({
|
||||||
return generationId === null ? {} : { [GENERATION_ID_HEADER]: generationId }
|
authenticate: requireUser,
|
||||||
}
|
providerKey: () => readProviderKey('ANTHROPIC_API_KEY'),
|
||||||
|
quotaStore: () => supabaseLlmQuotaStore(service()),
|
||||||
Deno.serve(async (req: Request) => {
|
issueGenerationReceipt: async (userId, purpose, model) => {
|
||||||
const preflight = handleCorsPreflightRequest(req)
|
const { data, error } = await service().rpc('issue_content_generation_receipt_v1', {
|
||||||
if (preflight) return preflight
|
p_actor_id: userId,
|
||||||
|
p_purpose: purpose,
|
||||||
if (req.method !== 'POST') {
|
p_model: model,
|
||||||
return new Response(JSON.stringify({ error: 'Method not allowed' }), {
|
|
||||||
status: 405,
|
|
||||||
headers: { ...corsHeaders, 'Content-Type': 'application/json' },
|
|
||||||
})
|
})
|
||||||
}
|
if (error) {
|
||||||
|
console.error('Generation receipt issuance failed', { code: error.code ?? 'unknown' })
|
||||||
try {
|
throw new GenerationReceiptError('generation_receipt_unavailable')
|
||||||
const user = await requireUser(req)
|
|
||||||
const serviceClient = createServiceRoleClient()
|
|
||||||
|
|
||||||
let rawBody: unknown
|
|
||||||
try {
|
|
||||||
rawBody = await req.json()
|
|
||||||
} catch {
|
|
||||||
throw new LlmRequestError('Invalid JSON body')
|
|
||||||
}
|
}
|
||||||
const body = parseLlmRequest(rawBody)
|
return parseGenerationReceiptId(data)
|
||||||
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<ReturnType<typeof consumeQuota>>
|
|
||||||
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' },
|
|
||||||
})
|
|
||||||
}
|
|
||||||
})
|
|
||||||
|
|
|
||||||
226
server/supabase/functions/llm-proxy/provider-deadline.ts
Normal file
226
server/supabase/functions/llm-proxy/provider-deadline.ts
Normal file
|
|
@ -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<LlmProviderTimeouts> = 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<typeof setTimeout> | 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<void>
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* 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<Uint8Array>,
|
||||||
|
options: StreamRelayOptions,
|
||||||
|
): ReadableStream<Uint8Array> {
|
||||||
|
const reader = upstream.getReader()
|
||||||
|
const tracker = createSseCompletionTracker()
|
||||||
|
let settled = false
|
||||||
|
let totalTimer: ReturnType<typeof setTimeout> | null = null
|
||||||
|
let idleTimer: ReturnType<typeof setTimeout> | null = null
|
||||||
|
let rejectIdle: ((reason: unknown) => void) | null = null
|
||||||
|
let downstream: ReadableStreamDefaultController<Uint8Array> | 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<boolean> => {
|
||||||
|
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<void> => {
|
||||||
|
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<ReadableStreamReadResult<Uint8Array>> =>
|
||||||
|
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<Uint8Array>({
|
||||||
|
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<Uint8Array>
|
||||||
|
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')
|
||||||
|
},
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
@ -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.';
|
||||||
Loading…
Add table
Add a link
Reference in a new issue