@neoline/hostbridge 0.2.2 → 1.1.0

package/src/server_control_mcp/manager.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import base64
+import hashlib
 import re
 import shlex
 import threading
@@ -8,7 +10,8 @@ import uuid
 from collections.abc import Callable
 from contextlib import suppress
 from dataclasses import dataclass, field
-from typing import Protocol
+from pathlib import Path
+from typing import Protocol, cast
 
 import pexpect
 
@@ -21,6 +24,10 @@ DEFAULT_REAPER_INTERVAL_SECONDS = 60
 MAX_OUTPUT_CHARS = 120_000
 REMOTE_JOB_ROOT = "$HOME/.hostbridge/jobs"
 REMOTE_SHELL = "bash"
+MAX_SHELL_TRANSFER_BYTES = 100 * 1024 * 1024
+DEFAULT_MAX_TRANSFER_BYTES = MAX_SHELL_TRANSFER_BYTES
+UPLOAD_BASE64_CHUNK_CHARS = 256 * 1024
+DOWNLOAD_CHUNK_BYTES = 5 * 1024 * 1024
 
 class PolicyDeniedError(RuntimeError):
     """Raised when a command is blocked by the configured policy."""
@@ -62,6 +69,7 @@ class RemoteSession:
     created_at: float = field(default_factory=time.time)
     last_used_at: float = field(default_factory=time.time)
     busy: bool = False
+    owner_agent_id: str | None = None
     last_output: str = ""
 
 
@@ -83,6 +91,7 @@ class RemoteTask:
     command: str
     job_dir: str
     remote_pid: int | None
+    owner_agent_id: str | None = None
     created_at: float = field(default_factory=time.time)
 
 
@@ -146,6 +155,7 @@ class SessionManager:
         remote_shell: str = REMOTE_SHELL,
         policy: PolicyConfig | None = None,
         audit_sink: AuditSink | None = None,
+        max_sessions_per_host: int = 0,
     ):
         self.hosts = hosts
         self.child_factory = child_factory
@@ -154,6 +164,7 @@ class SessionManager:
         self._tasks: dict[str, RemoteTask] = {}
         self._policy = policy
         self._audit = audit_sink
+        self._max_sessions_per_host = max(0, int(max_sessions_per_host))
         self._reaper_stop = threading.Event()
         self._reaper_thread: threading.Thread | None = None
 
@@ -223,8 +234,10 @@ class SessionManager:
         *,
         connect_timeout: int = DEFAULT_CONNECT_TIMEOUT,
         token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
     ) -> RemoteSession:
         host = self.hosts.get(host_id)
+        self._check_host_session_limit(host.id)
         child = self.child_factory(host.command, list(host.args), connect_timeout)
         token = token_factory()
         marker = f"__SCM_READY_{token}__"
@@ -247,6 +260,7 @@ class SessionManager:
             command=host.command,
             args=list(host.args),
             child=child,
+            owner_agent_id=owner_agent_id,
             last_output=_tail(str(child.before)),
         )
         self._sessions[session.id] = session
@@ -286,14 +300,23 @@ class SessionManager:
             closed.append(session_id)
         return {"closed": closed, "count": len(closed)}
 
-    def close_session(self, session_id: str) -> dict[str, object]:
+    def close_session(
+        self,
+        session_id: str,
+        *,
+        owner_agent_id: str | None = None,
+        force: bool = False,
+    ) -> dict[str, object]:
         session = self._get_session(session_id)
+        self._check_owner(session, owner_agent_id, force=force)
+        if session.busy and not force:
+            raise RemoteCommandError(f"session {session_id} is busy; retry later or close with force")
         session.child.close(force=True)
         self._sessions.pop(session_id, None)
-        self._emit("session_closed", session_id=session_id, host_id=session.host_id, outcome="user")
+        self._emit("session_closed", session_id=session_id, host_id=session.host_id, outcome="forced" if force else "user")
         return {"session_id": session_id, "closed": True}
 
-    def stop_task(self, session_id: str, task_id: str) -> dict[str, object]:
+    def stop_task(self, session_id: str, task_id: str, *, owner_agent_id: str | None = None, force: bool = False) -> dict[str, object]:
         """Best-effort terminate a background task by killing its remote PID."""
 
         task = self._tasks.get(task_id)
@@ -301,10 +324,11 @@ class SessionManager:
             raise RemoteCommandError(f"unknown task {task_id}")
         if task.session_id != session_id:
             raise RemoteCommandError(f"task {task_id} does not belong to session {session_id}")
+        self._check_task_owner(task, owner_agent_id, force=force)
         if task.remote_pid is None:
             raise RemoteCommandError(f"task {task_id} has no recorded remote pid")
         kill_script = f'if [ -f "{task.job_dir}/exit_code" ]; then echo already; else kill -TERM {task.remote_pid} 2>/dev/null && echo killed || echo noop; fi'
-        result = self.run_command(session_id, kill_script, timeout=10)
+        result = self.run_command(session_id, kill_script, timeout=10, owner_agent_id=owner_agent_id)
         outcome = result.output.strip().splitlines()[-1] if result.output.strip() else "unknown"
         self._emit(
             "task_stopped",
@@ -322,6 +346,8 @@ class SessionManager:
         *,
         timeout: int = DEFAULT_COMMAND_TIMEOUT,
         token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
+        output_limit: int = MAX_OUTPUT_CHARS,
     ) -> CommandResult:
         if not command.strip():
             raise RemoteCommandError("command must not be empty")
@@ -336,6 +362,7 @@ class SessionManager:
             )
             raise PolicyDeniedError(decision.reason)
         session = self._get_session(session_id)
+        self._check_owner(session, owner_agent_id)
         if session.busy:
             raise RemoteCommandError(f"session {session_id} is busy")
         if not session.child.isalive():
@@ -351,7 +378,7 @@ class SessionManager:
             session.child.expect(exit_pattern, timeout=timeout)
             output = _extract_between_markers(str(session.child.before), start_marker)
             exit_code = _parse_exit_marker(session.child.after, token)
-            result = CommandResult(session.id, session.host_id, command, exit_code, _tail(output))
+            result = CommandResult(session.id, session.host_id, command, exit_code, _tail(output, max(1, output_limit)))
             self._emit(
                 "command_run",
                 session_id=session.id,
@@ -381,6 +408,7 @@ class SessionManager:
         command: str,
         *,
         token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
     ) -> RemoteTask:
         if not command.strip():
             raise RemoteCommandError("command must not be empty")
@@ -406,7 +434,7 @@ class SessionManager:
             'date -u +%Y-%m-%dT%H:%M:%SZ > "$job_dir/finished_at") & '
             'pid=$!; printf "%s\\n" "$pid" > "$job_dir/pid"; printf "%s" "$pid"'
         )
-        result = self.run_command(session_id, script, timeout=10, token_factory=lambda: f"task_{task_id}")
+        result = self.run_command(session_id, script, timeout=10, token_factory=lambda: f"task_{task_id}", owner_agent_id=owner_agent_id)
         pid_match = re.search(r"(\d+)", result.output)
         task = RemoteTask(
             id=task_id,
@@ -415,17 +443,26 @@ class SessionManager:
             command=command,
             job_dir=job_dir,
             remote_pid=int(pid_match.group(1)) if pid_match else None,
+            owner_agent_id=owner_agent_id,
         )
         self._tasks[task.id] = task
         self._emit("task_started", session_id=session_id, task_id=task.id, host_id=task.host_id, command=command, outcome="ok")
         return task
 
-    def task_status(self, session_id: str, task_id: str, *, tail_lines: int = 80) -> dict[str, object]:
+    def task_status(
+        self,
+        session_id: str,
+        task_id: str,
+        *,
+        tail_lines: int = 80,
+        owner_agent_id: str | None = None,
+    ) -> dict[str, object]:
         task = self._tasks.get(task_id)
         if task is None:
             raise RemoteCommandError(f"unknown task {task_id}")
         if task.session_id != session_id:
             raise RemoteCommandError(f"task {task_id} does not belong to session {session_id}")
+        self._check_task_owner(task, owner_agent_id)
         safe_tail_lines = max(1, min(int(tail_lines), 500))
         script = (
             f'job_dir="{task.job_dir}"; '
@@ -436,7 +473,7 @@ class SessionManager:
             'printf "STATE=%s\\nPID=%s\\nEXIT=%s\\n---LOG---\\n" "$state" "$pid" "$exit_code"; '
             f'tail -n {safe_tail_lines} "$job_dir/stdout.log" 2>/dev/null || true'
         )
-        result = self.run_command(session_id, script, timeout=15)
+        result = self.run_command(session_id, script, timeout=15, owner_agent_id=owner_agent_id)
         metadata, _, log = result.output.partition("---LOG---\n")
         fields: dict[str, str] = {}
         for line in metadata.splitlines():
@@ -456,14 +493,320 @@ class SessionManager:
             "log_tail": log,
         }
 
-    def write(self, session_id: str, text: str) -> dict[str, object]:
+
+    def file_write_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        content: str,
+        *,
+        mode: str = "0644",
+        max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+        token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
+    ) -> dict[str, object]:
+        data = content.encode("utf-8")
+        return self._upload_bytes(
+            session_id,
+            data,
+            remote_path,
+            mode=mode,
+            max_bytes=max_bytes,
+            token_factory=token_factory,
+            owner_agent_id=owner_agent_id,
+        )
+
+    def file_read_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+        token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
+    ) -> dict[str, object]:
+        result = self._download_bytes(
+            session_id,
+            remote_path,
+            max_bytes=max_bytes,
+            token_factory=token_factory,
+            owner_agent_id=owner_agent_id,
+        )
+        data = cast(bytes, result.pop("data_bytes"))
+        result["content"] = data.decode("utf-8")
+        return result
+
+    def file_upload(
+        self,
+        session_id: str,
+        local_path: Path | str,
+        remote_path: str,
+        *,
+        mode: str = "0644",
+        max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+        token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
+    ) -> dict[str, object]:
+        path = Path(local_path).expanduser()
+        if not path.is_absolute():
+            raise RemoteCommandError("local_path must be absolute")
+        if not path.is_file():
+            raise RemoteCommandError(f"local file not found: {path}")
+        size = path.stat().st_size
+        self._check_transfer_size(size, max_bytes)
+        return self._upload_bytes(
+            session_id,
+            path.read_bytes(),
+            remote_path,
+            mode=mode,
+            max_bytes=max_bytes,
+            token_factory=token_factory,
+            owner_agent_id=owner_agent_id,
+            local_path=str(path),
+        )
+
+    def file_download(
+        self,
+        session_id: str,
+        remote_path: str,
+        local_path: Path | str,
+        *,
+        max_bytes: int = DEFAULT_MAX_TRANSFER_BYTES,
+        token_factory: TokenFactory = _new_token,
+        owner_agent_id: str | None = None,
+    ) -> dict[str, object]:
+        path = Path(local_path).expanduser()
+        if not path.is_absolute():
+            raise RemoteCommandError("local_path must be absolute")
+        result = self._download_bytes(
+            session_id,
+            remote_path,
+            max_bytes=max_bytes,
+            token_factory=token_factory,
+            owner_agent_id=owner_agent_id,
+        )
+        data = cast(bytes, result.pop("data_bytes"))
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(data)
+        result["local_path"] = str(path)
+        return result
+
+    def _upload_bytes(
+        self,
+        session_id: str,
+        data: bytes,
+        remote_path: str,
+        *,
+        mode: str,
+        max_bytes: int,
+        token_factory: TokenFactory,
+        owner_agent_id: str | None,
+        local_path: str | None = None,
+    ) -> dict[str, object]:
+        self._check_transfer_size(len(data), max_bytes)
+        self._validate_remote_path(remote_path)
+        b64 = base64.b64encode(data).decode("ascii")
+        sha256 = hashlib.sha256(data).hexdigest()
+        token = token_factory()
+        remote_q = shlex.quote(remote_path)
+        mode_q = shlex.quote(mode)
+        tmp_prefix = f"/tmp/hostbridge.upload.{token}"
+        tmp_b64 = f"{tmp_prefix}.b64"
+        tmp_out = f"{tmp_prefix}.out"
+        cleanup_script = f"rm -f {shlex.quote(tmp_b64)} {shlex.quote(tmp_out)}"
+        init_script = (
+            "set -e; "
+            f"target={remote_q}; "
+            'dir=$(dirname "$target"); mkdir -p "$dir"; '
+            f"rm -f {shlex.quote(tmp_b64)} {shlex.quote(tmp_out)}; "
+            f": > {shlex.quote(tmp_b64)}"
+        )
+        result = self.run_command(session_id, init_script, timeout=60, token_factory=lambda: token, owner_agent_id=owner_agent_id)
+        if result.exit_code != 0:
+            raise RemoteCommandError(result.output or f"failed to prepare upload for {remote_path}")
+        try:
+            self._stream_base64_to_remote_file(
+                session_id,
+                b64,
+                tmp_b64,
+                token=token,
+                owner_agent_id=owner_agent_id,
+            )
+            finalize_script = (
+                "set -e; "
+                f"target={remote_q}; mode={mode_q}; tmp_b64={shlex.quote(tmp_b64)}; tmp_out={shlex.quote(tmp_out)}; "
+                'base64 -d "$tmp_b64" > "$tmp_out"; '
+                'chmod "$mode" "$tmp_out"; mv "$tmp_out" "$target"; rm -f "$tmp_b64"; '
+                f"printf 'BYTES=%s\\nSHA256=%s\\n' {len(data)} {shlex.quote(sha256)}"
+            )
+            result = self.run_command(
+                session_id,
+                finalize_script,
+                timeout=60,
+                token_factory=lambda: token,
+                owner_agent_id=owner_agent_id,
+            )
+            if result.exit_code != 0:
+                raise RemoteCommandError(result.output or f"failed to finalize upload for {remote_path}")
+        except Exception:
+            with suppress(Exception):
+                self.run_command(session_id, cleanup_script, timeout=10, token_factory=lambda: token, owner_agent_id=owner_agent_id)
+            raise
+        payload: dict[str, object] = {
+            "session_id": session_id,
+            "remote_path": remote_path,
+            "bytes": len(data),
+            "sha256": sha256,
+            "strategy": "shell_base64_chunked",
+        }
+        if local_path is not None:
+            payload["local_path"] = local_path
+        self._emit("file_uploaded", session_id=session_id, command=remote_path, outcome=f"bytes={len(data)}")
+        return payload
+
+    def _stream_base64_to_remote_file(
+        self,
+        session_id: str,
+        encoded: str,
+        remote_tmp_b64: str,
+        *,
+        token: str,
+        owner_agent_id: str | None,
+    ) -> None:
+        session = self._get_session(session_id)
+        self._check_owner(session, owner_agent_id)
+        if session.busy:
+            raise RemoteCommandError(f"session {session_id} is busy")
+        if not session.child.isalive():
+            raise RemoteCommandError(f"session {session_id} is not alive")
+
+        heredoc = f"__HOSTBRIDGE_UPLOAD_{token}__"
+        done_pattern = re.compile(rf"__HOSTBRIDGE_UPLOAD_DONE_{re.escape(token)}__:(\d+)")
+        chunk_size = max(4, int(UPLOAD_BASE64_CHUNK_CHARS))
+        session.busy = True
+        try:
+            session.child.send(f"cat >> {shlex.quote(remote_tmp_b64)} <<'{heredoc}'\n")
+            for offset in range(0, len(encoded), chunk_size):
+                session.child.send(encoded[offset : offset + chunk_size])
+            session.child.send(f"\n{heredoc}\nprintf '\\n__HOSTBRIDGE_UPLOAD_DONE_{token}__:%s\\n' \"$?\"\n")
+            session.child.expect(done_pattern, timeout=60)
+            marker = session.child.after
+            text = marker.group(0) if hasattr(marker, "group") else str(marker)
+            match = re.search(r":(\d+)", text)
+            if match is None or int(match.group(1)) != 0:
+                raise RemoteCommandError(f"failed to stream upload chunk for {session_id}")
+        finally:
+            session.busy = False
+            session.last_used_at = time.time()
+            session.last_output = _tail(str(session.child.before))
+
+    def _download_bytes(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        max_bytes: int,
+        token_factory: TokenFactory,
+        owner_agent_id: str | None,
+    ) -> dict[str, object]:
+        self._validate_remote_path(remote_path)
+        token = token_factory()
+        remote_q = shlex.quote(remote_path)
+        metadata_script = (
+            "set -e; "
+            f"target={remote_q}; "
+            'bytes=$(wc -c < "$target" | tr -d " "); '
+            f"if [ \"$bytes\" -gt {int(max_bytes)} ]; then echo 'FILE_TOO_LARGE remote pull required' >&2; exit 73; fi; "
+            "if command -v sha256sum >/dev/null 2>&1; then sha=$(sha256sum \"$target\" | awk '{print $1}'); "
+            "else sha=$(shasum -a 256 \"$target\" | awk '{print $1}'); fi; "
+            "printf 'BYTES=%s\\nSHA256=%s\\n' \"$bytes\" \"$sha\""
+        )
+        result = self.run_command(
+            session_id,
+            metadata_script,
+            timeout=60,
+            token_factory=lambda: token,
+            owner_agent_id=owner_agent_id,
+        )
+        if result.exit_code == 73:
+            self._raise_remote_pull_required(max_bytes)
+        if result.exit_code != 0:
+            raise RemoteCommandError(result.output or f"failed to download {remote_path}")
+        fields: dict[str, str] = {}
+        for line in result.output.splitlines():
+            key, _, value = line.partition("=")
+            if key:
+                fields[key] = value
+        byte_count = int(fields.get("BYTES") or 0)
+        self._check_transfer_size(byte_count, max_bytes)
+        chunks: list[bytes] = []
+        chunk_size = max(1, int(DOWNLOAD_CHUNK_BYTES))
+        for index in range((byte_count + chunk_size - 1) // chunk_size):
+            chunk_script = (
+                "set -e; "
+                f"dd if={remote_q} bs={chunk_size} skip={index} count=1 2>/dev/null | base64"
+            )
+            output_limit = max(MAX_OUTPUT_CHARS, chunk_size * 2 + 4096)
+            result = self.run_command(
+                session_id,
+                chunk_script,
+                timeout=60,
+                token_factory=lambda: token,
+                owner_agent_id=owner_agent_id,
+                output_limit=output_limit,
+            )
+            if result.exit_code != 0:
+                raise RemoteCommandError(result.output or f"failed to download chunk {index} from {remote_path}")
+            chunks.append(base64.b64decode("".join(result.output.split())))
+        data = b"".join(chunks)
+        sha256 = hashlib.sha256(data).hexdigest()
+        remote_sha = fields.get("SHA256") or sha256
+        if remote_sha != "ignored" and sha256 != remote_sha:
+            raise RemoteCommandError("download checksum mismatch")
+        self._emit("file_downloaded", session_id=session_id, command=remote_path, outcome=f"bytes={len(data)}")
+        return {
+            "session_id": session_id,
+            "remote_path": remote_path,
+            "bytes": len(data),
+            "sha256": remote_sha,
+            "strategy": "shell_base64_chunked",
+            "data_bytes": data,
+        }
+
+    def _validate_remote_path(self, remote_path: str) -> None:
+        if not remote_path or remote_path.strip() in {"", "/"}:
+            raise RemoteCommandError("remote_path must not be empty or root")
+
+    def _check_transfer_size(self, size: int, max_bytes: int) -> None:
+        if size > int(max_bytes):
+            self._raise_remote_pull_required(max_bytes)
+
+    def _raise_remote_pull_required(self, max_bytes: int) -> None:
+        raise RemoteCommandError(
+            f"file exceeds HostBridge shell transfer limit ({int(max_bytes)} bytes); use remote pull (curl/wget/modelscope/git) for large files"
+        )
+
+    def write(self, session_id: str, text: str, *, owner_agent_id: str | None = None) -> dict[str, object]:
         session = self._get_session(session_id)
+        self._check_owner(session, owner_agent_id)
+        if session.busy:
+            raise RemoteCommandError(f"session {session_id} is busy")
+        if not session.child.isalive():
+            raise RemoteCommandError(f"session {session_id} is not alive")
         session.child.sendline(text)
         session.last_used_at = time.time()
         return {"session_id": session_id, "written": True}
 
-    def read(self, session_id: str, *, timeout: float = 1.0, max_chars: int = 20_000) -> dict[str, object]:
+    def read(
+        self,
+        session_id: str,
+        *,
+        timeout: float = 1.0,
+        max_chars: int = 20_000,
+        owner_agent_id: str | None = None,
+    ) -> dict[str, object]:
         session = self._get_session(session_id)
+        self._check_owner(session, owner_agent_id)
         chunks: list[str] = []
         deadline = time.time() + max(0.0, timeout)
         while len("".join(chunks)) < max_chars:
@@ -488,6 +831,25 @@ class SessionManager:
         except KeyError as exc:
             raise RemoteCommandError(f"unknown session {session_id}") from exc
 
+    def _check_owner(self, session: RemoteSession, owner_agent_id: str | None, *, force: bool = False) -> None:
+        if force or owner_agent_id is None or session.owner_agent_id is None:
+            return
+        if owner_agent_id != session.owner_agent_id:
+            raise RemoteCommandError(f"session {session.id} is owned by another agent")
+
+    def _check_task_owner(self, task: RemoteTask, owner_agent_id: str | None, *, force: bool = False) -> None:
+        if force or owner_agent_id is None or task.owner_agent_id is None:
+            return
+        if owner_agent_id != task.owner_agent_id:
+            raise RemoteCommandError(f"task {task.id} is owned by another agent")
+
+    def _check_host_session_limit(self, host_id: str) -> None:
+        if self._max_sessions_per_host <= 0:
+            return
+        open_count = sum(1 for session in self._sessions.values() if session.host_id == host_id and session.child.isalive())
+        if open_count >= self._max_sessions_per_host:
+            raise RemoteCommandError(f"host {host_id} reached session limit {self._max_sessions_per_host}")
+
     def _build_remote_line(self, command: str, token: str, start_marker: str) -> str:
         quoted_command = shlex.quote(command)
         return (
@@ -507,6 +869,7 @@ class SessionManager:
             "args": session.args,
             "alive": session.child.isalive(),
             "busy": session.busy,
+            "owner_agent_id": session.owner_agent_id,
             "created_at": session.created_at,
             "last_used_at": session.last_used_at,
         }