diff --git a/UnityMcpBridge/Editor/UnityMcpBridge.cs b/UnityMcpBridge/Editor/UnityMcpBridge.cs index fc06dd2..b7e8ef0 100644 --- a/UnityMcpBridge/Editor/UnityMcpBridge.cs +++ b/UnityMcpBridge/Editor/UnityMcpBridge.cs @@ -395,80 +395,22 @@ namespace UnityMcpBridge.Editor using (client) using (NetworkStream stream = client.GetStream()) { - const int MaxMessageBytes = 64 * 1024 * 1024; // 64 MB safety cap byte[] buffer = new byte[8192]; while (isRunning) { try { - // Read message with optional length prefix (8-byte big-endian) - bool usedFraming = false; - string commandText = null; - - // First, attempt to read an 8-byte header - byte[] header = new byte[8]; - int headerFilled = 0; - while (headerFilled < 8) + int bytesRead = await stream.ReadAsync(buffer, 0, buffer.Length); + if (bytesRead == 0) { - int r = await stream.ReadAsync(header, headerFilled, 8 - headerFilled); - if (r == 0) - { - // Disconnected - return; - } - headerFilled += r; - } - - // Interpret header as big-endian payload length, with plausibility check - ulong payloadLen = ReadUInt64BigEndian(header); - if (payloadLen > 0 && payloadLen <= (ulong)MaxMessageBytes) - { - // Framed message path - usedFraming = true; - byte[] payload = await ReadExactAsync(stream, (int)payloadLen); - commandText = System.Text.Encoding.UTF8.GetString(payload); - } - else - { - // Legacy path: treat header bytes as the beginning of a JSON/plain message and read until we have a full JSON - usedFraming = false; - using var ms = new MemoryStream(); - ms.Write(header, 0, header.Length); - - // Read available data in chunks; stop when we have valid JSON or ping, or when no more data available for now - while (true) - { - // If we already have enough text, try to interpret - string currentText = System.Text.Encoding.UTF8.GetString(ms.ToArray()); - string trimmed = currentText.Trim(); - if (trimmed == "ping") - { - commandText = trimmed; - break; - } - if (IsValidJson(trimmed)) - { - commandText = trimmed; - break; - } - - // Read next chunk - int r = await stream.ReadAsync(buffer, 0, buffer.Length); - if (r == 0) - { - // Disconnected mid-message; fall back to whatever we have - commandText = currentText; - break; - } - ms.Write(buffer, 0, r); - - if (ms.Length > MaxMessageBytes) - { - throw new IOException($"Incoming message exceeded {MaxMessageBytes} bytes cap"); - } - } + break; // Client disconnected } + string commandText = System.Text.Encoding.UTF8.GetString( + buffer, + 0, + bytesRead + ); string commandId = Guid.NewGuid().ToString(); TaskCompletionSource tcs = new(); @@ -480,14 +422,6 @@ namespace UnityMcpBridge.Editor /*lang=json,strict*/ "{\"status\":\"success\",\"result\":{\"message\":\"pong\"}}" ); - - if (usedFraming) - { - // Mirror framing for response - byte[] outHeader = new byte[8]; - WriteUInt64BigEndian(outHeader, (ulong)pingResponseBytes.Length); - await stream.WriteAsync(outHeader, 0, outHeader.Length); - } await stream.WriteAsync(pingResponseBytes, 0, pingResponseBytes.Length); continue; } @@ -499,12 +433,6 @@ namespace UnityMcpBridge.Editor string response = await tcs.Task; byte[] responseBytes = System.Text.Encoding.UTF8.GetBytes(response); - if (usedFraming) - { - byte[] outHeader = new byte[8]; - WriteUInt64BigEndian(outHeader, (ulong)responseBytes.Length); - await stream.WriteAsync(outHeader, 0, outHeader.Length); - } await stream.WriteAsync(responseBytes, 0, responseBytes.Length); } catch (Exception ex) @@ -516,55 +444,6 @@ namespace UnityMcpBridge.Editor } } - // Read exactly count bytes or throw if stream closes prematurely - private static async Task ReadExactAsync(NetworkStream stream, int count) - { - byte[] data = new byte[count]; - int offset = 0; - while (offset < count) - { - int r = await stream.ReadAsync(data, offset, count - offset); - if (r == 0) - { - throw new IOException("Connection closed before reading expected bytes"); - } - offset += r; - } - return data; - } - - private static ulong ReadUInt64BigEndian(byte[] buffer) - { - if (buffer == null || buffer.Length < 8) - { - return 0UL; - } - return ((ulong)buffer[0] << 56) - | ((ulong)buffer[1] << 48) - | ((ulong)buffer[2] << 40) - | ((ulong)buffer[3] << 32) - | ((ulong)buffer[4] << 24) - | ((ulong)buffer[5] << 16) - | ((ulong)buffer[6] << 8) - | buffer[7]; - } - - private static void WriteUInt64BigEndian(byte[] dest, ulong value) - { - if (dest == null || dest.Length < 8) - { - throw new ArgumentException("Destination buffer too small for UInt64"); - } - dest[0] = (byte)(value >> 56); - dest[1] = (byte)(value >> 48); - dest[2] = (byte)(value >> 40); - dest[3] = (byte)(value >> 32); - dest[4] = (byte)(value >> 24); - dest[5] = (byte)(value >> 16); - dest[6] = (byte)(value >> 8); - dest[7] = (byte)(value); - } - private static void ProcessCommands() { List processedIds = new(); diff --git a/UnityMcpBridge/UnityMcpServer~/src/unity_connection.py b/UnityMcpBridge/UnityMcpServer~/src/unity_connection.py index bf030d0..9bad736 100644 --- a/UnityMcpBridge/UnityMcpServer~/src/unity_connection.py +++ b/UnityMcpBridge/UnityMcpServer~/src/unity_connection.py @@ -9,7 +9,6 @@ import errno from typing import Dict, Any from config import config from port_discovery import PortDiscovery -import struct # Configure logging using settings from config logging.basicConfig( @@ -54,52 +53,60 @@ class UnityConnection: finally: self.sock = None - def receive_full_response(self, sock) -> bytes: - """Receive a complete response from Unity using 8-byte length-prefixed framing, with legacy fallback.""" - sock.settimeout(config.connection_timeout) - # Try framed first + def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes: + """Receive a complete response from Unity, handling chunked data.""" + chunks = [] + sock.settimeout(config.connection_timeout) # Use timeout from config try: - header = self._read_exact(sock, 8) - (payload_len,) = struct.unpack('>Q', header) - if 0 < payload_len <= (64 * 1024 * 1024): - return self._read_exact(sock, payload_len) - # Implausible length -> treat as legacy stream; fall through - legacy_prefix = header - except Exception: - # Could not read header — treat as legacy - legacy_prefix = b'' - - # Legacy: read until parses as JSON or times out - chunks: list[bytes] = [] - if legacy_prefix: - chunks.append(legacy_prefix) - while True: - chunk = sock.recv(config.buffer_size) - if not chunk: + while True: + chunk = sock.recv(buffer_size) + if not chunk: + if not chunks: + raise Exception("Connection closed before receiving data") + break + chunks.append(chunk) + + # Process the data received so far data = b''.join(chunks) - if not data: - raise Exception("Connection closed before receiving data") - return data - chunks.append(chunk) - data = b''.join(chunks) - try: - if data.strip() == b'ping': + decoded_data = data.decode('utf-8') + + # Check if we've received a complete response + try: + # Special case for ping-pong + if decoded_data.strip().startswith('{"status":"success","result":{"message":"pong"'): + logger.debug("Received ping response") + return data + + # Handle escaped quotes in the content + if '"content":' in decoded_data: + # Find the content field and its value + content_start = decoded_data.find('"content":') + 9 + content_end = decoded_data.rfind('"', content_start) + if content_end > content_start: + # Replace escaped quotes in content with regular quotes + content = decoded_data[content_start:content_end] + content = content.replace('\\"', '"') + decoded_data = decoded_data[:content_start] + content + decoded_data[content_end:] + + # Validate JSON format + json.loads(decoded_data) + + # If we get here, we have valid JSON + logger.info(f"Received complete response ({len(data)} bytes)") return data - json.loads(data.decode('utf-8')) - return data - except Exception: - continue - - def _read_exact(self, sock: socket.socket, n: int) -> bytes: - buf = bytearray(n) - view = memoryview(buf) - read = 0 - while read < n: - r = sock.recv_into(view[read:]) - if r == 0: - raise Exception("Connection closed during read") - read += r - return bytes(buf) + except json.JSONDecodeError: + # We haven't received a complete valid JSON response yet + continue + except Exception as e: + logger.warning(f"Error processing response chunk: {str(e)}") + # Continue reading more chunks as this might not be the complete response + continue + except socket.timeout: + logger.warning("Socket timeout during receive") + raise Exception("Timeout receiving Unity response") + except Exception as e: + logger.error(f"Error during receive: {str(e)}") + raise def send_command(self, command_type: str, params: Dict[str, Any] = None) -> Dict[str, Any]: """Send a command with retry/backoff and port rediscovery. Pings only when requested.""" @@ -153,14 +160,13 @@ class UnityConnection: # Build payload if command_type == 'ping': - body = b'ping' + payload = b'ping' else: command = {"type": command_type, "params": params or {}} - body = json.dumps(command, ensure_ascii=False).encode('utf-8') + payload = json.dumps(command, ensure_ascii=False).encode('utf-8') - # Send with 8-byte big-endian length prefix for robustness - header = struct.pack('>Q', len(body)) - self.sock.sendall(header + body) + # Send + self.sock.sendall(payload) # During retry bursts use a short receive timeout if attempt > 0 and last_short_timeout is None: