revert: remove protocol framing changes from config-stability PR (keep config-only changes)
David Sarno
parent
616d3998b5
commit
fae347b03a
|
|
@ -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<string> 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<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()
|
||||
{
|
||||
List<string> processedIds = new();
|
||||
|
|
|
|||
|
|
@ -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''
|
||||
while True:
|
||||
chunk = sock.recv(buffer_size)
|
||||
if not chunk:
|
||||
if not chunks:
|
||||
raise Exception("Connection closed before receiving data")
|
||||
break
|
||||
chunks.append(chunk)
|
||||
|
||||
# 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:
|
||||
# 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':
|
||||
return data
|
||||
json.loads(data.decode('utf-8'))
|
||||
return data
|
||||
except Exception:
|
||||
continue
|
||||
decoded_data = data.decode('utf-8')
|
||||
|
||||
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)
|
||||
# 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
|
||||
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:
|
||||
|
|
|
|||
Loading…
Reference in New Issue