mirror of
https://github.com/NanmiCoder/cc-haha
synced 2026-07-16 13:03:31 +08:00
1659 lines
54 KiB
TypeScript
1659 lines
54 KiB
TypeScript
/**
|
||
* WebSocket connection handler
|
||
*
|
||
* 管理 WebSocket 连接生命周期,处理消息路由。
|
||
* 用户消息通过 CLI 子进程(stream-json 模式)处理,
|
||
* CLI stdout 消息被转换为 ServerMessage 并转发到 WebSocket。
|
||
*/
|
||
|
||
import type { ServerWebSocket } from 'bun'
|
||
import type { ClientMessage, ServerMessage } from './events.js'
|
||
import * as os from 'node:os'
|
||
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 { diagnosticsService } from '../services/diagnosticsService.js'
|
||
import { deriveTitle, generateTitle, saveAiTitle } from '../services/titleService.js'
|
||
import { parseSlashCommand } from '../../utils/slashCommandParsing.js'
|
||
import {
|
||
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()
|
||
|
||
/**
|
||
* Cache slash commands from CLI init messages, keyed by sessionId.
|
||
*/
|
||
export type SessionSlashCommand = {
|
||
name: string
|
||
description: string
|
||
argumentHint?: string
|
||
}
|
||
|
||
const sessionSlashCommands = new Map<string, SessionSlashCommand[]>()
|
||
|
||
/**
|
||
* Timers for delayed session cleanup after client disconnect.
|
||
* If a client reconnects within 5 minutes, the timer is cancelled.
|
||
*/
|
||
const sessionCleanupTimers = new Map<string, ReturnType<typeof setTimeout>>()
|
||
|
||
/**
|
||
* Track sessions where user requested stop — suppress the CLI_ERROR that
|
||
* follows an interrupt so the frontend doesn't show "处理过程中发生错误".
|
||
*/
|
||
const sessionStopRequested = new Set<string>()
|
||
|
||
/**
|
||
* Track user message count and title state per session for auto-title generation.
|
||
*/
|
||
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
|
||
}
|
||
|
||
// Active WebSocket sessions
|
||
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}`)
|
||
|
||
// Cancel pending cleanup timer if client reconnects
|
||
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)
|
||
|
||
// Schedule delayed cleanup: if the client doesn't reconnect within 30 seconds,
|
||
// stop the CLI subprocess to avoid leaking resources.
|
||
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>) {
|
||
// Backpressure handling - called when the socket is ready to receive more data
|
||
},
|
||
}
|
||
|
||
// ============================================================================
|
||
// Message handlers
|
||
// ============================================================================
|
||
|
||
async function handleUserMessage(
|
||
ws: ServerWebSocket<WebSocketData>,
|
||
message: Extract<ClientMessage, { type: 'user_message' }>
|
||
) {
|
||
const { sessionId } = ws.data
|
||
|
||
// Clear any stale stop flag from a previous turn
|
||
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
|
||
}
|
||
|
||
// Send thinking status
|
||
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' })
|
||
}
|
||
|
||
// Track and emit the first placeholder title before CLI startup/streaming.
|
||
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)
|
||
}
|
||
titleState.userMessageCount++
|
||
titleState.allUserMessages.push(message.content)
|
||
if (titleState.userMessageCount === 1) {
|
||
titleState.firstUserMessage = message.content
|
||
}
|
||
triggerTitleGeneration(ws, sessionId)
|
||
|
||
// 启动 CLI 子进程(如果还没有)
|
||
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
|
||
}
|
||
|
||
// Register the callback before sending the turn so startup errors are not lost.
|
||
// Keep output muted until the current user turn is enqueued to avoid forwarding
|
||
// any pre-turn SDK chatter as fresh chat history.
|
||
let userMessageSent = false
|
||
|
||
rebindSessionOutput(sessionId, ws, {
|
||
shouldForward: (cliMsg) => userMessageSent || (cliMsg.type === 'result' && cliMsg.is_error),
|
||
})
|
||
|
||
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
|
||
|
||
// Switching to/from bypassPermissions requires the CLI to be (re)started with
|
||
// --dangerously-skip-permissions. The CLI rejects a runtime set_permission_mode
|
||
// to bypassPermissions if it wasn't launched with that flag. Rather than just
|
||
// sending the SDK message (which would silently fail), restart the CLI subprocess
|
||
// with the correct arguments so the new permission mode takes effect.
|
||
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 {
|
||
// Persist the new mode first so it's read on restart
|
||
await settingsService.setPermissionMode(mode)
|
||
|
||
const workDir = conversationService.getSessionWorkDir(sessionId)
|
||
conversationService.stopSession(sessionId)
|
||
|
||
// Rebuild runtime settings (will pick up the persisted mode)
|
||
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)) {
|
||
// First try graceful interrupt via SDK control message
|
||
conversationService.sendInterrupt(sessionId)
|
||
|
||
// Force-kill if still running after 3 seconds
|
||
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' })
|
||
}
|
||
|
||
// ============================================================================
|
||
// Title generation
|
||
// ============================================================================
|
||
|
||
function triggerTitleGeneration(ws: ServerWebSocket<WebSocketData>, sessionId: string): void {
|
||
const state = sessionTitleState.get(sessionId)
|
||
if (!state || state.hasCustomTitle) return
|
||
|
||
const count = state.userMessageCount
|
||
|
||
// Generate on count 1 (first response) and count 3 (with more context)
|
||
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
|
||
|
||
// Fire-and-forget: derive quick title, then upgrade with AI
|
||
void (async () => {
|
||
try {
|
||
// Stage 1: quick placeholder (only on first message)
|
||
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 })
|
||
}
|
||
}
|
||
|
||
// Stage 2: AI-generated title
|
||
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)
|
||
}
|
||
})()
|
||
}
|
||
|
||
// ============================================================================
|
||
// CLI message translation
|
||
// ============================================================================
|
||
|
||
/**
|
||
* Per-session streaming state to avoid cross-session interference.
|
||
* Each session tracks its own dedup flag, active block types, and tool blocks.
|
||
*/
|
||
type SessionStreamState = {
|
||
hasReceivedStreamEvents: boolean
|
||
activeBlockTypes: Map<number, 'text' | 'tool_use' | 'thinking'>
|
||
activeToolBlocks: Map<number, { toolName: string; toolUseId: string; inputJson: string }>
|
||
/** Tool blocks whose input JSON failed to parse in content_block_stop.
|
||
* The assistant message carries the complete input — defer to that. */
|
||
pendingToolBlocks: Map<string, { toolName: string; toolUseId: string; parentToolUseId?: 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(),
|
||
pendingToolBlocks: new Map(),
|
||
lastApiError: undefined,
|
||
}
|
||
sessionStreamStates.set(sessionId, state)
|
||
}
|
||
return state
|
||
}
|
||
|
||
/** Clean up stream state when session disconnects */
|
||
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 we already received stream_events, text/thinking were already sent.
|
||
// Only extract tool_use blocks (stream_event's content_block_stop lacks complete tool info).
|
||
if (cliMsg.message?.content && Array.isArray(cliMsg.message.content)) {
|
||
const messages: ServerMessage[] = []
|
||
|
||
for (const block of cliMsg.message.content) {
|
||
if (streamState.hasReceivedStreamEvents) {
|
||
// Stream events handled most blocks — but any tool_use whose
|
||
// input JSON failed to parse in content_block_stop was deferred.
|
||
// Emit those now with the complete input from the assistant message.
|
||
if (block.type === 'tool_use' && streamState.pendingToolBlocks.has(block.id)) {
|
||
const pending = streamState.pendingToolBlocks.get(block.id)!
|
||
streamState.pendingToolBlocks.delete(block.id)
|
||
messages.push({
|
||
type: 'tool_use_complete',
|
||
toolName: pending.toolName || block.name,
|
||
toolUseId: block.id,
|
||
input: block.input,
|
||
parentToolUseId: pending.parentToolUseId,
|
||
})
|
||
}
|
||
} else {
|
||
// No stream events received — this is the only source, process everything
|
||
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') {
|
||
messages.push({
|
||
type: 'tool_use_complete',
|
||
toolName: block.name,
|
||
toolUseId: block.id,
|
||
input: block.input,
|
||
parentToolUseId:
|
||
typeof cliMsg.parent_tool_use_id === 'string'
|
||
? cliMsg.parent_tool_use_id
|
||
: undefined,
|
||
})
|
||
}
|
||
}
|
||
}
|
||
|
||
// Reset flags for next turn
|
||
streamState.hasReceivedStreamEvents = false
|
||
streamState.pendingToolBlocks.clear()
|
||
return messages
|
||
}
|
||
return []
|
||
}
|
||
|
||
case 'user': {
|
||
// Bug #1: 处理 tool_result 消息
|
||
// CLI 发送 type:'user' 消息,其中 content 包含 tool_result 块
|
||
const messages: ServerMessage[] = []
|
||
|
||
const localCommandOutput = extractLocalCommandOutput(
|
||
cliMsg.message?.content,
|
||
)
|
||
if (localCommandOutput) {
|
||
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') {
|
||
messages.push({
|
||
type: 'tool_result',
|
||
toolUseId: block.tool_use_id,
|
||
content: block.content,
|
||
isError: !!block.is_error,
|
||
parentToolUseId:
|
||
typeof cliMsg.parent_tool_use_id === 'string'
|
||
? cliMsg.parent_tool_use_id
|
||
: undefined,
|
||
})
|
||
}
|
||
}
|
||
}
|
||
|
||
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') {
|
||
streamState.activeBlockTypes.set(index, 'tool_use')
|
||
// Track tool info so content_block_stop can emit complete data
|
||
streamState.activeToolBlocks.set(index, {
|
||
toolName: contentBlock.name || '',
|
||
toolUseId: contentBlock.id || '',
|
||
inputJson: '',
|
||
})
|
||
return [{
|
||
type: 'content_start',
|
||
blockType: 'tool_use',
|
||
toolName: contentBlock.name,
|
||
toolUseId: contentBlock.id,
|
||
parentToolUseId:
|
||
typeof cliMsg.parent_tool_use_id === 'string'
|
||
? cliMsg.parent_tool_use_id
|
||
: undefined,
|
||
}]
|
||
}
|
||
|
||
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) {
|
||
// Accumulate tool input 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 =
|
||
typeof cliMsg.parent_tool_use_id === 'string'
|
||
? cliMsg.parent_tool_use_id
|
||
: undefined
|
||
let parsedInput = null
|
||
try { parsedInput = JSON.parse(toolBlock.inputJson) } catch {}
|
||
|
||
if (parsedInput !== null) {
|
||
return [{
|
||
type: 'tool_use_complete',
|
||
toolName: toolBlock.toolName,
|
||
toolUseId: toolBlock.toolUseId,
|
||
input: parsedInput,
|
||
parentToolUseId,
|
||
}]
|
||
}
|
||
|
||
// JSON parse failed — defer to the assistant message which
|
||
// carries the complete, already-parsed tool input.
|
||
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': {
|
||
// message_stop is handled by the 'result' message
|
||
return []
|
||
}
|
||
|
||
case 'message_delta': {
|
||
// message_delta may contain stop_reason or usage updates
|
||
return []
|
||
}
|
||
|
||
default:
|
||
return []
|
||
}
|
||
}
|
||
|
||
case 'control_request': {
|
||
// 权限请求 — CLI 需要用户授权才能执行工具
|
||
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 the user requested stop, this "error" is just the interrupt
|
||
// result — don't show it as an error in the chat UI.
|
||
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 },
|
||
]
|
||
}
|
||
|
||
// Clear stop flag on successful completion too
|
||
sessionStopRequested.delete(sessionId)
|
||
streamState.lastApiError = undefined
|
||
return [{ type: 'message_complete', usage }]
|
||
}
|
||
|
||
case 'system': {
|
||
// 区分不同的 system 子类型
|
||
const subtype = cliMsg.subtype
|
||
if (subtype === 'init') {
|
||
// CLI 初始化完成 — 缓存 slash commands 并发送模型信息
|
||
// NOTE: Do NOT send status:idle here — the CLI init fires while
|
||
// processing the first user message, and sending idle would reset
|
||
// the frontend's streaming state prematurely.
|
||
cacheSessionInitMetadata(sessionId, cliMsg)
|
||
const messages: ServerMessage[] = [
|
||
// Send model info as a system notification, not a status change
|
||
{ type: 'system_notification', subtype: 'init', message: `Model: ${cliMsg.model || 'unknown'}`, data: { model: cliMsg.model } },
|
||
]
|
||
// Send slash commands to frontend
|
||
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 === 'hook_started' || subtype === 'hook_response') {
|
||
// Hook 执行中 — 不转发给前端
|
||
return []
|
||
}
|
||
if (subtype === 'local_command' || subtype === 'local_command_output') {
|
||
const localCommandOutput = extractLocalCommandOutput(
|
||
cliMsg.content ?? cliMsg.message,
|
||
{ allowUntagged: subtype === 'local_command_output' },
|
||
)
|
||
if (!localCommandOutput) return []
|
||
return [
|
||
{ type: 'content_start', blockType: 'text' },
|
||
{ type: 'content_delta', text: localCommandOutput },
|
||
]
|
||
}
|
||
// Bug #7: 处理 task/team system 消息
|
||
if (subtype === 'task_notification') {
|
||
return [{
|
||
type: 'system_notification',
|
||
subtype: 'task_notification',
|
||
message: cliMsg.message || cliMsg.title,
|
||
data: cliMsg,
|
||
}]
|
||
}
|
||
if (subtype === 'task_started') {
|
||
return [{
|
||
type: 'status',
|
||
state: 'tool_executing',
|
||
verb: cliMsg.message || 'Task started',
|
||
}]
|
||
}
|
||
if (subtype === 'task_progress') {
|
||
return [{
|
||
type: 'status',
|
||
state: 'tool_executing',
|
||
verb: cliMsg.message || '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,
|
||
}]
|
||
}
|
||
// 其他 system 消息
|
||
return []
|
||
}
|
||
|
||
default:
|
||
// 未知类型 — 调试输出但不转发
|
||
console.log(`[WS] Unknown CLI message type: ${cliMsg.type}`, JSON.stringify(cliMsg).substring(0, 200))
|
||
return []
|
||
}
|
||
}
|
||
|
||
// ============================================================================
|
||
// Helpers
|
||
// ============================================================================
|
||
|
||
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 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 extractTaggedContent(raw: string, tag: string): string | null {
|
||
const match = raw.match(new RegExp(`<${tag}>([\\s\\S]*?)</${tag}>`))
|
||
return match?.[1]?.trim() ?? null
|
||
}
|
||
|
||
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 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
|
||
}
|
||
|
||
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 = providers.some((provider) => provider.id === runtimeOverride.providerId)
|
||
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> {
|
||
// Check if a custom provider is active
|
||
const { providers, activeId } = await providerService.listProviders()
|
||
let resolvedActiveId = activeId
|
||
if (activeId && !providers.some((provider) => provider.id === activeId)) {
|
||
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) {
|
||
// Provider is active — only consult provider-managed cc-haha settings.
|
||
// Global ~/.claude/settings.json model values must not bleed into provider mode.
|
||
const baseModel =
|
||
typeof modelSettings.model === 'string' && modelSettings.model.trim()
|
||
? modelSettings.model
|
||
: ''
|
||
if (baseModel) {
|
||
model = baseModel
|
||
if (modelContext) model += `:${modelContext}`
|
||
}
|
||
} else {
|
||
// No provider — pass model normally
|
||
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
|
||
}
|
||
|
||
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 }
|
||
}
|
||
|
||
/**
|
||
* Send a message to a specific session's WebSocket (for use by services)
|
||
*/
|
||
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()
|
||
}
|