feat(#672): Add Keep Server Running toggle in Unity Editor Window and HTTP mode support #725
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -133,11 +133,15 @@ def doRollover(self): | |
| # In-memory custom tool service initialized after MCP construction | ||
| custom_tool_service: CustomToolService | None = None | ||
|
|
||
| # Per-session keep_server_running flag, indexed by session_id | ||
| # Set by middleware when Unity registers with the flag enabled | ||
| keep_server_running: dict[str, bool] = {} | ||
|
Comment on lines
+136
to
+138
Contributor
There was a problem hiding this comment.
The global As a result, 🔧 Option: query the registry directly at shutdown instead of maintaining a separate dict-# Per-session keep_server_running flag, indexed by session_id
-# Set by middleware when Unity registers with the flag enabled
-keep_server_running: dict[str, bool] = {}Then at shutdown (lines 256-271), query the registry: - has_keep_running = any(keep_server_running.values()) if keep_server_running else False
+ has_keep_running = False
+ if _plugin_registry is not None:
+ try:
+ sessions = await _plugin_registry.list_sessions()
+ has_keep_running = any(s.keep_server_running for s in sessions.values())
+ except Exception:
+ pass
if _unity_connection_pool and not has_keep_running:
_unity_connection_pool.disconnect_all()
shutdown_msg = "MCP for Unity Server shut down"
if has_keep_running:
shutdown_msg += " (Keep-Running mode: maintaining server for Unity reconnection)"
logger.info(shutdown_msg)
-
- # Note: keep_server_running state is per-session, not server-global.
- # Cleared when sessions expire, used by middleware to pass state.
- keep_server_running.clear()🤖 Prompt for AI Agents |
||
|
|
||
|
|
||
| @asynccontextmanager | ||
| async def server_lifespan(server: FastMCP) -> AsyncIterator[dict[str, Any]]: | ||
| """Handle server startup and shutdown.""" | ||
| global _unity_connection_pool, _server_version, keep_server_running | ||
| _server_version = get_package_version() | ||
| logger.info(f"MCP for Unity Server v{_server_version} starting up") | ||
|
|
||
|
|
@@ -183,9 +187,17 @@ def _emit_startup(): | |
| logger.info( | ||
| "Skipping Unity connection on startup (UNITY_MCP_SKIP_STARTUP_CONNECT=1)") | ||
| else: | ||
| # Initialize connection pool and discover instances (stdio mode only) | ||
| # In HTTP mode, PluginHub handles connections directly | ||
| instances = [] # Initialize to avoid UnboundLocalError in HTTP mode | ||
| if enable_http_server: | ||
| _unity_connection_pool = None # No legacy pool in HTTP mode | ||
| logger.info("HTTP mode enabled - PluginHub will manage Unity connections") | ||
| # Skip stdio-specific connection logic in HTTP mode. | ||
| # Lifespan must still continue to reach the context yield. | ||
| else: | ||
| _unity_connection_pool = get_unity_connection_pool() | ||
| instances = _unity_connection_pool.discover_all_instances() | ||
|
|
||
| if instances: | ||
| logger.info( | ||
|
|
@@ -244,9 +256,23 @@ def _emit_startup(): | |
| "plugin_registry": _plugin_registry, | ||
| } | ||
| finally: | ||
| # In HTTP mode, PluginHub manages connections independently. | ||
| # Only disconnect legacy pool if it was initialized (stdio mode). | ||
|
Comment on lines
258
to
+260
There was a problem hiding this comment.
|
||
| # Check if any session has keep_server_running enabled. | ||
| has_keep_running = any(keep_server_running.values()) if keep_server_running else False | ||
|
There was a problem hiding this comment.
The shutdown path now evaluates Useful? React with 👍 / 👎. |
||
|
|
||
| if _unity_connection_pool and not has_keep_running: | ||
| _unity_connection_pool.disconnect_all() | ||
|
|
||
| shutdown_msg = "MCP for Unity Server shut down" | ||
| if has_keep_running: | ||
| shutdown_msg += " (Keep-Running mode: maintaining server for Unity reconnection)" | ||
| logger.info(shutdown_msg) | ||
|
|
||
| # Note: keep_server_running state is per-session, not server-global. | ||
| # Cleared when sessions expire, used by middleware to pass state. | ||
| keep_server_running.clear() | ||
|
|
||
|
|
||
|
|
||
| def _build_instructions(project_scoped_tools: bool) -> str: | ||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -191,6 +191,13 @@ async def on_disconnect(self, websocket: WebSocket, close_code: int) -> None: | |||||
| ping_task.cancel() | ||||||
| # Clean up last pong tracking | ||||||
| cls._last_pong.pop(session_id, None) | ||||||
| # Clean up keep_server_running flag to avoid stale entries | ||||||
| try: | ||||||
| import main as main_module | ||||||
| main_module.keep_server_running.pop(session_id, None) | ||||||
| except (ImportError, AttributeError, KeyError): | ||||||
| # main module not yet imported or key already removed | ||||||
| pass | ||||||
| # Fail-fast any in-flight commands for this session to avoid waiting for COMMAND_TIMEOUT. | ||||||
| pending_ids = [ | ||||||
| command_id | ||||||
|
|
@@ -364,6 +371,7 @@ async def _handle_register(self, websocket: WebSocket, payload: RegisterMessage) | |||||
| project_hash = payload.project_hash | ||||||
| unity_version = payload.unity_version | ||||||
| project_path = payload.project_path | ||||||
| keep_server_running = getattr(payload, 'keep_server_running', False) | ||||||
|
There was a problem hiding this comment.
Useful? React with 👍 / 👎. There was a problem hiding this comment.
Suggested change
|
||||||
|
|
||||||
| if not project_hash: | ||||||
| await websocket.close(code=4400) | ||||||
|
|
@@ -378,7 +386,15 @@ async def _handle_register(self, websocket: WebSocket, payload: RegisterMessage) | |||||
| response = RegisteredMessage(session_id=session_id) | ||||||
| await websocket.send_json(response.model_dump()) | ||||||
|
|
||||||
| session = await registry.register( | ||||||
| session_id, project_name, project_hash, unity_version, project_path, | ||||||
| user_id=user_id, keep_server_running=keep_server_running | ||||||
| ) | ||||||
|
|
||||||
| # Store keep_server_running flag in global dict for shutdown logic. | ||||||
| # Import here to avoid circular dependency at module import time. | ||||||
| import main as main_module | ||||||
| main_module.keep_server_running[session.session_id] = keep_server_running | ||||||
| async with lock: | ||||||
| cls._connections[session.session_id] = websocket | ||||||
| # Initialize last pong time and start ping loop for this session | ||||||
|
|
||||||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
PR description mentions a Unity Editor UI toggle and sending a
keep_server_runningflag during WebSocket registration, but the only Unity-side change in this PR appears to be adding the EditorPrefs key. A repo-wide search shows no other references toEditorPrefKeys.KeepServerRunningor any C# code emittingkeep_server_runningin a registration message, so the feature looks incomplete or the description is ahead of the code.