From 6921bda30f793c6d1db736fcb81adb6eaa9979cc Mon Sep 17 00:00:00 2001 From: Raza Rauf Date: Mon, 26 Jan 2026 17:51:42 +0500 Subject: [PATCH 01/16] refactor: extracted process listeners from main/index.ts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - types.ts: ProcessListenerDependencies interface and type re-exports. - data-listener.ts: Process output streaming with group chat buffering. - exit-listener.ts: Process exit handling with routing, recovery, synthesis. - session-id-listener.ts: Agent session ID tracking for group chat. - usage-listener.ts: Token/cost statistics for participants and moderator. - error-listener.ts: Agent error logging and forwarding. - stats-listener.ts: Query-complete events for stats database. - forwarding-listeners.ts: Simple IPC forwarding (slash-commands, etc.). This reduces main/index.ts by ~600 lines (1270 → 670) and adds 15 new tests covering forwarding, error, and stats listeners. --- src/main/index.ts | 695 ++---------------- .../__tests__/error-listener.test.ts | 118 +++ .../__tests__/forwarding-listeners.test.ts | 106 +++ .../__tests__/stats-listener.test.ts | 180 +++++ src/main/process-listeners/data-listener.ts | 102 +++ src/main/process-listeners/error-listener.ts | 31 + src/main/process-listeners/exit-listener.ts | 411 +++++++++++ .../process-listeners/forwarding-listeners.ts | 48 ++ src/main/process-listeners/index.ts | 55 ++ .../process-listeners/session-id-listener.ts | 77 ++ src/main/process-listeners/stats-listener.ts | 50 ++ src/main/process-listeners/types.ts | 150 ++++ src/main/process-listeners/usage-listener.ts | 114 +++ 13 files changed, 1490 insertions(+), 647 deletions(-) create mode 100644 src/main/process-listeners/__tests__/error-listener.test.ts create mode 100644 src/main/process-listeners/__tests__/forwarding-listeners.test.ts create mode 100644 src/main/process-listeners/__tests__/stats-listener.test.ts create mode 100644 src/main/process-listeners/data-listener.ts create mode 100644 src/main/process-listeners/error-listener.ts create mode 100644 src/main/process-listeners/exit-listener.ts create mode 100644 src/main/process-listeners/forwarding-listeners.ts create mode 100644 src/main/process-listeners/index.ts create mode 100644 src/main/process-listeners/session-id-listener.ts create mode 100644 src/main/process-listeners/stats-listener.ts create mode 100644 src/main/process-listeners/types.ts create mode 100644 src/main/process-listeners/usage-listener.ts diff --git a/src/main/index.ts b/src/main/index.ts index 58da5424..4f2838cf 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -98,6 +98,8 @@ import { createWindowManager, createQuitHandler, } from './app-lifecycle'; +// Phase 3 refactoring - process listeners +import { setupProcessListeners as setupProcessListenersModule } from './process-listeners'; // ============================================================================ // Data Directory Configuration (MUST happen before any Store initialization) @@ -614,655 +616,54 @@ function setupIpcHandlers() { } // Handle process output streaming (set up after initialization) +// Phase 3 refactoring - delegates to extracted process-listeners module function setupProcessListeners() { if (processManager) { - processManager.on('data', (sessionId: string, data: string) => { - // Handle group chat moderator output - buffer it - // Session ID format: group-chat-{groupChatId}-moderator-{uuid} or group-chat-{groupChatId}-moderator-synthesis-{uuid} - const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION); - if (moderatorMatch) { - const groupChatId = moderatorMatch[1]; - debugLog('GroupChat:Debug', `MODERATOR DATA received for chat ${groupChatId}`); - debugLog('GroupChat:Debug', `Session ID: ${sessionId}`); - debugLog('GroupChat:Debug', `Data length: ${data.length}`); - // Buffer the output - will be routed on process exit - const totalLength = appendToGroupChatBuffer(sessionId, data); - debugLog('GroupChat:Debug', `Buffered total: ${totalLength} chars`); - return; // Don't send to regular process:data handler - } - - // Handle group chat participant output - buffer it - // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} - const participantInfo = parseParticipantSessionId(sessionId); - if (participantInfo) { - debugLog('GroupChat:Debug', 'PARTICIPANT DATA received'); - debugLog( - 'GroupChat:Debug', - `Chat: ${participantInfo.groupChatId}, Participant: ${participantInfo.participantName}` - ); - debugLog('GroupChat:Debug', `Session ID: ${sessionId}`); - debugLog('GroupChat:Debug', `Data length: ${data.length}`); - // Buffer the output - will be routed on process exit - const totalLength = appendToGroupChatBuffer(sessionId, data); - debugLog('GroupChat:Debug', `Buffered total: ${totalLength} chars`); - return; // Don't send to regular process:data handler - } - - safeSend('process:data', sessionId, data); - - // Broadcast to web clients - extract base session ID (remove -ai or -terminal suffix) - // IMPORTANT: Skip PTY terminal output (-terminal suffix) as it contains raw ANSI codes. - // Web interface terminal commands use runCommand() which emits with plain session IDs. - if (webServer) { - // Don't broadcast raw PTY terminal output to web clients - if (sessionId.endsWith('-terminal')) { - debugLog('WebBroadcast', `SKIPPING PTY terminal output for web: session=${sessionId}`); - return; - } - - // Don't broadcast background batch/synopsis output to web clients - // These are internal Auto Run operations that should only appear in history, not as chat messages - if (sessionId.includes('-batch-') || sessionId.includes('-synopsis-')) { - debugLog('WebBroadcast', `SKIPPING batch/synopsis output for web: session=${sessionId}`); - return; - } - - // Extract base session ID and tab ID from format: {id}-ai-{tabId} - const baseSessionId = sessionId.replace(REGEX_AI_SUFFIX, ''); - const isAiOutput = sessionId.includes('-ai-'); - - // Extract tab ID from session ID format: {id}-ai-{tabId} - const tabIdMatch = sessionId.match(REGEX_AI_TAB_ID); - const tabId = tabIdMatch ? tabIdMatch[1] : undefined; - - const msgId = `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; - debugLog( - 'WebBroadcast', - `Broadcasting session_output: msgId=${msgId}, session=${baseSessionId}, tabId=${tabId || 'none'}, source=${isAiOutput ? 'ai' : 'terminal'}, dataLen=${data.length}` - ); - webServer.broadcastToSessionClients(baseSessionId, { - type: 'session_output', - sessionId: baseSessionId, - tabId, - data, - source: isAiOutput ? 'ai' : 'terminal', - timestamp: Date.now(), - msgId, - }); - } + setupProcessListenersModule(processManager, { + getProcessManager: () => processManager, + getWebServer: () => webServer, + getAgentDetector: () => agentDetector, + safeSend, + powerManager, + groupChatEmitters, + groupChatRouter: { + routeModeratorResponse, + routeAgentResponse, + markParticipantResponded, + spawnModeratorSynthesis, + getGroupChatReadOnlyState, + respawnParticipantWithRecovery, + }, + groupChatStorage: { + loadGroupChat, + updateGroupChat, + updateParticipant, + }, + sessionRecovery: { + needsSessionRecovery, + initiateSessionRecovery, + }, + outputBuffer: { + appendToGroupChatBuffer, + getGroupChatBufferedOutput, + clearGroupChatBuffer, + }, + outputParser: { + extractTextFromStreamJson, + parseParticipantSessionId, + }, + usageAggregator: { + calculateContextTokens, + }, + getStatsDB, + debugLog, + patterns: { + REGEX_MODERATOR_SESSION, + REGEX_MODERATOR_SESSION_TIMESTAMP, + REGEX_AI_SUFFIX, + REGEX_AI_TAB_ID, + }, + logger, }); - - processManager.on('exit', (sessionId: string, code: number) => { - // Remove power block reason for this session - // This allows system sleep when no AI sessions are active - powerManager.removeBlockReason(`session:${sessionId}`); - - // Handle group chat moderator exit - route buffered output and set state back to idle - // Session ID format: group-chat-{groupChatId}-moderator-{uuid} - // This handles BOTH initial moderator responses AND synthesis responses. - // The routeModeratorResponse function will check for @mentions: - // - If @mentions present: route to agents (continue conversation) - // - If no @mentions: final response to user (conversation complete for this turn) - const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION); - if (moderatorMatch) { - const groupChatId = moderatorMatch[1]; - debugLog('GroupChat:Debug', ` ========== MODERATOR PROCESS EXIT ==========`); - debugLog('GroupChat:Debug', ` Group Chat ID: ${groupChatId}`); - debugLog('GroupChat:Debug', ` Session ID: ${sessionId}`); - debugLog('GroupChat:Debug', ` Exit code: ${code}`); - logger.debug(`[GroupChat] Moderator exit: groupChatId=${groupChatId}`, 'ProcessListener', { - sessionId, - }); - // Route the buffered output now that process is complete - const bufferedOutput = getGroupChatBufferedOutput(sessionId); - debugLog('GroupChat:Debug', ` Buffered output length: ${bufferedOutput?.length ?? 0}`); - if (bufferedOutput) { - debugLog( - 'GroupChat:Debug', - ` Raw buffered output preview: "${bufferedOutput.substring(0, 300)}${bufferedOutput.length > 300 ? '...' : ''}"` - ); - logger.debug( - `[GroupChat] Moderator has buffered output (${bufferedOutput.length} chars)`, - 'ProcessListener', - { groupChatId } - ); - void (async () => { - try { - const chat = await loadGroupChat(groupChatId); - debugLog('GroupChat:Debug', ` Chat loaded for parsing: ${chat?.name || 'null'}`); - const agentType = chat?.moderatorAgentId; - debugLog('GroupChat:Debug', ` Agent type for parsing: ${agentType}`); - const parsedText = extractTextFromStreamJson(bufferedOutput, agentType); - debugLog('GroupChat:Debug', ` Parsed text length: ${parsedText.length}`); - debugLog( - 'GroupChat:Debug', - ` Parsed text preview: "${parsedText.substring(0, 300)}${parsedText.length > 300 ? '...' : ''}"` - ); - if (parsedText.trim()) { - debugLog('GroupChat:Debug', ` Routing moderator response...`); - logger.info( - `[GroupChat] Routing moderator response (${parsedText.length} chars)`, - 'ProcessListener', - { groupChatId } - ); - const readOnly = getGroupChatReadOnlyState(groupChatId); - debugLog('GroupChat:Debug', ` Read-only state: ${readOnly}`); - routeModeratorResponse( - groupChatId, - parsedText, - processManager ?? undefined, - agentDetector ?? undefined, - readOnly - ).catch((err) => { - debugLog('GroupChat:Debug', ` ERROR routing moderator response:`, err); - logger.error( - '[GroupChat] Failed to route moderator response', - 'ProcessListener', - { error: String(err) } - ); - }); - } else { - debugLog('GroupChat:Debug', ` WARNING: Parsed text is empty!`); - logger.warn( - '[GroupChat] Moderator output parsed to empty string', - 'ProcessListener', - { groupChatId, bufferedLength: bufferedOutput.length } - ); - } - } catch (err) { - debugLog('GroupChat:Debug', ` ERROR loading chat:`, err); - logger.error( - '[GroupChat] Failed to load chat for moderator output parsing', - 'ProcessListener', - { error: String(err) } - ); - const parsedText = extractTextFromStreamJson(bufferedOutput); - if (parsedText.trim()) { - const readOnly = getGroupChatReadOnlyState(groupChatId); - routeModeratorResponse( - groupChatId, - parsedText, - processManager ?? undefined, - agentDetector ?? undefined, - readOnly - ).catch((routeErr) => { - debugLog( - 'GroupChat:Debug', - ` ERROR routing moderator response (fallback):`, - routeErr - ); - logger.error( - '[GroupChat] Failed to route moderator response', - 'ProcessListener', - { error: String(routeErr) } - ); - }); - } - } - })().finally(() => { - clearGroupChatBuffer(sessionId); - debugLog('GroupChat:Debug', ` Cleared output buffer for session`); - }); - } else { - debugLog('GroupChat:Debug', ` WARNING: No buffered output!`); - logger.warn('[GroupChat] Moderator exit with no buffered output', 'ProcessListener', { - groupChatId, - sessionId, - }); - } - groupChatEmitters.emitStateChange?.(groupChatId, 'idle'); - debugLog('GroupChat:Debug', ` Emitted state change: idle`); - debugLog('GroupChat:Debug', ` =============================================`); - // Don't send to regular exit handler - return; - } - - // Handle group chat participant exit - route buffered output and update participant state - // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} - const participantExitInfo = parseParticipantSessionId(sessionId); - if (participantExitInfo) { - const { groupChatId, participantName } = participantExitInfo; - debugLog('GroupChat:Debug', ` ========== PARTICIPANT PROCESS EXIT ==========`); - debugLog('GroupChat:Debug', ` Group Chat ID: ${groupChatId}`); - debugLog('GroupChat:Debug', ` Participant: ${participantName}`); - debugLog('GroupChat:Debug', ` Session ID: ${sessionId}`); - debugLog('GroupChat:Debug', ` Exit code: ${code}`); - logger.debug( - `[GroupChat] Participant exit: ${participantName} (groupChatId=${groupChatId})`, - 'ProcessListener', - { sessionId } - ); - - // Emit participant state change to show this participant is done working - groupChatEmitters.emitParticipantState?.(groupChatId, participantName, 'idle'); - debugLog('GroupChat:Debug', ` Emitted participant state: idle`); - - // Route the buffered output now that process is complete - // IMPORTANT: We must wait for the response to be logged before triggering synthesis - // to avoid a race condition where synthesis reads the log before the response is written - const bufferedOutput = getGroupChatBufferedOutput(sessionId); - debugLog('GroupChat:Debug', ` Buffered output length: ${bufferedOutput?.length ?? 0}`); - - // Helper function to mark participant and potentially trigger synthesis - const markAndMaybeSynthesize = () => { - const isLastParticipant = markParticipantResponded(groupChatId, participantName); - debugLog('GroupChat:Debug', ` Is last participant to respond: ${isLastParticipant}`); - if (isLastParticipant && processManager && agentDetector) { - // All participants have responded - spawn moderator synthesis round - debugLog( - 'GroupChat:Debug', - ` All participants responded - spawning synthesis round...` - ); - logger.info( - '[GroupChat] All participants responded, spawning moderator synthesis', - 'ProcessListener', - { groupChatId } - ); - spawnModeratorSynthesis(groupChatId, processManager, agentDetector).catch((err) => { - debugLog('GroupChat:Debug', ` ERROR spawning synthesis:`, err); - logger.error('[GroupChat] Failed to spawn moderator synthesis', 'ProcessListener', { - error: String(err), - groupChatId, - }); - }); - } else if (!isLastParticipant) { - // More participants pending - debugLog('GroupChat:Debug', ` Waiting for more participants to respond...`); - } - }; - - if (bufferedOutput) { - debugLog( - 'GroupChat:Debug', - ` Raw buffered output preview: "${bufferedOutput.substring(0, 300)}${bufferedOutput.length > 300 ? '...' : ''}"` - ); - - // Handle session recovery and normal processing in an async IIFE - void (async () => { - // Check if this is a session_not_found error - if so, recover and retry - const chat = await loadGroupChat(groupChatId); - const agentType = chat?.participants.find((p) => p.name === participantName)?.agentId; - - if (needsSessionRecovery(bufferedOutput, agentType)) { - debugLog( - 'GroupChat:Debug', - ` Session not found error detected for ${participantName} - initiating recovery` - ); - logger.info('[GroupChat] Session recovery needed', 'ProcessListener', { - groupChatId, - participantName, - }); - - // Clear the buffer first - clearGroupChatBuffer(sessionId); - - // Initiate recovery (clears agentSessionId) - await initiateSessionRecovery(groupChatId, participantName); - - // Re-spawn the participant with recovery context - if (processManager && agentDetector) { - debugLog( - 'GroupChat:Debug', - ` Re-spawning ${participantName} with recovery context...` - ); - try { - await respawnParticipantWithRecovery( - groupChatId, - participantName, - processManager, - agentDetector - ); - debugLog( - 'GroupChat:Debug', - ` Successfully re-spawned ${participantName} for recovery` - ); - // Don't mark as responded yet - the recovery spawn will complete and trigger this - } catch (respawnErr) { - debugLog('GroupChat:Debug', ` Failed to respawn ${participantName}:`, respawnErr); - logger.error( - '[GroupChat] Failed to respawn participant for recovery', - 'ProcessListener', - { - error: String(respawnErr), - participant: participantName, - } - ); - // Mark as responded since recovery failed - markAndMaybeSynthesize(); - } - } else { - debugLog( - 'GroupChat:Debug', - ` Cannot respawn - processManager or agentDetector not available` - ); - markAndMaybeSynthesize(); - } - debugLog('GroupChat:Debug', ` ===============================================`); - return; - } - - // Normal processing - parse and route the response - try { - debugLog( - 'GroupChat:Debug', - ` Chat loaded for participant parsing: ${chat?.name || 'null'}` - ); - debugLog('GroupChat:Debug', ` Agent type for parsing: ${agentType}`); - const parsedText = extractTextFromStreamJson(bufferedOutput, agentType); - debugLog('GroupChat:Debug', ` Parsed text length: ${parsedText.length}`); - debugLog( - 'GroupChat:Debug', - ` Parsed text preview: "${parsedText.substring(0, 200)}${parsedText.length > 200 ? '...' : ''}"` - ); - if (parsedText.trim()) { - debugLog('GroupChat:Debug', ` Routing agent response from ${participantName}...`); - // Await the response logging before marking participant as responded - await routeAgentResponse( - groupChatId, - participantName, - parsedText, - processManager ?? undefined - ); - debugLog( - 'GroupChat:Debug', - ` Successfully routed agent response from ${participantName}` - ); - } else { - debugLog( - 'GroupChat:Debug', - ` WARNING: Parsed text is empty for ${participantName}!` - ); - } - } catch (err) { - debugLog('GroupChat:Debug', ` ERROR loading chat for participant:`, err); - logger.error( - '[GroupChat] Failed to load chat for participant output parsing', - 'ProcessListener', - { error: String(err), participant: participantName } - ); - try { - const parsedText = extractTextFromStreamJson(bufferedOutput); - if (parsedText.trim()) { - await routeAgentResponse( - groupChatId, - participantName, - parsedText, - processManager ?? undefined - ); - } - } catch (routeErr) { - debugLog('GroupChat:Debug', ` ERROR routing agent response (fallback):`, routeErr); - logger.error('[GroupChat] Failed to route agent response', 'ProcessListener', { - error: String(routeErr), - participant: participantName, - }); - } - } - })().finally(() => { - clearGroupChatBuffer(sessionId); - debugLog('GroupChat:Debug', ` Cleared output buffer for participant session`); - // Mark participant and trigger synthesis AFTER logging is complete - markAndMaybeSynthesize(); - }); - } else { - debugLog( - 'GroupChat:Debug', - ` WARNING: No buffered output for participant ${participantName}!` - ); - // No output to log, so mark participant as responded immediately - markAndMaybeSynthesize(); - } - debugLog('GroupChat:Debug', ` ===============================================`); - // Don't send to regular exit handler - return; - } - - safeSend('process:exit', sessionId, code); - - // Broadcast exit to web clients - if (webServer) { - // Extract base session ID from formats: {id}-ai-{tabId}, {id}-terminal, {id}-batch-{timestamp}, {id}-synopsis-{timestamp} - const baseSessionId = sessionId.replace( - /-ai-[^-]+$|-terminal$|-batch-\d+$|-synopsis-\d+$/, - '' - ); - webServer.broadcastToSessionClients(baseSessionId, { - type: 'session_exit', - sessionId: baseSessionId, - exitCode: code, - timestamp: Date.now(), - }); - } - }); - - processManager.on('session-id', (sessionId: string, agentSessionId: string) => { - // Handle group chat participant session ID - store the agent's session ID - // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} - const participantSessionInfo = parseParticipantSessionId(sessionId); - if (participantSessionInfo) { - const { groupChatId, participantName } = participantSessionInfo; - // Update the participant with the agent's session ID - updateParticipant(groupChatId, participantName, { agentSessionId }) - .then(async () => { - // Emit participants changed so UI updates with the new session ID - const chat = await loadGroupChat(groupChatId); - if (chat) { - groupChatEmitters.emitParticipantsChanged?.(groupChatId, chat.participants); - } - }) - .catch((err) => { - logger.error( - '[GroupChat] Failed to update participant agentSessionId', - 'ProcessListener', - { error: String(err), participant: participantName } - ); - }); - // Don't return - still send to renderer for logging purposes - } - - // Handle group chat moderator session ID - store the real agent session ID - // Session ID format: group-chat-{groupChatId}-moderator-{timestamp} - const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION_TIMESTAMP); - if (moderatorMatch) { - const groupChatId = moderatorMatch[1]; - // Update the group chat with the moderator's real agent session ID - // Store in moderatorAgentSessionId (not moderatorSessionId which is the routing prefix) - updateGroupChat(groupChatId, { moderatorAgentSessionId: agentSessionId }) - .then(() => { - // Emit session ID change event so UI updates with the new session ID - groupChatEmitters.emitModeratorSessionIdChanged?.(groupChatId, agentSessionId); - }) - .catch((err: unknown) => { - logger.error( - '[GroupChat] Failed to update moderator agent session ID', - 'ProcessListener', - { error: String(err), groupChatId } - ); - }); - // Don't return - still send to renderer for logging purposes - } - - safeSend('process:session-id', sessionId, agentSessionId); - }); - - // Handle slash commands from Claude Code init message - processManager.on('slash-commands', (sessionId: string, slashCommands: string[]) => { - safeSend('process:slash-commands', sessionId, slashCommands); - }); - - // Handle thinking/streaming content chunks from AI agents - // Emitted when agents produce partial text events (isPartial: true) - // Renderer decides whether to display based on tab's showThinking setting - processManager.on('thinking-chunk', (sessionId: string, content: string) => { - safeSend('process:thinking-chunk', sessionId, content); - }); - - // Handle tool execution events (OpenCode, Codex) - processManager.on( - 'tool-execution', - (sessionId: string, toolEvent: { toolName: string; state?: unknown; timestamp: number }) => { - safeSend('process:tool-execution', sessionId, toolEvent); - } - ); - - // Handle stderr separately from runCommand (for clean command execution) - processManager.on('stderr', (sessionId: string, data: string) => { - safeSend('process:stderr', sessionId, data); - }); - - // Handle command exit (from runCommand - separate from PTY exit) - processManager.on('command-exit', (sessionId: string, code: number) => { - safeSend('process:command-exit', sessionId, code); - }); - - // Handle usage statistics from AI responses - processManager.on( - 'usage', - ( - sessionId: string, - usageStats: { - inputTokens: number; - outputTokens: number; - cacheReadInputTokens: number; - cacheCreationInputTokens: number; - totalCostUsd: number; - contextWindow: number; - reasoningTokens?: number; // Separate reasoning tokens (Codex o3/o4-mini) - } - ) => { - // Handle group chat participant usage - update participant stats - const participantUsageInfo = parseParticipantSessionId(sessionId); - if (participantUsageInfo) { - const { groupChatId, participantName } = participantUsageInfo; - - // Calculate context usage percentage using agent-specific logic - // Note: For group chat, we don't have agent type here, defaults to Claude behavior - const totalContextTokens = calculateContextTokens(usageStats); - const contextUsage = - usageStats.contextWindow > 0 - ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) - : 0; - - // Update participant with usage stats - updateParticipant(groupChatId, participantName, { - contextUsage, - tokenCount: totalContextTokens, - totalCost: usageStats.totalCostUsd, - }) - .then(async () => { - // Emit participants changed so UI updates - const chat = await loadGroupChat(groupChatId); - if (chat) { - groupChatEmitters.emitParticipantsChanged?.(groupChatId, chat.participants); - } - }) - .catch((err) => { - logger.error('[GroupChat] Failed to update participant usage', 'ProcessListener', { - error: String(err), - participant: participantName, - }); - }); - // Still send to renderer for consistency - } - - // Handle group chat moderator usage - emit for UI - const moderatorUsageMatch = sessionId.match(REGEX_MODERATOR_SESSION); - if (moderatorUsageMatch) { - const groupChatId = moderatorUsageMatch[1]; - // Calculate context usage percentage using agent-specific logic - // Note: Moderator is typically Claude, defaults to Claude behavior - const totalContextTokens = calculateContextTokens(usageStats); - const contextUsage = - usageStats.contextWindow > 0 - ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) - : 0; - - // Emit moderator usage for the moderator card - groupChatEmitters.emitModeratorUsage?.(groupChatId, { - contextUsage, - totalCost: usageStats.totalCostUsd, - tokenCount: totalContextTokens, - }); - } - - safeSend('process:usage', sessionId, usageStats); - } - ); - - // Handle agent errors (auth expired, token exhaustion, rate limits, etc.) - processManager.on( - 'agent-error', - ( - sessionId: string, - agentError: { - type: string; - message: string; - recoverable: boolean; - agentId: string; - sessionId?: string; - timestamp: number; - raw?: { - exitCode?: number; - stderr?: string; - stdout?: string; - errorLine?: string; - }; - } - ) => { - logger.info(`Agent error detected: ${agentError.type}`, 'AgentError', { - sessionId, - agentId: agentError.agentId, - errorType: agentError.type, - message: agentError.message, - recoverable: agentError.recoverable, - }); - safeSend('agent:error', sessionId, agentError); - } - ); - - // Handle query-complete events for stats tracking - // This is emitted when a batch mode AI query completes (user or auto) - processManager.on( - 'query-complete', - ( - _sessionId: string, - queryData: { - sessionId: string; - agentType: string; - source: 'user' | 'auto'; - startTime: number; - duration: number; - projectPath?: string; - tabId?: string; - } - ) => { - try { - const db = getStatsDB(); - if (db.isReady()) { - const id = db.insertQueryEvent({ - sessionId: queryData.sessionId, - agentType: queryData.agentType, - source: queryData.source, - startTime: queryData.startTime, - duration: queryData.duration, - projectPath: queryData.projectPath, - tabId: queryData.tabId, - }); - logger.debug(`Recorded query event: ${id}`, '[Stats]', { - sessionId: queryData.sessionId, - agentType: queryData.agentType, - source: queryData.source, - duration: queryData.duration, - }); - // Broadcast stats update to renderer for real-time dashboard refresh - safeSend('stats:updated'); - } - } catch (error) { - logger.error(`Failed to record query event: ${error}`, '[Stats]', { - sessionId: queryData.sessionId, - }); - } - } - ); } } diff --git a/src/main/process-listeners/__tests__/error-listener.test.ts b/src/main/process-listeners/__tests__/error-listener.test.ts new file mode 100644 index 00000000..51c8754c --- /dev/null +++ b/src/main/process-listeners/__tests__/error-listener.test.ts @@ -0,0 +1,118 @@ +/** + * Tests for error listener. + * Handles agent errors (auth expired, token exhaustion, rate limits, etc.). + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupErrorListener } from '../error-listener'; +import type { ProcessManager } from '../../process-manager'; +import type { SafeSendFn } from '../../utils/safe-send'; +import type { AgentError } from '../../../shared/types'; +import type { ProcessListenerDependencies } from '../types'; + +describe('Error Listener', () => { + let mockProcessManager: ProcessManager; + let mockSafeSend: SafeSendFn; + let mockLogger: ProcessListenerDependencies['logger']; + let eventHandlers: Map void>; + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockSafeSend = vi.fn(); + mockLogger = { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }; + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + }); + + it('should register the agent-error event listener', () => { + setupErrorListener(mockProcessManager, { safeSend: mockSafeSend, logger: mockLogger }); + + expect(mockProcessManager.on).toHaveBeenCalledWith('agent-error', expect.any(Function)); + }); + + it('should log agent error and forward to renderer', () => { + setupErrorListener(mockProcessManager, { safeSend: mockSafeSend, logger: mockLogger }); + + const handler = eventHandlers.get('agent-error'); + const testSessionId = 'test-session-123'; + const testAgentError: AgentError = { + type: 'auth_expired', + agentId: 'claude-code', + message: 'Authentication token has expired', + recoverable: true, + timestamp: Date.now(), + }; + + handler?.(testSessionId, testAgentError); + + expect(mockLogger.info).toHaveBeenCalledWith( + 'Agent error detected: auth_expired', + 'AgentError', + expect.objectContaining({ + sessionId: testSessionId, + agentId: 'claude-code', + errorType: 'auth_expired', + message: 'Authentication token has expired', + recoverable: true, + }) + ); + + expect(mockSafeSend).toHaveBeenCalledWith('agent:error', testSessionId, testAgentError); + }); + + it('should handle token exhaustion errors', () => { + setupErrorListener(mockProcessManager, { safeSend: mockSafeSend, logger: mockLogger }); + + const handler = eventHandlers.get('agent-error'); + const testSessionId = 'session-456'; + const testAgentError: AgentError = { + type: 'token_exhaustion', + agentId: 'codex', + message: 'Token limit exceeded', + recoverable: false, + timestamp: Date.now(), + }; + + handler?.(testSessionId, testAgentError); + + expect(mockSafeSend).toHaveBeenCalledWith('agent:error', testSessionId, testAgentError); + }); + + it('should handle rate limit errors', () => { + setupErrorListener(mockProcessManager, { safeSend: mockSafeSend, logger: mockLogger }); + + const handler = eventHandlers.get('agent-error'); + const testSessionId = 'session-789'; + const testAgentError: AgentError = { + type: 'rate_limited', + agentId: 'opencode', + message: 'Rate limit exceeded, retry after 60 seconds', + recoverable: true, + timestamp: Date.now(), + }; + + handler?.(testSessionId, testAgentError); + + expect(mockLogger.info).toHaveBeenCalledWith( + 'Agent error detected: rate_limited', + 'AgentError', + expect.objectContaining({ + sessionId: testSessionId, + errorType: 'rate_limited', + }) + ); + + expect(mockSafeSend).toHaveBeenCalledWith('agent:error', testSessionId, testAgentError); + }); +}); diff --git a/src/main/process-listeners/__tests__/forwarding-listeners.test.ts b/src/main/process-listeners/__tests__/forwarding-listeners.test.ts new file mode 100644 index 00000000..5fcbcf53 --- /dev/null +++ b/src/main/process-listeners/__tests__/forwarding-listeners.test.ts @@ -0,0 +1,106 @@ +/** + * Tests for forwarding listeners. + * These listeners simply forward process events to the renderer via IPC. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupForwardingListeners } from '../forwarding-listeners'; +import type { ProcessManager } from '../../process-manager'; +import type { SafeSendFn } from '../../utils/safe-send'; + +describe('Forwarding Listeners', () => { + let mockProcessManager: ProcessManager; + let mockSafeSend: SafeSendFn; + let eventHandlers: Map void>; + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockSafeSend = vi.fn(); + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + }); + + it('should register all forwarding event listeners', () => { + setupForwardingListeners(mockProcessManager, { safeSend: mockSafeSend }); + + expect(mockProcessManager.on).toHaveBeenCalledWith('slash-commands', expect.any(Function)); + expect(mockProcessManager.on).toHaveBeenCalledWith('thinking-chunk', expect.any(Function)); + expect(mockProcessManager.on).toHaveBeenCalledWith('tool-execution', expect.any(Function)); + expect(mockProcessManager.on).toHaveBeenCalledWith('stderr', expect.any(Function)); + expect(mockProcessManager.on).toHaveBeenCalledWith('command-exit', expect.any(Function)); + }); + + it('should forward slash-commands events to renderer', () => { + setupForwardingListeners(mockProcessManager, { safeSend: mockSafeSend }); + + const handler = eventHandlers.get('slash-commands'); + const testSessionId = 'test-session-123'; + const testCommands = ['/help', '/clear']; + + handler?.(testSessionId, testCommands); + + expect(mockSafeSend).toHaveBeenCalledWith( + 'process:slash-commands', + testSessionId, + testCommands + ); + }); + + it('should forward thinking-chunk events to renderer', () => { + setupForwardingListeners(mockProcessManager, { safeSend: mockSafeSend }); + + const handler = eventHandlers.get('thinking-chunk'); + const testSessionId = 'test-session-123'; + const testChunk = { content: 'thinking...' }; + + handler?.(testSessionId, testChunk); + + expect(mockSafeSend).toHaveBeenCalledWith('process:thinking-chunk', testSessionId, testChunk); + }); + + it('should forward tool-execution events to renderer', () => { + setupForwardingListeners(mockProcessManager, { safeSend: mockSafeSend }); + + const handler = eventHandlers.get('tool-execution'); + const testSessionId = 'test-session-123'; + const testToolExecution = { tool: 'read_file', status: 'completed' }; + + handler?.(testSessionId, testToolExecution); + + expect(mockSafeSend).toHaveBeenCalledWith( + 'process:tool-execution', + testSessionId, + testToolExecution + ); + }); + + it('should forward stderr events to renderer', () => { + setupForwardingListeners(mockProcessManager, { safeSend: mockSafeSend }); + + const handler = eventHandlers.get('stderr'); + const testSessionId = 'test-session-123'; + const testStderr = 'Error: something went wrong'; + + handler?.(testSessionId, testStderr); + + expect(mockSafeSend).toHaveBeenCalledWith('process:stderr', testSessionId, testStderr); + }); + + it('should forward command-exit events to renderer', () => { + setupForwardingListeners(mockProcessManager, { safeSend: mockSafeSend }); + + const handler = eventHandlers.get('command-exit'); + const testSessionId = 'test-session-123'; + const testExitCode = 0; + + handler?.(testSessionId, testExitCode); + + expect(mockSafeSend).toHaveBeenCalledWith('process:command-exit', testSessionId, testExitCode); + }); +}); diff --git a/src/main/process-listeners/__tests__/stats-listener.test.ts b/src/main/process-listeners/__tests__/stats-listener.test.ts new file mode 100644 index 00000000..41e3836d --- /dev/null +++ b/src/main/process-listeners/__tests__/stats-listener.test.ts @@ -0,0 +1,180 @@ +/** + * Tests for stats listener. + * Handles query-complete events for usage statistics tracking. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupStatsListener } from '../stats-listener'; +import type { ProcessManager } from '../../process-manager'; +import type { SafeSendFn } from '../../utils/safe-send'; +import type { QueryCompleteData } from '../../process-manager/types'; +import type { StatsDB } from '../../stats-db'; +import type { ProcessListenerDependencies } from '../types'; + +describe('Stats Listener', () => { + let mockProcessManager: ProcessManager; + let mockSafeSend: SafeSendFn; + let mockStatsDB: StatsDB; + let mockLogger: ProcessListenerDependencies['logger']; + let eventHandlers: Map void>; + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockSafeSend = vi.fn(); + mockLogger = { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }; + + mockStatsDB = { + isReady: vi.fn(() => true), + insertQueryEvent: vi.fn(() => 1), + } as unknown as StatsDB; + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + }); + + it('should register the query-complete event listener', () => { + setupStatsListener(mockProcessManager, { + safeSend: mockSafeSend, + getStatsDB: () => mockStatsDB, + logger: mockLogger, + }); + + expect(mockProcessManager.on).toHaveBeenCalledWith('query-complete', expect.any(Function)); + }); + + it('should record query event to stats database when ready', () => { + setupStatsListener(mockProcessManager, { + safeSend: mockSafeSend, + getStatsDB: () => mockStatsDB, + logger: mockLogger, + }); + + const handler = eventHandlers.get('query-complete'); + const testSessionId = 'test-session-123'; + const testQueryData: QueryCompleteData = { + sessionId: testSessionId, + agentType: 'claude-code', + source: 'user', + startTime: Date.now() - 5000, + duration: 5000, + projectPath: '/test/project', + tabId: 'tab-123', + }; + + handler?.(testSessionId, testQueryData); + + expect(mockStatsDB.isReady).toHaveBeenCalled(); + expect(mockStatsDB.insertQueryEvent).toHaveBeenCalledWith({ + sessionId: testQueryData.sessionId, + agentType: testQueryData.agentType, + source: testQueryData.source, + startTime: testQueryData.startTime, + duration: testQueryData.duration, + projectPath: testQueryData.projectPath, + tabId: testQueryData.tabId, + }); + expect(mockSafeSend).toHaveBeenCalledWith('stats:updated'); + }); + + it('should not record event when stats database is not ready', () => { + vi.mocked(mockStatsDB.isReady).mockReturnValue(false); + + setupStatsListener(mockProcessManager, { + safeSend: mockSafeSend, + getStatsDB: () => mockStatsDB, + logger: mockLogger, + }); + + const handler = eventHandlers.get('query-complete'); + const testQueryData: QueryCompleteData = { + sessionId: 'session-456', + agentType: 'codex', + source: 'auto', + startTime: Date.now(), + duration: 1000, + projectPath: '/test/project', + tabId: 'tab-456', + }; + + handler?.('session-456', testQueryData); + + expect(mockStatsDB.isReady).toHaveBeenCalled(); + expect(mockStatsDB.insertQueryEvent).not.toHaveBeenCalled(); + expect(mockSafeSend).not.toHaveBeenCalled(); + }); + + it('should log error when recording fails', () => { + vi.mocked(mockStatsDB.insertQueryEvent).mockImplementation(() => { + throw new Error('Database error'); + }); + + setupStatsListener(mockProcessManager, { + safeSend: mockSafeSend, + getStatsDB: () => mockStatsDB, + logger: mockLogger, + }); + + const handler = eventHandlers.get('query-complete'); + const testQueryData: QueryCompleteData = { + sessionId: 'session-789', + agentType: 'opencode', + source: 'user', + startTime: Date.now(), + duration: 2000, + projectPath: '/test/project', + tabId: 'tab-789', + }; + + handler?.('session-789', testQueryData); + + expect(mockLogger.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to record query event'), + '[Stats]', + expect.objectContaining({ + sessionId: 'session-789', + }) + ); + }); + + it('should log debug info when recording succeeds', () => { + setupStatsListener(mockProcessManager, { + safeSend: mockSafeSend, + getStatsDB: () => mockStatsDB, + logger: mockLogger, + }); + + const handler = eventHandlers.get('query-complete'); + const testQueryData: QueryCompleteData = { + sessionId: 'session-abc', + agentType: 'claude-code', + source: 'user', + startTime: Date.now(), + duration: 3000, + projectPath: '/test/project', + tabId: 'tab-abc', + }; + + handler?.('session-abc', testQueryData); + + expect(mockLogger.debug).toHaveBeenCalledWith( + expect.stringContaining('Recorded query event'), + '[Stats]', + expect.objectContaining({ + sessionId: 'session-abc', + agentType: 'claude-code', + source: 'user', + duration: 3000, + }) + ); + }); +}); diff --git a/src/main/process-listeners/data-listener.ts b/src/main/process-listeners/data-listener.ts new file mode 100644 index 00000000..71b2570d --- /dev/null +++ b/src/main/process-listeners/data-listener.ts @@ -0,0 +1,102 @@ +/** + * Data output listener. + * Handles process output data, including group chat buffering and web broadcasting. + */ + +import type { ProcessManager } from '../process-manager'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up the data listener for process output. + * Handles: + * - Group chat moderator/participant output buffering + * - Regular process data forwarding to renderer + * - Web broadcast to connected clients + */ +export function setupDataListener( + processManager: ProcessManager, + deps: Pick< + ProcessListenerDependencies, + 'safeSend' | 'getWebServer' | 'outputBuffer' | 'outputParser' | 'debugLog' | 'patterns' + > +): void { + const { safeSend, getWebServer, outputBuffer, outputParser, debugLog, patterns } = deps; + const { REGEX_MODERATOR_SESSION, REGEX_AI_SUFFIX, REGEX_AI_TAB_ID } = patterns; + + processManager.on('data', (sessionId: string, data: string) => { + // Handle group chat moderator output - buffer it + // Session ID format: group-chat-{groupChatId}-moderator-{uuid} or group-chat-{groupChatId}-moderator-synthesis-{uuid} + const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION); + if (moderatorMatch) { + const groupChatId = moderatorMatch[1]; + debugLog('GroupChat:Debug', `MODERATOR DATA received for chat ${groupChatId}`); + debugLog('GroupChat:Debug', `Session ID: ${sessionId}`); + debugLog('GroupChat:Debug', `Data length: ${data.length}`); + // Buffer the output - will be routed on process exit + const totalLength = outputBuffer.appendToGroupChatBuffer(sessionId, data); + debugLog('GroupChat:Debug', `Buffered total: ${totalLength} chars`); + return; // Don't send to regular process:data handler + } + + // Handle group chat participant output - buffer it + // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} + const participantInfo = outputParser.parseParticipantSessionId(sessionId); + if (participantInfo) { + debugLog('GroupChat:Debug', 'PARTICIPANT DATA received'); + debugLog( + 'GroupChat:Debug', + `Chat: ${participantInfo.groupChatId}, Participant: ${participantInfo.participantName}` + ); + debugLog('GroupChat:Debug', `Session ID: ${sessionId}`); + debugLog('GroupChat:Debug', `Data length: ${data.length}`); + // Buffer the output - will be routed on process exit + const totalLength = outputBuffer.appendToGroupChatBuffer(sessionId, data); + debugLog('GroupChat:Debug', `Buffered total: ${totalLength} chars`); + return; // Don't send to regular process:data handler + } + + safeSend('process:data', sessionId, data); + + // Broadcast to web clients - extract base session ID (remove -ai or -terminal suffix) + // IMPORTANT: Skip PTY terminal output (-terminal suffix) as it contains raw ANSI codes. + // Web interface terminal commands use runCommand() which emits with plain session IDs. + const webServer = getWebServer(); + if (webServer) { + // Don't broadcast raw PTY terminal output to web clients + if (sessionId.endsWith('-terminal')) { + debugLog('WebBroadcast', `SKIPPING PTY terminal output for web: session=${sessionId}`); + return; + } + + // Don't broadcast background batch/synopsis output to web clients + // These are internal Auto Run operations that should only appear in history, not as chat messages + if (sessionId.includes('-batch-') || sessionId.includes('-synopsis-')) { + debugLog('WebBroadcast', `SKIPPING batch/synopsis output for web: session=${sessionId}`); + return; + } + + // Extract base session ID and tab ID from format: {id}-ai-{tabId} + const baseSessionId = sessionId.replace(REGEX_AI_SUFFIX, ''); + const isAiOutput = sessionId.includes('-ai-'); + + // Extract tab ID from session ID format: {id}-ai-{tabId} + const tabIdMatch = sessionId.match(REGEX_AI_TAB_ID); + const tabId = tabIdMatch ? tabIdMatch[1] : undefined; + + const msgId = `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; + debugLog( + 'WebBroadcast', + `Broadcasting session_output: msgId=${msgId}, session=${baseSessionId}, tabId=${tabId || 'none'}, source=${isAiOutput ? 'ai' : 'terminal'}, dataLen=${data.length}` + ); + webServer.broadcastToSessionClients(baseSessionId, { + type: 'session_output', + sessionId: baseSessionId, + tabId, + data, + source: isAiOutput ? 'ai' : 'terminal', + timestamp: Date.now(), + msgId, + }); + } + }); +} diff --git a/src/main/process-listeners/error-listener.ts b/src/main/process-listeners/error-listener.ts new file mode 100644 index 00000000..04ce341c --- /dev/null +++ b/src/main/process-listeners/error-listener.ts @@ -0,0 +1,31 @@ +/** + * Agent error listener. + * Handles agent errors (auth expired, token exhaustion, rate limits, etc.). + */ + +import type { ProcessManager } from '../process-manager'; +import type { AgentError } from '../../shared/types'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up the agent-error listener. + * Handles logging and forwarding of agent errors to renderer. + */ +export function setupErrorListener( + processManager: ProcessManager, + deps: Pick +): void { + const { safeSend, logger } = deps; + + // Handle agent errors (auth expired, token exhaustion, rate limits, etc.) + processManager.on('agent-error', (sessionId: string, agentError: AgentError) => { + logger.info(`Agent error detected: ${agentError.type}`, 'AgentError', { + sessionId, + agentId: agentError.agentId, + errorType: agentError.type, + message: agentError.message, + recoverable: agentError.recoverable, + }); + safeSend('agent:error', sessionId, agentError); + }); +} diff --git a/src/main/process-listeners/exit-listener.ts b/src/main/process-listeners/exit-listener.ts new file mode 100644 index 00000000..2f20b54f --- /dev/null +++ b/src/main/process-listeners/exit-listener.ts @@ -0,0 +1,411 @@ +/** + * Process exit listener. + * Handles process exit events, including group chat moderator/participant exits. + * This is the largest and most complex listener with routing, recovery, and synthesis logic. + */ + +import type { ProcessManager } from '../process-manager'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up the exit listener for process termination. + * Handles: + * - Power management cleanup + * - Group chat moderator exit (routing buffered output) + * - Group chat participant exit (routing, recovery, synthesis triggering) + * - Regular process exit forwarding + * - Web broadcast of exit events + */ +export function setupExitListener( + processManager: ProcessManager, + deps: Pick< + ProcessListenerDependencies, + | 'safeSend' + | 'getProcessManager' + | 'getAgentDetector' + | 'getWebServer' + | 'powerManager' + | 'outputBuffer' + | 'outputParser' + | 'groupChatEmitters' + | 'groupChatRouter' + | 'groupChatStorage' + | 'sessionRecovery' + | 'debugLog' + | 'logger' + | 'patterns' + > +): void { + const { + safeSend, + getProcessManager, + getAgentDetector, + getWebServer, + powerManager, + outputBuffer, + outputParser, + groupChatEmitters, + groupChatRouter, + groupChatStorage, + sessionRecovery, + debugLog, + logger, + patterns, + } = deps; + const { REGEX_MODERATOR_SESSION } = patterns; + + processManager.on('exit', (sessionId: string, code: number) => { + // Remove power block reason for this session + // This allows system sleep when no AI sessions are active + powerManager.removeBlockReason(`session:${sessionId}`); + + // Handle group chat moderator exit - route buffered output and set state back to idle + // Session ID format: group-chat-{groupChatId}-moderator-{uuid} + // This handles BOTH initial moderator responses AND synthesis responses. + // The routeModeratorResponse function will check for @mentions: + // - If @mentions present: route to agents (continue conversation) + // - If no @mentions: final response to user (conversation complete for this turn) + const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION); + if (moderatorMatch) { + const groupChatId = moderatorMatch[1]; + debugLog('GroupChat:Debug', ` ========== MODERATOR PROCESS EXIT ==========`); + debugLog('GroupChat:Debug', ` Group Chat ID: ${groupChatId}`); + debugLog('GroupChat:Debug', ` Session ID: ${sessionId}`); + debugLog('GroupChat:Debug', ` Exit code: ${code}`); + logger.debug(`[GroupChat] Moderator exit: groupChatId=${groupChatId}`, 'ProcessListener', { + sessionId, + }); + // Route the buffered output now that process is complete + const bufferedOutput = outputBuffer.getGroupChatBufferedOutput(sessionId); + debugLog('GroupChat:Debug', ` Buffered output length: ${bufferedOutput?.length ?? 0}`); + if (bufferedOutput) { + debugLog( + 'GroupChat:Debug', + ` Raw buffered output preview: "${bufferedOutput.substring(0, 300)}${bufferedOutput.length > 300 ? '...' : ''}"` + ); + logger.debug( + `[GroupChat] Moderator has buffered output (${bufferedOutput.length} chars)`, + 'ProcessListener', + { groupChatId } + ); + void (async () => { + try { + const chat = await groupChatStorage.loadGroupChat(groupChatId); + debugLog('GroupChat:Debug', ` Chat loaded for parsing: ${chat?.name || 'null'}`); + const agentType = chat?.moderatorAgentId; + debugLog('GroupChat:Debug', ` Agent type for parsing: ${agentType}`); + const parsedText = outputParser.extractTextFromStreamJson(bufferedOutput, agentType); + debugLog('GroupChat:Debug', ` Parsed text length: ${parsedText.length}`); + debugLog( + 'GroupChat:Debug', + ` Parsed text preview: "${parsedText.substring(0, 300)}${parsedText.length > 300 ? '...' : ''}"` + ); + if (parsedText.trim()) { + debugLog('GroupChat:Debug', ` Routing moderator response...`); + logger.info( + `[GroupChat] Routing moderator response (${parsedText.length} chars)`, + 'ProcessListener', + { groupChatId } + ); + const readOnly = groupChatRouter.getGroupChatReadOnlyState(groupChatId); + debugLog('GroupChat:Debug', ` Read-only state: ${readOnly}`); + const pm = getProcessManager(); + const ad = getAgentDetector(); + groupChatRouter + .routeModeratorResponse( + groupChatId, + parsedText, + pm ?? undefined, + ad ?? undefined, + readOnly + ) + .catch((err) => { + debugLog('GroupChat:Debug', ` ERROR routing moderator response:`, err); + logger.error( + '[GroupChat] Failed to route moderator response', + 'ProcessListener', + { error: String(err) } + ); + }); + } else { + debugLog('GroupChat:Debug', ` WARNING: Parsed text is empty!`); + logger.warn( + '[GroupChat] Moderator output parsed to empty string', + 'ProcessListener', + { groupChatId, bufferedLength: bufferedOutput.length } + ); + } + } catch (err) { + debugLog('GroupChat:Debug', ` ERROR loading chat:`, err); + logger.error( + '[GroupChat] Failed to load chat for moderator output parsing', + 'ProcessListener', + { error: String(err) } + ); + const parsedText = outputParser.extractTextFromStreamJson(bufferedOutput); + if (parsedText.trim()) { + const readOnly = groupChatRouter.getGroupChatReadOnlyState(groupChatId); + const pm = getProcessManager(); + const ad = getAgentDetector(); + groupChatRouter + .routeModeratorResponse( + groupChatId, + parsedText, + pm ?? undefined, + ad ?? undefined, + readOnly + ) + .catch((routeErr) => { + debugLog( + 'GroupChat:Debug', + ` ERROR routing moderator response (fallback):`, + routeErr + ); + logger.error( + '[GroupChat] Failed to route moderator response', + 'ProcessListener', + { error: String(routeErr) } + ); + }); + } + } + })().finally(() => { + outputBuffer.clearGroupChatBuffer(sessionId); + debugLog('GroupChat:Debug', ` Cleared output buffer for session`); + }); + } else { + debugLog('GroupChat:Debug', ` WARNING: No buffered output!`); + logger.warn('[GroupChat] Moderator exit with no buffered output', 'ProcessListener', { + groupChatId, + sessionId, + }); + } + groupChatEmitters.emitStateChange?.(groupChatId, 'idle'); + debugLog('GroupChat:Debug', ` Emitted state change: idle`); + debugLog('GroupChat:Debug', ` =============================================`); + // Don't send to regular exit handler + return; + } + + // Handle group chat participant exit - route buffered output and update participant state + // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} + const participantExitInfo = outputParser.parseParticipantSessionId(sessionId); + if (participantExitInfo) { + const { groupChatId, participantName } = participantExitInfo; + debugLog('GroupChat:Debug', ` ========== PARTICIPANT PROCESS EXIT ==========`); + debugLog('GroupChat:Debug', ` Group Chat ID: ${groupChatId}`); + debugLog('GroupChat:Debug', ` Participant: ${participantName}`); + debugLog('GroupChat:Debug', ` Session ID: ${sessionId}`); + debugLog('GroupChat:Debug', ` Exit code: ${code}`); + logger.debug( + `[GroupChat] Participant exit: ${participantName} (groupChatId=${groupChatId})`, + 'ProcessListener', + { sessionId } + ); + + // Emit participant state change to show this participant is done working + groupChatEmitters.emitParticipantState?.(groupChatId, participantName, 'idle'); + debugLog('GroupChat:Debug', ` Emitted participant state: idle`); + + // Route the buffered output now that process is complete + // IMPORTANT: We must wait for the response to be logged before triggering synthesis + // to avoid a race condition where synthesis reads the log before the response is written + const bufferedOutput = outputBuffer.getGroupChatBufferedOutput(sessionId); + debugLog('GroupChat:Debug', ` Buffered output length: ${bufferedOutput?.length ?? 0}`); + + // Helper function to mark participant and potentially trigger synthesis + const markAndMaybeSynthesize = () => { + const isLastParticipant = groupChatRouter.markParticipantResponded( + groupChatId, + participantName + ); + debugLog('GroupChat:Debug', ` Is last participant to respond: ${isLastParticipant}`); + const pm = getProcessManager(); + const ad = getAgentDetector(); + if (isLastParticipant && pm && ad) { + // All participants have responded - spawn moderator synthesis round + debugLog('GroupChat:Debug', ` All participants responded - spawning synthesis round...`); + logger.info( + '[GroupChat] All participants responded, spawning moderator synthesis', + 'ProcessListener', + { groupChatId } + ); + groupChatRouter.spawnModeratorSynthesis(groupChatId, pm, ad).catch((err) => { + debugLog('GroupChat:Debug', ` ERROR spawning synthesis:`, err); + logger.error('[GroupChat] Failed to spawn moderator synthesis', 'ProcessListener', { + error: String(err), + groupChatId, + }); + }); + } else if (!isLastParticipant) { + // More participants pending + debugLog('GroupChat:Debug', ` Waiting for more participants to respond...`); + } + }; + + if (bufferedOutput) { + debugLog( + 'GroupChat:Debug', + ` Raw buffered output preview: "${bufferedOutput.substring(0, 300)}${bufferedOutput.length > 300 ? '...' : ''}"` + ); + + // Handle session recovery and normal processing in an async IIFE + void (async () => { + // Check if this is a session_not_found error - if so, recover and retry + const chat = await groupChatStorage.loadGroupChat(groupChatId); + const agentType = chat?.participants.find((p) => p.name === participantName)?.agentId; + + if (sessionRecovery.needsSessionRecovery(bufferedOutput, agentType)) { + debugLog( + 'GroupChat:Debug', + ` Session not found error detected for ${participantName} - initiating recovery` + ); + logger.info('[GroupChat] Session recovery needed', 'ProcessListener', { + groupChatId, + participantName, + }); + + // Clear the buffer first + outputBuffer.clearGroupChatBuffer(sessionId); + + // Initiate recovery (clears agentSessionId) + await sessionRecovery.initiateSessionRecovery(groupChatId, participantName); + + // Re-spawn the participant with recovery context + const pm = getProcessManager(); + const ad = getAgentDetector(); + if (pm && ad) { + debugLog( + 'GroupChat:Debug', + ` Re-spawning ${participantName} with recovery context...` + ); + try { + await groupChatRouter.respawnParticipantWithRecovery( + groupChatId, + participantName, + pm, + ad + ); + debugLog( + 'GroupChat:Debug', + ` Successfully re-spawned ${participantName} for recovery` + ); + // Don't mark as responded yet - the recovery spawn will complete and trigger this + } catch (respawnErr) { + debugLog('GroupChat:Debug', ` Failed to respawn ${participantName}:`, respawnErr); + logger.error( + '[GroupChat] Failed to respawn participant for recovery', + 'ProcessListener', + { + error: String(respawnErr), + participant: participantName, + } + ); + // Mark as responded since recovery failed + markAndMaybeSynthesize(); + } + } else { + debugLog( + 'GroupChat:Debug', + ` Cannot respawn - processManager or agentDetector not available` + ); + markAndMaybeSynthesize(); + } + debugLog('GroupChat:Debug', ` ===============================================`); + return; + } + + // Normal processing - parse and route the response + try { + debugLog( + 'GroupChat:Debug', + ` Chat loaded for participant parsing: ${chat?.name || 'null'}` + ); + debugLog('GroupChat:Debug', ` Agent type for parsing: ${agentType}`); + const parsedText = outputParser.extractTextFromStreamJson(bufferedOutput, agentType); + debugLog('GroupChat:Debug', ` Parsed text length: ${parsedText.length}`); + debugLog( + 'GroupChat:Debug', + ` Parsed text preview: "${parsedText.substring(0, 200)}${parsedText.length > 200 ? '...' : ''}"` + ); + if (parsedText.trim()) { + debugLog('GroupChat:Debug', ` Routing agent response from ${participantName}...`); + // Await the response logging before marking participant as responded + const pm = getProcessManager(); + await groupChatRouter.routeAgentResponse( + groupChatId, + participantName, + parsedText, + pm ?? undefined + ); + debugLog( + 'GroupChat:Debug', + ` Successfully routed agent response from ${participantName}` + ); + } else { + debugLog('GroupChat:Debug', ` WARNING: Parsed text is empty for ${participantName}!`); + } + } catch (err) { + debugLog('GroupChat:Debug', ` ERROR loading chat for participant:`, err); + logger.error( + '[GroupChat] Failed to load chat for participant output parsing', + 'ProcessListener', + { error: String(err), participant: participantName } + ); + try { + const parsedText = outputParser.extractTextFromStreamJson(bufferedOutput); + if (parsedText.trim()) { + const pm = getProcessManager(); + await groupChatRouter.routeAgentResponse( + groupChatId, + participantName, + parsedText, + pm ?? undefined + ); + } + } catch (routeErr) { + debugLog('GroupChat:Debug', ` ERROR routing agent response (fallback):`, routeErr); + logger.error('[GroupChat] Failed to route agent response', 'ProcessListener', { + error: String(routeErr), + participant: participantName, + }); + } + } + })().finally(() => { + outputBuffer.clearGroupChatBuffer(sessionId); + debugLog('GroupChat:Debug', ` Cleared output buffer for participant session`); + // Mark participant and trigger synthesis AFTER logging is complete + markAndMaybeSynthesize(); + }); + } else { + debugLog( + 'GroupChat:Debug', + ` WARNING: No buffered output for participant ${participantName}!` + ); + // No output to log, so mark participant as responded immediately + markAndMaybeSynthesize(); + } + debugLog('GroupChat:Debug', ` ===============================================`); + // Don't send to regular exit handler + return; + } + + safeSend('process:exit', sessionId, code); + + // Broadcast exit to web clients + const webServer = getWebServer(); + if (webServer) { + // Extract base session ID from formats: {id}-ai-{tabId}, {id}-terminal, {id}-batch-{timestamp}, {id}-synopsis-{timestamp} + const baseSessionId = sessionId.replace( + /-ai-[^-]+$|-terminal$|-batch-\d+$|-synopsis-\d+$/, + '' + ); + webServer.broadcastToSessionClients(baseSessionId, { + type: 'session_exit', + sessionId: baseSessionId, + exitCode: code, + timestamp: Date.now(), + }); + } + }); +} diff --git a/src/main/process-listeners/forwarding-listeners.ts b/src/main/process-listeners/forwarding-listeners.ts new file mode 100644 index 00000000..3ec6bb4c --- /dev/null +++ b/src/main/process-listeners/forwarding-listeners.ts @@ -0,0 +1,48 @@ +/** + * Simple IPC forwarding listeners. + * These listeners just forward events from ProcessManager to the renderer. + */ + +import type { ProcessManager } from '../process-manager'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up simple forwarding listeners that pass events directly to renderer. + * These are lightweight handlers that don't require any processing logic. + */ +export function setupForwardingListeners( + processManager: ProcessManager, + deps: Pick +): void { + const { safeSend } = deps; + + // Handle slash commands from Claude Code init message + processManager.on('slash-commands', (sessionId: string, slashCommands: string[]) => { + safeSend('process:slash-commands', sessionId, slashCommands); + }); + + // Handle thinking/streaming content chunks from AI agents + // Emitted when agents produce partial text events (isPartial: true) + // Renderer decides whether to display based on tab's showThinking setting + processManager.on('thinking-chunk', (sessionId: string, content: string) => { + safeSend('process:thinking-chunk', sessionId, content); + }); + + // Handle tool execution events (OpenCode, Codex) + processManager.on( + 'tool-execution', + (sessionId: string, toolEvent: { toolName: string; state?: unknown; timestamp: number }) => { + safeSend('process:tool-execution', sessionId, toolEvent); + } + ); + + // Handle stderr separately from runCommand (for clean command execution) + processManager.on('stderr', (sessionId: string, data: string) => { + safeSend('process:stderr', sessionId, data); + }); + + // Handle command exit (from runCommand - separate from PTY exit) + processManager.on('command-exit', (sessionId: string, code: number) => { + safeSend('process:command-exit', sessionId, code); + }); +} diff --git a/src/main/process-listeners/index.ts b/src/main/process-listeners/index.ts new file mode 100644 index 00000000..06882f3d --- /dev/null +++ b/src/main/process-listeners/index.ts @@ -0,0 +1,55 @@ +/** + * Process event listeners module. + * Handles all events emitted by ProcessManager and routes them appropriately. + * + * This module extracts the setupProcessListeners() logic from main/index.ts + * into smaller, focused modules for better maintainability. + */ + +import type { ProcessManager } from '../process-manager'; +import type { ProcessListenerDependencies } from './types'; + +// Import individual listener setup functions +import { setupForwardingListeners } from './forwarding-listeners'; +import { setupDataListener } from './data-listener'; +import { setupUsageListener } from './usage-listener'; +import { setupSessionIdListener } from './session-id-listener'; +import { setupErrorListener } from './error-listener'; +import { setupStatsListener } from './stats-listener'; +import { setupExitListener } from './exit-listener'; + +// Re-export types for consumers +export type { ProcessListenerDependencies, ParticipantInfo } from './types'; + +/** + * Sets up all process event listeners. + * This is the main entry point that orchestrates all listener modules. + * + * @param processManager - The ProcessManager instance to attach listeners to + * @param deps - Dependencies for the listeners + */ +export function setupProcessListeners( + processManager: ProcessManager, + deps: ProcessListenerDependencies +): void { + // Simple forwarding listeners (slash-commands, thinking-chunk, tool-execution, stderr, command-exit) + setupForwardingListeners(processManager, deps); + + // Data output listener (with group chat buffering and web broadcast) + setupDataListener(processManager, deps); + + // Usage statistics listener (with group chat participant/moderator updates) + setupUsageListener(processManager, deps); + + // Session ID listener (with group chat participant/moderator storage) + setupSessionIdListener(processManager, deps); + + // Agent error listener + setupErrorListener(processManager, deps); + + // Stats/query-complete listener + setupStatsListener(processManager, deps); + + // Exit listener (with group chat routing, recovery, and synthesis) + setupExitListener(processManager, deps); +} diff --git a/src/main/process-listeners/session-id-listener.ts b/src/main/process-listeners/session-id-listener.ts new file mode 100644 index 00000000..7d70e394 --- /dev/null +++ b/src/main/process-listeners/session-id-listener.ts @@ -0,0 +1,77 @@ +/** + * Session ID listener. + * Handles agent session ID events for group chat participant/moderator tracking. + */ + +import type { ProcessManager } from '../process-manager'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up the session-id listener. + * Handles: + * - Group chat participant session ID storage + * - Group chat moderator session ID storage + * - Regular session ID forwarding to renderer + */ +export function setupSessionIdListener( + processManager: ProcessManager, + deps: Pick< + ProcessListenerDependencies, + 'safeSend' | 'outputParser' | 'groupChatEmitters' | 'groupChatStorage' | 'logger' | 'patterns' + > +): void { + const { safeSend, outputParser, groupChatEmitters, groupChatStorage, logger, patterns } = deps; + const { REGEX_MODERATOR_SESSION_TIMESTAMP } = patterns; + + processManager.on('session-id', (sessionId: string, agentSessionId: string) => { + // Handle group chat participant session ID - store the agent's session ID + // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} + const participantSessionInfo = outputParser.parseParticipantSessionId(sessionId); + if (participantSessionInfo) { + const { groupChatId, participantName } = participantSessionInfo; + // Update the participant with the agent's session ID + groupChatStorage + .updateParticipant(groupChatId, participantName, { agentSessionId }) + .then(async () => { + // Emit participants changed so UI updates with the new session ID + const chat = await groupChatStorage.loadGroupChat(groupChatId); + if (chat) { + groupChatEmitters.emitParticipantsChanged?.(groupChatId, chat.participants); + } + }) + .catch((err) => { + logger.error( + '[GroupChat] Failed to update participant agentSessionId', + 'ProcessListener', + { error: String(err), participant: participantName } + ); + }); + // Don't return - still send to renderer for logging purposes + } + + // Handle group chat moderator session ID - store the real agent session ID + // Session ID format: group-chat-{groupChatId}-moderator-{timestamp} + const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION_TIMESTAMP); + if (moderatorMatch) { + const groupChatId = moderatorMatch[1]; + // Update the group chat with the moderator's real agent session ID + // Store in moderatorAgentSessionId (not moderatorSessionId which is the routing prefix) + groupChatStorage + .updateGroupChat(groupChatId, { moderatorAgentSessionId: agentSessionId }) + .then(() => { + // Emit session ID change event so UI updates with the new session ID + groupChatEmitters.emitModeratorSessionIdChanged?.(groupChatId, agentSessionId); + }) + .catch((err: unknown) => { + logger.error( + '[GroupChat] Failed to update moderator agent session ID', + 'ProcessListener', + { error: String(err), groupChatId } + ); + }); + // Don't return - still send to renderer for logging purposes + } + + safeSend('process:session-id', sessionId, agentSessionId); + }); +} diff --git a/src/main/process-listeners/stats-listener.ts b/src/main/process-listeners/stats-listener.ts new file mode 100644 index 00000000..f831ef20 --- /dev/null +++ b/src/main/process-listeners/stats-listener.ts @@ -0,0 +1,50 @@ +/** + * Stats listener. + * Handles query-complete events for usage statistics tracking. + */ + +import type { ProcessManager } from '../process-manager'; +import type { QueryCompleteData } from '../process-manager/types'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up the query-complete listener for stats tracking. + * Records AI query events to the stats database. + */ +export function setupStatsListener( + processManager: ProcessManager, + deps: Pick +): void { + const { safeSend, getStatsDB, logger } = deps; + + // Handle query-complete events for stats tracking + // This is emitted when a batch mode AI query completes (user or auto) + processManager.on('query-complete', (_sessionId: string, queryData: QueryCompleteData) => { + try { + const db = getStatsDB(); + if (db.isReady()) { + const id = db.insertQueryEvent({ + sessionId: queryData.sessionId, + agentType: queryData.agentType, + source: queryData.source, + startTime: queryData.startTime, + duration: queryData.duration, + projectPath: queryData.projectPath, + tabId: queryData.tabId, + }); + logger.debug(`Recorded query event: ${id}`, '[Stats]', { + sessionId: queryData.sessionId, + agentType: queryData.agentType, + source: queryData.source, + duration: queryData.duration, + }); + // Broadcast stats update to renderer for real-time dashboard refresh + safeSend('stats:updated'); + } + } catch (error) { + logger.error(`Failed to record query event: ${error}`, '[Stats]', { + sessionId: queryData.sessionId, + }); + } + }); +} diff --git a/src/main/process-listeners/types.ts b/src/main/process-listeners/types.ts new file mode 100644 index 00000000..b230d739 --- /dev/null +++ b/src/main/process-listeners/types.ts @@ -0,0 +1,150 @@ +/** + * Type definitions for process event listeners. + * Re-exports existing types and defines the dependency interface. + */ + +import type { ProcessManager } from '../process-manager'; +import type { WebServer } from '../web-server'; +import type { AgentDetector } from '../agent-detector'; +import type { SafeSendFn } from '../utils/safe-send'; +import type { StatsDB } from '../stats-db'; +import type { GroupChat, GroupChatParticipant } from '../group-chat/group-chat-storage'; +import type { GroupChatState } from '../../shared/group-chat-types'; +import type { ParticipantState } from '../ipc/handlers/groupChat'; + +// Re-export types from their canonical locations +export type { UsageStats, QueryCompleteData, ToolExecution } from '../process-manager/types'; +export type { AgentError } from '../../shared/types'; +export type { GroupChat, GroupChatParticipant }; +export type { SafeSendFn } from '../utils/safe-send'; +export type { GroupChatState }; +export type { ParticipantState }; + +// Import emitters and state types from groupChat handlers +export type { groupChatEmitters, ModeratorUsage } from '../ipc/handlers/groupChat'; + +/** + * Participant info parsed from session ID. + * Matches return type of parseParticipantSessionId. + */ +export interface ParticipantInfo { + groupChatId: string; + participantName: string; +} + +/** + * Dependencies for process event listeners. + * All external dependencies are injected to enable testing and modularity. + */ +export interface ProcessListenerDependencies { + /** Function to get the process manager */ + getProcessManager: () => ProcessManager | null; + /** Function to get the web server (may be null if not started) */ + getWebServer: () => WebServer | null; + /** Function to get the agent detector */ + getAgentDetector: () => AgentDetector | null; + /** Safe send function for IPC messages */ + safeSend: SafeSendFn; + /** Power manager instance */ + powerManager: { + addBlockReason: (reason: string) => void; + removeBlockReason: (reason: string) => void; + }; + /** Group chat event emitters */ + groupChatEmitters: { + emitStateChange?: (groupChatId: string, state: GroupChatState) => void; + emitParticipantState?: ( + groupChatId: string, + participantName: string, + state: ParticipantState + ) => void; + emitParticipantsChanged?: (groupChatId: string, participants: GroupChatParticipant[]) => void; + emitModeratorSessionIdChanged?: (groupChatId: string, agentSessionId: string) => void; + emitModeratorUsage?: ( + groupChatId: string, + usage: { contextUsage: number; totalCost: number; tokenCount: number } + ) => void; + }; + /** Group chat router functions */ + groupChatRouter: { + routeModeratorResponse: ( + groupChatId: string, + text: string, + processManager: ProcessManager | undefined, + agentDetector: AgentDetector | undefined, + readOnly: boolean + ) => Promise; + routeAgentResponse: ( + groupChatId: string, + participantName: string, + text: string, + processManager: ProcessManager | undefined + ) => Promise; + markParticipantResponded: (groupChatId: string, participantName: string) => boolean; + spawnModeratorSynthesis: ( + groupChatId: string, + processManager: ProcessManager, + agentDetector: AgentDetector + ) => Promise; + getGroupChatReadOnlyState: (groupChatId: string) => boolean; + respawnParticipantWithRecovery: ( + groupChatId: string, + participantName: string, + processManager: ProcessManager, + agentDetector: AgentDetector + ) => Promise; + }; + /** Group chat storage functions */ + groupChatStorage: { + loadGroupChat: (groupChatId: string) => Promise; + updateGroupChat: (groupChatId: string, updates: Record) => Promise; + updateParticipant: ( + groupChatId: string, + participantName: string, + updates: Record + ) => Promise; + }; + /** Session recovery functions */ + sessionRecovery: { + needsSessionRecovery: (output: string, agentType?: string) => boolean; + initiateSessionRecovery: (groupChatId: string, participantName: string) => Promise; + }; + /** Output buffer functions */ + outputBuffer: { + appendToGroupChatBuffer: (sessionId: string, data: string) => number; + getGroupChatBufferedOutput: (sessionId: string) => string | undefined; + clearGroupChatBuffer: (sessionId: string) => void; + }; + /** Output parser functions */ + outputParser: { + extractTextFromStreamJson: (output: string, agentType?: string) => string; + parseParticipantSessionId: (sessionId: string) => ParticipantInfo | null; + }; + /** Usage aggregator functions */ + usageAggregator: { + calculateContextTokens: (usageStats: { + inputTokens: number; + outputTokens: number; + cacheReadInputTokens: number; + cacheCreationInputTokens: number; + }) => number; + }; + /** Stats database getter */ + getStatsDB: () => StatsDB; + /** Debug log function */ + debugLog: (prefix: string, message: string, ...args: unknown[]) => void; + /** Regex patterns */ + patterns: { + REGEX_MODERATOR_SESSION: RegExp; + REGEX_MODERATOR_SESSION_TIMESTAMP: RegExp; + REGEX_AI_SUFFIX: RegExp; + REGEX_AI_TAB_ID: RegExp; + }; + /** Logger instance */ + logger: { + info: (message: string, context: string, data?: Record) => void; + error: (message: string, context: string, data?: Record) => void; + warn: (message: string, context: string, data?: Record) => void; + debug: (message: string, context: string, data?: Record) => void; + }; +} diff --git a/src/main/process-listeners/usage-listener.ts b/src/main/process-listeners/usage-listener.ts new file mode 100644 index 00000000..1330bd42 --- /dev/null +++ b/src/main/process-listeners/usage-listener.ts @@ -0,0 +1,114 @@ +/** + * Usage statistics listener. + * Handles usage stats from AI responses, including group chat participant/moderator updates. + */ + +import type { ProcessManager } from '../process-manager'; +import type { ProcessListenerDependencies } from './types'; + +/** + * Sets up the usage listener for token/cost statistics. + * Handles: + * - Group chat participant usage updates + * - Group chat moderator usage updates + * - Regular process usage forwarding to renderer + */ +export function setupUsageListener( + processManager: ProcessManager, + deps: Pick< + ProcessListenerDependencies, + | 'safeSend' + | 'outputParser' + | 'groupChatEmitters' + | 'groupChatStorage' + | 'usageAggregator' + | 'logger' + | 'patterns' + > +): void { + const { + safeSend, + outputParser, + groupChatEmitters, + groupChatStorage, + usageAggregator, + logger, + patterns, + } = deps; + const { REGEX_MODERATOR_SESSION } = patterns; + + // Handle usage statistics from AI responses + processManager.on( + 'usage', + ( + sessionId: string, + usageStats: { + inputTokens: number; + outputTokens: number; + cacheReadInputTokens: number; + cacheCreationInputTokens: number; + totalCostUsd: number; + contextWindow: number; + reasoningTokens?: number; // Separate reasoning tokens (Codex o3/o4-mini) + } + ) => { + // Handle group chat participant usage - update participant stats + const participantUsageInfo = outputParser.parseParticipantSessionId(sessionId); + if (participantUsageInfo) { + const { groupChatId, participantName } = participantUsageInfo; + + // Calculate context usage percentage using agent-specific logic + // Note: For group chat, we don't have agent type here, defaults to Claude behavior + const totalContextTokens = usageAggregator.calculateContextTokens(usageStats); + const contextUsage = + usageStats.contextWindow > 0 + ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) + : 0; + + // Update participant with usage stats + groupChatStorage + .updateParticipant(groupChatId, participantName, { + contextUsage, + tokenCount: totalContextTokens, + totalCost: usageStats.totalCostUsd, + }) + .then(async () => { + // Emit participants changed so UI updates + const chat = await groupChatStorage.loadGroupChat(groupChatId); + if (chat) { + groupChatEmitters.emitParticipantsChanged?.(groupChatId, chat.participants); + } + }) + .catch((err) => { + logger.error('[GroupChat] Failed to update participant usage', 'ProcessListener', { + error: String(err), + participant: participantName, + }); + }); + // Still send to renderer for consistency + } + + // Handle group chat moderator usage - emit for UI + const moderatorUsageMatch = sessionId.match(REGEX_MODERATOR_SESSION); + if (moderatorUsageMatch) { + const groupChatId = moderatorUsageMatch[1]; + // Calculate context usage percentage using agent-specific logic + // Note: Moderator is typically Claude, defaults to Claude behavior + const totalContextTokens = usageAggregator.calculateContextTokens(usageStats); + const contextUsage = + usageStats.contextWindow > 0 + ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) + : 0; + + // Emit moderator usage for the moderator card + groupChatEmitters.emitModeratorUsage?.(groupChatId, { + contextUsage, + totalCost: usageStats.totalCostUsd, + tokenCount: totalContextTokens, + }); + } + + safeSend('process:usage', sessionId, usageStats); + } + ); +} From 84ed2d4b56cccfb1d5fe21189f20422f56ee8a21 Mon Sep 17 00:00:00 2001 From: Raza Rauf Date: Tue, 27 Jan 2026 00:21:47 +0500 Subject: [PATCH 02/16] fix: address PR review feedback and add DB caching optimization PR Review Fixes: - Fix race condition in exit-listener by moving markAndMaybeSynthesize to explicit code paths instead of finally() block - Add buffer size limits (MAX_BUFFER_SIZE 10MB) with warning logs - Add REGEX_BATCH_SESSION and REGEX_SYNOPSIS_SESSION for proper filtering - Fix type safety using canonical ToolExecution and UsageStats imports - Fix usage-listener indentation bug where safeSend was inside wrong block Performance Optimizations: - Add GROUP_CHAT_PREFIX constant for fast string prefix checks - Skip expensive regex matching for non-group-chat sessions - Eliminate redundant loadGroupChat calls by using updateParticipant return value directly (DB caching) - Add MSG_ID_RANDOM_LENGTH constant for web broadcast deduplication Test Coverage: - Add 4 new test files (exit, data, usage, session-id listeners) - Total 93 tests covering edge cases, DB caching, and performance - Verify exact participants data flow from updateParticipant - Test optional emitter handling and empty participants arrays --- src/main/constants.ts | 5 + src/main/index.ts | 4 + .../__tests__/data-listener.test.ts | 324 +++++++++++++ .../__tests__/exit-listener.test.ts | 420 +++++++++++++++++ .../__tests__/session-id-listener.test.ts | 402 ++++++++++++++++ .../__tests__/usage-listener.test.ts | 431 ++++++++++++++++++ src/main/process-listeners/data-listener.ts | 59 ++- src/main/process-listeners/exit-listener.ts | 33 +- .../process-listeners/forwarding-listeners.ts | 11 +- .../process-listeners/session-id-listener.ts | 25 +- src/main/process-listeners/types.ts | 4 + src/main/process-listeners/usage-listener.ts | 131 +++--- 12 files changed, 1759 insertions(+), 90 deletions(-) create mode 100644 src/main/process-listeners/__tests__/data-listener.test.ts create mode 100644 src/main/process-listeners/__tests__/exit-listener.test.ts create mode 100644 src/main/process-listeners/__tests__/session-id-listener.test.ts create mode 100644 src/main/process-listeners/__tests__/usage-listener.test.ts diff --git a/src/main/constants.ts b/src/main/constants.ts index cda32029..f5006fd5 100644 --- a/src/main/constants.ts +++ b/src/main/constants.ts @@ -26,6 +26,11 @@ export const REGEX_PARTICIPANT_FALLBACK = /^group-chat-(.+)-participant-([^-]+)- export const REGEX_AI_SUFFIX = /-ai-[^-]+$/; export const REGEX_AI_TAB_ID = /-ai-([^-]+)$/; +// Auto Run session ID patterns (batch and synopsis operations) +// Format: {sessionId}-batch-{timestamp} or {sessionId}-synopsis-{timestamp} +export const REGEX_BATCH_SESSION = /-batch-\d+$/; +export const REGEX_SYNOPSIS_SESSION = /-synopsis-\d+$/; + // ============================================================================ // Buffer Size Limits // ============================================================================ diff --git a/src/main/index.ts b/src/main/index.ts index 4f2838cf..a9029f48 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -76,6 +76,8 @@ import { REGEX_MODERATOR_SESSION_TIMESTAMP, REGEX_AI_SUFFIX, REGEX_AI_TAB_ID, + REGEX_BATCH_SESSION, + REGEX_SYNOPSIS_SESSION, debugLog, } from './constants'; // initAutoUpdater is now used by window-manager.ts (Phase 4 refactoring) @@ -662,6 +664,8 @@ function setupProcessListeners() { REGEX_MODERATOR_SESSION_TIMESTAMP, REGEX_AI_SUFFIX, REGEX_AI_TAB_ID, + REGEX_BATCH_SESSION, + REGEX_SYNOPSIS_SESSION, }, logger, }); diff --git a/src/main/process-listeners/__tests__/data-listener.test.ts b/src/main/process-listeners/__tests__/data-listener.test.ts new file mode 100644 index 00000000..d3950d89 --- /dev/null +++ b/src/main/process-listeners/__tests__/data-listener.test.ts @@ -0,0 +1,324 @@ +/** + * Tests for data listener. + * Handles process output data including group chat buffering and web broadcasting. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupDataListener } from '../data-listener'; +import type { ProcessManager } from '../../process-manager'; +import type { SafeSendFn } from '../../utils/safe-send'; +import type { ProcessListenerDependencies } from '../types'; + +describe('Data Listener', () => { + let mockProcessManager: ProcessManager; + let mockSafeSend: SafeSendFn; + let mockGetWebServer: ProcessListenerDependencies['getWebServer']; + let mockWebServer: { broadcastToSessionClients: ReturnType }; + let mockOutputBuffer: ProcessListenerDependencies['outputBuffer']; + let mockOutputParser: ProcessListenerDependencies['outputParser']; + let mockDebugLog: ProcessListenerDependencies['debugLog']; + let mockPatterns: ProcessListenerDependencies['patterns']; + let eventHandlers: Map void>; + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockSafeSend = vi.fn(); + mockWebServer = { + broadcastToSessionClients: vi.fn(), + }; + mockGetWebServer = vi.fn().mockReturnValue(mockWebServer); + mockOutputBuffer = { + appendToGroupChatBuffer: vi.fn().mockReturnValue(100), + getGroupChatBufferedOutput: vi.fn().mockReturnValue('test output'), + clearGroupChatBuffer: vi.fn(), + }; + mockOutputParser = { + extractTextFromStreamJson: vi.fn().mockReturnValue('parsed response'), + parseParticipantSessionId: vi.fn().mockReturnValue(null), + }; + mockDebugLog = vi.fn(); + mockPatterns = { + REGEX_MODERATOR_SESSION: /^group-chat-(.+)-moderator-/, + REGEX_MODERATOR_SESSION_TIMESTAMP: /^group-chat-(.+)-moderator-\d+$/, + REGEX_AI_SUFFIX: /-ai-[^-]+$/, + REGEX_AI_TAB_ID: /-ai-([^-]+)$/, + REGEX_BATCH_SESSION: /-batch-\d+$/, + REGEX_SYNOPSIS_SESSION: /-synopsis-\d+$/, + }; + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + }); + + const setupListener = () => { + setupDataListener(mockProcessManager, { + safeSend: mockSafeSend, + getWebServer: mockGetWebServer, + outputBuffer: mockOutputBuffer, + outputParser: mockOutputParser, + debugLog: mockDebugLog, + patterns: mockPatterns, + }); + }; + + describe('Event Registration', () => { + it('should register the data event listener', () => { + setupListener(); + expect(mockProcessManager.on).toHaveBeenCalledWith('data', expect.any(Function)); + }); + }); + + describe('Regular Process Data', () => { + it('should forward data to renderer for non-group-chat sessions', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('regular-session-123', 'test output'); + + expect(mockSafeSend).toHaveBeenCalledWith( + 'process:data', + 'regular-session-123', + 'test output' + ); + }); + + it('should broadcast to web clients for AI sessions', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-ai-tab1', 'test output'); + + expect(mockWebServer.broadcastToSessionClients).toHaveBeenCalledWith( + 'session-123', + expect.objectContaining({ + type: 'session_output', + sessionId: 'session-123', + tabId: 'tab1', + data: 'test output', + source: 'ai', + }) + ); + }); + + it('should extract base session ID correctly', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('my-session-ai-mytab', 'test output'); + + expect(mockWebServer.broadcastToSessionClients).toHaveBeenCalledWith( + 'my-session', + expect.objectContaining({ + sessionId: 'my-session', + tabId: 'mytab', + }) + ); + }); + }); + + describe('Moderator Output Buffering', () => { + it('should buffer moderator output instead of forwarding', () => { + setupListener(); + const handler = eventHandlers.get('data'); + const sessionId = 'group-chat-test-chat-123-moderator-abc123'; + + handler?.(sessionId, 'moderator output'); + + expect(mockOutputBuffer.appendToGroupChatBuffer).toHaveBeenCalledWith( + sessionId, + 'moderator output' + ); + expect(mockSafeSend).not.toHaveBeenCalled(); + }); + + it('should extract group chat ID from moderator session', () => { + setupListener(); + const handler = eventHandlers.get('data'); + const sessionId = 'group-chat-my-chat-id-moderator-12345'; + + handler?.(sessionId, 'test'); + + expect(mockDebugLog).toHaveBeenCalledWith( + 'GroupChat:Debug', + expect.stringContaining('my-chat-id') + ); + }); + + it('should warn when buffer size exceeds limit', () => { + mockOutputBuffer.appendToGroupChatBuffer = vi.fn().mockReturnValue(15 * 1024 * 1024); // 15MB + setupListener(); + const handler = eventHandlers.get('data'); + const sessionId = 'group-chat-test-chat-123-moderator-abc123'; + + handler?.(sessionId, 'large output'); + + expect(mockDebugLog).toHaveBeenCalledWith( + 'GroupChat:Debug', + expect.stringContaining('WARNING: Buffer size') + ); + }); + }); + + describe('Participant Output Buffering', () => { + beforeEach(() => { + mockOutputParser.parseParticipantSessionId = vi.fn().mockReturnValue({ + groupChatId: 'test-chat-123', + participantName: 'TestAgent', + }); + }); + + it('should buffer participant output instead of forwarding', () => { + setupListener(); + const handler = eventHandlers.get('data'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 'participant output'); + + expect(mockOutputBuffer.appendToGroupChatBuffer).toHaveBeenCalledWith( + sessionId, + 'participant output' + ); + expect(mockSafeSend).not.toHaveBeenCalled(); + }); + + it('should warn when participant buffer size exceeds limit', () => { + mockOutputBuffer.appendToGroupChatBuffer = vi.fn().mockReturnValue(15 * 1024 * 1024); // 15MB + setupListener(); + const handler = eventHandlers.get('data'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 'large output'); + + expect(mockDebugLog).toHaveBeenCalledWith( + 'GroupChat:Debug', + expect.stringContaining('WARNING: Buffer size') + ); + }); + }); + + describe('Web Broadcast Filtering', () => { + it('should skip PTY terminal output', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-terminal', 'terminal output'); + + expect(mockWebServer.broadcastToSessionClients).not.toHaveBeenCalled(); + // But should still forward to renderer + expect(mockSafeSend).toHaveBeenCalledWith( + 'process:data', + 'session-123-terminal', + 'terminal output' + ); + }); + + it('should skip batch session output using regex pattern', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-batch-1234567890', 'batch output'); + + expect(mockWebServer.broadcastToSessionClients).not.toHaveBeenCalled(); + }); + + it('should skip synopsis session output using regex pattern', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-synopsis-1234567890', 'synopsis output'); + + expect(mockWebServer.broadcastToSessionClients).not.toHaveBeenCalled(); + }); + + it('should NOT skip sessions with "batch" in UUID (false positive prevention)', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + // Session ID with "batch" in the UUID but not matching the pattern -batch-{digits} + handler?.('session-batch-uuid-ai-tab1', 'output'); + + // Should broadcast because it doesn't match the -batch-\d+$ pattern + expect(mockWebServer.broadcastToSessionClients).toHaveBeenCalled(); + }); + + it('should broadcast when no web server is available', () => { + mockGetWebServer = vi.fn().mockReturnValue(null); + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-ai-tab1', 'test output'); + + // Should still forward to renderer + expect(mockSafeSend).toHaveBeenCalledWith( + 'process:data', + 'session-123-ai-tab1', + 'test output' + ); + // But not broadcast (no web server) + expect(mockWebServer.broadcastToSessionClients).not.toHaveBeenCalled(); + }); + }); + + describe('Message ID Generation', () => { + it('should generate unique message IDs for broadcasts', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-ai-tab1', 'output 1'); + handler?.('session-123-ai-tab1', 'output 2'); + + const calls = mockWebServer.broadcastToSessionClients.mock.calls; + const msgId1 = calls[0][1].msgId; + const msgId2 = calls[1][1].msgId; + + expect(msgId1).toBeDefined(); + expect(msgId2).toBeDefined(); + expect(msgId1).not.toBe(msgId2); + }); + + it('should include timestamp in message ID', () => { + const beforeTime = Date.now(); + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-ai-tab1', 'test output'); + + const msgId = mockWebServer.broadcastToSessionClients.mock.calls[0][1].msgId; + const timestamp = parseInt(msgId.split('-')[0], 10); + + expect(timestamp).toBeGreaterThanOrEqual(beforeTime); + expect(timestamp).toBeLessThanOrEqual(Date.now()); + }); + }); + + describe('Source Detection', () => { + it('should identify AI source from session ID', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123-ai-tab1', 'ai output'); + + expect(mockWebServer.broadcastToSessionClients).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ source: 'ai' }) + ); + }); + + it('should identify terminal source for non-AI sessions', () => { + setupListener(); + const handler = eventHandlers.get('data'); + + handler?.('session-123', 'terminal output'); + + expect(mockWebServer.broadcastToSessionClients).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ source: 'terminal' }) + ); + }); + }); +}); diff --git a/src/main/process-listeners/__tests__/exit-listener.test.ts b/src/main/process-listeners/__tests__/exit-listener.test.ts new file mode 100644 index 00000000..793e1233 --- /dev/null +++ b/src/main/process-listeners/__tests__/exit-listener.test.ts @@ -0,0 +1,420 @@ +/** + * Tests for exit listener. + * Handles process exit events including group chat moderator/participant exits. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupExitListener } from '../exit-listener'; +import type { ProcessManager } from '../../process-manager'; +import type { ProcessListenerDependencies } from '../types'; + +describe('Exit Listener', () => { + let mockProcessManager: ProcessManager; + let mockDeps: Parameters[1]; + let eventHandlers: Map void>; + + // Create a minimal mock group chat + const createMockGroupChat = () => ({ + id: 'test-chat-123', + name: 'Test Chat', + moderatorAgentId: 'claude-code', + moderatorSessionId: 'group-chat-test-chat-123-moderator', + participants: [ + { + name: 'TestAgent', + agentId: 'claude-code', + sessionId: 'group-chat-test-chat-123-participant-TestAgent-abc123', + addedAt: Date.now(), + }, + ], + createdAt: Date.now(), + updatedAt: Date.now(), + logPath: '/tmp/test-chat.log', + imagesDir: '/tmp/test-chat-images', + }); + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + + mockDeps = { + safeSend: vi.fn(), + powerManager: { + addBlockReason: vi.fn(), + removeBlockReason: vi.fn(), + }, + groupChatEmitters: { + emitStateChange: vi.fn(), + emitParticipantState: vi.fn(), + emitParticipantsChanged: vi.fn(), + emitModeratorUsage: vi.fn(), + }, + groupChatRouter: { + routeModeratorResponse: vi.fn().mockResolvedValue(undefined), + routeAgentResponse: vi.fn().mockResolvedValue(undefined), + markParticipantResponded: vi.fn().mockResolvedValue(undefined), + spawnModeratorSynthesis: vi.fn().mockResolvedValue(undefined), + getGroupChatReadOnlyState: vi.fn().mockReturnValue(false), + respawnParticipantWithRecovery: vi.fn().mockResolvedValue(undefined), + }, + groupChatStorage: { + loadGroupChat: vi.fn().mockResolvedValue(createMockGroupChat()), + updateGroupChat: vi.fn().mockResolvedValue(createMockGroupChat()), + updateParticipant: vi.fn().mockResolvedValue(createMockGroupChat()), + }, + sessionRecovery: { + needsSessionRecovery: vi.fn().mockReturnValue(false), + initiateSessionRecovery: vi.fn().mockResolvedValue(true), + }, + outputBuffer: { + appendToGroupChatBuffer: vi.fn().mockReturnValue(100), + getGroupChatBufferedOutput: vi.fn().mockReturnValue('{"type":"text","text":"test output"}'), + clearGroupChatBuffer: vi.fn(), + }, + outputParser: { + extractTextFromStreamJson: vi.fn().mockReturnValue('parsed response'), + parseParticipantSessionId: vi.fn().mockReturnValue(null), + }, + getProcessManager: () => mockProcessManager, + getAgentDetector: () => + ({ + detectAgents: vi.fn(), + }) as unknown as ReturnType, + getWebServer: () => null, + logger: { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }, + debugLog: vi.fn(), + patterns: { + REGEX_MODERATOR_SESSION: /^group-chat-(.+)-moderator-/, + REGEX_MODERATOR_SESSION_TIMESTAMP: /^group-chat-(.+)-moderator-\d+$/, + REGEX_AI_SUFFIX: /-ai-[^-]+$/, + REGEX_AI_TAB_ID: /-ai-([^-]+)$/, + REGEX_BATCH_SESSION: /-batch-\d+$/, + REGEX_SYNOPSIS_SESSION: /-synopsis-\d+$/, + }, + }; + }); + + const setupListener = () => { + setupExitListener(mockProcessManager, mockDeps); + }; + + describe('Event Registration', () => { + it('should register the exit event listener', () => { + setupListener(); + expect(mockProcessManager.on).toHaveBeenCalledWith('exit', expect.any(Function)); + }); + }); + + describe('Regular Process Exit', () => { + it('should forward exit event to renderer for non-group-chat sessions', () => { + setupListener(); + const handler = eventHandlers.get('exit'); + + handler?.('regular-session-123', 0); + + expect(mockDeps.safeSend).toHaveBeenCalledWith('process:exit', 'regular-session-123', 0); + }); + + it('should remove power block for non-group-chat sessions', () => { + setupListener(); + const handler = eventHandlers.get('exit'); + + handler?.('regular-session-123', 0); + + expect(mockDeps.powerManager.removeBlockReason).toHaveBeenCalledWith( + 'session:regular-session-123' + ); + }); + }); + + describe('Participant Exit', () => { + beforeEach(() => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue({ + groupChatId: 'test-chat-123', + participantName: 'TestAgent', + }); + }); + + it('should parse and route participant response on exit', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.routeAgentResponse).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent', + 'parsed response', + expect.anything() + ); + }); + }); + + it('should mark participant as responded after successful routing', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.markParticipantResponded).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent' + ); + }); + }); + + it('should clear output buffer after processing', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.outputBuffer.clearGroupChatBuffer).toHaveBeenCalledWith(sessionId); + }); + }); + + it('should not route when buffered output is empty', async () => { + mockDeps.outputBuffer.getGroupChatBufferedOutput = vi.fn().mockReturnValue(''); + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + // Give async operations time to complete + await new Promise((resolve) => setTimeout(resolve, 50)); + + expect(mockDeps.groupChatRouter.routeAgentResponse).not.toHaveBeenCalled(); + }); + + it('should not route when parsed text is empty', async () => { + mockDeps.outputParser.extractTextFromStreamJson = vi.fn().mockReturnValue(' '); + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + // Give async operations time to complete + await new Promise((resolve) => setTimeout(resolve, 50)); + + expect(mockDeps.groupChatRouter.routeAgentResponse).not.toHaveBeenCalled(); + }); + }); + + describe('Session Recovery', () => { + beforeEach(() => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue({ + groupChatId: 'test-chat-123', + participantName: 'TestAgent', + }); + mockDeps.sessionRecovery.needsSessionRecovery = vi.fn().mockReturnValue(true); + }); + + it('should initiate session recovery when needed', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.sessionRecovery.initiateSessionRecovery).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent' + ); + }); + }); + + it('should respawn participant after recovery initiation', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.respawnParticipantWithRecovery).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent', + expect.anything(), + expect.anything() + ); + }); + }); + + it('should clear buffer before initiating recovery', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.outputBuffer.clearGroupChatBuffer).toHaveBeenCalledWith(sessionId); + }); + }); + + it('should not mark participant as responded when recovery succeeds', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + // Wait for async operations + await new Promise((resolve) => setTimeout(resolve, 50)); + + // When recovery succeeds, markParticipantResponded should NOT be called + // because the recovery spawn will handle that + expect(mockDeps.groupChatRouter.markParticipantResponded).not.toHaveBeenCalled(); + }); + + it('should mark participant as responded when recovery fails', async () => { + mockDeps.groupChatRouter.respawnParticipantWithRecovery = vi + .fn() + .mockRejectedValue(new Error('Recovery failed')); + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.markParticipantResponded).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent' + ); + }); + }); + }); + + describe('Moderator Exit', () => { + it('should route moderator response on exit', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-moderator-1234567890'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.routeModeratorResponse).toHaveBeenCalledWith( + 'test-chat-123', + 'parsed response', + expect.anything(), + expect.anything(), + false + ); + }); + }); + + it('should clear moderator buffer after processing', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-moderator-1234567890'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.outputBuffer.clearGroupChatBuffer).toHaveBeenCalledWith(sessionId); + }); + }); + + it('should handle synthesis sessions correctly', async () => { + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-moderator-synthesis-1234567890'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.routeModeratorResponse).toHaveBeenCalled(); + }); + }); + }); + + describe('Error Handling', () => { + beforeEach(() => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue({ + groupChatId: 'test-chat-123', + participantName: 'TestAgent', + }); + }); + + it('should log error when routing fails', async () => { + mockDeps.groupChatRouter.routeAgentResponse = vi + .fn() + .mockRejectedValue(new Error('Route failed')); + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.logger.error).toHaveBeenCalled(); + }); + }); + + it('should attempt fallback parsing when primary parsing fails', async () => { + // First call throws, second call (fallback) succeeds + mockDeps.outputParser.extractTextFromStreamJson = vi + .fn() + .mockImplementationOnce(() => { + throw new Error('Parse error'); + }) + .mockReturnValueOnce('fallback parsed response'); + + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + // Should have been called twice: once with agentType, once without (fallback) + expect(mockDeps.outputParser.extractTextFromStreamJson).toHaveBeenCalledTimes(2); + }); + }); + + it('should still mark participant as responded after routing error', async () => { + mockDeps.groupChatRouter.routeAgentResponse = vi + .fn() + .mockRejectedValue(new Error('Route failed')); + mockDeps.outputParser.extractTextFromStreamJson = vi + .fn() + .mockReturnValueOnce('parsed response') + .mockReturnValueOnce('fallback response'); + + setupListener(); + const handler = eventHandlers.get('exit'); + const sessionId = 'group-chat-test-chat-123-participant-TestAgent-abc123'; + + handler?.(sessionId, 0); + + await vi.waitFor(() => { + expect(mockDeps.groupChatRouter.markParticipantResponded).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent' + ); + }); + }); + }); +}); diff --git a/src/main/process-listeners/__tests__/session-id-listener.test.ts b/src/main/process-listeners/__tests__/session-id-listener.test.ts new file mode 100644 index 00000000..c0f6773c --- /dev/null +++ b/src/main/process-listeners/__tests__/session-id-listener.test.ts @@ -0,0 +1,402 @@ +/** + * Tests for session ID listener. + * Handles agent session ID storage for conversation resume. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupSessionIdListener } from '../session-id-listener'; +import type { ProcessManager } from '../../process-manager'; + +describe('Session ID Listener', () => { + let mockProcessManager: ProcessManager; + let mockDeps: Parameters[1]; + let eventHandlers: Map void>; + + // Create a minimal mock group chat + const createMockGroupChat = () => ({ + id: 'test-chat-123', + name: 'Test Chat', + moderatorAgentId: 'claude-code', + moderatorSessionId: 'group-chat-test-chat-123-moderator', + participants: [ + { + name: 'TestAgent', + agentId: 'claude-code', + sessionId: 'group-chat-test-chat-123-participant-TestAgent-abc123', + addedAt: Date.now(), + }, + ], + createdAt: Date.now(), + updatedAt: Date.now(), + logPath: '/tmp/test-chat.log', + imagesDir: '/tmp/test-chat-images', + }); + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + + mockDeps = { + safeSend: vi.fn(), + logger: { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }, + groupChatEmitters: { + emitParticipantsChanged: vi.fn(), + emitModeratorSessionIdChanged: vi.fn(), + }, + groupChatStorage: { + loadGroupChat: vi.fn().mockResolvedValue(createMockGroupChat()), + updateGroupChat: vi.fn().mockResolvedValue(createMockGroupChat()), + updateParticipant: vi.fn().mockResolvedValue(createMockGroupChat()), + }, + outputParser: { + extractTextFromStreamJson: vi.fn().mockReturnValue('parsed response'), + parseParticipantSessionId: vi.fn().mockReturnValue(null), + }, + patterns: { + REGEX_MODERATOR_SESSION: /^group-chat-(.+)-moderator-/, + REGEX_MODERATOR_SESSION_TIMESTAMP: /^group-chat-(.+)-moderator-\d+$/, + REGEX_AI_SUFFIX: /-ai-[^-]+$/, + REGEX_AI_TAB_ID: /-ai-([^-]+)$/, + REGEX_BATCH_SESSION: /-batch-\d+$/, + REGEX_SYNOPSIS_SESSION: /-synopsis-\d+$/, + }, + }; + }); + + const setupListener = () => { + setupSessionIdListener(mockProcessManager, mockDeps); + }; + + describe('Event Registration', () => { + it('should register the session-id event listener', () => { + setupListener(); + expect(mockProcessManager.on).toHaveBeenCalledWith('session-id', expect.any(Function)); + }); + }); + + describe('Regular Process Session ID', () => { + it('should forward session ID to renderer', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('regular-session-123', 'agent-session-abc'); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:session-id', + 'regular-session-123', + 'agent-session-abc' + ); + }); + }); + + describe('Participant Session ID Storage', () => { + beforeEach(() => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue({ + groupChatId: 'test-chat-123', + participantName: 'TestAgent', + }); + }); + + it('should store agent session ID for participant', async () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateParticipant).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent', + { agentSessionId: 'agent-session-xyz' } + ); + }); + }); + + it('should emit participants changed after storage', async () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalledWith( + 'test-chat-123', + expect.any(Array) + ); + }); + }); + + it('should use updateParticipant return value instead of loading chat again (DB caching)', async () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalled(); + }); + + // Verify we didn't make a redundant loadGroupChat call + // The code should use the return value from updateParticipant directly + expect(mockDeps.groupChatStorage.loadGroupChat).not.toHaveBeenCalled(); + }); + + it('should pass exact participants from updateParticipant return value', async () => { + const specificParticipants = [ + { name: 'Agent1', agentId: 'claude-code', sessionId: 'session-1', addedAt: 1000 }, + { name: 'Agent2', agentId: 'codex', sessionId: 'session-2', addedAt: 2000 }, + ]; + mockDeps.groupChatStorage.updateParticipant = vi.fn().mockResolvedValue({ + ...createMockGroupChat(), + participants: specificParticipants, + }); + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalledWith( + 'test-chat-123', + specificParticipants + ); + }); + }); + + it('should handle empty participants array from updateParticipant', async () => { + mockDeps.groupChatStorage.updateParticipant = vi.fn().mockResolvedValue({ + ...createMockGroupChat(), + participants: [], + }); + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalledWith( + 'test-chat-123', + [] + ); + }); + }); + + it('should handle undefined emitParticipantsChanged gracefully (optional chaining)', async () => { + mockDeps.groupChatEmitters.emitParticipantsChanged = undefined; + setupListener(); + const handler = eventHandlers.get('session-id'); + + // Should not throw + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateParticipant).toHaveBeenCalled(); + }); + // No error should be logged for the optional emitter + expect(mockDeps.logger.error).not.toHaveBeenCalled(); + }); + + it('should log error when storage fails', async () => { + mockDeps.groupChatStorage.updateParticipant = vi + .fn() + .mockRejectedValue(new Error('DB error')); + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.logger.error).toHaveBeenCalledWith( + '[GroupChat] Failed to update participant agentSessionId', + 'ProcessListener', + expect.objectContaining({ + error: 'Error: DB error', + participant: 'TestAgent', + }) + ); + }); + }); + + it('should still forward to renderer after storage', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', 'agent-session-xyz'); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:session-id', + 'group-chat-test-chat-123-participant-TestAgent-abc123', + 'agent-session-xyz' + ); + }); + }); + + describe('Moderator Session ID Storage', () => { + it('should store agent session ID for moderator', async () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-moderator-1234567890', 'moderator-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateGroupChat).toHaveBeenCalledWith('test-chat-123', { + moderatorAgentSessionId: 'moderator-session-xyz', + }); + }); + }); + + it('should emit moderator session ID changed after storage', async () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-moderator-1234567890', 'moderator-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitModeratorSessionIdChanged).toHaveBeenCalledWith( + 'test-chat-123', + 'moderator-session-xyz' + ); + }); + }); + + it('should log error when moderator storage fails', async () => { + mockDeps.groupChatStorage.updateGroupChat = vi.fn().mockRejectedValue(new Error('DB error')); + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-moderator-1234567890', 'moderator-session-xyz'); + + await vi.waitFor(() => { + expect(mockDeps.logger.error).toHaveBeenCalledWith( + '[GroupChat] Failed to update moderator agent session ID', + 'ProcessListener', + expect.objectContaining({ + error: 'Error: DB error', + groupChatId: 'test-chat-123', + }) + ); + }); + }); + + it('should still forward to renderer for moderator sessions', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('group-chat-test-chat-123-moderator-1234567890', 'moderator-session-xyz'); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:session-id', + 'group-chat-test-chat-123-moderator-1234567890', + 'moderator-session-xyz' + ); + }); + + it('should NOT store for synthesis moderator sessions (different pattern)', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + // Synthesis session ID doesn't match REGEX_MODERATOR_SESSION_TIMESTAMP + // because it has 'synthesis' in it: group-chat-xxx-moderator-synthesis-timestamp + handler?.('group-chat-test-chat-123-moderator-synthesis-1234567890', 'synthesis-session-xyz'); + + // Should NOT call updateGroupChat for synthesis sessions (doesn't match timestamp pattern) + expect(mockDeps.groupChatStorage.updateGroupChat).not.toHaveBeenCalled(); + }); + }); + + describe('Session ID Format Handling', () => { + it('should handle empty agent session ID', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('regular-session-123', ''); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:session-id', + 'regular-session-123', + '' + ); + }); + + it('should handle UUID format session IDs', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + handler?.('regular-session-123', 'a1b2c3d4-e5f6-7890-abcd-ef1234567890'); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:session-id', + 'regular-session-123', + 'a1b2c3d4-e5f6-7890-abcd-ef1234567890' + ); + }); + + it('should handle long session IDs', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + const longSessionId = 'a'.repeat(500); + + handler?.('regular-session-123', longSessionId); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:session-id', + 'regular-session-123', + longSessionId + ); + }); + }); + + describe('Performance Optimization', () => { + it('should skip participant parsing for non-group-chat sessions (prefix check)', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + // Regular session ID doesn't start with 'group-chat-' + handler?.('regular-session-123', 'agent-session-abc'); + + // parseParticipantSessionId should NOT be called for non-group-chat sessions + expect(mockDeps.outputParser.parseParticipantSessionId).not.toHaveBeenCalled(); + }); + + it('should only parse participant session ID for group-chat sessions', () => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue(null); + setupListener(); + const handler = eventHandlers.get('session-id'); + + // Group chat session ID starts with 'group-chat-' + handler?.('group-chat-test-123-participant-Agent-abc', 'agent-session-xyz'); + + // parseParticipantSessionId SHOULD be called for group-chat sessions + expect(mockDeps.outputParser.parseParticipantSessionId).toHaveBeenCalledWith( + 'group-chat-test-123-participant-Agent-abc' + ); + }); + + it('should skip moderator regex for non-group-chat sessions', () => { + setupListener(); + const handler = eventHandlers.get('session-id'); + + // Process many non-group-chat sessions - should be fast since regex is skipped + for (let i = 0; i < 100; i++) { + handler?.(`regular-session-${i}`, `agent-session-${i}`); + } + + // Neither storage method should be called for regular sessions + expect(mockDeps.groupChatStorage.updateParticipant).not.toHaveBeenCalled(); + expect(mockDeps.groupChatStorage.updateGroupChat).not.toHaveBeenCalled(); + // But all should still forward to renderer + expect(mockDeps.safeSend).toHaveBeenCalledTimes(100); + }); + }); +}); diff --git a/src/main/process-listeners/__tests__/usage-listener.test.ts b/src/main/process-listeners/__tests__/usage-listener.test.ts new file mode 100644 index 00000000..16c52d61 --- /dev/null +++ b/src/main/process-listeners/__tests__/usage-listener.test.ts @@ -0,0 +1,431 @@ +/** + * Tests for usage listener. + * Handles token/cost statistics from AI responses. + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { setupUsageListener } from '../usage-listener'; +import type { ProcessManager } from '../../process-manager'; +import type { UsageStats } from '../types'; + +describe('Usage Listener', () => { + let mockProcessManager: ProcessManager; + let mockDeps: Parameters[1]; + let eventHandlers: Map void>; + + const createMockUsageStats = (overrides: Partial = {}): UsageStats => ({ + inputTokens: 1000, + outputTokens: 500, + cacheReadInputTokens: 200, + cacheCreationInputTokens: 100, + totalCostUsd: 0.05, + contextWindow: 100000, + ...overrides, + }); + + // Create a minimal mock group chat + const createMockGroupChat = () => ({ + id: 'test-chat-123', + name: 'Test Chat', + moderatorAgentId: 'claude-code', + moderatorSessionId: 'group-chat-test-chat-123-moderator', + participants: [ + { + name: 'TestAgent', + agentId: 'claude-code', + sessionId: 'group-chat-test-chat-123-participant-TestAgent-abc123', + addedAt: Date.now(), + }, + ], + createdAt: Date.now(), + updatedAt: Date.now(), + logPath: '/tmp/test-chat.log', + imagesDir: '/tmp/test-chat-images', + }); + + beforeEach(() => { + vi.clearAllMocks(); + eventHandlers = new Map(); + + mockProcessManager = { + on: vi.fn((event: string, handler: (...args: unknown[]) => void) => { + eventHandlers.set(event, handler); + }), + } as unknown as ProcessManager; + + mockDeps = { + safeSend: vi.fn(), + logger: { + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }, + groupChatEmitters: { + emitParticipantsChanged: vi.fn(), + emitModeratorUsage: vi.fn(), + }, + groupChatStorage: { + loadGroupChat: vi.fn().mockResolvedValue(createMockGroupChat()), + updateGroupChat: vi.fn().mockResolvedValue(createMockGroupChat()), + updateParticipant: vi.fn().mockResolvedValue(createMockGroupChat()), + }, + outputParser: { + extractTextFromStreamJson: vi.fn().mockReturnValue('parsed response'), + parseParticipantSessionId: vi.fn().mockReturnValue(null), + }, + usageAggregator: { + calculateContextTokens: vi.fn().mockReturnValue(1800), + }, + patterns: { + REGEX_MODERATOR_SESSION: /^group-chat-(.+)-moderator-/, + REGEX_MODERATOR_SESSION_TIMESTAMP: /^group-chat-(.+)-moderator-\d+$/, + REGEX_AI_SUFFIX: /-ai-[^-]+$/, + REGEX_AI_TAB_ID: /-ai-([^-]+)$/, + REGEX_BATCH_SESSION: /-batch-\d+$/, + REGEX_SYNOPSIS_SESSION: /-synopsis-\d+$/, + }, + }; + }); + + const setupListener = () => { + setupUsageListener(mockProcessManager, mockDeps); + }; + + describe('Event Registration', () => { + it('should register the usage event listener', () => { + setupListener(); + expect(mockProcessManager.on).toHaveBeenCalledWith('usage', expect.any(Function)); + }); + }); + + describe('Regular Process Usage', () => { + it('should forward usage stats to renderer', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('regular-session-123', usageStats); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:usage', + 'regular-session-123', + usageStats + ); + }); + }); + + describe('Participant Usage', () => { + beforeEach(() => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue({ + groupChatId: 'test-chat-123', + participantName: 'TestAgent', + }); + }); + + it('should update participant with usage stats', async () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateParticipant).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent', + expect.objectContaining({ + contextUsage: expect.any(Number), + tokenCount: 1800, + totalCost: 0.05, + }) + ); + }); + }); + + it('should calculate context usage percentage correctly', async () => { + mockDeps.usageAggregator.calculateContextTokens = vi.fn().mockReturnValue(50000); // 50% of 100000 + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats({ contextWindow: 100000 }); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateParticipant).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent', + expect.objectContaining({ + contextUsage: 50, + }) + ); + }); + }); + + it('should handle zero context window gracefully', async () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats({ contextWindow: 0 }); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateParticipant).toHaveBeenCalledWith( + 'test-chat-123', + 'TestAgent', + expect.objectContaining({ + contextUsage: 0, + }) + ); + }); + }); + + it('should emit participants changed after update', async () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalledWith( + 'test-chat-123', + expect.any(Array) + ); + }); + }); + + it('should use updateParticipant return value instead of loading chat again (DB caching)', async () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalled(); + }); + + // Verify we didn't make a redundant loadGroupChat call + // The code should use the return value from updateParticipant directly + expect(mockDeps.groupChatStorage.loadGroupChat).not.toHaveBeenCalled(); + }); + + it('should pass exact participants from updateParticipant return value', async () => { + const specificParticipants = [ + { name: 'Agent1', agentId: 'claude-code', sessionId: 'session-1', addedAt: 1000 }, + { name: 'Agent2', agentId: 'codex', sessionId: 'session-2', addedAt: 2000 }, + ]; + mockDeps.groupChatStorage.updateParticipant = vi.fn().mockResolvedValue({ + ...createMockGroupChat(), + participants: specificParticipants, + }); + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalledWith( + 'test-chat-123', + specificParticipants + ); + }); + }); + + it('should handle empty participants array from updateParticipant', async () => { + mockDeps.groupChatStorage.updateParticipant = vi.fn().mockResolvedValue({ + ...createMockGroupChat(), + participants: [], + }); + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatEmitters.emitParticipantsChanged).toHaveBeenCalledWith( + 'test-chat-123', + [] + ); + }); + }); + + it('should handle undefined emitParticipantsChanged gracefully (optional chaining)', async () => { + mockDeps.groupChatEmitters.emitParticipantsChanged = undefined; + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + // Should not throw + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.groupChatStorage.updateParticipant).toHaveBeenCalled(); + }); + // No error should be logged for the optional emitter + expect(mockDeps.logger.error).not.toHaveBeenCalled(); + }); + + it('should log error when participant update fails', async () => { + mockDeps.groupChatStorage.updateParticipant = vi + .fn() + .mockRejectedValue(new Error('DB error')); + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + await vi.waitFor(() => { + expect(mockDeps.logger.error).toHaveBeenCalledWith( + '[GroupChat] Failed to update participant usage', + 'ProcessListener', + expect.objectContaining({ + error: 'Error: DB error', + participant: 'TestAgent', + }) + ); + }); + }); + + it('should still forward to renderer for participant usage', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-participant-TestAgent-abc123', usageStats); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:usage', + 'group-chat-test-chat-123-participant-TestAgent-abc123', + usageStats + ); + }); + }); + + describe('Moderator Usage', () => { + it('should emit moderator usage for moderator sessions', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-moderator-1234567890', usageStats); + + expect(mockDeps.groupChatEmitters.emitModeratorUsage).toHaveBeenCalledWith( + 'test-chat-123', + expect.objectContaining({ + contextUsage: expect.any(Number), + totalCost: 0.05, + tokenCount: 1800, + }) + ); + }); + + it('should calculate moderator context usage correctly', () => { + mockDeps.usageAggregator.calculateContextTokens = vi.fn().mockReturnValue(25000); // 25% of 100000 + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats({ contextWindow: 100000 }); + + handler?.('group-chat-test-chat-123-moderator-1234567890', usageStats); + + expect(mockDeps.groupChatEmitters.emitModeratorUsage).toHaveBeenCalledWith( + 'test-chat-123', + expect.objectContaining({ + contextUsage: 25, + }) + ); + }); + + it('should still forward to renderer for moderator usage', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-moderator-1234567890', usageStats); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:usage', + 'group-chat-test-chat-123-moderator-1234567890', + usageStats + ); + }); + + it('should handle synthesis moderator sessions', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + handler?.('group-chat-test-chat-123-moderator-synthesis-1234567890', usageStats); + + expect(mockDeps.groupChatEmitters.emitModeratorUsage).toHaveBeenCalledWith( + 'test-chat-123', + expect.any(Object) + ); + }); + }); + + describe('Usage with Reasoning Tokens', () => { + it('should handle usage stats with reasoning tokens', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats({ reasoningTokens: 1000 }); + + handler?.('regular-session-123', usageStats); + + expect(mockDeps.safeSend).toHaveBeenCalledWith( + 'process:usage', + 'regular-session-123', + expect.objectContaining({ reasoningTokens: 1000 }) + ); + }); + }); + + describe('Performance Optimization', () => { + it('should skip participant parsing for non-group-chat sessions (prefix check)', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + // Regular session ID doesn't start with 'group-chat-' + handler?.('regular-session-123', usageStats); + + // parseParticipantSessionId should NOT be called for non-group-chat sessions + expect(mockDeps.outputParser.parseParticipantSessionId).not.toHaveBeenCalled(); + }); + + it('should only parse participant session ID for group-chat sessions', () => { + mockDeps.outputParser.parseParticipantSessionId = vi.fn().mockReturnValue(null); + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + // Group chat session ID starts with 'group-chat-' + handler?.('group-chat-test-123-participant-Agent-abc', usageStats); + + // parseParticipantSessionId SHOULD be called for group-chat sessions + expect(mockDeps.outputParser.parseParticipantSessionId).toHaveBeenCalledWith( + 'group-chat-test-123-participant-Agent-abc' + ); + }); + + it('should skip moderator regex for non-group-chat sessions', () => { + setupListener(); + const handler = eventHandlers.get('usage'); + const usageStats = createMockUsageStats(); + + // Process many non-group-chat sessions - should be fast since regex is skipped + for (let i = 0; i < 100; i++) { + handler?.(`regular-session-${i}`, usageStats); + } + + // Moderator usage should NOT be emitted for any regular sessions + expect(mockDeps.groupChatEmitters.emitModeratorUsage).not.toHaveBeenCalled(); + // But all should still forward to renderer + expect(mockDeps.safeSend).toHaveBeenCalledTimes(100); + }); + }); +}); diff --git a/src/main/process-listeners/data-listener.ts b/src/main/process-listeners/data-listener.ts index 71b2570d..642f73c6 100644 --- a/src/main/process-listeners/data-listener.ts +++ b/src/main/process-listeners/data-listener.ts @@ -6,6 +6,24 @@ import type { ProcessManager } from '../process-manager'; import type { ProcessListenerDependencies } from './types'; +/** + * Maximum buffer size per session (10MB). + * Prevents unbounded memory growth from long-running or misbehaving processes. + */ +const MAX_BUFFER_SIZE = 10 * 1024 * 1024; + +/** + * Prefix for group chat session IDs. + * Used for fast string check before expensive regex matching. + */ +const GROUP_CHAT_PREFIX = 'group-chat-'; + +/** + * Length of random suffix in message IDs (9 characters of base36). + * Combined with timestamp provides uniqueness for web broadcast deduplication. + */ +const MSG_ID_RANDOM_LENGTH = 9; + /** * Sets up the data listener for process output. * Handles: @@ -21,12 +39,22 @@ export function setupDataListener( > ): void { const { safeSend, getWebServer, outputBuffer, outputParser, debugLog, patterns } = deps; - const { REGEX_MODERATOR_SESSION, REGEX_AI_SUFFIX, REGEX_AI_TAB_ID } = patterns; + const { + REGEX_MODERATOR_SESSION, + REGEX_AI_SUFFIX, + REGEX_AI_TAB_ID, + REGEX_BATCH_SESSION, + REGEX_SYNOPSIS_SESSION, + } = patterns; processManager.on('data', (sessionId: string, data: string) => { + // Fast path: skip regex for non-group-chat sessions (performance optimization) + // Most sessions don't start with 'group-chat-', so this avoids expensive regex matching + const isGroupChatSession = sessionId.startsWith(GROUP_CHAT_PREFIX); + // Handle group chat moderator output - buffer it // Session ID format: group-chat-{groupChatId}-moderator-{uuid} or group-chat-{groupChatId}-moderator-synthesis-{uuid} - const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION); + const moderatorMatch = isGroupChatSession ? sessionId.match(REGEX_MODERATOR_SESSION) : null; if (moderatorMatch) { const groupChatId = moderatorMatch[1]; debugLog('GroupChat:Debug', `MODERATOR DATA received for chat ${groupChatId}`); @@ -35,12 +63,22 @@ export function setupDataListener( // Buffer the output - will be routed on process exit const totalLength = outputBuffer.appendToGroupChatBuffer(sessionId, data); debugLog('GroupChat:Debug', `Buffered total: ${totalLength} chars`); + // Warn if buffer is growing too large (potential memory leak) + if (totalLength > MAX_BUFFER_SIZE) { + debugLog( + 'GroupChat:Debug', + `WARNING: Buffer size ${totalLength} exceeds ${MAX_BUFFER_SIZE} bytes for moderator session ${sessionId}` + ); + } return; // Don't send to regular process:data handler } // Handle group chat participant output - buffer it // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} - const participantInfo = outputParser.parseParticipantSessionId(sessionId); + // Only parse if it's a group chat session (performance optimization) + const participantInfo = isGroupChatSession + ? outputParser.parseParticipantSessionId(sessionId) + : null; if (participantInfo) { debugLog('GroupChat:Debug', 'PARTICIPANT DATA received'); debugLog( @@ -52,6 +90,13 @@ export function setupDataListener( // Buffer the output - will be routed on process exit const totalLength = outputBuffer.appendToGroupChatBuffer(sessionId, data); debugLog('GroupChat:Debug', `Buffered total: ${totalLength} chars`); + // Warn if buffer is growing too large (potential memory leak) + if (totalLength > MAX_BUFFER_SIZE) { + debugLog( + 'GroupChat:Debug', + `WARNING: Buffer size ${totalLength} exceeds ${MAX_BUFFER_SIZE} bytes for participant ${participantInfo.participantName}` + ); + } return; // Don't send to regular process:data handler } @@ -70,7 +115,8 @@ export function setupDataListener( // Don't broadcast background batch/synopsis output to web clients // These are internal Auto Run operations that should only appear in history, not as chat messages - if (sessionId.includes('-batch-') || sessionId.includes('-synopsis-')) { + // Use proper regex patterns to avoid false positives from UUIDs containing "batch" or "synopsis" + if (REGEX_BATCH_SESSION.test(sessionId) || REGEX_SYNOPSIS_SESSION.test(sessionId)) { debugLog('WebBroadcast', `SKIPPING batch/synopsis output for web: session=${sessionId}`); return; } @@ -83,7 +129,10 @@ export function setupDataListener( const tabIdMatch = sessionId.match(REGEX_AI_TAB_ID); const tabId = tabIdMatch ? tabIdMatch[1] : undefined; - const msgId = `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`; + // Generate unique message ID: timestamp + random suffix for deduplication + const msgId = `${Date.now()}-${Math.random() + .toString(36) + .substring(2, 2 + MSG_ID_RANDOM_LENGTH)}`; debugLog( 'WebBroadcast', `Broadcasting session_output: msgId=${msgId}, session=${baseSessionId}, tabId=${tabId || 'none'}, source=${isAiOutput ? 'ai' : 'terminal'}, dataLen=${data.length}` diff --git a/src/main/process-listeners/exit-listener.ts b/src/main/process-listeners/exit-listener.ts index 2f20b54f..8c0ea4c1 100644 --- a/src/main/process-listeners/exit-listener.ts +++ b/src/main/process-listeners/exit-listener.ts @@ -7,6 +7,12 @@ import type { ProcessManager } from '../process-manager'; import type { ProcessListenerDependencies } from './types'; +/** + * Prefix for group chat session IDs. + * Used for fast string check before expensive regex matching. + */ +const GROUP_CHAT_PREFIX = 'group-chat-'; + /** * Sets up the exit listener for process termination. * Handles: @@ -59,13 +65,17 @@ export function setupExitListener( // This allows system sleep when no AI sessions are active powerManager.removeBlockReason(`session:${sessionId}`); + // Fast path: skip regex for non-group-chat sessions (performance optimization) + // Most sessions don't start with 'group-chat-', so this avoids expensive regex matching + const isGroupChatSession = sessionId.startsWith(GROUP_CHAT_PREFIX); + // Handle group chat moderator exit - route buffered output and set state back to idle // Session ID format: group-chat-{groupChatId}-moderator-{uuid} // This handles BOTH initial moderator responses AND synthesis responses. // The routeModeratorResponse function will check for @mentions: // - If @mentions present: route to agents (continue conversation) // - If no @mentions: final response to user (conversation complete for this turn) - const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION); + const moderatorMatch = isGroupChatSession ? sessionId.match(REGEX_MODERATOR_SESSION) : null; if (moderatorMatch) { const groupChatId = moderatorMatch[1]; debugLog('GroupChat:Debug', ` ========== MODERATOR PROCESS EXIT ==========`); @@ -189,7 +199,10 @@ export function setupExitListener( // Handle group chat participant exit - route buffered output and update participant state // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} - const participantExitInfo = outputParser.parseParticipantSessionId(sessionId); + // Only parse if it's a group chat session (performance optimization) + const participantExitInfo = isGroupChatSession + ? outputParser.parseParticipantSessionId(sessionId) + : null; if (participantExitInfo) { const { groupChatId, participantName } = participantExitInfo; debugLog('GroupChat:Debug', ` ========== PARTICIPANT PROCESS EXIT ==========`); @@ -342,8 +355,12 @@ export function setupExitListener( 'GroupChat:Debug', ` Successfully routed agent response from ${participantName}` ); + // Mark participant AFTER routing completes successfully + markAndMaybeSynthesize(); } else { debugLog('GroupChat:Debug', ` WARNING: Parsed text is empty for ${participantName}!`); + // No response to route, mark participant as done + markAndMaybeSynthesize(); } } catch (err) { debugLog('GroupChat:Debug', ` ERROR loading chat for participant:`, err); @@ -362,6 +379,11 @@ export function setupExitListener( parsedText, pm ?? undefined ); + // Mark participant AFTER routing completes successfully + markAndMaybeSynthesize(); + } else { + // No response to route, mark participant as done + markAndMaybeSynthesize(); } } catch (routeErr) { debugLog('GroupChat:Debug', ` ERROR routing agent response (fallback):`, routeErr); @@ -369,13 +391,16 @@ export function setupExitListener( error: String(routeErr), participant: participantName, }); + // Mark participant as done even after error (can't retry) + markAndMaybeSynthesize(); } } })().finally(() => { outputBuffer.clearGroupChatBuffer(sessionId); debugLog('GroupChat:Debug', ` Cleared output buffer for participant session`); - // Mark participant and trigger synthesis AFTER logging is complete - markAndMaybeSynthesize(); + // Note: markAndMaybeSynthesize() is called explicitly in each code path above + // to ensure proper sequencing - NOT in finally() which would cause race conditions + // with session recovery (where we DON'T want to mark until recovery completes) }); } else { debugLog( diff --git a/src/main/process-listeners/forwarding-listeners.ts b/src/main/process-listeners/forwarding-listeners.ts index 3ec6bb4c..90121b8b 100644 --- a/src/main/process-listeners/forwarding-listeners.ts +++ b/src/main/process-listeners/forwarding-listeners.ts @@ -4,7 +4,7 @@ */ import type { ProcessManager } from '../process-manager'; -import type { ProcessListenerDependencies } from './types'; +import type { ProcessListenerDependencies, ToolExecution } from './types'; /** * Sets up simple forwarding listeners that pass events directly to renderer. @@ -29,12 +29,9 @@ export function setupForwardingListeners( }); // Handle tool execution events (OpenCode, Codex) - processManager.on( - 'tool-execution', - (sessionId: string, toolEvent: { toolName: string; state?: unknown; timestamp: number }) => { - safeSend('process:tool-execution', sessionId, toolEvent); - } - ); + processManager.on('tool-execution', (sessionId: string, toolEvent: ToolExecution) => { + safeSend('process:tool-execution', sessionId, toolEvent); + }); // Handle stderr separately from runCommand (for clean command execution) processManager.on('stderr', (sessionId: string, data: string) => { diff --git a/src/main/process-listeners/session-id-listener.ts b/src/main/process-listeners/session-id-listener.ts index 7d70e394..67245e46 100644 --- a/src/main/process-listeners/session-id-listener.ts +++ b/src/main/process-listeners/session-id-listener.ts @@ -6,6 +6,12 @@ import type { ProcessManager } from '../process-manager'; import type { ProcessListenerDependencies } from './types'; +/** + * Prefix for group chat session IDs. + * Used for fast string check before expensive regex matching. + */ +const GROUP_CHAT_PREFIX = 'group-chat-'; + /** * Sets up the session-id listener. * Handles: @@ -24,20 +30,23 @@ export function setupSessionIdListener( const { REGEX_MODERATOR_SESSION_TIMESTAMP } = patterns; processManager.on('session-id', (sessionId: string, agentSessionId: string) => { + // Fast path: skip regex for non-group-chat sessions (performance optimization) + const isGroupChatSession = sessionId.startsWith(GROUP_CHAT_PREFIX); + // Handle group chat participant session ID - store the agent's session ID // Session ID format: group-chat-{groupChatId}-participant-{name}-{uuid|timestamp} - const participantSessionInfo = outputParser.parseParticipantSessionId(sessionId); + const participantSessionInfo = isGroupChatSession + ? outputParser.parseParticipantSessionId(sessionId) + : null; if (participantSessionInfo) { const { groupChatId, participantName } = participantSessionInfo; // Update the participant with the agent's session ID groupChatStorage .updateParticipant(groupChatId, participantName, { agentSessionId }) - .then(async () => { + .then((updatedChat) => { // Emit participants changed so UI updates with the new session ID - const chat = await groupChatStorage.loadGroupChat(groupChatId); - if (chat) { - groupChatEmitters.emitParticipantsChanged?.(groupChatId, chat.participants); - } + // Note: updateParticipant returns the updated chat, avoiding extra DB read + groupChatEmitters.emitParticipantsChanged?.(groupChatId, updatedChat.participants); }) .catch((err) => { logger.error( @@ -51,7 +60,9 @@ export function setupSessionIdListener( // Handle group chat moderator session ID - store the real agent session ID // Session ID format: group-chat-{groupChatId}-moderator-{timestamp} - const moderatorMatch = sessionId.match(REGEX_MODERATOR_SESSION_TIMESTAMP); + const moderatorMatch = isGroupChatSession + ? sessionId.match(REGEX_MODERATOR_SESSION_TIMESTAMP) + : null; if (moderatorMatch) { const groupChatId = moderatorMatch[1]; // Update the group chat with the moderator's real agent session ID diff --git a/src/main/process-listeners/types.ts b/src/main/process-listeners/types.ts index b230d739..dd0b7256 100644 --- a/src/main/process-listeners/types.ts +++ b/src/main/process-listeners/types.ts @@ -139,6 +139,10 @@ export interface ProcessListenerDependencies { REGEX_MODERATOR_SESSION_TIMESTAMP: RegExp; REGEX_AI_SUFFIX: RegExp; REGEX_AI_TAB_ID: RegExp; + /** Matches batch session IDs: {id}-batch-{timestamp} */ + REGEX_BATCH_SESSION: RegExp; + /** Matches synopsis session IDs: {id}-synopsis-{timestamp} */ + REGEX_SYNOPSIS_SESSION: RegExp; }; /** Logger instance */ logger: { diff --git a/src/main/process-listeners/usage-listener.ts b/src/main/process-listeners/usage-listener.ts index 1330bd42..9fe32fec 100644 --- a/src/main/process-listeners/usage-listener.ts +++ b/src/main/process-listeners/usage-listener.ts @@ -4,7 +4,13 @@ */ import type { ProcessManager } from '../process-manager'; -import type { ProcessListenerDependencies } from './types'; +import type { ProcessListenerDependencies, UsageStats } from './types'; + +/** + * Prefix for group chat session IDs. + * Used for fast string check before expensive regex matching. + */ +const GROUP_CHAT_PREFIX = 'group-chat-'; /** * Sets up the usage listener for token/cost statistics. @@ -38,77 +44,68 @@ export function setupUsageListener( const { REGEX_MODERATOR_SESSION } = patterns; // Handle usage statistics from AI responses - processManager.on( - 'usage', - ( - sessionId: string, - usageStats: { - inputTokens: number; - outputTokens: number; - cacheReadInputTokens: number; - cacheCreationInputTokens: number; - totalCostUsd: number; - contextWindow: number; - reasoningTokens?: number; // Separate reasoning tokens (Codex o3/o4-mini) - } - ) => { - // Handle group chat participant usage - update participant stats - const participantUsageInfo = outputParser.parseParticipantSessionId(sessionId); - if (participantUsageInfo) { - const { groupChatId, participantName } = participantUsageInfo; + processManager.on('usage', (sessionId: string, usageStats: UsageStats) => { + // Fast path: skip regex for non-group-chat sessions (performance optimization) + const isGroupChatSession = sessionId.startsWith(GROUP_CHAT_PREFIX); - // Calculate context usage percentage using agent-specific logic - // Note: For group chat, we don't have agent type here, defaults to Claude behavior - const totalContextTokens = usageAggregator.calculateContextTokens(usageStats); - const contextUsage = - usageStats.contextWindow > 0 - ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) - : 0; + // Handle group chat participant usage - update participant stats + const participantUsageInfo = isGroupChatSession + ? outputParser.parseParticipantSessionId(sessionId) + : null; + if (participantUsageInfo) { + const { groupChatId, participantName } = participantUsageInfo; - // Update participant with usage stats - groupChatStorage - .updateParticipant(groupChatId, participantName, { - contextUsage, - tokenCount: totalContextTokens, - totalCost: usageStats.totalCostUsd, - }) - .then(async () => { - // Emit participants changed so UI updates - const chat = await groupChatStorage.loadGroupChat(groupChatId); - if (chat) { - groupChatEmitters.emitParticipantsChanged?.(groupChatId, chat.participants); - } - }) - .catch((err) => { - logger.error('[GroupChat] Failed to update participant usage', 'ProcessListener', { - error: String(err), - participant: participantName, - }); - }); - // Still send to renderer for consistency - } + // Calculate context usage percentage using agent-specific logic + // Note: For group chat, we don't have agent type here, defaults to Claude behavior + const totalContextTokens = usageAggregator.calculateContextTokens(usageStats); + const contextUsage = + usageStats.contextWindow > 0 + ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) + : 0; - // Handle group chat moderator usage - emit for UI - const moderatorUsageMatch = sessionId.match(REGEX_MODERATOR_SESSION); - if (moderatorUsageMatch) { - const groupChatId = moderatorUsageMatch[1]; - // Calculate context usage percentage using agent-specific logic - // Note: Moderator is typically Claude, defaults to Claude behavior - const totalContextTokens = usageAggregator.calculateContextTokens(usageStats); - const contextUsage = - usageStats.contextWindow > 0 - ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) - : 0; - - // Emit moderator usage for the moderator card - groupChatEmitters.emitModeratorUsage?.(groupChatId, { + // Update participant with usage stats + groupChatStorage + .updateParticipant(groupChatId, participantName, { contextUsage, - totalCost: usageStats.totalCostUsd, tokenCount: totalContextTokens, + totalCost: usageStats.totalCostUsd, + }) + .then((updatedChat) => { + // Emit participants changed so UI updates + // Note: updateParticipant returns the updated chat, avoiding extra DB read + groupChatEmitters.emitParticipantsChanged?.(groupChatId, updatedChat.participants); + }) + .catch((err) => { + logger.error('[GroupChat] Failed to update participant usage', 'ProcessListener', { + error: String(err), + participant: participantName, + }); }); - } - - safeSend('process:usage', sessionId, usageStats); + // Still send to renderer for consistency } - ); + + // Handle group chat moderator usage - emit for UI + const moderatorUsageMatch = isGroupChatSession + ? sessionId.match(REGEX_MODERATOR_SESSION) + : null; + if (moderatorUsageMatch) { + const groupChatId = moderatorUsageMatch[1]; + // Calculate context usage percentage using agent-specific logic + // Note: Moderator is typically Claude, defaults to Claude behavior + const totalContextTokens = usageAggregator.calculateContextTokens(usageStats); + const contextUsage = + usageStats.contextWindow > 0 + ? Math.round((totalContextTokens / usageStats.contextWindow) * 100) + : 0; + + // Emit moderator usage for the moderator card + groupChatEmitters.emitModeratorUsage?.(groupChatId, { + contextUsage, + totalCost: usageStats.totalCostUsd, + tokenCount: totalContextTokens, + }); + } + + safeSend('process:usage', sessionId, usageStats); + }); } From 8c56831ccc8af04bc6939ca6a2184decb41fa650 Mon Sep 17 00:00:00 2001 From: Raza Rauf Date: Tue, 27 Jan 2026 00:48:22 +0500 Subject: [PATCH 03/16] refactor: extract agent error handlers and centralize GROUP_CHAT_PREFIX - Extract agent:clearError and agent:retryAfterError handlers from main/index.ts to dedicated handlers/agent-error.ts module - Add comprehensive test coverage for agent error handlers (29 tests) - Centralize GROUP_CHAT_PREFIX constant in process-listeners/types.ts to eliminate duplication across 4 listener files - Remove unused ipcMain import from main/index.ts (all IPC handlers now registered through handlers module) --- .../main/ipc/handlers/agent-error.test.ts | 350 ++++++++++++++++++ src/main/index.ts | 40 +- src/main/ipc/handlers/agent-error.ts | 73 ++++ src/main/ipc/handlers/index.ts | 4 + src/main/process-listeners/data-listener.ts | 8 +- src/main/process-listeners/exit-listener.ts | 8 +- .../process-listeners/session-id-listener.ts | 8 +- src/main/process-listeners/types.ts | 11 + src/main/process-listeners/usage-listener.ts | 8 +- 9 files changed, 444 insertions(+), 66 deletions(-) create mode 100644 src/__tests__/main/ipc/handlers/agent-error.test.ts create mode 100644 src/main/ipc/handlers/agent-error.ts diff --git a/src/__tests__/main/ipc/handlers/agent-error.test.ts b/src/__tests__/main/ipc/handlers/agent-error.test.ts new file mode 100644 index 00000000..33df7238 --- /dev/null +++ b/src/__tests__/main/ipc/handlers/agent-error.test.ts @@ -0,0 +1,350 @@ +/** + * Tests for agent error IPC handlers + * + * Tests the agent:clearError and agent:retryAfterError handlers + * which manage error state transitions and recovery operations. + */ + +import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; +import { ipcMain } from 'electron'; + +// Create hoisted mocks for more reliable mocking +const mocks = vi.hoisted(() => ({ + mockLogger: { + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, +})); + +// Mock electron +vi.mock('electron', () => ({ + ipcMain: { + handle: vi.fn(), + }, +})); + +// Mock logger +vi.mock('../../../../main/utils/logger', () => ({ + logger: mocks.mockLogger, +})); + +// Alias for easier access in tests +const mockLogger = mocks.mockLogger; + +import { registerAgentErrorHandlers } from '../../../../main/ipc/handlers/agent-error'; + +describe('Agent Error IPC Handlers', () => { + let handlers: Map; + + beforeEach(() => { + vi.clearAllMocks(); + handlers = new Map(); + + // Capture registered handlers + vi.mocked(ipcMain.handle).mockImplementation((channel: string, handler: Function) => { + handlers.set(channel, handler); + }); + + registerAgentErrorHandlers(); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + describe('handler registration', () => { + it('should register agent:clearError handler', () => { + expect(handlers.has('agent:clearError')).toBe(true); + }); + + it('should register agent:retryAfterError handler', () => { + expect(handlers.has('agent:retryAfterError')).toBe(true); + }); + + it('should register exactly 2 handlers', () => { + expect(handlers.size).toBe(2); + }); + }); + + describe('agent:clearError', () => { + it('should return success for valid session ID', async () => { + const handler = handlers.get('agent:clearError')!; + const result = await handler({}, 'session-123'); + + expect(result).toEqual({ success: true }); + }); + + it('should log debug message with session ID', async () => { + const handler = handlers.get('agent:clearError')!; + await handler({}, 'test-session-abc'); + + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Clearing agent error for session', + 'AgentError', + { sessionId: 'test-session-abc' } + ); + }); + + it('should handle empty session ID', async () => { + const handler = handlers.get('agent:clearError')!; + const result = await handler({}, ''); + + expect(result).toEqual({ success: true }); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Clearing agent error for session', + 'AgentError', + { sessionId: '' } + ); + }); + + it('should handle UUID format session ID', async () => { + const handler = handlers.get('agent:clearError')!; + const uuid = 'a1b2c3d4-e5f6-7890-abcd-ef1234567890'; + const result = await handler({}, uuid); + + expect(result).toEqual({ success: true }); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Clearing agent error for session', + 'AgentError', + { sessionId: uuid } + ); + }); + + it('should handle long session ID', async () => { + const handler = handlers.get('agent:clearError')!; + const longSessionId = 'session-' + 'a'.repeat(500); + const result = await handler({}, longSessionId); + + expect(result).toEqual({ success: true }); + }); + + it('should handle special characters in session ID', async () => { + const handler = handlers.get('agent:clearError')!; + const specialId = 'session_with-special.chars:123'; + const result = await handler({}, specialId); + + expect(result).toEqual({ success: true }); + }); + + it('should handle group chat session ID format', async () => { + const handler = handlers.get('agent:clearError')!; + const groupChatId = 'group-chat-test-123-participant-Agent-abc'; + const result = await handler({}, groupChatId); + + expect(result).toEqual({ success: true }); + expect(mockLogger.debug).toHaveBeenCalledWith( + 'Clearing agent error for session', + 'AgentError', + { sessionId: groupChatId } + ); + }); + }); + + describe('agent:retryAfterError', () => { + it('should return success without options', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123'); + + expect(result).toEqual({ success: true }); + }); + + it('should return success with empty options', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123', {}); + + expect(result).toEqual({ success: true }); + }); + + it('should return success with prompt option', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123', { prompt: 'Retry with this prompt' }); + + expect(result).toEqual({ success: true }); + }); + + it('should return success with newSession option', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123', { newSession: true }); + + expect(result).toEqual({ success: true }); + }); + + it('should return success with both options', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123', { + prompt: 'Retry prompt', + newSession: true, + }); + + expect(result).toEqual({ success: true }); + }); + + it('should log info with session ID and no options', async () => { + const handler = handlers.get('agent:retryAfterError')!; + await handler({}, 'test-session-xyz'); + + expect(mockLogger.info).toHaveBeenCalledWith('Retrying after agent error', 'AgentError', { + sessionId: 'test-session-xyz', + hasPrompt: false, + newSession: false, + }); + }); + + it('should log info with hasPrompt=true when prompt provided', async () => { + const handler = handlers.get('agent:retryAfterError')!; + await handler({}, 'session-abc', { prompt: 'Some prompt text' }); + + expect(mockLogger.info).toHaveBeenCalledWith('Retrying after agent error', 'AgentError', { + sessionId: 'session-abc', + hasPrompt: true, + newSession: false, + }); + }); + + it('should log info with newSession=true when specified', async () => { + const handler = handlers.get('agent:retryAfterError')!; + await handler({}, 'session-def', { newSession: true }); + + expect(mockLogger.info).toHaveBeenCalledWith('Retrying after agent error', 'AgentError', { + sessionId: 'session-def', + hasPrompt: false, + newSession: true, + }); + }); + + it('should handle empty prompt string as no prompt', async () => { + const handler = handlers.get('agent:retryAfterError')!; + await handler({}, 'session-123', { prompt: '' }); + + expect(mockLogger.info).toHaveBeenCalledWith('Retrying after agent error', 'AgentError', { + sessionId: 'session-123', + hasPrompt: false, + newSession: false, + }); + }); + + it('should handle undefined newSession as false', async () => { + const handler = handlers.get('agent:retryAfterError')!; + await handler({}, 'session-123', { prompt: 'test' }); + + expect(mockLogger.info).toHaveBeenCalledWith( + 'Retrying after agent error', + 'AgentError', + expect.objectContaining({ newSession: false }) + ); + }); + + it('should handle newSession=false explicitly', async () => { + const handler = handlers.get('agent:retryAfterError')!; + await handler({}, 'session-123', { newSession: false }); + + expect(mockLogger.info).toHaveBeenCalledWith( + 'Retrying after agent error', + 'AgentError', + expect.objectContaining({ newSession: false }) + ); + }); + + it('should handle very long prompt', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const longPrompt = 'A'.repeat(10000); + const result = await handler({}, 'session-123', { prompt: longPrompt }); + + expect(result).toEqual({ success: true }); + expect(mockLogger.info).toHaveBeenCalledWith( + 'Retrying after agent error', + 'AgentError', + expect.objectContaining({ hasPrompt: true }) + ); + }); + + it('should handle unicode in prompt', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123', { prompt: '请重试这个操作 🔄' }); + + expect(result).toEqual({ success: true }); + expect(mockLogger.info).toHaveBeenCalledWith( + 'Retrying after agent error', + 'AgentError', + expect.objectContaining({ hasPrompt: true }) + ); + }); + }); + + describe('handler idempotency', () => { + it('should handle multiple clearError calls for same session', async () => { + const handler = handlers.get('agent:clearError')!; + + const result1 = await handler({}, 'session-123'); + const result2 = await handler({}, 'session-123'); + const result3 = await handler({}, 'session-123'); + + expect(result1).toEqual({ success: true }); + expect(result2).toEqual({ success: true }); + expect(result3).toEqual({ success: true }); + expect(mockLogger.debug).toHaveBeenCalledTimes(3); + }); + + it('should handle multiple retryAfterError calls for same session', async () => { + const handler = handlers.get('agent:retryAfterError')!; + + const result1 = await handler({}, 'session-123', { prompt: 'First retry' }); + const result2 = await handler({}, 'session-123', { prompt: 'Second retry' }); + + expect(result1).toEqual({ success: true }); + expect(result2).toEqual({ success: true }); + expect(mockLogger.info).toHaveBeenCalledTimes(2); + }); + }); + + describe('concurrent handler calls', () => { + it('should handle concurrent clearError calls for different sessions', async () => { + const handler = handlers.get('agent:clearError')!; + + const results = await Promise.all([ + handler({}, 'session-1'), + handler({}, 'session-2'), + handler({}, 'session-3'), + ]); + + expect(results).toEqual([{ success: true }, { success: true }, { success: true }]); + }); + + it('should handle concurrent retryAfterError calls', async () => { + const handler = handlers.get('agent:retryAfterError')!; + + const results = await Promise.all([ + handler({}, 'session-1', { prompt: 'Prompt 1' }), + handler({}, 'session-2', { newSession: true }), + handler({}, 'session-3', { prompt: 'Prompt 3', newSession: false }), + ]); + + expect(results).toEqual([{ success: true }, { success: true }, { success: true }]); + }); + }); + + describe('edge cases', () => { + it('should handle null-ish session ID', async () => { + const handler = handlers.get('agent:clearError')!; + + // TypeScript would normally prevent this, but testing runtime behavior + const result = await handler({}, null as unknown as string); + + expect(result).toEqual({ success: true }); + }); + + it('should handle undefined options in retryAfterError', async () => { + const handler = handlers.get('agent:retryAfterError')!; + const result = await handler({}, 'session-123', undefined); + + expect(result).toEqual({ success: true }); + expect(mockLogger.info).toHaveBeenCalledWith('Retrying after agent error', 'AgentError', { + sessionId: 'session-123', + hasPrompt: false, + newSession: false, + }); + }); + }); +}); diff --git a/src/main/index.ts b/src/main/index.ts index a9029f48..2e9bd077 100644 --- a/src/main/index.ts +++ b/src/main/index.ts @@ -1,4 +1,4 @@ -import { app, BrowserWindow, ipcMain } from 'electron'; +import { app, BrowserWindow } from 'electron'; import path from 'path'; import crypto from 'crypto'; // Sentry is imported dynamically below to avoid module-load-time access to electron.app @@ -566,43 +566,7 @@ function setupIpcHandlers() { // Claude Code sessions - extracted to src/main/ipc/handlers/claude.ts - // ========================================================================== - // Agent Error Handling API - // ========================================================================== - - // Clear an error state for a session (called after recovery action) - ipcMain.handle('agent:clearError', async (_event, sessionId: string) => { - logger.debug('Clearing agent error for session', 'AgentError', { sessionId }); - // Note: The actual error state is managed in the renderer. - // This handler is used to log the clear action and potentially - // perform any main process cleanup needed. - return { success: true }; - }); - - // Retry the last operation after an error (optionally with modified parameters) - ipcMain.handle( - 'agent:retryAfterError', - async ( - _event, - sessionId: string, - options?: { - prompt?: string; - newSession?: boolean; - } - ) => { - logger.info('Retrying after agent error', 'AgentError', { - sessionId, - hasPrompt: !!options?.prompt, - newSession: options?.newSession || false, - }); - // Note: The actual retry logic is handled in the renderer, which will: - // 1. Clear the error state - // 2. Optionally start a new session - // 3. Re-send the last command or the provided prompt - // This handler exists for logging and potential future main process coordination. - return { success: true }; - } - ); + // Agent Error Handling API - extracted to src/main/ipc/handlers/agent-error.ts // Register notification handlers (extracted to handlers/notifications.ts) registerNotificationsHandlers(); diff --git a/src/main/ipc/handlers/agent-error.ts b/src/main/ipc/handlers/agent-error.ts new file mode 100644 index 00000000..108a1aa8 --- /dev/null +++ b/src/main/ipc/handlers/agent-error.ts @@ -0,0 +1,73 @@ +/** + * Agent Error Handling IPC Handlers + * + * Handles agent error state management: + * - Clearing error states after recovery + * - Retrying operations after errors + * + * Note: The actual error state is managed in the renderer. These handlers + * provide logging and potential future main process coordination. + */ + +import { ipcMain } from 'electron'; +import { logger } from '../../utils/logger'; + +// ========================================================================== +// Types +// ========================================================================== + +/** + * Options for retrying after an error + */ +export interface RetryOptions { + /** Optional prompt to use for the retry */ + prompt?: string; + /** Whether to start a new session for the retry */ + newSession?: boolean; +} + +/** + * Response from agent error operations + */ +export interface AgentErrorResponse { + success: boolean; +} + +// ========================================================================== +// Handler Registration +// ========================================================================== + +/** + * Register all agent error-related IPC handlers + */ +export function registerAgentErrorHandlers(): void { + // Clear an error state for a session (called after recovery action) + ipcMain.handle( + 'agent:clearError', + async (_event, sessionId: string): Promise => { + logger.debug('Clearing agent error for session', 'AgentError', { sessionId }); + // Note: The actual error state is managed in the renderer. + // This handler is used to log the clear action and potentially + // perform any main process cleanup needed. + return { success: true }; + } + ); + + // Retry the last operation after an error (optionally with modified parameters) + ipcMain.handle( + 'agent:retryAfterError', + async (_event, sessionId: string, options?: RetryOptions): Promise => { + logger.info('Retrying after agent error', 'AgentError', { + sessionId, + hasPrompt: !!options?.prompt, + newSession: options?.newSession || false, + }); + // Note: The actual retry logic is handled in the renderer, which will: + // 1. Clear the error state + // 2. Optionally start a new session + // 3. Re-send the last command or the provided prompt + // This handler exists for logging and potential future main process coordination. + return { success: true }; + } + ); +} diff --git a/src/main/ipc/handlers/index.ts b/src/main/ipc/handlers/index.ts index f06b442b..9a7025f0 100644 --- a/src/main/ipc/handlers/index.ts +++ b/src/main/ipc/handlers/index.ts @@ -48,6 +48,7 @@ import { registerAttachmentsHandlers, AttachmentsHandlerDependencies } from './a import { registerWebHandlers, WebHandlerDependencies } from './web'; import { registerLeaderboardHandlers, LeaderboardHandlerDependencies } from './leaderboard'; import { registerNotificationsHandlers } from './notifications'; +import { registerAgentErrorHandlers } from './agent-error'; import { AgentDetector } from '../../agent-detector'; import { ProcessManager } from '../../process-manager'; import { WebServer } from '../../web-server'; @@ -85,6 +86,7 @@ export type { WebHandlerDependencies }; export { registerLeaderboardHandlers }; export type { LeaderboardHandlerDependencies }; export { registerNotificationsHandlers }; +export { registerAgentErrorHandlers }; export type { AgentsHandlerDependencies }; export type { ProcessHandlerDependencies }; export type { PersistenceHandlerDependencies }; @@ -246,6 +248,8 @@ export function registerAllHandlers(deps: HandlerDependencies): void { }); // Register notification handlers (OS notifications and TTS) registerNotificationsHandlers(); + // Register agent error handlers (error state management) + registerAgentErrorHandlers(); // Setup logger event forwarding to renderer setupLoggerEventForwarding(deps.getMainWindow); } diff --git a/src/main/process-listeners/data-listener.ts b/src/main/process-listeners/data-listener.ts index 642f73c6..b9580294 100644 --- a/src/main/process-listeners/data-listener.ts +++ b/src/main/process-listeners/data-listener.ts @@ -4,7 +4,7 @@ */ import type { ProcessManager } from '../process-manager'; -import type { ProcessListenerDependencies } from './types'; +import { GROUP_CHAT_PREFIX, type ProcessListenerDependencies } from './types'; /** * Maximum buffer size per session (10MB). @@ -12,12 +12,6 @@ import type { ProcessListenerDependencies } from './types'; */ const MAX_BUFFER_SIZE = 10 * 1024 * 1024; -/** - * Prefix for group chat session IDs. - * Used for fast string check before expensive regex matching. - */ -const GROUP_CHAT_PREFIX = 'group-chat-'; - /** * Length of random suffix in message IDs (9 characters of base36). * Combined with timestamp provides uniqueness for web broadcast deduplication. diff --git a/src/main/process-listeners/exit-listener.ts b/src/main/process-listeners/exit-listener.ts index 8c0ea4c1..0b5152b1 100644 --- a/src/main/process-listeners/exit-listener.ts +++ b/src/main/process-listeners/exit-listener.ts @@ -5,13 +5,7 @@ */ import type { ProcessManager } from '../process-manager'; -import type { ProcessListenerDependencies } from './types'; - -/** - * Prefix for group chat session IDs. - * Used for fast string check before expensive regex matching. - */ -const GROUP_CHAT_PREFIX = 'group-chat-'; +import { GROUP_CHAT_PREFIX, type ProcessListenerDependencies } from './types'; /** * Sets up the exit listener for process termination. diff --git a/src/main/process-listeners/session-id-listener.ts b/src/main/process-listeners/session-id-listener.ts index 67245e46..9b2c1cc7 100644 --- a/src/main/process-listeners/session-id-listener.ts +++ b/src/main/process-listeners/session-id-listener.ts @@ -4,13 +4,7 @@ */ import type { ProcessManager } from '../process-manager'; -import type { ProcessListenerDependencies } from './types'; - -/** - * Prefix for group chat session IDs. - * Used for fast string check before expensive regex matching. - */ -const GROUP_CHAT_PREFIX = 'group-chat-'; +import { GROUP_CHAT_PREFIX, type ProcessListenerDependencies } from './types'; /** * Sets up the session-id listener. diff --git a/src/main/process-listeners/types.ts b/src/main/process-listeners/types.ts index dd0b7256..ec8559c2 100644 --- a/src/main/process-listeners/types.ts +++ b/src/main/process-listeners/types.ts @@ -4,6 +4,17 @@ */ import type { ProcessManager } from '../process-manager'; + +// ========================================================================== +// Constants +// ========================================================================== + +/** + * Prefix for group chat session IDs. + * Used for fast string check before expensive regex matching. + * Session IDs starting with this prefix belong to group chat sessions. + */ +export const GROUP_CHAT_PREFIX = 'group-chat-'; import type { WebServer } from '../web-server'; import type { AgentDetector } from '../agent-detector'; import type { SafeSendFn } from '../utils/safe-send'; diff --git a/src/main/process-listeners/usage-listener.ts b/src/main/process-listeners/usage-listener.ts index 9fe32fec..757bd0c9 100644 --- a/src/main/process-listeners/usage-listener.ts +++ b/src/main/process-listeners/usage-listener.ts @@ -4,13 +4,7 @@ */ import type { ProcessManager } from '../process-manager'; -import type { ProcessListenerDependencies, UsageStats } from './types'; - -/** - * Prefix for group chat session IDs. - * Used for fast string check before expensive regex matching. - */ -const GROUP_CHAT_PREFIX = 'group-chat-'; +import { GROUP_CHAT_PREFIX, type ProcessListenerDependencies, type UsageStats } from './types'; /** * Sets up the usage listener for token/cost statistics. From 31f2a93fea94232d371fe9178d37f8afc13dff28 Mon Sep 17 00:00:00 2001 From: Raza Rauf Date: Tue, 27 Jan 2026 00:59:08 +0500 Subject: [PATCH 04/16] fix: address code review feedback - retry logic and import order - Add retry logic with exponential backoff for stats DB insertions (100ms, 200ms, 400ms delays, max 3 attempts) - Fix import order in types.ts (move constant after imports) - Update stats-listener tests for async retry behavior - Add new test for retry success on transient failure Addresses review recommendations: - High: Stats database error handling with retry logic - Low: Import order consistency --- .../__tests__/stats-listener.test.ts | 115 +++++++++++++----- src/main/process-listeners/stats-listener.ts | 91 +++++++++++--- src/main/process-listeners/types.ts | 14 +-- 3 files changed, 167 insertions(+), 53 deletions(-) diff --git a/src/main/process-listeners/__tests__/stats-listener.test.ts b/src/main/process-listeners/__tests__/stats-listener.test.ts index 41e3836d..ff2a0a30 100644 --- a/src/main/process-listeners/__tests__/stats-listener.test.ts +++ b/src/main/process-listeners/__tests__/stats-listener.test.ts @@ -32,7 +32,7 @@ describe('Stats Listener', () => { mockStatsDB = { isReady: vi.fn(() => true), - insertQueryEvent: vi.fn(() => 1), + insertQueryEvent: vi.fn(() => 'event-id-123'), } as unknown as StatsDB; mockProcessManager = { @@ -52,7 +52,7 @@ describe('Stats Listener', () => { expect(mockProcessManager.on).toHaveBeenCalledWith('query-complete', expect.any(Function)); }); - it('should record query event to stats database when ready', () => { + it('should record query event to stats database when ready', async () => { setupStatsListener(mockProcessManager, { safeSend: mockSafeSend, getStatsDB: () => mockStatsDB, @@ -73,17 +73,20 @@ describe('Stats Listener', () => { handler?.(testSessionId, testQueryData); - expect(mockStatsDB.isReady).toHaveBeenCalled(); - expect(mockStatsDB.insertQueryEvent).toHaveBeenCalledWith({ - sessionId: testQueryData.sessionId, - agentType: testQueryData.agentType, - source: testQueryData.source, - startTime: testQueryData.startTime, - duration: testQueryData.duration, - projectPath: testQueryData.projectPath, - tabId: testQueryData.tabId, + // Wait for async processing + await vi.waitFor(() => { + expect(mockStatsDB.isReady).toHaveBeenCalled(); + expect(mockStatsDB.insertQueryEvent).toHaveBeenCalledWith({ + sessionId: testQueryData.sessionId, + agentType: testQueryData.agentType, + source: testQueryData.source, + startTime: testQueryData.startTime, + duration: testQueryData.duration, + projectPath: testQueryData.projectPath, + tabId: testQueryData.tabId, + }); + expect(mockSafeSend).toHaveBeenCalledWith('stats:updated'); }); - expect(mockSafeSend).toHaveBeenCalledWith('stats:updated'); }); it('should not record event when stats database is not ready', () => { @@ -113,7 +116,7 @@ describe('Stats Listener', () => { expect(mockSafeSend).not.toHaveBeenCalled(); }); - it('should log error when recording fails', () => { + it('should log error when recording fails after retries', async () => { vi.mocked(mockStatsDB.insertQueryEvent).mockImplementation(() => { throw new Error('Database error'); }); @@ -137,16 +140,26 @@ describe('Stats Listener', () => { handler?.('session-789', testQueryData); - expect(mockLogger.error).toHaveBeenCalledWith( - expect.stringContaining('Failed to record query event'), - '[Stats]', - expect.objectContaining({ - sessionId: 'session-789', - }) + // Wait for all retries to complete (100ms + 200ms + final attempt) + await vi.waitFor( + () => { + expect(mockLogger.error).toHaveBeenCalledWith( + expect.stringContaining('Failed to record query event after 3 attempts'), + '[Stats]', + expect.objectContaining({ + sessionId: 'session-789', + }) + ); + }, + { timeout: 1000 } ); + // Should have tried 3 times + expect(mockStatsDB.insertQueryEvent).toHaveBeenCalledTimes(3); + // Should not have broadcasted update on failure + expect(mockSafeSend).not.toHaveBeenCalled(); }); - it('should log debug info when recording succeeds', () => { + it('should log debug info when recording succeeds', async () => { setupStatsListener(mockProcessManager, { safeSend: mockSafeSend, getStatsDB: () => mockStatsDB, @@ -166,15 +179,61 @@ describe('Stats Listener', () => { handler?.('session-abc', testQueryData); - expect(mockLogger.debug).toHaveBeenCalledWith( - expect.stringContaining('Recorded query event'), - '[Stats]', - expect.objectContaining({ - sessionId: 'session-abc', - agentType: 'claude-code', - source: 'user', - duration: 3000, + // Wait for async processing + await vi.waitFor(() => { + expect(mockLogger.debug).toHaveBeenCalledWith( + expect.stringContaining('Recorded query event'), + '[Stats]', + expect.objectContaining({ + sessionId: 'session-abc', + agentType: 'claude-code', + source: 'user', + duration: 3000, + }) + ); + }); + }); + + it('should retry on transient failure and succeed', async () => { + // First call fails, second succeeds + vi.mocked(mockStatsDB.insertQueryEvent) + .mockImplementationOnce(() => { + throw new Error('Transient error'); }) + .mockImplementationOnce(() => 'event-id-456'); + + setupStatsListener(mockProcessManager, { + safeSend: mockSafeSend, + getStatsDB: () => mockStatsDB, + logger: mockLogger, + }); + + const handler = eventHandlers.get('query-complete'); + const testQueryData: QueryCompleteData = { + sessionId: 'session-retry', + agentType: 'claude-code', + source: 'user', + startTime: Date.now(), + duration: 1000, + projectPath: '/test/project', + tabId: 'tab-retry', + }; + + handler?.('session-retry', testQueryData); + + // Wait for retry to complete + await vi.waitFor( + () => { + expect(mockStatsDB.insertQueryEvent).toHaveBeenCalledTimes(2); + expect(mockSafeSend).toHaveBeenCalledWith('stats:updated'); + }, + { timeout: 500 } + ); + // Should have logged warning for first failure + expect(mockLogger.warn).toHaveBeenCalledWith( + expect.stringContaining('Stats DB insert failed'), + '[Stats]', + expect.any(Object) ); }); }); diff --git a/src/main/process-listeners/stats-listener.ts b/src/main/process-listeners/stats-listener.ts index f831ef20..b37a8c18 100644 --- a/src/main/process-listeners/stats-listener.ts +++ b/src/main/process-listeners/stats-listener.ts @@ -7,9 +7,70 @@ import type { ProcessManager } from '../process-manager'; import type { QueryCompleteData } from '../process-manager/types'; import type { ProcessListenerDependencies } from './types'; +/** + * Maximum number of retry attempts for transient database failures. + */ +const MAX_RETRY_ATTEMPTS = 3; + +/** + * Base delay in milliseconds for exponential backoff (doubles each retry). + */ +const RETRY_BASE_DELAY_MS = 100; + +/** + * Attempts to insert a query event with retry logic for transient failures. + * Uses exponential backoff: 100ms, 200ms, 400ms delays between retries. + */ +async function insertQueryEventWithRetry( + db: ReturnType, + queryData: QueryCompleteData, + logger: ProcessListenerDependencies['logger'] +): Promise { + for (let attempt = 1; attempt <= MAX_RETRY_ATTEMPTS; attempt++) { + try { + const id = db.insertQueryEvent({ + sessionId: queryData.sessionId, + agentType: queryData.agentType, + source: queryData.source, + startTime: queryData.startTime, + duration: queryData.duration, + projectPath: queryData.projectPath, + tabId: queryData.tabId, + }); + return id; + } catch (error) { + const isLastAttempt = attempt === MAX_RETRY_ATTEMPTS; + + if (isLastAttempt) { + logger.error( + `Failed to record query event after ${MAX_RETRY_ATTEMPTS} attempts`, + '[Stats]', + { + error: String(error), + sessionId: queryData.sessionId, + } + ); + } else { + const delay = RETRY_BASE_DELAY_MS * Math.pow(2, attempt - 1); + logger.warn( + `Stats DB insert failed (attempt ${attempt}/${MAX_RETRY_ATTEMPTS}), retrying in ${delay}ms`, + '[Stats]', + { + error: String(error), + sessionId: queryData.sessionId, + } + ); + await new Promise((resolve) => setTimeout(resolve, delay)); + } + } + } + + return null; +} + /** * Sets up the query-complete listener for stats tracking. - * Records AI query events to the stats database. + * Records AI query events to the stats database with retry logic for transient failures. */ export function setupStatsListener( processManager: ProcessManager, @@ -20,18 +81,16 @@ export function setupStatsListener( // Handle query-complete events for stats tracking // This is emitted when a batch mode AI query completes (user or auto) processManager.on('query-complete', (_sessionId: string, queryData: QueryCompleteData) => { - try { - const db = getStatsDB(); - if (db.isReady()) { - const id = db.insertQueryEvent({ - sessionId: queryData.sessionId, - agentType: queryData.agentType, - source: queryData.source, - startTime: queryData.startTime, - duration: queryData.duration, - projectPath: queryData.projectPath, - tabId: queryData.tabId, - }); + const db = getStatsDB(); + if (!db.isReady()) { + return; + } + + // Use async IIFE to handle retry logic without blocking + void (async () => { + const id = await insertQueryEventWithRetry(db, queryData, logger); + + if (id !== null) { logger.debug(`Recorded query event: ${id}`, '[Stats]', { sessionId: queryData.sessionId, agentType: queryData.agentType, @@ -41,10 +100,6 @@ export function setupStatsListener( // Broadcast stats update to renderer for real-time dashboard refresh safeSend('stats:updated'); } - } catch (error) { - logger.error(`Failed to record query event: ${error}`, '[Stats]', { - sessionId: queryData.sessionId, - }); - } + })(); }); } diff --git a/src/main/process-listeners/types.ts b/src/main/process-listeners/types.ts index ec8559c2..b80f267d 100644 --- a/src/main/process-listeners/types.ts +++ b/src/main/process-listeners/types.ts @@ -4,6 +4,13 @@ */ import type { ProcessManager } from '../process-manager'; +import type { WebServer } from '../web-server'; +import type { AgentDetector } from '../agent-detector'; +import type { SafeSendFn } from '../utils/safe-send'; +import type { StatsDB } from '../stats-db'; +import type { GroupChat, GroupChatParticipant } from '../group-chat/group-chat-storage'; +import type { GroupChatState } from '../../shared/group-chat-types'; +import type { ParticipantState } from '../ipc/handlers/groupChat'; // ========================================================================== // Constants @@ -15,13 +22,6 @@ import type { ProcessManager } from '../process-manager'; * Session IDs starting with this prefix belong to group chat sessions. */ export const GROUP_CHAT_PREFIX = 'group-chat-'; -import type { WebServer } from '../web-server'; -import type { AgentDetector } from '../agent-detector'; -import type { SafeSendFn } from '../utils/safe-send'; -import type { StatsDB } from '../stats-db'; -import type { GroupChat, GroupChatParticipant } from '../group-chat/group-chat-storage'; -import type { GroupChatState } from '../../shared/group-chat-types'; -import type { ParticipantState } from '../ipc/handlers/groupChat'; // Re-export types from their canonical locations export type { UsageStats, QueryCompleteData, ToolExecution } from '../process-manager/types'; From b37b8fb91baed287eb28d83378026ff589937ff0 Mon Sep 17 00:00:00 2001 From: Raza Rauf Date: Thu, 29 Jan 2026 01:19:56 +0500 Subject: [PATCH 05/16] refactor: restructure web-server module with security and memory leak fixes - Move WebServer class to dedicated file and add module index - Extract shared types to centralized types.ts - Fix XSS vulnerability by sanitizing sessionId/tabId in URL parameters - Fix IPC listener memory leak with proper cleanup on timeout - Add autoRunStates cleanup when sessions go offline - Refactor message handlers with send() and sendError() helpers - Add XSS sanitization tests and e2e test configuration --- .../main/web-server/routes/apiRoutes.test.ts | 457 ++++++++++++++++ .../web-server/routes/staticRoutes.test.ts | 239 +++++++++ .../main/web-server/routes/wsRoute.test.ts | 497 ++++++++++++++++++ .../web-server/web-server-factory.test.ts | 5 +- .../WebServer.ts} | 245 +++------ .../web-server/handlers/messageHandlers.ts | 313 +++-------- src/main/web-server/index.ts | 62 +++ src/main/web-server/routes/apiRoutes.ts | 121 +---- src/main/web-server/routes/staticRoutes.ts | 25 +- src/main/web-server/routes/wsRoute.ts | 50 +- .../web-server/services/broadcastService.ts | 103 +--- src/main/web-server/types.ts | 302 +++++++++++ src/main/web-server/web-server-factory.ts | 24 +- vitest.e2e.config.ts | 31 ++ 14 files changed, 1809 insertions(+), 665 deletions(-) create mode 100644 src/__tests__/main/web-server/routes/apiRoutes.test.ts create mode 100644 src/__tests__/main/web-server/routes/staticRoutes.test.ts create mode 100644 src/__tests__/main/web-server/routes/wsRoute.test.ts rename src/main/{web-server.ts => web-server/WebServer.ts} (79%) create mode 100644 src/main/web-server/index.ts create mode 100644 src/main/web-server/types.ts create mode 100644 vitest.e2e.config.ts diff --git a/src/__tests__/main/web-server/routes/apiRoutes.test.ts b/src/__tests__/main/web-server/routes/apiRoutes.test.ts new file mode 100644 index 00000000..82bc6a9a --- /dev/null +++ b/src/__tests__/main/web-server/routes/apiRoutes.test.ts @@ -0,0 +1,457 @@ +/** + * Tests for ApiRoutes + * + * API Routes handle REST API requests from web clients. + * Routes are protected by a security token prefix. + * + * Endpoints tested: + * - GET /api/sessions - List all sessions with live info + * - GET /api/session/:id - Get single session detail + * - POST /api/session/:id/send - Send command to session + * - GET /api/theme - Get current theme + * - POST /api/session/:id/interrupt - Interrupt session + * - GET /api/history - Get history entries + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { + ApiRoutes, + type ApiRouteCallbacks, + type RateLimitConfig, +} from '../../../../main/web-server/routes/apiRoutes'; + +// Mock the logger +vi.mock('../../../../main/utils/logger', () => ({ + logger: { + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, +})); + +/** + * Create mock callbacks with all methods as vi.fn() + */ +function createMockCallbacks(): ApiRouteCallbacks { + return { + getSessions: vi.fn().mockReturnValue([ + { + id: 'session-1', + name: 'Session 1', + toolType: 'claude-code', + state: 'idle', + inputMode: 'ai', + cwd: '/test/project', + groupId: null, + }, + { + id: 'session-2', + name: 'Session 2', + toolType: 'codex', + state: 'busy', + inputMode: 'terminal', + cwd: '/test/project2', + groupId: 'group-1', + }, + ]), + getSessionDetail: vi.fn().mockReturnValue({ + id: 'session-1', + name: 'Session 1', + toolType: 'claude-code', + state: 'idle', + inputMode: 'ai', + cwd: '/test/project', + aiTabs: [{ id: 'tab-1', name: 'Tab 1', logs: [] }], + activeAITabId: 'tab-1', + }), + getTheme: vi.fn().mockReturnValue({ + name: 'dark', + background: '#1a1a1a', + foreground: '#ffffff', + }), + writeToSession: vi.fn().mockReturnValue(true), + interruptSession: vi.fn().mockResolvedValue(true), + getHistory: vi + .fn() + .mockReturnValue([{ id: '1', command: 'test command', timestamp: Date.now() }]), + getLiveSessionInfo: vi.fn().mockReturnValue({ + sessionId: 'session-1', + agentSessionId: 'claude-agent-123', + enabledAt: Date.now(), + }), + isSessionLive: vi.fn().mockReturnValue(true), + }; +} + +/** + * Mock Fastify instance with route registration tracking + */ +function createMockFastify() { + const routes: Map = new Map(); + + return { + get: vi.fn((path: string, options: any, handler?: Function) => { + const h = handler || options; + const config = handler ? options?.config : undefined; + routes.set(`GET:${path}`, { handler: h, config }); + }), + post: vi.fn((path: string, options: any, handler?: Function) => { + const h = handler || options; + const config = handler ? options?.config : undefined; + routes.set(`POST:${path}`, { handler: h, config }); + }), + getRoute: (method: string, path: string) => routes.get(`${method}:${path}`), + routes, + }; +} + +/** + * Mock reply object + */ +function createMockReply() { + const reply: any = { + code: vi.fn().mockReturnThis(), + send: vi.fn().mockReturnThis(), + type: vi.fn().mockReturnThis(), + }; + return reply; +} + +describe('ApiRoutes', () => { + const securityToken = 'test-token-123'; + const rateLimitConfig: RateLimitConfig = { + max: 100, + maxPost: 30, + timeWindow: 60000, // 1 minute in milliseconds + enabled: true, + }; + + let apiRoutes: ApiRoutes; + let callbacks: ApiRouteCallbacks; + let mockFastify: ReturnType; + + beforeEach(() => { + apiRoutes = new ApiRoutes(securityToken, rateLimitConfig); + callbacks = createMockCallbacks(); + apiRoutes.setCallbacks(callbacks); + mockFastify = createMockFastify(); + apiRoutes.registerRoutes(mockFastify as any); + }); + + describe('Route Registration', () => { + it('should register all API routes', () => { + expect(mockFastify.get).toHaveBeenCalledTimes(4); // sessions, session/:id, theme, history + expect(mockFastify.post).toHaveBeenCalledTimes(2); // send, interrupt + }); + + it('should register routes with correct token prefix', () => { + expect(mockFastify.routes.has(`GET:/${securityToken}/api/sessions`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/api/session/:id`)).toBe(true); + expect(mockFastify.routes.has(`POST:/${securityToken}/api/session/:id/send`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/api/theme`)).toBe(true); + expect(mockFastify.routes.has(`POST:/${securityToken}/api/session/:id/interrupt`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/api/history`)).toBe(true); + }); + + it('should configure rate limiting for GET routes', () => { + const sessionsRoute = mockFastify.getRoute('GET', `/${securityToken}/api/sessions`); + expect(sessionsRoute?.config?.rateLimit?.max).toBe(rateLimitConfig.max); + expect(sessionsRoute?.config?.rateLimit?.timeWindow).toBe(rateLimitConfig.timeWindow); + }); + + it('should configure stricter rate limiting for POST routes', () => { + const sendRoute = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/send`); + expect(sendRoute?.config?.rateLimit?.max).toBe(rateLimitConfig.maxPost); + }); + }); + + describe('GET /api/sessions', () => { + it('should return all sessions with live info', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/sessions`); + const result = await route!.handler(); + + expect(result.sessions).toHaveLength(2); + expect(result.count).toBe(2); + expect(result.timestamp).toBeDefined(); + expect(callbacks.getSessions).toHaveBeenCalled(); + }); + + it('should enrich sessions with live info', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/sessions`); + const result = await route!.handler(); + + expect(result.sessions[0].agentSessionId).toBe('claude-agent-123'); + expect(result.sessions[0].isLive).toBe(true); + expect(result.sessions[0].liveEnabledAt).toBeDefined(); + }); + + it('should return empty array when no callbacks configured', async () => { + const emptyRoutes = new ApiRoutes(securityToken, rateLimitConfig); + const emptyFastify = createMockFastify(); + emptyRoutes.registerRoutes(emptyFastify as any); + + const route = emptyFastify.getRoute('GET', `/${securityToken}/api/sessions`); + const result = await route!.handler(); + + expect(result.sessions).toEqual([]); + expect(result.count).toBe(0); + }); + }); + + describe('GET /api/session/:id', () => { + it('should return session detail', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/session/:id`); + const reply = createMockReply(); + const result = await route!.handler({ params: { id: 'session-1' }, query: {} }, reply); + + expect(result.session.id).toBe('session-1'); + expect(result.session.agentSessionId).toBe('claude-agent-123'); + expect(result.session.isLive).toBe(true); + expect(callbacks.getSessionDetail).toHaveBeenCalledWith('session-1', undefined); + }); + + it('should pass tabId query param to callback', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/session/:id`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' }, query: { tabId: 'tab-5' } }, reply); + + expect(callbacks.getSessionDetail).toHaveBeenCalledWith('session-1', 'tab-5'); + }); + + it('should return 404 for non-existent session', async () => { + (callbacks.getSessionDetail as any).mockReturnValue(null); + + const route = mockFastify.getRoute('GET', `/${securityToken}/api/session/:id`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'nonexistent' }, query: {} }, reply); + + expect(reply.code).toHaveBeenCalledWith(404); + expect(reply.send).toHaveBeenCalledWith(expect.objectContaining({ error: 'Not Found' })); + }); + + it('should return 503 when getSessionDetail callback not configured', async () => { + const emptyRoutes = new ApiRoutes(securityToken, rateLimitConfig); + const emptyFastify = createMockFastify(); + emptyRoutes.registerRoutes(emptyFastify as any); + + const route = emptyFastify.getRoute('GET', `/${securityToken}/api/session/:id`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' }, query: {} }, reply); + + expect(reply.code).toHaveBeenCalledWith(503); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ error: 'Service Unavailable' }) + ); + }); + }); + + describe('POST /api/session/:id/send', () => { + it('should send command to session', async () => { + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/send`); + const reply = createMockReply(); + const result = await route!.handler( + { params: { id: 'session-1' }, body: { command: 'ls -la' } }, + reply + ); + + expect(result.success).toBe(true); + expect(result.sessionId).toBe('session-1'); + expect(callbacks.writeToSession).toHaveBeenCalledWith('session-1', 'ls -la\n'); + }); + + it('should return 400 for missing command', async () => { + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/send`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' }, body: {} }, reply); + + expect(reply.code).toHaveBeenCalledWith(400); + expect(reply.send).toHaveBeenCalledWith(expect.objectContaining({ error: 'Bad Request' })); + }); + + it('should return 400 for non-string command', async () => { + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/send`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' }, body: { command: 123 } }, reply); + + expect(reply.code).toHaveBeenCalledWith(400); + }); + + it('should return 500 when writeToSession fails', async () => { + (callbacks.writeToSession as any).mockReturnValue(false); + + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/send`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' }, body: { command: 'test' } }, reply); + + expect(reply.code).toHaveBeenCalledWith(500); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ error: 'Internal Server Error' }) + ); + }); + + it('should return 503 when writeToSession callback not configured', async () => { + const emptyRoutes = new ApiRoutes(securityToken, rateLimitConfig); + const emptyFastify = createMockFastify(); + emptyRoutes.registerRoutes(emptyFastify as any); + + const route = emptyFastify.getRoute('POST', `/${securityToken}/api/session/:id/send`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' }, body: { command: 'test' } }, reply); + + expect(reply.code).toHaveBeenCalledWith(503); + }); + }); + + describe('GET /api/theme', () => { + it('should return current theme', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/theme`); + const reply = createMockReply(); + const result = await route!.handler({}, reply); + + expect(result.theme.name).toBe('dark'); + expect(result.timestamp).toBeDefined(); + expect(callbacks.getTheme).toHaveBeenCalled(); + }); + + it('should return 404 when no theme configured', async () => { + (callbacks.getTheme as any).mockReturnValue(null); + + const route = mockFastify.getRoute('GET', `/${securityToken}/api/theme`); + const reply = createMockReply(); + await route!.handler({}, reply); + + expect(reply.code).toHaveBeenCalledWith(404); + expect(reply.send).toHaveBeenCalledWith(expect.objectContaining({ error: 'Not Found' })); + }); + + it('should return 503 when getTheme callback not configured', async () => { + const emptyRoutes = new ApiRoutes(securityToken, rateLimitConfig); + const emptyFastify = createMockFastify(); + emptyRoutes.registerRoutes(emptyFastify as any); + + const route = emptyFastify.getRoute('GET', `/${securityToken}/api/theme`); + const reply = createMockReply(); + await route!.handler({}, reply); + + expect(reply.code).toHaveBeenCalledWith(503); + }); + }); + + describe('POST /api/session/:id/interrupt', () => { + it('should interrupt session successfully', async () => { + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/interrupt`); + const reply = createMockReply(); + const result = await route!.handler({ params: { id: 'session-1' } }, reply); + + expect(result.success).toBe(true); + expect(result.sessionId).toBe('session-1'); + expect(callbacks.interruptSession).toHaveBeenCalledWith('session-1'); + }); + + it('should return 500 when interrupt fails', async () => { + (callbacks.interruptSession as any).mockResolvedValue(false); + + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/interrupt`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' } }, reply); + + expect(reply.code).toHaveBeenCalledWith(500); + }); + + it('should return 500 when interrupt throws error', async () => { + (callbacks.interruptSession as any).mockRejectedValue(new Error('Session not found')); + + const route = mockFastify.getRoute('POST', `/${securityToken}/api/session/:id/interrupt`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' } }, reply); + + expect(reply.code).toHaveBeenCalledWith(500); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining('Session not found'), + }) + ); + }); + + it('should return 503 when interruptSession callback not configured', async () => { + const emptyRoutes = new ApiRoutes(securityToken, rateLimitConfig); + const emptyFastify = createMockFastify(); + emptyRoutes.registerRoutes(emptyFastify as any); + + const route = emptyFastify.getRoute('POST', `/${securityToken}/api/session/:id/interrupt`); + const reply = createMockReply(); + await route!.handler({ params: { id: 'session-1' } }, reply); + + expect(reply.code).toHaveBeenCalledWith(503); + }); + }); + + describe('GET /api/history', () => { + it('should return history entries', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/history`); + const reply = createMockReply(); + const result = await route!.handler({ query: {} }, reply); + + expect(result.entries).toHaveLength(1); + expect(result.count).toBe(1); + expect(callbacks.getHistory).toHaveBeenCalledWith(undefined, undefined); + }); + + it('should pass projectPath and sessionId to callback', async () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/api/history`); + const reply = createMockReply(); + await route!.handler({ query: { projectPath: '/test', sessionId: 'session-1' } }, reply); + + expect(callbacks.getHistory).toHaveBeenCalledWith('/test', 'session-1'); + }); + + it('should return 500 when getHistory throws error', async () => { + (callbacks.getHistory as any).mockImplementation(() => { + throw new Error('Database error'); + }); + + const route = mockFastify.getRoute('GET', `/${securityToken}/api/history`); + const reply = createMockReply(); + await route!.handler({ query: {} }, reply); + + expect(reply.code).toHaveBeenCalledWith(500); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ + message: expect.stringContaining('Database error'), + }) + ); + }); + + it('should return 503 when getHistory callback not configured', async () => { + const emptyRoutes = new ApiRoutes(securityToken, rateLimitConfig); + const emptyFastify = createMockFastify(); + emptyRoutes.registerRoutes(emptyFastify as any); + + const route = emptyFastify.getRoute('GET', `/${securityToken}/api/history`); + const reply = createMockReply(); + await route!.handler({ query: {} }, reply); + + expect(reply.code).toHaveBeenCalledWith(503); + }); + }); + + describe('Rate Limit Configuration', () => { + it('should update rate limit config', () => { + const newConfig: RateLimitConfig = { + max: 200, + maxPost: 50, + timeWindow: 120000, // 2 minutes in milliseconds + enabled: true, + }; + apiRoutes.updateRateLimitConfig(newConfig); + + // Re-register routes to see new config + const newFastify = createMockFastify(); + apiRoutes.registerRoutes(newFastify as any); + + const sessionsRoute = newFastify.getRoute('GET', `/${securityToken}/api/sessions`); + expect(sessionsRoute?.config?.rateLimit?.max).toBe(200); + }); + }); +}); diff --git a/src/__tests__/main/web-server/routes/staticRoutes.test.ts b/src/__tests__/main/web-server/routes/staticRoutes.test.ts new file mode 100644 index 00000000..642f2c79 --- /dev/null +++ b/src/__tests__/main/web-server/routes/staticRoutes.test.ts @@ -0,0 +1,239 @@ +/** + * Tests for StaticRoutes + * + * Static Routes handle dashboard views, PWA files, and security redirects. + * Routes are protected by a security token prefix. + * + * Note: Tests that require fs mocking are skipped due to ESM module limitations. + * The fs-dependent functionality is tested via integration tests. + * + * Routes tested: + * - / - Redirect to website (no access without token) + * - /health - Health check endpoint + * - /:token - Invalid token catch-all, redirect to website + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { StaticRoutes } from '../../../../main/web-server/routes/staticRoutes'; + +// Mock the logger +vi.mock('../../../../main/utils/logger', () => ({ + logger: { + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, +})); + +/** + * Mock Fastify instance with route registration tracking + */ +function createMockFastify() { + const routes: Map = new Map(); + + return { + get: vi.fn((path: string, handler: Function) => { + routes.set(`GET:${path}`, { handler }); + }), + getRoute: (method: string, path: string) => routes.get(`${method}:${path}`), + routes, + }; +} + +/** + * Mock reply object + */ +function createMockReply() { + const reply: any = { + code: vi.fn().mockReturnThis(), + send: vi.fn().mockReturnThis(), + type: vi.fn().mockReturnThis(), + redirect: vi.fn().mockReturnThis(), + }; + return reply; +} + +describe('StaticRoutes', () => { + const securityToken = 'test-token-123'; + const webAssetsPath = '/path/to/web/assets'; + + let staticRoutes: StaticRoutes; + let mockFastify: ReturnType; + + beforeEach(() => { + vi.clearAllMocks(); + staticRoutes = new StaticRoutes(securityToken, webAssetsPath); + mockFastify = createMockFastify(); + staticRoutes.registerRoutes(mockFastify as any); + }); + + describe('Route Registration', () => { + it('should register all static routes', () => { + // 8 routes: /, /health, manifest.json, sw.js, dashboard, dashboard/, session/:id, /:token + expect(mockFastify.get).toHaveBeenCalledTimes(8); + }); + + it('should register routes with correct paths', () => { + expect(mockFastify.routes.has('GET:/')).toBe(true); + expect(mockFastify.routes.has('GET:/health')).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/manifest.json`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/sw.js`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/`)).toBe(true); + expect(mockFastify.routes.has(`GET:/${securityToken}/session/:sessionId`)).toBe(true); + expect(mockFastify.routes.has('GET:/:token')).toBe(true); + }); + }); + + describe('GET / (Root Redirect)', () => { + it('should redirect to website', async () => { + const route = mockFastify.getRoute('GET', '/'); + const reply = createMockReply(); + await route!.handler({}, reply); + + expect(reply.redirect).toHaveBeenCalledWith(302, 'https://runmaestro.ai'); + }); + }); + + describe('GET /health', () => { + it('should return health status', async () => { + const route = mockFastify.getRoute('GET', '/health'); + const result = await route!.handler(); + + expect(result.status).toBe('ok'); + expect(result.timestamp).toBeDefined(); + }); + }); + + describe('Null webAssetsPath handling', () => { + it('should return 404 for manifest.json when webAssetsPath is null', async () => { + const noAssetsRoutes = new StaticRoutes(securityToken, null); + const noAssetsFastify = createMockFastify(); + noAssetsRoutes.registerRoutes(noAssetsFastify as any); + + const route = noAssetsFastify.getRoute('GET', `/${securityToken}/manifest.json`); + const reply = createMockReply(); + await route!.handler({}, reply); + + expect(reply.code).toHaveBeenCalledWith(404); + }); + + it('should return 404 for sw.js when webAssetsPath is null', async () => { + const noAssetsRoutes = new StaticRoutes(securityToken, null); + const noAssetsFastify = createMockFastify(); + noAssetsRoutes.registerRoutes(noAssetsFastify as any); + + const route = noAssetsFastify.getRoute('GET', `/${securityToken}/sw.js`); + const reply = createMockReply(); + await route!.handler({}, reply); + + expect(reply.code).toHaveBeenCalledWith(404); + }); + + it('should return 503 for dashboard when webAssetsPath is null', async () => { + const noAssetsRoutes = new StaticRoutes(securityToken, null); + const noAssetsFastify = createMockFastify(); + noAssetsRoutes.registerRoutes(noAssetsFastify as any); + + const route = noAssetsFastify.getRoute('GET', `/${securityToken}`); + const reply = createMockReply(); + await route!.handler({}, reply); + + expect(reply.code).toHaveBeenCalledWith(503); + expect(reply.send).toHaveBeenCalledWith( + expect.objectContaining({ error: 'Service Unavailable' }) + ); + }); + }); + + describe('GET /:token (Invalid Token Catch-all)', () => { + it('should redirect to website for invalid token', async () => { + const route = mockFastify.getRoute('GET', '/:token'); + const reply = createMockReply(); + await route!.handler({ params: { token: 'invalid-token' } }, reply); + + expect(reply.redirect).toHaveBeenCalledWith(302, 'https://runmaestro.ai'); + }); + }); + + describe('Security Token Validation', () => { + it('should use provided security token in routes', () => { + const customToken = 'custom-secure-token-456'; + const customRoutes = new StaticRoutes(customToken, webAssetsPath); + const customFastify = createMockFastify(); + customRoutes.registerRoutes(customFastify as any); + + expect(customFastify.routes.has(`GET:/${customToken}`)).toBe(true); + expect(customFastify.routes.has(`GET:/${customToken}/manifest.json`)).toBe(true); + expect(customFastify.routes.has(`GET:/${customToken}/sw.js`)).toBe(true); + expect(customFastify.routes.has(`GET:/${customToken}/session/:sessionId`)).toBe(true); + }); + }); + + describe('XSS Sanitization (sanitizeId)', () => { + // Access private method via type casting for testing + const getSanitizeId = (routes: StaticRoutes) => { + return (routes as any).sanitizeId.bind(routes); + }; + + it('should allow valid UUID-style IDs', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('abc123')).toBe('abc123'); + expect(sanitizeId('session-1')).toBe('session-1'); + expect(sanitizeId('tab_abc_123')).toBe('tab_abc_123'); + expect(sanitizeId('a1b2c3d4-e5f6-7890-abcd-ef1234567890')).toBe( + 'a1b2c3d4-e5f6-7890-abcd-ef1234567890' + ); + }); + + it('should return null for null/undefined input', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId(null)).toBeNull(); + expect(sanitizeId(undefined)).toBeNull(); + expect(sanitizeId('')).toBeNull(); + }); + + it('should reject XSS payloads with script tags', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('')).toBeNull(); + expect(sanitizeId('session')).toBeNull(); + }); + + it('should reject XSS payloads with JavaScript URLs', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('javascript:alert(1)')).toBeNull(); + expect(sanitizeId('session:javascript')).toBeNull(); + }); + + it('should reject XSS payloads with HTML entities', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('<script>')).toBeNull(); + expect(sanitizeId('session<')).toBeNull(); + }); + + it('should reject special characters that could break HTML/JS', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('"onload="alert(1)')).toBeNull(); + expect(sanitizeId("'onclick='alert(1)")).toBeNull(); + expect(sanitizeId('session;alert(1)')).toBeNull(); + expect(sanitizeId('session&alert=1')).toBeNull(); + expect(sanitizeId('session?alert=1')).toBeNull(); + expect(sanitizeId('session#alert')).toBeNull(); + }); + + it('should reject whitespace', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('session 1')).toBeNull(); + expect(sanitizeId('tab\t1')).toBeNull(); + expect(sanitizeId('tab\n1')).toBeNull(); + }); + + it('should reject path traversal attempts', () => { + const sanitizeId = getSanitizeId(staticRoutes); + expect(sanitizeId('../../../etc/passwd')).toBeNull(); + expect(sanitizeId('..%2F..%2F')).toBeNull(); + }); + }); +}); diff --git a/src/__tests__/main/web-server/routes/wsRoute.test.ts b/src/__tests__/main/web-server/routes/wsRoute.test.ts new file mode 100644 index 00000000..b17eb34b --- /dev/null +++ b/src/__tests__/main/web-server/routes/wsRoute.test.ts @@ -0,0 +1,497 @@ +/** + * Tests for WsRoute + * + * WebSocket Route handles WebSocket connections, initial state sync, and message delegation. + * Route: /$TOKEN/ws + * + * Connection Flow: + * 1. Client connects with optional ?sessionId= query param + * 2. Server sends 'connected' message with client ID + * 3. Server sends 'sessions_list' with all sessions (enriched with live info) + * 4. Server sends 'theme' with current theme + * 5. Server sends 'custom_commands' with available commands + * 6. Server sends 'autorun_state' for active AutoRun sessions + * 7. Client can send messages which are delegated to WebSocketMessageHandler + */ + +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { WebSocket } from 'ws'; +import { WsRoute, type WsRouteCallbacks } from '../../../../main/web-server/routes/wsRoute'; + +// Mock the logger +vi.mock('../../../../main/utils/logger', () => ({ + logger: { + info: vi.fn(), + debug: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + }, +})); + +/** + * Create mock callbacks with all methods as vi.fn() + */ +function createMockCallbacks(): WsRouteCallbacks { + return { + getSessions: vi.fn().mockReturnValue([ + { + id: 'session-1', + name: 'Session 1', + toolType: 'claude-code', + state: 'idle', + inputMode: 'ai', + cwd: '/test/project', + groupId: null, + }, + { + id: 'session-2', + name: 'Session 2', + toolType: 'codex', + state: 'busy', + inputMode: 'terminal', + cwd: '/test/project2', + groupId: 'group-1', + }, + ]), + getTheme: vi.fn().mockReturnValue({ + name: 'dark', + background: '#1a1a1a', + foreground: '#ffffff', + }), + getCustomCommands: vi + .fn() + .mockReturnValue([{ id: 'cmd-1', name: 'Test Command', prompt: 'Do something' }]), + getAutoRunStates: vi.fn().mockReturnValue( + new Map([ + [ + 'session-1', + { + isRunning: true, + totalTasks: 5, + completedTasks: 2, + currentTask: 'Task 3', + }, + ], + ]) + ), + getLiveSessionInfo: vi.fn().mockReturnValue({ + sessionId: 'session-1', + agentSessionId: 'claude-agent-123', + enabledAt: Date.now(), + }), + isSessionLive: vi.fn().mockReturnValue(true), + onClientConnect: vi.fn(), + onClientDisconnect: vi.fn(), + onClientError: vi.fn(), + handleMessage: vi.fn(), + }; +} + +/** + * Create mock WebSocket + */ +function createMockSocket() { + const eventHandlers: Map = new Map(); + return { + readyState: WebSocket.OPEN, + send: vi.fn(), + on: vi.fn((event: string, handler: Function) => { + if (!eventHandlers.has(event)) { + eventHandlers.set(event, []); + } + eventHandlers.get(event)!.push(handler); + }), + emit: (event: string, ...args: any[]) => { + const handlers = eventHandlers.get(event) || []; + handlers.forEach((h) => h(...args)); + }, + eventHandlers, + }; +} + +/** + * Create mock Fastify connection + */ +function createMockConnection() { + return { + socket: createMockSocket(), + }; +} + +/** + * Create mock Fastify request + */ +function createMockRequest(sessionId?: string) { + const queryString = sessionId ? `?sessionId=${sessionId}` : ''; + return { + url: `/test-token/ws${queryString}`, + headers: { + host: 'localhost:3000', + }, + }; +} + +/** + * Mock Fastify instance with route registration tracking + */ +function createMockFastify() { + const routes: Map = new Map(); + + return { + get: vi.fn((path: string, options: any, handler?: Function) => { + const h = handler || options; + const opts = handler ? options : undefined; + routes.set(`GET:${path}`, { handler: h, options: opts }); + }), + getRoute: (method: string, path: string) => routes.get(`${method}:${path}`), + routes, + }; +} + +describe('WsRoute', () => { + const securityToken = 'test-token-123'; + + let wsRoute: WsRoute; + let callbacks: WsRouteCallbacks; + let mockFastify: ReturnType; + + beforeEach(() => { + wsRoute = new WsRoute(securityToken); + callbacks = createMockCallbacks(); + wsRoute.setCallbacks(callbacks); + mockFastify = createMockFastify(); + wsRoute.registerRoute(mockFastify as any); + }); + + describe('Route Registration', () => { + it('should register WebSocket route with correct path', () => { + expect(mockFastify.get).toHaveBeenCalledTimes(1); + expect(mockFastify.routes.has(`GET:/${securityToken}/ws`)).toBe(true); + }); + + it('should register route with websocket option', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + expect(route?.options?.websocket).toBe(true); + }); + }); + + describe('Connection Handling', () => { + it('should generate unique client IDs', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + + // Connect first client + const conn1 = createMockConnection(); + route!.handler(conn1, createMockRequest()); + + // Connect second client + const conn2 = createMockConnection(); + route!.handler(conn2, createMockRequest()); + + // Verify unique IDs + expect(callbacks.onClientConnect).toHaveBeenCalledTimes(2); + const client1 = (callbacks.onClientConnect as any).mock.calls[0][0]; + const client2 = (callbacks.onClientConnect as any).mock.calls[1][0]; + expect(client1.id).not.toBe(client2.id); + expect(client1.id).toMatch(/^web-client-\d+$/); + expect(client2.id).toMatch(/^web-client-\d+$/); + }); + + it('should notify parent on client connect', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + expect(callbacks.onClientConnect).toHaveBeenCalledWith( + expect.objectContaining({ + id: expect.stringMatching(/^web-client-/), + socket: connection.socket, + connectedAt: expect.any(Number), + }) + ); + }); + + it('should extract sessionId from query string', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest('session-123')); + + expect(callbacks.onClientConnect).toHaveBeenCalledWith( + expect.objectContaining({ + subscribedSessionId: 'session-123', + }) + ); + }); + + it('should set subscribedSessionId to undefined when not in query', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + expect(callbacks.onClientConnect).toHaveBeenCalledWith( + expect.objectContaining({ + subscribedSessionId: undefined, + }) + ); + }); + }); + + describe('Initial Sync Messages', () => { + it('should send connected message', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest('session-123')); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const connectedMsg = sentMessages.find((m: any) => m.type === 'connected'); + expect(connectedMsg).toBeDefined(); + expect(connectedMsg.clientId).toMatch(/^web-client-/); + expect(connectedMsg.subscribedSessionId).toBe('session-123'); + expect(connectedMsg.timestamp).toBeDefined(); + }); + + it('should send sessions_list with enriched live info', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const sessionsMsg = sentMessages.find((m: any) => m.type === 'sessions_list'); + expect(sessionsMsg).toBeDefined(); + expect(sessionsMsg.sessions).toHaveLength(2); + expect(sessionsMsg.sessions[0].agentSessionId).toBe('claude-agent-123'); + expect(sessionsMsg.sessions[0].isLive).toBe(true); + expect(sessionsMsg.sessions[0].liveEnabledAt).toBeDefined(); + }); + + it('should send theme', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const themeMsg = sentMessages.find((m: any) => m.type === 'theme'); + expect(themeMsg).toBeDefined(); + expect(themeMsg.theme.name).toBe('dark'); + }); + + it('should not send theme when null', () => { + (callbacks.getTheme as any).mockReturnValue(null); + + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const themeMsg = sentMessages.find((m: any) => m.type === 'theme'); + expect(themeMsg).toBeUndefined(); + }); + + it('should send custom_commands', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const commandsMsg = sentMessages.find((m: any) => m.type === 'custom_commands'); + expect(commandsMsg).toBeDefined(); + expect(commandsMsg.commands).toHaveLength(1); + expect(commandsMsg.commands[0].name).toBe('Test Command'); + }); + + it('should send autorun_state for running sessions', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const autoRunMsg = sentMessages.find((m: any) => m.type === 'autorun_state'); + expect(autoRunMsg).toBeDefined(); + expect(autoRunMsg.sessionId).toBe('session-1'); + expect(autoRunMsg.state.isRunning).toBe(true); + expect(autoRunMsg.state.completedTasks).toBe(2); + expect(autoRunMsg.state.totalTasks).toBe(5); + }); + + it('should not send autorun_state for non-running sessions', () => { + (callbacks.getAutoRunStates as any).mockReturnValue( + new Map([ + [ + 'session-1', + { + isRunning: false, + totalTasks: 5, + completedTasks: 5, + }, + ], + ]) + ); + + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const autoRunMsg = sentMessages.find((m: any) => m.type === 'autorun_state'); + expect(autoRunMsg).toBeUndefined(); + }); + }); + + describe('Message Handling', () => { + it('should delegate messages to handleMessage callback', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + // Simulate incoming message + const message = JSON.stringify({ type: 'ping' }); + connection.socket.emit('message', message); + + expect(callbacks.handleMessage).toHaveBeenCalledWith(expect.stringMatching(/^web-client-/), { + type: 'ping', + }); + }); + + it('should send error for invalid JSON messages', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + // Clear previous sends + (connection.socket.send as any).mockClear(); + + // Simulate invalid message + connection.socket.emit('message', 'not valid json'); + + const lastSend = (connection.socket.send as any).mock.calls[0]; + const errorMsg = JSON.parse(lastSend[0]); + expect(errorMsg.type).toBe('error'); + expect(errorMsg.message).toBe('Invalid message format'); + }); + }); + + describe('Disconnection Handling', () => { + it('should notify parent on client disconnect', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const clientId = (callbacks.onClientConnect as any).mock.calls[0][0].id; + + // Simulate close event + connection.socket.emit('close'); + + expect(callbacks.onClientDisconnect).toHaveBeenCalledWith(clientId); + }); + }); + + describe('Error Handling', () => { + it('should notify parent on client error', () => { + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const clientId = (callbacks.onClientConnect as any).mock.calls[0][0].id; + const error = new Error('Connection lost'); + + // Simulate error event + connection.socket.emit('error', error); + + expect(callbacks.onClientError).toHaveBeenCalledWith(clientId, error); + }); + }); + + describe('Callback Resilience', () => { + it('should handle missing callbacks gracefully', () => { + const emptyWsRoute = new WsRoute(securityToken); + // Don't set any callbacks + const emptyFastify = createMockFastify(); + emptyWsRoute.registerRoute(emptyFastify as any); + + const route = emptyFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + + // Should not throw + expect(() => { + route!.handler(connection, createMockRequest()); + }).not.toThrow(); + + // Should still send connected message + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + const connectedMsg = sentMessages.find((m: any) => m.type === 'connected'); + expect(connectedMsg).toBeDefined(); + }); + + it('should handle partial callbacks', () => { + const partialWsRoute = new WsRoute(securityToken); + partialWsRoute.setCallbacks({ + getSessions: vi.fn().mockReturnValue([]), + getTheme: vi.fn().mockReturnValue(null), + getCustomCommands: vi.fn().mockReturnValue([]), + getAutoRunStates: vi.fn().mockReturnValue(new Map()), + getLiveSessionInfo: vi.fn().mockReturnValue(undefined), + isSessionLive: vi.fn().mockReturnValue(false), + onClientConnect: vi.fn(), + onClientDisconnect: vi.fn(), + onClientError: vi.fn(), + handleMessage: vi.fn(), + }); + const partialFastify = createMockFastify(); + partialWsRoute.registerRoute(partialFastify as any); + + const route = partialFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + + // Should not throw + expect(() => { + route!.handler(connection, createMockRequest()); + }).not.toThrow(); + }); + }); + + describe('Multiple AutoRun States', () => { + it('should send autorun_state for all running sessions', () => { + (callbacks.getAutoRunStates as any).mockReturnValue( + new Map([ + ['session-1', { isRunning: true, totalTasks: 5, completedTasks: 2 }], + ['session-2', { isRunning: true, totalTasks: 3, completedTasks: 1 }], + ['session-3', { isRunning: false, totalTasks: 2, completedTasks: 2 }], + ]) + ); + + const route = mockFastify.getRoute('GET', `/${securityToken}/ws`); + const connection = createMockConnection(); + route!.handler(connection, createMockRequest()); + + const sentMessages = (connection.socket.send as any).mock.calls.map((call: any[]) => + JSON.parse(call[0]) + ); + + const autoRunMsgs = sentMessages.filter((m: any) => m.type === 'autorun_state'); + expect(autoRunMsgs).toHaveLength(2); // Only running sessions + expect(autoRunMsgs.map((m: any) => m.sessionId)).toEqual(['session-1', 'session-2']); + }); + }); +}); diff --git a/src/__tests__/main/web-server/web-server-factory.test.ts b/src/__tests__/main/web-server/web-server-factory.test.ts index 619f56d6..2bb09a76 100644 --- a/src/__tests__/main/web-server/web-server-factory.test.ts +++ b/src/__tests__/main/web-server/web-server-factory.test.ts @@ -14,7 +14,8 @@ vi.mock('electron', () => ({ })); // Mock WebServer - use class syntax to make it a proper constructor -vi.mock('../../../main/web-server', () => { +// Note: Mock the specific file path that web-server-factory.ts imports from +vi.mock('../../../main/web-server/WebServer', () => { return { WebServer: class MockWebServer { port: number; @@ -68,7 +69,7 @@ import { createWebServerFactory, type WebServerFactoryDependencies, } from '../../../main/web-server/web-server-factory'; -import { WebServer } from '../../../main/web-server'; +import { WebServer } from '../../../main/web-server/WebServer'; import { getThemeById } from '../../../main/themes'; import { getHistoryManager } from '../../../main/history-manager'; import { logger } from '../../../main/utils/logger'; diff --git a/src/main/web-server.ts b/src/main/web-server/WebServer.ts similarity index 79% rename from src/main/web-server.ts rename to src/main/web-server/WebServer.ts index c03b26eb..a067b7c4 100644 --- a/src/main/web-server.ts +++ b/src/main/web-server/WebServer.ts @@ -1,49 +1,3 @@ -import Fastify from 'fastify'; -import cors from '@fastify/cors'; -import websocket from '@fastify/websocket'; -import rateLimit from '@fastify/rate-limit'; -import fastifyStatic from '@fastify/static'; -import { FastifyInstance, FastifyRequest } from 'fastify'; -import { randomUUID } from 'crypto'; -import path from 'path'; -import { existsSync } from 'fs'; -import type { Theme } from '../shared/theme-types'; -import { HistoryEntry } from '../shared/types'; -import { getLocalIpAddressSync } from './utils/networkUtils'; -import { logger } from './utils/logger'; -import { WebSocketMessageHandler, WebClient, WebClientMessage } from './web-server/handlers'; -import { - BroadcastService, - AITabData as BroadcastAITabData, - CustomAICommand as BroadcastCustomAICommand, - AutoRunState, - CliActivity, - SessionBroadcastData, -} from './web-server/services'; -import { ApiRoutes, StaticRoutes, WsRoute } from './web-server/routes'; - -// Logger context for all web server logs -const LOG_CONTEXT = 'WebServer'; - -// Live session info -interface LiveSessionInfo { - sessionId: string; - agentSessionId?: string; - enabledAt: number; -} - -// Rate limiting configuration -export interface RateLimitConfig { - // Maximum requests per time window - max: number; - // Time window in milliseconds - timeWindow: number; - // Maximum requests for POST endpoints (typically lower) - maxPost: number; - // Enable/disable rate limiting - enabled: boolean; -} - /** * WebServer - HTTP and WebSocket server for remote access * @@ -62,141 +16,55 @@ export interface RateLimitConfig { * * Security: * - Token regenerated on each app restart - * - Invalid/missing token redirects to GitHub + * - Invalid/missing token redirects to website * - No access without knowing the token */ -// Usage stats type for session cost/token tracking -export interface SessionUsageStats { - inputTokens?: number; - outputTokens?: number; - cacheReadInputTokens?: number; - cacheCreationInputTokens?: number; - totalCostUsd?: number; - contextWindow?: number; -} -// Last response type for mobile preview (truncated to save bandwidth) -export interface LastResponsePreview { - text: string; // First 3 lines or ~500 chars of the last AI response - timestamp: number; - source: 'stdout' | 'stderr' | 'system'; - fullLength: number; // Total length of the original response -} +import Fastify from 'fastify'; +import cors from '@fastify/cors'; +import websocket from '@fastify/websocket'; +import rateLimit from '@fastify/rate-limit'; +import fastifyStatic from '@fastify/static'; +import { FastifyInstance, FastifyRequest } from 'fastify'; +import { randomUUID } from 'crypto'; +import path from 'path'; +import { existsSync } from 'fs'; +import { getLocalIpAddressSync } from '../utils/networkUtils'; +import { logger } from '../utils/logger'; +import { WebSocketMessageHandler } from './handlers'; +import { BroadcastService } from './services'; +import { ApiRoutes, StaticRoutes, WsRoute } from './routes'; -// AI Tab type for multi-tab support within a Maestro session -export interface AITabData { - id: string; - agentSessionId: string | null; - name: string | null; - starred: boolean; - inputValue: string; - usageStats?: SessionUsageStats | null; - createdAt: number; - state: 'idle' | 'busy'; - thinkingStartTime?: number | null; -} +// Import shared types from canonical location +import type { + Theme, + LiveSessionInfo, + RateLimitConfig, + AITabData, + CustomAICommand, + AutoRunState, + CliActivity, + SessionBroadcastData, + WebClient, + WebClientMessage, + GetSessionsCallback, + GetSessionDetailCallback, + WriteToSessionCallback, + ExecuteCommandCallback, + InterruptSessionCallback, + SwitchModeCallback, + SelectSessionCallback, + SelectTabCallback, + NewTabCallback, + CloseTabCallback, + RenameTabCallback, + GetThemeCallback, + GetCustomCommandsCallback, + GetHistoryCallback, +} from './types'; -// Callback type for fetching sessions data -export type GetSessionsCallback = () => Array<{ - id: string; - name: string; - toolType: string; - state: string; - inputMode: string; - cwd: string; - groupId: string | null; - groupName: string | null; - groupEmoji: string | null; - usageStats?: SessionUsageStats | null; - lastResponse?: LastResponsePreview | null; - agentSessionId?: string | null; - thinkingStartTime?: number | null; // Timestamp when AI started thinking (for elapsed time display) - aiTabs?: AITabData[]; - activeTabId?: string; - bookmarked?: boolean; // Whether session is bookmarked (shows in Bookmarks group) -}>; - -// Session detail type for single session endpoint -export interface SessionDetail { - id: string; - name: string; - toolType: string; - state: string; - inputMode: string; - cwd: string; - aiLogs?: Array<{ timestamp: number; content: string; type?: string }>; - shellLogs?: Array<{ timestamp: number; content: string; type?: string }>; - usageStats?: { - inputTokens?: number; - outputTokens?: number; - totalCost?: number; - }; - agentSessionId?: string; - isGitRepo?: boolean; - activeTabId?: string; -} - -// Callback type for fetching single session details -// Optional tabId allows fetching logs for a specific tab (avoids race conditions) -export type GetSessionDetailCallback = (sessionId: string, tabId?: string) => SessionDetail | null; - -// Callback type for sending commands to a session -// Returns true if successful, false if session not found or write failed -export type WriteToSessionCallback = (sessionId: string, data: string) => boolean; - -// Callback type for executing a command through the desktop's existing logic -// This forwards the command to the renderer which handles spawn, state, and broadcasts -// Returns true if command was accepted (session not busy) -// inputMode is optional - if provided, the renderer will use it instead of querying session state -export type ExecuteCommandCallback = ( - sessionId: string, - command: string, - inputMode?: 'ai' | 'terminal' -) => Promise; - -// Callback type for interrupting a session through the desktop's existing logic -// This forwards to the renderer which handles state updates and broadcasts -export type InterruptSessionCallback = (sessionId: string) => Promise; - -// Callback type for switching session input mode through the desktop's existing logic -// This forwards to the renderer which handles state updates and broadcasts -export type SwitchModeCallback = (sessionId: string, mode: 'ai' | 'terminal') => Promise; - -// Callback type for selecting/switching to a session in the desktop app -// This forwards to the renderer which handles state updates and broadcasts -// Optional tabId to also switch to a specific tab within the session -export type SelectSessionCallback = (sessionId: string, tabId?: string) => Promise; - -// Tab operation callbacks for multi-tab support -export type SelectTabCallback = (sessionId: string, tabId: string) => Promise; -export type NewTabCallback = (sessionId: string) => Promise<{ tabId: string } | null>; -export type CloseTabCallback = (sessionId: string, tabId: string) => Promise; -export type RenameTabCallback = ( - sessionId: string, - tabId: string, - newName: string -) => Promise; - -// Re-export Theme type from shared for backwards compatibility -export type { Theme } from '../shared/theme-types'; - -// Callback type for fetching current theme -export type GetThemeCallback = () => Theme | null; - -// Custom AI command definition (matches renderer's CustomAICommand) -export interface CustomAICommand { - id: string; - command: string; - description: string; - prompt: string; -} - -// Callback type for fetching custom AI commands -export type GetCustomCommandsCallback = () => CustomAICommand[]; - -// Callback type for fetching history entries -// Uses HistoryEntry from shared/types.ts as the canonical type -export type GetHistoryCallback = (projectPath?: string, sessionId?: string) => HistoryEntry[]; +// Logger context for all web server logs +const LOG_CONTEXT = 'WebServer'; // Default rate limit configuration const DEFAULT_RATE_LIMIT_CONFIG: RateLimitConfig = { @@ -291,11 +159,11 @@ export class WebServer { // Try multiple locations for the web assets const possiblePaths = [ // Production: relative to the compiled main process - path.join(__dirname, '..', 'web'), + path.join(__dirname, '..', '..', 'web'), // Development: from project root path.join(process.cwd(), 'dist', 'web'), // Alternative: relative to __dirname going up to dist - path.join(__dirname, 'web'), + path.join(__dirname, '..', 'web'), ]; for (const p of possiblePaths) { @@ -343,6 +211,12 @@ export class WebServer { LOG_CONTEXT ); + // Clean up any associated AutoRun state to prevent memory leaks + if (this.autoRunStates.has(sessionId)) { + this.autoRunStates.delete(sessionId); + logger.debug(`Cleaned up AutoRun state for offline session ${sessionId}`, LOG_CONTEXT); + } + // Broadcast to all connected clients this.broadcastService.broadcastSessionOffline(sessionId); } @@ -705,7 +579,7 @@ export class WebServer { * Broadcast tab change to all connected web clients * Called when the tabs array or active tab changes in a session */ - broadcastTabsChange(sessionId: string, aiTabs: BroadcastAITabData[], activeTabId: string): void { + broadcastTabsChange(sessionId: string, aiTabs: AITabData[], activeTabId: string): void { this.broadcastService.broadcastTabsChange(sessionId, aiTabs, activeTabId); } @@ -721,7 +595,7 @@ export class WebServer { * Broadcast custom commands update to all connected web clients * Called when the user modifies custom AI commands in the desktop app */ - broadcastCustomCommands(commands: BroadcastCustomAICommand[]): void { + broadcastCustomCommands(commands: CustomAICommand[]): void { this.broadcastService.broadcastCustomCommands(commands); } @@ -862,11 +736,20 @@ export class WebServer { return; } - // Mark all live sessions as offline + // Mark all live sessions as offline (this also cleans up autoRunStates) for (const sessionId of this.liveSessions.keys()) { this.setSessionOffline(sessionId); } + // Clear any remaining autoRunStates as a safety measure + if (this.autoRunStates.size > 0) { + logger.debug( + `Clearing ${this.autoRunStates.size} remaining AutoRun states on server stop`, + LOG_CONTEXT + ); + this.autoRunStates.clear(); + } + try { await this.server.close(); this.isRunning = false; diff --git a/src/main/web-server/handlers/messageHandlers.ts b/src/main/web-server/handlers/messageHandlers.ts index 7d198bbc..e227a8d3 100644 --- a/src/main/web-server/handlers/messageHandlers.ts +++ b/src/main/web-server/handlers/messageHandlers.ts @@ -90,7 +90,6 @@ export interface MessageHandlerCallbacks { inputMode: string; cwd: string; agentSessionId?: string | null; - [key: string]: unknown; }>; getLiveSessionInfo: (sessionId: string) => LiveSessionInfo | undefined; isSessionLive: (sessionId: string) => boolean; @@ -112,6 +111,20 @@ export class WebSocketMessageHandler { this.callbacks = { ...this.callbacks, ...callbacks }; } + /** + * Helper to send a JSON message to a client with timestamp + */ + private send(client: WebClient, data: Record): void { + client.socket.send(JSON.stringify({ ...data, timestamp: Date.now() })); + } + + /** + * Helper to send an error message to a client + */ + private sendError(client: WebClient, message: string, extra?: Record): void { + this.send(client, { type: 'error', message, ...extra }); + } + /** * Handle incoming WebSocket message from a web client * @@ -175,12 +188,7 @@ export class WebSocketMessageHandler { * Handle ping message - respond with pong */ private handlePing(client: WebClient): void { - client.socket.send( - JSON.stringify({ - type: 'pong', - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'pong' }); } /** @@ -190,13 +198,7 @@ export class WebSocketMessageHandler { if (message.sessionId) { client.subscribedSessionId = message.sessionId as string; } - client.socket.send( - JSON.stringify({ - type: 'subscribed', - sessionId: message.sessionId, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'subscribed', sessionId: message.sessionId }); } /** @@ -218,38 +220,25 @@ export class WebSocketMessageHandler { `[Web Command] Missing sessionId or command: sessionId=${sessionId}, commandLen=${command?.length}`, LOG_CONTEXT ); - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId or command', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId or command'); return; } // Get session details to check state and determine how to handle const sessionDetail = this.callbacks.getSessionDetail?.(sessionId); if (!sessionDetail) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Session not found', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Session not found'); return; } // Check if session is busy - prevent race conditions between desktop and web if (sessionDetail.state === 'busy') { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Session is busy - please wait for the current operation to complete', + this.sendError( + client, + 'Session is busy - please wait for the current operation to complete', + { sessionId, - timestamp: Date.now(), - }) + } ); logger.debug(`Command rejected - session ${sessionId} is busy`, LOG_CONTEXT); return; @@ -274,14 +263,7 @@ export class WebSocketMessageHandler { this.callbacks .executeCommand(sessionId, command, clientInputMode) .then((success) => { - client.socket.send( - JSON.stringify({ - type: 'command_result', - success, - sessionId, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'command_result', success, sessionId }); if (!success) { logger.warn( `[Web Command] ${mode} command rejected for session ${sessionId}`, @@ -294,22 +276,10 @@ export class WebSocketMessageHandler { `[Web Command] ${mode} command failed for session ${sessionId}: ${error.message}`, LOG_CONTEXT ); - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to execute command: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to execute command: ${error.message}`); }); } else { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Command execution not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Command execution not configured'); } } @@ -325,25 +295,13 @@ export class WebSocketMessageHandler { ); if (!sessionId || !mode) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId or mode', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId or mode'); return; } if (!this.callbacks.switchMode) { logger.warn(`[Web] switchModeCallback is not set!`, LOG_CONTEXT); - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Mode switching not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Mode switching not configured'); return; } @@ -353,28 +311,14 @@ export class WebSocketMessageHandler { this.callbacks .switchMode(sessionId, mode) .then((success) => { - client.socket.send( - JSON.stringify({ - type: 'mode_switch_result', - success, - sessionId, - mode, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'mode_switch_result', success, sessionId, mode }); logger.debug( `Mode switch for session ${sessionId} to ${mode}: ${success ? 'success' : 'failed'}`, LOG_CONTEXT ); }) .catch((error) => { - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to switch mode: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to switch mode: ${error.message}`); }); } @@ -390,25 +334,13 @@ export class WebSocketMessageHandler { ); if (!sessionId) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId'); return; } if (!this.callbacks.selectSession) { logger.warn(`[Web] selectSessionCallback is not set!`, LOG_CONTEXT); - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Session selection not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Session selection not configured'); return; } @@ -420,14 +352,7 @@ export class WebSocketMessageHandler { this.callbacks .selectSession(sessionId, tabId) .then((success) => { - client.socket.send( - JSON.stringify({ - type: 'select_session_result', - success, - sessionId, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'select_session_result', success, sessionId }); if (success) { logger.debug(`Session ${sessionId} selected in desktop`, LOG_CONTEXT); } else { @@ -435,13 +360,7 @@ export class WebSocketMessageHandler { } }) .catch((error) => { - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to select session: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to select session: ${error.message}`); }); } @@ -465,13 +384,7 @@ export class WebSocketMessageHandler { isLive: this.callbacks.isSessionLive!(s.id), }; }); - client.socket.send( - JSON.stringify({ - type: 'sessions_list', - sessions: sessionsWithLiveInfo, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'sessions_list', sessions: sessionsWithLiveInfo }); } } @@ -487,48 +400,22 @@ export class WebSocketMessageHandler { ); if (!sessionId || !tabId) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId or tabId', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId or tabId'); return; } if (!this.callbacks.selectTab) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Tab selection not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Tab selection not configured'); return; } this.callbacks .selectTab(sessionId, tabId) .then((success) => { - client.socket.send( - JSON.stringify({ - type: 'select_tab_result', - success, - sessionId, - tabId, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'select_tab_result', success, sessionId, tabId }); }) .catch((error) => { - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to select tab: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to select tab: ${error.message}`); }); } @@ -540,48 +427,27 @@ export class WebSocketMessageHandler { logger.info(`[Web] Received new_tab message: session=${sessionId}`, LOG_CONTEXT); if (!sessionId) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId'); return; } if (!this.callbacks.newTab) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Tab creation not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Tab creation not configured'); return; } this.callbacks .newTab(sessionId) .then((result) => { - client.socket.send( - JSON.stringify({ - type: 'new_tab_result', - success: !!result, - sessionId, - tabId: result?.tabId, - timestamp: Date.now(), - }) - ); + this.send(client, { + type: 'new_tab_result', + success: !!result, + sessionId, + tabId: result?.tabId, + }); }) .catch((error) => { - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to create tab: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to create tab: ${error.message}`); }); } @@ -597,48 +463,22 @@ export class WebSocketMessageHandler { ); if (!sessionId || !tabId) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId or tabId', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId or tabId'); return; } if (!this.callbacks.closeTab) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Tab closing not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Tab closing not configured'); return; } this.callbacks .closeTab(sessionId, tabId) .then((success) => { - client.socket.send( - JSON.stringify({ - type: 'close_tab_result', - success, - sessionId, - tabId, - timestamp: Date.now(), - }) - ); + this.send(client, { type: 'close_tab_result', success, sessionId, tabId }); }) .catch((error) => { - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to close tab: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to close tab: ${error.message}`); }); } @@ -655,24 +495,12 @@ export class WebSocketMessageHandler { ); if (!sessionId || !tabId) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Missing sessionId or tabId', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Missing sessionId or tabId'); return; } if (!this.callbacks.renameTab) { - client.socket.send( - JSON.stringify({ - type: 'error', - message: 'Tab renaming not configured', - timestamp: Date.now(), - }) - ); + this.sendError(client, 'Tab renaming not configured'); return; } @@ -680,25 +508,16 @@ export class WebSocketMessageHandler { this.callbacks .renameTab(sessionId, tabId, newName || '') .then((success) => { - client.socket.send( - JSON.stringify({ - type: 'rename_tab_result', - success, - sessionId, - tabId, - newName: newName || '', - timestamp: Date.now(), - }) - ); + this.send(client, { + type: 'rename_tab_result', + success, + sessionId, + tabId, + newName: newName || '', + }); }) .catch((error) => { - client.socket.send( - JSON.stringify({ - type: 'error', - message: `Failed to rename tab: ${error.message}`, - timestamp: Date.now(), - }) - ); + this.sendError(client, `Failed to rename tab: ${error.message}`); }); } @@ -707,12 +526,6 @@ export class WebSocketMessageHandler { */ private handleUnknown(client: WebClient, message: WebClientMessage): void { logger.debug(`Unknown message type: ${message.type}`, LOG_CONTEXT); - client.socket.send( - JSON.stringify({ - type: 'echo', - originalType: message.type, - data: message, - }) - ); + this.send(client, { type: 'echo', originalType: message.type, data: message }); } } diff --git a/src/main/web-server/index.ts b/src/main/web-server/index.ts new file mode 100644 index 00000000..b00d2c51 --- /dev/null +++ b/src/main/web-server/index.ts @@ -0,0 +1,62 @@ +/** + * Web Server Module Index + * + * Main entry point for the web server module. + * Re-exports all public types, classes, and utilities. + * + * Import from this module: + * import { WebServer } from './web-server'; + * import type { Theme, LiveSessionInfo } from './web-server'; + */ + +// ============ Main Export ============ +// Export the WebServer class +export { WebServer } from './WebServer'; + +// ============ Shared Types ============ +// Export all shared types (canonical location for all web server types) +export type { + // Core types + Theme, + LiveSessionInfo, + RateLimitConfig, + SessionUsageStats, + LastResponsePreview, + AITabData, + SessionDetail, + CustomAICommand, + AutoRunState, + CliActivity, + SessionBroadcastData, + SessionData, + WebClient, + WebClientMessage, + // Callback types + GetSessionsCallback, + GetSessionDetailCallback, + WriteToSessionCallback, + ExecuteCommandCallback, + InterruptSessionCallback, + SwitchModeCallback, + SelectSessionCallback, + SelectTabCallback, + NewTabCallback, + CloseTabCallback, + RenameTabCallback, + GetThemeCallback, + GetCustomCommandsCallback, + GetHistoryCallback, + GetWebClientsCallback, +} from './types'; + +// ============ Handlers ============ +export { WebSocketMessageHandler } from './handlers'; +export type { SessionDetailForHandler, MessageHandlerCallbacks } from './handlers'; + +// ============ Services ============ +export { BroadcastService } from './services'; +export type { WebClientInfo } from './services'; + +// ============ Routes ============ +export { ApiRoutes, StaticRoutes, WsRoute } from './routes'; +export type { ApiRouteCallbacks, WsRouteCallbacks, WsSessionData } from './routes'; diff --git a/src/main/web-server/routes/apiRoutes.ts b/src/main/web-server/routes/apiRoutes.ts index 5f03b5c3..ee720781 100644 --- a/src/main/web-server/routes/apiRoutes.ts +++ b/src/main/web-server/routes/apiRoutes.ts @@ -16,118 +16,23 @@ import { FastifyInstance } from 'fastify'; import { HistoryEntry } from '../../../shared/types'; import { logger } from '../../utils/logger'; +import type { Theme, SessionData, SessionDetail, LiveSessionInfo, RateLimitConfig } from '../types'; + +// Re-export types for backwards compatibility +export type { + Theme, + SessionUsageStats, + LastResponsePreview, + AITabData, + SessionData, + SessionDetail, + LiveSessionInfo, + RateLimitConfig, +} from '../types'; // Logger context for all API route logs const LOG_CONTEXT = 'WebServer:API'; -/** - * Usage stats type for session cost/token tracking - */ -export interface SessionUsageStats { - inputTokens?: number; - outputTokens?: number; - cacheReadInputTokens?: number; - cacheCreationInputTokens?: number; - totalCostUsd?: number; - contextWindow?: number; -} - -/** - * Last response type for mobile preview (truncated to save bandwidth) - */ -export interface LastResponsePreview { - text: string; // First 3 lines or ~500 chars of the last AI response - timestamp: number; - source: 'stdout' | 'stderr' | 'system'; - fullLength: number; // Total length of the original response -} - -/** - * AI Tab type for multi-tab support within a Maestro session - */ -export interface AITabData { - id: string; - agentSessionId: string | null; - name: string | null; - starred: boolean; - inputValue: string; - usageStats?: SessionUsageStats | null; - createdAt: number; - state: 'idle' | 'busy'; - thinkingStartTime?: number | null; -} - -/** - * Session data returned by getSessions callback - */ -export interface SessionData { - id: string; - name: string; - toolType: string; - state: string; - inputMode: string; - cwd: string; - groupId: string | null; - groupName: string | null; - groupEmoji: string | null; - usageStats?: SessionUsageStats | null; - lastResponse?: LastResponsePreview | null; - agentSessionId?: string | null; - thinkingStartTime?: number | null; - aiTabs?: AITabData[]; - activeTabId?: string; - bookmarked?: boolean; -} - -/** - * Session detail type for single session endpoint - */ -export interface SessionDetail { - id: string; - name: string; - toolType: string; - state: string; - inputMode: string; - cwd: string; - aiLogs?: Array<{ timestamp: number; content: string; type?: string }>; - shellLogs?: Array<{ timestamp: number; content: string; type?: string }>; - usageStats?: { - inputTokens?: number; - outputTokens?: number; - totalCost?: number; - }; - agentSessionId?: string; - isGitRepo?: boolean; - activeTabId?: string; -} - -// HistoryEntry is imported from shared/types.ts as the canonical type - -/** - * Live session info for enriching sessions - */ -export interface LiveSessionInfo { - sessionId: string; - agentSessionId?: string; - enabledAt: number; -} - -/** - * Rate limit configuration - */ -export interface RateLimitConfig { - max: number; - timeWindow: number; - maxPost: number; - enabled: boolean; -} - -/** - * Theme type (imported from shared, re-exported for convenience) - */ -export type { Theme } from '../../../shared/theme-types'; -import type { Theme } from '../../../shared/theme-types'; - /** * Callbacks required by API routes */ diff --git a/src/main/web-server/routes/staticRoutes.ts b/src/main/web-server/routes/staticRoutes.ts index b50957fb..135be5b7 100644 --- a/src/main/web-server/routes/staticRoutes.ts +++ b/src/main/web-server/routes/staticRoutes.ts @@ -47,6 +47,22 @@ export class StaticRoutes { return token === this.securityToken; } + /** + * Sanitize a string for safe injection into HTML/JavaScript + * Only allows alphanumeric characters, hyphens, and underscores (valid for UUIDs and IDs) + * Returns null if the input contains invalid characters + */ + private sanitizeId(input: string | undefined | null): string | null { + if (!input) return null; + // Only allow characters that are safe for UUID-style IDs + // This prevents XSS attacks via malicious sessionId/tabId parameters + if (!/^[a-zA-Z0-9_-]+$/.test(input)) { + logger.warn(`Rejected potentially unsafe ID: ${input.substring(0, 50)}`, LOG_CONTEXT); + return null; + } + return input; + } + /** * Serve the index.html file for SPA routes * Rewrites asset paths to include the security token @@ -79,12 +95,17 @@ export class StaticRoutes { html = html.replace(/\.\/icons\//g, `/${this.securityToken}/icons/`); html = html.replace(/\.\/sw\.js/g, `/${this.securityToken}/sw.js`); + // Sanitize sessionId and tabId to prevent XSS attacks + // Only allow safe characters (alphanumeric, hyphens, underscores) + const safeSessionId = this.sanitizeId(sessionId); + const safeTabId = this.sanitizeId(tabId); + // Inject config for the React app to know the token and session context const configScript = `