fix AI chat not working above limit and with key

WorldBrain · May 18, 2024 · 9848ac6 · 9848ac6
1 parent d49a91d
commit 9848ac6
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 14 deletions.
diff --git a/external/@worldbrain/memex-common b/external/@worldbrain/memex-common
diff --git a/package.json b/package.json
@@ -1,6 +1,6 @@
 {
     "name": "worldbrain-extension",
-    "version": "3.19.27",
+    "version": "3.19.28",
     "homepage": "https://memex.garden",
     "repository": "https://github.com/WorldBrain/Memex",
     "scripts": {

diff --git a/src/sidebar/annotations-sidebar/containers/logic.ts b/src/sidebar/annotations-sidebar/containers/logic.ts
@@ -905,6 +905,18 @@ export class SidebarContainerLogic extends UILogic<
             (await browser.storage.local.get('@Sidebar-reading_view')) ?? true
         // this.readingViewStorageListener(true)
 
+        const openAIKey = await this.syncSettings.openAI?.get('apiKey')
+        const hasAPIKey = openAIKey && openAIKey?.trim().startsWith('sk-')
+
+        const selectedModel = await this.syncSettings.openAI.get(
+            'selectedModel',
+        )
+
+        this.emitMutation({
+            hasKey: { $set: hasAPIKey },
+            AImodel: { $set: selectedModel ?? 'claude-3-haiku' },
+        })
+
         await loadInitial<SidebarContainerState>(this, async () => {
             this.showState = initialState ?? 'hidden'
             this.emitMutation({
@@ -971,18 +983,6 @@ export class SidebarContainerLogic extends UILogic<
             pageAlreadySaved: { $set: pageAlreadySaved },
         })
 
-        const openAIKey = await this.syncSettings.openAI?.get('apiKey')
-        const hasAPIKey = openAIKey && openAIKey?.trim().startsWith('sk-')
-
-        const selectedModel = await this.syncSettings.openAI.get(
-            'selectedModel',
-        )
-
-        this.emitMutation({
-            hasKey: { $set: hasAPIKey },
-            AImodel: { $set: selectedModel ?? 'claude-3-haiku' },
-        })
-
         const highlightColors = await this.fetchHighlightColors()
 
         this.emitMutation({ highlightColors: { $set: highlightColors } })

diff --git a/src/tests/subscriptionManagement/subscription-storage.test.ts b/src/tests/subscriptionManagement/subscription-storage.test.ts
@@ -321,6 +321,38 @@ describe('Ribbon logic', () => {
 
             expect(result).toBe(false)
         })
+
+        it('should allow action if user has AIpowerupOwnKey powerup AND has key AND using GPT-3 but over free tier limit', async ({
+            device,
+        }) => {
+            const { browserAPIs, analytics, collectionsBG } = await setupTest(
+                device,
+            )
+
+            const fakeStorageEntry = {
+                ...DEFAULT_COUNTER_STORAGE_VALUE,
+                cQ: DEFAULT_POWERUP_LIMITS.AIpowerup + 10,
+                pU: {
+                    ...DEFAULT_COUNTER_STORAGE_VALUE.pU,
+                    AIpowerupOwnKey: true,
+                },
+            }
+
+            await browserAPIs.storage.local.set({
+                [COUNTER_STORAGE_KEY]: fakeStorageEntry,
+            })
+
+            const result = await AIActionAllowed(
+                browserAPIs,
+                analytics,
+                true,
+                false,
+                'gpt-3',
+            )
+
+            expect(result).toBe(true)
+        })
+
         it('should prevent action if user has NO AI powerup AND no Key AND hit the limit', async ({
             device,
         }) => {