Update main.py

main.py

@@ -69,7 +69,7 @@ async def kill_kernel(session_token: str):
 # Function to execute code in a kernel
 async def execute_code(session_token: str, code: str):
     setup_code = "%matplotlib inline"
-
+    
     if session_token not in kernel_managers:
         await create_kernel(session_token)
 
@@ -85,15 +85,12 @@ async def execute_code(session_token: str, code: str):
         output = []
         timeout = datetime.now() + timedelta(seconds=TIMEOUT_DURATION)
 
+        output = []
         while True:
             if datetime.now() > timeout:
                 raise TimeoutError("Code execution timed out.")
-
-            try:
-                msg = kc.get_iopub_msg(timeout=1)
-            except Empty:
-                continue
-
+                
+            msg = st.session_state.kernel_client.get_iopub_msg()
             if msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
                 break
             elif msg['msg_type'] == 'error':
@@ -103,15 +100,16 @@ async def execute_code(session_token: str, code: str):
                     "traceback": msg['content']['traceback']
                 }
                 output.append({"error": error_output})
-            elif 'data' in msg['content']:
+            if 'data' in msg['content']:
                 output.append({"data": msg['content']['data']})
             elif 'text' in msg['content']:
                 output.append({"text": msg['content']['text']})
-
+            
         last_access_times[session_token] = datetime.now()
         return {'status': 'success', 'value': output}
 
     except Exception as e:
+        last_access_times[session_token] = datetime.now()
         return {'status': 'error', 'value': str(e)}