feat: Complete remaining improvements

MEDIUM PRIORITY:
4. restore operations - use haproxy_cmd_checked()
   - restore_servers_from_config() now validates commands
   - haproxy_restore_state() logs warnings for individual failures

5. container health check - add exception logging
   - logger.warning() for unexpected exceptions

6. haproxy_add_domain - optimize map file reading
   - Read map contents once, reuse for domain check and pool lookup
   - Eliminates redundant file reads

LOW PRIORITY:
7. Remove unused IP_PATTERN constant
   - Was unused since validate_ip() uses ipaddress module

8. haproxy_add_server - auto-select slot
   - slot=0 or slot=-1 finds first available slot
   - Returns error if all slots in use

9. haproxy_wait_drain - new tool
   - Wait for connections to drain before maintenance
   - Polls HAProxy stats until scur=0 or timeout

MCP Tools: 21 → 22

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

mcp/server.py

@@ -69,7 +69,6 @@ DOMAIN_PATTERN = re.compile(
     r'^[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?'
     r'(\.[a-zA-Z0-9]([a-zA-Z0-9\-]{0,61}[a-zA-Z0-9])?)*$'
 )
-IP_PATTERN = re.compile(r'^(\d{1,3}\.){3}\d{1,3}$')
 # Backend/server names: alphanumeric, underscore, hyphen only
 BACKEND_NAME_PATTERN = re.compile(r'^[a-zA-Z0-9_-]+$')
 # Pattern for converting domain to backend name
@@ -673,14 +672,14 @@ def restore_servers_from_config() -> int:
                 logger.warning("Invalid port for %s slot %d, skipping", domain, slot)
                 continue
 
-            try:
             for suffix, port in get_server_suffixes(http_port):
                 server = f"{server_prefix}{suffix}_{slot}"
-                    haproxy_cmd(f"set server {backend}/{server} addr {ip} port {port}")
-                    haproxy_cmd(f"set server {backend}/{server} state ready")
+                try:
+                    haproxy_cmd_checked(f"set server {backend}/{server} addr {ip} port {port}")
+                    haproxy_cmd_checked(f"set server {backend}/{server} state ready")
                     restored += 1
                 except HaproxyError as e:
-                logger.warning("Failed to restore %s slot %d: %s", domain, slot, e)
+                    logger.warning("Failed to restore %s/%s: %s", backend, server, e)
 
     return restored
 
@@ -777,21 +776,33 @@ def haproxy_add_domain(domain: str, ip: str = "", http_port: int = 80) -> str:
     if not (1 <= http_port <= 65535):
         return "Error: Port must be between 1 and 65535"
 
-    # Check if domain already exists
-    existing_backend = get_domain_backend(domain)
-    if existing_backend:
-        return f"Error: Domain {domain} already exists (mapped to {existing_backend})"
+    # Read map contents once for both existence check and pool lookup
+    entries = get_map_contents()
 
-    # Find available pool
-    try:
-        pool = find_available_pool()
-    except NoAvailablePoolError as e:
-        return f"Error: {e}"
+    # Check if domain already exists (using cached entries)
+    for domain_entry, backend in entries:
+        if domain_entry == domain:
+            return f"Error: Domain {domain} already exists (mapped to {backend})"
+
+    # Find available pool (using cached entries)
+    used_pools: Set[str] = set()
+    for _, backend in entries:
+        if backend.startswith("pool_"):
+            used_pools.add(backend)
+
+    pool = None
+    for i in range(1, POOL_COUNT + 1):
+        pool_name = f"pool_{i}"
+        if pool_name not in used_pools:
+            pool = pool_name
+            break
+    if not pool:
+        return f"Error: All {POOL_COUNT} pool backends are in use"
 
     try:
         # Save to disk first (atomic write for persistence)
         # If HAProxy update fails after this, state will be correct on restart
-        entries = get_map_contents()
+        # Note: We already have 'entries' from the map contents read above
         entries.append((domain, pool))
         entries.append((f".{domain}", pool))
         try:
@@ -948,7 +959,7 @@ def haproxy_add_server(domain: str, slot: int, ip: str, http_port: int = 80) ->
 
     Args:
         domain: The domain name to add the server to
-        slot: Server slot number (1-10). Use different slots for multiple servers
+        slot: Server slot number (1-10), or 0/-1 for auto-select first available slot
         ip: IP address of the server (required)
         http_port: HTTP port (default: 80)
 
@@ -962,6 +973,9 @@ def haproxy_add_server(domain: str, slot: int, ip: str, http_port: int = 80) ->
 
         # Add a third server
         haproxy_add_server("api.example.com", slot=3, ip="10.0.0.3", http_port=8080)
+
+        # Auto-select next available slot
+        haproxy_add_server("api.example.com", slot=0, ip="10.0.0.4", http_port=8080)
     """
     if not validate_domain(domain):
         return "Error: Invalid domain format"
@@ -969,14 +983,35 @@ def haproxy_add_server(domain: str, slot: int, ip: str, http_port: int = 80) ->
         return "Error: IP address is required"
     if not validate_ip(ip):
         return "Error: Invalid IP address format"
-    if not (1 <= slot <= MAX_SLOTS):
-        return f"Error: Slot must be between 1 and {MAX_SLOTS}"
     if not (1 <= http_port <= 65535):
         return "Error: Port must be between 1 and 65535"
 
     try:
         backend, server_prefix = get_backend_and_prefix(domain)
 
+        # Auto-select slot if slot <= 0
+        if slot <= 0:
+            state = haproxy_cmd("show servers state")
+            used_slots: Set[int] = set()
+            for line in state.split("\n"):
+                parts = line.split()
+                if len(parts) >= STATE_MIN_COLUMNS and parts[StateField.BE_NAME] == backend:
+                    if parts[StateField.SRV_ADDR] != "0.0.0.0":
+                        # Extract slot number from server name (e.g., pool_1_3 -> 3)
+                        server_name = parts[StateField.SRV_NAME]
+                        try:
+                            used_slots.add(int(server_name.rsplit("_", 1)[1]))
+                        except (ValueError, IndexError):
+                            pass
+            for s in range(1, MAX_SLOTS + 1):
+                if s not in used_slots:
+                    slot = s
+                    break
+            else:
+                return f"Error: No available slots (all {MAX_SLOTS} slots in use)"
+        elif not (1 <= slot <= MAX_SLOTS):
+            return f"Error: Slot must be between 1 and {MAX_SLOTS}, or 0/-1 for auto-select"
+
         # Save to persistent config FIRST (disk-first pattern)
         add_server_to_config(domain, slot, ip, http_port)
 
@@ -1242,6 +1277,60 @@ def haproxy_set_domain_state(domain: str, state: str) -> str:
     return result
 
 
+@mcp.tool()
+def haproxy_wait_drain(domain: str, timeout: int = 30) -> str:
+    """Wait for all active connections to drain from a domain's servers.
+
+    Useful after setting servers to 'drain' state before maintenance.
+    Polls every second until all servers have 0 current connections or timeout.
+
+    Args:
+        domain: The domain name to wait for
+        timeout: Maximum seconds to wait (default: 30, max: 300)
+
+    Returns:
+        Success message or timeout error
+
+    Example:
+        haproxy_set_domain_state("api.example.com", "drain")
+        haproxy_wait_drain("api.example.com", timeout=60)
+        # Now safe to perform maintenance
+    """
+    if not validate_domain(domain):
+        return "Error: Invalid domain format"
+    if not (1 <= timeout <= 300):
+        return "Error: Timeout must be between 1 and 300 seconds"
+
+    try:
+        backend, _ = get_backend_and_prefix(domain)
+    except ValueError as e:
+        return f"Error: {e}"
+
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            stats = haproxy_cmd("show stat")
+            total_connections = 0
+            for line in stats.split("\n"):
+                parts = line.split(",")
+                if len(parts) > StatField.SCUR and parts[0] == backend and parts[1] not in ["FRONTEND", "BACKEND", ""]:
+                    try:
+                        scur = int(parts[StatField.SCUR]) if parts[StatField.SCUR] else 0
+                        total_connections += scur
+                    except ValueError:
+                        pass
+
+            if total_connections == 0:
+                elapsed = int(time.time() - start_time)
+                return f"All connections drained for {domain} ({elapsed}s)"
+
+            time.sleep(1)
+        except HaproxyError as e:
+            return f"Error checking connections: {e}"
+
+    return f"Timeout: Connections still active after {timeout}s"
+
+
 @mcp.tool()
 def haproxy_stats() -> str:
     """Get HAProxy status and statistics.
@@ -1327,6 +1416,7 @@ def haproxy_health() -> str:
         result["components"]["container"] = {"status": "timeout"}
         result["status"] = "unhealthy"
     except Exception as e:
+        logger.warning("Container health check failed: %s", e)
         result["components"]["container"] = {"status": "error", "error": str(e)}
 
     # Check configuration files
@@ -1730,9 +1820,12 @@ def haproxy_restore_state() -> str:
                     skipped += 1
                     continue
 
-                haproxy_cmd(f"set server {backend}/{server} addr {addr} port {port}")
-                haproxy_cmd(f"set server {backend}/{server} state ready")
+                try:
+                    haproxy_cmd_checked(f"set server {backend}/{server} addr {addr} port {port}")
+                    haproxy_cmd_checked(f"set server {backend}/{server} state ready")
                     restored += 1
+                except HaproxyError as e:
+                    logger.warning("Failed to restore %s/%s: %s", backend, server, e)
 
         result = f"Server state restored ({restored} servers)"
         if skipped: