nroxy/packages/domain/src/generation.ts

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.",
    );
  }
}