fixes based on comments from code review

Author: abhishek-anand

api/rest_server.py

@@ -37,6 +37,7 @@
 )
 
 # Async task storage (simple in-memory for now)
+# TODO: Replace with persistent storage (Redis/database) for production to survive server restarts
 async_tasks: Dict[str, Dict[str, Any]] = {}
 
 
@@ -69,7 +70,7 @@ async def health_check():
         session_count = len(session_manager.sessions)
 
         # Check kernel pool
-        from core.kernel_pool import kernel_pool
+        from ..core.kernel_pool import kernel_pool
         kernel_status = {
             "total_kernels": len(kernel_pool.kernels),
             "available_kernels": len([k for k in kernel_pool.kernels.values() if k.is_available()]),

container_manager.py

@@ -18,7 +18,6 @@ class ContainerManager:
     """Manage CodeRunner Docker container lifecycle"""
 
     CONTAINER_NAME = "coderunner"
-    REST_PORT = 8223
     MCP_PORT = 8222
     JUPYTER_PORT = 8888
     PLAYWRIGHT_PORT = 3000
@@ -175,7 +174,6 @@ def start_container(cls):
             subprocess.run([
                 "docker", "run", "-d",
                 "--name", cls.CONTAINER_NAME,
-                "-p", f"{cls.REST_PORT}:8223",
                 "-p", f"{cls.MCP_PORT}:8222", 
                 "-p", f"{cls.JUPYTER_PORT}:8888",
                 "-p", f"{cls.PLAYWRIGHT_PORT}:3000",

core/session_manager.py

@@ -172,7 +172,7 @@ async def close_session(self, session_id: str) -> bool:
 
         try:
             # Import here to avoid circular imports
-            from server import kernel_pool
+            from .kernel_pool import kernel_pool
 
             # Release kernel back to pool
             await kernel_pool.release_kernel(session.kernel_id)

main.py

@@ -11,24 +11,11 @@
 main_app.mount("/api", rest_app)
 
 # Get the existing startup/shutdown events from REST API
-rest_startup = None
-rest_shutdown = None
-for event_handler in rest_app.router.on_startup:
-    rest_startup = event_handler
+for handler in rest_app.router.on_startup:
+    main_app.add_event_handler("startup", handler)
 
-for event_handler in rest_app.router.on_shutdown:
-    rest_shutdown = event_handler
-
-# Add REST API startup/shutdown to main app
-if rest_startup:
-    @main_app.on_event("startup")
-    async def combined_startup():
-        await rest_startup()
-
-if rest_shutdown:
-    @main_app.on_event("shutdown") 
-    async def combined_shutdown():
-        await rest_shutdown()
+for handler in rest_app.router.on_shutdown:
+    main_app.add_event_handler("shutdown", handler)
 
 # Export the unified app
 app = main_app