From d3b2f1c8ac8bc87afc6f51b009785a70bc1bdee8 Mon Sep 17 00:00:00 2001
From: Andrew Liu <andrewliu278@gmail.com>
Date: Wed, 29 Jan 2025 06:45:32 +0000
Subject: [PATCH] address comments

---
 modal/_utils/shell_utils.py | 33 ++++++++++++++++++---------------
 modal/container_process.py  |  6 ------
 modal/io_streams.py         |  3 ++-
 3 files changed, 20 insertions(+), 22 deletions(-)

diff --git a/modal/_utils/shell_utils.py b/modal/_utils/shell_utils.py
index 051d8bff8a..f1f81d8f83 100644
--- a/modal/_utils/shell_utils.py
+++ b/modal/_utils/shell_utils.py
@@ -3,16 +3,13 @@
 import asyncio
 import contextlib
 import errno
-import fcntl
 import os
 import select
+import shutil
 import signal
-import struct
 import sys
-import termios
 import threading
 from collections.abc import Coroutine
-from queue import Empty, Queue
 from types import FrameType
 from typing import Callable, Optional
 
@@ -90,21 +87,28 @@ class WindowSizeHandler:
     """Handles terminal window resize events."""
 
     def __init__(self):
-        """Initialize window size handler. Must be called from the main thread to set signals properly.
+        """Initialize window size handler. It is a system limitation that we can only set signals from the main thread.
+
         In case this is invoked from a thread that is not the main thread, e.g. in tests, the context manager
-        becomes a no-op."""
+        becomes a no-op.
+        """
         self._is_main_thread = threading.current_thread() is threading.main_thread()
-        self._event_queue: Queue[tuple[int, int]] = Queue()
+
+        self._current_size: tuple[int, int] = self._get_terminal_size()
+        self._size_changed = asyncio.Event()
 
         if self._is_main_thread and hasattr(signal, "SIGWINCH"):
             signal.signal(signal.SIGWINCH, self._queue_resize_event)
 
+    def _get_terminal_size(self) -> tuple[int, int]:
+        c, r = shutil.get_terminal_size()
+        return r, c
+
     def _queue_resize_event(self, signum: Optional[int] = None, frame: Optional[FrameType] = None) -> None:
         """Signal handler for SIGWINCH that queues events."""
         try:
-            hw = struct.unpack("hh", fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, b"1234"))
-            rows, cols = hw
-            self._event_queue.put((rows, cols))
+            self._current_size = self._get_terminal_size()
+            self._size_changed.set()
         except Exception:
             # ignore failed window size reads
             pass
@@ -112,6 +116,7 @@ def _queue_resize_event(self, signum: Optional[int] = None, frame: Optional[Fram
     @contextlib.asynccontextmanager
     async def watch_window_size(self, handler: Callable[[int, int], Coroutine]):
         """Context manager that processes window resize events from the queue.
+
         Can be run from any thread. If the window manager was initialized from a thread that is not the main thread,
         e.g. in tests, this context manager is a no-op.
         """
@@ -121,11 +126,9 @@ async def watch_window_size(self, handler: Callable[[int, int], Coroutine]):
 
         async def process_events():
             while True:
-                try:
-                    rows, cols = self._event_queue.get_nowait()
-                    await handler(rows, cols)
-                except Empty:
-                    await asyncio.sleep(0.1)
+                await self._size_changed.wait()
+                await handler(*self._current_size)
+                self._size_changed.clear()
 
         event_task = asyncio.create_task(process_events())
         try:
diff --git a/modal/container_process.py b/modal/container_process.py
index ab575f56e4..94ba44c442 100644
--- a/modal/container_process.py
+++ b/modal/container_process.py
@@ -1,7 +1,6 @@
 # Copyright Modal Labs 2024
 import asyncio
 import platform
-import struct
 from typing import Generic, Optional, TypeVar
 
 from modal_proto import api_pb2
@@ -153,11 +152,6 @@ async def _handle_input(data: bytes, message_index: int):
             await self.stdin.drain()
 
         async def _send_window_resize(rows: int, cols: int):
-            # resize sequence:
-            # - 2 bytes for the number of rows (big-endian)
-            # - 2 bytes for the number of columns (big-endian)
-            dims = struct.pack(">HH", rows, cols)
-            self.stdin.write(dims)
             await self.stdin.drain(_terminal_size=(rows, cols))
 
         async with TaskContext() as tc:
diff --git a/modal/io_streams.py b/modal/io_streams.py
index d4dca8064c..911f9200a4 100644
--- a/modal/io_streams.py
+++ b/modal/io_streams.py
@@ -386,7 +386,8 @@ def write_eof(self) -> None:
 
     async def drain(
         self,
-        _terminal_size: Optional[tuple[int, int]] = None,  # internal option to send terminal window resize events
+        # internal option to send terminal resize events to container exec
+        _terminal_size: Optional[tuple[int, int]] = None,
     ) -> None:
         """Flush the write buffer and send data to the running process.