diff --git a/npm-app/src/cli.ts b/npm-app/src/cli.ts index 64f9dfd3a..e37439b90 100644 --- a/npm-app/src/cli.ts +++ b/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() } diff --git a/npm-app/src/client.ts b/npm-app/src/client.ts index fa014f687..0fb8fdca7 100644 --- a/npm-app/src/client.ts +++ b/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 diff --git a/npm-app/src/rage-detectors.ts b/npm-app/src/rage-detectors.ts index 30087ef8b..a055171a6 100644 --- a/npm-app/src/rage-detectors.ts +++ b/npm-app/src/rage-detectors.ts @@ -12,18 +12,17 @@ export interface RageDetectors { keyMashingDetector: ReturnType repeatInputDetector: ReturnType exitAfterErrorDetector: ReturnType - webSocketHangDetector: ReturnType< - typeof createTimeoutDetector + responseHangDetector: ReturnType< + typeof createTimeoutDetector > startupTimeDetector: ReturnType exitTimeDetector: ReturnType } -// 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({ - reason: 'websocket_persistent_failure', + + + responseHangDetector: createTimeoutDetector({ + 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() }, }), diff --git a/npm-app/src/utils/__tests__/rage-detector.test.ts b/npm-app/src/utils/__tests__/rage-detector.test.ts index b53b7496f..f6680d7e1 100644 --- a/npm-app/src/utils/__tests__/rage-detector.test.ts +++ b/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() }) })