Update app.py

app.py

@@ -1,48 +1,43 @@
-import os, json, time, base64, uuid, threading, subprocess, sys
-from flask import Flask, request, jsonify, Response, send_file
+import os, json, time, base64, uuid, threading, queue, io
+from flask import Flask, request, jsonify, Response
 from flask_cors import CORS
 import jupyter_client
-import queue, io
 
 app = Flask(__name__)
 CORS(app, resources={r"/*": {"origins": "*"}})
 
-# ── Auth ──────────────────────────────────────────────────────────────────────
-VALID_API_KEY = os.getenv("SANDBOX_API_KEY", "your-secret-key")
+VALID_API_KEY = os.getenv("SANDBOX_API_KEY", "")
 
 def check_auth():
-    key = request.headers.get("X-API-KEY") or request.headers.get("Authorization", "").replace("Bearer ", "")
+    if not VALID_API_KEY:
+        return True
+    key = (
+        request.headers.get("X-API-KEY", "") or
+        request.headers.get("Authorization", "").replace("Bearer ", "")
+    )
     return key == VALID_API_KEY
 
-# ── Kernel Session Manager ────────────────────────────────────────────────────
-# Each session gets its own isolated Jupyter kernel
-sessions = {}
+sessions      = {}
 sessions_lock = threading.Lock()
 
 def get_or_create_kernel(session_id="default"):
-    """Get existing kernel or create a new isolated one."""
     with sessions_lock:
         if session_id not in sessions:
-            print(f"[SANDBOX] Creating new kernel for session: {session_id}")
             km = jupyter_client.KernelManager(kernel_name='python3')
             km.start_kernel()
             kc = km.client()
             kc.start_channels()
             kc.wait_for_ready(timeout=30)
             sessions[session_id] = {
-                "km": km,
-                "kc": kc,
+                "km": km, "kc": kc,
                 "created_at": time.time(),
-                "last_used": time.time()
+                "last_used":  time.time()
             }
-            print(f"[SANDBOX] Kernel ready for session: {session_id}")
         else:
             sessions[session_id]["last_used"] = time.time()
-        
         return sessions[session_id]["kc"]
 
 def kill_kernel(session_id):
-    """Cleanly shut down a kernel session."""
     with sessions_lock:
         if session_id in sessions:
             try:
@@ -51,172 +46,166 @@ def kill_kernel(session_id):
             except:
                 pass
             del sessions[session_id]
-            print(f"[SANDBOX] Kernel killed for session: {session_id}")
 
 def execute_code(kc, code, timeout=60):
-    """
-    Execute code on a Jupyter kernel.
-    Returns: dict with stdout, stderr, result, artifacts (base64 images)
-    """
-    msg_id = kc.execute(code)
-    
-    stdout_parts = []
-    stderr_parts = []
-    result_parts = []
-    artifacts = []
-    error = None
-
-    deadline = time.time() + timeout
+    msg_id      = kc.execute(code)
+    stdout      = []
+    stderr      = []
+    result      = []
+    artifacts   = []
+    error       = None
+    deadline    = time.time() + timeout
 
     while time.time() < deadline:
         try:
-            msg = kc.get_iopub_msg(timeout=1)
+            msg      = kc.get_iopub_msg(timeout=1)
             msg_type = msg['msg_type']
-            content = msg['content']
+            content  = msg['content']
 
             if msg_type == 'stream':
-                if content['name'] == 'stdout':
-                    stdout_parts.append(content['text'])
-                elif content['name'] == 'stderr':
-                    stderr_parts.append(content['text'])
+                if content['name'] == 'stdout': stdout.append(content['text'])
+                else: stderr.append(content['text'])
 
             elif msg_type == 'execute_result':
-                result_parts.append(content['data'].get('text/plain', ''))
-                # Capture HTML output too
+                result.append(content['data'].get('text/plain', ''))
                 if 'text/html' in content['data']:
-                    result_parts.append(content['data']['text/html'])
+                    result.append(content['data']['text/html'])
 
             elif msg_type == 'display_data':
-                # ── Artifact capture — images, charts, plots ──────────────────
                 data = content.get('data', {})
                 if 'image/png' in data:
                     artifacts.append({
-                        "type": "image/png",
-                        "data": data['image/png'],  # already base64
+                        "type":     "image/png",
+                        "data":     data['image/png'],
                         "filename": f"artifact_{len(artifacts)+1}.png"
                     })
                 if 'image/jpeg' in data:
                     artifacts.append({
-                        "type": "image/jpeg",
-                        "data": data['image/jpeg'],
+                        "type":     "image/jpeg",
+                        "data":     data['image/jpeg'],
                         "filename": f"artifact_{len(artifacts)+1}.jpg"
                     })
-                if 'text/html' in data:
-                    result_parts.append(data['text/html'])
                 if 'text/plain' in data:
-                    result_parts.append(data['text/plain'])
+                    result.append(data['text/plain'])
 
             elif msg_type == 'error':
                 error = {
-                    "ename": content.get('ename', 'Error'),
-                    "evalue": content.get('evalue', ''),
+                    "ename":     content.get('ename', 'Error'),
+                    "evalue":    content.get('evalue', ''),
                     "traceback": "\n".join(
-                        # Strip ANSI escape codes
                         line.encode('ascii', errors='ignore').decode()
                         for line in content.get('traceback', [])
                     )
                 }
 
             elif msg_type == 'status' and content.get('execution_state') == 'idle':
-                # Kernel finished executing
                 break
 
         except queue.Empty:
             continue
-        except Exception as e:
+        except Exception:
             break
 
     return {
-        "stdout": "".join(stdout_parts),
-        "stderr": "".join(stderr_parts),
-        "result": "\n".join(result_parts),
+        "stdout":    "".join(stdout),
+        "stderr":    "".join(stderr),
+        "result":    "\n".join(result),
         "artifacts": artifacts,
-        "error": error
+        "error":     error
     }
 
-# ── Cleanup old idle kernels every 30 mins ────────────────────────────────────
-def cleanup_idle_kernels():
+# ── Cleanup idle kernels ──────────────────────────────────────────────────────
+def cleanup_loop():
     while True:
         time.sleep(1800)
-        now = time.time()
+        now     = time.time()
         to_kill = []
         with sessions_lock:
             for sid, data in sessions.items():
-                if now - data["last_used"] > 3600:  # 1 hour idle
+                if now - data["last_used"] > 3600:
                     to_kill.append(sid)
         for sid in to_kill:
             kill_kernel(sid)
-            print(f"[SANDBOX] Cleaned up idle kernel: {sid}")
 
-threading.Thread(target=cleanup_idle_kernels, daemon=True).start()
+threading.Thread(target=cleanup_loop, daemon=True).start()
 
-# ==========================================
-# ENDPOINTS
-# ==========================================
+# ── Endpoints ─────────────────────────────────────────────────────────────────
 
 @app.route('/')
 def status():
     return jsonify({
-        "status": "online",
+        "status":          "online",
         "active_sessions": len(sessions),
-        "message": "Jupyter Kernel Sandbox is running"
+        "message":         "Obsidian Jupyter Kernel Sandbox"
     })
 
-# ── Main execute endpoint (compatible with your ollama-tools) ─────────────────
 @app.route('/execute', methods=['POST'])
 def execute():
     if not check_auth():
-        return jsonify({"output": "Unauthorized: Invalid API Key"}), 401
+        return jsonify({"output": "Unauthorized"}), 401
 
-    data = request.json or {}
-    code = data.get("code", "")
+    data       = request.json or {}
+    code       = data.get("code", "")
     session_id = data.get("session_id", "default")
-    timeout = data.get("timeout", 60)
+    timeout    = data.get("timeout", 60)
 
     if not code:
-        return jsonify({"output": "Error: No code provided"}), 400
-
-    print(f"[SANDBOX] Executing code in session: {session_id}")
+        return jsonify({"output": "No code provided"}), 400
 
     try:
-        kc = get_or_create_kernel(session_id)
+        kc     = get_or_create_kernel(session_id)
         result = execute_code(kc, code, timeout=timeout)
 
-        # Build readable output
-        output_parts = []
-        if result["stdout"]: output_parts.append(result["stdout"])
-        if result["result"]: output_parts.append(result["result"])
+        parts = []
+        if result["stdout"]:  parts.append(result["stdout"])
+        if result["result"]:  parts.append(result["result"])
         if result["error"]:
-            output_parts.append(
-                f"Error: {result['error']['ename']}: {result['error']['evalue']}\n{result['error']['traceback']}"
+            parts.append(
+                f"Error: {result['error']['ename']}: {result['error']['evalue']}\n"
+                f"{result['error']['traceback']}"
             )
-        if result["stderr"] and not result["error"]:
-            output_parts.append(result["stderr"])
+        elif result["stderr"] and not result["error"]:
+            parts.append(result["stderr"])
 
         return jsonify({
-            "output": "\n".join(output_parts) or "Code executed successfully (no output)",
-            "artifacts": result["artifacts"],  # list of base64 images/charts
-            "error": result["error"]
+            "output":    "\n".join(parts) or "Executed successfully (no output)",
+            "artifacts": result["artifacts"],
+            "error":     result["error"]
         })
-
     except Exception as e:
         return jsonify({"output": f"Kernel Error: {str(e)}"}), 500
 
-# ── Session management ────────────────────────────────────────────────────────
+@app.route('/install', methods=['POST'])
+def install():
+    if not check_auth():
+        return jsonify({"error": "Unauthorized"}), 401
+
+    data       = request.json or {}
+    package    = data.get("package", "")
+    session_id = data.get("session_id", "default")
+
+    if not package:
+        return jsonify({"error": "No package specified"}), 400
+
+    kc     = get_or_create_kernel(session_id)
+    result = execute_code(
+        kc,
+        f"import subprocess; subprocess.run(['pip', 'install', '{package}', '-q'])",
+        timeout=120
+    )
+    return jsonify({"message": f"{package} installed", "output": result["stdout"]})
+
 @app.route('/sessions', methods=['GET'])
 def list_sessions():
     if not check_auth():
         return jsonify({"error": "Unauthorized"}), 401
-    return jsonify({
-        "sessions": [
-            {
-                "id": sid,
-                "created_at": data["created_at"],
-                "last_used": data["last_used"]
-            }
-            for sid, data in sessions.items()
-        ]
-    })
+    with sessions_lock:
+        return jsonify({
+            "sessions": [
+                {"id": sid, "created_at": d["created_at"], "last_used": d["last_used"]}
+                for sid, d in sessions.items()
+            ]
+        })
 
 @app.route('/sessions/<session_id>', methods=['DELETE'])
 def delete_session(session_id):
@@ -225,7 +214,6 @@ def delete_session(session_id):
     kill_kernel(session_id)
     return jsonify({"message": f"Session {session_id} terminated"})
 
-# ── New session ───────────────────────────────────────────────────────────────
 @app.route('/sessions', methods=['POST'])
 def create_session():
     if not check_auth():
@@ -234,44 +222,19 @@ def create_session():
     get_or_create_kernel(session_id)
     return jsonify({"session_id": session_id})
 
-# ── Install packages into a session kernel ────────────────────────────────────
-@app.route('/install', methods=['POST'])
-def install_package():
-    if not check_auth():
-        return jsonify({"error": "Unauthorized"}), 401
-    
-    data = request.json or {}
-    package = data.get("package", "")
-    session_id = data.get("session_id", "default")
-    
-    if not package:
-        return jsonify({"error": "No package specified"}), 400
-    
-    kc = get_or_create_kernel(session_id)
-    result = execute_code(kc, f"import subprocess; subprocess.run(['pip', 'install', '{package}', '-q'])", timeout=120)
-    
-    return jsonify({
-        "message": f"Package {package} installed",
-        "output": result["stdout"]
-    })
-
-# ── File upload into kernel environment ──────────────────────────────────────
 @app.route('/upload', methods=['POST'])
 def upload_file():
     if not check_auth():
         return jsonify({"error": "Unauthorized"}), 401
-    
-    data = request.json or {}
-    filename = data.get("filename", "uploaded_file")
-    file_b64 = data.get("base64", "")
+    data       = request.json or {}
+    filename   = data.get("filename", "uploaded_file")
+    file_b64   = data.get("base64", "")
     session_id = data.get("session_id", "default")
-    
+
     file_bytes = base64.b64decode(file_b64)
-    
-    # Write file via kernel
-    kc = get_or_create_kernel(session_id)
-    encoded = base64.b64encode(file_bytes).decode()
-    code = f"""
+    kc         = get_or_create_kernel(session_id)
+    encoded    = base64.b64encode(file_bytes).decode()
+    code       = f"""
 import base64
 with open('/tmp/{filename}', 'wb') as f:
     f.write(base64.b64decode('{encoded}'))
@@ -281,8 +244,7 @@ print('File written to /tmp/{filename}')
     return jsonify({"output": result["stdout"], "path": f"/tmp/{filename}"})
 
 if __name__ == "__main__":
-    print("[SANDBOX] Jupyter Kernel Sandbox starting on port 7860...")
-    # Pre-warm default kernel
+    print("[SANDBOX] Pre-warming default kernel...")
     get_or_create_kernel("default")
-    print("[SANDBOX] Default kernel ready!")
+    print("[SANDBOX] Ready on port 7860")
     app.run(host="0.0.0.0", port=7860)