Update Consumer.py

Consumer.py

@@ -30,8 +30,8 @@ def update_server(uid, output, status):
         print(f"Error updating server: {e}")
 
 def stop(uid):
-    global loop
-    global active_tasks
+    #global loop
+    #global active_tasks
     """Stops the execution of the cell with the given uid."""
     if uid in active_tasks:
         print(f"Stopping task {uid}...")
@@ -45,8 +45,8 @@ def stop(uid):
 
 def stopall():
     """Stops all running executions."""
-    global loop
-    global active_tasks
+    #global loop
+    #global active_tasks
     count = 0
     for uid in list(active_tasks.keys()):
         if active_tasks != asyncio.current_task():
@@ -77,16 +77,16 @@ def stopall():
 
 def execute_code_sync(code, uid):
     # Capture stdout
+    #global loop
+    #global active_tasks
     stdout_capture = i_o.StringIO()
     status = "completed"
     output = ""
-    global loop
-    global active_tasks
     
     # Inject stop functions into locals
     local_scope = locals()
-    #local_scope['stop'] = stop
-    #local_scope['stopall'] = stopall
+    local_scope['stop'] = stop
+    local_scope['stopall'] = stopall
 
     try:
         with contextlib.redirect_stdout(stdout_capture):
@@ -102,8 +102,8 @@ def execute_code_sync(code, uid):
     return output, status
 
 async def process_task(uid, code):
-    global loop
-    global active_tasks
+    #global loop
+    #global active_tasks
     print(f"Processing {uid}...")
     
     try:
@@ -146,8 +146,6 @@ async def process_task(uid, code):
 
 async def main():
     print(f"Starting FoxDot Consumer (Async)... Polling {API_URL}")
-    global loop
-    global active_tasks
     
     while True:
         # Move polling to thread to avoid blocking loop
@@ -159,15 +157,14 @@ async def main():
             print(f"Received task {uid}")
             
             # Create async task
-            task = active_tasks[uid]
-            active_tasks[uid] = asyncio.create_task(process_task(uid, code))
-            if task:
-                task.cancel()
-            
+            task = asyncio.create_task(process_task(uid, code))
+            active_tasks[uid] = task
+            if uid in active_tasks:
+                active_tasks[uid].cancel() 
+
         await asyncio.sleep(POLL_INTERVAL)
 
 if __name__ == "__main__":
-    loop = None
     try:
         # Check if an event loop is already running
         try:
@@ -178,9 +175,12 @@ if __name__ == "__main__":
         if loop and loop.is_running():
             print("Event loop already running. Scheduling main task...")
             loop.create_task(main())
+            # Note: We assume the existing loop will keep running.
+            # If this is a script that just exits, we might need loop.run_forever() if not already happening.
         else:
             print("Starting new event loop...")
-            loop=asyncio.run(main())
-        loop.run_forever()    
+            loop = asyncio.run(main())
+            
+        loop.run_forever()      
     except KeyboardInterrupt:
         print("Stopping consumer...")