File size: 5,560 Bytes
5e4b249 4e105ef 9a3a32e 5e4b249 4e105ef 5e4b249 4e105ef 5e4b249 4e105ef 9a3a32e 4e105ef 9a3a32e 4e105ef 5e4b249 965bbbc 5e4b249 4e105ef 5e4b249 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | import { Server as SocketServer, Socket } from 'socket.io';
import { WS_EVENTS } from './events.js';
import { asi1 } from '../services/asi1-client.js';
import { orchestrator } from '../agents/orchestrator.js';
import { logger } from '../utils/logger.js';
import { handleCodeExecution, sendProcessInput, cancelGraceKill, scheduleGraceKill } from './execution.js';
import { getStableSessionId } from './session-id.js';
import type { WSCompletionRequest, WSReviewRequest, WSChatMessage, AgentType, WSExecuteRequest } from '@asipilot/shared';
const activeAbortControllers = new Map<string, AbortController>();
export function setupWebSocketHandlers(io: SocketServer) {
io.on('connection', (socket: Socket) => {
const sessionId = getStableSessionId(socket);
socket.join(sessionId);
cancelGraceKill(sessionId);
logger.info(`WebSocket connected: ${socket.id} (session: ${sessionId})`);
// Code completion
socket.on(WS_EVENTS.COMPLETION_REQUEST, async (data: WSCompletionRequest) => {
try {
const contextStr = data.contextFiles
?.slice(0, 3)
.map(f => f.content.split('\n').slice(0, 100).join('\n'))
.join('\n\n') || '';
const suggestion = await asi1.completeCode(data.prefix, data.suffix, data.language, contextStr);
socket.emit(WS_EVENTS.COMPLETION_RESULT, {
suggestion,
filePath: data.filePath,
cursorOffset: data.prefix.length,
});
} catch (error) {
logger.error('Completion error', { error: (error as Error).message });
socket.emit(WS_EVENTS.COMPLETION_RESULT, { suggestion: '', filePath: data.filePath });
}
});
// Full review
socket.on(WS_EVENTS.REVIEW_START, async (data: WSReviewRequest) => {
try {
const files = new Map(Object.entries(data.files));
const agents = data.agents as AgentType[] | undefined;
orchestrator.on('progress', (event) => {
socket.emit(WS_EVENTS.REVIEW_PROGRESS, event);
});
const result = await orchestrator.reviewCodebase(files, { agents });
socket.emit(WS_EVENTS.REVIEW_COMPLETE, result);
orchestrator.removeAllListeners('progress');
} catch (error) {
logger.error('Review error', { error: (error as Error).message });
socket.emit(WS_EVENTS.REVIEW_COMPLETE, { error: (error as Error).message });
orchestrator.removeAllListeners('progress');
}
});
// Single file review
socket.on(WS_EVENTS.REVIEW_FILE, async (data: { filePath: string; content: string }) => {
try {
const result = await orchestrator.reviewSingleFile(data.filePath, data.content);
socket.emit(WS_EVENTS.REVIEW_FILE_RESULT, result);
} catch (error) {
logger.error('File review error', { error: (error as Error).message });
socket.emit(WS_EVENTS.REVIEW_FILE_RESULT, { error: (error as Error).message });
}
});
// Code execution (use stable sessionId so stdin/output stay matched after socket.io reconnects)
socket.on(WS_EVENTS.EXECUTE_REQUEST, async (data: WSExecuteRequest) => {
await handleCodeExecution(io, sessionId, data);
});
socket.on(
WS_EVENTS.EXECUTE_INPUT,
(data: { runId: string; input: string }, ack?: (r: { ok: boolean }) => void) => {
logger.debug(`EXECUTE_INPUT run ${data.runId}: ${JSON.stringify(data.input)}`);
const success = sendProcessInput(data.runId, data.input);
if (typeof ack === 'function') ack({ ok: success });
if (!success) {
logger.warn(`Failed stdin for runId ${data.runId}`);
}
}
);
// Chat with streaming
socket.on(WS_EVENTS.CHAT_MESSAGE, async (data: WSChatMessage) => {
const abortController = new AbortController();
activeAbortControllers.set(socket.id, abortController);
try {
const messages = [
{
role: 'system' as const,
content: 'You are AsiPilot AI, an expert frontend development assistant. You help with HTML, CSS, JavaScript, TypeScript, React, Vue, Svelte, Next.js, Tailwind CSS, and all frontend technologies. Provide clear, concise, and accurate code with explanations. When showing code, use appropriate language tags in code blocks.',
},
...data.history,
{ role: 'user' as const, content: data.message },
];
const stream = asi1.chatStream(messages, { temperature: 0.7 }, abortController.signal);
for await (const token of stream) {
socket.emit(WS_EVENTS.CHAT_TOKEN, { token });
}
socket.emit(WS_EVENTS.CHAT_COMPLETE, {});
} catch (error) {
if (!abortController.signal.aborted) {
logger.error('Chat error', { error: (error as Error).message });
socket.emit(WS_EVENTS.CHAT_ERROR, { error: (error as Error).message });
}
} finally {
activeAbortControllers.delete(socket.id);
}
});
// Stop chat generation
socket.on(WS_EVENTS.CHAT_STOP, () => {
const controller = activeAbortControllers.get(socket.id);
if (controller) {
controller.abort();
activeAbortControllers.delete(socket.id);
}
});
socket.on('disconnect', () => {
const controller = activeAbortControllers.get(socket.id);
if (controller) controller.abort();
activeAbortControllers.delete(socket.id);
scheduleGraceKill(sessionId);
logger.info(`WebSocket disconnected: ${socket.id} (session: ${sessionId})`);
});
});
}
|