feat: Add subscription checks for enhancing prompt and Posthog tracking

app/ai/components/ai-chat-form.tsx

@@ -9,6 +9,7 @@ import { Button } from "@/components/ui/button";
 import { useAIChatForm } from "@/hooks/use-ai-chat-form";
 import { useAIEnhancePrompt } from "@/hooks/use-ai-enhance-prompt";
 import { useGuards } from "@/hooks/use-guards";
+import { useSubscription } from "@/hooks/use-subscription";
 import { MAX_IMAGE_FILES } from "@/lib/constants";
 import { cn } from "@/lib/utils";
 import { AIPromptData } from "@/types/ai";
@@ -39,13 +40,15 @@ export function AIChatForm({
   } = useAIChatForm();
 
   const { checkValidSession, checkValidSubscription } = useGuards();
+  const { subscriptionStatus } = useSubscription();
+  const isPro = subscriptionStatus?.isSubscribed ?? false;
+  const hasFreeRequestsLeft = (subscriptionStatus?.requestsRemaining ?? 0) > 0;
 
   const { startEnhance, stopEnhance, enhancedPromptAsJsonContent, isEnhancingPrompt } =
     useAIEnhancePrompt();
 
   const handleEnhancePrompt = () => {
-    // TODO: Add subscription check, this should be a Pro only feature
-    if (!checkValidSession() || !checkValidSubscription()) return; // Act as an early return;
+    if (!checkValidSession() || !checkValidSubscription()) return;
 
     // Only send images that are not loading, and strip loading property
     const images = uploadedImages.filter((img) => !img.loading).map(({ url }) => ({ url }));
@@ -109,8 +112,7 @@ export function AIChatForm({
           </div>
 
           <div className="flex items-center gap-2">
-            {/* TODO: This should be a Pro only feature */}
-            {promptData?.content ? (
+            {(isPro || hasFreeRequestsLeft) && promptData?.content ? (
               <EnhancePromptButton
                 isEnhancing={isEnhancingPrompt}
                 onStart={handleEnhancePrompt}

app/api/enhance-prompt/route.ts

@@ -1,32 +1,38 @@
 import { ENHANCE_PROMPT_SYSTEM } from "@/lib/ai/prompts";
 import { baseProviderOptions, myProvider } from "@/lib/ai/providers";
+import { handleError } from "@/lib/error-response";
+import { requireSubscriptionOrFreeUsage } from "@/lib/subscription";
 import { AIPromptData } from "@/types/ai";
 import { buildUserContentPartsFromPromptData } from "@/utils/ai/message-converter";
 import { smoothStream, streamText } from "ai";
+import { NextRequest } from "next/server";
 
-export async function POST(req: Request) {
-  // TODO: Add session and subscription check, this should be a Pro only feature
-  // TODO: Record AI usage, providing the model id to `recordAIUsage` function
+export async function POST(req: NextRequest) {
+  try {
+    await requireSubscriptionOrFreeUsage(req);
 
-  const body = await req.json();
-  const { prompt: _prompt, promptData }: { prompt: string; promptData: AIPromptData } = body;
-  const userContentParts = buildUserContentPartsFromPromptData(promptData);
+    const body = await req.json();
+    const { prompt: _prompt, promptData }: { prompt: string; promptData: AIPromptData } = body;
+    const userContentParts = buildUserContentPartsFromPromptData(promptData);
 
-  const result = streamText({
-    system: ENHANCE_PROMPT_SYSTEM,
-    messages: [
-      {
-        role: "user",
-        content: userContentParts,
-      },
-    ],
-    model: myProvider.languageModel("prompt-enhancement"),
-    providerOptions: baseProviderOptions,
-    experimental_transform: smoothStream({
-      delayInMs: 10,
-      chunking: "word",
-    }),
-  });
+    const result = streamText({
+      system: ENHANCE_PROMPT_SYSTEM,
+      messages: [
+        {
+          role: "user",
+          content: userContentParts,
+        },
+      ],
+      model: myProvider.languageModel("prompt-enhancement"),
+      providerOptions: baseProviderOptions,
+      experimental_transform: smoothStream({
+        delayInMs: 10,
+        chunking: "word",
+      }),
+    });
 
-  return result.toUIMessageStreamResponse();
+    return result.toUIMessageStreamResponse();
+  } catch (error) {
+    return handleError(error, { route: "/api/enhance-prompt" });
+  }
 }

components/editor/ai/chat-input.tsx

@@ -7,6 +7,7 @@ import { useAIChatForm } from "@/hooks/use-ai-chat-form";
 import { useAIEnhancePrompt } from "@/hooks/use-ai-enhance-prompt";
 import { useChatContext } from "@/hooks/use-chat-context";
 import { useGuards } from "@/hooks/use-guards";
+import { useSubscription } from "@/hooks/use-subscription";
 import { usePostLoginAction } from "@/hooks/use-post-login-action";
 import { MAX_IMAGE_FILES } from "@/lib/constants";
 import { cn } from "@/lib/utils";
@@ -37,6 +38,9 @@ export function ChatInput({
 }: ChatInputProps) {
   const { messages, startNewChat } = useChatContext();
   const { checkValidSession, checkValidSubscription } = useGuards();
+  const { subscriptionStatus } = useSubscription();
+  const isPro = subscriptionStatus?.isSubscribed ?? false;
+  const hasFreeRequestsLeft = (subscriptionStatus?.requestsRemaining ?? 0) > 0;
 
   const {
     editorContentDraft,
@@ -64,8 +68,7 @@ export function ChatInput({
     useAIEnhancePrompt();
 
   const handleEnhancePrompt = () => {
-    // TODO: Add subscription check, this should be a Pro only feature
-    if (!checkValidSession() || !checkValidSubscription()) return; // Act as an early return;
+    if (!checkValidSession() || !checkValidSubscription()) return;
 
     // Only send images that are not loading, and strip loading property
     const images = uploadedImages.filter((img) => !img.loading).map(({ url }) => ({ url }));
@@ -158,8 +161,7 @@ export function ChatInput({
           </TooltipWrapper>
 
           <div className="flex items-center gap-2">
-            {/* TODO: This should be a Pro only feature */}
-            {promptData?.content ? (
+            {(isPro || hasFreeRequestsLeft) && promptData?.content ? (
               <EnhancePromptButton
                 isEnhancing={isEnhancingPrompt}
                 onStart={handleEnhancePrompt}

hooks/use-ai-enhance-prompt.ts

@@ -1,21 +1,36 @@
 "use client";
 
 import { toast } from "@/hooks/use-toast";
+import { parseAiSdkTransportError } from "@/lib/ai/parse-ai-sdk-transport-error";
 import { useAILocalDraftStore } from "@/store/ai-local-draft-store";
+import { useGetProDialogStore } from "@/store/get-pro-dialog-store";
 import { AIPromptData } from "@/types/ai";
 import { convertPromptDataToJSONContent } from "@/utils/ai/ai-prompt";
-import { parseAiSdkTransportError } from "@/lib/ai/parse-ai-sdk-transport-error";
 import { useCompletion } from "@ai-sdk/react";
 import { JSONContent } from "@tiptap/react";
+import posthog from "posthog-js";
 import { useCallback, useMemo, useRef } from "react";
 
 export function useAIEnhancePrompt() {
+  const { openGetProDialog } = useGetProDialogStore();
   const { complete, completion, isLoading, stop, setCompletion } = useCompletion({
     api: "/api/enhance-prompt",
     onError: (error) => {
       const defaultMessage = "Failed to enhance prompt. Please try again.";
       const normalized = parseAiSdkTransportError(error, defaultMessage);
 
+      try {
+        posthog.capture("ENHANCE_PROMPT_ERROR", {
+          message: normalized.message,
+          code: normalized.code,
+          status: normalized.status,
+        });
+      } catch {}
+
+      if (normalized.code === "SUBSCRIPTION_REQUIRED") {
+        openGetProDialog();
+      }
+
       toast({
         title: "An error occurred",
         description: normalized.message,
@@ -23,6 +38,14 @@ export function useAIEnhancePrompt() {
       });
     },
     onFinish: (_prompt, finalCompletion) => {
+      try {
+        const durationMs = startTimeRef.current ? Date.now() - startTimeRef.current : undefined;
+        posthog.capture("ENHANCE_PROMPT_FINISH", {
+          durationMs,
+          finalLength: finalCompletion?.length ?? 0,
+        });
+      } catch {}
+
       const promptData: AIPromptData = {
         content: finalCompletion,
         mentions: activeMentionsRef.current.map((m) => ({
@@ -37,6 +60,7 @@ export function useAIEnhancePrompt() {
   });
 
   const activeMentionsRef = useRef<Array<{ id: string; label: string }>>([]);
+  const startTimeRef = useRef<number | null>(null);
 
   const enhancedPromptAsJsonContent: JSONContent | undefined = useMemo(() => {
     if (!completion) return undefined;
@@ -55,10 +79,22 @@ export function useAIEnhancePrompt() {
     async (promptData: AIPromptData) => {
       const prompt = promptData?.content ?? "";
       if (!prompt?.trim()) return;
+
       if (isLoading) stop();
       setCompletion("");
+
+      startTimeRef.current = Date.now();
       activeMentionsRef.current =
         promptData?.mentions?.map((m) => ({ id: m.id, label: m.label })) ?? [];
+
+      try {
+        posthog.capture("ENHANCE_PROMPT_START", {
+          contentLength: prompt.length,
+          mentionCount: promptData?.mentions?.length ?? 0,
+          imageCount: promptData?.images?.length ?? 0,
+        });
+      } catch {}
+
       await complete(prompt, { body: { promptData } });
     },
     [complete, isLoading, stop, setCompletion]
@@ -66,9 +102,15 @@ export function useAIEnhancePrompt() {
 
   const stopEnhance = useCallback(() => {
     stop();
+
     if (enhancedPromptAsJsonContent) {
       useAILocalDraftStore.getState().setEditorContentDraft(enhancedPromptAsJsonContent);
     }
+
+    try {
+      const durationMs = startTimeRef.current ? Date.now() - startTimeRef.current : undefined;
+      posthog.capture("ENHANCE_PROMPT_CANCEL", { durationMs });
+    } catch {}
   }, [stop, enhancedPromptAsJsonContent]);
 
   return {

hooks/use-guards.ts

@@ -6,11 +6,12 @@ import { useSubscription } from "./use-subscription";
 
 export function useGuards() {
   const { checkValidSession } = useSessionGuard();
-  const { checkValidSubscription } = useSubscriptionGuard();
+  const { checkValidSubscription, checkValidProSubscription } = useSubscriptionGuard();
 
   return {
     checkValidSession,
     checkValidSubscription,
+    checkValidProSubscription,
   };
 }
 
@@ -57,7 +58,22 @@ export function useSubscriptionGuard() {
     return true; // Allow if not subscribed but still has requests left
   };
 
+  // Use this guard for features that are Pro only, not including free pro usage
+  const checkValidProSubscription = () => {
+    if (isPending) return false;
+
+    if (!subscriptionStatus) return false;
+
+    const { isSubscribed } = subscriptionStatus;
+
+    if (isSubscribed) return true;
+
+    openGetProDialog();
+    return false;
+  };
+
   return {
     checkValidSubscription,
+    checkValidProSubscription,
   };
 }