Initial import

packages/domain/src/generation.ts

@@ -0,0 +1,221 @@
+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<GenerationRequestRecord | null>;
+  findActiveSubscriptionContext(userId: string): Promise<ActiveSubscriptionContext | null>;
+  createGenerationRequest(
+    input: CreateGenerationRequestInput,
+  ): Promise<GenerationRequestRecord>;
+}
+
+export interface SuccessfulGenerationRecord {
+  request: GenerationRequestRecord;
+  quotaConsumed: boolean;
+}
+
+export interface MarkGenerationSucceededDeps {
+  getGenerationRequest(requestId: string): Promise<GenerationRequestRecord | null>;
+  markGenerationSucceeded(requestId: string): Promise<SuccessfulGenerationRecord>;
+}
+
+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<CreateGenerationRequestResult> {
+  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<SuccessfulGenerationRecord> {
+  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.",
+    );
+  }
+}