diff --git a/src/cli.ts b/src/cli.ts index 22bf1f8..3b403e0 100755 --- a/src/cli.ts +++ b/src/cli.ts @@ -60,7 +60,9 @@ async function main() { let pty: IPty | null = null; let ws: WebSocket | null = null; let isExiting = false; + let isAuthenticated = false; let reconnectTimer: Timer | null = null; + const disposables: Array<{ dispose: () => void }> = []; const cleanup = () => { if (isExiting) return; @@ -69,6 +71,12 @@ async function main() { if (reconnectTimer) { clearTimeout(reconnectTimer); } + for (const d of disposables) { + d.dispose(); + } + if (process.stdin.isTTY) { + process.stdin.setRawMode(false); + } if (pty) { pty.kill(); } @@ -102,19 +110,28 @@ async function main() { } }); - // Forward PTY output to stdout - pty.onData((data: string) => { - process.stdout.write(data); - }); + // Forward PTY output to stdout AND WebSocket (if authenticated) + disposables.push( + pty.onData((data: string) => { + process.stdout.write(data); + + if (ws && ws.readyState === WebSocket.OPEN && isAuthenticated) { + const msg: ClientMessage = { type: "output", data }; + ws.send(JSON.stringify(msg)); + } + }), + ); // Handle PTY exit - pty.onExit((event) => { - if (ws && !isExiting) { - const msg: ClientMessage = { type: "exit", code: event.exitCode }; - ws.send(JSON.stringify(msg)); - } - cleanup(); - }); + disposables.push( + pty.onExit((event) => { + if (ws && !isExiting) { + const msg: ClientMessage = { type: "exit", code: event.exitCode }; + ws.send(JSON.stringify(msg)); + } + cleanup(); + }), + ); // Handle terminal resize process.stdout.on("resize", () => { @@ -154,15 +171,20 @@ async function main() { }; ws.onmessage = (event) => { - if (!pty) return; - try { const msg: ServerMessage = JSON.parse(event.data); - if (msg.type === "input") { - pty.write(msg.data); + if (msg.type === "authenticated") { + isAuthenticated = true; + console.debug(`Authenticated with session ID: ${msg.session_id}`); + } else if (msg.type === "input") { + if (pty) { + pty.write(msg.data); + } } else if (msg.type === "resize") { - pty.resize(msg.cols, msg.rows); + if (pty) { + pty.resize(msg.cols, msg.rows); + } } else if (msg.type === "ping") { // Acknowledge ping (keep-alive) } @@ -171,13 +193,16 @@ async function main() { } }; - ws.onerror = () => { - console.error("WebSocket error"); + ws.onerror = (event) => { + console.error("WebSocket error:", event); }; ws.onclose = () => { if (isExiting) return; + // Reset auth flag on disconnect + isAuthenticated = false; + // Try to reconnect after 2 seconds reconnectTimer = setTimeout(() => { console.debug("Reconnecting to server..."); @@ -186,16 +211,6 @@ async function main() { }; }; - // Forward PTY output to server - if (pty) { - pty.onData((data: string) => { - if (ws && ws.readyState === WebSocket.OPEN) { - const msg: ClientMessage = { type: "output", data }; - ws.send(JSON.stringify(msg)); - } - }); - } - connect(); } diff --git a/src/types.ts b/src/types.ts index e35be2a..eca8ef3 100644 --- a/src/types.ts +++ b/src/types.ts @@ -45,6 +45,7 @@ export type ClientMessage = | { type: "prompt"; session_id: number; prompt_text: string }; export type ServerMessage = + | { type: "authenticated"; session_id: number } | { type: "input"; data: string } | { type: "resize"; cols: number; rows: number } | { type: "ping" };