revert: remove protocol framing changes from config-stability PR (keep config-only changes)
parent
616d3998b5
commit
fae347b03a
|
|
@ -395,80 +395,22 @@ namespace UnityMcpBridge.Editor
|
||||||
using (client)
|
using (client)
|
||||||
using (NetworkStream stream = client.GetStream())
|
using (NetworkStream stream = client.GetStream())
|
||||||
{
|
{
|
||||||
const int MaxMessageBytes = 64 * 1024 * 1024; // 64 MB safety cap
|
|
||||||
byte[] buffer = new byte[8192];
|
byte[] buffer = new byte[8192];
|
||||||
while (isRunning)
|
while (isRunning)
|
||||||
{
|
{
|
||||||
try
|
try
|
||||||
{
|
{
|
||||||
// Read message with optional length prefix (8-byte big-endian)
|
int bytesRead = await stream.ReadAsync(buffer, 0, buffer.Length);
|
||||||
bool usedFraming = false;
|
if (bytesRead == 0)
|
||||||
string commandText = null;
|
|
||||||
|
|
||||||
// First, attempt to read an 8-byte header
|
|
||||||
byte[] header = new byte[8];
|
|
||||||
int headerFilled = 0;
|
|
||||||
while (headerFilled < 8)
|
|
||||||
{
|
{
|
||||||
int r = await stream.ReadAsync(header, headerFilled, 8 - headerFilled);
|
break; // Client disconnected
|
||||||
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");
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
string commandText = System.Text.Encoding.UTF8.GetString(
|
||||||
|
buffer,
|
||||||
|
0,
|
||||||
|
bytesRead
|
||||||
|
);
|
||||||
string commandId = Guid.NewGuid().ToString();
|
string commandId = Guid.NewGuid().ToString();
|
||||||
TaskCompletionSource<string> tcs = new();
|
TaskCompletionSource<string> tcs = new();
|
||||||
|
|
||||||
|
|
@ -480,14 +422,6 @@ namespace UnityMcpBridge.Editor
|
||||||
/*lang=json,strict*/
|
/*lang=json,strict*/
|
||||||
"{\"status\":\"success\",\"result\":{\"message\":\"pong\"}}"
|
"{\"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);
|
await stream.WriteAsync(pingResponseBytes, 0, pingResponseBytes.Length);
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
@ -499,12 +433,6 @@ namespace UnityMcpBridge.Editor
|
||||||
|
|
||||||
string response = await tcs.Task;
|
string response = await tcs.Task;
|
||||||
byte[] responseBytes = System.Text.Encoding.UTF8.GetBytes(response);
|
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);
|
await stream.WriteAsync(responseBytes, 0, responseBytes.Length);
|
||||||
}
|
}
|
||||||
catch (Exception ex)
|
catch (Exception ex)
|
||||||
|
|
@ -516,55 +444,6 @@ namespace UnityMcpBridge.Editor
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Read exactly count bytes or throw if stream closes prematurely
|
|
||||||
private static async Task<byte[]> 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()
|
private static void ProcessCommands()
|
||||||
{
|
{
|
||||||
List<string> processedIds = new();
|
List<string> processedIds = new();
|
||||||
|
|
|
||||||
|
|
@ -9,7 +9,6 @@ import errno
|
||||||
from typing import Dict, Any
|
from typing import Dict, Any
|
||||||
from config import config
|
from config import config
|
||||||
from port_discovery import PortDiscovery
|
from port_discovery import PortDiscovery
|
||||||
import struct
|
|
||||||
|
|
||||||
# Configure logging using settings from config
|
# Configure logging using settings from config
|
||||||
logging.basicConfig(
|
logging.basicConfig(
|
||||||
|
|
@ -54,52 +53,60 @@ class UnityConnection:
|
||||||
finally:
|
finally:
|
||||||
self.sock = None
|
self.sock = None
|
||||||
|
|
||||||
def receive_full_response(self, sock) -> bytes:
|
def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes:
|
||||||
"""Receive a complete response from Unity using 8-byte length-prefixed framing, with legacy fallback."""
|
"""Receive a complete response from Unity, handling chunked data."""
|
||||||
sock.settimeout(config.connection_timeout)
|
chunks = []
|
||||||
# Try framed first
|
sock.settimeout(config.connection_timeout) # Use timeout from config
|
||||||
try:
|
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:
|
while True:
|
||||||
chunk = sock.recv(config.buffer_size)
|
chunk = sock.recv(buffer_size)
|
||||||
if not chunk:
|
if not chunk:
|
||||||
data = b''.join(chunks)
|
if not chunks:
|
||||||
if not data:
|
|
||||||
raise Exception("Connection closed before receiving data")
|
raise Exception("Connection closed before receiving data")
|
||||||
return data
|
break
|
||||||
chunks.append(chunk)
|
chunks.append(chunk)
|
||||||
data = b''.join(chunks)
|
|
||||||
try:
|
|
||||||
if data.strip() == b'ping':
|
|
||||||
return data
|
|
||||||
json.loads(data.decode('utf-8'))
|
|
||||||
return data
|
|
||||||
except Exception:
|
|
||||||
continue
|
|
||||||
|
|
||||||
def _read_exact(self, sock: socket.socket, n: int) -> bytes:
|
# Process the data received so far
|
||||||
buf = bytearray(n)
|
data = b''.join(chunks)
|
||||||
view = memoryview(buf)
|
decoded_data = data.decode('utf-8')
|
||||||
read = 0
|
|
||||||
while read < n:
|
# Check if we've received a complete response
|
||||||
r = sock.recv_into(view[read:])
|
try:
|
||||||
if r == 0:
|
# Special case for ping-pong
|
||||||
raise Exception("Connection closed during read")
|
if decoded_data.strip().startswith('{"status":"success","result":{"message":"pong"'):
|
||||||
read += r
|
logger.debug("Received ping response")
|
||||||
return bytes(buf)
|
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
|
||||||
|
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]:
|
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."""
|
"""Send a command with retry/backoff and port rediscovery. Pings only when requested."""
|
||||||
|
|
@ -153,14 +160,13 @@ class UnityConnection:
|
||||||
|
|
||||||
# Build payload
|
# Build payload
|
||||||
if command_type == 'ping':
|
if command_type == 'ping':
|
||||||
body = b'ping'
|
payload = b'ping'
|
||||||
else:
|
else:
|
||||||
command = {"type": command_type, "params": params or {}}
|
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
|
# Send
|
||||||
header = struct.pack('>Q', len(body))
|
self.sock.sendall(payload)
|
||||||
self.sock.sendall(header + body)
|
|
||||||
|
|
||||||
# During retry bursts use a short receive timeout
|
# During retry bursts use a short receive timeout
|
||||||
if attempt > 0 and last_short_timeout is None:
|
if attempt > 0 and last_short_timeout is None:
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue