|
4 | 4 | from logging.handlers import RotatingFileHandler |
5 | 5 | import os |
6 | 6 | from contextlib import asynccontextmanager |
| 7 | +import signal |
| 8 | +import sys |
| 9 | +import threading |
7 | 10 | from typing import AsyncIterator, Dict, Any |
8 | 11 | from config import config |
9 | 12 | from tools import register_all_tools |
|
64 | 67 | # Global connection state |
65 | 68 | _unity_connection: UnityConnection = None |
66 | 69 |
|
| 70 | +# Global shutdown coordination |
| 71 | +_shutdown_flag = threading.Event() |
| 72 | +_exit_timer_scheduled = threading.Event() |
| 73 | + |
67 | 74 |
|
68 | 75 | @asynccontextmanager |
69 | 76 | async def server_lifespan(server: FastMCP) -> AsyncIterator[Dict[str, Any]]: |
@@ -186,9 +193,98 @@ def _emit_startup(): |
186 | 193 | register_all_resources(mcp) |
187 | 194 |
|
188 | 195 |
|
| 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 | + |
189 | 253 | def main(): |
190 | 254 | """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() |
192 | 288 |
|
193 | 289 |
|
194 | 290 | # Run the server |
|
0 commit comments