feat: Improve autocorrelation error handling (#933)

Author: e-fisher

src/handlers/ai/preload.ts

@@ -1,12 +1,11 @@
 import { ipcRenderer } from 'electron'
 
 import {
+  AbortStreamChatRequest,
   AiHandler,
-  StreamChatRequest,
   StreamChatChunk,
   StreamChatEnd,
-  StreamChatError,
-  AbortStreamChatRequest,
+  StreamChatRequest,
   TokenUsage,
 } from './types'
 
@@ -43,20 +42,6 @@ export function streamChat(request: StreamChatRequest) {
       return () => ipcRenderer.removeListener(AiHandler.StreamChatEnd, handler)
     },
 
-    onError: (callback: (error: string) => void) => {
-      const handler = (
-        _event: Electron.IpcRendererEvent,
-        data: StreamChatError
-      ) => {
-        if (data.id === request.id) {
-          callback(data.error)
-        }
-      }
-      ipcRenderer.on(AiHandler.StreamChatError, handler)
-      return () =>
-        ipcRenderer.removeListener(AiHandler.StreamChatError, handler)
-    },
-
     abort: () => {
       const abortRequest: AbortStreamChatRequest = { id: request.id }
       ipcRenderer.send(AiHandler.AbortStreamChat, abortRequest)

src/handlers/ai/streamMessages.ts

@@ -1,57 +1,25 @@
 import { StreamTextResult, ToolSet } from 'ai'
-import log from 'electron-log/main'
 
-import {
-  AiHandler,
-  StreamChatChunk,
-  StreamChatEnd,
-  StreamChatError,
-} from './types'
+import { AiHandler } from './types'
 
 export async function streamMessages<Tools extends ToolSet, PARTIAL_OUTPUT>(
   webContents: Electron.WebContents,
   response: StreamTextResult<Tools, PARTIAL_OUTPUT>,
   requestId: string
 ) {
-  try {
-    const uiStream = response.toUIMessageStream({
-      onError: (error) => {
-        // Throw tools errors, without this they get silenced
-        throw error
-      },
-    })
-
-    // Process the stream and send chunks via IPC
-    const reader = uiStream.getReader()
-    let done = false
-
-    while (!done) {
-      const result = await reader.read()
-      done = result.done || false
-
-      if (done) {
-        const usageData = await response.usage
-
-        webContents.send(AiHandler.StreamChatEnd, {
-          id: requestId,
-          usage: usageData,
-        } satisfies StreamChatEnd)
-        break
-      }
+  const stream = response.toUIMessageStream({})
 
-      // Send chunk event
-      webContents.send(AiHandler.StreamChatChunk, {
-        id: requestId,
-        chunk: result.value,
-      } satisfies StreamChatChunk)
-    }
-  } catch (error) {
-    log.error('Error in handleStreamChat:', error)
-
-    // Send error event
-    webContents.send(AiHandler.StreamChatError, {
+  for await (const part of stream) {
+    webContents.send(AiHandler.StreamChatChunk, {
       id: requestId,
-      error: error instanceof Error ? error.message : 'Unknown error',
-    } satisfies StreamChatError)
+      chunk: part,
+    })
   }
+
+  const usageData = await response.usage
+
+  webContents.send(AiHandler.StreamChatEnd, {
+    id: requestId,
+    usage: usageData,
+  })
 }

src/handlers/ai/types.ts

@@ -4,7 +4,6 @@ export enum AiHandler {
   StreamChat = 'ai:streamChat',
   StreamChatChunk = 'ai:streamChatChunk',
   StreamChatEnd = 'ai:streamChatEnd',
-  StreamChatError = 'ai:streamChatError',
   AbortStreamChat = 'ai:abortStreamChat',
 }
 
@@ -29,11 +28,6 @@ export interface StreamChatEnd {
   usage?: TokenUsage
 }
 
-export interface StreamChatError {
-  id: string
-  error: string
-}
-
 export interface AbortStreamChatRequest {
   id: string
 }

src/hooks/useSettings.ts

@@ -1,4 +1,5 @@
 import { useMutation, useQuery } from '@tanstack/react-query'
+import { useEffect, useRef } from 'react'
 
 import { queryClient } from '@/utils/query'
 
@@ -32,3 +33,19 @@ export function useBrowserCheck() {
     queryFn: window.studio.ui.detectBrowser,
   })
 }
+
+export function useSettingsChanged(onChanged: () => void) {
+  const { data } = useSettings()
+  const prevDataRef = useRef(data)
+
+  useEffect(() => {
+    if (!data) {
+      return
+    }
+
+    if (prevDataRef.current !== data) {
+      onChanged()
+      prevDataRef.current = data
+    }
+  }, [data, onChanged])
+}

src/views/Generator/AutoCorrelation/AutoCorrelation.tsx

@@ -4,6 +4,7 @@ import { useEffect, useState } from 'react'
 import { useListenProxyData } from '@/hooks/useListenProxyData'
 import { useGeneratorStore } from '@/store/generator'
 
+import { ErrorMessage } from './ErrorMessage'
 import { IntroductionMessage } from './IntroductionMessage'
 import { SuggestedRules } from './SuggestedRules'
 import { TokenUsageIndicator } from './TokenUsageIndicator'
@@ -34,7 +35,9 @@ export function AutoCorrelation({
     correlationStatus,
     outcomeReason,
     tokenUsage,
+    error,
     stop,
+    restart,
   } = useGenerateRules({
     clearValidation: clearValidation,
   })
@@ -71,6 +74,10 @@ export function AutoCorrelation({
     return <IntroductionMessage onStart={start} />
   }
 
+  if (error) {
+    return <ErrorMessage error={error} onRetry={restart} />
+  }
+
   return (
     <Flex
       css={{
@@ -90,7 +97,6 @@ export function AutoCorrelation({
                     isLoading={isLoading}
                     onCheckRules={setCheckedRuleIds}
                     checkedRuleIds={checkedRuleIds}
-                    correlationStatus={correlationStatus}
                   />
                   {outcomeReason !== '' && (
                     <Box px="3" pt="2">

src/views/Generator/AutoCorrelation/ErrorMessage.tsx

@@ -0,0 +1,128 @@
+import { Button, Flex, Text } from '@radix-ui/themes'
+import { ExternalLink, KeyIcon, RefreshCw } from 'lucide-react'
+
+import grotCrashed from '@/assets/grot-crashed.svg'
+import { useSettingsChanged } from '@/hooks/useSettings'
+import { useStudioUIStore } from '@/store/ui'
+
+interface AutoCorrelationErrorProps {
+  error: Error
+  onRetry: () => void
+}
+
+export function ErrorMessage({ error, onRetry }: AutoCorrelationErrorProps) {
+  const errorMessage = error.message.toLowerCase()
+  const openSettingsDialog = useStudioUIStore(
+    (state) => state.openSettingsDialog
+  )
+
+  useSettingsChanged(() => {
+    onRetry()
+  })
+
+  const retryButton = (
+    <Button onClick={onRetry}>
+      <RefreshCw />
+      Retry
+    </Button>
+  )
+
+  const openSettingsButton = (
+    <Button onClick={() => openSettingsDialog('ai')}>
+      <KeyIcon />
+      Settings
+    </Button>
+  )
+
+  const reportIssueButton = (
+    <Button onClick={() => window.studio.ui.reportIssue()} variant="outline">
+      <ExternalLink />
+      Report issue
+    </Button>
+  )
+
+  if (errorMessage.includes('incorrect api key')) {
+    return (
+      <MessageContent
+        title="Incorrect API key"
+        message="The OpenAI API key is incorrect or has been revoked. Check your API key in settings."
+      >
+        {openSettingsButton}
+      </MessageContent>
+    )
+  }
+
+  if (errorMessage.includes('context window')) {
+    return (
+      <MessageContent
+        title="Token usage limit exceeded"
+        message="The recording exceeds the token limit. Try reducing the number of allowed hosts in your recording or work with a smaller recording."
+      />
+    )
+  }
+
+  if (errorMessage.includes('rate limit')) {
+    return (
+      <MessageContent
+        title="Too many requests"
+        message="You have exceeded the API rate limit. Wait a moment and try again."
+      >
+        {retryButton}
+      </MessageContent>
+    )
+  }
+
+  return (
+    <MessageContent
+      title="Something went wrong"
+      message="An unexpected error occurred during auto-correlation. Click retry to try again or report an issue if problem persists."
+    >
+      {retryButton}
+      {reportIssueButton}
+    </MessageContent>
+  )
+}
+
+function MessageContent({
+  title,
+  message,
+  children,
+}: {
+  title: string
+  message: React.ReactNode
+  children?: React.ReactNode
+}) {
+  return (
+    <Flex
+      direction="column"
+      align="center"
+      gap="6"
+      justify="center"
+      height="100%"
+      p="6"
+    >
+      <img
+        css={{
+          maxWidth: '200px',
+          transform: 'scaleX(-1)',
+        }}
+        src={grotCrashed}
+        aria-label="Error illustration"
+      />
+
+      <Flex direction="column" align="center" gap="3" maxWidth="400px">
+        <Text size="4" weight="medium" align="center">
+          {title}
+        </Text>
+
+        <Text size="2" color="gray" align="center">
+          {message}
+        </Text>
+
+        <Flex gap="3" mt="4">
+          {children}
+        </Flex>
+      </Flex>
+    </Flex>
+  )
+}

src/views/Generator/AutoCorrelation/SuggestedRules.tsx

@@ -7,22 +7,18 @@ import { fadeIn } from '@/utils/animations'
 import { TestRuleInlineContent } from '../TestRuleContainer/TestRule/TestRuleInlineContent'
 import { TestRuleTypeBadge } from '../TestRuleContainer/TestRule/TestRuleTypeBadge'
 
-import { CorrelationStatus } from './types'
-
 interface SuggestedRulesProps {
   suggestedRules: CorrelationRule[]
   isLoading: boolean
   checkedRuleIds: string[]
   onCheckRules: (ruleIds: string[]) => void
-  correlationStatus: CorrelationStatus
 }
 
 export function SuggestedRules({
   suggestedRules,
   isLoading,
   checkedRuleIds,
   onCheckRules,
-  correlationStatus,
 }: SuggestedRulesProps) {
   const allChecked =
     suggestedRules.length > 0 && checkedRuleIds.length === suggestedRules.length
@@ -35,22 +31,18 @@ export function SuggestedRules({
     }
   }
 
-  if (
-    ['not-started', 'correlation-not-needed', 'error'].includes(
-      correlationStatus
-    )
-  ) {
-    return null
-  }
-
-  if (suggestedRules.length === 0) {
+  if (isLoading && suggestedRules.length === 0) {
     return (
       <Flex height="100%" align="center" justify="center">
         <Text color="gray">Your rules will appear here</Text>
       </Flex>
     )
   }
 
+  if (!isLoading && suggestedRules.length === 0) {
+    return null
+  }
+
   return (
     <Box p="3">
       <Label mb="2">