diff --git a/strix/interface/tui/runtime.py b/strix/interface/tui/runtime.py index b9e5d895f..9688109b0 100644 --- a/strix/interface/tui/runtime.py +++ b/strix/interface/tui/runtime.py @@ -49,7 +49,7 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) # A cold go run can download a toolchain and modules, then compile, before the -# child answers the handshake. +# child connects and answers the handshake. _GO_RUN_HANDSHAKE_TIMEOUT = 300.0 @@ -435,11 +435,9 @@ class GoTuiRuntime: file=original_stdout, flush=True, ) - process, backend_socket = await launch_tui_process(command, env, cwd) - await self.server.start( - backend_socket, - handshake_timeout=_GO_RUN_HANDSHAKE_TIMEOUT if cwd is not None else None, - ) + startup_timeout = _GO_RUN_HANDSHAKE_TIMEOUT if cwd is not None else None + process, backend_socket = await launch_tui_process(command, env, cwd, startup_timeout) + await self.server.start(backend_socket, handshake_timeout=startup_timeout) prepare_task = self._start_preparation() sync_task = asyncio.create_task(self.sync_state()) return_code = await wait_process(process) diff --git a/strix/interface/tui/sidecar.py b/strix/interface/tui/sidecar.py index c0d0f986d..f3326f5c2 100644 --- a/strix/interface/tui/sidecar.py +++ b/strix/interface/tui/sidecar.py @@ -9,12 +9,14 @@ import os import secrets import socket import subprocess +import time from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Any _WINDOWS_AUTH_TIMEOUT = 10.0 +_WINDOWS_ACCEPT_POLL_INTERVAL = 0.2 _PROCESS_EXIT_TIMEOUT = 5.0 _SENSITIVE_ENV_SUFFIXES = ("_API_KEY", "_ACCESS_KEY") _SENSITIVE_ENV_PARTS = frozenset( @@ -86,10 +88,28 @@ def _authenticate_connection( def _accept_authenticated_connection( listener: socket.socket, expected_token: str, + process: subprocess.Popen[bytes] | None = None, + connect_timeout: float | None = None, ) -> socket.socket: - """Accept and authenticate the one Windows loopback connection.""" - listener.settimeout(_WINDOWS_AUTH_TIMEOUT) - connection, address = listener.accept() + """Accept and authenticate the one Windows loopback connection. + + Unlike the POSIX socketpair, a listener gets no EOF when the child dies, so + poll for its exit instead of waiting out the whole limit. + """ + timeout = _WINDOWS_AUTH_TIMEOUT if connect_timeout is None else connect_timeout + deadline = time.monotonic() + timeout + listener.settimeout(min(timeout, _WINDOWS_ACCEPT_POLL_INTERVAL)) + while True: + try: + connection, address = listener.accept() + break + except TimeoutError: + if process is not None and process.poll() is not None: + raise ConnectionError( + f"TUI exited with status {process.returncode} before connecting" + ) from None + if time.monotonic() >= deadline: + raise TimeoutError("Timed out waiting for the TUI to connect") from None try: _authenticate_connection(connection, address, expected_token) except BaseException: @@ -126,9 +146,10 @@ async def launch_tui_process( command: list[str], env: dict[str, str], cwd: str | None, + connect_timeout: float | None = None, ) -> tuple[asyncio.subprocess.Process | subprocess.Popen[bytes], socket.socket]: if os.name == "nt": - return await _launch_windows_tui_process(command, env, cwd) + return await _launch_windows_tui_process(command, env, cwd, connect_timeout) return await _launch_posix_tui_process(command, env, cwd) @@ -155,6 +176,7 @@ async def _launch_windows_tui_process( command: list[str], env: dict[str, str], cwd: str | None, + connect_timeout: float | None = None, ) -> tuple[subprocess.Popen[bytes], socket.socket]: listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM) windows_process: subprocess.Popen[bytes] | None = None @@ -167,7 +189,9 @@ async def _launch_windows_tui_process( host, port = listener.getsockname()[:2] env.update({"STRIX_TUI_ADDR": f"{host}:{port}", "STRIX_TUI_TOKEN": token}) windows_process = subprocess.Popen(command, env=env, cwd=cwd) # noqa: S603 - connection = await asyncio.to_thread(_accept_authenticated_connection, listener, token) + connection = await asyncio.to_thread( + _accept_authenticated_connection, listener, token, windows_process, connect_timeout + ) except BaseException: await terminate_process(windows_process) raise diff --git a/tests/test_go_tui_runtime.py b/tests/test_go_tui_runtime.py index 0d669a9f9..8d8f1e892 100644 --- a/tests/test_go_tui_runtime.py +++ b/tests/test_go_tui_runtime.py @@ -9,6 +9,7 @@ import socket import struct import sys import threading +import time from pathlib import Path from types import SimpleNamespace from typing import Any, cast @@ -177,6 +178,76 @@ with socket.create_connection((host, int(port))) as connection: assert await sidecar.wait_process(process) == 0 +_CONNECT_AFTER_DELAY = """ +import os +import socket +import sys +import time + +time.sleep(float(sys.argv[1])) +host, port = os.environ["STRIX_TUI_ADDR"].rsplit(":", 1) +with socket.create_connection((host, int(port))) as connection: + connection.sendall(os.environ["STRIX_TUI_TOKEN"].encode("ascii")) +""" + + +@pytest.mark.asyncio +async def test_windows_transport_waits_for_a_child_that_is_still_compiling( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(sidecar, "_WINDOWS_AUTH_TIMEOUT", 0.1) + + process, connection = await sidecar._launch_windows_tui_process( + [sys.executable, "-c", _CONNECT_AFTER_DELAY, "1"], os.environ.copy(), None, 30.0 + ) + connection.close() + + assert await sidecar.wait_process(process) == 0 + + +@pytest.mark.asyncio +async def test_windows_transport_keeps_the_short_limit_by_default( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(sidecar, "_WINDOWS_AUTH_TIMEOUT", 0.1) + + with pytest.raises(TimeoutError, match="Timed out waiting for the TUI to connect"): + await sidecar._launch_windows_tui_process( + [sys.executable, "-c", _CONNECT_AFTER_DELAY, "30"], os.environ.copy(), None + ) + + +@pytest.mark.asyncio +async def test_windows_transport_fails_fast_when_the_child_exits_before_connecting() -> None: + launch = sidecar._launch_windows_tui_process( + [sys.executable, "-c", "raise SystemExit(3)"], os.environ.copy(), None, 120.0 + ) + + with pytest.raises(ConnectionError, match="status 3"): + await asyncio.wait_for(launch, timeout=30) + + +def test_cancelled_windows_launch_does_not_hold_up_shutdown() -> None: + async def cancel_launch() -> None: + launch = asyncio.create_task( + sidecar._launch_windows_tui_process( + [sys.executable, "-c", _CONNECT_AFTER_DELAY, "120"], + os.environ.copy(), + None, + 120.0, + ) + ) + await asyncio.sleep(0.5) + launch.cancel() + with pytest.raises(asyncio.CancelledError): + await launch + + started = time.monotonic() + asyncio.run(cancel_launch()) + + assert time.monotonic() - started < 30 + + async def _receive_exactly(connection: socket.socket, size: int) -> bytes: result = b"" while len(result) < size: @@ -211,7 +282,10 @@ async def test_runtime_prepares_state_before_launch_and_scans_only_after_ready( scan_started = asyncio.Event() async def launch( - _command: list[str], _env: dict[str, str], _cwd: str | None + _command: list[str], + _env: dict[str, str], + _cwd: str | None, + _connect_timeout: float | None, ) -> tuple[SimpleNamespace, socket.socket]: return SimpleNamespace(returncode=None), backend @@ -272,10 +346,16 @@ def _stub_launch( runtime: GoTuiRuntime, command: list[str], backend: socket.socket, -) -> None: +) -> list[float | None]: + connect_timeouts: list[float | None] = [] + async def launch( - _command: list[str], _env: dict[str, str], _cwd: str | None + _command: list[str], + _env: dict[str, str], + _cwd: str | None, + connect_timeout: float | None, ) -> tuple[SimpleNamespace, socket.socket]: + connect_timeouts.append(connect_timeout) return SimpleNamespace(returncode=0), backend async def wait_process(_process: object) -> int: @@ -288,6 +368,7 @@ def _stub_launch( monkeypatch.setattr(go_tui, "wait_process", wait_process) monkeypatch.setattr(go_tui, "persist_current", lambda: None) monkeypatch.setattr(server_module, "_HANDSHAKE_TIMEOUT", 0.05) + return connect_timeouts async def _reply_ready_after(child: socket.socket, delay: float) -> None: @@ -308,13 +389,14 @@ async def test_go_run_child_may_compile_past_the_packaged_handshake_timeout( runtime = GoTuiRuntime(runtime_args) backend, child = socket.socketpair() child.setblocking(False) # noqa: FBT003 - _stub_launch(monkeypatch, runtime, ["go", "run", "./cmd/strix-tui"], backend) + connect_timeouts = _stub_launch(monkeypatch, runtime, ["go", "run", "./cmd/strix-tui"], backend) run_task = asyncio.create_task(runtime.run()) try: - await _reply_ready_after(child, 0.2) + await asyncio.wait_for(_reply_ready_after(child, 0.2), timeout=2) await asyncio.wait_for(run_task, timeout=2) assert runtime.server.activated is True + assert connect_timeouts == [go_tui._GO_RUN_HANDSHAKE_TIMEOUT] finally: child.close() if not run_task.done(): @@ -330,11 +412,12 @@ async def test_packaged_child_slower_than_handshake_timeout_fails( runtime = GoTuiRuntime(runtime_args) backend, child = socket.socketpair() child.setblocking(False) # noqa: FBT003 - _stub_launch(monkeypatch, runtime, ["test-sidecar"], backend) + connect_timeouts = _stub_launch(monkeypatch, runtime, ["test-sidecar"], backend) try: with pytest.raises(go_tui.GoTuiPreActivationError, match="Timed out"): await asyncio.wait_for(runtime.run(), timeout=2) + assert connect_timeouts == [None] finally: child.close()