diff --git a/python_files/python_server.py b/python_files/python_server.py index 77b43c692dc3..e7ee92794a21 100644 --- a/python_files/python_server.py +++ b/python_files/python_server.py @@ -64,6 +64,9 @@ def custom_input(prompt=""): message_text = STDIN.buffer.read(content_length).decode() message_json = json.loads(message_text) return message_json["result"]["userInput"] + except EOFError: + # Input stream closed, exit gracefully + sys.exit(0) except Exception: print_log(traceback.format_exc()) @@ -74,7 +77,7 @@ def custom_input(prompt=""): def handle_response(request_id): - while not STDIN.closed: + while True: try: headers = get_headers() # Content-Length is the data size in bytes. @@ -88,8 +91,10 @@ def handle_response(request_id): send_response(our_user_input, message_json["id"]) elif message_json["method"] == "exit": sys.exit(0) - - except Exception: # noqa: PERF203 + except EOFError: # noqa: PERF203 + # Input stream closed, exit gracefully + sys.exit(0) + except Exception: print_log(traceback.format_exc()) @@ -164,7 +169,11 @@ def get_value(self) -> str: def get_headers(): headers = {} while True: - line = STDIN.buffer.readline().decode().strip() + raw = STDIN.buffer.readline() + # Detect EOF: readline() returns empty bytes when input stream is closed + if raw == b"": + raise EOFError("EOF reached while reading headers") + line = raw.decode().strip() if not line: break name, value = line.split(":", 1) @@ -183,7 +192,7 @@ def get_headers(): while "" in sys.path: sys.path.remove("") sys.path.insert(0, "") - while not STDIN.closed: + while True: try: headers = get_headers() # Content-Length is the data size in bytes. @@ -198,6 +207,8 @@ def get_headers(): check_valid_command(request_json) elif request_json["method"] == "exit": sys.exit(0) - - except Exception: # noqa: PERF203 + except EOFError: # noqa: PERF203 + # Input stream closed (VS Code terminated), exit gracefully + sys.exit(0) + except Exception: print_log(traceback.format_exc()) diff --git a/python_files/tests/test_python_server.py b/python_files/tests/test_python_server.py new file mode 100644 index 000000000000..ca542b8ea292 --- /dev/null +++ b/python_files/tests/test_python_server.py @@ -0,0 +1,162 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +"""Tests for python_server.py, specifically EOF handling to prevent infinite loops.""" + +import io +from unittest import mock + +import pytest + + +class TestGetHeaders: + """Tests for the get_headers function.""" + + def test_get_headers_normal(self): + """Test get_headers with valid headers.""" + # Arrange: Import the module + import python_server + + # Create a mock stdin with valid headers + mock_input = b"Content-Length: 100\r\nContent-Type: application/json\r\n\r\n" + mock_stdin = io.BytesIO(mock_input) + + # Act + with mock.patch.object(python_server, "STDIN", mock.Mock(buffer=mock_stdin)): + headers = python_server.get_headers() + + # Assert + assert headers == {"Content-Length": "100", "Content-Type": "application/json"} + + def test_get_headers_eof_raises_error(self): + """Test that get_headers raises EOFError when stdin is closed (EOF).""" + # Arrange: Import the module + import python_server + + # Create a mock stdin that returns empty bytes (EOF) + mock_stdin = io.BytesIO(b"") + + # Act & Assert + with mock.patch.object(python_server, "STDIN", mock.Mock(buffer=mock_stdin)), pytest.raises( + EOFError, match="EOF reached while reading headers" + ): + python_server.get_headers() + + def test_get_headers_eof_mid_headers_raises_error(self): + """Test that get_headers raises EOFError when EOF occurs mid-headers.""" + # Arrange: Import the module + import python_server + + # Create a mock stdin with partial headers then EOF + mock_input = b"Content-Length: 100\r\n" # No terminating empty line + mock_stdin = io.BytesIO(mock_input) + + # Act & Assert + with mock.patch.object(python_server, "STDIN", mock.Mock(buffer=mock_stdin)), pytest.raises( + EOFError, match="EOF reached while reading headers" + ): + python_server.get_headers() + + def test_get_headers_empty_line_terminates(self): + """Test that an empty line (not EOF) properly terminates header reading.""" + # Arrange: Import the module + import python_server + + # Create a mock stdin with headers followed by empty line + mock_input = b"Content-Length: 50\r\n\r\nsome body content" + mock_stdin = io.BytesIO(mock_input) + + # Act + with mock.patch.object(python_server, "STDIN", mock.Mock(buffer=mock_stdin)): + headers = python_server.get_headers() + + # Assert + assert headers == {"Content-Length": "50"} + + +class TestEOFHandling: + """Tests for EOF handling in various functions that use get_headers.""" + + def test_custom_input_exits_on_eof(self): + """Test that custom_input exits gracefully on EOF.""" + # Arrange: Import the module + import python_server + + # Create a mock stdin that returns empty bytes (EOF) + mock_stdin = io.BytesIO(b"") + mock_stdout = io.BytesIO() + + # Act & Assert + with mock.patch.object( + python_server, "STDIN", mock.Mock(buffer=mock_stdin) + ), mock.patch.object(python_server, "STDOUT", mock.Mock(buffer=mock_stdout)), pytest.raises( + SystemExit + ) as exc_info: + python_server.custom_input("prompt> ") + + # Should exit with code 0 (graceful exit) + assert exc_info.value.code == 0 + + def test_handle_response_exits_on_eof(self): + """Test that handle_response exits gracefully on EOF.""" + # Arrange: Import the module + import python_server + + # Create a mock stdin that returns empty bytes (EOF) + mock_stdin = io.BytesIO(b"") + + # Act & Assert + with mock.patch.object(python_server, "STDIN", mock.Mock(buffer=mock_stdin)), pytest.raises( + SystemExit + ) as exc_info: + python_server.handle_response("test-request-id") + + # Should exit with code 0 (graceful exit) + assert exc_info.value.code == 0 + + +class TestMainLoopEOFHandling: + """Tests that simulate the main loop EOF scenario.""" + + def test_main_loop_exits_on_eof(self): + """Test that the main loop pattern exits gracefully on EOF. + + This test verifies the fix for GitHub issue #25620 where the server + would spin at 100% CPU instead of exiting when VS Code closes. + """ + # Arrange: Import the module + import python_server + + # Create a mock stdin that returns empty bytes (EOF) + mock_stdin = io.BytesIO(b"") + + # Simulate what happens in the main loop + with mock.patch.object(python_server, "STDIN", mock.Mock(buffer=mock_stdin)): + try: + python_server.get_headers() + # If we get here without raising EOFError, the fix isn't working + pytest.fail("Expected EOFError to be raised on EOF") + except EOFError: + # This is the expected behavior - the fix is working + pass + + def test_readline_eof_vs_empty_line(self): + """Test that we correctly distinguish between EOF and empty line. + + EOF: readline() returns b'' (empty bytes) + Empty line: readline() returns b'\\r\\n' or b'\\n' (newline bytes) + """ + # Test EOF case + eof_stream = io.BytesIO(b"") + result = eof_stream.readline() + assert result == b"", "EOF should return empty bytes" + + # Test empty line case + empty_line_stream = io.BytesIO(b"\r\n") + result = empty_line_stream.readline() + assert result == b"\r\n", "Empty line should return newline bytes" + + # Test empty line with just newline + empty_line_stream2 = io.BytesIO(b"\n") + result = empty_line_stream2.readline() + assert result == b"\n", "Empty line should return newline bytes"