Skip to content

Commit ca01fc7

Browse files
authored
Server: Robust shutdown on stdio detach (signals, stdin/parent monitor, forced exit) (#363)
* Server: robust shutdown on stdio detach (signals, stdin/parent monitor, forced exit)\nTests: move telemetry tests to tests/ and convert to asserts * Server: simplify _force_exit to os._exit; guard exit timers to avoid duplicates; fix Windows ValueError in parent monitor; tests: add autouse cwd fixture for telemetry to locate pyproject.toml * Server: add DEBUG logs for transient stdin checks and monitor thread errors * Mirror shutdown improvements: signal handlers, stdin/parent monitor, guarded exit timers, and os._exit force-exit in UnityMcpServer~ entry points
1 parent 040eb6d commit ca01fc7

File tree

5 files changed

+357
-164
lines changed

5 files changed

+357
-164
lines changed

MCPForUnity/UnityMcpServer~/src/server.py

Lines changed: 97 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44
from logging.handlers import RotatingFileHandler
55
import os
66
from contextlib import asynccontextmanager
7+
import signal
8+
import sys
9+
import threading
710
from typing import AsyncIterator, Dict, Any
811
from config import config
912
from tools import register_all_tools
@@ -64,6 +67,10 @@
6467
# Global connection state
6568
_unity_connection: UnityConnection = None
6669

70+
# Global shutdown coordination
71+
_shutdown_flag = threading.Event()
72+
_exit_timer_scheduled = threading.Event()
73+
6774

6875
@asynccontextmanager
6976
async def server_lifespan(server: FastMCP) -> AsyncIterator[Dict[str, Any]]:
@@ -186,9 +193,98 @@ def _emit_startup():
186193
register_all_resources(mcp)
187194

188195

196+
def _force_exit(code: int = 0):
197+
"""Force process exit, bypassing any background threads that might linger."""
198+
os._exit(code)
199+
200+
201+
def _signal_handler(signum, frame):
202+
logger.info(f"Received signal {signum}, initiating shutdown...")
203+
_shutdown_flag.set()
204+
if not _exit_timer_scheduled.is_set():
205+
_exit_timer_scheduled.set()
206+
threading.Timer(1.0, _force_exit, args=(0,)).start()
207+
208+
209+
def _monitor_stdin():
210+
"""Background thread to detect stdio detach (stdin EOF) or parent exit."""
211+
try:
212+
parent_pid = os.getppid() if hasattr(os, "getppid") else None
213+
while not _shutdown_flag.is_set():
214+
if _shutdown_flag.wait(0.5):
215+
break
216+
217+
if parent_pid is not None:
218+
try:
219+
os.kill(parent_pid, 0)
220+
except ValueError:
221+
# Signal 0 unsupported on this platform (e.g., Windows); disable parent probing
222+
parent_pid = None
223+
except (ProcessLookupError, OSError):
224+
logger.info(f"Parent process {parent_pid} no longer exists; shutting down")
225+
break
226+
227+
try:
228+
if sys.stdin.closed:
229+
logger.info("stdin.closed is True; client disconnected")
230+
break
231+
fd = sys.stdin.fileno()
232+
if fd < 0:
233+
logger.info("stdin fd invalid; client disconnected")
234+
break
235+
except (ValueError, OSError, AttributeError):
236+
# Closed pipe or unavailable stdin
237+
break
238+
except Exception:
239+
# Ignore transient errors
240+
logger.debug("Transient error checking stdin", exc_info=True)
241+
242+
if not _shutdown_flag.is_set():
243+
logger.info("Client disconnected (stdin or parent), initiating shutdown...")
244+
_shutdown_flag.set()
245+
if not _exit_timer_scheduled.is_set():
246+
_exit_timer_scheduled.set()
247+
threading.Timer(0.5, _force_exit, args=(0,)).start()
248+
except Exception:
249+
# Never let monitor thread crash the process
250+
logger.debug("Monitor thread error", exc_info=True)
251+
252+
189253
def main():
190254
"""Entry point for uvx and console scripts."""
191-
mcp.run(transport='stdio')
255+
try:
256+
signal.signal(signal.SIGTERM, _signal_handler)
257+
signal.signal(signal.SIGINT, _signal_handler)
258+
if hasattr(signal, "SIGPIPE"):
259+
signal.signal(signal.SIGPIPE, signal.SIG_IGN)
260+
if hasattr(signal, "SIGBREAK"):
261+
signal.signal(signal.SIGBREAK, _signal_handler)
262+
except Exception:
263+
# Signals can fail in some environments
264+
pass
265+
266+
t = threading.Thread(target=_monitor_stdin, daemon=True)
267+
t.start()
268+
269+
try:
270+
mcp.run(transport='stdio')
271+
logger.info("FastMCP run() returned (stdin EOF or disconnect)")
272+
except (KeyboardInterrupt, SystemExit):
273+
logger.info("Server interrupted; shutting down")
274+
_shutdown_flag.set()
275+
except BrokenPipeError:
276+
logger.info("Broken pipe; shutting down")
277+
_shutdown_flag.set()
278+
except Exception as e:
279+
logger.error(f"Server error: {e}", exc_info=True)
280+
_shutdown_flag.set()
281+
_force_exit(1)
282+
finally:
283+
_shutdown_flag.set()
284+
logger.info("Server main loop exited")
285+
if not _exit_timer_scheduled.is_set():
286+
_exit_timer_scheduled.set()
287+
threading.Timer(0.5, _force_exit, args=(0,)).start()
192288

193289

194290
# Run the server

Server/server.py

Lines changed: 97 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,9 @@
44
from logging.handlers import RotatingFileHandler
55
import os
66
from contextlib import asynccontextmanager
7+
import sys
8+
import signal
9+
import threading
710
from typing import AsyncIterator, Dict, Any
811
from config import config
912
from tools import register_all_tools
@@ -64,6 +67,10 @@
6467
# Global connection state
6568
_unity_connection: UnityConnection = None
6669

70+
# Global shutdown coordination
71+
_shutdown_flag = threading.Event()
72+
_exit_timer_scheduled = threading.Event()
73+
6774

6875
@asynccontextmanager
6976
async def server_lifespan(server: FastMCP) -> AsyncIterator[Dict[str, Any]]:
@@ -186,9 +193,98 @@ def _emit_startup():
186193
register_all_resources(mcp)
187194

188195

196+
def _force_exit(code: int = 0):
197+
"""Force process exit, bypassing any background threads that might linger."""
198+
os._exit(code)
199+
200+
201+
def _signal_handler(signum, frame):
202+
logger.info(f"Received signal {signum}, initiating shutdown...")
203+
_shutdown_flag.set()
204+
if not _exit_timer_scheduled.is_set():
205+
_exit_timer_scheduled.set()
206+
threading.Timer(1.0, _force_exit, args=(0,)).start()
207+
208+
209+
def _monitor_stdin():
210+
"""Background thread to detect stdio detach (stdin EOF) or parent exit."""
211+
try:
212+
parent_pid = os.getppid() if hasattr(os, "getppid") else None
213+
while not _shutdown_flag.is_set():
214+
if _shutdown_flag.wait(0.5):
215+
break
216+
217+
if parent_pid is not None:
218+
try:
219+
os.kill(parent_pid, 0)
220+
except ValueError:
221+
# Signal 0 unsupported on this platform (e.g., Windows); disable parent probing
222+
parent_pid = None
223+
except (ProcessLookupError, OSError):
224+
logger.info(f"Parent process {parent_pid} no longer exists; shutting down")
225+
break
226+
227+
try:
228+
if sys.stdin.closed:
229+
logger.info("stdin.closed is True; client disconnected")
230+
break
231+
fd = sys.stdin.fileno()
232+
if fd < 0:
233+
logger.info("stdin fd invalid; client disconnected")
234+
break
235+
except (ValueError, OSError, AttributeError):
236+
# Closed pipe or unavailable stdin
237+
break
238+
except Exception:
239+
# Ignore transient errors
240+
logger.debug("Transient error checking stdin", exc_info=True)
241+
242+
if not _shutdown_flag.is_set():
243+
logger.info("Client disconnected (stdin or parent), initiating shutdown...")
244+
_shutdown_flag.set()
245+
if not _exit_timer_scheduled.is_set():
246+
_exit_timer_scheduled.set()
247+
threading.Timer(0.5, _force_exit, args=(0,)).start()
248+
except Exception:
249+
# Never let monitor thread crash the process
250+
logger.debug("Monitor thread error", exc_info=True)
251+
252+
189253
def main():
190254
"""Entry point for uvx and console scripts."""
191-
mcp.run(transport='stdio')
255+
try:
256+
signal.signal(signal.SIGTERM, _signal_handler)
257+
signal.signal(signal.SIGINT, _signal_handler)
258+
if hasattr(signal, "SIGPIPE"):
259+
signal.signal(signal.SIGPIPE, signal.SIG_IGN)
260+
if hasattr(signal, "SIGBREAK"):
261+
signal.signal(signal.SIGBREAK, _signal_handler)
262+
except Exception:
263+
# Signals can fail in some environments
264+
pass
265+
266+
t = threading.Thread(target=_monitor_stdin, daemon=True)
267+
t.start()
268+
269+
try:
270+
mcp.run(transport='stdio')
271+
logger.info("FastMCP run() returned (stdin EOF or disconnect)")
272+
except (KeyboardInterrupt, SystemExit):
273+
logger.info("Server interrupted; shutting down")
274+
_shutdown_flag.set()
275+
except BrokenPipeError:
276+
logger.info("Broken pipe; shutting down")
277+
_shutdown_flag.set()
278+
except Exception as e:
279+
logger.error(f"Server error: {e}", exc_info=True)
280+
_shutdown_flag.set()
281+
_force_exit(1)
282+
finally:
283+
_shutdown_flag.set()
284+
logger.info("Server main loop exited")
285+
if not _exit_timer_scheduled.is_set():
286+
_exit_timer_scheduled.set()
287+
threading.Timer(0.5, _force_exit, args=(0,)).start()
192288

193289

194290
# Run the server

0 commit comments

Comments
 (0)