freebuff: Block non-US/CA IPs

Author: jahooma

bun.lock

@@ -1,55 +1,37 @@
 import { afterEach, beforeEach, describe, expect, mock, test } from 'bun:test'
 
+import { useChatStore } from '../../state/chat-store'
+
 import type { RouterParams } from '../command-registry'
-import * as analytics from '../../utils/analytics'
 
-const setInputMode = mock(() => {})
-const setMessages = mock(() => {})
 const saveToHistory = mock(() => {})
 const setInputValue = mock(() => {})
+const setMessages = mock(() => {})
 const handleChatGptAuthCode = mock(async () => ({
   success: true,
   message: 'ok',
 }))
 
-mock.module('../../state/chat-store', () => ({
-  useChatStore: {
-    getState: () => ({
-      inputMode: 'connect:chatgpt',
-      setInputMode,
-      pendingAttachments: [],
-    }),
-  },
-}))
-
 mock.module('../../components/chatgpt-connect-banner', () => ({
   handleChatGptAuthCode,
 }))
 
-mock.module('../../utils/analytics', () => ({
-  ...analytics,
-  trackEvent: () => {},
-}))
-
 mock.module('@codebuff/common/constants/chatgpt-oauth', () => ({
   CHATGPT_OAUTH_ENABLED: true,
 }))
 
 describe('routeUserPrompt connect:chatgpt mode', () => {
   beforeEach(() => {
-    setInputMode.mockClear()
-    setMessages.mockClear()
+    useChatStore.getState().reset()
+    useChatStore.getState().setInputMode('connect:chatgpt')
     saveToHistory.mockClear()
     setInputValue.mockClear()
+    setMessages.mockClear()
     handleChatGptAuthCode.mockClear()
   })
 
   afterEach(() => {
-    setInputMode.mockClear()
-    setMessages.mockClear()
-    saveToHistory.mockClear()
-    setInputValue.mockClear()
-    handleChatGptAuthCode.mockClear()
+    useChatStore.getState().reset()
   })
 
   test('when in connect:chatgpt mode, it exchanges the auth code and updates messages', async () => {
@@ -82,6 +64,6 @@ describe('routeUserPrompt connect:chatgpt mode', () => {
 
     expect(handleChatGptAuthCode).toHaveBeenCalledWith('auth-code-123')
     expect(setMessages).toHaveBeenCalled()
-    expect(setInputMode).toHaveBeenCalledWith('default')
+    expect(useChatStore.getState().inputMode).toBe('default')
   })
 })

cli/src/commands/__tests__/router-connect-chatgpt.test.ts

@@ -6,7 +6,9 @@ import { processBashContext } from '../../utils/bash-context-processor'
 import { markRunningAgentsAsCancelled } from '../../utils/block-operations'
 import {
   isOutOfCreditsError,
+  isFreeModeUnavailableError,
   OUT_OF_CREDITS_MESSAGE,
+  FREE_MODE_UNAVAILABLE_MESSAGE,
 } from '../../utils/error-handling'
 import { formatElapsedTime } from '../../utils/format-elapsed-time'
 import { processImagesForMessage } from '../../utils/image-processor'
@@ -336,6 +338,12 @@ export const handleRunCompletion = (params: {
       return
     }
 
+    if (isFreeModeUnavailableError(output)) {
+      updater.setError(FREE_MODE_UNAVAILABLE_MESSAGE)
+      finalizeAfterError()
+      return
+    }
+
     // Pass the raw error message to setError (displayed in UserErrorBanner without additional wrapper formatting)
     updater.setError(output.message ?? DEFAULT_RUN_OUTPUT_ERROR_MESSAGE)
 
@@ -418,6 +426,11 @@ export const handleRunError = (params: {
     return
   }
 
+  if (isFreeModeUnavailableError(error)) {
+    updater.setError(FREE_MODE_UNAVAILABLE_MESSAGE)
+    return
+  }
+
   // Use setError for all errors so they display in UserErrorBanner consistently
   const errorMessage = errorInfo.message || 'An unexpected error occurred'
   updater.setError(errorMessage)

cli/src/hooks/helpers/send-message.ts

@@ -37,8 +37,29 @@ export const isOutOfCreditsError = (error: unknown): boolean => {
   return false
 }
 
+/**
+ * Check if an error indicates free mode is not available in the user's country.
+ * Standardized on statusCode === 403 + error === 'free_mode_unavailable'.
+ */
+export const isFreeModeUnavailableError = (error: unknown): boolean => {
+  if (
+    error &&
+    typeof error === 'object' &&
+    'statusCode' in error &&
+    (error as { statusCode: unknown }).statusCode === 403 &&
+    'error' in error &&
+    (error as { error: unknown }).error === 'free_mode_unavailable'
+  ) {
+    return true
+  }
+  return false
+}
+
 export const OUT_OF_CREDITS_MESSAGE = `Out of credits. Please add credits at ${defaultAppUrl}/usage`
 
+export const FREE_MODE_UNAVAILABLE_MESSAGE =
+  'Free mode is not available outside of the United States and Canada. Please upgrade to a paid plan to use Codebuff outside the US and Canada.'
+
 export const createErrorMessage = (
   error: unknown,
   aiMessageId: string,

cli/src/utils/error-handling.ts

@@ -36,6 +36,7 @@ const nextConfig = {
       'encoding',
       'perf_hooks',
       'async_hooks',
+      'geoip-lite',
     )
 
     // Externalize code-map package to avoid bundling tree-sitter WASM files