Fix JSON-RPC pipe reads >64KB by handling short reads

python/copilot/jsonrpc.py

@@ -160,6 +160,29 @@ def _read_loop(self):
             if self._running:
                 print(f"JSON-RPC read loop error: {e}")
 
+    def _read_exact(self, num_bytes: int) -> bytes:
+        """
+        Read exactly num_bytes, handling partial/short reads from pipes.
+
+        Args:
+            num_bytes: Number of bytes to read
+
+        Returns:
+            Bytes read from stream
+
+        Raises:
+            EOFError: If stream ends before reading all bytes
+        """
+        chunks = []
+        remaining = num_bytes
+        while remaining > 0:
+            chunk = self.process.stdout.read(remaining)
+            if not chunk:
+                raise EOFError("Unexpected end of stream while reading JSON-RPC message")
+            chunks.append(chunk)
+            remaining -= len(chunk)
+        return b"".join(chunks)
+
     def _read_message(self) -> Optional[dict]:
         """
         Read a single JSON-RPC message with Content-Length header (blocking)
@@ -182,8 +205,8 @@ def _read_message(self) -> Optional[dict]:
         # Read empty line
         self.process.stdout.readline()
 
-        # Read exact content
-        content_bytes = self.process.stdout.read(content_length)
+        # Read exact content using loop to handle short reads
+        content_bytes = self._read_exact(content_length)
         content = content_bytes.decode("utf-8")
 
         return json.loads(content)

python/test_jsonrpc.py

@@ -0,0 +1,267 @@
+"""
+JsonRpcClient Unit Tests
+
+Tests for the JSON-RPC client implementation, focusing on proper handling
+of large payloads and short reads from pipes.
+"""
+
+import io
+import json
+
+import pytest
+
+from copilot.jsonrpc import JsonRpcClient
+
+
+class MockProcess:
+    """Mock subprocess.Popen for testing JSON-RPC client"""
+
+    def __init__(self):
+        self.stdin = io.BytesIO()
+        self.stdout = None  # Will be set per test
+        self.returncode = None
+
+    def poll(self):
+        return self.returncode
+
+
+class ShortReadStream:
+    """
+    Mock stream that simulates short reads from a pipe.
+
+    This simulates the behavior of Unix pipes when reading data larger than
+    the pipe buffer (typically 64KB). The read() method will return fewer
+    bytes than requested, requiring multiple read calls.
+    """
+
+    def __init__(self, data: bytes, chunk_size: int = 32768):
+        """
+        Args:
+            data: Complete data to be read
+            chunk_size: Maximum bytes to return per read() call (simulates pipe buffer)
+        """
+        self.data = data
+        self.chunk_size = chunk_size
+        self.pos = 0
+
+    def readline(self):
+        """Read until newline"""
+        end = self.data.find(b"\n", self.pos) + 1
+        if end == 0:  # Not found
+            result = self.data[self.pos:]
+            self.pos = len(self.data)
+        else:
+            result = self.data[self.pos:end]
+            self.pos = end
+        return result
+
+    def read(self, n: int) -> bytes:
+        """
+        Read at most n bytes, but may return fewer (short read).
+
+        This simulates the behavior of pipes when data exceeds buffer size.
+        """
+        # Calculate how much we can return (limited by chunk_size)
+        available = len(self.data) - self.pos
+        to_read = min(n, available, self.chunk_size)
+
+        result = self.data[self.pos:self.pos + to_read]
+        self.pos += to_read
+        return result
+
+
+class TestReadExact:
+    """Tests for the _read_exact() method that handles short reads"""
+
+    def test_read_exact_single_chunk(self):
+        """Test reading data that fits in a single chunk"""
+        content = b"Hello, World!"
+        mock_stream = ShortReadStream(content, chunk_size=1024)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_exact(len(content))
+
+        assert result == content
+
+    def test_read_exact_multiple_chunks(self):
+        """Test reading data that requires multiple chunks (short reads)"""
+        # Create 100KB of data
+        content = b"x" * 100000
+        # Simulate 32KB chunks (typical pipe behavior)
+        mock_stream = ShortReadStream(content, chunk_size=32768)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_exact(len(content))
+
+        assert result == content
+        assert len(result) == 100000
+
+    def test_read_exact_at_64kb_boundary(self):
+        """Test reading exactly 64KB (common pipe buffer size)"""
+        content = b"y" * 65536  # Exactly 64KB
+        mock_stream = ShortReadStream(content, chunk_size=65536)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_exact(len(content))
+
+        assert result == content
+        assert len(result) == 65536
+
+    def test_read_exact_exceeds_64kb(self):
+        """Test reading data that exceeds 64KB (triggers the bug without fix)"""
+        # 80KB - larger than typical pipe buffer
+        content = b"z" * 81920
+        # Simulate reading with 64KB limit (macOS pipe buffer)
+        mock_stream = ShortReadStream(content, chunk_size=65536)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_exact(len(content))
+
+        assert result == content
+        assert len(result) == 81920
+
+    def test_read_exact_empty_stream_raises_eof(self):
+        """Test that reading from closed stream raises EOFError"""
+        mock_stream = ShortReadStream(b"", chunk_size=1024)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+
+        with pytest.raises(EOFError, match="Unexpected end of stream"):
+            client._read_exact(10)
+
+    def test_read_exact_partial_data_raises_eof(self):
+        """Test that stream ending mid-message raises EOFError"""
+        # Only 50 bytes available, but we request 100
+        content = b"a" * 50
+        mock_stream = ShortReadStream(content, chunk_size=1024)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+
+        with pytest.raises(EOFError, match="Unexpected end of stream"):
+            client._read_exact(100)
+
+
+class TestReadMessageWithLargePayloads:
+    """Tests for _read_message() with large JSON-RPC messages"""
+
+    def create_jsonrpc_message(self, content_dict: dict) -> bytes:
+        """Create a complete JSON-RPC message with Content-Length header"""
+        content = json.dumps(content_dict, separators=(",", ":"))
+        content_bytes = content.encode("utf-8")
+        header = f"Content-Length: {len(content_bytes)}\r\n\r\n"
+        return header.encode("utf-8") + content_bytes
+
+    def test_read_message_small_payload(self):
+        """Test reading a small JSON-RPC message"""
+        message = {"jsonrpc": "2.0", "id": "1", "result": {"status": "ok"}}
+        full_data = self.create_jsonrpc_message(message)
+
+        mock_stream = ShortReadStream(full_data, chunk_size=1024)
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_message()
+
+        assert result == message
+
+    def test_read_message_large_payload_70kb(self):
+        """Test reading a 70KB JSON-RPC message (exceeds typical pipe buffer)"""
+        # Simulate a large response with context echo (common pattern)
+        large_content = "x" * 70000  # 70KB of data
+        message = {
+            "jsonrpc": "2.0",
+            "id": "1",
+            "result": {"content": large_content, "status": "complete"},
+        }
+
+        full_data = self.create_jsonrpc_message(message)
+        # Simulate 64KB pipe buffer limit
+        mock_stream = ShortReadStream(full_data, chunk_size=65536)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_message()
+
+        assert result == message
+        assert len(result["result"]["content"]) == 70000
+
+    def test_read_message_large_payload_100kb(self):
+        """Test reading a 100KB JSON-RPC message"""
+        large_content = "y" * 100000  # 100KB
+        message = {
+            "jsonrpc": "2.0",
+            "id": "2",
+            "result": {"data": large_content, "metadata": {"size": 100000}},
+        }
+
+        full_data = self.create_jsonrpc_message(message)
+        # Simulate short reads with 32KB chunks
+        mock_stream = ShortReadStream(full_data, chunk_size=32768)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_message()
+
+        assert result == message
+        assert len(result["result"]["data"]) == 100000
+
+    def test_read_message_exactly_64kb_content(self):
+        """Test reading message with exactly 64KB of content"""
+        content_64kb = "z" * 65536  # Exactly 64KB
+        message = {"jsonrpc": "2.0", "id": "3", "result": {"content": content_64kb}}
+
+        full_data = self.create_jsonrpc_message(message)
+        mock_stream = ShortReadStream(full_data, chunk_size=65536)
+
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+        result = client._read_message()
+
+        assert result == message
+        assert len(result["result"]["content"]) == 65536
+
+    def test_read_message_multiple_messages_in_sequence(self):
+        """Test reading multiple large messages in sequence"""
+        message1 = {"jsonrpc": "2.0", "id": "1", "result": {"data": "a" * 50000}}
+        message2 = {"jsonrpc": "2.0", "id": "2", "result": {"data": "b" * 80000}}
+
+        data1 = self.create_jsonrpc_message(message1)
+        data2 = self.create_jsonrpc_message(message2)
+        full_data = data1 + data2
+
+        mock_stream = ShortReadStream(full_data, chunk_size=32768)
+        process = MockProcess()
+        process.stdout = mock_stream
+
+        client = JsonRpcClient(process)
+
+        result1 = client._read_message()
+        assert result1 == message1
+
+        result2 = client._read_message()
+        assert result2 == message2