1 files changed, 38 insertions, 3 deletions
diff --git a/continuedev/src/continuedev/server/main.py b/continuedev/src/continuedev/server/main.py
index f4d82903..42dc0cc1 100644
--- a/continuedev/src/continuedev/server/main.py
+++ b/continuedev/src/continuedev/server/main.py
@@ -1,10 +1,12 @@
+import time
+import psutil
 import os
-import sys
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from .ide import router as ide_router
 from .gui import router as gui_router
-import logging
+from .session_manager import session_manager
+import atexit
 import uvicorn
 import argparse
 
@@ -44,5 +46,38 @@ def run_server():
     uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 
+def cleanup():
+    print("Cleaning up sessions")
+    for session_id in session_manager.sessions:
+        session_manager.persist_session(session_id)
+
+
+def cpu_usage_report():
+    process = psutil.Process(os.getpid())
+    # Call cpu_percent once to start measurement, but ignore the result
+    process.cpu_percent(interval=None)
+    # Wait for a short period of time
+    time.sleep(1)
+    # Call cpu_percent again to get the CPU usage over the interval
+    cpu_usage = process.cpu_percent(interval=None)
+    print(f"CPU usage: {cpu_usage}%")
+
+
+atexit.register(cleanup)
+
 if __name__ == "__main__":
-    run_server()
+    try:
+        # import threading
+
+        # def cpu_usage_loop():
+        #     while True:
+        #         cpu_usage_report()
+        #         time.sleep(2)
+
+        # cpu_thread = threading.Thread(target=cpu_usage_loop)
+        # cpu_thread.start()
+
+        run_server()
+    except Exception as e:
+        cleanup()
+        raise e