@neoline/hostbridge 1.4.3 → 2.0.0

package/src/server_control_mcp/services.py

@@ -0,0 +1,508 @@
+"""Daemon-owned HostBridge session and execution services."""
+
+from __future__ import annotations
+
+import hashlib
+import re
+import shlex
+import threading
+import time
+import uuid
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass, field
+
+from .hosts import HostConfig, HostRegistry
+from .policy import AuditSink, PolicyConfig, audit_event
+from .transports.base import ExecResult, ShellReadResult, TransferResult, Transport, TransportBusy, TransportError
+
+TransportFactory = Callable[[HostConfig], Transport]
+
+
+class ServiceError(RuntimeError):
+    def __init__(
+        self,
+        code: str,
+        message: str,
+        *,
+        retryable: bool = False,
+        details: dict[str, object] | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.retryable = retryable
+        self.details = details or {}
+
+
+@dataclass(slots=True)
+class ServiceSession:
+    session_id: str
+    host_id: str
+    owner: str | None
+    transport: Transport = field(repr=False)
+    created_at: float = field(default_factory=time.time)
+    last_used_at: float = field(default_factory=time.time)
+    active_operations: int = 0
+
+    def describe(self) -> dict[str, object]:
+        return {
+            "session_id": self.session_id,
+            "host_id": self.host_id,
+            "owner": self.owner,
+            "created_at": self.created_at,
+            "last_used_at": self.last_used_at,
+            "active_operations": self.active_operations,
+            "capabilities": {
+                "separate_stderr": self.transport.capabilities.separate_stderr,
+                "native_binary_transfer": self.transport.capabilities.native_binary_transfer,
+                "parallel_channels": self.transport.capabilities.parallel_channels,
+            },
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class ServiceExecResult:
+    session_id: str
+    exit_code: int | None
+    stdout: bytes
+    stderr: bytes
+    timed_out: bool
+
+
+@dataclass(frozen=True, slots=True)
+class ServiceTask:
+    task_id: str
+    session_id: str
+    host_id: str
+    command: str
+    job_dir: str
+    remote_pid: int | None
+    owner: str | None
+
+
+class HostBridgeServices:
+    def __init__(
+        self,
+        hosts: HostRegistry,
+        *,
+        transport_factory: TransportFactory,
+        policy: PolicyConfig | None = None,
+        audit_sink: AuditSink | None = None,
+    ) -> None:
+        self.hosts = hosts
+        self._transport_factory = transport_factory
+        self._policy = policy
+        self._audit = audit_sink
+        self._sessions: dict[str, ServiceSession] = {}
+        self._tasks: dict[str, ServiceTask] = {}
+        self._sessions_lock = threading.RLock()
+        self._reaper_stop = threading.Event()
+        self._reaper_thread: threading.Thread | None = None
+
+    def open_session(self, host_id: str, *, owner: str | None = None) -> ServiceSession:
+        try:
+            host = self.hosts.get(host_id)
+        except KeyError as exc:
+            raise ServiceError("not_found", str(exc)) from exc
+        try:
+            transport = self._transport_factory(host)
+        except Exception as exc:
+            raise ServiceError("connection_lost", f"failed to connect to host {host_id}", retryable=True) from exc
+        session = ServiceSession(uuid.uuid4().hex, host.id, owner, transport)
+        with self._sessions_lock:
+            self._sessions[session.session_id] = session
+        self._emit("session_opened", session_id=session.session_id, host_id=host.id, owner_agent_id=owner, outcome="ok")
+        return session
+
+    def list_sessions(self) -> list[dict[str, object]]:
+        with self._sessions_lock:
+            return [session.describe() for session in self._sessions.values()]
+
+    def exec(
+        self,
+        session_id: str,
+        command: str,
+        *,
+        stdin: Iterable[bytes] | None = None,
+        timeout: float = 60,
+        output_limit: int = 1_000_000,
+        owner: str | None = None,
+    ) -> ServiceExecResult:
+        if not command.strip():
+            raise ServiceError("invalid_request", "command must not be empty")
+        if self._policy is not None:
+            decision = self._policy.evaluate(command)
+            if not decision.allowed:
+                raise ServiceError("denied", decision.reason, details={"matched_pattern": decision.matched_pattern})
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+        with self._sessions_lock:
+            session.active_operations += 1
+        try:
+            result: ExecResult = session.transport.exec(
+                command,
+                stdin=stdin,
+                timeout=timeout,
+                output_limit=output_limit,
+            )
+        except TransportBusy as exc:
+            raise ServiceError("busy", str(exc), retryable=True) from exc
+        except TransportError as exc:
+            raise ServiceError("remote_error", str(exc), retryable=True) from exc
+        finally:
+            with self._sessions_lock:
+                session.active_operations -= 1
+                session.last_used_at = time.time()
+        self._emit(
+            "command_run",
+            session_id=session.session_id,
+            host_id=session.host_id,
+            owner_agent_id=owner,
+            command=command,
+            outcome="timeout" if result.timed_out else f"exit={result.exit_code}",
+        )
+        return ServiceExecResult(
+            session.session_id,
+            result.exit_code,
+            result.stdout,
+            result.stderr,
+            result.timed_out,
+        )
+
+    def close_session(self, session_id: str, *, owner: str | None = None, force: bool = False) -> None:
+        session = self._get_session(session_id)
+        if not force:
+            self._check_owner(session, owner)
+        with self._sessions_lock:
+            removed = self._sessions.pop(session_id, None)
+        if removed is not None:
+            with self._sessions_lock:
+                self._tasks = {task_id: task for task_id, task in self._tasks.items() if task.session_id != session_id}
+            removed.transport.close()
+            self._emit(
+                "session_closed",
+                session_id=session_id,
+                host_id=removed.host_id,
+                owner_agent_id=owner,
+                outcome="forced" if force else "user",
+            )
+
+    def close_all(self) -> dict[str, object]:
+        with self._sessions_lock:
+            sessions = list(self._sessions.values())
+            self._sessions.clear()
+            self._tasks.clear()
+        for session in sessions:
+            session.transport.close()
+            self._emit("session_closed", session_id=session.session_id, host_id=session.host_id, outcome="shutdown")
+        return {"closed": [session.session_id for session in sessions], "count": len(sessions)}
+
+    def reap_idle(self, timeout: float) -> int:
+        if timeout <= 0:
+            return 0
+        now = time.time()
+        with self._sessions_lock:
+            expired = [
+                session
+                for session in self._sessions.values()
+                if session.active_operations == 0 and now - session.last_used_at > timeout
+            ]
+            for session in expired:
+                self._sessions.pop(session.session_id, None)
+                self._tasks = {
+                    task_id: task
+                    for task_id, task in self._tasks.items()
+                    if task.session_id != session.session_id
+                }
+        for session in expired:
+            session.transport.close()
+            self._emit("session_idle_closed", session_id=session.session_id, host_id=session.host_id, outcome="expired")
+        return len(expired)
+
+    def start_idle_reaper(self, *, timeout: float, interval: float = 30) -> None:
+        if timeout <= 0 or interval <= 0:
+            return
+        if self._reaper_thread is not None and self._reaper_thread.is_alive():
+            return
+        self._reaper_stop.clear()
+
+        def reap_loop() -> None:
+            while not self._reaper_stop.wait(interval):
+                self.reap_idle(timeout)
+
+        self._reaper_thread = threading.Thread(target=reap_loop, daemon=True, name="hostbridge-v1-reaper")
+        self._reaper_thread.start()
+
+    def stop_idle_reaper(self) -> None:
+        self._reaper_stop.set()
+        thread = self._reaper_thread
+        if thread is not None and thread.is_alive():
+            thread.join(timeout=2)
+        self._reaper_thread = None
+
+    def start_task(self, session_id: str, command: str, *, owner: str | None = None) -> ServiceTask:
+        if not command.strip():
+            raise ServiceError("invalid_request", "command must not be empty")
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+        task_id = uuid.uuid4().hex
+        job_dir = f"/tmp/hostbridge/jobs/{task_id}"
+        quoted_command = shlex.quote(command)
+        script = (
+            f"job_dir={shlex.quote(job_dir)}; mkdir -p \"$job_dir\"; "
+            f"({quoted_command} > \"$job_dir/stdout.log\" 2>&1; "
+            "status=$?; printf '%s\n' \"$status\" > \"$job_dir/exit_code\") & "
+            "pid=$!; printf '%s\n' \"$pid\" > \"$job_dir/pid\"; printf '%s' \"$pid\""
+        )
+        result = self.exec(session_id, script, timeout=10, output_limit=4096, owner=owner)
+        pid_match = re.search(rb"(\d+)", result.stdout)
+        task = ServiceTask(
+            task_id,
+            session_id,
+            session.host_id,
+            command,
+            job_dir,
+            int(pid_match.group(1)) if pid_match else None,
+            owner,
+        )
+        with self._sessions_lock:
+            self._tasks[task_id] = task
+        self._emit(
+            "task_started",
+            session_id=session_id,
+            task_id=task_id,
+            host_id=session.host_id,
+            owner_agent_id=owner,
+            command=command,
+            outcome="ok",
+        )
+        return task
+
+    def task_status(
+        self,
+        session_id: str,
+        task_id: str,
+        *,
+        owner: str | None = None,
+        tail_lines: int = 80,
+    ) -> dict[str, object]:
+        task = self._get_task(session_id, task_id, owner)
+        safe_tail = max(1, min(int(tail_lines), 500))
+        script = (
+            f"job_dir={shlex.quote(task.job_dir)}; pid=$(cat \"$job_dir/pid\" 2>/dev/null || true); "
+            "if [ -f \"$job_dir/exit_code\" ]; then state=finished; exit_code=$(cat \"$job_dir/exit_code\"); "
+            "elif [ -n \"$pid\" ] && kill -0 \"$pid\" 2>/dev/null; then state=running; exit_code=''; "
+            "else state=unknown; exit_code=''; fi; "
+            "printf 'STATE=%s\nPID=%s\nEXIT=%s\n---LOG---\n' \"$state\" \"$pid\" \"$exit_code\"; "
+            f"tail -n {safe_tail} \"$job_dir/stdout.log\" 2>/dev/null || true"
+        )
+        result = self.exec(session_id, script, timeout=15, output_limit=1_000_000, owner=owner)
+        metadata, _, log = result.stdout.decode("utf-8", errors="replace").partition("---LOG---\n")
+        fields: dict[str, str] = {}
+        for line in metadata.splitlines():
+            key, separator, value = line.partition("=")
+            if separator:
+                fields[key] = value
+        exit_text = fields.get("EXIT", "")
+        return {
+            "task_id": task.task_id,
+            "session_id": session_id,
+            "host_id": task.host_id,
+            "command": task.command,
+            "job_dir": task.job_dir,
+            "state": fields.get("STATE", "unknown"),
+            "pid": fields.get("PID") or task.remote_pid,
+            "exit_code": int(exit_text) if exit_text.isdigit() else None,
+            "log_tail": log,
+        }
+
+    def cancel_task(self, session_id: str, task_id: str, *, owner: str | None = None) -> dict[str, object]:
+        task = self._get_task(session_id, task_id, owner)
+        if task.remote_pid is None:
+            raise ServiceError("remote_error", f"task {task_id} has no recorded remote pid")
+        script = (
+            f"job_dir={shlex.quote(task.job_dir)}; "
+            "if [ -f \"$job_dir/exit_code\" ]; then printf already; "
+            f"elif kill -TERM {task.remote_pid} 2>/dev/null; then printf killed; else printf noop; fi"
+        )
+        result = self.exec(session_id, script, timeout=10, output_limit=4096, owner=owner)
+        outcome = result.stdout.decode("utf-8", errors="replace").strip().splitlines()[-1] if result.stdout.strip() else "unknown"
+        self._emit("task_cancelled", session_id=session_id, task_id=task_id, outcome=outcome)
+        return {"task_id": task_id, "stopped": outcome in {"killed", "already"}, "outcome": outcome}
+
+    def upload(
+        self,
+        session_id: str,
+        chunks: Iterable[bytes],
+        remote_path: str,
+        *,
+        mode: int,
+        expected_size: int | None,
+        expected_sha256: str | None,
+        owner: str | None = None,
+    ) -> TransferResult:
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+        with self._sessions_lock:
+            session.active_operations += 1
+        try:
+            result = session.transport.upload(
+                chunks,
+                remote_path,
+                mode=mode,
+                expected_size=expected_size,
+            )
+        except TransportError as exc:
+            raise ServiceError("remote_error", str(exc), retryable=True) from exc
+        finally:
+            with self._sessions_lock:
+                session.active_operations -= 1
+                session.last_used_at = time.time()
+        if expected_size is not None and result.bytes_transferred != expected_size:
+            raise ServiceError(
+                "integrity_error",
+                f"transfer size mismatch: expected {expected_size}, received {result.bytes_transferred}",
+            )
+        if expected_sha256 is not None and result.sha256.lower() != expected_sha256.lower():
+            raise ServiceError("integrity_error", "transfer SHA-256 mismatch")
+        self._emit(
+            "file_uploaded",
+            session_id=session_id,
+            host_id=session.host_id,
+            owner_agent_id=owner,
+            command=remote_path,
+            outcome=f"bytes={result.bytes_transferred}",
+        )
+        return result
+
+    def download(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        chunk_size: int,
+        owner: str | None = None,
+    ) -> Iterator[bytes]:
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+
+        def stream() -> Iterator[bytes]:
+            with self._sessions_lock:
+                session.active_operations += 1
+            transferred = 0
+            try:
+                for chunk in session.transport.download(remote_path, chunk_size=chunk_size):
+                    if not isinstance(chunk, bytes):
+                        raise ServiceError("internal", "transport returned a non-bytes download chunk")
+                    transferred += len(chunk)
+                    yield chunk
+            except TransportError as exc:
+                raise ServiceError("remote_error", str(exc), retryable=True) from exc
+            finally:
+                with self._sessions_lock:
+                    session.active_operations -= 1
+                    session.last_used_at = time.time()
+                self._emit(
+                    "file_downloaded",
+                    session_id=session_id,
+                    host_id=session.host_id,
+                    owner_agent_id=owner,
+                    command=remote_path,
+                    outcome=f"bytes={transferred}",
+                )
+
+        return stream()
+
+    def write_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        content: str,
+        *,
+        mode: int = 0o644,
+        max_bytes: int = 4 * 1024 * 1024,
+        owner: str | None = None,
+    ) -> TransferResult:
+        data = content.encode("utf-8")
+        if len(data) > max_bytes:
+            raise ServiceError("resource_limit", f"text content exceeds {max_bytes} bytes")
+        return self.upload(
+            session_id,
+            iter([data]),
+            remote_path,
+            mode=mode,
+            expected_size=len(data),
+            expected_sha256=hashlib.sha256(data).hexdigest(),
+            owner=owner,
+        )
+
+    def read_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        max_bytes: int = 4 * 1024 * 1024,
+        owner: str | None = None,
+    ) -> str:
+        data = bytearray()
+        for chunk in self.download(session_id, remote_path, chunk_size=256 * 1024, owner=owner):
+            data.extend(chunk)
+            if len(data) > max_bytes:
+                raise ServiceError("resource_limit", f"text file exceeds {max_bytes} bytes")
+        try:
+            return bytes(data).decode("utf-8")
+        except UnicodeDecodeError as exc:
+            raise ServiceError("invalid_request", f"remote file is not valid UTF-8 at byte {exc.start}") from exc
+
+    def shell_write(self, session_id: str, data: bytes, *, owner: str | None = None) -> int:
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+        try:
+            session.transport.shell_write(data)
+        except TransportError as exc:
+            raise ServiceError("remote_error", str(exc), retryable=True) from exc
+        session.last_used_at = time.time()
+        return len(data)
+
+    def shell_read(
+        self,
+        session_id: str,
+        *,
+        timeout: float,
+        max_bytes: int,
+        owner: str | None = None,
+    ) -> ShellReadResult:
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+        try:
+            result = session.transport.shell_read(timeout=timeout, max_bytes=max_bytes)
+        except TransportError as exc:
+            raise ServiceError("remote_error", str(exc), retryable=True) from exc
+        session.last_used_at = time.time()
+        return result
+
+    def _get_session(self, session_id: str) -> ServiceSession:
+        with self._sessions_lock:
+            session = self._sessions.get(session_id)
+        if session is None:
+            raise ServiceError("not_found", f"unknown session {session_id}")
+        return session
+
+    def _get_task(self, session_id: str, task_id: str, owner: str | None) -> ServiceTask:
+        session = self._get_session(session_id)
+        self._check_owner(session, owner)
+        with self._sessions_lock:
+            task = self._tasks.get(task_id)
+        if task is None or task.session_id != session_id:
+            raise ServiceError("not_found", f"unknown task {task_id}")
+        if task.owner is not None and task.owner != owner:
+            raise ServiceError("denied", f"task {task_id} is owned by another agent")
+        return task
+
+    @staticmethod
+    def _check_owner(session: ServiceSession, owner: str | None) -> None:
+        if session.owner is not None and session.owner != owner:
+            raise ServiceError("denied", f"session {session.session_id} is owned by another agent")
+
+    def _emit(self, event: str, **fields: object) -> None:
+        if self._audit is not None:
+            audit_event(self._audit, event=event, **{key: value for key, value in fields.items() if value is not None})  # type: ignore[arg-type]

package/src/server_control_mcp/transports/__init__.py

@@ -0,0 +1,17 @@
+"""Remote transport implementations used by HostBridge daemon services."""
+
+from .base import ExecResult, ShellReadResult, TransferResult, Transport, TransportBusy, TransportCapabilities
+from .pty import PtyTransport
+from .ssh import SshConnectionConfig, SshTransport
+
+__all__ = [
+    "ExecResult",
+    "PtyTransport",
+    "SshConnectionConfig",
+    "SshTransport",
+    "ShellReadResult",
+    "TransferResult",
+    "Transport",
+    "TransportBusy",
+    "TransportCapabilities",
+]

package/src/server_control_mcp/transports/base.py

@@ -0,0 +1,78 @@
+"""Shared transport contracts and result models."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass
+from typing import Protocol
+
+
+class TransportError(RuntimeError):
+    """Base error raised by a remote transport."""
+
+
+class TransportBusy(TransportError):
+    """Raised when an ordered transport already has an active operation."""
+
+
+@dataclass(frozen=True, slots=True)
+class TransportCapabilities:
+    separate_stderr: bool
+    native_binary_transfer: bool
+    parallel_channels: int
+
+
+@dataclass(frozen=True, slots=True)
+class ExecResult:
+    exit_code: int | None
+    stdout: bytes
+    stderr: bytes = b""
+    timed_out: bool = False
+
+
+@dataclass(frozen=True, slots=True)
+class TransferResult:
+    bytes_transferred: int
+    sha256: str
+
+
+@dataclass(frozen=True, slots=True)
+class ShellReadResult:
+    data: bytes
+    alive: bool
+
+
+class Transport(Protocol):
+    capabilities: TransportCapabilities
+
+    def exec(
+        self,
+        command: str,
+        *,
+        stdin: Iterable[bytes] | None,
+        timeout: float,
+        output_limit: int,
+    ) -> ExecResult:
+        raise NotImplementedError
+
+    def upload(
+        self,
+        chunks: Iterable[bytes],
+        remote_path: str,
+        *,
+        mode: int,
+        expected_size: int | None,
+    ) -> TransferResult:
+        raise NotImplementedError
+
+    def download(self, remote_path: str, *, chunk_size: int) -> Iterator[bytes]:
+        raise NotImplementedError
+
+    def close(self) -> None:
+        raise NotImplementedError
+
+    def shell_write(self, data: bytes) -> None:
+        raise NotImplementedError
+
+    def shell_read(self, *, timeout: float, max_bytes: int) -> ShellReadResult:
+        raise NotImplementedError