Skip to content

Commit fae347b

Browse files
committed
revert: remove protocol framing changes from config-stability PR (keep config-only changes)
1 parent 616d399 commit fae347b

File tree

2 files changed

+63
-178
lines changed

2 files changed

+63
-178
lines changed

UnityMcpBridge/Editor/UnityMcpBridge.cs

Lines changed: 8 additions & 129 deletions
Original file line numberDiff line numberDiff line change
@@ -395,80 +395,22 @@ private static async Task HandleClientAsync(TcpClient client)
395395
using (client)
396396
using (NetworkStream stream = client.GetStream())
397397
{
398-
const int MaxMessageBytes = 64 * 1024 * 1024; // 64 MB safety cap
399398
byte[] buffer = new byte[8192];
400399
while (isRunning)
401400
{
402401
try
403402
{
404-
// Read message with optional length prefix (8-byte big-endian)
405-
bool usedFraming = false;
406-
string commandText = null;
407-
408-
// First, attempt to read an 8-byte header
409-
byte[] header = new byte[8];
410-
int headerFilled = 0;
411-
while (headerFilled < 8)
403+
int bytesRead = await stream.ReadAsync(buffer, 0, buffer.Length);
404+
if (bytesRead == 0)
412405
{
413-
int r = await stream.ReadAsync(header, headerFilled, 8 - headerFilled);
414-
if (r == 0)
415-
{
416-
// Disconnected
417-
return;
418-
}
419-
headerFilled += r;
420-
}
421-
422-
// Interpret header as big-endian payload length, with plausibility check
423-
ulong payloadLen = ReadUInt64BigEndian(header);
424-
if (payloadLen > 0 && payloadLen <= (ulong)MaxMessageBytes)
425-
{
426-
// Framed message path
427-
usedFraming = true;
428-
byte[] payload = await ReadExactAsync(stream, (int)payloadLen);
429-
commandText = System.Text.Encoding.UTF8.GetString(payload);
430-
}
431-
else
432-
{
433-
// Legacy path: treat header bytes as the beginning of a JSON/plain message and read until we have a full JSON
434-
usedFraming = false;
435-
using var ms = new MemoryStream();
436-
ms.Write(header, 0, header.Length);
437-
438-
// Read available data in chunks; stop when we have valid JSON or ping, or when no more data available for now
439-
while (true)
440-
{
441-
// If we already have enough text, try to interpret
442-
string currentText = System.Text.Encoding.UTF8.GetString(ms.ToArray());
443-
string trimmed = currentText.Trim();
444-
if (trimmed == "ping")
445-
{
446-
commandText = trimmed;
447-
break;
448-
}
449-
if (IsValidJson(trimmed))
450-
{
451-
commandText = trimmed;
452-
break;
453-
}
454-
455-
// Read next chunk
456-
int r = await stream.ReadAsync(buffer, 0, buffer.Length);
457-
if (r == 0)
458-
{
459-
// Disconnected mid-message; fall back to whatever we have
460-
commandText = currentText;
461-
break;
462-
}
463-
ms.Write(buffer, 0, r);
464-
465-
if (ms.Length > MaxMessageBytes)
466-
{
467-
throw new IOException($"Incoming message exceeded {MaxMessageBytes} bytes cap");
468-
}
469-
}
406+
break; // Client disconnected
470407
}
471408

409+
string commandText = System.Text.Encoding.UTF8.GetString(
410+
buffer,
411+
0,
412+
bytesRead
413+
);
472414
string commandId = Guid.NewGuid().ToString();
473415
TaskCompletionSource<string> tcs = new();
474416

@@ -480,14 +422,6 @@ private static async Task HandleClientAsync(TcpClient client)
480422
/*lang=json,strict*/
481423
"{\"status\":\"success\",\"result\":{\"message\":\"pong\"}}"
482424
);
483-
484-
if (usedFraming)
485-
{
486-
// Mirror framing for response
487-
byte[] outHeader = new byte[8];
488-
WriteUInt64BigEndian(outHeader, (ulong)pingResponseBytes.Length);
489-
await stream.WriteAsync(outHeader, 0, outHeader.Length);
490-
}
491425
await stream.WriteAsync(pingResponseBytes, 0, pingResponseBytes.Length);
492426
continue;
493427
}
@@ -499,12 +433,6 @@ private static async Task HandleClientAsync(TcpClient client)
499433

500434
string response = await tcs.Task;
501435
byte[] responseBytes = System.Text.Encoding.UTF8.GetBytes(response);
502-
if (usedFraming)
503-
{
504-
byte[] outHeader = new byte[8];
505-
WriteUInt64BigEndian(outHeader, (ulong)responseBytes.Length);
506-
await stream.WriteAsync(outHeader, 0, outHeader.Length);
507-
}
508436
await stream.WriteAsync(responseBytes, 0, responseBytes.Length);
509437
}
510438
catch (Exception ex)
@@ -516,55 +444,6 @@ private static async Task HandleClientAsync(TcpClient client)
516444
}
517445
}
518446

519-
// Read exactly count bytes or throw if stream closes prematurely
520-
private static async Task<byte[]> ReadExactAsync(NetworkStream stream, int count)
521-
{
522-
byte[] data = new byte[count];
523-
int offset = 0;
524-
while (offset < count)
525-
{
526-
int r = await stream.ReadAsync(data, offset, count - offset);
527-
if (r == 0)
528-
{
529-
throw new IOException("Connection closed before reading expected bytes");
530-
}
531-
offset += r;
532-
}
533-
return data;
534-
}
535-
536-
private static ulong ReadUInt64BigEndian(byte[] buffer)
537-
{
538-
if (buffer == null || buffer.Length < 8)
539-
{
540-
return 0UL;
541-
}
542-
return ((ulong)buffer[0] << 56)
543-
| ((ulong)buffer[1] << 48)
544-
| ((ulong)buffer[2] << 40)
545-
| ((ulong)buffer[3] << 32)
546-
| ((ulong)buffer[4] << 24)
547-
| ((ulong)buffer[5] << 16)
548-
| ((ulong)buffer[6] << 8)
549-
| buffer[7];
550-
}
551-
552-
private static void WriteUInt64BigEndian(byte[] dest, ulong value)
553-
{
554-
if (dest == null || dest.Length < 8)
555-
{
556-
throw new ArgumentException("Destination buffer too small for UInt64");
557-
}
558-
dest[0] = (byte)(value >> 56);
559-
dest[1] = (byte)(value >> 48);
560-
dest[2] = (byte)(value >> 40);
561-
dest[3] = (byte)(value >> 32);
562-
dest[4] = (byte)(value >> 24);
563-
dest[5] = (byte)(value >> 16);
564-
dest[6] = (byte)(value >> 8);
565-
dest[7] = (byte)(value);
566-
}
567-
568447
private static void ProcessCommands()
569448
{
570449
List<string> processedIds = new();

UnityMcpBridge/UnityMcpServer~/src/unity_connection.py

Lines changed: 55 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@
99
from typing import Dict, Any
1010
from config import config
1111
from port_discovery import PortDiscovery
12-
import struct
1312

1413
# Configure logging using settings from config
1514
logging.basicConfig(
@@ -54,52 +53,60 @@ def disconnect(self):
5453
finally:
5554
self.sock = None
5655

57-
def receive_full_response(self, sock) -> bytes:
58-
"""Receive a complete response from Unity using 8-byte length-prefixed framing, with legacy fallback."""
59-
sock.settimeout(config.connection_timeout)
60-
# Try framed first
56+
def receive_full_response(self, sock, buffer_size=config.buffer_size) -> bytes:
57+
"""Receive a complete response from Unity, handling chunked data."""
58+
chunks = []
59+
sock.settimeout(config.connection_timeout) # Use timeout from config
6160
try:
62-
header = self._read_exact(sock, 8)
63-
(payload_len,) = struct.unpack('>Q', header)
64-
if 0 < payload_len <= (64 * 1024 * 1024):
65-
return self._read_exact(sock, payload_len)
66-
# Implausible length -> treat as legacy stream; fall through
67-
legacy_prefix = header
68-
except Exception:
69-
# Could not read header — treat as legacy
70-
legacy_prefix = b''
71-
72-
# Legacy: read until parses as JSON or times out
73-
chunks: list[bytes] = []
74-
if legacy_prefix:
75-
chunks.append(legacy_prefix)
76-
while True:
77-
chunk = sock.recv(config.buffer_size)
78-
if not chunk:
61+
while True:
62+
chunk = sock.recv(buffer_size)
63+
if not chunk:
64+
if not chunks:
65+
raise Exception("Connection closed before receiving data")
66+
break
67+
chunks.append(chunk)
68+
69+
# Process the data received so far
7970
data = b''.join(chunks)
80-
if not data:
81-
raise Exception("Connection closed before receiving data")
82-
return data
83-
chunks.append(chunk)
84-
data = b''.join(chunks)
85-
try:
86-
if data.strip() == b'ping':
71+
decoded_data = data.decode('utf-8')
72+
73+
# Check if we've received a complete response
74+
try:
75+
# Special case for ping-pong
76+
if decoded_data.strip().startswith('{"status":"success","result":{"message":"pong"'):
77+
logger.debug("Received ping response")
78+
return data
79+
80+
# Handle escaped quotes in the content
81+
if '"content":' in decoded_data:
82+
# Find the content field and its value
83+
content_start = decoded_data.find('"content":') + 9
84+
content_end = decoded_data.rfind('"', content_start)
85+
if content_end > content_start:
86+
# Replace escaped quotes in content with regular quotes
87+
content = decoded_data[content_start:content_end]
88+
content = content.replace('\\"', '"')
89+
decoded_data = decoded_data[:content_start] + content + decoded_data[content_end:]
90+
91+
# Validate JSON format
92+
json.loads(decoded_data)
93+
94+
# If we get here, we have valid JSON
95+
logger.info(f"Received complete response ({len(data)} bytes)")
8796
return data
88-
json.loads(data.decode('utf-8'))
89-
return data
90-
except Exception:
91-
continue
92-
93-
def _read_exact(self, sock: socket.socket, n: int) -> bytes:
94-
buf = bytearray(n)
95-
view = memoryview(buf)
96-
read = 0
97-
while read < n:
98-
r = sock.recv_into(view[read:])
99-
if r == 0:
100-
raise Exception("Connection closed during read")
101-
read += r
102-
return bytes(buf)
97+
except json.JSONDecodeError:
98+
# We haven't received a complete valid JSON response yet
99+
continue
100+
except Exception as e:
101+
logger.warning(f"Error processing response chunk: {str(e)}")
102+
# Continue reading more chunks as this might not be the complete response
103+
continue
104+
except socket.timeout:
105+
logger.warning("Socket timeout during receive")
106+
raise Exception("Timeout receiving Unity response")
107+
except Exception as e:
108+
logger.error(f"Error during receive: {str(e)}")
109+
raise
103110

104111
def send_command(self, command_type: str, params: Dict[str, Any] = None) -> Dict[str, Any]:
105112
"""Send a command with retry/backoff and port rediscovery. Pings only when requested."""
@@ -153,14 +160,13 @@ def read_status_file() -> dict | None:
153160

154161
# Build payload
155162
if command_type == 'ping':
156-
body = b'ping'
163+
payload = b'ping'
157164
else:
158165
command = {"type": command_type, "params": params or {}}
159-
body = json.dumps(command, ensure_ascii=False).encode('utf-8')
166+
payload = json.dumps(command, ensure_ascii=False).encode('utf-8')
160167

161-
# Send with 8-byte big-endian length prefix for robustness
162-
header = struct.pack('>Q', len(body))
163-
self.sock.sendall(header + body)
168+
# Send
169+
self.sock.sendall(payload)
164170

165171
# During retry bursts use a short receive timeout
166172
if attempt > 0 and last_short_timeout is None:

0 commit comments

Comments
 (0)