From 24c1f733b388a63632f6b1413180bf120eaae1c4 Mon Sep 17 00:00:00 2001 From: Riccardo Giorato Date: Mon, 5 Jan 2026 23:06:30 +0100 Subject: [PATCH] Add content policy checks and rate limiting --- app/api/add-page/route.ts | 57 ++++++++++++++-------- app/api/generate-comic/route.ts | 57 ++++++++++++++-------- components/editor/generate-page-modal.tsx | 17 +++++-- components/landing/comic-creation-form.tsx | 16 ++++-- lib/db-actions.ts | 4 ++ lib/utils.ts | 14 ++++++ 6 files changed, 117 insertions(+), 48 deletions(-) diff --git a/app/api/add-page/route.ts b/app/api/add-page/route.ts index f79a6d9..1a5f2fc 100644 --- a/app/api/add-page/route.ts +++ b/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, diff --git a/app/api/generate-comic/route.ts b/app/api/generate-comic/route.ts index 522fe64..3270c62 100644 --- a/app/api/generate-comic/route.ts +++ b/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 } : { diff --git a/components/editor/generate-page-modal.tsx b/components/editor/generate-page-modal.tsx index 3d68852..fce75cb 100644 --- a/components/editor/generate-page-modal.tsx +++ b/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", - description: - error instanceof Error - ? error.message - : "Failed to generate page. Please try again.", + title, + description: errorMessage, variant: "destructive", duration: 4000, }); setIsGenerating(false); + throw error; // Re-throw so the parent handler knows generation failed } }; diff --git a/components/landing/comic-creation-form.tsx b/components/landing/comic-creation-form.tsx index 297aa2e..03d3ef8 100644 --- a/components/landing/comic-creation-form.tsx +++ b/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", - description: - error instanceof Error - ? error.message - : "Failed to create comic. Please try again.", + title, + description: errorMessage, variant: "destructive", duration: 4000, }); diff --git a/lib/db-actions.ts b/lib/db-actions.ts index 993d563..d21745e 100644 --- a/lib/db-actions.ts +++ b/lib/db-actions.ts @@ -145,4 +145,8 @@ export async function getNextPageNumber(storyId: string): Promise { export async function deletePage(pageId: string): Promise { await db.delete(pages).where(eq(pages.id, pageId)); +} + +export async function deleteStory(storyId: string): Promise { + await db.delete(stories).where(eq(stories.id, storyId)); } \ No newline at end of file diff --git a/lib/utils.ts b/lib/utils.ts index df18e5a..145aded 100644 --- a/lib/utils.ts +++ b/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."; +}