@neoline/hostbridge 2.0.3 → 2.0.5

package/src/server_control_mcp/services.py

@@ -1,21 +1,10 @@
-"""Daemon-owned HostBridge session and execution services."""
+"""Shared service errors and result models."""
 
 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 dataclasses import dataclass
 
-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]
+MAX_EXEC_OUTPUT_BYTES = 3 * 1024 * 1024 - 64 * 1024
 
 
 class ServiceError(RuntimeError):
@@ -33,32 +22,6 @@ class ServiceError(RuntimeError):
         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
@@ -77,432 +40,3 @@ class ServiceTask:
     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

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

package/src/server_control_mcp/transports/base.py

@@ -2,17 +2,40 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable, Iterator
+import asyncio
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Awaitable, Iterable, Iterator, Mapping
 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."""
+ByteStream = AsyncIterable[bytes] | Iterable[bytes]
+
+
+async def iterate_byte_stream(chunks: ByteStream, *, label: str) -> AsyncGenerator[bytes, None]:
+    if isinstance(chunks, AsyncIterable):
+        async for chunk in chunks:
+            if not isinstance(chunk, bytes):
+                raise TypeError(f"{label} chunks must be bytes")
+            yield chunk
+        return
+    iterator = iter(chunks)
+    while True:
+        present, chunk = await asyncio.to_thread(_next_chunk, iterator)
+        if not present:
+            return
+        if not isinstance(chunk, bytes):
+            raise TypeError(f"{label} chunks must be bytes")
+        yield chunk
+
+
+def _next_chunk(iterator: Iterator[bytes]) -> tuple[bool, bytes]:
+    try:
+        return True, next(iterator)
+    except StopIteration:
+        return False, b""
 
 
 @dataclass(frozen=True, slots=True)
@@ -42,37 +65,36 @@ class ShellReadResult:
     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
+@dataclass(frozen=True, slots=True)
+class ShellTerminalResult:
+    exit_code: int | None
+    signal: str | None = None
+
+    def __post_init__(self) -> None:
+        if self.exit_code is not None and (
+            not isinstance(self.exit_code, int) or isinstance(self.exit_code, bool) or not 0 <= self.exit_code <= 255
+        ):
+            raise ValueError("shell exit_code must be between 0 and 255 or null")
+        if self.signal is not None and (not isinstance(self.signal, str) or not self.signal.strip()):
+            raise ValueError("shell signal must be a non-empty string or null")
+        if self.exit_code is not None and self.signal is not None:
+            raise ValueError("shell terminal result cannot contain both exit_code and signal")
+
+    @classmethod
+    def from_mapping(cls, value: Mapping[str, object]) -> ShellTerminalResult:
+        return cls(value.get("exit_code"), value.get("signal"))  # type: ignore[arg-type]
+
+    def to_mapping(self) -> dict[str, object]:
+        return {"exit_code": self.exit_code, "signal": self.signal}
+
+
+class ShellAttachment:
+    def __init__(self, output: AsyncIterator[bytes], terminal: Awaitable[ShellTerminalResult]) -> None:
+        self._output = output
+        self._terminal = terminal
+
+    def __aiter__(self) -> AsyncIterator[bytes]:
+        return self._output
+
+    async def finish(self) -> ShellTerminalResult:
+        return await self._terminal