| 1 | import asyncio, os, sys, platform, errno, signal |
| 2 | |
| 3 | _IS_WIN = platform.system() == "Windows" |
| 4 | if _IS_WIN: |
| 5 | import winpty # pip install pywinpty # type: ignore |
| 6 | import msvcrt |
| 7 | |
| 8 | _CLOSE_TIMEOUT_SECONDS = 2 |
| 9 | |
| 10 | |
| 11 | def _reconfigure_stream_errors(stream) -> None: |
| 12 | reconfigure = getattr(stream, "reconfigure", None) |
| 13 | if not callable(reconfigure): |
| 14 | return |
| 15 | reconfigure(errors="replace") |
| 16 | |
| 17 | |
| 18 | # Make stdin / stdout tolerant to broken UTF-8 so input() never aborts |
| 19 | _reconfigure_stream_errors(sys.stdin) |
| 20 | _reconfigure_stream_errors(sys.stdout) |
| 21 | |
| 22 | |
| 23 | # ──────────────────────────── PUBLIC CLASS ──────────────────────────── |
| 24 | |
| 25 | |
| 26 | class TTYSession: |
| 27 | def __init__(self, cmd, *, cwd=None, env=None, encoding="utf-8", echo=False): |
| 28 | self.cmd = cmd if isinstance(cmd, str) else " ".join(cmd) |
| 29 | self.cwd = cwd |
| 30 | self.env = env or os.environ.copy() |
| 31 | self.encoding = encoding |
| 32 | self.echo = echo # ← store preference |
| 33 | self._proc = None |
| 34 | self._buf: asyncio.Queue = None # type: ignore |
| 35 | self._pump_task = None |
| 36 | self._pty_master = None |
| 37 | self._pty_master_ref = None |
| 38 | |
| 39 | def __del__(self): |
| 40 | # Simple cleanup on object destruction |
| 41 | import nest_asyncio |
| 42 | |
| 43 | nest_asyncio.apply() |
| 44 | if hasattr(self, "close"): |
| 45 | try: |
| 46 | asyncio.run(self.close()) |
| 47 | except Exception: |
| 48 | pass |
| 49 | |
| 50 | # ── user-facing coroutines ──────────────────────────────────────── |
| 51 | async def start(self): |
| 52 | self._buf = asyncio.Queue() |
| 53 | if _IS_WIN: |
| 54 | self._proc = await _spawn_winpty( |
| 55 | self.cmd, self.cwd, self.env, self.echo |
| 56 | ) # ← pass echo |
| 57 | else: |
| 58 | self._proc = await _spawn_posix_pty( |
| 59 | self.cmd, self.cwd, self.env, self.echo |
| 60 | ) # ← pass echo |
| 61 | self._pty_master_ref = getattr(self._proc, "_pty_master_ref", None) |
| 62 | self._pty_master = ( |
| 63 | self._pty_master_ref.get("fd") |
| 64 | if self._pty_master_ref is not None |
| 65 | else getattr(self._proc, "_pty_master", None) |
| 66 | ) |
| 67 | self._pump_task = asyncio.create_task(self._pump_stdout()) |
| 68 | |
| 69 | async def close(self): |
| 70 | # Cancel the pump task if it exists |
| 71 | if self._pump_task: |
| 72 | self._pump_task.cancel() |
| 73 | try: |
| 74 | await self._pump_task |
| 75 | except asyncio.CancelledError: |
| 76 | pass |
| 77 | except Exception: |
| 78 | pass |
| 79 | |
| 80 | # Terminate the process if it exists |
| 81 | if self._proc: |
| 82 | if getattr(self._proc, "returncode", None) is None: |
| 83 | self._signal_process(signal.SIGTERM) |
| 84 | try: |
| 85 | await asyncio.wait_for(self._proc.wait(), _CLOSE_TIMEOUT_SECONDS) |
| 86 | except asyncio.TimeoutError: |
| 87 | self._signal_process(signal.SIGKILL) |
| 88 | try: |
| 89 | await asyncio.wait_for(self._proc.wait(), _CLOSE_TIMEOUT_SECONDS) |
| 90 | except Exception: |
| 91 | pass |
| 92 | except Exception: |
| 93 | pass |
| 94 | |
| 95 | self._release_pty_master() |
| 96 | self._proc = None |
| 97 | self._pump_task = None |
| 98 | |
| 99 | def _signal_process(self, sig): |
| 100 | if self._proc is None: |
| 101 | return |
| 102 | try: |
| 103 | if _IS_WIN: |
| 104 | if sig == signal.SIGKILL: |
| 105 | self._proc.kill() |
| 106 | else: |
| 107 | self._proc.terminate() |
| 108 | return |
| 109 | os.killpg(self._proc.pid, sig) |
| 110 | except ProcessLookupError: |
| 111 | pass |
| 112 | except Exception: |
| 113 | try: |
| 114 | if sig == signal.SIGKILL: |
| 115 | self._proc.kill() |
| 116 | else: |
| 117 | self._proc.terminate() |
| 118 | except Exception: |
| 119 | pass |
| 120 | |
| 121 | def _release_pty_master(self): |
| 122 | """Release the POSIX PTY master exactly once. |
| 123 | |
| 124 | The fd number is invalidated before os.close() so that a concurrent or |
| 125 | later cleanup path cannot close the same integer after the OS has reused |
| 126 | it for another file/socket. |
| 127 | """ |
| 128 | ref = self._pty_master_ref |
| 129 | master = ref.get("fd") if ref is not None else self._pty_master |
| 130 | if master is None: |
| 131 | self._pty_master = None |
| 132 | return |
| 133 | if ref is not None: |
| 134 | ref["fd"] = None |
| 135 | self._pty_master = None |
| 136 | try: |
| 137 | loop = asyncio.get_running_loop() |
| 138 | loop.remove_reader(master) |
| 139 | except Exception: |
| 140 | pass |
| 141 | try: |
| 142 | os.close(master) |
| 143 | except OSError: |
| 144 | pass |
| 145 | self._pty_master_ref = None |
| 146 | |
| 147 | async def send(self, data: str | bytes): |
| 148 | if self._proc is None: |
| 149 | raise RuntimeError("TTYSpawn is not started") |
| 150 | if not _IS_WIN: |
| 151 | master = ( |
| 152 | self._pty_master_ref.get("fd") |
| 153 | if self._pty_master_ref is not None |
| 154 | else self._pty_master |
| 155 | ) |
| 156 | if master is None: |
| 157 | raise RuntimeError("TTYSpawn PTY is closed") |
| 158 | if getattr(self._proc, "returncode", None) is not None: |
| 159 | raise RuntimeError("TTYSpawn process has exited") |
| 160 | if isinstance(data, str): |
| 161 | data = data.encode(self.encoding) |
| 162 | try: |
| 163 | self._proc.stdin.write(data) # type: ignore |
| 164 | await self._proc.stdin.drain() # type: ignore |
| 165 | except OSError as e: |
| 166 | if e.errno in (errno.EBADF, errno.EIO, errno.EINVAL): |
| 167 | self._release_pty_master() |
| 168 | raise RuntimeError("TTYSpawn PTY is closed") from e |
| 169 | raise |
| 170 | |
| 171 | async def sendline(self, line: str): |
| 172 | await self.send(line + "\n") |
| 173 | |
| 174 | async def wait(self): |
| 175 | if self._proc is None: |
| 176 | raise RuntimeError("TTYSpawn is not started") |
| 177 | return await self._proc.wait() |
| 178 | |
| 179 | def is_terminated(self) -> bool: |
| 180 | """Return whether the managed shell process has exited.""" |
| 181 | return self._proc is None or getattr(self._proc, "returncode", None) is not None |
| 182 | |
| 183 | def get_exit_code(self) -> int | None: |
| 184 | """Return the managed shell exit code when it is already available.""" |
| 185 | if self._proc is None: |
| 186 | return None |
| 187 | return getattr(self._proc, "returncode", None) |
| 188 | |
| 189 | def kill(self): |
| 190 | """Force-kill the running child process. |
| 191 | |
| 192 | This is best-effort: if the process has already terminated (which can |
| 193 | happen if *close()* was called elsewhere or the child exited by |
| 194 | itself) we silently ignore the *ProcessLookupError* raised by |
| 195 | *asyncio.subprocess.Process.kill()*. This prevents race conditions |
| 196 | where multiple coroutines attempt to close the same session. |
| 197 | """ |
| 198 | if self._proc is None: |
| 199 | # Already closed or never started – nothing to do |
| 200 | return |
| 201 | |
| 202 | # Only attempt to kill if the process is still running |
| 203 | if getattr(self._proc, "returncode", None) is None: |
| 204 | self._signal_process(signal.SIGKILL) |
| 205 | self._release_pty_master() |
| 206 | |
| 207 | async def read(self, timeout=None): |
| 208 | # Return any decoded text the child produced, or None on timeout |
| 209 | try: |
| 210 | return await asyncio.wait_for(self._buf.get(), timeout) |
| 211 | except asyncio.TimeoutError: |
| 212 | return None |
| 213 | |
| 214 | # backward-compat alias: |
| 215 | readline = read |
| 216 | |
| 217 | async def read_full_until_idle(self, idle_timeout, total_timeout): |
| 218 | # Collect child output using iter_until_idle to avoid duplicate logic |
| 219 | return "".join( |
| 220 | [ |
| 221 | chunk |
| 222 | async for chunk in self.read_chunks_until_idle( |
| 223 | idle_timeout, total_timeout |
| 224 | ) |
| 225 | ] |
| 226 | ) |
| 227 | |
| 228 | async def read_chunks_until_idle(self, idle_timeout, total_timeout): |
| 229 | # Yield each chunk as soon as it arrives until idle or total timeout |
| 230 | import time |
| 231 | |
| 232 | start = time.monotonic() |
| 233 | while True: |
| 234 | if time.monotonic() - start > total_timeout: |
| 235 | break |
| 236 | chunk = await self.read(timeout=idle_timeout) |
| 237 | if chunk is None: |
| 238 | break |
| 239 | yield chunk |
| 240 | |
| 241 | # ── internal: stream raw output into the queue ──────────────────── |
| 242 | async def _pump_stdout(self): |
| 243 | if self._proc is None: |
| 244 | raise RuntimeError("TTYSpawn is not started") |
| 245 | reader = self._proc.stdout |
| 246 | while True: |
| 247 | chunk = await reader.read(4096) # grab whatever is ready # type: ignore |
| 248 | if not chunk: |
| 249 | break |
| 250 | self._buf.put_nowait(chunk.decode(self.encoding, "replace")) |
| 251 | |
| 252 | |
| 253 | # ──────────────────────────── POSIX IMPLEMENTATION ──────────────────── |
| 254 | |
| 255 | |
| 256 | async def _spawn_posix_pty(cmd, cwd, env, echo): |
| 257 | import pty, asyncio, os, termios |
| 258 | |
| 259 | master, slave = pty.openpty() |
| 260 | |
| 261 | # ── Disable ECHO on the slave side if requested ── |
| 262 | if not echo: |
| 263 | attrs = termios.tcgetattr(slave) |
| 264 | attrs[3] &= ~termios.ECHO # lflag |
| 265 | termios.tcsetattr(slave, termios.TCSANOW, attrs) |
| 266 | |
| 267 | proc = await asyncio.create_subprocess_shell( |
| 268 | cmd, |
| 269 | stdin=slave, |
| 270 | stdout=slave, |
| 271 | stderr=slave, |
| 272 | cwd=cwd, |
| 273 | env=env, |
| 274 | close_fds=True, |
| 275 | start_new_session=True, |
| 276 | ) |
| 277 | os.close(slave) |
| 278 | |
| 279 | loop = asyncio.get_running_loop() |
| 280 | reader = asyncio.StreamReader() |
| 281 | master_ref = {"fd": master} |
| 282 | |
| 283 | def _release_master_fd(): |
| 284 | cur = master_ref.get("fd") |
| 285 | if cur is None: |
| 286 | return |
| 287 | # Invalidate before close so later cleanup cannot close a reused fd. |
| 288 | master_ref["fd"] = None |
| 289 | try: |
| 290 | proc._pty_master = None # type: ignore[attr-defined] |
| 291 | except Exception: |
| 292 | pass |
| 293 | try: |
| 294 | loop.remove_reader(cur) |
| 295 | except Exception: |
| 296 | pass |
| 297 | try: |
| 298 | os.close(cur) |
| 299 | except OSError: |
| 300 | pass |
| 301 | |
| 302 | def _on_data(): |
| 303 | cur = master_ref.get("fd") |
| 304 | if cur is None: |
| 305 | reader.feed_eof() |
| 306 | return |
| 307 | try: |
| 308 | data = os.read(cur, 1 << 16) |
| 309 | except OSError as e: |
| 310 | if e.errno != errno.EIO: # EIO == EOF on some systems |
| 311 | raise |
| 312 | data = b"" |
| 313 | if data: |
| 314 | reader.feed_data(data) |
| 315 | else: |
| 316 | reader.feed_eof() |
| 317 | _release_master_fd() |
| 318 | |
| 319 | loop.add_reader(master, _on_data) |
| 320 | |
| 321 | class _Stdin: |
| 322 | def write(self, d): |
| 323 | cur = master_ref.get("fd") |
| 324 | if cur is None: |
| 325 | raise OSError(errno.EBADF, "PTY master closed") |
| 326 | os.write(cur, d) |
| 327 | |
| 328 | async def drain(self): |
| 329 | await asyncio.sleep(0) |
| 330 | |
| 331 | proc.stdin = _Stdin() # type: ignore |
| 332 | proc.stdout = reader |
| 333 | proc._pty_master = master # type: ignore[attr-defined] |
| 334 | proc._pty_master_ref = master_ref # type: ignore[attr-defined] |
| 335 | return proc |
| 336 | |
| 337 | |
| 338 | # ──────────────────────────── WINDOWS IMPLEMENTATION ────────────────── |
| 339 | |
| 340 | |
| 341 | async def _spawn_winpty(cmd, cwd, env, echo): |
| 342 | # Clean PowerShell startup: no logo, no profile, bypass execution policy for deterministic behavior |
| 343 | if cmd.strip().lower().startswith("powershell"): |
| 344 | if "-nolog" not in cmd.lower(): |
| 345 | cmd = cmd.replace("powershell.exe", "powershell.exe -NoLogo -NoProfile -ExecutionPolicy Bypass", 1) |
| 346 | |
| 347 | cols, rows = 80, 25 |
| 348 | child = winpty.PtyProcess.spawn(cmd, dimensions=(rows, cols), cwd=cwd or os.getcwd(), env=env) # type: ignore |
| 349 | |
| 350 | loop = asyncio.get_running_loop() |
| 351 | reader = asyncio.StreamReader() |
| 352 | |
| 353 | async def _on_data(): |
| 354 | while child.isalive(): |
| 355 | try: |
| 356 | # Run blocking read in executor to not block event loop |
| 357 | data = await loop.run_in_executor(None, child.read, 1 << 16) |
| 358 | if data: |
| 359 | reader.feed_data(data.encode('utf-8') if isinstance(data, str) else data) |
| 360 | except EOFError: |
| 361 | break |
| 362 | except Exception: |
| 363 | await asyncio.sleep(0.01) |
| 364 | reader.feed_eof() |
| 365 | |
| 366 | # Start pumping output in background |
| 367 | asyncio.create_task(_on_data()) |
| 368 | |
| 369 | class _Stdin: |
| 370 | def write(self, d): |
| 371 | # Use winpty's write method, not os.write |
| 372 | if isinstance(d, bytes): |
| 373 | d = d.decode('utf-8', errors='replace') |
| 374 | # Windows needs \r\n for proper line endings |
| 375 | if _IS_WIN: |
| 376 | d = d.replace('\n', '\r\n') |
| 377 | child.write(d) |
| 378 | |
| 379 | async def drain(self): |
| 380 | await asyncio.sleep(0.01) # Give write time to complete |
| 381 | |
| 382 | class _Proc: |
| 383 | def __init__(self): |
| 384 | self.stdin = _Stdin() # type: ignore |
| 385 | self.stdout = reader |
| 386 | self.pid = child.pid |
| 387 | self.returncode = None |
| 388 | |
| 389 | async def wait(self): |
| 390 | while child.isalive(): |
| 391 | await asyncio.sleep(0.2) |
| 392 | self.returncode = 0 |
| 393 | return 0 |
| 394 | |
| 395 | def terminate(self): |
| 396 | if child.isalive(): |
| 397 | child.terminate() |
| 398 | |
| 399 | def kill(self): |
| 400 | if child.isalive(): |
| 401 | child.kill() |
| 402 | |
| 403 | return _Proc() |
| 404 | |
| 405 | |
| 406 | # ───────────────────────── INTERACTIVE DRIVER ───────────────────────── |
| 407 | if __name__ == "__main__": |
| 408 | |
| 409 | async def interactive_shell(): |
| 410 | shell_cmd, prompt_hint = ("powershell.exe", ">") if _IS_WIN else ("/bin/bash", "$") |
| 411 | |
| 412 | # echo=False → suppress the shell’s own echo of commands |
| 413 | term = TTYSession(shell_cmd) |
| 414 | await term.start() |
| 415 | |
| 416 | timeout = 1.0 |
| 417 | |
| 418 | print(f"Connected to {shell_cmd}.") |
| 419 | print("Type commands for the shell.") |
| 420 | print("• /t=<seconds> → change idle timeout") |
| 421 | print("• /exit → quit helper\n") |
| 422 | |
| 423 | await term.sendline(" ") |
| 424 | print(await term.read_full_until_idle(timeout, timeout), end="", flush=True) |
| 425 | |
| 426 | while True: |
| 427 | try: |
| 428 | user = input(f"(timeout={timeout}) {prompt_hint} ") |
| 429 | except (EOFError, KeyboardInterrupt): |
| 430 | print("\nLeaving…") |
| 431 | break |
| 432 | |
| 433 | if user.lower() == "/exit": |
| 434 | break |
| 435 | if user.startswith("/t="): |
| 436 | try: |
| 437 | timeout = float(user.split("=", 1)[1]) |
| 438 | print(f"[helper] idle timeout set to {timeout}s") |
| 439 | except ValueError: |
| 440 | print("[helper] invalid number") |
| 441 | continue |
| 442 | |
| 443 | idle_timeout = timeout |
| 444 | total_timeout = 10 * idle_timeout |
| 445 | if user == "": |
| 446 | # Just read output, do not send empty line |
| 447 | async for chunk in term.read_chunks_until_idle( |
| 448 | idle_timeout, total_timeout |
| 449 | ): |
| 450 | print(chunk, end="", flush=True) |
| 451 | else: |
| 452 | await term.sendline(user) |
| 453 | async for chunk in term.read_chunks_until_idle( |
| 454 | idle_timeout, total_timeout |
| 455 | ): |
| 456 | print(chunk, end="", flush=True) |
| 457 | |
| 458 | await term.sendline("exit") |
| 459 | await term.wait() |
| 460 | |
| 461 | asyncio.run(interactive_shell()) |