fix(quotas): unifier le décompte IA (BYOK, rollback) et combler les fuites
Centralise la réserve via ai-quota, corrige admin unavailable (-1), brancher les routes sans quota et le host-pays brainstorm, avec usage-meter élargi, noms de clusters, MCP et ajustements dashboard/insights. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
135
memento-note/lib/ai-quota.ts
Normal file
135
memento-note/lib/ai-quota.ts
Normal file
@@ -0,0 +1,135 @@
|
||||
import { willUseByokForLane } from '@/lib/ai/provider-for-user'
|
||||
import type { AiFeatureLane } from '@/lib/ai/router'
|
||||
import { getSystemConfig } from '@/lib/config'
|
||||
import {
|
||||
reserveUsageOrThrow,
|
||||
releaseUsage,
|
||||
checkSessionEntitlementOrThrow,
|
||||
QuotaExceededError,
|
||||
QuotaServiceUnavailableError,
|
||||
} from '@/lib/entitlements'
|
||||
import type { FeatureName } from '@/lib/quota-utils'
|
||||
import { emitAiUsageChanged } from '@/lib/ai-usage-sync'
|
||||
|
||||
export type ReserveAiUsageOptions = {
|
||||
lane?: AiFeatureLane
|
||||
/** Defaults to userId — host-pays brainstorm bills this account. */
|
||||
billingUserId?: string
|
||||
}
|
||||
|
||||
/** Reserve one AI unit; skips when BYOK will be used for the lane. */
|
||||
export async function reserveAiUsageOrThrow(
|
||||
userId: string,
|
||||
feature: FeatureName,
|
||||
options?: ReserveAiUsageOptions,
|
||||
): Promise<{ usedByok: boolean; billedUserId: string }> {
|
||||
const billedUserId = options?.billingUserId ?? userId
|
||||
const lane = options?.lane ?? 'chat'
|
||||
const config = await getSystemConfig()
|
||||
const { usedByok } = await willUseByokForLane(lane, config, billedUserId)
|
||||
if (!usedByok) {
|
||||
await reserveUsageOrThrow(billedUserId, feature)
|
||||
}
|
||||
return { usedByok, billedUserId }
|
||||
}
|
||||
|
||||
/** Host-pays: reserve on session owner with guest metadata on 402. */
|
||||
export async function reserveSessionAiUsageOrThrow(
|
||||
billingOwnerId: string,
|
||||
triggeredByUserId: string,
|
||||
isGuestActor: boolean,
|
||||
feature: FeatureName,
|
||||
options?: { lane?: AiFeatureLane },
|
||||
): Promise<{ usedByok: boolean }> {
|
||||
const lane = options?.lane ?? 'chat'
|
||||
const config = await getSystemConfig()
|
||||
const { usedByok } = await willUseByokForLane(lane, config, billingOwnerId)
|
||||
if (!usedByok) {
|
||||
await checkSessionEntitlementOrThrow(
|
||||
billingOwnerId,
|
||||
triggeredByUserId,
|
||||
isGuestActor,
|
||||
feature,
|
||||
)
|
||||
}
|
||||
return { usedByok }
|
||||
}
|
||||
|
||||
/** Run fn under quota; releases the unit if fn throws (unless BYOK). */
|
||||
export async function withAiQuota<T>(
|
||||
userId: string,
|
||||
feature: FeatureName,
|
||||
fn: () => Promise<T>,
|
||||
options?: ReserveAiUsageOptions,
|
||||
): Promise<T> {
|
||||
const { usedByok, billedUserId } = await reserveAiUsageOrThrow(userId, feature, options)
|
||||
try {
|
||||
const result = await fn()
|
||||
if (!usedByok) {
|
||||
emitAiUsageChanged()
|
||||
}
|
||||
return result
|
||||
} catch (err) {
|
||||
if (!usedByok) {
|
||||
await releaseUsage(billedUserId, feature)
|
||||
}
|
||||
throw err
|
||||
}
|
||||
}
|
||||
|
||||
/** Host-pays collaborative flows (brainstorm) with rollback on failure. */
|
||||
export async function withSessionAiQuota<T>(
|
||||
billingOwnerId: string,
|
||||
triggeredByUserId: string,
|
||||
isGuestActor: boolean,
|
||||
feature: FeatureName,
|
||||
fn: () => Promise<T>,
|
||||
options?: { lane?: AiFeatureLane },
|
||||
): Promise<T> {
|
||||
const { usedByok } = await reserveSessionAiUsageOrThrow(
|
||||
billingOwnerId,
|
||||
triggeredByUserId,
|
||||
isGuestActor,
|
||||
feature,
|
||||
options,
|
||||
)
|
||||
try {
|
||||
const result = await fn()
|
||||
if (!usedByok) {
|
||||
emitAiUsageChanged()
|
||||
}
|
||||
return result
|
||||
} catch (err) {
|
||||
if (!usedByok) {
|
||||
await releaseUsage(billingOwnerId, feature)
|
||||
}
|
||||
throw err
|
||||
}
|
||||
}
|
||||
|
||||
export function isQuotaError(err: unknown): err is QuotaExceededError {
|
||||
return err instanceof QuotaExceededError
|
||||
}
|
||||
|
||||
export function quotaExceededJson(err: QuotaExceededError) {
|
||||
return err.toJSON()
|
||||
}
|
||||
|
||||
export function quotaErrorStatus() {
|
||||
return 402
|
||||
}
|
||||
|
||||
/** Standard Next.js / Response handlers for quota errors. */
|
||||
export function handleQuotaHttpError(err: unknown): Response | null {
|
||||
if (err instanceof QuotaExceededError) {
|
||||
return Response.json(err.toJSON(), { status: 402 })
|
||||
}
|
||||
if (err instanceof QuotaServiceUnavailableError) {
|
||||
return Response.json({ error: err.code }, { status: 503 })
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
export function notifyAiUsageChanged(): void {
|
||||
emitAiUsageChanged()
|
||||
}
|
||||
Reference in New Issue
Block a user