feat: refactor subscription backend logic (#2426)

* cleaning up the webhook lifecyle side

* Add indexes to optimize count/sum queries and wire frontend to rate limits

* Adding rate limit decrement logic

* remove console.log

* adjusting logic from self-review

* fixing suggested comments

* great catch from auto-review

* adjusting base on comment

* small formatting

* refactor

* update seed and use db from context

* decrement of rate limit should only if user has active subscription

* Fix for login button loading states (#2481)

* Fixes both auth login buttons from showing loading animation when either one is actually loading

* Fix decimal input support, add VH units. (#2469)

* Fix decimal input support, add VH units.

* fix: preserve arrow key behavior

* attempt at fixing free/paid increment

* Populate schedule data when a subscription is scheduled to be downgraded

* Handle downgrade -> upgrade gracefully

* add backfill script

* count previous day and month for free plan

---------

Co-authored-by: Thomas Potaire <[email protected]>

Author: http-teapot

apps/web/client/messages/en.d.json.ts

@@ -182,7 +182,7 @@ declare const messages: {
             }
         },
         "footer": {
-            "unusedMessages": "Unused chat messages don't rollover to the next month"
+            "unusedMessages": "Unused chat messages will roll over to the next month."
         }
     },
     "editor": {

apps/web/client/messages/en.json

@@ -179,7 +179,7 @@
             }
         },
         "footer": {
-            "unusedMessages": "Unused chat messages don't rollover to the next month"
+            "unusedMessages": "Unused chat messages will roll over to the next month."
         }
     },
     "editor": {

apps/web/client/messages/es.json

@@ -179,7 +179,7 @@
             }
         },
         "footer": {
-            "unusedMessages": "Los mensajes de chat no utilizados no se transfieren al siguiente mes"
+            "unusedMessages": "Los mensajes de chat no utilizados se transferirán al siguiente mes"
         }
     },
     "editor": {

apps/web/client/messages/ja.json

@@ -78,7 +78,7 @@
             }
         },
         "footer": {
-            "unusedMessages": "未使用のチャットメッセージは翌月に繰り越されません。"
+            "unusedMessages": "未使用のチャットメッセージは翌月に繰り越されます。"
         }
     },
     "editor": {

apps/web/client/messages/ko.json

@@ -179,7 +179,7 @@
             }
         },
         "footer": {
-            "unusedMessages": "사용하지 않은 채팅 메시지는 다음 달로 이월되지 않습니다."
+            "unusedMessages": "사용하지 않은 채팅 메시지는 다음 달로 이월됩니다."
         }
     },
     "editor": {

apps/web/client/messages/zh.json

@@ -179,7 +179,7 @@
             }
         },
         "footer": {
-            "unusedMessages": "未使用的聊天消息不会结转到下个月"
+            "unusedMessages": "未使用的聊天消息会结转到下个月"
         }
     },
     "editor": {

apps/web/client/package.json

@@ -62,6 +62,7 @@
         "class-variance-authority": "^0.7.1",
         "clsx": "^2.1.1",
         "culori": "^4.0.1",
+        "date-fns": "^4.1.0",
         "dayjs": "^1.11.13",
         "embla-carousel-react": "^8.6.0",
         "flexsearch": "^0.8.160",

apps/web/client/public/onlook-preload-script.js

@@ -104,6 +104,24 @@ export const streamResponse = async (req: NextRequest) => {
     const { messages, maxSteps, chatType } = await req.json();
     const { model, providerOptions, headers } = await initModel(MainModelConfig);
 
+    // Updating the usage record and rate limit is done here to avoid
+    // abuse in the case where a single user sends many concurrent requests.
+    // If the call below fails, the user will not be penalized.
+    let usageRecordId: string | undefined;
+    let rateLimitId: string | undefined;
+    if (chatType === ChatType.EDIT) {
+        const user = await getSupabaseUser(req);
+        if (!user) {
+            throw new Error('User not found');
+        }
+        const { api } = await createTRPCClient(req);
+        const incrementRes = await api.usage.increment({
+            type: UsageType.MESSAGE,
+        });
+        usageRecordId = incrementRes?.usageRecordId;
+        rateLimitId = incrementRes?.rateLimitId;
+    }
+
     let systemPrompt: string;
     switch (chatType) {
         case ChatType.CREATE:
@@ -160,26 +178,17 @@ export const streamResponse = async (req: NextRequest) => {
 
             return { ...toolCall, args: JSON.stringify(repairedArgs) };
         },
-        onError: (error) => {
+        onError: async (error) => {
             console.error('Error in chat', error);
+            // if there was an error with the API, do not penalize the user
+            if (usageRecordId && rateLimitId) {
+                await createTRPCClient(req)
+                    .then(({ api }) => api.usage.revertIncrement({ usageRecordId, rateLimitId }))
+                    .catch(error => console.error('Error in chat usage decrement', error));
+            }
             throw error;
         },
     });
 
-    try {
-        if (chatType === ChatType.EDIT) {
-            const user = await getSupabaseUser(req);
-            if (!user) {
-                throw new Error('User not found');
-            }
-            const { api } = await createTRPCClient(req);
-            await api.usage.increment({
-                type: UsageType.MESSAGE,
-            });
-        }
-    } catch (error) {
-        console.error('Error in chat usage increment', error);
-    }
-
     return result.toDataStreamResponse();
 }

apps/web/client/src/app/api/chat/route.ts

@@ -1,17 +1,39 @@
 import { useEditorEngine } from '@/components/store/editor';
 import { useStateManager } from '@/components/store/state';
+import { api } from '@/trpc/react';
 import { Button } from '@onlook/ui/button';
 import { Icons } from '@onlook/ui/icons';
 import { observer } from 'mobx-react-lite';
+import { useEffect } from 'react';
 
 export const ErrorMessage = observer(() => {
     const editorEngine = useEditorEngine();
     const stateManager = useStateManager();
     const error = editorEngine.chat.error.message;
     const usage = editorEngine.chat.error.usage;
+    const isUsageError = usage && usage.usageCount >= usage.limitCount;
 
-    if (usage && usage.usageCount >= usage.limitCount) {
+    const { data: usageData, refetch: refetchUsage } = api.usage.get.useQuery();
 
+    useEffect(() => {
+        if (isUsageError) {
+            // Start polling for subscription changes.
+            const interval = setInterval(() => {
+                refetchUsage();
+            }, 3000);
+            return () => clearInterval(interval);
+        }
+    }, [isUsageError]);
+
+    // Clear error when usage data shows user is no longer over limit
+    useEffect(() => {
+        const isUnderLimit = usageData && usageData.daily.usageCount < usageData.daily.limitCount && usageData.monthly.usageCount < usageData.monthly.limitCount;
+        if (isUsageError && isUnderLimit) {
+            editorEngine.chat.error.usage = null;
+        }
+    }, [usageData, isUsageError]);
+
+    if (isUsageError) {
         return (
             <div className="flex w-full flex-col items-center justify-center gap-2 text-small px-4 pb-4">
                 <p className="text-foreground-secondary text-mini my-1 text-blue-300 select-none">