[fix] convert hang detector from websocket monitoring to system response wait monitoring

npm-app/src/cli.ts

@@ -844,18 +844,11 @@ export class CLI {
       'WebSocket connection error'
     )
 
-    // Start hang detection for persistent connection issues
-    rageDetectors.webSocketHangDetector.start({
-      connectionIssue: 'websocket_persistent_failure',
-      url: websocketUrl,
-      getWebsocketState: () => Client.getInstance().webSocket.state,
-    })
+    // Stop response hang detector on error
+    rageDetectors.responseHangDetector.stop()
   }
 
   private onWebSocketReconnect() {
-    // Stop hang detection on successful reconnection
-    rageDetectors.webSocketHangDetector.stop()
-
     console.log('\n' + green('Reconnected!'))
     this.freshPrompt()
   }

npm-app/src/client.ts

@@ -18,7 +18,6 @@ import {
 import os from 'os'
 
 import { ApiKeyType, READABLE_NAME } from '@codebuff/common/api-keys/constants'
-import { AGENT_NAME_TO_TYPES, UNIQUE_AGENT_NAMES } from '@codebuff/common/constants/agents'
 import {
   ASKED_CONFIG,
   CostMode,
@@ -29,6 +28,10 @@ import {
   SHOULD_ASK_CONFIG,
   UserState,
 } from '@codebuff/common/constants'
+import {
+  AGENT_NAME_TO_TYPES,
+  UNIQUE_AGENT_NAMES,
+} from '@codebuff/common/constants/agents'
 import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
 import { codebuffConfigFile as CONFIG_FILE_NAME } from '@codebuff/common/json-config/constants'
 import {
@@ -75,6 +78,7 @@ import {
   getWorkingDirectory,
   startNewChat,
 } from './project-files'
+import { rageDetectors } from './rage-detectors'
 import { logAndHandleStartup } from './startup-process-handler'
 import { handleToolCall } from './tool-handlers'
 import { GitCommand, MakeNullable } from './types'
@@ -146,6 +150,7 @@ export class Client {
   private costMode: CostMode
   private responseComplete: boolean = false
   private userInputId: string | undefined
+  private isReceivingResponse: boolean = false
 
   public usageData: UsageData = {
     usage: 0,
@@ -956,7 +961,7 @@ export class Client {
 
     // Parse agent references from the prompt
     const { cleanPrompt, preferredAgents } = this.parseAgentReferences(prompt)
-    
+
     const urls = parseUrlsFromContent(cleanPrompt)
     const scrapedBlocks = await getScrapedContentBlocks(urls)
     const scrapedContent =
@@ -1000,37 +1005,46 @@ export class Client {
     }
   }
 
-  private parseAgentReferences(prompt: string): { cleanPrompt: string; preferredAgents: string[] } {
+  private parseAgentReferences(prompt: string): {
+    cleanPrompt: string
+    preferredAgents: string[]
+  } {
     let cleanPrompt = prompt
     const preferredAgents: string[] = []
-    
+
     // Create a regex pattern that matches any of the known agent names
-    const agentNamePattern = UNIQUE_AGENT_NAMES.map(name => 
-      name.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') // Escape special regex chars
+    const agentNamePattern = UNIQUE_AGENT_NAMES.map(
+      (name) => name.replace(/[.*+?^${}()|[\]\\]/g, '\\$&') // Escape special regex chars
     ).join('|')
-
-    const agentRegex = new RegExp(`@(${agentNamePattern})(?=\\s|$|[,.!?])`, 'gi')
+
+    const agentRegex = new RegExp(
+      `@(${agentNamePattern})(?=\\s|$|[,.!?])`,
+      'gi'
+    )
     const matches = prompt.match(agentRegex) || []
-    
+
     for (const match of matches) {
       const agentName = match.substring(1).trim() // Remove @ and trim
       // Find the exact agent name (case-insensitive)
-      const exactAgentName = UNIQUE_AGENT_NAMES.find(name => 
-        name.toLowerCase() === agentName.toLowerCase()
+      const exactAgentName = UNIQUE_AGENT_NAMES.find(
+        (name) => name.toLowerCase() === agentName.toLowerCase()
       )
-      
+
       if (exactAgentName) {
         const agentTypes = AGENT_NAME_TO_TYPES[exactAgentName]
         if (agentTypes && agentTypes.length > 0) {
           // Use the first matching agent type
           preferredAgents.push(agentTypes[0])
           // Remove ALL occurrences of this @ reference from the prompt using global replace
-          const matchRegex = new RegExp(match.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'g')
+          const matchRegex = new RegExp(
+            match.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'),
+            'g'
+          )
           cleanPrompt = cleanPrompt.replace(matchRegex, '').trim()
         }
       }
     }
-    
+
     return { cleanPrompt, preferredAgents }
   }
 
@@ -1149,6 +1163,13 @@ export class Client {
 
       rawChunkBuffer.push(chunk)
 
+      // Reset the response hang detector on each chunk received
+      rageDetectors.responseHangDetector.stop()
+      rageDetectors.responseHangDetector.start({
+        promptId: userInputId,
+        isReceivingResponse: () => this.isReceivingResponse,
+      })
+
       const trimmed = chunk.trim()
       for (const tag of ONE_TIME_TAGS) {
         if (trimmed.startsWith(`<${tag}>`) && trimmed.endsWith(closeXml(tag))) {
@@ -1306,6 +1327,16 @@ Go to https://www.codebuff.com/config for more information.`) +
 
     // Reset flags at the start of each response
     this.responseComplete = false
+    this.isReceivingResponse = true
+
+    // Start the response hang detector
+    rageDetectors.responseHangDetector.start({
+      promptId: userInputId,
+      isReceivingResponse: () => this.isReceivingResponse,
+    })
+
+    // Stop the response hang detector when user stops the response
+    rageDetectors.responseHangDetector.stop()
 
     return {
       responsePromise,
@@ -1331,6 +1362,9 @@ Go to https://www.codebuff.com/config for more information.`) +
         }),
       })
 
+      // Stop the response hang detector when response is complete
+      this.isReceivingResponse = false
+      rageDetectors.responseHangDetector.stop()
       const data = await response.json()
 
       // Use zod schema to validate response

npm-app/src/rage-detectors.ts

@@ -12,18 +12,17 @@ export interface RageDetectors {
   keyMashingDetector: ReturnType<typeof createCountDetector>
   repeatInputDetector: ReturnType<typeof createCountDetector>
   exitAfterErrorDetector: ReturnType<typeof createTimeBetweenDetector>
-  webSocketHangDetector: ReturnType<
-    typeof createTimeoutDetector<WebSocketHangDetectorContext>
+  responseHangDetector: ReturnType<
+    typeof createTimeoutDetector<ResponseHangDetectorContext>
   >
   startupTimeDetector: ReturnType<typeof createTimeBetweenDetector>
   exitTimeDetector: ReturnType<typeof createTimeBetweenDetector>
 }
 
-// Define the specific context type for WebSocket hang detector
-interface WebSocketHangDetectorContext {
-  connectionIssue?: string
-  url?: string
-  getWebsocketState: () => ReadyState
+// Define the specific context type for Response hang detector
+interface ResponseHangDetectorContext {
+  promptId?: string
+  isReceivingResponse: () => boolean
 }
 
 export function createRageDetectors(): RageDetectors {
@@ -77,20 +76,18 @@ export function createRageDetectors(): RageDetectors {
       operator: 'lt',
     }),
 
-    webSocketHangDetector: createTimeoutDetector<WebSocketHangDetectorContext>({
-      reason: 'websocket_persistent_failure',
+
+
+    responseHangDetector: createTimeoutDetector<ResponseHangDetectorContext>({
+      reason: 'response_hang',
       timeoutMs: 60_000,
       shouldFire: async (context) => {
-        if (!context || !context.getWebsocketState) {
+        if (!context || !context.isReceivingResponse) {
           return false
         }
 
-        // Add a 2-second grace period for reconnection
-        await sleep(2000)
-
-        // Only fire if the websocket is still not connected.
-        // This prevents firing if the connection is restored right before the timeout.
-        return context.getWebsocketState() !== WebSocket.OPEN
+        // Only fire if we're still expecting a response
+        return context.isReceivingResponse()
       },
     }),
 

npm-app/src/utils/__tests__/rage-detector.test.ts

@@ -584,14 +584,10 @@ describe('Rage Detectors', () => {
       expect(detectors.keyMashingDetector).toBeDefined()
       expect(detectors.repeatInputDetector).toBeDefined()
       expect(detectors.exitAfterErrorDetector).toBeDefined()
-      expect(detectors.webSocketHangDetector).toBeDefined()
-
       expect(typeof detectors.keyMashingDetector.recordEvent).toBe('function')
       expect(typeof detectors.repeatInputDetector.recordEvent).toBe('function')
       expect(typeof detectors.exitAfterErrorDetector.start).toBe('function')
       expect(typeof detectors.exitAfterErrorDetector.end).toBe('function')
-      expect(typeof detectors.webSocketHangDetector.start).toBe('function')
-      expect(typeof detectors.webSocketHangDetector.stop).toBe('function')
     })
 
     test('keyMashingDetector should work with realistic scenario', () => {
@@ -632,33 +628,7 @@ describe('Rage Detectors', () => {
       })
     })
 
-    test('webSocketHangDetector should work with realistic scenario', async () => {
-      // Test the timeout detector directly without the complex shouldFire logic
-      const detector = createTimeoutDetector({
-        reason: 'websocket_persistent_failure',
-        timeoutMs: 60000,
-      })
 
-      // Simulate WebSocket connection hanging
-      detector.start({
-        connectionIssue: 'websocket_persistent_failure',
-        url: 'ws://localhost:3000',
-      })
-
-      // Advance time to trigger the timeout
-      advanceTime(60000)
-
-      expect(mockTrackEvent).toHaveBeenCalledWith(
-        AnalyticsEvent.RAGE,
-        expect.objectContaining({
-          reason: 'websocket_persistent_failure',
-          durationMs: 60000,
-          timeoutMs: 60000,
-          connectionIssue: 'websocket_persistent_failure',
-          url: 'ws://localhost:3000',
-        })
-      )
-    })
 
     test('should not produce false positives for normal usage', () => {
       const detectors = createRageDetectors()
@@ -675,11 +645,6 @@ describe('Rage Detectors', () => {
       detectors.repeatInputDetector.recordEvent('fix the bug')
       detectors.repeatInputDetector.recordEvent('add tests')
 
-      // Normal WebSocket operation - connection succeeds quickly
-      detectors.webSocketHangDetector.start()
-      advanceTime(5000)
-      detectors.webSocketHangDetector.stop()
-
       expect(mockTrackEvent).not.toHaveBeenCalled()
     })
   })