Add file transfer (push/pull, <=100MB); dropdown allow_custom_value

app.py

@@ -23,6 +23,9 @@ second so a leaked HF read-token alone can't drive your shells.
 import os
 import time
 import uuid
+import shutil
+import tempfile
+import pathlib
 import threading
 from collections import defaultdict
 
@@ -32,17 +35,29 @@ SECRET = os.environ.get("RELAY_SECRET", "")          # set this as a Space secre
 CLIENT_STALE_AFTER = 30                                # seconds w/o a poll => "offline"
 MAX_OUTPUT_CHARS = 200_000                             # cap stored output per command
 MAX_HISTORY_PER_CLIENT = 200                           # ring-buffer of commands
+MAX_FILE_BYTES = 100 * 1024 * 1024                     # 100 MB cap on transfers
+
+FILES_DIR = pathlib.Path(tempfile.gettempdir()) / "relay_files"
+FILES_DIR.mkdir(exist_ok=True)
 
 LOCK = threading.Lock()
 CLIENTS: dict[str, dict] = {}                          # cid -> {last_seen, meta}
 COMMANDS: dict[str, dict] = {}                         # command_id -> record
 BY_CLIENT: dict[str, list[str]] = defaultdict(list)    # cid -> [command_id, ...]
+PUT_FILES: dict[str, dict] = {}                        # file_id -> {path, filename} (server->client payloads)
 
 
 def _now() -> float:
     return time.time()
 
 
+def _human(n: float) -> str:
+    for unit in ("B", "KB", "MB", "GB"):
+        if n < 1024 or unit == "GB":
+            return f"{n:.0f}{unit}" if unit == "B" else f"{n:.1f}{unit}"
+        n /= 1024
+
+
 def _check(secret: str):
     if SECRET and secret != SECRET:
         raise gr.Error("bad secret")
@@ -82,9 +97,15 @@ def api_poll(client_id: str, secret: str):
                 rec["acknowledged_at"] = _now()
                 out.append({
                     "id": rec["id"],
-                    "session": rec["session"],
-                    "command": rec["command"],
-                    "timeout": rec["timeout"],
+                    "type": rec.get("type", "exec"),
+                    "session": rec.get("session", "main"),
+                    "command": rec.get("command", ""),
+                    "timeout": rec.get("timeout", 120),
+                    # file-transfer fields (None for plain exec commands)
+                    "file_id": rec.get("file_id"),
+                    "dest": rec.get("dest"),
+                    "remote_path": rec.get("remote_path"),
+                    "filename": rec.get("filename"),
                 })
     return {"commands": out, "ts": _now()}
 
@@ -110,6 +131,43 @@ def api_result(command_id: str, client_id: str, secret: str,
     return {"ok": True}
 
 
+def api_fetch_file(file_id: str, secret: str):
+    """Daemon calls this to download a server->client (push) payload."""
+    _check(secret)
+    info = PUT_FILES.get(file_id)
+    if not info or not os.path.isfile(info["path"]):
+        raise gr.Error("unknown or expired file_id")
+    return info["path"]
+
+
+def api_upload_file(command_id: str, client_id: str, secret: str, fileobj, error: str):
+    """Daemon calls this to post the result of a client->server (pull) request.
+    `fileobj` is the uploaded file path on success, or None when `error` is set."""
+    _check(secret)
+    with LOCK:
+        rec = COMMANDS.get(command_id)
+        if rec is None:  # relay restarted — keep the file anyway
+            rec = COMMANDS[command_id] = {
+                "id": command_id, "client_id": client_id, "session": "file",
+                "type": "get", "command": "(file fetch — relay restarted)",
+                "created": _now(), "status": "pending",
+            }
+            BY_CLIENT[client_id].append(command_id)
+        rec["status"] = "done"
+        rec["finished_at"] = _now()
+        _touch_client(client_id)
+        if fileobj:
+            fname = os.path.basename(fileobj)
+            saved = FILES_DIR / f"{command_id}__{fname}"
+            shutil.copy(fileobj, saved)
+            sz = saved.stat().st_size
+            rec.update(stored_path=str(saved), filename=fname, size=sz,
+                       exit_code=0, stdout=f"fetched {fname} ({_human(sz)})")
+        else:
+            rec.update(exit_code=1, stdout=(error or "fetch failed"))
+    return {"ok": True}
+
+
 # --------------------------------------------------------------------------- #
 # Human UI helpers                                                             #
 # --------------------------------------------------------------------------- #
@@ -151,6 +209,65 @@ def ui_send(client_label: str, session: str, command: str, timeout: int):
     return f"➡️ sent to `{cid}` [{session}]", ""
 
 
+def ui_put_file(client_label: str, fileobj, dest: str):
+    """Queue a push: send an uploaded file to the selected container."""
+    cid = _strip_dot(client_label)
+    if not cid:
+        return "⚠️ pick a container first", None
+    if not fileobj:
+        return "⚠️ choose a file to send", None
+    sz = os.path.getsize(fileobj)
+    if sz > MAX_FILE_BYTES:
+        return f"⚠️ file is {_human(sz)} — over the {_human(MAX_FILE_BYTES)} limit", None
+    fname = os.path.basename(fileobj)
+    file_id = uuid.uuid4().hex
+    saved = FILES_DIR / f"{file_id}__{fname}"
+    shutil.copy(fileobj, saved)
+    dest = (dest or "").strip() or fname
+    cmd_id = uuid.uuid4().hex
+    with LOCK:
+        PUT_FILES[file_id] = {"path": str(saved), "filename": fname}
+        COMMANDS[cmd_id] = {
+            "id": cmd_id, "client_id": cid, "session": "file", "type": "put",
+            "file_id": file_id, "dest": dest, "filename": fname, "size": sz,
+            "command": f"📤 send {fname} ({_human(sz)}) → {dest}",
+            "timeout": 300, "status": "pending", "created": _now(),
+        }
+        BY_CLIENT[cid].append(cmd_id)
+    return f"➡️ sending `{fname}` to `{cid}`:`{dest}`", None
+
+
+def ui_get_file(client_label: str, remote_path: str):
+    """Queue a pull: ask the container for a file at `remote_path`."""
+    cid = _strip_dot(client_label)
+    if not cid:
+        return "⚠️ pick a container first"
+    remote_path = (remote_path or "").strip()
+    if not remote_path:
+        return "⚠️ enter a path on the container"
+    cmd_id = uuid.uuid4().hex
+    with LOCK:
+        COMMANDS[cmd_id] = {
+            "id": cmd_id, "client_id": cid, "session": "file", "type": "get",
+            "remote_path": remote_path, "command": f"📥 fetch {remote_path}",
+            "timeout": 300, "status": "pending", "created": _now(),
+        }
+        BY_CLIENT[cid].append(cmd_id)
+    return f"➡️ requested `{remote_path}` from `{cid}`"
+
+
+def _fetched_files(client_label: str) -> list[str]:
+    cid = _strip_dot(client_label)
+    paths = []
+    with LOCK:
+        for cmd_id in BY_CLIENT.get(cid, []):
+            r = COMMANDS.get(cmd_id)
+            if r and r.get("type") == "get" and r.get("stored_path") \
+                    and os.path.isfile(r["stored_path"]):
+                paths.append(r["stored_path"])
+    return paths
+
+
 _STATUS_ICON = {"pending": "⏳", "acknowledged": "📨", "done": "✅"}
 
 
@@ -191,9 +308,10 @@ def ui_refresh(client_label: str):
     return gr.update(choices=choices, value=keep), ui_history(keep or "")
 
 
-def ui_tick(client_label: str, last_choices: list[str], last_hist: str):
+def ui_tick(client_label: str, last_choices: list[str], last_hist: str,
+            last_files: list[str]):
     """Timer handler that only emits updates when something actually changed,
-    so the dropdown and output don't repaint (flicker) every tick."""
+    so the dropdown, output, and downloads don't repaint (flicker) every tick."""
     choices = _client_choices()
     if choices == last_choices:
         dd_out = gr.skip()                       # dropdown unchanged -> leave it alone
@@ -204,7 +322,9 @@ def ui_tick(client_label: str, last_choices: list[str], last_hist: str):
         label_for_hist = keep or ""
     hist = ui_history(label_for_hist)
     hist_out = gr.skip() if hist == last_hist else hist
-    return dd_out, hist_out, choices, hist
+    files = _fetched_files(label_for_hist)
+    files_out = gr.skip() if files == last_files else gr.update(value=files)
+    return dd_out, hist_out, files_out, choices, hist, files
 
 
 # --------------------------------------------------------------------------- #
@@ -220,7 +340,7 @@ with gr.Blocks(title="NII Relay", theme=gr.themes.Soft()) as demo:
 
     with gr.Row():
         client_dd = gr.Dropdown(label="Container", choices=_client_choices(),
-                                interactive=True, scale=3)
+                                interactive=True, allow_custom_value=True, scale=3)
         refresh_btn = gr.Button("↻ Refresh", scale=1)
 
     with gr.Row():
@@ -235,9 +355,28 @@ with gr.Blocks(title="NII Relay", theme=gr.themes.Soft()) as demo:
 
     history_md = gr.Markdown("_Select a container to see its command history._")
 
+    with gr.Accordion("📁 File transfer (≤ 100 MB)", open=False):
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("**Send to container** (push)")
+                put_file = gr.File(label="file", type="filepath")
+                put_dest = gr.Textbox(label="destination path on container",
+                                      placeholder="/mnt/data/  or  /mnt/data/foo.bin")
+                put_btn = gr.Button("Send ⬆", variant="primary")
+                put_status = gr.Markdown("")
+            with gr.Column():
+                gr.Markdown("**Fetch from container** (pull)")
+                get_path = gr.Textbox(label="path on container",
+                                      placeholder="/mnt/data/train.log")
+                get_btn = gr.Button("Fetch ⬇", variant="primary")
+                get_status = gr.Markdown("")
+                downloads = gr.Files(label="fetched files (click to download)",
+                                     interactive=False)
+
     # remember what we last pushed, so the timer can skip no-op repaints
     last_choices = gr.State([])
     last_hist = gr.State("")
+    last_files = gr.State([])
 
     timer = gr.Timer(3.0)
 
@@ -248,8 +387,11 @@ with gr.Blocks(title="NII Relay", theme=gr.themes.Soft()) as demo:
                       [send_status, command_tb])
     refresh_btn.click(ui_refresh, [client_dd], [client_dd, history_md])
     client_dd.change(ui_history, [client_dd], history_md)
-    timer.tick(ui_tick, [client_dd, last_choices, last_hist],
-               [client_dd, history_md, last_choices, last_hist])
+    client_dd.change(_fetched_files, [client_dd], downloads)
+    put_btn.click(ui_put_file, [client_dd, put_file, put_dest], [put_status, put_file])
+    get_btn.click(ui_get_file, [client_dd, get_path], get_status)
+    timer.tick(ui_tick, [client_dd, last_choices, last_hist, last_files],
+               [client_dd, history_md, downloads, last_choices, last_hist, last_files])
 
     # --- hidden machine API endpoints (called by the daemon) ---
     with gr.Row(visible=False):
@@ -260,17 +402,27 @@ with gr.Blocks(title="NII Relay", theme=gr.themes.Soft()) as demo:
         a_session = gr.Textbox()
         a_exit = gr.Number()
         a_stdout = gr.Textbox()
+        a_fileid = gr.Textbox()
+        a_err = gr.Textbox()
+        a_file_in = gr.File()
+        a_file_out = gr.File()
         a_out = gr.JSON()
         b_reg = gr.Button()
         b_poll = gr.Button()
         b_res = gr.Button()
+        b_fetch = gr.Button()
+        b_upload = gr.Button()
         b_reg.click(api_register, [a_cid, a_secret, a_meta], a_out, api_name="register")
         b_poll.click(api_poll, [a_cid, a_secret], a_out, api_name="poll")
         b_res.click(api_result, [a_cmdid, a_cid, a_secret, a_exit, a_stdout, a_session],
                     a_out, api_name="result")
+        b_fetch.click(api_fetch_file, [a_fileid, a_secret], a_file_out, api_name="fetch_file")
+        b_upload.click(api_upload_file, [a_cmdid, a_cid, a_secret, a_file_in, a_err],
+                       a_out, api_name="upload_result")
 
 
 if __name__ == "__main__":
     demo.queue(default_concurrency_limit=16).launch(
-        server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860))
+        server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)),
+        max_file_size="150mb",
     )