fix: restart sandbox after cancelled docker exec

.pre-commit-config.yaml

@@ -1,28 +1,10 @@
 repos:
-  - repo: https://github.com/pycqa/flake8.git
-    rev: 7.1.1
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.16.4
     hooks:
-      - id: flake8
-        exclude: |
-            (?x)^(
-                examples/|
-                docs/|
-                tests/
-            )
-  - repo: https://github.com/PyCQA/isort.git
-    rev: 6.0.1
-    hooks:
-      - id: isort
-  - repo: https://github.com/google/yapf
-    rev: v0.43.0
-    hooks:
-      - id: yapf
-        exclude: |
-            (?x)^(
-                examples/|
-                docs/|
-                tests/
-            )
+      - id: ruff-check
+        args: [--fix]
+      - id: ruff-format
   - repo: https://github.com/pre-commit/pre-commit-hooks.git
     rev: v5.0.0
     hooks:
@@ -35,8 +17,6 @@ repos:
         exclude: thirdparty/|docs/
       - id: requirements-txt-fixer
         exclude: thirdparty/
-      - id: double-quote-string-fixer
-        exclude: thirdparty/
       - id: check-merge-conflict
         exclude: thirdparty/
       - id: mixed-line-ending

ms_enclave/cli/start_server.py

@@ -10,8 +10,7 @@
 
 
 def subparser_func(args):
-    """ Function which will be called for a specific sub parser.
-    """
+    """Function which will be called for a specific sub parser."""
     return ServerCMD(args)
 
 
@@ -23,8 +22,7 @@ def __init__(self, args):
 
     @staticmethod
     def define_args(parsers: ArgumentParser):
-        """Define args for the server command.
-        """
+        """Define args for the server command."""
         parser = parsers.add_parser(ServerCMD.name, help='Start the MS-Enclave sandbox HTTP server')
         add_argument(parser)
         parser.set_defaults(func=subparser_func)
@@ -67,18 +65,18 @@ def add_argument(parser: ArgumentParser) -> None:
         type=str,
         choices=['critical', 'error', 'warning', 'info', 'debug'],
         default='info',
-        help='Log level for the server (default: info)'
+        help='Log level for the server (default: info)',
     )
     parser.add_argument(
         '--cleanup-interval',
         type=int,
         default=300,
         metavar='SECONDS',
-        help='Background cleanup interval in seconds (default: 300)'
+        help='Background cleanup interval in seconds (default: 300)',
     )
     parser.add_argument(
         '--api-key',
         type=str,
         default=None,
-        help='Optional API key to protect endpoints. If omitted, no authentication is enforced.'
+        help='Optional API key to protect endpoints. If omitted, no authentication is enforced.',
     )

ms_enclave/sandbox/boxes/__init__.py

@@ -12,7 +12,6 @@
     'SandboxFactory',
     'register_sandbox',
     'StatelessSandbox',
-
     # Implementations
     'DockerSandbox',
     'DockerNotebookSandbox',

ms_enclave/sandbox/boxes/base.py

@@ -206,7 +206,7 @@ def get_info(self) -> SandboxInfo:
             created_at=self.created_at,
             updated_at=self.updated_at,
             metadata=self.metadata,
-            available_tools=self.get_available_tools()
+            available_tools=self.get_available_tools(),
         )
 
     async def __aenter__(self):
@@ -239,7 +239,7 @@ def create_sandbox(
         cls,
         sandbox_type: SandboxType,
         config: Optional[Union[SandboxConfig, Dict]] = None,
-        sandbox_id: Optional[str] = None
+        sandbox_id: Optional[str] = None,
     ) -> Sandbox:
         """Create a sandbox instance.
 

ms_enclave/sandbox/boxes/docker_notebook.py

@@ -58,6 +58,7 @@ async def start(self) -> None:
 
             # Initialize Docker client first
             import docker
+
             self.client = docker.from_env()
 
             # Build Jupyter image if needed before creating container
@@ -161,9 +162,9 @@ def build_image():
                 # Process and log build output
                 for log in build_logs[1]:  # build_logs[1] contains the build log generator
                     if 'stream' in log:
-                        logger.info(f"[📦 {self.id}] {log['stream'].strip()}")
+                        logger.info(f'[📦 {self.id}] {log["stream"].strip()}')
                     elif 'error' in log:
-                        logger.error(f"[📦 {self.id}] {log['error']}")
+                        logger.error(f'[📦 {self.id}] {log["error"]}')
                 return build_logs[0]  # Return the built image
 
             await self._run_blocking(build_image)
@@ -191,34 +192,50 @@ async def _create_kernel(self) -> None:
         # Establish websocket connection
         try:
             from websocket import create_connection
+
             ws_url = f'ws://{self.host}:{self.port}/api/kernels/{self.kernel_id}/channels'
             self.ws = create_connection(ws_url)
             logger.info(f'Kernel {self.kernel_id} created and connected')
         except ImportError:
             raise RuntimeError('websocket-client package is required. Install with: pip install websocket-client')
 
+    def _clear_jupyter_session(self) -> None:
+        """Drop local references to the current Jupyter session."""
+        if self.ws:
+            try:
+                self.ws.close()
+            except Exception:
+                pass
+        self.ws = None
+        self.kernel_id = None
+        self.base_url = None
+
+    async def _reset_after_interrupted_exec(self) -> None:
+        """Restart the container and recreate the Jupyter session."""
+        self.update_status(SandboxStatus.INITIALIZING)
+        self._clear_jupyter_session()
+        try:
+            await self._restart_container()
+            await self._setup_jupyter()
+        except BaseException as exc:
+            self.update_status(SandboxStatus.ERROR)
+            self.metadata['error'] = str(exc)
+            logger.error(f'Failed to reset notebook sandbox after interrupted exec: {exc}')
+            raise
+        self.update_status(SandboxStatus.RUNNING)
+
     async def cleanup(self) -> None:
         """Clean up Jupyter resources and Docker container."""
-        try:
-            # Close websocket connection
-            if self.ws:
-                try:
-                    self.ws.close()
-                except Exception:
-                    pass
-                self.ws = None
-
-            # Delete kernel
-            if self.kernel_id and self.base_url:
-                try:
-                    import requests
-                    requests.delete(f'{self.base_url}/api/kernels/{self.kernel_id}')
-                except Exception:
-                    pass
-                self.kernel_id = None
+        kernel_id = self.kernel_id
+        base_url = self.base_url
+        self._clear_jupyter_session()
 
-        except Exception as e:
-            logger.error(f'Error during Jupyter cleanup: {e}')
+        if kernel_id and base_url:
+            try:
+                import requests
+
+                requests.delete(f'{base_url}/api/kernels/{kernel_id}')
+            except Exception:
+                pass
 
-        # Call parent cleanup
         await super().cleanup()