ws-handler.ts
1import type { Server as HttpServer } from 'http';
2import { WebSocketServer, type WebSocket } from 'ws';
3import {
4 startSession,
5 sendInput,
6 sendReviewPrompt,
7 resize,
8 killSession,
9 isAlive,
10 getBufferedOutput,
11 onData,
12 onExit
13} from './pty-session.ts';
14
15export function attachWebSocket(server: HttpServer, appPort?: number): void {
16 const wss = new WebSocketServer({ noServer: true });
17
18 server.on('upgrade', (req, socket, head) => {
19 if (req.url === '/ws/terminal') {
20 wss.handleUpgrade(req, socket, head, (ws) => {
21 wss.emit('connection', ws, req);
22 });
23 }
24 // Don't destroy other upgrade requests — Vite HMR uses WebSocket too
25 });
26
27 wss.on('connection', async (ws: WebSocket) => {
28 // Don't start session yet — wait for first resize message with correct dimensions.
29 // If session is already running, send buffered output immediately.
30 const buffered = isAlive() ? getBufferedOutput() : null;
31 if (buffered) {
32 ws.send(JSON.stringify({ type: 'output', data: buffered }));
33 }
34
35 // Forward PTY output to WebSocket
36 const unsubData = onData((data) => {
37 if (ws.readyState === ws.OPEN) {
38 ws.send(JSON.stringify({ type: 'output', data }));
39 }
40 });
41
42 const unsubExit = onExit((code) => {
43 if (ws.readyState === ws.OPEN) {
44 ws.send(JSON.stringify({ type: 'exit', code }));
45 }
46 });
47
48 // Handle messages from client
49 ws.on('message', async (raw) => {
50 try {
51 const msg = JSON.parse(raw.toString());
52
53 switch (msg.type) {
54 case 'input':
55 sendInput(msg.data);
56 break;
57
58 case 'review': {
59 // Fetch the review prompt via HTTP to ensure we use the same
60 // module instance as the SvelteKit API routes
61 const port = appPort || (() => { const a = server.address(); return typeof a === 'object' && a ? a.port : 5173; })();
62 try {
63 const res = await fetch(`http://localhost:${port}/api/export`, {
64 method: 'POST',
65 headers: { 'Content-Type': 'application/json' },
66 body: JSON.stringify({ format: 'clipboard' })
67 });
68 const result = await res.json();
69 if (!result.success || !result.text || result.text === 'No review comments.') {
70 ws.send(JSON.stringify({
71 type: 'output',
72 data: '\r\n\x1b[33mâš No review comments to send. Add comments first.\x1b[0m\r\n'
73 }));
74 break;
75 }
76 if (!isAlive()) {
77 // Session should already be started by resize on connect.
78 // If somehow not, start with defaults and wait.
79 await startSession(80, 24);
80 await new Promise((r) => setTimeout(r, 2000));
81 }
82 sendReviewPrompt(result.text);
83 } catch (err) {
84 ws.send(JSON.stringify({
85 type: 'output',
86 data: `\r\n\x1b[31mFailed to fetch review: ${err}\x1b[0m\r\n`
87 }));
88 }
89 break;
90 }
91
92 case 'resize':
93 if (msg.cols && msg.rows) {
94 if (!isAlive()) {
95 // First resize — start session with correct dimensions
96 try {
97 await startSession(msg.cols, msg.rows);
98 const buf = getBufferedOutput();
99 if (buf) ws.send(JSON.stringify({ type: 'output', data: buf }));
100 } catch (err) {
101 ws.send(JSON.stringify({
102 type: 'output',
103 data: `\r\n\x1b[31mFailed to start terminal: ${err}\x1b[0m\r\n`
104 }));
105 }
106 } else {
107 resize(msg.cols, msg.rows);
108 }
109 }
110 break;
111
112 case 'stop':
113 killSession();
114 break;
115
116 case 'restart':
117 killSession();
118 // Ask client to send dimensions so we restart with correct size
119 ws.send(JSON.stringify({ type: 'request-resize' }));
120 break;
121 }
122 } catch (err) {
123 console.error('WS message error:', err);
124 }
125 });
126
127 ws.on('close', () => {
128 unsubData();
129 unsubExit();
130 });
131 });
132}
