Fix CLI resource cleanup and auth flow

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
+  // 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
+  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") {
+        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") {
+          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();
 }
 

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" };