| |
| |
| |
| |
| |
| |
| |
|
|
| import type { ServerWebSocket } from 'bun' |
| import type { ClientMessage, ServerMessage } from './events.js' |
| import * as os from 'node:os' |
| import * as fs from 'node:fs' |
| import * as path from 'node:path' |
| import { |
| ConversationStartupError, |
| conversationService, |
| } from '../services/conversationService.js' |
| import { computerUseApprovalService } from '../services/computerUseApprovalService.js' |
| import { sessionService } from '../services/sessionService.js' |
| import { SettingsService } from '../services/settingsService.js' |
| import { ProviderService } from '../services/providerService.js' |
| import { isOpenAIOfficialProviderId } from '../services/openaiOfficialProvider.js' |
| import { diagnosticsService } from '../services/diagnosticsService.js' |
| import { deriveTitle, generateTitle, saveAiTitle } from '../services/titleService.js' |
| import { parseSlashCommand } from '../../utils/slashCommandParsing.js' |
| import { |
| COMMAND_NAME_TAG, |
| LOCAL_COMMAND_STDERR_TAG, |
| LOCAL_COMMAND_STDOUT_TAG, |
| } from '../../constants/xml.js' |
| import { shouldCreateWorktreeForSessionLaunch } from '../services/repositoryLaunchService.js' |
| const settingsService = new SettingsService() |
| const providerService = new ProviderService() |
|
|
| |
| |
| |
| export type SessionSlashCommand = { |
| name: string |
| description: string |
| argumentHint?: string |
| } |
|
|
| const sessionSlashCommands = new Map<string, SessionSlashCommand[]>() |
|
|
| |
| |
| |
| |
| const sessionCleanupTimers = new Map<string, ReturnType<typeof setTimeout>>() |
|
|
| |
| |
| |
| |
| const sessionStopRequested = new Set<string>() |
|
|
| |
| |
| |
| const sessionTitleState = new Map<string, { |
| userMessageCount: number |
| hasCustomTitle: boolean |
| firstUserMessage: string |
| allUserMessages: string[] |
| startedGenerationCounts: Set<number> |
| }>() |
|
|
| const runtimeOverrides = new Map<string, { |
| providerId: string | null |
| modelId: string |
| }>() |
|
|
| const runtimeTransitionPromises = new Map<string, Promise<void>>() |
| const sessionStartupPromises = new Map<string, Promise<void>>() |
| const lastResolvedStartupWorkDirs = new Map<string, string>() |
| const prewarmPendingSessions = new Set<string>() |
| const prewarmedSessions = new Set<string>() |
| const prewarmIdleTimers = new Map<string, ReturnType<typeof setTimeout>>() |
| const DEFAULT_PREWARM_IDLE_TIMEOUT_MS = 5 * 60_000 |
|
|
| async function sendRepositoryStartupStatus( |
| ws: ServerWebSocket<WebSocketData>, |
| sessionId: string, |
| reason: 'user_message' | 'prewarm_session', |
| ): Promise<void> { |
| if (reason !== 'user_message') return |
|
|
| const launchInfo = await sessionService.getSessionLaunchInfo(sessionId).catch(() => null) |
| const repository = launchInfo?.repository |
| if (!repository) return |
|
|
| if (shouldCreateWorktreeForSessionLaunch(launchInfo)) { |
| sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Creating worktree' }) |
| } |
| } |
|
|
| export function getSlashCommands(sessionId: string): SessionSlashCommand[] { |
| return sessionSlashCommands.get(sessionId) || [] |
| } |
|
|
| export type WebSocketData = { |
| sessionId: string |
| connectedAt: number |
| channel: 'client' | 'sdk' |
| sdkToken: string | null |
| serverPort: number |
| serverHost: string |
| } |
|
|
| |
| const activeSessions = new Map<string, ServerWebSocket<WebSocketData>>() |
|
|
| export const handleWebSocket = { |
| open(ws: ServerWebSocket<WebSocketData>) { |
| const { sessionId, channel, sdkToken } = ws.data |
|
|
| if (channel === 'sdk') { |
| if (!conversationService.authorizeSdkConnection(sessionId, sdkToken)) { |
| console.warn(`[WS] Rejected SDK connection for session: ${sessionId}`) |
| ws.close(1008, 'Invalid SDK token') |
| return |
| } |
|
|
| conversationService.attachSdkConnection(sessionId, ws) |
| console.log(`[WS] SDK connected for session: ${sessionId}`) |
| return |
| } |
|
|
| console.log(`[WS] Client connected for session: ${sessionId}`) |
|
|
| |
| const pendingTimer = sessionCleanupTimers.get(sessionId) |
| if (pendingTimer) { |
| clearTimeout(pendingTimer) |
| sessionCleanupTimers.delete(sessionId) |
| } |
|
|
| activeSessions.set(sessionId, ws) |
| if (prewarmedSessions.has(sessionId)) { |
| bindPrewarmMetadataCapture(sessionId) |
| } else { |
| rebindSessionOutput(sessionId, ws) |
| } |
|
|
| const msg: ServerMessage = { type: 'connected', sessionId } |
| ws.send(JSON.stringify(msg)) |
| }, |
|
|
| message(ws: ServerWebSocket<WebSocketData>, rawMessage: string | Buffer) { |
| if (ws.data.channel === 'sdk') { |
| const payload = typeof rawMessage === 'string' ? rawMessage : rawMessage.toString() |
| conversationService.handleSdkPayload(ws.data.sessionId, payload) |
| return |
| } |
|
|
| try { |
| const message = JSON.parse( |
| typeof rawMessage === 'string' ? rawMessage : rawMessage.toString() |
| ) as ClientMessage |
|
|
| switch (message.type) { |
| case 'user_message': |
| handleUserMessage(ws, message).catch((err) => { |
| void diagnosticsService.recordEvent({ |
| type: 'ws_user_message_failed', |
| severity: 'error', |
| sessionId: ws.data.sessionId, |
| summary: err instanceof Error ? err.message : String(err), |
| details: err, |
| }) |
| console.error(`[WS] Unhandled error in handleUserMessage:`, err) |
| }) |
| break |
|
|
| case 'permission_response': |
| handlePermissionResponse(ws, message) |
| break |
|
|
| case 'computer_use_permission_response': |
| handleComputerUsePermissionResponse(ws, message) |
| break |
|
|
| case 'set_permission_mode': |
| handleSetPermissionMode(ws, message) |
| break |
|
|
| case 'set_runtime_config': |
| void handleSetRuntimeConfig(ws, message) |
| break |
|
|
| case 'prewarm_session': |
| void handlePrewarmSession(ws) |
| break |
|
|
| case 'stop_generation': |
| handleStopGeneration(ws) |
| break |
|
|
| case 'ping': |
| ws.send(JSON.stringify({ type: 'pong' } satisfies ServerMessage)) |
| break |
|
|
| default: |
| sendError(ws, `Unknown message type: ${(message as any).type}`, 'UNKNOWN_TYPE') |
| } |
| } catch (error) { |
| sendError(ws, `Invalid message format: ${error}`, 'PARSE_ERROR') |
| } |
| }, |
|
|
| close(ws: ServerWebSocket<WebSocketData>, code: number, reason: string) { |
| const { sessionId, channel } = ws.data |
|
|
| if (channel === 'sdk') { |
| console.log(`[WS] SDK disconnected from session: ${sessionId} (${code}: ${reason})`) |
| conversationService.detachSdkConnection(sessionId) |
| return |
| } |
|
|
| console.log(`[WS] Client disconnected from session: ${sessionId} (${code}: ${reason})`) |
| if (activeSessions.get(sessionId) !== ws) { |
| console.log(`[WS] Ignoring stale client disconnect for session: ${sessionId}`) |
| return |
| } |
| computerUseApprovalService.cancelSession(sessionId) |
| activeSessions.delete(sessionId) |
| conversationService.clearOutputCallbacks(sessionId) |
|
|
| |
| |
| const cleanupTimer = setTimeout(() => { |
| sessionCleanupTimers.delete(sessionId) |
| if (!activeSessions.has(sessionId)) { |
| console.log(`[WS] Session ${sessionId} not reconnected after 30s, stopping CLI subprocess`) |
| conversationService.stopSession(sessionId) |
| cleanupSessionRuntimeState(sessionId) |
| } |
| }, 30_000) |
| sessionCleanupTimers.set(sessionId, cleanupTimer) |
| }, |
|
|
| drain(ws: ServerWebSocket<WebSocketData>) { |
| |
| }, |
| } |
|
|
| |
| |
| |
|
|
| async function handleUserMessage( |
| ws: ServerWebSocket<WebSocketData>, |
| message: Extract<ClientMessage, { type: 'user_message' }> |
| ) { |
| const { sessionId } = ws.data |
|
|
| |
| sessionStopRequested.delete(sessionId) |
| clearPrewarmState(sessionId) |
|
|
| const desktopSlashCommand = getDesktopSlashCommand(message.content) |
| if (desktopSlashCommand?.commandName === 'clear' && desktopSlashCommand.args.trim()) { |
| sendMessage(ws, { |
| type: 'error', |
| message: 'The /clear command does not accept arguments.', |
| code: 'INVALID_SLASH_COMMAND_ARGS', |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| return |
| } |
|
|
| if (desktopSlashCommand?.commandName === 'clear') { |
| await handleDesktopClearCommand(ws) |
| return |
| } |
|
|
| |
| sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) |
|
|
| const initialRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId) |
| if (!initialRuntimeTransition.ok) return |
| if (initialRuntimeTransition.waited) { |
| sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) |
| } |
|
|
| |
| let titleState = sessionTitleState.get(sessionId) |
| if (!titleState) { |
| titleState = { |
| userMessageCount: 0, |
| hasCustomTitle: !!(await sessionService.getCustomTitle(sessionId)), |
| firstUserMessage: '', |
| allUserMessages: [], |
| startedGenerationCounts: new Set<number>(), |
| } |
| sessionTitleState.set(sessionId, titleState) |
| } |
| const titleInput = getTitleInputForUserMessage(message.content, desktopSlashCommand) |
| if (titleInput) { |
| titleState.userMessageCount++ |
| titleState.allUserMessages.push(titleInput) |
| if (titleState.userMessageCount === 1) { |
| titleState.firstUserMessage = titleInput |
| } |
| triggerTitleGeneration(ws, sessionId) |
| } |
|
|
| |
| try { |
| await ensureCliSessionStarted(ws, sessionId, 'user_message') |
| } catch (err) { |
| const errMsg = err instanceof Error ? err.message : String(err) |
| const code = |
| err instanceof ConversationStartupError ? err.code : 'CLI_START_FAILED' |
| console.error(`[WS] CLI start failed for ${sessionId}: ${errMsg}`) |
| sendMessage(ws, { |
| type: 'error', |
| message: await buildSessionStartupDiagnosticMessage(sessionId, errMsg), |
| code, |
| retryable: |
| err instanceof ConversationStartupError ? err.retryable : false, |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| return |
| } |
|
|
| const startupRuntimeTransition = await waitForRuntimeTransitionBeforeUserTurn(ws, sessionId) |
| if (startupRuntimeTransition.ok) { |
| if (startupRuntimeTransition.waited) { |
| sendMessage(ws, { type: 'status', state: 'thinking', verb: 'Thinking' }) |
| } |
| } else { |
| return |
| } |
|
|
| |
| |
| |
| let userMessageSent = false |
| const shouldForwardCurrentTurnLocalCommand = |
| createCurrentTurnLocalCommandForwarder(desktopSlashCommand) |
|
|
| rebindSessionOutput(sessionId, ws, { |
| shouldForward: (cliMsg) => { |
| if (userMessageSent || (cliMsg.type === 'result' && cliMsg.is_error)) { |
| return true |
| } |
| return shouldForwardCurrentTurnLocalCommand(cliMsg) |
| }, |
| }) |
|
|
| const sent = conversationService.sendMessage( |
| sessionId, |
| message.content, |
| message.attachments |
| ) |
| if (!sent) { |
| sendMessage(ws, { |
| type: 'error', |
| message: 'CLI process is not running. The session may have ended or the process crashed.', |
| code: 'CLI_NOT_RUNNING', |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| return |
| } |
|
|
| userMessageSent = true |
| } |
|
|
| async function handleDesktopClearCommand( |
| ws: ServerWebSocket<WebSocketData>, |
| ) { |
| const { sessionId } = ws.data |
|
|
| const workDir = conversationService.getSessionWorkDir(sessionId) |
| conversationService.stopSession(sessionId) |
| conversationService.clearOutputCallbacks(sessionId) |
| sessionSlashCommands.delete(sessionId) |
| sessionTitleState.delete(sessionId) |
| cleanupStreamState(sessionId) |
|
|
| try { |
| await sessionService.clearSessionTranscript(sessionId, workDir || undefined) |
| } catch (err) { |
| const errMsg = err instanceof Error ? err.message : String(err) |
| sendMessage(ws, { |
| type: 'error', |
| message: errMsg, |
| code: 'SESSION_CLEAR_FAILED', |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| return |
| } |
|
|
| sendMessage(ws, { |
| type: 'system_notification', |
| subtype: 'session_cleared', |
| message: 'Conversation cleared', |
| }) |
| sendMessage(ws, { |
| type: 'message_complete', |
| usage: { input_tokens: 0, output_tokens: 0 }, |
| }) |
| } |
|
|
| async function handlePrewarmSession(ws: ServerWebSocket<WebSocketData>) { |
| const { sessionId } = ws.data |
| if (conversationService.hasSession(sessionId) || sessionStartupPromises.has(sessionId)) { |
| return |
| } |
|
|
| const launchInfo = await sessionService.getSessionLaunchInfo(sessionId).catch(() => null) |
| if (launchInfo?.repository) { |
| console.log(`[WS] Skipping prewarm for pending repository launch session ${sessionId}`) |
| return |
| } |
|
|
| prewarmPendingSessions.add(sessionId) |
| void ensureCliSessionStarted(ws, sessionId, 'prewarm_session') |
| .then(() => { |
| if (!prewarmPendingSessions.delete(sessionId)) return |
| bindPrewarmMetadataCapture(sessionId) |
| markPrewarmed(sessionId) |
| }) |
| .catch((err) => { |
| prewarmPendingSessions.delete(sessionId) |
| console.warn( |
| `[WS] Prewarm failed for ${sessionId}: ${ |
| err instanceof Error ? err.message : String(err) |
| }`, |
| ) |
| }) |
| } |
|
|
| function handlePermissionResponse( |
| ws: ServerWebSocket<WebSocketData>, |
| message: Extract<ClientMessage, { type: 'permission_response' }> |
| ) { |
| const { sessionId } = ws.data |
| conversationService.respondToPermission( |
| sessionId, |
| message.requestId, |
| message.allowed, |
| message.rule, |
| message.updatedInput, |
| ) |
| console.log(`[WS] Permission response for ${message.requestId}: ${message.allowed}`) |
| } |
|
|
| function handleComputerUsePermissionResponse( |
| ws: ServerWebSocket<WebSocketData>, |
| message: Extract<ClientMessage, { type: 'computer_use_permission_response' }> |
| ) { |
| const { sessionId } = ws.data |
| const ok = computerUseApprovalService.resolveApproval( |
| message.requestId, |
| message.response, |
| ) |
| if (!ok) { |
| console.warn( |
| `[WS] Ignored Computer Use permission response for unknown request ${message.requestId} from ${sessionId}` |
| ) |
| } |
| } |
|
|
| function handleSetPermissionMode( |
| ws: ServerWebSocket<WebSocketData>, |
| message: Extract<ClientMessage, { type: 'set_permission_mode' }> |
| ) { |
| const { sessionId } = ws.data |
|
|
| |
| |
| |
| |
| |
| const needsRestart = |
| conversationService.hasSession(sessionId) && |
| (message.mode === 'bypassPermissions' || conversationService.getSessionPermissionMode(sessionId) === 'bypassPermissions') |
|
|
| if (needsRestart) { |
| void enqueueRuntimeTransition(sessionId, () => |
| restartSessionWithPermissionMode(ws, sessionId, message.mode), |
| ) |
| return |
| } |
|
|
| const ok = conversationService.setPermissionMode(sessionId, message.mode) |
| if (!ok) { |
| console.warn(`[WS] Ignored permission mode update for inactive session ${sessionId}`) |
| } |
| } |
|
|
| async function handleSetRuntimeConfig( |
| ws: ServerWebSocket<WebSocketData>, |
| message: Extract<ClientMessage, { type: 'set_runtime_config' }> |
| ) { |
| const { sessionId } = ws.data |
| const modelId = typeof message.modelId === 'string' ? message.modelId.trim() : '' |
| if (!modelId) { |
| sendMessage(ws, { |
| type: 'error', |
| message: 'Runtime model selection is invalid.', |
| code: 'RUNTIME_CONFIG_INVALID', |
| }) |
| return |
| } |
|
|
| const nextOverride = { |
| providerId: message.providerId ?? null, |
| modelId, |
| } |
| const prevOverride = runtimeOverrides.get(sessionId) |
| runtimeOverrides.set(sessionId, nextOverride) |
|
|
| if ( |
| prevOverride && |
| prevOverride.providerId === nextOverride.providerId && |
| prevOverride.modelId === nextOverride.modelId |
| ) { |
| return |
| } |
|
|
| if (!conversationService.hasSession(sessionId)) { |
| const pendingStartup = sessionStartupPromises.get(sessionId) |
| if (pendingStartup) { |
| await enqueueRuntimeTransition(sessionId, async () => { |
| await pendingStartup.catch(() => undefined) |
| const currentOverride = runtimeOverrides.get(sessionId) |
| if ( |
| currentOverride?.providerId !== nextOverride.providerId || |
| currentOverride.modelId !== nextOverride.modelId || |
| !conversationService.hasSession(sessionId) |
| ) { |
| return |
| } |
| await restartSessionWithRuntimeConfig(ws, sessionId) |
| }) |
| } |
| return |
| } |
|
|
| await enqueueRuntimeTransition(sessionId, () => |
| restartSessionWithRuntimeConfig(ws, sessionId), |
| ) |
| } |
|
|
| async function restartSessionWithPermissionMode( |
| ws: ServerWebSocket<WebSocketData>, |
| sessionId: string, |
| mode: string, |
| ): Promise<void> { |
| try { |
| |
| await settingsService.setPermissionMode(mode) |
|
|
| const workDir = conversationService.getSessionWorkDir(sessionId) |
| conversationService.stopSession(sessionId) |
|
|
| |
| const runtimeSettings = await getRuntimeSettings(sessionId) |
| const sdkUrl = |
| `ws://${ws.data.serverHost}:${ws.data.serverPort}/sdk/${sessionId}` + |
| `?token=${encodeURIComponent(crypto.randomUUID())}` |
| await conversationService.startSession(sessionId, workDir, sdkUrl, runtimeSettings) |
|
|
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| console.log(`[WS] Restarted CLI for ${sessionId} with permission mode: ${mode}`) |
| } catch (err) { |
| const errMsg = err instanceof Error ? err.message : String(err) |
| void diagnosticsService.recordEvent({ |
| type: 'permission_restart_failed', |
| severity: 'error', |
| sessionId, |
| summary: errMsg, |
| details: { mode, error: err }, |
| }) |
| console.error(`[WS] Failed to restart CLI for ${sessionId}: ${errMsg}`) |
| sendMessage(ws, { |
| type: 'error', |
| message: await buildSessionStartupDiagnosticMessage( |
| sessionId, |
| `Failed to restart session with new permission mode: ${errMsg}`, |
| ), |
| code: 'CLI_RESTART_FAILED', |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| } |
| } |
|
|
| async function restartSessionWithRuntimeConfig( |
| ws: ServerWebSocket<WebSocketData>, |
| sessionId: string, |
| ): Promise<void> { |
| try { |
| const workDir = conversationService.getSessionWorkDir(sessionId) |
| conversationService.stopSession(sessionId) |
|
|
| const runtimeSettings = await getRuntimeSettings(sessionId) |
| const sdkUrl = |
| `ws://${ws.data.serverHost}:${ws.data.serverPort}/sdk/${sessionId}` + |
| `?token=${encodeURIComponent(crypto.randomUUID())}` |
| await conversationService.startSession(sessionId, workDir, sdkUrl, runtimeSettings) |
|
|
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| console.log(`[WS] Restarted CLI for ${sessionId} with runtime override`) |
| } catch (err) { |
| const errMsg = err instanceof Error ? err.message : String(err) |
| void diagnosticsService.recordEvent({ |
| type: 'runtime_config_restart_failed', |
| severity: 'error', |
| sessionId, |
| summary: errMsg, |
| details: { runtimeOverride: runtimeOverrides.get(sessionId), error: err }, |
| }) |
| console.error(`[WS] Failed to restart CLI for ${sessionId} after runtime override: ${errMsg}`) |
| sendMessage(ws, { |
| type: 'error', |
| message: await buildSessionStartupDiagnosticMessage( |
| sessionId, |
| `Failed to switch provider/model: ${errMsg}`, |
| ), |
| code: 'CLI_RESTART_FAILED', |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| } |
| } |
|
|
| function handleStopGeneration(ws: ServerWebSocket<WebSocketData>) { |
| const { sessionId } = ws.data |
| console.log(`[WS] Stop generation requested for session: ${sessionId}`) |
|
|
| sessionStopRequested.add(sessionId) |
|
|
| if (conversationService.hasSession(sessionId)) { |
| |
| conversationService.sendInterrupt(sessionId) |
|
|
| |
| setTimeout(() => { |
| if (conversationService.hasSession(sessionId)) { |
| console.log(`[WS] Force-killing CLI subprocess for session: ${sessionId}`) |
| conversationService.stopSession(sessionId) |
| } |
| }, 3_000) |
| } |
|
|
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| } |
|
|
| |
| |
| |
|
|
| function triggerTitleGeneration(ws: ServerWebSocket<WebSocketData>, sessionId: string): void { |
| const state = sessionTitleState.get(sessionId) |
| if (!state || state.hasCustomTitle) return |
|
|
| const count = state.userMessageCount |
|
|
| |
| if (count !== 1 && count !== 3) return |
| if (state.startedGenerationCounts.has(count)) return |
| state.startedGenerationCounts.add(count) |
|
|
| const text = count === 1 |
| ? state.firstUserMessage |
| : state.allUserMessages.join('\n') |
| const runtimeProviderId = runtimeOverrides.get(sessionId)?.providerId |
|
|
| |
| void (async () => { |
| try { |
| |
| if (count === 1) { |
| const placeholder = deriveTitle(text) |
| if (placeholder) { |
| const saved = await saveAiTitle(sessionId, placeholder) |
| if (!saved) { |
| state.hasCustomTitle = true |
| return |
| } |
| sendMessage(ws, { type: 'session_title_updated', sessionId, title: placeholder }) |
| } |
| } |
|
|
| |
| const aiTitle = await generateTitle(text, runtimeProviderId) |
| if (aiTitle) { |
| const saved = await saveAiTitle(sessionId, aiTitle) |
| if (!saved) { |
| state.hasCustomTitle = true |
| return |
| } |
| sendMessage(ws, { type: 'session_title_updated', sessionId, title: aiTitle }) |
| } |
| } catch (err) { |
| console.error(`[Title] Failed to generate title for ${sessionId}:`, err) |
| } |
| })() |
| } |
|
|
| |
| |
| |
|
|
| |
| |
| |
| |
| type SessionStreamState = { |
| hasReceivedStreamEvents: boolean |
| activeBlockTypes: Map<number, 'text' | 'tool_use' | 'thinking'> |
| activeToolBlocks: Map<number, { toolName: string; toolUseId: string; inputJson: string; parentToolUseId?: string }> |
| pendingLocalCommand?: { name: string; args: string } |
| |
| |
| pendingToolBlocks: Map<string, { toolName: string; toolUseId: string; parentToolUseId?: string }> |
| toolParentUseIds: Map<string, string> |
| lastApiError?: { |
| message: string |
| code: string |
| } |
| } |
|
|
| const sessionStreamStates = new Map<string, SessionStreamState>() |
|
|
| function getStreamState(sessionId: string): SessionStreamState { |
| let state = sessionStreamStates.get(sessionId) |
| if (!state) { |
| state = { |
| hasReceivedStreamEvents: false, |
| activeBlockTypes: new Map(), |
| activeToolBlocks: new Map(), |
| pendingLocalCommand: undefined, |
| pendingToolBlocks: new Map(), |
| toolParentUseIds: new Map(), |
| lastApiError: undefined, |
| } |
| sessionStreamStates.set(sessionId, state) |
| } |
| return state |
| } |
|
|
| function cliParentToolUseId(cliMsg: any): string | undefined { |
| return typeof cliMsg.parent_tool_use_id === 'string' && cliMsg.parent_tool_use_id.length > 0 |
| ? cliMsg.parent_tool_use_id |
| : undefined |
| } |
|
|
| function rememberToolParentUseId( |
| streamState: SessionStreamState, |
| toolUseId: string | undefined, |
| parentToolUseId: string | undefined, |
| ): void { |
| if (!toolUseId || !parentToolUseId) return |
| streamState.toolParentUseIds.set(toolUseId, parentToolUseId) |
| } |
|
|
| function consumeToolParentUseId( |
| streamState: SessionStreamState, |
| toolUseId: string | undefined, |
| ): string | undefined { |
| if (!toolUseId) return undefined |
| const parentToolUseId = streamState.toolParentUseIds.get(toolUseId) |
| streamState.toolParentUseIds.delete(toolUseId) |
| return parentToolUseId |
| } |
|
|
| |
| function cleanupStreamState(sessionId: string) { |
| sessionStreamStates.delete(sessionId) |
| } |
|
|
| function cleanupSessionRuntimeState(sessionId: string) { |
| cleanupStreamState(sessionId) |
| sessionSlashCommands.delete(sessionId) |
| sessionTitleState.delete(sessionId) |
| runtimeOverrides.delete(sessionId) |
| runtimeTransitionPromises.delete(sessionId) |
| sessionStartupPromises.delete(sessionId) |
| lastResolvedStartupWorkDirs.delete(sessionId) |
| clearPrewarmState(sessionId) |
| } |
|
|
| function getPrewarmIdleTimeoutMs(): number { |
| const raw = process.env.CC_HAHA_PREWARM_IDLE_TIMEOUT_MS |
| if (!raw) return DEFAULT_PREWARM_IDLE_TIMEOUT_MS |
| const parsed = Number.parseInt(raw, 10) |
| return Number.isFinite(parsed) && parsed >= 0 |
| ? parsed |
| : DEFAULT_PREWARM_IDLE_TIMEOUT_MS |
| } |
|
|
| function clearPrewarmState(sessionId: string) { |
| prewarmPendingSessions.delete(sessionId) |
| prewarmedSessions.delete(sessionId) |
| const timer = prewarmIdleTimers.get(sessionId) |
| if (timer) { |
| clearTimeout(timer) |
| prewarmIdleTimers.delete(sessionId) |
| } |
| } |
|
|
| function markPrewarmed(sessionId: string) { |
| prewarmedSessions.add(sessionId) |
| const timeoutMs = getPrewarmIdleTimeoutMs() |
| if (timeoutMs === 0) return |
|
|
| const existingTimer = prewarmIdleTimers.get(sessionId) |
| if (existingTimer) clearTimeout(existingTimer) |
|
|
| const timer = setTimeout(() => { |
| prewarmIdleTimers.delete(sessionId) |
| if (!prewarmedSessions.has(sessionId)) return |
| console.log(`[WS] Prewarmed session ${sessionId} idle for ${timeoutMs}ms, stopping CLI subprocess`) |
| conversationService.stopSession(sessionId) |
| prewarmedSessions.delete(sessionId) |
| }, timeoutMs) |
| prewarmIdleTimers.set(sessionId, timer) |
| } |
|
|
| function cacheSessionInitMetadata(sessionId: string, cliMsg: any) { |
| if (cliMsg?.type !== 'system' || cliMsg.subtype !== 'init') return |
| if (typeof cliMsg.cwd === 'string' && cliMsg.cwd.trim()) { |
| conversationService.updateSessionWorkDir(sessionId, cliMsg.cwd) |
| void (async () => { |
| await sessionService.appendSessionMetadata(sessionId, { |
| workDir: cliMsg.cwd, |
| }) |
| await sessionService.deletePlaceholderSessionFiles(sessionId, cliMsg.cwd) |
| })() |
| } |
| if (cliMsg.slash_commands && Array.isArray(cliMsg.slash_commands)) { |
| updateSessionSlashCommands(sessionId, cliMsg.slash_commands, { notifyClient: false }) |
| } |
| } |
|
|
| function extractAssistantText(cliMsg: any): string { |
| const content = cliMsg?.message?.content |
| if (!Array.isArray(content)) return '' |
| const textBlock = content.find( |
| (block: unknown): block is { type: string; text: string } => |
| !!block && |
| typeof block === 'object' && |
| (block as { type?: unknown }).type === 'text' && |
| typeof (block as { text?: unknown }).text === 'string', |
| ) |
| return textBlock?.text || '' |
| } |
|
|
| function isDuplicateOfLastApiError( |
| lastApiError: SessionStreamState['lastApiError'], |
| resultMessage: string, |
| ): boolean { |
| if (!lastApiError?.message) return false |
| if (resultMessage === lastApiError.message) return true |
| return ( |
| resultMessage.includes(lastApiError.message) && |
| /CLI (?:process exited unexpectedly|exited during startup)/i.test(resultMessage) |
| ) |
| } |
|
|
| function bindPrewarmMetadataCapture(sessionId: string) { |
| for (const msg of conversationService.getRecentSdkMessages(sessionId)) { |
| cacheSessionInitMetadata(sessionId, msg) |
| } |
| if (!conversationService.hasSession(sessionId)) return |
|
|
| conversationService.clearOutputCallbacks(sessionId) |
| conversationService.onOutput(sessionId, (cliMsg) => { |
| cacheSessionInitMetadata(sessionId, cliMsg) |
| }) |
| } |
|
|
| async function resolveSessionWorkDir(sessionId: string, fallback = os.homedir()): Promise<string> { |
| let workDir = fallback |
| try { |
| const resolved = await sessionService.getSessionWorkDir(sessionId) |
| if (resolved) workDir = resolved |
| console.log( |
| `[WS] resolveSessionWorkDir: sessionId=${sessionId}, resolved workDir=${JSON.stringify( |
| resolved, |
| )}, will spawn CLI with workDir=${workDir}`, |
| ) |
| } catch (resolveErr) { |
| console.warn( |
| `[WS] resolveSessionWorkDir: failed to resolve workDir for ${sessionId}, using fallback=${workDir}: ${ |
| resolveErr instanceof Error ? resolveErr.message : String(resolveErr) |
| }`, |
| ) |
| } |
| return workDir |
| } |
|
|
| async function ensureCliSessionStarted( |
| ws: ServerWebSocket<WebSocketData>, |
| sessionId: string, |
| reason: 'user_message' | 'prewarm_session', |
| ): Promise<void> { |
| const pendingStartup = sessionStartupPromises.get(sessionId) |
| if (pendingStartup) { |
| await pendingStartup |
| return |
| } |
|
|
| if (conversationService.hasSession(sessionId)) return |
|
|
| const startup = (async () => { |
| const workDir = await resolveSessionWorkDir(sessionId) |
| lastResolvedStartupWorkDirs.set(sessionId, workDir) |
| const runtimeSettings = await getRuntimeSettings(sessionId) |
| const sdkUrl = |
| `ws://${ws.data.serverHost}:${ws.data.serverPort}/sdk/${sessionId}` + |
| `?token=${encodeURIComponent(crypto.randomUUID())}` |
| await sendRepositoryStartupStatus(ws, sessionId, reason) |
| console.log(`[WS] Starting CLI for ${sessionId} due to ${reason}`) |
| await conversationService.startSession(sessionId, workDir, sdkUrl, runtimeSettings) |
| })() |
|
|
| sessionStartupPromises.set(sessionId, startup) |
| try { |
| await startup |
| } finally { |
| if (sessionStartupPromises.get(sessionId) === startup) { |
| sessionStartupPromises.delete(sessionId) |
| } |
| } |
| } |
|
|
| export function translateCliMessage(cliMsg: any, sessionId: string): ServerMessage[] { |
| const streamState = getStreamState(sessionId) |
| switch (cliMsg.type) { |
| case 'assistant': { |
| if (cliMsg.error || cliMsg.isApiErrorMessage) { |
| const message = extractAssistantText(cliMsg) || cliMsg.error || 'Unknown API error' |
| const code = typeof cliMsg.error === 'string' ? cliMsg.error : 'API_ERROR' |
| streamState.lastApiError = { message, code } |
| return [{ |
| type: 'error', |
| message, |
| code, |
| }] |
| } |
|
|
| |
| |
| if (cliMsg.message?.content && Array.isArray(cliMsg.message.content)) { |
| const messages: ServerMessage[] = [] |
|
|
| for (const block of cliMsg.message.content) { |
| if (streamState.hasReceivedStreamEvents) { |
| |
| |
| |
| if (block.type === 'tool_use' && streamState.pendingToolBlocks.has(block.id)) { |
| const pending = streamState.pendingToolBlocks.get(block.id)! |
| streamState.pendingToolBlocks.delete(block.id) |
| rememberToolParentUseId(streamState, block.id, pending.parentToolUseId) |
| messages.push({ |
| type: 'tool_use_complete', |
| toolName: pending.toolName || block.name, |
| toolUseId: block.id, |
| input: block.input, |
| parentToolUseId: pending.parentToolUseId, |
| }) |
| } |
| } else { |
| |
| if (block.type === 'thinking' && block.thinking) { |
| messages.push({ type: 'thinking', text: block.thinking }) |
| } else if (block.type === 'text' && block.text) { |
| messages.push({ type: 'content_start', blockType: 'text' }) |
| messages.push({ type: 'content_delta', text: block.text }) |
| } else if (block.type === 'tool_use') { |
| const parentToolUseId = cliParentToolUseId(cliMsg) |
| rememberToolParentUseId(streamState, block.id, parentToolUseId) |
| messages.push({ |
| type: 'tool_use_complete', |
| toolName: block.name, |
| toolUseId: block.id, |
| input: block.input, |
| parentToolUseId, |
| }) |
| } |
| } |
| } |
|
|
| |
| streamState.hasReceivedStreamEvents = false |
| streamState.pendingToolBlocks.clear() |
| return messages |
| } |
| return [] |
| } |
|
|
| case 'user': { |
| |
| |
| const messages: ServerMessage[] = [] |
|
|
| if (isCompactSummaryMessageContent(cliMsg.message?.content)) { |
| messages.push({ |
| type: 'system_notification', |
| subtype: 'compact_summary', |
| message: cliMsg.message.content, |
| data: { |
| isSynthetic: cliMsg.isSynthetic, |
| }, |
| }) |
| } |
|
|
| const localCommandOutput = extractLocalCommandOutput( |
| cliMsg.message?.content, |
| ) |
| if (localCommandOutput) { |
| const pendingLocalCommand = streamState.pendingLocalCommand |
| streamState.pendingLocalCommand = undefined |
| if (!isCompactLocalCommandOutput(localCommandOutput)) { |
| const goalEvent = extractGoalEvent( |
| localCommandOutput, |
| pendingLocalCommand, |
| ) |
| if (goalEvent) { |
| messages.push({ |
| type: 'system_notification', |
| subtype: 'goal_event', |
| message: goalEvent.message, |
| data: goalEvent, |
| }) |
| } else { |
| messages.push({ type: 'content_start', blockType: 'text' }) |
| messages.push({ type: 'content_delta', text: localCommandOutput }) |
| } |
| } |
| } |
|
|
| if (cliMsg.message?.content && Array.isArray(cliMsg.message.content)) { |
| for (const block of cliMsg.message.content) { |
| if (block.type === 'tool_result') { |
| const rememberedParentToolUseId = consumeToolParentUseId(streamState, block.tool_use_id) |
| const parentToolUseId = |
| cliParentToolUseId(cliMsg) ?? rememberedParentToolUseId |
| messages.push({ |
| type: 'tool_result', |
| toolUseId: block.tool_use_id, |
| content: block.content, |
| isError: !!block.is_error, |
| parentToolUseId, |
| }) |
| } |
| } |
| } |
|
|
| return messages |
| } |
|
|
| case 'stream_event': { |
| streamState.hasReceivedStreamEvents = true |
| const event = cliMsg.event |
| if (!event) return [] |
|
|
| switch (event.type) { |
| case 'message_start': { |
| return [{ type: 'status', state: 'thinking' }] |
| } |
|
|
| case 'content_block_start': { |
| const contentBlock = event.content_block |
| if (!contentBlock) return [] |
|
|
| const index = event.index ?? 0 |
|
|
| if (contentBlock.type === 'tool_use') { |
| const parentToolUseId = cliParentToolUseId(cliMsg) |
| streamState.activeBlockTypes.set(index, 'tool_use') |
| |
| streamState.activeToolBlocks.set(index, { |
| toolName: contentBlock.name || '', |
| toolUseId: contentBlock.id || '', |
| inputJson: '', |
| parentToolUseId, |
| }) |
| return [{ |
| type: 'content_start', |
| blockType: 'tool_use', |
| toolName: contentBlock.name, |
| toolUseId: contentBlock.id, |
| parentToolUseId, |
| }] |
| } |
|
|
| if (contentBlock.type === 'thinking' || contentBlock.type === 'redacted_thinking') { |
| streamState.activeBlockTypes.set(index, 'thinking') |
| return [{ type: 'status', state: 'thinking', verb: 'Thinking' }] |
| } |
|
|
| streamState.activeBlockTypes.set(index, 'text') |
| return [{ type: 'content_start', blockType: 'text' }] |
| } |
|
|
| case 'content_block_delta': { |
| const delta = event.delta |
| if (!delta) return [] |
|
|
| if (delta.type === 'text_delta' && delta.text) { |
| return [{ type: 'content_delta', text: delta.text }] |
| } |
| if (delta.type === 'input_json_delta' && delta.partial_json) { |
| |
| const index = event.index ?? 0 |
| const toolBlock = streamState.activeToolBlocks.get(index) |
| if (toolBlock) toolBlock.inputJson += delta.partial_json |
| return [{ type: 'content_delta', toolInput: delta.partial_json }] |
| } |
| if (delta.type === 'thinking_delta' && delta.thinking) { |
| return [{ type: 'thinking', text: delta.thinking }] |
| } |
| return [] |
| } |
|
|
| case 'content_block_stop': { |
| const index = event.index ?? 0 |
| const blockType = streamState.activeBlockTypes.get(index) |
| streamState.activeBlockTypes.delete(index) |
|
|
| if (blockType === 'tool_use') { |
| const toolBlock = streamState.activeToolBlocks.get(index) |
| streamState.activeToolBlocks.delete(index) |
| if (toolBlock) { |
| const parentToolUseId = |
| cliParentToolUseId(cliMsg) ?? toolBlock.parentToolUseId |
| let parsedInput = null |
| try { parsedInput = JSON.parse(toolBlock.inputJson) } catch {} |
|
|
| if (parsedInput !== null) { |
| rememberToolParentUseId(streamState, toolBlock.toolUseId, parentToolUseId) |
| return [{ |
| type: 'tool_use_complete', |
| toolName: toolBlock.toolName, |
| toolUseId: toolBlock.toolUseId, |
| input: parsedInput, |
| parentToolUseId, |
| }] |
| } |
|
|
| |
| |
| console.warn( |
| `[WS] Tool input JSON parse failed for ${toolBlock.toolName} (${toolBlock.toolUseId}), deferring to assistant message`, |
| ) |
| streamState.pendingToolBlocks.set(toolBlock.toolUseId, { |
| toolName: toolBlock.toolName, |
| toolUseId: toolBlock.toolUseId, |
| parentToolUseId, |
| }) |
| } |
| } |
| return [] |
| } |
|
|
| case 'message_stop': { |
| |
| return [] |
| } |
|
|
| case 'message_delta': { |
| |
| return [] |
| } |
|
|
| default: |
| return [] |
| } |
| } |
|
|
| case 'control_request': { |
| |
| if (cliMsg.request?.subtype === 'can_use_tool') { |
| return [{ |
| type: 'permission_request', |
| requestId: cliMsg.request_id, |
| toolName: cliMsg.request.tool_name || 'Unknown', |
| toolUseId: |
| typeof cliMsg.request.tool_use_id === 'string' |
| ? cliMsg.request.tool_use_id |
| : undefined, |
| input: cliMsg.request.input || {}, |
| description: cliMsg.request.description, |
| }] |
| } |
| return [] |
| } |
|
|
| case 'control_response': |
| return [] |
|
|
| case 'result': { |
| |
| const usage = { |
| input_tokens: cliMsg.usage?.input_tokens || 0, |
| output_tokens: cliMsg.usage?.output_tokens || 0, |
| } |
|
|
| if (cliMsg.is_error) { |
| |
| |
| if (sessionStopRequested.has(sessionId)) { |
| sessionStopRequested.delete(sessionId) |
| return [{ type: 'message_complete', usage }] |
| } |
|
|
| const resultMessage = |
| (typeof cliMsg.result === 'string' && cliMsg.result) || |
| (Array.isArray(cliMsg.errors) && cliMsg.errors.length > 0 |
| ? cliMsg.errors.join('\n') |
| : 'Unknown error') |
| if (isDuplicateOfLastApiError(streamState.lastApiError, resultMessage)) { |
| streamState.lastApiError = undefined |
| return [{ type: 'message_complete', usage }] |
| } |
| |
| return [ |
| { |
| type: 'error', |
| message: resultMessage, |
| code: 'CLI_ERROR', |
| }, |
| { type: 'message_complete', usage }, |
| ] |
| } |
|
|
| |
| sessionStopRequested.delete(sessionId) |
| streamState.lastApiError = undefined |
| return [{ type: 'message_complete', usage }] |
| } |
|
|
| case 'system': { |
| |
| const subtype = cliMsg.subtype |
| if (subtype === 'api_retry') { |
| const apiRetryMessage = toApiRetryServerMessage(cliMsg) |
| return apiRetryMessage ? [apiRetryMessage] : [] |
| } |
| if (subtype === 'init') { |
| |
| |
| |
| |
| cacheSessionInitMetadata(sessionId, cliMsg) |
| const messages: ServerMessage[] = [ |
| |
| { type: 'system_notification', subtype: 'init', message: `Model: ${cliMsg.model || 'unknown'}`, data: { model: cliMsg.model } }, |
| ] |
| |
| const cmds = sessionSlashCommands.get(sessionId) |
| if (cmds && cmds.length > 0) { |
| messages.push({ |
| type: 'system_notification', |
| subtype: 'slash_commands', |
| data: cmds, |
| }) |
| } |
| return messages |
| } |
| if (subtype === 'memory_saved') { |
| return [{ |
| type: 'system_notification', |
| subtype: 'memory_saved', |
| message: cliMsg.message, |
| data: { |
| writtenPaths: Array.isArray(cliMsg.writtenPaths) ? cliMsg.writtenPaths : [], |
| teamCount: typeof cliMsg.teamCount === 'number' ? cliMsg.teamCount : undefined, |
| verb: typeof cliMsg.verb === 'string' ? cliMsg.verb : undefined, |
| }, |
| }] |
| } |
| if (subtype === 'status') { |
| if (cliMsg.status === 'compacting') { |
| return [{ |
| type: 'status', |
| state: 'compacting', |
| verb: 'Compacting conversation', |
| }] |
| } |
| if (cliMsg.status == null) { |
| return [{ type: 'status', state: 'thinking', verb: 'Thinking' }] |
| } |
| return [] |
| } |
| if (subtype === 'hook_started' || subtype === 'hook_response') { |
| |
| return [] |
| } |
| if (subtype === 'local_command' || subtype === 'local_command_output') { |
| const localCommand = extractLocalCommand(cliMsg.content ?? cliMsg.message) |
| if (localCommand) { |
| streamState.pendingLocalCommand = localCommand |
| return [] |
| } |
|
|
| const localCommandOutput = extractLocalCommandOutput( |
| cliMsg.content ?? cliMsg.message, |
| { allowUntagged: subtype === 'local_command_output' }, |
| ) |
| if (!localCommandOutput) return [] |
| const goalEvent = extractGoalEvent( |
| localCommandOutput, |
| streamState.pendingLocalCommand, |
| ) |
| streamState.pendingLocalCommand = undefined |
| if (goalEvent) { |
| return [{ |
| type: 'system_notification', |
| subtype: 'goal_event', |
| message: goalEvent.message, |
| data: goalEvent, |
| }] |
| } |
| return [ |
| { type: 'content_start', blockType: 'text' }, |
| { type: 'content_delta', text: localCommandOutput }, |
| ] |
| } |
| |
| if (subtype === 'task_notification') { |
| return [{ |
| type: 'system_notification', |
| subtype: 'task_notification', |
| message: cliMsg.message || cliMsg.title, |
| data: cliMsg, |
| }] |
| } |
| if (subtype === 'task_started') { |
| return [ |
| { |
| type: 'system_notification', |
| subtype: 'task_started', |
| message: cliMsg.message || cliMsg.description || 'Task started', |
| data: cliMsg, |
| }, |
| { |
| type: 'status', |
| state: 'tool_executing', |
| verb: cliMsg.message || cliMsg.description || 'Task started', |
| }, |
| ] |
| } |
| if (subtype === 'task_progress') { |
| return [ |
| { |
| type: 'system_notification', |
| subtype: 'task_progress', |
| message: cliMsg.message || cliMsg.summary || cliMsg.description || 'Task in progress', |
| data: cliMsg, |
| }, |
| { |
| type: 'status', |
| state: 'tool_executing', |
| verb: cliMsg.message || cliMsg.summary || cliMsg.description || 'Task in progress', |
| }, |
| ] |
| } |
| if (subtype === 'session_state_changed') { |
| return [{ |
| type: 'system_notification', |
| subtype: 'session_state_changed', |
| message: cliMsg.message, |
| data: cliMsg, |
| }] |
| } |
| if (subtype === 'compact_boundary') { |
| return [{ |
| type: 'system_notification', |
| subtype: 'compact_boundary', |
| message: getCompactBoundaryMessage(cliMsg), |
| data: cliMsg.compact_metadata ?? cliMsg, |
| }] |
| } |
| |
| return [] |
| } |
|
|
| default: |
| |
| console.log(`[WS] Unknown CLI message type: ${cliMsg.type}`, JSON.stringify(cliMsg).substring(0, 200)) |
| return [] |
| } |
| } |
|
|
| |
| |
| |
|
|
| function finiteNumber(value: unknown): number | null { |
| return typeof value === 'number' && Number.isFinite(value) ? value : null |
| } |
|
|
| function normalizeRetryCount(value: unknown): number | null { |
| const numeric = finiteNumber(value) |
| if (numeric === null) return null |
| return Math.max(0, Math.trunc(numeric)) |
| } |
|
|
| function readRetryErrorRecord(value: unknown): Record<string, unknown> | null { |
| if (!value || typeof value !== 'object' || Array.isArray(value)) return null |
| return value as Record<string, unknown> |
| } |
|
|
| function readRetryErrorString(value: unknown, keys: string[]): string | undefined { |
| const record = readRetryErrorRecord(value) |
| if (!record) return undefined |
| for (const key of keys) { |
| const candidate = record[key] |
| if (typeof candidate === 'string' && candidate.trim()) return candidate.trim() |
| } |
| return undefined |
| } |
|
|
| function toApiRetryServerMessage(cliMsg: any): ServerMessage | null { |
| const attempt = normalizeRetryCount(cliMsg.attempt) |
| const maxRetries = normalizeRetryCount(cliMsg.max_retries) |
| const retryDelayMs = normalizeRetryCount(cliMsg.retry_delay_ms) |
| if (attempt === null || maxRetries === null || retryDelayMs === null) return null |
|
|
| const embeddedError = readRetryErrorRecord(cliMsg.error) |
| const embeddedStatus = embeddedError ? finiteNumber(embeddedError.status) : null |
| const rawStatus = cliMsg.error_status === null |
| ? null |
| : finiteNumber(cliMsg.error_status) ?? embeddedStatus |
| const errorType = typeof cliMsg.error === 'string' && cliMsg.error.trim() |
| ? cliMsg.error.trim() |
| : readRetryErrorString(cliMsg.error, ['type', 'code', 'name']) |
| const errorMessage = readRetryErrorString(cliMsg.error, ['message', 'error']) |
|
|
| return { |
| type: 'api_retry', |
| attempt, |
| maxRetries, |
| retryDelayMs, |
| errorStatus: rawStatus === null ? null : Math.trunc(rawStatus), |
| ...(errorType ? { errorType } : {}), |
| ...(errorMessage ? { errorMessage } : {}), |
| } |
| } |
|
|
| function sendMessage(ws: ServerWebSocket<WebSocketData>, message: ServerMessage) { |
| ws.send(JSON.stringify(message)) |
| } |
|
|
| function sendError(ws: ServerWebSocket<WebSocketData>, message: string, code: string) { |
| sendMessage(ws, { type: 'error', message, code }) |
| } |
|
|
| function getDesktopSlashCommand(content: string): ReturnType<typeof parseSlashCommand> { |
| const parsed = parseSlashCommand(content.trim()) |
| if (!parsed || parsed.isMcp) return null |
| return parsed |
| } |
|
|
| function getTitleInputForUserMessage( |
| content: string, |
| command: ReturnType<typeof parseSlashCommand>, |
| ): string | null { |
| if (command?.commandName !== 'goal') return content |
|
|
| const args = command.args.trim() |
| if (!args || args === 'clear') return null |
| return args |
| } |
|
|
| export function createCurrentTurnLocalCommandForwarder( |
| command: ReturnType<typeof parseSlashCommand>, |
| ): (cliMsg: any) => boolean { |
| let awaitingCurrentTurnLocalCommandOutput = false |
|
|
| return (cliMsg: any) => { |
| if (command && isMatchingCurrentTurnLocalCommand(cliMsg, command)) { |
| awaitingCurrentTurnLocalCommandOutput = true |
| return true |
| } |
| if (command?.commandName === 'goal' && isLocalCommandOutputMessage(cliMsg)) { |
| const output = extractLocalCommandOutput( |
| cliMsg.content ?? cliMsg.message, |
| { allowUntagged: cliMsg.subtype === 'local_command_output' }, |
| ) |
| if (output && looksLikeGoalCommandOutput(output)) { |
| awaitingCurrentTurnLocalCommandOutput = false |
| return true |
| } |
| } |
| if ( |
| awaitingCurrentTurnLocalCommandOutput && |
| isLocalCommandOutputMessage(cliMsg) |
| ) { |
| awaitingCurrentTurnLocalCommandOutput = false |
| return true |
| } |
| return false |
| } |
| } |
|
|
| function isMatchingCurrentTurnLocalCommand( |
| cliMsg: any, |
| command: NonNullable<ReturnType<typeof parseSlashCommand>>, |
| ): boolean { |
| if (cliMsg?.type !== 'system' || cliMsg?.subtype !== 'local_command') { |
| return false |
| } |
| const localCommand = extractLocalCommand(cliMsg.content ?? cliMsg.message) |
| if (!localCommand) return false |
| return ( |
| localCommand.name === command.commandName && |
| localCommand.args.trim() === command.args.trim() |
| ) |
| } |
|
|
| function isLocalCommandOutputMessage(cliMsg: any): boolean { |
| if ( |
| cliMsg?.type !== 'system' || |
| (cliMsg?.subtype !== 'local_command' && |
| cliMsg?.subtype !== 'local_command_output') |
| ) { |
| return false |
| } |
| return extractLocalCommandOutput( |
| cliMsg.content ?? cliMsg.message, |
| { allowUntagged: cliMsg.subtype === 'local_command_output' }, |
| ) !== null |
| } |
|
|
| function extractLocalCommandOutput( |
| content: unknown, |
| options: { allowUntagged?: boolean } = {}, |
| ): string | null { |
| const raw = typeof content === 'string' |
| ? content |
| : Array.isArray(content) |
| ? content |
| .flatMap((block) => { |
| if (!block || typeof block !== 'object') return [] |
| const text = (block as { text?: unknown }).text |
| return typeof text === 'string' ? [text] : [] |
| }) |
| .join('\n') |
| : '' |
|
|
| if (!raw) return null |
|
|
| const stdout = extractTaggedContent(raw, LOCAL_COMMAND_STDOUT_TAG) |
| if (stdout !== null) return stdout |
|
|
| const stderr = extractTaggedContent(raw, LOCAL_COMMAND_STDERR_TAG) |
| if (stderr !== null) return stderr |
|
|
| if (options.allowUntagged) { |
| const normalized = raw.trim() |
| return normalized || null |
| } |
|
|
| return null |
| } |
|
|
| function isCompactLocalCommandOutput(output: string): boolean { |
| return output.trim() === 'Compacted' |
| } |
|
|
| function extractTaggedContent(raw: string, tag: string): string | null { |
| const match = raw.match(new RegExp(`<${tag}>([\\s\\S]*?)</${tag}>`)) |
| return match?.[1]?.trim() ?? null |
| } |
|
|
| function extractLocalCommand(content: unknown): { name: string; args: string } | null { |
| const raw = typeof content === 'string' |
| ? content |
| : Array.isArray(content) |
| ? content |
| .flatMap((block) => { |
| if (!block || typeof block !== 'object') return [] |
| const text = (block as { text?: unknown }).text |
| return typeof text === 'string' ? [text] : [] |
| }) |
| .join('\n') |
| : '' |
|
|
| const name = extractTaggedContent(raw, COMMAND_NAME_TAG) |
| if (!name) return null |
| return { |
| name: name.replace(/^\//, ''), |
| args: extractTaggedContent(raw, 'command-args') ?? '', |
| } |
| } |
|
|
| type GoalEventData = { |
| action: 'created' | 'replaced' | 'status' | 'paused' | 'resumed' | 'completed' | 'cleared' | 'message' |
| status?: string |
| objective?: string |
| budget?: string |
| elapsed?: string |
| continuations?: string |
| message?: string |
| } |
|
|
| function extractGoalEvent( |
| output: string, |
| command?: { name: string; args: string }, |
| ): GoalEventData | null { |
| if (command && command.name !== 'goal') return null |
|
|
| const trimmed = output.trim() |
| if (!trimmed) return null |
|
|
| if (trimmed === 'Goal cleared.' || trimmed.startsWith('Goal cleared:')) { |
| return { action: 'cleared', message: trimmed } |
| } |
| if (trimmed === 'Goal marked complete.') { |
| return { action: 'completed', message: trimmed } |
| } |
| if (trimmed === 'No active goal.') { |
| return { action: 'message', message: trimmed } |
| } |
|
|
| if (trimmed.startsWith('Goal set:')) { |
| const objective = trimmed.slice('Goal set:'.length).trim() |
| return { |
| action: 'created', |
| status: 'active', |
| objective: objective || undefined, |
| message: trimmed, |
| } |
| } |
|
|
| return command?.name === 'goal' ? { action: 'message', message: trimmed } : null |
| } |
|
|
| function looksLikeGoalCommandOutput(output: string): boolean { |
| const trimmed = output.trim() |
| return ( |
| trimmed.startsWith('Goal set:') || |
| trimmed.startsWith('Goal cleared:') || |
| trimmed === 'Goal cleared.' || |
| trimmed === 'Goal marked complete.' || |
| trimmed === 'No active goal.' |
| ) |
| } |
|
|
| function getCompactBoundaryMessage(cliMsg: any): string { |
| const message = typeof cliMsg?.message === 'string' ? cliMsg.message.trim() : '' |
| if (message) return message |
|
|
| const content = typeof cliMsg?.content === 'string' ? cliMsg.content.trim() : '' |
| if (content) return content |
|
|
| return 'Context compacted' |
| } |
|
|
| function isCompactSummaryMessageContent(content: unknown): content is string { |
| return ( |
| typeof content === 'string' && |
| content.trim().startsWith( |
| 'This session is being continued from a previous conversation that ran out of context. The summary below covers the earlier portion of the conversation.', |
| ) |
| ) |
| } |
|
|
| function rebindSessionOutput( |
| sessionId: string, |
| ws: ServerWebSocket<WebSocketData>, |
| options?: { |
| shouldForward?: (cliMsg: any) => boolean |
| }, |
| ) { |
| if (!conversationService.hasSession(sessionId)) return |
|
|
| conversationService.clearOutputCallbacks(sessionId) |
| conversationService.onOutput(sessionId, (cliMsg) => { |
| if (options?.shouldForward && !options.shouldForward(cliMsg)) { |
| return |
| } |
|
|
| const serverMsgs = translateCliMessage(cliMsg, sessionId) |
| for (const msg of serverMsgs) { |
| sendMessage(ws, msg) |
| } |
|
|
| if (cliMsg.type === 'result') { |
| triggerTitleGeneration(ws, sessionId) |
| } |
| }) |
| } |
|
|
| type RuntimeSettings = { |
| permissionMode?: string |
| model?: string |
| effort?: string |
| thinking?: 'disabled' |
| providerId?: string | null |
| } |
|
|
| function isKnownRuntimeProviderId( |
| providerId: string, |
| providers: Array<{ id: string }>, |
| ): boolean { |
| return ( |
| isOpenAIOfficialProviderId(providerId) || |
| providers.some((provider) => provider.id === providerId) |
| ) |
| } |
|
|
| async function getRuntimeSettings(sessionId?: string): Promise<RuntimeSettings> { |
| const runtimeOverride = sessionId ? runtimeOverrides.get(sessionId) : undefined |
| if (runtimeOverride) { |
| if (typeof runtimeOverride.providerId === 'string') { |
| const { providers } = await providerService.listProviders() |
| const providerExists = isKnownRuntimeProviderId(runtimeOverride.providerId, providers) |
| if (!providerExists) { |
| console.warn( |
| `[WS] Ignoring stale runtime provider id for ${sessionId}: ${runtimeOverride.providerId}`, |
| ) |
| runtimeOverrides.delete(sessionId!) |
| return getDefaultRuntimeSettings() |
| } |
| } |
|
|
| const userSettings = await settingsService.getUserSettings() |
| const effort = |
| typeof userSettings.effort === 'string' && userSettings.effort.trim() |
| ? userSettings.effort |
| : undefined |
| const thinking = resolveDesktopThinkingMode(userSettings) |
|
|
| return { |
| permissionMode: await settingsService.getPermissionMode().catch(() => undefined), |
| model: runtimeOverride.modelId, |
| effort, |
| thinking, |
| providerId: runtimeOverride.providerId, |
| } |
| } |
|
|
| return getDefaultRuntimeSettings() |
| } |
|
|
| async function getDefaultRuntimeSettings(): Promise<RuntimeSettings> { |
| |
| |
| |
| const cliConfig = await readCliGlobalConfig() |
| if (cliConfig?.authProvider && cliConfig.authProvider !== 'anthropic' && cliConfig.authProvider !== 'openai') { |
| const userSettings = await settingsService.getUserSettings() |
| const baseModel = |
| typeof userSettings.model === 'string' && userSettings.model.trim() |
| ? userSettings.model |
| : undefined |
| const modelContext = |
| typeof userSettings.modelContext === 'string' && userSettings.modelContext.trim() |
| ? userSettings.modelContext |
| : undefined |
| const effort = |
| typeof userSettings.effort === 'string' && userSettings.effort.trim() |
| ? userSettings.effort |
| : undefined |
| const thinking = resolveDesktopThinkingMode(userSettings) |
|
|
| return { |
| permissionMode: await settingsService.getPermissionMode().catch(() => undefined), |
| model: baseModel ? (modelContext ? `${baseModel}:${modelContext}` : baseModel) : undefined, |
| effort, |
| thinking, |
| providerId: undefined, |
| } |
| } |
|
|
| |
| const { providers, activeId } = await providerService.listProviders() |
| let resolvedActiveId = activeId |
| if (activeId && !isKnownRuntimeProviderId(activeId, providers)) { |
| console.warn(`[WS] Active provider id is stale, falling back to official provider: ${activeId}`) |
| resolvedActiveId = null |
| await providerService.activateOfficial() |
| } |
|
|
| const userSettings = await settingsService.getUserSettings() |
| const providerSettings = resolvedActiveId |
| ? await providerService.getManagedSettings() |
| : undefined |
| const modelSettings = providerSettings ?? userSettings |
| const modelContext = |
| typeof modelSettings.modelContext === 'string' && modelSettings.modelContext.trim() |
| ? modelSettings.modelContext |
| : undefined |
| const effort = |
| typeof userSettings.effort === 'string' && userSettings.effort.trim() |
| ? userSettings.effort |
| : undefined |
| const thinking = resolveDesktopThinkingMode(userSettings) |
|
|
| let model: string | undefined |
| if (resolvedActiveId) { |
| |
| |
| const baseModel = |
| typeof modelSettings.model === 'string' && modelSettings.model.trim() |
| ? modelSettings.model |
| : '' |
| if (baseModel) { |
| model = baseModel |
| if (modelContext) model += `:${modelContext}` |
| } |
| } else { |
| |
| const baseModel = |
| typeof userSettings.model === 'string' && userSettings.model.trim() |
| ? userSettings.model |
| : undefined |
| model = baseModel ? (modelContext ? `${baseModel}:${modelContext}` : baseModel) : undefined |
| } |
|
|
| return { |
| permissionMode: await settingsService.getPermissionMode().catch(() => undefined), |
| model, |
| effort, |
| thinking, |
| providerId: resolvedActiveId, |
| } |
| } |
|
|
| function resolveDesktopThinkingMode( |
| settings: Record<string, unknown>, |
| ): 'disabled' | undefined { |
| return settings.alwaysThinkingEnabled === false ? 'disabled' : undefined |
| } |
|
|
| |
| |
| |
| |
| |
| function readCliGlobalConfig(): { |
| authProvider?: 'anthropic' | 'openai' | 'openrouter' | 'local' | 'opencode' | 'nvidia' |
| openCodeApiKey?: string |
| openCodeModelName?: string |
| nvidiaApiKey?: string |
| openRouterApiKey?: string |
| localBaseUrl?: string |
| localModelName?: string |
| } | null { |
| const configDir = process.env.CLAUDE_CONFIG_DIR || path.join(os.homedir(), '.claude') |
| const configPath = path.join(configDir, '.claude.json') |
| try { |
| const raw = fs.readFileSync(configPath, 'utf-8') |
| return JSON.parse(raw) |
| } catch { |
| return null |
| } |
| } |
|
|
| async function buildSessionStartupDiagnosticMessage( |
| sessionId: string, |
| cause: string, |
| ): Promise<string> { |
| const lines = [ |
| cause, |
| '', |
| 'Desktop service diagnostics:', |
| `- sessionId: ${sessionId}`, |
| ] |
|
|
| try { |
| const recentWorkDir = lastResolvedStartupWorkDirs.get(sessionId) |
| const workDir = |
| recentWorkDir || |
| conversationService.getSessionWorkDir(sessionId) || |
| await sessionService.getSessionWorkDir(sessionId) |
| lines.push(`- workDir: ${workDir ?? '(unknown)'}`) |
| } catch (err) { |
| lines.push(`- workDir: failed to resolve (${err instanceof Error ? err.message : String(err)})`) |
| } |
|
|
| const runtimeOverride = runtimeOverrides.get(sessionId) |
| if (runtimeOverride) { |
| lines.push(`- runtimeOverride.providerId: ${runtimeOverride.providerId ?? '(official)'}`) |
| lines.push(`- runtimeOverride.modelId: ${runtimeOverride.modelId}`) |
| } else { |
| lines.push('- runtimeOverride: (none)') |
| } |
|
|
| try { |
| const { providers, activeId } = await providerService.listProviders() |
| lines.push(`- activeProviderId: ${activeId ?? '(official)'}`) |
| lines.push(`- configuredProviders: ${providers.length}`) |
| if (providers.length > 0) { |
| lines.push( |
| `- providerIndex: ${providers |
| .map((provider) => `${provider.name} (${provider.id})`) |
| .join(', ')}`, |
| ) |
| } |
| } catch (err) { |
| lines.push(`- providers: failed to read (${err instanceof Error ? err.message : String(err)})`) |
| } |
|
|
| return lines.join('\n') |
| } |
|
|
| function enqueueRuntimeTransition( |
| sessionId: string, |
| transition: () => Promise<void>, |
| ): Promise<void> { |
| const previous = runtimeTransitionPromises.get(sessionId) ?? Promise.resolve() |
| const next = previous |
| .catch(() => {}) |
| .then(transition) |
| .finally(() => { |
| if (runtimeTransitionPromises.get(sessionId) === next) { |
| runtimeTransitionPromises.delete(sessionId) |
| } |
| }) |
| runtimeTransitionPromises.set(sessionId, next) |
| return next |
| } |
|
|
| async function waitForRuntimeTransitionBeforeUserTurn( |
| ws: ServerWebSocket<WebSocketData>, |
| sessionId: string, |
| ): Promise<{ ok: boolean; waited: boolean }> { |
| let waited = false |
| let pendingRuntimeTransition = runtimeTransitionPromises.get(sessionId) |
| while (pendingRuntimeTransition) { |
| waited = true |
| try { |
| await pendingRuntimeTransition |
| } catch (err) { |
| const errMsg = err instanceof Error ? err.message : String(err) |
| void diagnosticsService.recordEvent({ |
| type: 'runtime_transition_failed', |
| severity: 'error', |
| sessionId, |
| summary: errMsg, |
| details: err, |
| }) |
| console.error(`[WS] Runtime transition failed before handling user message for ${sessionId}: ${errMsg}`) |
| sendMessage(ws, { |
| type: 'error', |
| message: `Failed to switch provider/model: ${errMsg}`, |
| code: 'CLI_RESTART_FAILED', |
| }) |
| sendMessage(ws, { type: 'status', state: 'idle' }) |
| return { ok: false, waited } |
| } |
|
|
| const nextTransition = runtimeTransitionPromises.get(sessionId) |
| pendingRuntimeTransition = |
| nextTransition && nextTransition !== pendingRuntimeTransition |
| ? nextTransition |
| : undefined |
| } |
|
|
| return { ok: true, waited } |
| } |
|
|
| |
| |
| |
| export function sendToSession(sessionId: string, message: ServerMessage): boolean { |
| const ws = activeSessions.get(sessionId) |
| if (!ws) return false |
| ws.send(JSON.stringify(message)) |
| return true |
| } |
|
|
| export function updateSessionSlashCommands( |
| sessionId: string, |
| commands: unknown[], |
| options: { notifyClient?: boolean } = {}, |
| ): SessionSlashCommand[] { |
| const normalized = commands |
| .map(normalizeSessionSlashCommand) |
| .filter((command): command is SessionSlashCommand => command !== null) |
|
|
| sessionSlashCommands.set(sessionId, normalized) |
|
|
| if (options.notifyClient !== false) { |
| sendToSession(sessionId, { |
| type: 'system_notification', |
| subtype: 'slash_commands', |
| data: normalized, |
| }) |
| } |
|
|
| return normalized |
| } |
|
|
| function normalizeSessionSlashCommand(command: unknown): SessionSlashCommand | null { |
| if (typeof command === 'string') { |
| return command.trim() ? { name: command, description: '' } : null |
| } |
| if (!command || typeof command !== 'object') return null |
|
|
| const record = command as { |
| name?: unknown |
| command?: unknown |
| description?: unknown |
| argumentHint?: unknown |
| } |
| const name = |
| typeof record.name === 'string' |
| ? record.name |
| : typeof record.command === 'string' |
| ? record.command |
| : '' |
| if (!name.trim()) return null |
|
|
| return { |
| name, |
| description: typeof record.description === 'string' ? record.description : '', |
| ...(typeof record.argumentHint === 'string' ? { argumentHint: record.argumentHint } : {}), |
| } |
| } |
|
|
| export function closeSessionConnection(sessionId: string, reason = 'session closed'): boolean { |
| const cleanupTimer = sessionCleanupTimers.get(sessionId) |
| if (cleanupTimer) { |
| clearTimeout(cleanupTimer) |
| sessionCleanupTimers.delete(sessionId) |
| } |
| computerUseApprovalService.cancelSession(sessionId) |
| conversationService.clearOutputCallbacks(sessionId) |
| cleanupSessionRuntimeState(sessionId) |
|
|
| const ws = activeSessions.get(sessionId) |
| if (!ws) return false |
|
|
| activeSessions.delete(sessionId) |
| ws.close(1000, reason) |
| return true |
| } |
|
|
| export function getActiveSessionIds(): string[] { |
| return Array.from(activeSessions.keys()) |
| } |
|
|
| export function __resetWebSocketHandlerStateForTests(): void { |
| for (const timer of sessionCleanupTimers.values()) clearTimeout(timer) |
| for (const timer of prewarmIdleTimers.values()) clearTimeout(timer) |
| activeSessions.clear() |
| sessionCleanupTimers.clear() |
| prewarmIdleTimers.clear() |
| } |
|
|