Add content policy checks and rate limiting

app/api/add-page/route.ts

@@ -10,10 +10,12 @@ import {
   getNextPageNumber,
   getStoryWithPagesBySlug,
   getLastPageImage,
+  deletePage,
 } from "@/lib/db-actions";
 import { freeTierRateLimit } from "@/lib/rate-limit";
 import { uploadImageToS3 } from "@/lib/s3-upload";
 import { buildComicPrompt } from "@/lib/prompt";
+import { isContentPolicyViolation, getContentPolicyErrorMessage } from "@/lib/utils";
 
 const NEW_MODEL = false;
 
@@ -63,25 +65,7 @@ export async function POST(request: NextRequest) {
       return NextResponse.json({ error: "Unauthorized" }, { status: 403 });
     }
 
-    // Apply rate limiting for free tier
+
-    const hasApiKey = request.headers.get("x-api-key");
-    if (!hasApiKey) {
-      const { success, reset } = await freeTierRateLimit.limit(userId);
-      if (!success) {
-        const resetDate = new Date(reset);
-        const timeUntilReset = Math.ceil(
-          (reset - Date.now()) / (1000 * 60 * 60 * 24)
-        );
-        return NextResponse.json(
-          {
-            error: `Free tier limit reached. You can generate 1 comic per week. Try again in ${timeUntilReset} day(s), or provide your own API key.`,
-            resetDate: resetDate.toISOString(),
-            isRateLimited: true,
-          },
-          { status: 429 }
-        );
-      }
-    }
 
     let page;
     let pageNumber;
@@ -172,6 +156,30 @@ export async function POST(request: NextRequest) {
     } catch (error) {
       console.error("Together AI API error:", error);
 
+      // Clean up DB records if generation failed due to content policy
+      try {
+        if (error instanceof Error && error.message && error.message.includes("NO_IMAGE")) {
+          if (isRedraw) {
+            // For redraw, we don't delete the page, just don't update it
+          } else {
+            // For new page, delete the page that was created
+            await deletePage(page.id);
+          }
+        }
+      } catch (cleanupError) {
+        console.error("Error cleaning up DB on image generation failure:", cleanupError);
+      }
+
+      if (error instanceof Error && error.message && isContentPolicyViolation(error.message)) {
+        return NextResponse.json(
+          {
+            error: getContentPolicyErrorMessage(),
+            errorType: "content_policy",
+          },
+          { status: 400 }
+        );
+      }
+
       if (error instanceof Error && "status" in error) {
         const status = (error as any).status;
         if (status === 402) {
@@ -215,6 +223,17 @@ export async function POST(request: NextRequest) {
 
     await updatePage(page.id, s3ImageUrl);
 
+    // Apply rate limiting for free tier after successful generation
+    const hasApiKey = request.headers.get("x-api-key");
+    if (!hasApiKey) {
+      try {
+        await freeTierRateLimit.limit(userId);
+      } catch (rateLimitError) {
+        console.error("Error applying rate limit after successful generation:", rateLimitError);
+        // Don't fail the request if rate limiting fails, just log it
+      }
+    }
+
     return NextResponse.json({
       imageUrl: s3ImageUrl,
       pageId: page.id,

app/api/generate-comic/route.ts

@@ -10,11 +10,14 @@ import {
   getStoryById,
   getLastPageImage,
   getStoryCharacterImages,
+  deletePage,
+  deleteStory,
 } from "@/lib/db-actions";
 import { freeTierRateLimit } from "@/lib/rate-limit";
 import { COMIC_STYLES } from "@/lib/constants";
 import { uploadImageToS3 } from "@/lib/s3-upload";
 import { buildComicPrompt } from "@/lib/prompt";
+import { isContentPolicyViolation, getContentPolicyErrorMessage } from "@/lib/utils";
 
 const NEW_MODEL = false;
 
@@ -61,25 +64,6 @@ export async function POST(request: NextRequest) {
     const isUsingFreeTier = !apiKey;
 
     if (isUsingFreeTier) {
-      // Using free tier - apply rate limiting
-      const { success, reset } = await freeTierRateLimit.limit(userId);
-
-      if (!success) {
-        const resetDate = new Date(reset);
-        const timeUntilReset = Math.ceil(
-          (reset - Date.now()) / (1000 * 60 * 60 * 24)
-        ); // days
-
-        return NextResponse.json(
-          {
-            error: `Free tier limit reached. You can generate 1 comic per week. Try again in ${timeUntilReset} day(s), or provide your own API key for unlimited access.`,
-            resetDate: resetDate.toISOString(),
-            isRateLimited: true,
-          },
-          { status: 429 }
-        );
-      }
-
       // Use default API key for free tier
       finalApiKey = process.env.TOGETHER_API_KEY_DEFAULT;
       if (!finalApiKey) {
@@ -92,6 +76,8 @@ export async function POST(request: NextRequest) {
       }
     }
 
+
+
     let page;
     let story;
     let referenceImages: string[] = [];
@@ -251,6 +237,29 @@ Only return the JSON, no other text.`;
     } catch (error) {
       console.error("Together AI API error:", error);
 
+      // Clean up DB records if generation failed
+      try {
+        if (!storyId) {
+          // New story failed
+          await deleteStory(story!.id);
+        } else {
+          // Continuation failed
+          await deletePage(page.id);
+        }
+      } catch (cleanupError) {
+        console.error("Error cleaning up DB on image generation failure:", cleanupError);
+      }
+
+      if (error instanceof Error && error.message && isContentPolicyViolation(error.message)) {
+        return NextResponse.json(
+          {
+            error: getContentPolicyErrorMessage(),
+            errorType: "content_policy",
+          },
+          { status: 400 }
+        );
+      }
+
       if (error instanceof Error && "status" in error) {
         const status = (error as any).status;
         if (status === 402) {
@@ -334,6 +343,16 @@ Only return the JSON, no other text.`;
       );
     }
 
+    // Apply rate limiting for free tier after successful generation
+    if (isUsingFreeTier) {
+      try {
+        await freeTierRateLimit.limit(userId);
+      } catch (rateLimitError) {
+        console.error("Error applying rate limit after successful generation:", rateLimitError);
+        // Don't fail the request if rate limiting fails, just log it
+      }
+    }
+
     const responseData = storyId
       ? { imageUrl: s3ImageUrl, pageId: page.id, pageNumber: page.pageNumber }
       : {

components/editor/generate-page-modal.tsx

@@ -14,6 +14,7 @@ import { useToast } from "@/hooks/use-toast";
 import { useKeyboardShortcut } from "@/hooks/use-keyboard-shortcut";
 import { validateFileForUpload, generateFilePreview } from "@/lib/file-utils";
 import { useS3Upload } from "next-s3-upload";
+import { isContentPolicyViolation } from "@/lib/utils";
 
 interface CharacterItem {
   url: string;
@@ -268,16 +269,22 @@ export function GeneratePageModal({
       });
     } catch (error) {
       console.error("Error generating page:", error);
+      const errorMessage =
+        error instanceof Error
+          ? error.message
+          : "Failed to generate page. Please try again.";
+      let title = "Generation failed";
+      if (isContentPolicyViolation(errorMessage)) {
+        title = "Content policy violation";
+      }
       toast({
-        title: "Generation failed",
+        title,
-        description:
+        description: errorMessage,
-          error instanceof Error
-            ? error.message
-            : "Failed to generate page. Please try again.",
         variant: "destructive",
         duration: 4000,
       });
       setIsGenerating(false);
+      throw error; // Re-throw so the parent handler knows generation failed
     }
   };
 

components/landing/comic-creation-form.tsx

@@ -10,6 +10,7 @@ import { useAuth, SignInButton } from "@clerk/nextjs";
 import { COMIC_STYLES } from "@/lib/constants";
 import { useKeyboardShortcut } from "@/hooks/use-keyboard-shortcut";
 import { useApiKey } from "@/hooks/use-api-key";
+import { isContentPolicyViolation } from "@/lib/utils";
 
 interface ComicCreationFormProps {
   prompt: string;
@@ -171,12 +172,17 @@ export function ComicCreationForm({
       router.push(`/story/${result.storySlug}`);
     } catch (error) {
       console.error("Error creating comic:", error);
+      const errorMessage =
+        error instanceof Error
+          ? error.message
+          : "Failed to create comic. Please try again.";
+      let title = "Creation failed";
+      if (isContentPolicyViolation(errorMessage)) {
+        title = "Content policy violation";
+      }
       toast({
-        title: "Creation failed",
+        title,
-        description:
+        description: errorMessage,
-          error instanceof Error
-            ? error.message
-            : "Failed to create comic. Please try again.",
         variant: "destructive",
         duration: 4000,
       });

lib/db-actions.ts

@@ -145,4 +145,8 @@ export async function getNextPageNumber(storyId: string): Promise<number> {
 
 export async function deletePage(pageId: string): Promise<void> {
   await db.delete(pages).where(eq(pages.id, pageId));
+}
+
+export async function deleteStory(storyId: string): Promise<void> {
+  await db.delete(stories).where(eq(stories.id, storyId));
 }

lib/utils.ts

@@ -7,3 +7,17 @@ export function cn(...inputs: ClassValue[]) {
 
 export const TOGETHER_LINK =
   "https://togetherai.link/?utm_source=make-comics&utm_medium=referral&utm_campaign=example-app";
+
+export function isContentPolicyViolation(errorMessage: string): boolean {
+  return (
+    errorMessage.includes("content policy") ||
+    errorMessage.includes("Invalid content detected") ||
+    errorMessage.includes("content moderation") ||
+    errorMessage.includes("flagged and rejected") ||
+    errorMessage.includes("NO_IMAGE")
+  );
+}
+
+export function getContentPolicyErrorMessage(): string {
+  return "Unable to generate image due to content policy. Please try a different prompt.";
+}