fix: keep Ableton tools available across MCP clients

MCP_Server/cache/browser.py

@@ -11,7 +11,7 @@
 import time
 import logging
 import threading
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Optional
 from collections import deque
 
 import MCP_Server.state as state
@@ -166,7 +166,10 @@ def load_browser_cache_from_disk() -> bool:
 # Live browser scan
 # ---------------------------------------------------------------------------
 
-def populate_browser_cache(force: bool = False) -> bool:
+def populate_browser_cache(
+    force: bool = False,
+    stop_event: Optional[threading.Event] = None,
+) -> bool:
     """Scan Ableton's browser tree and cache all items for instant search.
 
     Uses a breadth-first walk up to depth 3 across 11 browser categories.
@@ -176,7 +179,7 @@ def populate_browser_cache(force: bool = False) -> bool:
     Uses a **dedicated TCP connection** to avoid corrupting the shared global
     connection when the BFS scan sends many rapid commands.
     """
-    from MCP_Server.connections.ableton import AbletonConnection
+    from MCP_Server.connections.ableton import AbletonConnection, CommandCancelled
 
     now = time.time()
     with state.browser_cache_lock:
@@ -191,6 +194,8 @@ def populate_browser_cache(force: bool = False) -> bool:
     ableton = AbletonConnection(host="localhost", port=9877)
 
     try:
+        if stop_event is not None and stop_event.is_set():
+            return False
         try:
             if not ableton.connect():
                 logger.warning("Browser cache: cannot connect to Ableton")
@@ -199,27 +204,48 @@ def populate_browser_cache(force: bool = False) -> bool:
             logger.warning("Browser cache: cannot connect to Ableton: %s", e)
             return False
 
+        if stop_event is not None and stop_event.is_set():
+            return False
+
         logger.info("Browser cache: starting scan...")
         flat_items: List[Dict[str, Any]] = []
         by_display: Dict[str, List[Dict[str, Any]]] = {}
         total = 0
 
         for path_root, display_name in BROWSER_CATEGORIES:
+            if stop_event is not None and stop_event.is_set():
+                return False
             category_items: List[Dict[str, Any]] = []
             cat_count = 0
 
             # BFS queue: (browser_path, depth)
             queue = deque([(path_root, 0)])
 
             while queue and cat_count < BROWSER_CACHE_MAX_ITEMS:
+                if stop_event is not None and stop_event.is_set():
+                    return False
                 current_path, depth = queue.popleft()
 
                 try:
-                    result = ableton.send_command("get_browser_items_at_path", {"path": current_path}, timeout=60.0)
+                    result = ableton.send_command(
+                        "get_browser_items_at_path",
+                        {"path": current_path},
+                        timeout=60.0,
+                        stop_event=stop_event,
+                    )
+                except CommandCancelled:
+                    logger.info("Browser cache scan cancelled during shutdown")
+                    return False
                 except Exception as e:
+                    if stop_event is not None and stop_event.is_set():
+                        return False
                     logger.warning("Browser cache: failed to read '%s': %s", current_path, e)
                     # Try to re-establish connection before continuing
-                    time.sleep(2)
+                    if stop_event is not None:
+                        if stop_event.wait(2.0):
+                            return False
+                    else:
+                        time.sleep(2.0)
                     try:
                         ableton.disconnect()
                         if not ableton.connect():
@@ -262,13 +288,19 @@ def populate_browser_cache(force: bool = False) -> bool:
                         queue.append((item_path, depth + 1))
 
                 # Rate-limit to avoid overwhelming Ableton's socket handler
-                time.sleep(0.01)
+                if stop_event is not None:
+                    if stop_event.wait(0.01):
+                        return False
+                else:
+                    time.sleep(0.01)
 
             by_display[display_name] = category_items
             logger.info("Browser cache: '%s' — %d items", display_name, len(category_items))
 
-        device_map = build_device_uri_map(flat_items)
+        if stop_event is not None and stop_event.is_set():
+            return False
 
+        device_map = build_device_uri_map(flat_items)
         with state.browser_cache_lock:
             state.browser_cache_flat = flat_items
             state.browser_cache_by_category = by_display

MCP_Server/connections/ableton.py

@@ -1,8 +1,9 @@
 """AbletonConnection — TCP socket connection to the Ableton Remote Script."""
 
-import socket
 import json
 import logging
+import select
+import socket
 import time
 import threading
 from dataclasses import dataclass
@@ -24,6 +25,10 @@
 ])
 
 
+class CommandCancelled(RuntimeError):
+    """Raised when cooperative shutdown cancels an in-flight command."""
+
+
 @dataclass
 class AbletonConnection:
     host: str
@@ -35,13 +40,15 @@ class AbletonConnection:
     def connect(self) -> bool:
         """Connect to the Ableton Remote Script socket server"""
         if self.sock:
+            self._last_socket_open = True
             return True
 
         try:
             self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
             self.sock.settimeout(5.0)
             self.sock.connect((self.host, self.port))
             self._recv_buffer = ""  # Clear buffer on new connection
+            self._last_socket_open = True
             logger.info("Connected to Ableton at %s:%s", self.host, self.port)
             return True
         except Exception as e:
@@ -52,6 +59,7 @@ def connect(self) -> bool:
                 except Exception:
                     pass
             self.sock = None
+            self._last_socket_open = False
             return False
 
     def disconnect(self):
@@ -63,6 +71,7 @@ def disconnect(self):
                 logger.error("Error disconnecting from Ableton: %s", e)
             finally:
                 self.sock = None
+                self._last_socket_open = False
         if self._udp_sock:
             try:
                 self._udp_sock.close()
@@ -72,8 +81,48 @@ def disconnect(self):
                 self._udp_sock = None
 
     def __post_init__(self):
+        """Initialize per-connection receive buffering and send serialization."""
         self._recv_buffer = ""
         self._send_lock = threading.Lock()
+        self._last_socket_open = self.sock is not None
+
+    def is_connected(self) -> bool:
+        """Passively check whether the Remote Script socket is still open.
+
+        The check never sends or consumes protocol data. If a command currently
+        owns the send lock, return the last verified socket result so a status
+        request cannot interfere with its response handling.
+        """
+        sock = self.sock
+        if sock is None:
+            self._last_socket_open = False
+            return False
+
+        try:
+            sock.getpeername()
+        except OSError:
+            self._last_socket_open = False
+            return False
+
+        if not self._send_lock.acquire(blocking=False):
+            return self._last_socket_open
+
+        try:
+            readable, _writable, _exceptional = select.select([sock], [], [], 0)
+            if not readable:
+                self._last_socket_open = True
+                return True
+            try:
+                self._last_socket_open = bool(sock.recv(1, socket.MSG_PEEK))
+            except (BlockingIOError, socket.timeout):
+                self._last_socket_open = True
+            except OSError:
+                self._last_socket_open = False
+        except (OSError, ValueError):
+            self._last_socket_open = False
+        finally:
+            self._send_lock.release()
+        return self._last_socket_open
 
     def _ensure_udp_socket(self):
         """Create a UDP socket for real-time parameter sending if not already open."""
@@ -95,12 +144,22 @@ def send_udp_command(self, command_type: str, params: Dict[str, Any] = None):
         sock.sendto(payload, (self.host, self._udp_port))
         logger.debug("Sent UDP command: %s", command_type)
 
-    def receive_full_response(self, sock, buffer_size=8192, timeout=15.0):
+    def receive_full_response(
+        self,
+        sock,
+        buffer_size=8192,
+        timeout=15.0,
+        stop_event: Optional[threading.Event] = None,
+    ):
         """Receive a complete newline-delimited JSON response and return the parsed object"""
-        sock.settimeout(timeout)
+        deadline = time.monotonic() + timeout
+        sock.settimeout(min(timeout, 0.25) if stop_event is not None else timeout)
 
         try:
             while True:
+                if stop_event is not None and stop_event.is_set():
+                    raise CommandCancelled("Ableton command cancelled during shutdown")
+
                 # Check if we already have a complete line in the buffer
                 if '\n' in self._recv_buffer:
                     line, self._recv_buffer = self._recv_buffer.split('\n', 1)
@@ -115,18 +174,30 @@ def receive_full_response(self, sock, buffer_size=8192, timeout=15.0):
                         return result
 
                 try:
+                    if stop_event is not None:
+                        remaining = deadline - time.monotonic()
+                        if remaining <= 0:
+                            raise socket.timeout()
+                        sock.settimeout(min(0.25, remaining))
                     chunk = sock.recv(buffer_size)
                     if not chunk:
                         raise Exception("Connection closed before receiving any data")
 
                     self._recv_buffer += chunk.decode('utf-8')
                 except socket.timeout:
+                    if stop_event is not None:
+                        if stop_event.is_set():
+                            raise CommandCancelled(
+                                "Ableton command cancelled during shutdown"
+                            ) from None
+                        if time.monotonic() < deadline:
+                            continue
                     logger.warning("Socket timeout during receive")
                     raise
                 except (ConnectionError, BrokenPipeError, ConnectionResetError) as e:
                     logger.error("Socket connection error during receive: %s", e)
                     raise
-        except (socket.timeout, json.JSONDecodeError):
+        except (socket.timeout, json.JSONDecodeError, CommandCancelled):
             raise
         except Exception as e:
             logger.error("Error during receive: %s", e)
@@ -139,7 +210,13 @@ def _reconnect(self) -> bool:
         self._recv_buffer = ""
         return self.connect()
 
-    def send_command(self, command_type: str, params: Dict[str, Any] = None, timeout: Optional[float] = None) -> Dict[str, Any]:
+    def send_command(
+        self,
+        command_type: str,
+        params: Dict[str, Any] = None,
+        timeout: Optional[float] = None,
+        stop_event: Optional[threading.Event] = None,
+    ) -> Dict[str, Any]:
         """Send a command to Ableton and return the response.
 
         Includes automatic retry: if the first attempt fails due to a
@@ -165,8 +242,13 @@ def send_command(self, command_type: str, params: Dict[str, Any] = None, timeout
 
         for attempt in range(1, max_attempts + 1):
             with self._send_lock:
+                if stop_event is not None and stop_event.is_set():
+                    raise CommandCancelled("Ableton command cancelled during shutdown")
                 if not self.sock and not self.connect():
                     raise ConnectionError("Not connected to Ableton")
+                if stop_event is not None and stop_event.is_set():
+                    self.disconnect()
+                    raise CommandCancelled("Ableton command cancelled during shutdown")
 
                 command = {
                     "type": command_type,
@@ -181,7 +263,13 @@ def send_command(self, command_type: str, params: Dict[str, Any] = None, timeout
 
                     # Pre-delay: give Ableton time to process before we read the response
                     if pre_delay:
-                        time.sleep(pre_delay)
+                        if stop_event is not None:
+                            if stop_event.wait(pre_delay):
+                                raise CommandCancelled(
+                                    "Ableton command cancelled during shutdown"
+                                )
+                        else:
+                            time.sleep(pre_delay)
 
                     # Set timeout based on command type (caller override takes priority)
                     if timeout is None:
@@ -190,7 +278,11 @@ def send_command(self, command_type: str, params: Dict[str, Any] = None, timeout
                             command_type, 15.0 if is_modifying else 10.0
                         )
                     # Receive the response (already parsed by receive_full_response)
-                    response = self.receive_full_response(self.sock, timeout=timeout)
+                    response = self.receive_full_response(
+                        self.sock,
+                        timeout=timeout,
+                        stop_event=stop_event,
+                    )
                     logger.debug("Response status: %s", response.get('status', 'unknown'))
 
                     if response.get("status") == "error":
@@ -199,10 +291,20 @@ def send_command(self, command_type: str, params: Dict[str, Any] = None, timeout
 
                     # Post-delay: let Ableton settle before the next command
                     if post_delay:
-                        time.sleep(post_delay)
+                        if stop_event is not None:
+                            if stop_event.wait(post_delay):
+                                raise CommandCancelled(
+                                    "Ableton command cancelled during shutdown"
+                                )
+                        else:
+                            time.sleep(post_delay)
 
                     return response.get("result", {})
 
+                except CommandCancelled:
+                    self.disconnect()
+                    self._recv_buffer = ""
+                    raise
                 except Exception as e:
                     logger.error("Command '%s' attempt %d failed: %s", command_type, attempt, e)
                     # Close the broken socket and clear buffer
@@ -211,24 +313,30 @@ def send_command(self, command_type: str, params: Dict[str, Any] = None, timeout
 
                     if attempt < max_attempts:
                         # Wait briefly then retry with a fresh connection
-                        time.sleep(0.1)
+                        if stop_event is not None:
+                            if stop_event.wait(0.1):
+                                raise CommandCancelled(
+                                    "Ableton command cancelled during shutdown"
+                                ) from None
+                        else:
+                            time.sleep(0.1)
                         if not self.connect():
-                            raise ConnectionError("Failed to reconnect to Ableton")
+                            raise ConnectionError("Failed to reconnect to Ableton") from e
                         logger.info("Reconnected, retrying command...")
                     else:
-                        raise Exception(f"Command '{command_type}' failed after {max_attempts} attempts: {e}")
+                        raise Exception(
+                            f"Command '{command_type}' failed after {max_attempts} attempts: {e}"
+                        ) from e
 
 
 def get_ableton_connection():
     """Get or create a persistent Ableton connection"""
 
     if state.ableton_connection is not None:
         try:
-            # Test if the socket is still connected
-            if state.ableton_connection.sock is None:
-                raise ConnectionError("Socket is None")
-            state.ableton_connection.sock.settimeout(1.0)
-            state.ableton_connection.sock.getpeername()  # raises if disconnected
+            if not state.ableton_connection.is_connected():
+                raise ConnectionError("Socket is no longer connected")
+            state.ableton_connected_event.set()
             return state.ableton_connection
         except Exception as e:
             logger.warning("Existing connection is no longer valid: %s", e)