From 1237d0d380e769c8755dd2e46a843192e267eae0 Mon Sep 17 00:00:00 2001
From: Joshen Lim <joshenlimek@gmail.com>
Date: Wed, 4 Jun 2025 17:36:07 +0800
Subject: [PATCH] Float errors from AI endpoionts as toasts

---
 .../interfaces/ProjectCreation/SchemaGenerator.tsx   |  2 ++
 .../components/ui/AIAssistantPanel/AIAssistant.tsx   |  5 +++--
 .../ui/AIAssistantPanel/AIAssistant.utils.ts         | 12 ++++++++++++
 3 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/apps/studio/components/interfaces/ProjectCreation/SchemaGenerator.tsx b/apps/studio/components/interfaces/ProjectCreation/SchemaGenerator.tsx
index d44f1d7fc24..120428fe922 100644
--- a/apps/studio/components/interfaces/ProjectCreation/SchemaGenerator.tsx
+++ b/apps/studio/components/interfaces/ProjectCreation/SchemaGenerator.tsx
@@ -2,6 +2,7 @@ import { useChat } from 'ai/react'
 import { useEffect, useState } from 'react'
 
 import { Markdown } from 'components/interfaces/Markdown'
+import { onErrorChat } from 'components/ui/AIAssistantPanel/AIAssistant.utils'
 import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
 import { BASE_PATH } from 'lib/constants'
 import { AiIconAnimation, Button, Label_Shadcn_, Textarea } from 'ui'
@@ -44,6 +45,7 @@ export const SchemaGenerator = ({
     api: `${BASE_PATH}/api/ai/onboarding/design`,
     id: 'schema-generator',
     maxSteps: 7,
+    onError: onErrorChat,
     onFinish: () => {
       setInput('')
     },
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
index 126d0cdb69e..d3f538c8da2 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.tsx
@@ -11,7 +11,7 @@ import { subscriptionHasHipaaAddon } from 'components/interfaces/Billing/Subscri
 import { Markdown } from 'components/interfaces/Markdown'
 import { SQL_TEMPLATES } from 'components/interfaces/SQLEditor/SQLEditor.queries'
 import { useCheckOpenAIKeyQuery } from 'data/ai/check-api-key-query'
-import { constructHeaders, handleError } from 'data/fetchers'
+import { constructHeaders } from 'data/fetchers'
 import { useOrgSubscriptionQuery } from 'data/subscriptions/org-subscription-query'
 import { useTablesQuery } from 'data/tables/tables-query'
 import { useSendEventMutation } from 'data/telemetry/send-event-mutation'
@@ -36,6 +36,7 @@ import { Admonition, AssistantChatForm, GenericSkeletonLoader } from 'ui-pattern
 import { ButtonTooltip } from '../ButtonTooltip'
 import { DotGrid } from '../DotGrid'
 import { ErrorBoundary } from '../ErrorBoundary'
+import { onErrorChat } from './AIAssistant.utils'
 import { AIAssistantChatSelector } from './AIAssistantChatSelector'
 import { AIOnboarding } from './AIOnboarding'
 import { AIOptInModal } from './AIOptInModal'
@@ -164,7 +165,7 @@ export const AIAssistant = ({ className }: AIAssistantProps) => {
       schema: currentSchema,
       table: currentTable?.name,
     },
-    onError: handleError,
+    onError: onErrorChat,
     onFinish: handleChatFinish,
   })
 
diff --git a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts
index 60827c734a0..b8d6bfe460d 100644
--- a/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts
+++ b/apps/studio/components/ui/AIAssistantPanel/AIAssistant.utils.ts
@@ -1,3 +1,5 @@
+import { toast } from 'sonner'
+
 import { authKeys } from 'data/auth/keys'
 import { databaseExtensionsKeys } from 'data/database-extensions/keys'
 import { databaseIndexesKeys } from 'data/database-indexes/keys'
@@ -6,6 +8,7 @@ import { databaseTriggerKeys } from 'data/database-triggers/keys'
 import { databaseKeys } from 'data/database/keys'
 import { enumeratedTypesKeys } from 'data/enumerated-types/keys'
 import { tableKeys } from 'data/tables/keys'
+import { tryParseJson } from 'lib/helpers'
 import { SAFE_FUNCTIONS } from './AiAssistant.constants'
 
 // [Joshen] This is just very basic identification, but possible can extend perhaps
@@ -99,3 +102,12 @@ export const getContextualInvalidationKeys = ({
     )[key] ?? []
   )
 }
+
+export const onErrorChat = (error: { message: string }) => {
+  const errorObject = tryParseJson(error.message)
+  if (errorObject) {
+    toast.error(errorObject.error)
+  } else {
+    toast.error(error.message)
+  }
+}