Add handshake fallback and logging checks #17
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the You can disable this status message by setting the ✨ Finishing Touches🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
Greptile Summary
This PR enhances the Unity MCP bridge's robustness and maintainability through three key improvements. The changes focus on socket handling, test reliability, and code quality enforcement.
The most significant change is in unity_connection.py, where the PR introduces thread-safe socket operations using a threading lock to prevent race conditions when multiple threads access the same socket connection. The handshake logic becomes more flexible by making FRAMING requirements configurable - the system can now fall back to legacy mode when Unity clients don't support the FRAMING=1 protocol, improving compatibility with different Unity setups. Enhanced debug logging with structured formatting and exception suppression prevents logging operations from crashing the main functionality while still providing valuable debugging information. The timeout handling logic is also fixed to properly restore original socket timeout values instead of using hardcoded configuration values.
In the test suite, test_transport_framing.py gains defensive programming with a 1-second socket timeout to prevent the unframed data disconnect test from hanging indefinitely. This ensures the test suite remains reliable in CI/CD environments. The test_logging_stdout.py test becomes more comprehensive by detecting builtins.print() calls in addition to regular print() statements, ensuring no stdout contamination can interfere with the JSON-RPC protocol communication between Unity and MCP clients.
These changes integrate well with the existing Unity MCP architecture by maintaining backward compatibility while adding essential production-ready features like thread safety and configurable handshake behavior. The improvements align with the project's goal of providing a stable bridge between Unity applications and Large Language Models via the Model Context Protocol.
Important Files Changed
File Changes Summary
| Filename | Score | Overview |
|---|---|---|
| UnityMcpBridge/UnityMcpServer~/src/unity_connection.py | 4/5 | Added thread-safe socket operations with locking, configurable handshake requirements, improved error handling, and enhanced debug logging |
| tests/test_transport_framing.py | 5/5 | Added 1-second socket timeout to prevent test hangs in unframed data disconnect test |
| tests/test_logging_stdout.py | 5/5 | Extended stdout detection to catch builtins.print() calls and refactored conditional logic for better readability |
Confidence score: 4/5
- This PR is safe to merge with minimal risk as it primarily adds defensive programming and robustness improvements
- Score reflects well-designed thread safety improvements and comprehensive test enhancements, with only minor complexity concerns around new configuration options
- Pay close attention to UnityMcpBridge/UnityMcpServer~/src/unity_connection.py for the new threading lock and configurable handshake behavior
Sequence Diagram
sequenceDiagram
participant User
participant MCP_Server as MCP Server
participant Unity_Connection as UnityConnection
participant Socket
participant Unity_Editor as Unity Editor
participant Port_Discovery as PortDiscovery
participant Status_File as Status File
User->>MCP_Server: "send_command(command_type, params)"
MCP_Server->>Unity_Connection: "get_unity_connection()"
alt Connection not established
Unity_Connection->>Port_Discovery: "discover_unity_port()"
Port_Discovery-->>Unity_Connection: "port number"
Unity_Connection->>Socket: "socket.connect(host, port)"
Socket->>Unity_Editor: "TCP connection request"
Unity_Editor-->>Socket: "connection established"
Note over Unity_Connection, Unity_Editor: Handshake Phase
Unity_Connection->>Socket: "settimeout(handshake_timeout)"
Unity_Editor->>Socket: "send greeting (FRAMING=1 or legacy)"
Socket-->>Unity_Connection: "recv(256) - greeting message"
alt Greeting contains FRAMING=1
Unity_Connection->>Unity_Connection: "use_framing = True"
Note right of Unity_Connection: Strict handshake successful
else require_framing is True
Unity_Connection->>Socket: "send advisory message"
Unity_Connection->>Unity_Connection: "raise ConnectionError"
else require_framing is False
Unity_Connection->>Unity_Connection: "use_framing = False (legacy mode)"
end
Unity_Connection->>Socket: "settimeout(connection_timeout)"
end
Unity_Connection-->>MCP_Server: "connection ready"
Note over MCP_Server, Status_File: Preflight Check
MCP_Server->>Status_File: "read_status_file()"
Status_File-->>MCP_Server: "status (reloading/ready)"
alt Unity is reloading
MCP_Server-->>User: "structured error: reloading state"
else Unity ready
loop Retry attempts (with backoff)
MCP_Server->>Unity_Connection: "send_command with _io_lock"
alt Ping command
Unity_Connection->>Socket: "send 'ping' payload"
else Regular command
Unity_Connection->>Socket: "send JSON command"
end
alt use_framing is True
Unity_Connection->>Socket: "send header (8-byte length) + payload"
Socket->>Unity_Editor: "framed message"
Unity_Editor->>Socket: "send header + response payload"
Socket-->>Unity_Connection: "recv header (8 bytes)"
Unity_Connection->>Socket: "_read_exact(payload_length)"
Socket-->>Unity_Connection: "complete framed response"
else Legacy mode
Unity_Connection->>Socket: "send raw payload"
Socket->>Unity_Editor: "raw message"
Unity_Editor->>Socket: "send raw response"
Unity_Connection->>Socket: "recv chunks until complete JSON"
Socket-->>Unity_Connection: "complete response"
end
Unity_Connection-->>MCP_Server: "response data"
MCP_Server->>MCP_Server: "parse JSON response"
alt Response successful
MCP_Server-->>User: "result data"
else Socket error or timeout
Unity_Connection->>Socket: "close()"
Unity_Connection->>Port_Discovery: "discover_unity_port() (rediscover)"
Port_Discovery-->>Unity_Connection: "new port (if changed)"
MCP_Server->>MCP_Server: "exponential backoff with jitter"
Note right of MCP_Server: Retry with new connection
end
end
end
3 files reviewed, 1 comment
| with contextlib.suppress(Exception): | ||
| msg = b'Unity MCP requires FRAMING=1' | ||
| header = struct.pack('>Q', len(msg)) | ||
| self.sock.sendall(header + msg) |
There was a problem hiding this comment.
style: The framed message sent here may not be properly handled by a non-framing peer, potentially causing connection issues. Consider logging when this advisory message is sent.
Summary
Testing
pytest tests/test_logging_stdout.py tests/test_transport_framing.py -qpytest -qhttps://chatgpt.com/codex/tasks/task_e_68a28f8913fc832789efbfb1ceb3af72