Update app.py

app.py

@@ -10,6 +10,7 @@ from typing import Dict, Any, Optional
 import threading
 import asyncio
 import signal
+from concurrent.futures import ThreadPoolExecutor # NEW: for loop.run_in_executor
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
 from pydantic import BaseModel
 from starlette.responses import HTMLResponse
@@ -77,6 +78,7 @@ class WebSocketInputOutput:
         """Synchronous wrapper for readline for use with exec."""
         try:
             print("DEBUG: Synchronous input wrapper called.")
+            # This is called from the executor thread, so asyncio.run is appropriate here
             return asyncio.run(self.readline())
         except Exception as e:
             print(f"ERROR: Synchronous input wrapper error: {e}")
@@ -387,17 +389,16 @@ class SecurePythonExecutor:
         original_stdout = sys.stdout
         original_stderr = sys.stderr
 
-        # Use a threading.Event to signal completion from the execution thread
-        execution_finished_event = threading.Event()
-
-        def target():
+        # Define the synchronous target function to be run in the executor thread
+        def _execute_in_thread():
+            # This function runs in a separate thread managed by the event loop's executor
             nonlocal result
             try:
                 # Redirect sys.stdin, sys.stdout, sys.stderr for this thread
                 sys.stdin = websocket_io
                 sys.stdout = websocket_io
                 sys.stderr = websocket_io
-                print(f"DEBUG: Thread '{threading.current_thread().name}' started. IO redirected.")
+                print(f"DEBUG: Executor thread '{threading.current_thread().name}' started. IO redirected.")
                 
                 # Create restricted execution environment
                 restricted_globals = {
@@ -405,65 +406,73 @@ class SecurePythonExecutor:
                     '__name__': '__main__',
                     '__doc__': None,
                 }
-                print("DEBUG: Restricted globals prepared.") # ADDED THIS LINE
+                print("DEBUG: Restricted globals prepared in executor thread.") # ADDED THIS LINE
                 
-                print(f"DEBUG: Preparing to execute code (thread): {code[:100]}...")
+                print(f"DEBUG: Preparing to execute code (executor thread): {code[:100]}...")
                 exec(code, restricted_globals)
-                print("DEBUG: Code execution completed successfully in thread.")
+                print("DEBUG: Code execution completed successfully in executor thread.")
                 result['success'] = True
                 
             except EOFError as e:
                 # This happens if the WebSocket disconnects during input()
                 result['error'] = f"Input stream closed: {str(e)}"
                 result['success'] = False
-                print(f"ERROR: EOFError during execution in thread: {e}")
+                print(f"ERROR: EOFError during execution in executor thread: {e}")
             except Exception as e:
                 error_msg = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
-                print(f"ERROR: Unhandled exception during execution in thread: {error_msg}")
-                # Write error to WebSocket and internal buffer
+                print(f"ERROR: Unhandled exception during execution in executor thread: {error_msg}")
                 try:
-                    # Use asyncio.run for writing within a synchronous thread, if necessary
-                    # This is generally safe if it's a new event loop per call or in a separate thread.
+                    # If this write fails, it's still good to log the error to stderr buffer
+                    # This asyncio.run call is crucial and runs a new event loop for this async call
                     asyncio.run(websocket_io.write(f"Error: {error_msg}\n"))
-                    print("DEBUG: Attempted to write error via websocket_io.write from thread.")
+                    print("DEBUG: Attempted to write error via websocket_io.write from executor thread.")
                 except Exception as write_e:
-                    print(f"ERROR: Failed to write error to WebSocket from thread (inner): {write_e}")
+                    print(f"ERROR: Failed to write error to WebSocket from executor thread (inner): {write_e}")
                 websocket_io.error_buffer.write(error_msg)
                 result['success'] = False
             finally:
-                # Restore original stdin/stdout/stderr for the thread pool (if applicable)
+                # Restore original stdin/stdout/stderr for the thread pool
                 sys.stdin = original_stdin
                 sys.stdout = original_stdout
                 sys.stderr = original_stderr
-                execution_finished_event.set() # Signal completion
-                print(f"DEBUG: Thread '{threading.current_thread().name}' finished. IO restored.")
+                print(f"DEBUG: Executor thread '{threading.current_thread().name}' finished. IO restored.")
 
-        # Run the execution in a separate thread to allow for timeout
-        thread = threading.Thread(target=target, name=f"CodeExecutorThread-{id(websocket_io)}")
-        thread.daemon = True # Allow the program to exit even if thread is running
-        thread.start()
-        print(f"DEBUG: Code execution thread '{thread.name}' started. Waiting for join.")
 
-        # Wait for the thread to finish or timeout
-        thread.join(timeout=self.timeout)
-        
-        if thread.is_alive():
-            print(f"DEBUG: Thread '{thread.name}' timed out after {self.timeout} seconds.")
+        # Use the event loop's default ThreadPoolExecutor to run the synchronous code
+        loop = asyncio.get_running_loop()
+        try:
+            print(f"DEBUG: Scheduling _execute_in_thread in executor with timeout {self.timeout}s.")
+            # Run the synchronous _execute_in_thread in a separate thread
+            await asyncio.wait_for(
+                loop.run_in_executor(None, _execute_in_thread), # None uses the default ThreadPoolExecutor
+                timeout=self.timeout
+            )
+            print("DEBUG: _execute_in_thread completed via executor.")
+        except asyncio.TimeoutError:
+            print(f"ERROR: Code execution timed out after {self.timeout} seconds in run_in_executor.")
             result['error'] = f"Code execution timed out after {self.timeout} seconds"
             result['success'] = False
-            # Attempt to send timeout message to client
             try:
                 await websocket_io.write(f"Error: {result['error']}\n")
             except Exception:
                 pass # Ignore if websocket is already closed
-        else:
-            print(f"DEBUG: Thread '{thread.name}' finished within timeout.")
+        except Exception as e:
+            # This catches exceptions from _execute_in_thread that propagate up
+            # (e.g., if the thread itself crashes in an unrecoverable way before our try/except)
+            print(f"ERROR: Unexpected exception from executor: {e}")
+            result['error'] = f"Unexpected executor error: {type(e).__name__}: {str(e)}"
+            result['success'] = False
+            try:
+                await websocket_io.write(f"Error: {result['error']}\n")
+            except Exception:
+                pass
+
+        # After execution (or timeout), collect output
+        if result['success']: # Only if execution was successful in _execute_in_thread
             result['output'] = websocket_io.output_buffer.getvalue()
-            if websocket_io.error_buffer.getvalue():
-                result['error'] = websocket_io.error_buffer.getvalue()
-                result['success'] = False
-            else:
-                result['success'] = True
+        if websocket_io.error_buffer.getvalue(): # Error buffer always checked
+            result['error'] = result['error'] + "\n" + websocket_io.error_buffer.getvalue() if result['error'] else websocket_io.error_buffer.getvalue()
+            result['success'] = False # Ensure success is false if there's any error output
         
         # Send a final message to the client indicating execution is complete
         try:
@@ -474,7 +483,7 @@ class SecurePythonExecutor:
             print("DEBUG: WebSocketDisconnect when sending execution_complete.")
             pass # Ignore if websocket is already closed
         except Exception as e:
-            print(f"ERROR: Error sending execution_complete message: {e}") # Warning if this fails
+            print(f"ERROR: Error sending execution_complete message: {e}")
 
         print(f"DEBUG: run_with_timeout returning result for client: {websocket_io.websocket.client}")
         return result
@@ -495,7 +504,7 @@ app = FastAPI(
 async def get_index():
     """Serves the interactive frontend."""
     # This assumes index.html is in the same directory as app.py
-    # [cite_start]Explicitly copy index.html to /app/index.html to avoid any ambiguity with '.' [cite: 5, 6]
+    # Explicitly copy index.html to /app/index.html to avoid any ambiguity with '.'
     with open("index.html", "r") as f:
         return HTMLResponse(content=f.read())
 
@@ -542,7 +551,7 @@ async def websocket_endpoint(websocket: WebSocket):
                 websocket_io.input_queue.get_nowait()
             except asyncio.QueueEmpty:
                 pass
-        print(f"WebSocket connection closed for {websocket.client}")
+        print(f"WebSocket connection closed for {websocket_io.websocket.client}") # Use websocket_io.websocket.client here too
 
 # This block is for local development using `uvicorn`
 # Hugging Face Spaces will use `uvicorn` directly via the Dockerfile CMD