crackerjack 0.31.10__py3-none-any.whl → 0.31.13__py3-none-any.whl

crackerjack/mcp/websocket/jobs.py

@@ -8,12 +8,17 @@ from typing import Any
 
 from rich.console import Console
 
+from crackerjack.core.timeout_manager import TimeoutStrategy, get_timeout_manager
+from crackerjack.services.input_validator import get_input_validator
+from crackerjack.services.secure_path_utils import SecurePathValidator
+
 console = Console()
 
 
 class JobManager:
     def __init__(self, progress_dir: Path) -> None:
-        self.progress_dir = progress_dir
+        # Validate and secure the progress directory path
+        self.progress_dir = SecurePathValidator.validate_safe_path(progress_dir)
         self.active_connections: dict[str, set[Any]] = {}
         self.known_jobs: set[str] = set()
         self.is_running = True
@@ -21,17 +26,18 @@ class JobManager:
         self.progress_dir.mkdir(exist_ok=True)
 
     def validate_job_id(self, job_id: str) -> bool:
+        """Validate job ID using secure input validator."""
         if not job_id:
             return False
 
+        # First check if it's a valid UUID
         with suppress(ValueError):
             uuid.UUID(job_id)
             return True
 
-        import re
-
-        # Performance: Use simpler regex pattern without whitespace
-        return bool(re.match(r"^[a-zA-Z0-9_-]+$", job_id) and len(job_id) <= 50)
+        # Use secure input validator for additional validation
+        result = get_input_validator().validate_job_id(job_id)
+        return result.valid
 
     def add_connection(self, job_id: str, websocket: Any) -> None:
         if job_id not in self.active_connections:
@@ -48,18 +54,79 @@ class JobManager:
         if job_id not in self.active_connections:
             return
 
+        timeout_manager = get_timeout_manager()
         connections = self.active_connections[job_id].copy()
+
+        # Create websocket send tasks
+        send_tasks = self._create_broadcast_tasks(connections, timeout_manager, data)
+
+        # Execute broadcast with timeout handling
+        if send_tasks:
+            await self._execute_broadcast_tasks(job_id, send_tasks)
+
+    def _create_broadcast_tasks(
+        self, connections: set, timeout_manager, data: dict
+    ) -> list:
+        """Create tasks for all websocket sends with timeout."""
+        send_tasks = []
         for websocket in connections:
-            try:
-                await websocket.send_json(data)
-            except Exception:
+            task = asyncio.create_task(
+                timeout_manager.with_timeout(
+                    "websocket_broadcast",
+                    websocket.send_json(data),
+                    timeout=2.0,  # Quick broadcast timeout
+                )
+            )
+            send_tasks.append((websocket, task))
+        return send_tasks
+
+    async def _execute_broadcast_tasks(self, job_id: str, send_tasks: list) -> None:
+        """Execute broadcast tasks with timeout and error handling."""
+        try:
+            # Use asyncio.wait with timeout for batch sending
+            done, pending = await asyncio.wait(
+                [task for _, task in send_tasks],
+                timeout=5.0,  # Overall timeout for all broadcasts
+                return_when=asyncio.ALL_COMPLETED,
+            )
+
+            # Handle completed and pending tasks
+            await self._handle_broadcast_results(job_id, send_tasks, done, pending)
+
+        except Exception as e:
+            console.print(f"[red]Broadcast error: {e}[/red]")
+            await self._cleanup_failed_broadcast(job_id, send_tasks)
+
+    async def _handle_broadcast_results(
+        self, job_id: str, send_tasks: list, done: set, pending: set
+    ) -> None:
+        """Handle results of broadcast tasks."""
+        # Cancel any pending tasks and remove failed connections
+        for websocket, task in send_tasks:
+            if task in pending:
+                task.cancel()
                 self.remove_connection(job_id, websocket)
+            elif task in done:
+                try:
+                    await task
+                except Exception:
+                    self.remove_connection(job_id, websocket)
+
+        # Wait for cancelled tasks to complete
+        if pending:
+            await asyncio.gather(*pending, return_exceptions=True)
+
+    async def _cleanup_failed_broadcast(self, job_id: str, send_tasks: list) -> None:
+        """Clean up connections after broadcast failure."""
+        for websocket, task in send_tasks:
+            if not task.done():
+                task.cancel()
+            self.remove_connection(job_id, websocket)
 
     def get_latest_job_id(self) -> str | None:
         if not self.progress_dir.exists():
             return None
 
-        # Performance: Use more specific glob pattern to reduce filesystem calls
         progress_files = list(self.progress_dir.glob("job-*.json"))
         if not progress_files:
             return None
@@ -68,88 +135,158 @@ class JobManager:
         return self.extract_job_id_from_file(latest_file)
 
     def extract_job_id_from_file(self, progress_file: Path) -> str | None:
-        # Performance: Use slice instead of replace for fixed prefix removal
         return (
-            progress_file.stem[4:]  # Remove "job-" prefix (4 chars)
-            if progress_file.stem.startswith("job-")
-            else None
+            progress_file.stem[4:] if progress_file.stem.startswith("job -") else None
         )
 
     def get_job_progress(self, job_id: str) -> dict | None:
         if not self.validate_job_id(job_id):
             return None
 
-        progress_file = self.progress_dir / f"job-{job_id}.json"
-        if not progress_file.exists():
-            return None
-
+        # Use secure path joining to create progress file path
         try:
+            progress_file = SecurePathValidator.secure_path_join(
+                self.progress_dir, f"job-{job_id}.json"
+            )
+            if not progress_file.exists():
+                return None
+
+            # Validate file size before reading
+            SecurePathValidator.validate_file_size(progress_file)
+
             return json.loads(progress_file.read_text())
         except (json.JSONDecodeError, OSError):
             return None
 
     async def _process_progress_file(self, progress_file: Path) -> None:
-        """Process a single progress file and handle new job detection."""
-        job_id = self.extract_job_id_from_file(progress_file)
+        # Validate the progress file path is within our allowed directory
+        try:
+            validated_file = SecurePathValidator.validate_safe_path(
+                progress_file, self.progress_dir
+            )
+        except Exception:
+            # If path validation fails, skip processing this file
+            return
+
+        job_id = self.extract_job_id_from_file(validated_file)
         if not (job_id and self.validate_job_id(job_id)):
             return
 
         progress_data = self.get_job_progress(job_id)
         if progress_data and job_id not in self.known_jobs:
             self.known_jobs.add(job_id)
-            console.print(f"[green]New job detected: {job_id}[/green]")
+            console.print(f"[green]New job detected: {job_id}[/ green]")
             await self.broadcast_to_job(job_id, progress_data)
 
     async def _monitor_directory_changes(self) -> None:
-        """Monitor the progress directory for new job files."""
+        timeout_manager = get_timeout_manager()
+        consecutive_errors = 0
+        max_consecutive_errors = 5
+
         while self.is_running:
             try:
-                if self.progress_dir.exists():
-                    for progress_file in self.progress_dir.glob("job-*.json"):
-                        await self._process_progress_file(progress_file)
+                # Monitor directory changes with timeout protection
+                async with timeout_manager.timeout_context(
+                    "file_operations",
+                    timeout=10.0,  # Timeout for directory monitoring cycle
+                    strategy=TimeoutStrategy.GRACEFUL_DEGRADATION,
+                ):
+                    if self.progress_dir.exists():
+                        # Process files with individual timeouts
+                        for progress_file in self.progress_dir.glob("job-*.json"):
+                            try:
+                                await timeout_manager.with_timeout(
+                                    "file_operations",
+                                    self._process_progress_file(progress_file),
+                                    timeout=5.0,  # Per-file timeout
+                                )
+                            except Exception as e:
+                                console.print(
+                                    f"[yellow]File processing error: {e}[/yellow]"
+                                )
+                                continue
+
+                    # Reset error count on successful cycle
+                    consecutive_errors = 0
+                    await asyncio.sleep(1)
 
-                await asyncio.sleep(1)  # Check every second
             except Exception as e:
+                consecutive_errors += 1
                 console.print(f"[red]Progress monitoring error: {e}[/red]")
-                await asyncio.sleep(5)  # Wait longer on error
+
+                # Implement exponential backoff for repeated errors
+                if consecutive_errors >= max_consecutive_errors:
+                    console.print(
+                        f"[red]Too many consecutive errors ({consecutive_errors}), stopping monitor[/red]"
+                    )
+                    break
+
+                # Exponential backoff with max delay
+                delay = min(5 * (2 ** (consecutive_errors - 1)), 60)
+                await asyncio.sleep(delay)
 
     async def monitor_progress_files(self) -> None:
         from crackerjack.mcp.file_monitor import create_progress_monitor
 
         console.print("[blue]Starting progress file monitoring...[/blue]")
+        timeout_manager = get_timeout_manager()
 
         try:
-            monitor = create_progress_monitor(self.progress_dir)
-            await monitor.start()
-
-            def on_progress_update(job_id: str, progress_data: dict) -> None:
-                """Callback for when progress files are updated."""
-                if job_id and self.validate_job_id(job_id):
-                    # Schedule the broadcast in the event loop
-                    asyncio.create_task(self.broadcast_to_job(job_id, progress_data))
-
-                    if job_id not in self.known_jobs:
-                        self.known_jobs.add(job_id)
-                        console.print(f"[green]New job detected: {job_id}[/green]")
-
-            # Monitor for new job files by checking the directory periodically
-            await self._monitor_directory_changes()
+            # Start file monitor with timeout protection
+            async with timeout_manager.timeout_context(
+                "file_operations",
+                timeout=30.0,  # Monitor startup timeout
+                strategy=TimeoutStrategy.GRACEFUL_DEGRADATION,
+            ):
+                monitor = create_progress_monitor(self.progress_dir)
+                await monitor.start()
+
+                def on_progress_update(job_id: str, progress_data: dict) -> None:
+                    if job_id and self.validate_job_id(job_id):
+                        # Create broadcast task with timeout handling
+                        async def safe_broadcast():
+                            try:
+                                await timeout_manager.with_timeout(
+                                    "websocket_broadcast",
+                                    self.broadcast_to_job(job_id, progress_data),
+                                    timeout=5.0,
+                                )
+                            except Exception as e:
+                                console.print(
+                                    f"[yellow]Broadcast failed for job {job_id}: {e}[/yellow]"
+                                )
+
+                        asyncio.create_task(safe_broadcast())
+
+                        if job_id not in self.known_jobs:
+                            self.known_jobs.add(job_id)
+                            console.print(f"[green]New job detected: {job_id}[/green]")
+
+                # Start directory monitoring with proper timeout handling
+                await self._monitor_directory_changes()
 
         except Exception as e:
             console.print(f"[red]Progress monitoring setup error: {e}[/red]")
 
     async def cleanup_old_jobs(self) -> None:
-        """Periodically clean up old job files."""
+        timeout_manager = get_timeout_manager()
+
         while self.is_running:
             try:
-                await self._perform_cleanup_cycle()
-                await asyncio.sleep(3600)
+                # Cleanup cycle with timeout protection
+                await timeout_manager.with_timeout(
+                    "file_operations",
+                    self._perform_cleanup_cycle(),
+                    timeout=30.0,  # Cleanup timeout
+                    strategy=TimeoutStrategy.GRACEFUL_DEGRADATION,
+                )
+                await asyncio.sleep(3600)  # 1 hour between cleanups
             except Exception as e:
                 console.print(f"[red]Cleanup error: {e}[/red]")
-                await asyncio.sleep(3600)
+                # Shorter sleep on error to retry sooner
+                await asyncio.sleep(1800)  # 30 minutes on error
 
     async def _perform_cleanup_cycle(self) -> None:
-        """Perform a single cleanup cycle for old jobs."""
         if not self.progress_dir.exists():
             return
 
@@ -160,44 +297,48 @@ class JobManager:
             self._cleanup_old_job_file(progress_file)
 
     def _calculate_cleanup_cutoff_time(self) -> float:
-        """Calculate cutoff time for job cleanup (24 hours ago)."""
         return time.time() - (24 * 60 * 60)
 
     def _find_old_job_files(self, cutoff_time: float) -> list[Path]:
-        """Find job files older than the cutoff time."""
         return [
             progress_file
-            for progress_file in self.progress_dir.glob("job-*.json")
+            for progress_file in self.progress_dir.glob("job - *.json")
             if progress_file.stat().st_mtime < cutoff_time
         ]
 
     def _cleanup_old_job_file(self, progress_file: Path) -> None:
-        """Clean up a single old job file if it's safe to do so."""
         job_id = self.extract_job_id_from_file(progress_file)
 
         if job_id not in self.active_connections:
             progress_file.unlink(missing_ok=True)
-            console.print(f"[yellow]Cleaned up old job: {job_id}[/yellow]")
+            console.print(f"[yellow]Cleaned up old job: {job_id}[/ yellow]")
 
     async def timeout_stuck_jobs(self) -> None:
-        """Monitor and timeout stuck jobs that haven't been updated."""
+        timeout_manager = get_timeout_manager()
+
         while self.is_running:
             try:
-                await self._check_and_timeout_stuck_jobs()
-                await asyncio.sleep(300)
+                # Timeout check with its own timeout protection
+                await timeout_manager.with_timeout(
+                    "file_operations",
+                    self._check_and_timeout_stuck_jobs(),
+                    timeout=60.0,  # Timeout check timeout
+                    strategy=TimeoutStrategy.GRACEFUL_DEGRADATION,
+                )
+                await asyncio.sleep(300)  # 5 minutes between checks
             except Exception as e:
                 console.print(f"[red]Timeout check error: {e}[/red]")
+                # Continue checking even on errors
                 await asyncio.sleep(300)
 
     async def _check_and_timeout_stuck_jobs(self) -> None:
-        """Check for stuck jobs and timeout those that are inactive."""
         if not self.progress_dir.exists():
             return
 
         current_time = time.time()
         timeout_seconds = 30 * 60
 
-        for progress_file in self.progress_dir.glob("job -* .json"):
+        for progress_file in self.progress_dir.glob("job-* .json"):
             await self._process_job_timeout_check(
                 progress_file,
                 current_time,
@@ -210,20 +351,28 @@ class JobManager:
         current_time: float,
         timeout_seconds: int,
     ) -> None:
-        """Process timeout check for a single job file."""
         try:
-            progress_data = json.loads(progress_file.read_text())
+            # Validate the progress file path is secure
+            validated_file = SecurePathValidator.validate_safe_path(
+                progress_file, self.progress_dir
+            )
+
+            # Validate file size before reading
+            SecurePathValidator.validate_file_size(validated_file)
+
+            progress_data = json.loads(validated_file.read_text())
 
             if self._should_timeout_job(
                 progress_data,
-                progress_file,
+                validated_file,
                 current_time,
                 timeout_seconds,
             ):
-                self._timeout_job(progress_data, progress_file)
+                self._timeout_job(progress_data, validated_file)
 
-        except (json.JSONDecodeError, OSError):
-            pass  # Skip files that can't be processed
+        except (json.JSONDecodeError, OSError, Exception):
+            # Catch validation errors as well as file errors
+            pass
 
     def _should_timeout_job(
         self,
@@ -232,22 +381,20 @@ class JobManager:
         current_time: float,
         timeout_seconds: int,
     ) -> bool:
-        """Determine if a job should be timed out."""
         return (
             progress_data.get("status") == "running"
             and current_time - progress_file.stat().st_mtime > timeout_seconds
         )
 
     def _timeout_job(self, progress_data: dict, progress_file: Path) -> None:
-        """Mark a job as failed due to timeout."""
         progress_data["status"] = "failed"
         progress_data["message"] = "Job timed out (no updates for 30 minutes)"
 
         progress_file.write_text(json.dumps(progress_data, indent=2))
 
         job_id = progress_data.get("job_id", "unknown")
-        console.print(f"[red]Job {job_id} timed out and marked as failed[/red]")
+        console.print(f"[red]Job {job_id} timed out and marked as failed[/ red]")
 
     def cleanup(self) -> None:
         self.is_running = False
-        console.print("[blue]Job manager cleanup completed[/blue]")
+        console.print("[blue]Job manager cleanup completed[/ blue]")

crackerjack/mcp/websocket/server.py

@@ -1,3 +1,5 @@
+import asyncio
+import contextlib
 import signal
 import subprocess
 import tempfile
@@ -7,6 +9,8 @@ from pathlib import Path
 import uvicorn
 from rich.console import Console
 
+from crackerjack.core.timeout_manager import get_timeout_manager
+
 from .app import create_websocket_app
 from .jobs import JobManager
 
@@ -20,6 +24,8 @@ class WebSocketServer:
         self.is_running = True
         self.job_manager: JobManager | None = None
         self.app = None
+        self.timeout_manager = get_timeout_manager()
+        self.server_task: asyncio.Task | None = None
 
     def setup(self) -> None:
         self.progress_dir.mkdir(exist_ok=True)
@@ -35,6 +41,31 @@ class WebSocketServer:
         console.print("\n[yellow]Shutting down WebSocket server...[/yellow]")
         self.is_running = False
 
+        # Cancel server task if running
+        if self.server_task and not self.server_task.done():
+            self.server_task.cancel()
+
+        # Clean up job manager connections
+        if self.job_manager:
+            with contextlib.suppress(Exception):
+                # Give existing connections 5 seconds to close
+                asyncio.create_task(self._graceful_shutdown())
+
+    async def _graceful_shutdown(self) -> None:
+        """Gracefully shutdown WebSocket connections."""
+        if self.job_manager:
+            try:
+                # Wait briefly for connections to close naturally
+                await asyncio.sleep(2.0)
+
+                # Force close any remaining connections
+                # Note: Implementation depends on JobManager API
+                console.print(
+                    "[yellow]Forcing remaining WebSocket connections to close[/yellow]"
+                )
+            except Exception as e:
+                console.print(f"[red]Error during graceful shutdown: {e}[/red]")
+
     def run(self) -> None:
         try:
             self.setup()
@@ -42,24 +73,73 @@ class WebSocketServer:
                 f"[green]Starting WebSocket server on port {self.port}[/green]",
             )
             console.print(f"Progress directory: {self.progress_dir}")
-            console.print("Press Ctrl + C to stop")
+            console.print("Press Ctrl+C to stop")
 
             config = uvicorn.Config(
                 app=self.app,
                 port=self.port,
                 host="127.0.0.1",
                 log_level="info",
+                # Add timeout configurations
+                timeout_keep_alive=30,  # Keep-alive timeout
+                timeout_graceful_shutdown=30,  # Graceful shutdown timeout
             )
 
             server = uvicorn.Server(config)
-            server.run()
+
+            # Use asyncio event loop for better control
+            try:
+                asyncio.run(self._run_with_timeout(server))
+            except KeyboardInterrupt:
+                console.print("\n[yellow]Server interrupted by user[/yellow]")
 
         except KeyboardInterrupt:
             console.print("\n[yellow]Server stopped by user[/yellow]")
         except Exception as e:
             console.print(f"[red]Server error: {e}[/red]")
         finally:
-            pass  # Cleanup handled by FastAPI lifespan
+            console.print("[green]WebSocket server shutdown complete[/green]")
+
+    async def _run_with_timeout(self, server: uvicorn.Server) -> None:
+        """Run the server with timeout protection."""
+        try:
+            # Start server as a background task
+            self.server_task = asyncio.create_task(server.serve())
+
+            # Monitor server health while running
+            while self.is_running and not self.server_task.done():
+                try:
+                    # Check server health periodically
+                    await asyncio.sleep(5.0)
+
+                    # Optional: Add health checks here
+                    # if not await self._server_health_check():
+                    #     console.print("[yellow]Server health check failed[/yellow]")
+                    #     break
+
+                except asyncio.CancelledError:
+                    console.print("[yellow]Server monitoring cancelled[/yellow]")
+                    break
+                except Exception as e:
+                    console.print(f"[red]Server monitoring error: {e}[/red]")
+                    break
+
+            # Wait for server task to complete
+            if self.server_task and not self.server_task.done():
+                try:
+                    await asyncio.wait_for(self.server_task, timeout=30.0)
+                except TimeoutError:
+                    console.print(
+                        "[yellow]Server shutdown timeout, forcing termination[/yellow]"
+                    )
+                    self.server_task.cancel()
+                    try:
+                        await self.server_task
+                    except asyncio.CancelledError:
+                        pass
+
+        except Exception as e:
+            console.print(f"[red]Server runtime error: {e}[/red]")
 
 
 def handle_websocket_server_command(
@@ -68,10 +148,9 @@ def handle_websocket_server_command(
     restart: bool = False,
     port: int = 8675,
 ) -> None:
-    """Handle WebSocket server start/stop/restart commands."""
     if stop or restart:
         console.print("[yellow]Stopping WebSocket servers...[/yellow]")
-        # Kill any existing uvicorn processes running on the port
+
         try:
             result = subprocess.run(
                 ["pkill", "-f", f"uvicorn.*:{port}"],
@@ -92,7 +171,6 @@ def handle_websocket_server_command(
         if stop:
             return
 
-        # For restart, wait a moment before starting again
         time.sleep(2)
 
     if start or restart: