Update main.py

main.py

@@ -90,7 +90,7 @@ async def execute_code(session_token: str, code: str):
             if datetime.now() > timeout:
                 raise TimeoutError("Code execution timed out.")
                 
-            msg = st.session_state.kernel_client.get_iopub_msg()
+            msg = kc.get_iopub_msg()
             if msg['msg_type'] == 'status' and msg['content']['execution_state'] == 'idle':
                 break
             elif msg['msg_type'] == 'error':