🔧 fix session_id handling

Author: shroominic

codeboxapi/box/codebox.py

@@ -63,28 +63,29 @@ def __new__(cls, *args, **kwargs):
 
         return super().__new__(cls)
 
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self, /, **kwargs) -> None:
+        super().__init__()
         self.session_id: Optional[UUID] = kwargs.pop("session_id", None)
         self.aiohttp_session: Optional[ClientSession] = None
 
     @classmethod
-    def from_id(cls, session_id: Union[int, UUID]) -> "CodeBox":
-        if isinstance(session_id, int):
-            session_id = UUID(int=session_id)
-        return cls(session_id=session_id)
+    def from_id(cls, session_id: Union[int, UUID], **kwargs) -> "CodeBox":
+        kwargs["session_id"] = (
+            UUID(int=session_id) if isinstance(session_id, int) else session_id
+        )
+        return cls(**kwargs)
 
     def _update(self) -> None:
         """Update last interaction time"""
-        if self.session_id is None:
-            raise RuntimeError("Make sure to start your CodeBox before using it.")
         self.last_interaction = datetime.now()
 
     def codebox_request(self, method, endpoint, *args, **kwargs) -> Dict[str, Any]:
         """Basic request to the CodeBox API"""
         self._update()
+        if self.session_id is None:
+            raise RuntimeError("Make sure to start your CodeBox before using it.")
         return base_request(
-            method, f"/codebox/{self.session_id}" + endpoint, *args, **kwargs
+            method, f"/codebox/{self.session_id.int}" + endpoint, *args, **kwargs
         )
 
     async def acodebox_request(
@@ -94,30 +95,46 @@ async def acodebox_request(
         self._update()
         if self.aiohttp_session is None:
             self.aiohttp_session = ClientSession()
+        if self.session_id is None:
+            raise RuntimeError("Make sure to start your CodeBox before using it.")
         return await abase_request(
             self.aiohttp_session,
             method,
-            f"/codebox/{self.session_id}" + endpoint,
+            f"/codebox/{self.session_id.int}" + endpoint,
             *args,
             **kwargs,
         )
 
     def start(self) -> CodeBoxStatus:
-        self.session_id = base_request(
-            method="GET",
-            endpoint="/codebox/start",
-        )["id"]
+        if self.session_id is not None:
+            if settings.VERBOSE:
+                print(f"{self} is already started!")
+            return CodeBoxStatus(status="started")
+        self.session_id = UUID(
+            int=base_request(
+                method="GET",
+                endpoint="/codebox/start",
+            )["id"]
+        )
+        if settings.VERBOSE:
+            print(f"{self} started!")
         return CodeBoxStatus(status="started")
 
     async def astart(self) -> CodeBoxStatus:
         self.aiohttp_session = ClientSession()
+        if self.session_id is not None:
+            if settings.VERBOSE:
+                print(f"{self} is already started!")
+            return CodeBoxStatus(status="started")
         self.session_id = (
             await abase_request(
                 self.aiohttp_session,
                 method="GET",
                 endpoint="/codebox/start",
             )
         )["id"]
+        if settings.VERBOSE:
+            print(f"{self} started!")
         return CodeBoxStatus(status="started")
 
     def status(self):

codeboxapi/box/localbox.py

@@ -8,6 +8,7 @@
 import asyncio
 import json
 import os
+import signal
 import subprocess
 import sys
 import time
@@ -38,10 +39,11 @@ class LocalBox(BaseBox):
     """
 
     _instance: Optional["LocalBox"] = None
+    _jupyter_pids: List[int] = []
 
     def __new__(cls, *args, **kwargs):
         if not cls._instance:
-            cls._instance = super().__new__(cls, *args, **kwargs)
+            cls._instance = super().__new__(cls)
         else:
             if settings.SHOW_INFO:
                 print(
@@ -52,15 +54,16 @@ def __new__(cls, *args, **kwargs):
                 )
         return cls._instance
 
-    def __init__(self) -> None:
-        super().__init__()
+    def __init__(self, /, **kwargs) -> None:
+        super().__init__(session_id=kwargs.pop("session_id", None))
         self.port: int = 8888
         self.kernel_id: Optional[dict] = None
         self.ws: Union[WebSocketClientProtocol, ClientConnection, None] = None
         self.jupyter: Union[Process, subprocess.Popen, None] = None
         self.aiohttp_session: Optional[aiohttp.ClientSession] = None
 
     def start(self) -> CodeBoxStatus:
+        self.session_id = uuid4()
         os.makedirs(".codebox", exist_ok=True)
         self._check_port()
         if settings.VERBOSE:
@@ -84,6 +87,7 @@ def start(self) -> CodeBoxStatus:
                 stderr=out,
                 cwd=".codebox",
             )
+            self._jupyter_pids.append(self.jupyter.pid)
         except FileNotFoundError:
             raise ModuleNotFoundError(
                 "Jupyter Kernel Gateway not found, please install it with:\n"
@@ -100,7 +104,10 @@ def start(self) -> CodeBoxStatus:
             if settings.VERBOSE:
                 print("Waiting for kernel to start...")
             time.sleep(1)
+        self._connect()
+        return CodeBoxStatus(status="started")
 
+    def _connect(self) -> None:
         response = requests.post(
             f"{self.kernel_url}/kernels",
             headers={"Content-Type": "application/json"},
@@ -112,20 +119,15 @@ def start(self) -> CodeBoxStatus:
 
         self.ws = ws_connect_sync(f"{self.ws_url}/kernels/{self.kernel_id}/channels")
 
-        return CodeBoxStatus(status="started")
-
     def _check_port(self) -> None:
         try:
             response = requests.get(f"http://localhost:{self.port}", timeout=90)
         except requests.exceptions.ConnectionError:
             pass
         else:
-            try:
-                requests.post(f"http://localhost:{self.port}/api/shutdown")
-            except requests.exceptions.ConnectionError:
-                if response.status_code == 200:
-                    self.port += 1
-                    self._check_port()
+            if response.status_code == 200:
+                self.port += 1
+                self._check_port()
 
     def _check_installed(self) -> None:
         try:
@@ -139,6 +141,7 @@ def _check_installed(self) -> None:
             raise
 
     async def astart(self) -> CodeBoxStatus:
+        self.session_id = uuid4()
         os.makedirs(".codebox", exist_ok=True)
         self.aiohttp_session = aiohttp.ClientSession()
         await self._acheck_port()
@@ -161,6 +164,7 @@ async def astart(self) -> CodeBoxStatus:
                 stderr=out,
                 cwd=".codebox",
             )
+            self._jupyter_pids.append(self.jupyter.pid)
         except Exception as e:
             print(e)
             raise ModuleNotFoundError(
@@ -180,7 +184,12 @@ async def astart(self) -> CodeBoxStatus:
             if settings.VERBOSE:
                 print("Waiting for kernel to start...")
             await asyncio.sleep(1)
+        await self._aconnect()
+        return CodeBoxStatus(status="started")
 
+    async def _aconnect(self) -> None:
+        if self.aiohttp_session is None:
+            self.aiohttp_session = aiohttp.ClientSession()
         response = await self.aiohttp_session.post(
             f"{self.kernel_url}/kernels", headers={"Content-Type": "application/json"}
         )
@@ -189,8 +198,6 @@ async def astart(self) -> CodeBoxStatus:
             raise Exception("Could not start kernel")
         self.ws = await ws_connect(f"{self.ws_url}/kernels/{self.kernel_id}/channels")
 
-        return CodeBoxStatus(status="started")
-
     async def _acheck_port(self) -> None:
         try:
             if self.aiohttp_session is None:
@@ -206,6 +213,9 @@ async def _acheck_port(self) -> None:
                 await self._acheck_port()
 
     def status(self) -> CodeBoxStatus:
+        if not self.kernel_id:
+            self._connect()
+
         return CodeBoxStatus(
             status="running"
             if self.kernel_id
@@ -214,6 +224,8 @@ def status(self) -> CodeBoxStatus:
         )
 
     async def astatus(self) -> CodeBoxStatus:
+        if not self.kernel_id:
+            await self._aconnect()
         return CodeBoxStatus(
             status="running"
             if self.kernel_id
@@ -242,9 +254,9 @@ def run(
         if retry <= 0:
             raise RuntimeError("Could not connect to kernel")
         if not self.ws:
-            self.start()
+            self._connect()
             if not self.ws:
-                raise RuntimeError("Could not connect to kernel")
+                raise RuntimeError("Jupyter not running. Make sure to start it first.")
 
         if settings.VERBOSE:
             print("Running code:\n", code)
@@ -354,9 +366,9 @@ async def arun(
         if retry <= 0:
             raise RuntimeError("Could not connect to kernel")
         if not self.ws:
-            await self.astart()
+            await self._aconnect()
             if not self.ws:
-                raise RuntimeError("Could not connect to kernel")
+                raise RuntimeError("Jupyter not running. Make sure to start it first.")
 
         if settings.VERBOSE:
             print("Running code:\n", code)
@@ -488,20 +500,16 @@ async def alist_files(self) -> List[CodeBoxFile]:
         return await asyncio.to_thread(self.list_files)
 
     def restart(self) -> CodeBoxStatus:
-        if self.jupyter is not None:
-            self.stop()
-        else:
-            self.start()
         return CodeBoxStatus(status="restarted")
 
     async def arestart(self) -> CodeBoxStatus:
-        if self.jupyter is not None:
-            await self.astop()
-        else:
-            await self.astart()
         return CodeBoxStatus(status="restarted")
 
     def stop(self) -> CodeBoxStatus:
+        for pid in self._jupyter_pids:
+            print(f"Killing {pid}")
+            os.kill(pid, signal.SIGTERM)
+
         if self.jupyter is not None:
             self.jupyter.terminate()
             self.jupyter.wait()