mirror of
https://github.com/siteboon/claudecodeui.git
synced 2026-05-28 23:15:33 +08:00
refactor(websocket): move websocket logic to its own module
This commit is contained in:
Haileyesus
271
server/modules/websocket/services/chat-websocket.service.ts
Normal file
271
server/modules/websocket/services/chat-websocket.service.ts
Normal file
@@ -0,0 +1,271 @@
|
||||
import type { WebSocket } from 'ws';
|
||||
|
||||
import { connectedClients } from '@/modules/websocket/services/websocket-state.service.js';
|
||||
import { WebSocketWriter } from '@/modules/websocket/services/websocket-writer.service.js';
|
||||
import type {
|
||||
AnyRecord,
|
||||
AuthenticatedWebSocketRequest,
|
||||
LLMProvider,
|
||||
} from '@/shared/types.js';
|
||||
import { createNormalizedMessage, parseIncomingJsonObject } from '@/shared/utils.js';
|
||||
|
||||
type ChatIncomingMessage = AnyRecord & {
|
||||
type?: string;
|
||||
command?: string;
|
||||
options?: AnyRecord;
|
||||
provider?: string;
|
||||
sessionId?: string;
|
||||
requestId?: string;
|
||||
allow?: unknown;
|
||||
updatedInput?: unknown;
|
||||
message?: unknown;
|
||||
rememberEntry?: unknown;
|
||||
};
|
||||
|
||||
const DEFAULT_PROVIDER: LLMProvider = 'claude';
|
||||
|
||||
type ChatWebSocketDependencies = {
|
||||
queryClaudeSDK: (command: string, options: unknown, writer: WebSocketWriter) => Promise<unknown>;
|
||||
spawnCursor: (command: string, options: unknown, writer: WebSocketWriter) => Promise<unknown>;
|
||||
queryCodex: (command: string, options: unknown, writer: WebSocketWriter) => Promise<unknown>;
|
||||
spawnGemini: (command: string, options: unknown, writer: WebSocketWriter) => Promise<unknown>;
|
||||
abortClaudeSDKSession: (sessionId: string) => Promise<boolean>;
|
||||
abortCursorSession: (sessionId: string) => boolean;
|
||||
abortCodexSession: (sessionId: string) => boolean;
|
||||
abortGeminiSession: (sessionId: string) => boolean;
|
||||
resolveToolApproval: (
|
||||
requestId: string,
|
||||
payload: {
|
||||
allow: boolean;
|
||||
updatedInput?: unknown;
|
||||
message?: string;
|
||||
rememberEntry?: unknown;
|
||||
}
|
||||
) => void;
|
||||
isClaudeSDKSessionActive: (sessionId: string) => boolean;
|
||||
isCursorSessionActive: (sessionId: string) => boolean;
|
||||
isCodexSessionActive: (sessionId: string) => boolean;
|
||||
isGeminiSessionActive: (sessionId: string) => boolean;
|
||||
reconnectSessionWriter: (sessionId: string, ws: WebSocket) => boolean;
|
||||
getPendingApprovalsForSession: (sessionId: string) => unknown[];
|
||||
getActiveClaudeSDKSessions: () => unknown;
|
||||
getActiveCursorSessions: () => unknown;
|
||||
getActiveCodexSessions: () => unknown;
|
||||
getActiveGeminiSessions: () => unknown;
|
||||
};
|
||||
|
||||
/**
|
||||
* Normalizes potentially invalid provider names coming from websocket payloads.
|
||||
*/
|
||||
function readProvider(value: unknown): LLMProvider {
|
||||
if (value === 'claude' || value === 'cursor' || value === 'codex' || value === 'gemini') {
|
||||
return value;
|
||||
}
|
||||
|
||||
return DEFAULT_PROVIDER;
|
||||
}
|
||||
|
||||
/**
|
||||
* Extracts the authenticated request user id in the formats currently produced
|
||||
* by platform and OSS auth code paths.
|
||||
*/
|
||||
function readRequestUserId(
|
||||
request: AuthenticatedWebSocketRequest | undefined
|
||||
): string | number | null {
|
||||
const user = request?.user;
|
||||
if (!user) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (typeof user.id === 'string' || typeof user.id === 'number') {
|
||||
return user.id;
|
||||
}
|
||||
|
||||
if (typeof user.userId === 'string' || typeof user.userId === 'number') {
|
||||
return user.userId;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles authenticated chat websocket messages used by the main chat panel.
|
||||
*/
|
||||
export function handleChatConnection(
|
||||
ws: WebSocket,
|
||||
request: AuthenticatedWebSocketRequest,
|
||||
dependencies: ChatWebSocketDependencies
|
||||
): void {
|
||||
console.log('[INFO] Chat WebSocket connected');
|
||||
connectedClients.add(ws);
|
||||
|
||||
const writer = new WebSocketWriter(ws, readRequestUserId(request));
|
||||
|
||||
ws.on('message', async (rawMessage) => {
|
||||
try {
|
||||
const parsed = parseIncomingJsonObject(rawMessage);
|
||||
if (!parsed) {
|
||||
throw new Error('Invalid websocket payload');
|
||||
}
|
||||
|
||||
const data = parsed as ChatIncomingMessage;
|
||||
const messageType = data.type;
|
||||
if (!messageType) {
|
||||
throw new Error('Message type is required');
|
||||
}
|
||||
|
||||
if (messageType === 'claude-command') {
|
||||
await dependencies.queryClaudeSDK(data.command ?? '', data.options, writer);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'cursor-command') {
|
||||
await dependencies.spawnCursor(data.command ?? '', data.options, writer);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'codex-command') {
|
||||
await dependencies.queryCodex(data.command ?? '', data.options, writer);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'gemini-command') {
|
||||
await dependencies.spawnGemini(data.command ?? '', data.options, writer);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'cursor-resume') {
|
||||
await dependencies.spawnCursor(
|
||||
'',
|
||||
{
|
||||
sessionId: data.sessionId,
|
||||
resume: true,
|
||||
cwd: data.options?.cwd,
|
||||
},
|
||||
writer
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'abort-session') {
|
||||
const provider = readProvider(data.provider);
|
||||
const sessionId = typeof data.sessionId === 'string' ? data.sessionId : '';
|
||||
let success = false;
|
||||
|
||||
if (provider === 'cursor') {
|
||||
success = dependencies.abortCursorSession(sessionId);
|
||||
} else if (provider === 'codex') {
|
||||
success = dependencies.abortCodexSession(sessionId);
|
||||
} else if (provider === 'gemini') {
|
||||
success = dependencies.abortGeminiSession(sessionId);
|
||||
} else {
|
||||
success = await dependencies.abortClaudeSDKSession(sessionId);
|
||||
}
|
||||
|
||||
writer.send(
|
||||
createNormalizedMessage({
|
||||
kind: 'complete',
|
||||
exitCode: success ? 0 : 1,
|
||||
aborted: true,
|
||||
success,
|
||||
sessionId,
|
||||
provider,
|
||||
})
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'claude-permission-response') {
|
||||
if (typeof data.requestId === 'string' && data.requestId.length > 0) {
|
||||
dependencies.resolveToolApproval(data.requestId, {
|
||||
allow: Boolean(data.allow),
|
||||
updatedInput: data.updatedInput,
|
||||
message: typeof data.message === 'string' ? data.message : undefined,
|
||||
rememberEntry: data.rememberEntry,
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'cursor-abort') {
|
||||
const sessionId = typeof data.sessionId === 'string' ? data.sessionId : '';
|
||||
const success = dependencies.abortCursorSession(sessionId);
|
||||
writer.send(
|
||||
createNormalizedMessage({
|
||||
kind: 'complete',
|
||||
exitCode: success ? 0 : 1,
|
||||
aborted: true,
|
||||
success,
|
||||
sessionId,
|
||||
provider: 'cursor',
|
||||
})
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'check-session-status') {
|
||||
const provider = readProvider(data.provider);
|
||||
const sessionId = typeof data.sessionId === 'string' ? data.sessionId : '';
|
||||
let isActive = false;
|
||||
|
||||
if (provider === 'cursor') {
|
||||
isActive = dependencies.isCursorSessionActive(sessionId);
|
||||
} else if (provider === 'codex') {
|
||||
isActive = dependencies.isCodexSessionActive(sessionId);
|
||||
} else if (provider === 'gemini') {
|
||||
isActive = dependencies.isGeminiSessionActive(sessionId);
|
||||
} else {
|
||||
isActive = dependencies.isClaudeSDKSessionActive(sessionId);
|
||||
if (isActive) {
|
||||
dependencies.reconnectSessionWriter(sessionId, ws);
|
||||
}
|
||||
}
|
||||
|
||||
writer.send({
|
||||
type: 'session-status',
|
||||
sessionId,
|
||||
provider,
|
||||
isProcessing: isActive,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'get-pending-permissions') {
|
||||
const sessionId = typeof data.sessionId === 'string' ? data.sessionId : '';
|
||||
if (sessionId && dependencies.isClaudeSDKSessionActive(sessionId)) {
|
||||
const pending = dependencies.getPendingApprovalsForSession(sessionId);
|
||||
writer.send({
|
||||
type: 'pending-permissions-response',
|
||||
sessionId,
|
||||
data: pending,
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (messageType === 'get-active-sessions') {
|
||||
writer.send({
|
||||
type: 'active-sessions',
|
||||
sessions: {
|
||||
claude: dependencies.getActiveClaudeSDKSessions(),
|
||||
cursor: dependencies.getActiveCursorSessions(),
|
||||
codex: dependencies.getActiveCodexSessions(),
|
||||
gemini: dependencies.getActiveGeminiSessions(),
|
||||
},
|
||||
});
|
||||
}
|
||||
} catch (error) {
|
||||
const message = error instanceof Error ? error.message : String(error);
|
||||
console.error('[ERROR] Chat WebSocket error:', message);
|
||||
writer.send({
|
||||
type: 'error',
|
||||
error: message,
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
ws.on('close', () => {
|
||||
console.log('[INFO] Chat client disconnected');
|
||||
connectedClients.delete(ws);
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user