import { getApproximateQuotaBucket, type QuotaBucket } from "./quota.js"; export type GenerationMode = "text_to_image" | "image_to_image"; export type GenerationRequestStatus = | "queued" | "running" | "succeeded" | "failed" | "canceled"; export interface ActiveSubscriptionContext { subscriptionId: string; planId: string; monthlyRequestLimit: number; usedSuccessfulRequests: number; } export interface GenerationRequestRecord { id: string; userId: string; mode: GenerationMode; status: GenerationRequestStatus; providerModel: string; prompt: string; sourceImageKey?: string; resolutionPreset: string; batchSize: number; imageStrength?: number; idempotencyKey?: string; terminalErrorCode?: string; terminalErrorText?: string; requestedAt: Date; startedAt?: Date; completedAt?: Date; createdAt: Date; updatedAt: Date; } export interface CreateGenerationRequestInput { userId: string; mode: GenerationMode; providerModel: string; prompt: string; sourceImageKey?: string; resolutionPreset: string; batchSize: number; imageStrength?: number; idempotencyKey?: string; } export interface CreateGenerationRequestResult { request: GenerationRequestRecord; reusedExistingRequest: boolean; approximateQuotaBucket: QuotaBucket; } export interface CreateGenerationRequestDeps { findReusableRequest( userId: string, idempotencyKey: string, ): Promise; findActiveSubscriptionContext(userId: string): Promise; createGenerationRequest( input: CreateGenerationRequestInput, ): Promise; } export interface SuccessfulGenerationRecord { request: GenerationRequestRecord; quotaConsumed: boolean; } export interface MarkGenerationSucceededDeps { getGenerationRequest(requestId: string): Promise; markGenerationSucceeded(requestId: string): Promise; } export class GenerationRequestError extends Error { readonly code: | "missing_active_subscription" | "quota_exhausted" | "invalid_prompt" | "invalid_batch_size" | "missing_source_image" | "unexpected_source_image" | "missing_image_strength" | "unexpected_image_strength" | "request_not_found" | "request_not_completable"; constructor(code: GenerationRequestError["code"], message: string) { super(message); this.code = code; } } export async function createGenerationRequest( deps: CreateGenerationRequestDeps, input: CreateGenerationRequestInput, ): Promise { validateGenerationRequestInput(input); if (input.idempotencyKey) { const existing = await deps.findReusableRequest(input.userId, input.idempotencyKey); if (existing) { const subscription = await deps.findActiveSubscriptionContext(input.userId); const approximateQuotaBucket = subscription ? getApproximateQuotaBucket({ used: subscription.usedSuccessfulRequests, limit: subscription.monthlyRequestLimit, }) : 0; return { request: existing, reusedExistingRequest: true, approximateQuotaBucket, }; } } const subscription = await deps.findActiveSubscriptionContext(input.userId); if (!subscription) { throw new GenerationRequestError( "missing_active_subscription", "An active subscription is required before creating generation requests.", ); } if (subscription.usedSuccessfulRequests >= subscription.monthlyRequestLimit) { throw new GenerationRequestError( "quota_exhausted", "The current billing cycle has no remaining successful generation quota.", ); } const request = await deps.createGenerationRequest(input); return { request, reusedExistingRequest: false, approximateQuotaBucket: getApproximateQuotaBucket({ used: subscription.usedSuccessfulRequests, limit: subscription.monthlyRequestLimit, }), }; } export async function markGenerationRequestSucceeded( deps: MarkGenerationSucceededDeps, requestId: string, ): Promise { const request = await deps.getGenerationRequest(requestId); if (!request) { throw new GenerationRequestError( "request_not_found", `Generation request ${requestId} was not found.`, ); } if (request.status === "failed" || request.status === "canceled") { throw new GenerationRequestError( "request_not_completable", `Generation request ${requestId} is terminal and cannot succeed.`, ); } return deps.markGenerationSucceeded(requestId); } function validateGenerationRequestInput(input: CreateGenerationRequestInput): void { if (input.prompt.trim().length === 0) { throw new GenerationRequestError( "invalid_prompt", "Prompt must not be empty after trimming.", ); } if (!Number.isInteger(input.batchSize) || input.batchSize <= 0) { throw new GenerationRequestError( "invalid_batch_size", "Batch size must be a positive integer.", ); } if (input.mode === "image_to_image") { if (!input.sourceImageKey) { throw new GenerationRequestError( "missing_source_image", "Image-to-image requests require a source image key.", ); } if (input.imageStrength === undefined) { throw new GenerationRequestError( "missing_image_strength", "Image-to-image requests require image strength.", ); } return; } if (input.sourceImageKey) { throw new GenerationRequestError( "unexpected_source_image", "Text-to-image requests must not include a source image key.", ); } if (input.imageStrength !== undefined) { throw new GenerationRequestError( "unexpected_image_strength", "Text-to-image requests must not include image strength.", ); } }