Update app.py

app.py

@@ -3,13 +3,15 @@ import gradio as gr
 import subprocess, os, sys, time
 import urllib.request, urllib.error
 from fastapi import Request
-from fastapi.responses import Response, PlainTextResponse, HTMLResponse
+from fastapi.responses import Response, PlainTextResponse, HTMLResponse, StreamingResponse
 
 PORT = 9001
 FM_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "file-manager")
 MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
 
-# --- Launch the file manager as a background child process ---
+# =================================================================
+# Launch the file manager as a background child process
+# =================================================================
 proc = {"p": None}
 def start_fm():
     if proc["p"] and proc["p"].poll() is None:
@@ -23,7 +25,9 @@ def start_fm():
 
 start_fm()
 
-# --- GPU info: runs nvidia-smi INSIDE the gate (GPU only attached here) ---
+# =================================================================
+# GPU info: runs nvidia-smi INSIDE the gate (GPU only attached here)
+# =================================================================
 @spaces.GPU
 def gpu_report():
     import torch
@@ -41,7 +45,9 @@ def gpu_report():
         lines.append(f"nvidia-smi error: {e}")
     return "\n".join(lines)
 
-# --- LLM inference: model loads + generates INSIDE the gate ---
+# =================================================================
+# LLM inference: model loads + generates INSIDE the gate
+# =================================================================
 _state = {"model": None, "tok": None}
 
 @spaces.GPU
@@ -60,9 +66,11 @@ def run_llm(prompt: str):
     reply = tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     return f"[GPU: {torch.cuda.get_device_name(0)}]\n\n{reply}"
 
-# --- Gradio owns the ONLY public port 7860 ---
+# =================================================================
+# Gradio owns the ONLY public port 7860
+# =================================================================
 with gr.Blocks() as demo:
-    gr.Markdown("Self-editing computer. File manager + terminal at /fm")
+    gr.Markdown("Self-editing computer. File manager + terminal + Python REPL at /fm")
     with gr.Row():
         prompt = gr.Textbox(label="Ask the LLM")
         out = gr.Textbox(label="Reply")
@@ -74,26 +82,36 @@ demo.queue()
 demo.launch(server_name="0.0.0.0", server_port=7860,
             ssr_mode=False, prevent_thread_lock=True)
 
-app = demo.app
+app = demo.app  # Gradio's FastAPI app on 7860
 
-# --- Hot-reload: respawn ONLY the file-manager process, no Space rebuild ---
+# =================================================================
+# Hot-reload: respawn ONLY the file-manager process, no Space rebuild
+# =================================================================
 @app.get("/reload")
 def reload_fm():
     start_fm()
     time.sleep(1)
     return HTMLResponse('Reloaded file-manager process. <a href="/fm/">back to /fm</a>')
 
-# --- GPU info route ---
+# =================================================================
+# GPU info route
+# =================================================================
 @app.get("/gpuinfo")
 def gpuinfo():
     return PlainTextResponse(gpu_report())
 
-# --- LLM route ---
+# =================================================================
+# LLM route
+# =================================================================
 @app.get("/llm")
 def llm(q: str = "Say hello and tell me which GPU you are running on."):
     return PlainTextResponse(run_llm(q))
 
-# --- Proxy: /fm/... -> 127.0.0.1:9001/fm/... (methods + body forwarded) ---
+# =================================================================
+# Proxy: /fm/... -> 127.0.0.1:9001/fm/...
+# Streams text/plain (live shell) chunk by chunk; buffers everything else.
+# Forwards method + body so POST forms and the JSON REPL work.
+# =================================================================
 @app.api_route("/fm", methods=["GET", "POST"])
 @app.api_route("/fm/{path:path}", methods=["GET", "POST"])
 async def proxy(request: Request, path: str = ""):
@@ -104,16 +122,33 @@ async def proxy(request: Request, path: str = ""):
     fwd = {k: v for k, v in request.headers.items() if k.lower() != "host"}
     req = urllib.request.Request(url, data=body or None, method=request.method, headers=fwd)
     try:
-        with urllib.request.urlopen(req) as r:
-            data, status, headers = r.read(), r.status, r.headers
+        r = urllib.request.urlopen(req)
     except urllib.error.HTTPError as e:
-        data, status, headers = e.read(), e.code, e.headers
+        return Response(
+            content=e.read(), status_code=e.code,
+            headers={"content-type": e.headers.get("content-type", "text/plain")})
     except urllib.error.URLError:
         return Response(content=b"file-manager not up yet on 9001", status_code=502)
-    out = {k: v for k, v in headers.items()
-           if k.lower() in ("set-cookie", "location", "content-type")}
-    return Response(content=data, status_code=status, headers=out)
 
-# --- Keep the process alive (launch returned control) ---
+    ctype = r.headers.get("content-type", "text/html")
+    # Stream the live shell output; buffer normal pages
+    if ctype.startswith("text/plain"):
+        def gen():
+            while True:
+                chunk = r.read(1024)
+                if not chunk:
+                    break
+                yield chunk
+        return StreamingResponse(gen(), media_type="text/plain")
+
+    out_headers = {"content-type": ctype}
+    for k, v in r.headers.items():
+        if k.lower() in ("set-cookie", "location"):
+            out_headers[k] = v
+    return Response(content=r.read(), status_code=r.status, headers=out_headers)
+
+# =================================================================
+# Keep the process alive (launch returned control)
+# =================================================================
 while True:
     time.sleep(3600)