@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/runtime_services.py

@@ -0,0 +1,862 @@
+"""Async session services backed by daemon-owned host runtime leases."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import logging
+import re
+import time
+import uuid
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager, suppress
+from dataclasses import dataclass, field
+from typing import Protocol
+
+from .hosts import HostRegistry
+from .policy import AuditSink, PolicyConfig, audit_event
+from .remote_task_agent import (
+    build_remote_task_command,
+    build_remote_tasks_cleanup_command,
+    remote_job_dir,
+    remote_job_dir_assignment,
+)
+from .services import MAX_EXEC_OUTPUT_BYTES, ServiceError, ServiceExecResult, ServiceTask
+from .transports.base import (
+    ByteStream,
+    ExecResult,
+    ShellAttachment,
+    ShellReadResult,
+    ShellTerminalResult,
+    TransferResult,
+    iterate_byte_stream,
+)
+from .tunnel_manager import HostRuntimeLease, TunnelError
+
+logger = logging.getLogger(__name__)
+
+
+class RuntimeLeaseManager(Protocol):
+    async def acquire(self, host_id: str) -> HostRuntimeLease: ...
+
+
+@dataclass(slots=True)
+class RuntimeServiceSession:
+    session_id: str
+    host_id: str
+    owner: str | None
+    lease: HostRuntimeLease = field(repr=False)
+    created_at: float = field(default_factory=time.time)
+    last_used_at: float = field(default_factory=time.time)
+    active_operations: int = 0
+    closing: bool = False
+    operation_tasks: set[asyncio.Task[object]] = field(default_factory=set, repr=False)
+
+    def describe(self) -> dict[str, object]:
+        capabilities = self.lease.capabilities
+        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": capabilities.separate_stderr,
+                "native_binary_transfer": capabilities.native_binary_transfer,
+                "parallel_channels": capabilities.parallel_channels,
+            },
+        }
+
+
+class AsyncHostBridgeServices:
+    def __init__(
+        self,
+        hosts: HostRegistry,
+        runtime_manager: RuntimeLeaseManager,
+        *,
+        default_connect_timeout: float = 30,
+        policy: PolicyConfig | None = None,
+        max_tasks_per_session: int = 32,
+        cleanup_timeout_seconds: float = 5,
+        idle_timeout_seconds: float | None = None,
+        reaper_interval_seconds: float = 60,
+        audit_sink: AuditSink | None = None,
+    ) -> None:
+        if default_connect_timeout <= 0:
+            raise ValueError("default_connect_timeout must be positive")
+        if max_tasks_per_session <= 0:
+            raise ValueError("max_tasks_per_session must be positive")
+        if cleanup_timeout_seconds <= 0:
+            raise ValueError("cleanup_timeout_seconds must be positive")
+        if idle_timeout_seconds is not None and idle_timeout_seconds < 0:
+            raise ValueError("idle_timeout_seconds must be non-negative")
+        if reaper_interval_seconds <= 0:
+            raise ValueError("reaper_interval_seconds must be positive")
+        self.hosts = hosts
+        self._runtime_manager = runtime_manager
+        self._default_connect_timeout = float(default_connect_timeout)
+        self._policy = policy
+        self._max_tasks_per_session = max_tasks_per_session
+        self._cleanup_timeout_seconds = float(cleanup_timeout_seconds)
+        self._idle_timeout_seconds = float(
+            policy.idle_timeout_seconds
+            if idle_timeout_seconds is None and policy is not None
+            else idle_timeout_seconds or 0
+        )
+        self._reaper_interval_seconds = float(reaper_interval_seconds)
+        self._audit = audit_sink
+        self._sessions: dict[str, RuntimeServiceSession] = {}
+        self._pending_sessions: dict[str, int] = {}
+        self._tasks: dict[str, ServiceTask] = {}
+        self._launching_tasks: dict[str, ServiceTask] = {}
+        self._lock = asyncio.Lock()
+        self._closing = False
+        self._reaper_task: asyncio.Task[None] | None = None
+
+    async def open_session(
+        self,
+        host_id: str,
+        *,
+        owner: str | None = None,
+        connect_timeout: float | None = None,
+    ) -> RuntimeServiceSession:
+        try:
+            host = self.hosts.get(host_id)
+        except KeyError as exc:
+            raise ServiceError("not_found", str(exc)) from exc
+        timeout = self._default_connect_timeout if connect_timeout is None else connect_timeout
+        if timeout <= 0:
+            raise ServiceError("invalid_request", "connect_timeout must be positive")
+        async with self._lock:
+            if self._closing:
+                raise ServiceError("connection_lost", "HostBridge services are shutting down", retryable=True)
+            current = sum(session.host_id == host.id for session in self._sessions.values())
+            pending = self._pending_sessions.get(host.id, 0)
+            if current + pending >= host.limits.max_sessions:
+                raise ServiceError("resource_limit", f"host {host_id} reached its session limit")
+            self._pending_sessions[host.id] = pending + 1
+        lease: HostRuntimeLease | None = None
+        registered = False
+        try:
+            async with asyncio.timeout(timeout):
+                lease = await self._runtime_manager.acquire(host.id)
+        except TimeoutError as exc:
+            raise ServiceError("timeout", f"failed to connect to host {host_id}", retryable=True) from exc
+        except TunnelError as exc:
+            raise ServiceError(exc.code, str(exc), retryable=exc.retryable) from exc
+        except BaseException as exc:
+            if isinstance(exc, asyncio.CancelledError):
+                raise
+            raise ServiceError("connection_lost", f"failed to connect to host {host_id}", retryable=True) from exc
+        else:
+            session = RuntimeServiceSession(uuid.uuid4().hex, host.id, owner, lease)
+            try:
+                async with self._lock:
+                    if self._closing:
+                        raise ServiceError(
+                            "connection_lost",
+                            "HostBridge services are shutting down",
+                            retryable=True,
+                        )
+                    self._sessions[session.session_id] = session
+                    self._ensure_idle_reaper_locked()
+                    registered = True
+                self._emit(
+                    "session_opened",
+                    session_id=session.session_id,
+                    host_id=host.id,
+                    owner_agent_id=owner,
+                    outcome="ok",
+                )
+                return session
+            except BaseException:
+                if not registered:
+                    await self._close_lease_quietly(lease)
+                raise
+        finally:
+            async with self._lock:
+                self._release_pending_locked(host.id)
+
+    async def list_sessions(self) -> list[dict[str, object]]:
+        async with self._lock:
+            return [session.describe() for session in self._sessions.values()]
+
+    async def close_session(
+        self,
+        session_id: str,
+        *,
+        owner: str | None = None,
+        force: bool = False,
+    ) -> None:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if session is None:
+                raise ServiceError("not_found", f"unknown session {session_id}")
+            if not force:
+                self._check_owner(session, owner)
+                if session.active_operations:
+                    raise ServiceError("busy", f"session {session_id} has active operations", retryable=True)
+            if session.closing:
+                raise ServiceError("busy", f"session {session_id} is already closing", retryable=True)
+            session.closing = True
+            operations = tuple(task for task in session.operation_tasks if task is not asyncio.current_task())
+            tasks = tuple(
+                task
+                for task in (*self._tasks.values(), *self._launching_tasks.values())
+                if task.session_id == session_id
+            )
+        finalizer = asyncio.create_task(
+            self._finalize_session(session, force=force, operations=operations, tasks=tasks),
+            name=f"hostbridge-session-finalize-{session_id}",
+        )
+        cancellation_requested = False
+        while True:
+            try:
+                errors = await asyncio.shield(finalizer)
+                break
+            except asyncio.CancelledError:
+                cancellation_requested = True
+                if finalizer.done():
+                    errors = finalizer.result()
+                    break
+        if cancellation_requested:
+            cancelled = asyncio.CancelledError(f"session {session_id} close was cancelled after cleanup started")
+            for error in errors:
+                cancelled.add_note(f"session cleanup also failed: {error!r}")
+            raise cancelled
+        if errors:
+            raise ServiceError(
+                "connection_lost",
+                f"failed to fully clean up session {session_id}: {errors[0]}",
+                retryable=True,
+            ) from errors[0]
+
+    async def close_mock_ssh_session(self, session_id: str) -> None:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if session is None:
+                raise ServiceError("not_found", f"unknown session {session_id}")
+            if session.owner != "mock-ssh":
+                raise ServiceError("denied", f"session {session_id} is not owned by mock-ssh")
+        await self.close_session(session_id, force=True)
+
+    async def _finalize_session(
+        self,
+        session: RuntimeServiceSession,
+        *,
+        force: bool,
+        operations: tuple[asyncio.Task[object], ...],
+        tasks: tuple[ServiceTask, ...],
+    ) -> list[BaseException]:
+        session_id = session.session_id
+        errors: list[BaseException] = []
+        if force and operations:
+            for operation in operations:
+                operation.cancel()
+            _, pending = await asyncio.wait(operations, timeout=self._cleanup_timeout_seconds)
+            if pending:
+                errors.append(TimeoutError(f"session {session_id} operations did not stop before cleanup deadline"))
+        if tasks:
+            try:
+                await self._cleanup_remote_tasks(session, tasks)
+            except BaseException as exc:
+                errors.append(exc)
+        try:
+            await session.lease.close()
+        except BaseException as exc:
+            errors.append(exc)
+        finally:
+            async with self._lock:
+                self._sessions.pop(session_id, None)
+                self._tasks = {task_id: task for task_id, task in self._tasks.items() if task.session_id != session_id}
+                self._launching_tasks = {
+                    task_id: task for task_id, task in self._launching_tasks.items() if task.session_id != session_id
+                }
+        self._emit(
+            "session_closed",
+            session_id=session_id,
+            host_id=session.host_id,
+            owner_agent_id=session.owner,
+            outcome="cleanup_failed" if errors else "forced" if force else "user",
+        )
+        return errors
+
+    async def close_all(self, *, owner: str | None = None, force: bool = False) -> dict[str, object]:
+        finalizer = asyncio.create_task(
+            self._close_all_owned(owner=owner, force=force),
+            name="hostbridge-services-close-all",
+        )
+        cancellation: asyncio.CancelledError | None = None
+        while not finalizer.done():
+            try:
+                await asyncio.shield(finalizer)
+            except asyncio.CancelledError as exc:
+                cancellation = exc
+        result = finalizer.result()
+        if cancellation is not None:
+            raise cancellation
+        return result
+
+    async def _close_all_owned(self, *, owner: str | None, force: bool) -> dict[str, object]:
+        async with self._lock:
+            if force:
+                self._closing = True
+                reaper = self._reaper_task
+                self._reaper_task = None
+            else:
+                reaper = None
+            candidates = [session for session in self._sessions.values() if force or session.owner == owner]
+        if reaper is not None and reaper is not asyncio.current_task():
+            reaper.cancel()
+            await asyncio.gather(reaper, return_exceptions=True)
+        closed: list[str] = []
+        failed: list[str] = []
+        for session in candidates:
+            try:
+                await self.close_session(session.session_id, owner=owner, force=force)
+            except ServiceError:
+                failed.append(session.session_id)
+            else:
+                closed.append(session.session_id)
+        return {"closed": closed, "count": len(closed), "busy": [], "failed": failed}
+
+    async def exec(
+        self,
+        session_id: str,
+        command: str,
+        *,
+        stdin: ByteStream | 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")
+        self._check_command_policy(command)
+        async with self._operation(session_id, owner) as session:
+            service_result = await self._exec_on_session(
+                session,
+                command,
+                stdin=stdin,
+                timeout=timeout,
+                output_limit=output_limit,
+            )
+        self._emit(
+            "command_run",
+            session_id=session_id,
+            host_id=session.host_id,
+            owner_agent_id=owner,
+            command=command,
+            outcome="timed_out"
+            if service_result.timed_out
+            else f"exit_{service_result.exit_code}"
+            if service_result.exit_code is not None
+            else "unknown",
+        )
+        if service_result.timed_out:
+            await self.close_session(session_id, force=True)
+        return service_result
+
+    async def upload(
+        self,
+        session_id: str,
+        chunks: ByteStream,
+        remote_path: str,
+        *,
+        mode: int,
+        expected_size: int | None,
+        expected_sha256: str | None,
+        owner: str | None = None,
+    ) -> TransferResult:
+        async with self._operation(session_id, owner) as session:
+            limit = self.hosts.get(session.host_id).limits.max_transfer_bytes
+            if expected_size is not None and expected_size > limit:
+                raise ServiceError("resource_limit", f"upload exceeds {limit} bytes")
+            result = await session.lease.upload(
+                self._bounded_chunks(chunks, limit, "upload"),
+                remote_path,
+                mode=mode,
+                expected_size=expected_size,
+                expected_sha256=expected_sha256,
+            )
+        if expected_size is not None and result.bytes_transferred != expected_size:
+            raise ServiceError("integrity_error", "transfer size mismatch")
+        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
+
+    async def download(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        chunk_size: int,
+        owner: str | None = None,
+    ) -> AsyncGenerator[bytes, None]:
+        async with self._operation(session_id, owner) as session:
+            limit = self.hosts.get(session.host_id).limits.max_transfer_bytes
+            transferred = 0
+            try:
+                async for chunk in session.lease.download(remote_path, chunk_size=chunk_size):
+                    if not isinstance(chunk, bytes):
+                        raise ServiceError("internal", "runtime returned a non-bytes download chunk")
+                    transferred += len(chunk)
+                    if transferred > limit:
+                        raise ServiceError("resource_limit", f"download exceeds {limit} bytes")
+                    yield chunk
+            finally:
+                self._emit(
+                    "file_downloaded",
+                    session_id=session_id,
+                    host_id=session.host_id,
+                    owner_agent_id=owner,
+                    command=remote_path,
+                    outcome=f"bytes={transferred}",
+                )
+
+    async 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")
+        session = await self._session_for_request(session_id, owner)
+        limit = min(max_bytes, self.hosts.get(session.host_id).limits.max_text_bytes)
+        if len(data) > limit:
+            raise ServiceError("resource_limit", f"text content exceeds {limit} bytes")
+        return await self.upload(
+            session_id,
+            [data],
+            remote_path,
+            mode=mode,
+            expected_size=len(data),
+            expected_sha256=hashlib.sha256(data).hexdigest(),
+            owner=owner,
+        )
+
+    async def read_text(
+        self,
+        session_id: str,
+        remote_path: str,
+        *,
+        max_bytes: int = 4 * 1024 * 1024,
+        owner: str | None = None,
+    ) -> str:
+        session = await self._session_for_request(session_id, owner)
+        limit = min(max_bytes, self.hosts.get(session.host_id).limits.max_text_bytes)
+        data = bytearray()
+        async for chunk in self.download(session_id, remote_path, chunk_size=256 * 1024, owner=owner):
+            data.extend(chunk)
+            if len(data) > limit:
+                raise ServiceError("resource_limit", f"text file exceeds {limit} 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
+
+    async 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")
+        self._check_command_policy(command)
+        async with self._operation(session_id, owner) as session:
+            task_id = uuid.uuid4().hex
+            launching = ServiceTask(
+                task_id,
+                session_id,
+                session.host_id,
+                command,
+                remote_job_dir(task_id),
+                None,
+                owner,
+            )
+            async with self._lock:
+                retained = sum(
+                    task.session_id == session_id for task in (*self._tasks.values(), *self._launching_tasks.values())
+                )
+                if retained >= self._max_tasks_per_session:
+                    raise ServiceError("resource_limit", f"session {session_id} reached its task retention limit")
+                self._launching_tasks[task_id] = launching
+            try:
+                result = await self._exec_on_session(
+                    session,
+                    build_remote_task_command(command, task_id),
+                    stdin=None,
+                    timeout=10,
+                    output_limit=4096,
+                )
+                pid_match = re.fullmatch(rb"\s*([1-9]\d*)\s*", result.stdout)
+                if result.timed_out or result.exit_code != 0 or pid_match is None:
+                    raise ServiceError("remote_error", f"failed to launch task {task_id}", retryable=True)
+                task = ServiceTask(
+                    task_id,
+                    session_id,
+                    session.host_id,
+                    command,
+                    launching.job_dir,
+                    int(pid_match.group(1)),
+                    owner,
+                )
+                async with self._lock:
+                    self._launching_tasks.pop(task_id, None)
+                    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
+            except BaseException:
+                async with self._lock:
+                    cleanup_here = task_id in self._launching_tasks and not session.closing
+                if cleanup_here:
+                    try:
+                        await self._cleanup_remote_tasks(session, (launching,))
+                    finally:
+                        async with self._lock:
+                            self._launching_tasks.pop(task_id, None)
+                else:
+                    async with self._lock:
+                        self._launching_tasks.pop(task_id, None)
+                raise
+
+    async def task_status(
+        self,
+        session_id: str,
+        task_id: str,
+        *,
+        owner: str | None = None,
+        tail_lines: int = 80,
+    ) -> dict[str, object]:
+        task = await self._task_for_request(session_id, task_id, owner)
+        safe_tail = max(1, min(int(tail_lines), 500))
+        script = (
+            f'{remote_job_dir_assignment(task.task_id)}; 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 = await 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(line.partition("=")[::2] for line in metadata.splitlines() if "=" in line)
+        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,
+        }
+
+    async def cancel_task(self, session_id: str, task_id: str, *, owner: str | None = None) -> dict[str, object]:
+        task = await self._task_for_request(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"{remote_job_dir_assignment(task.task_id)}; "
+            'if [ -f "$job_dir/exit_code" ]; then printf already; '
+            f"elif kill -TERM -- -{task.remote_pid} 2>/dev/null; then "
+            'attempt=0; while kill -0 -- -$(cat "$job_dir/pid") 2>/dev/null && [ "$attempt" -lt 20 ]; do '
+            "attempt=$((attempt + 1)); sleep 0.05; done; "
+            'if kill -0 -- -$(cat "$job_dir/pid") 2>/dev/null; then '
+            'kill -KILL -- -$(cat "$job_dir/pid") 2>/dev/null || true; fi; '
+            'printf \'143\n\' > "$job_dir/.exit_code.tmp"; mv "$job_dir/.exit_code.tmp" "$job_dir/exit_code"; '
+            "printf killed; else printf noop; fi"
+        )
+        result = await 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,
+            host_id=task.host_id,
+            owner_agent_id=owner,
+            outcome=outcome,
+        )
+        return {"task_id": task_id, "stopped": outcome in {"killed", "already"}, "outcome": outcome}
+
+    async def reap_idle(self) -> int:
+        if self._idle_timeout_seconds <= 0:
+            return 0
+        now = time.time()
+        async with self._lock:
+            expired = tuple(
+                session
+                for session in self._sessions.values()
+                if not session.closing
+                and session.active_operations == 0
+                and now - session.last_used_at > self._idle_timeout_seconds
+            )
+        closed = 0
+        for session in expired:
+            try:
+                await self.close_session(session.session_id, force=True)
+            except ServiceError:
+                logger.exception("failed to reap idle HostBridge session %s", session.session_id)
+            else:
+                closed += 1
+                self._emit(
+                    "session_idle_closed",
+                    session_id=session.session_id,
+                    host_id=session.host_id,
+                    owner_agent_id=session.owner,
+                    outcome="expired",
+                )
+        return closed
+
+    async def shell_write(self, session_id: str, data: bytes, *, owner: str | None = None) -> int:
+        if not isinstance(data, bytes):
+            raise ServiceError("invalid_request", "shell data must be bytes")
+        async with self._operation(session_id, owner) as session:
+            limit = self.hosts.get(session.host_id).limits.max_shell_bytes
+            if len(data) > limit:
+                raise ServiceError("resource_limit", f"shell write exceeds {limit} bytes")
+            try:
+                await session.lease.shell_write(data)
+            except TunnelError as exc:
+                raise ServiceError(exc.code, str(exc), retryable=exc.retryable) from exc
+        return len(data)
+
+    async def shell_read(
+        self,
+        session_id: str,
+        *,
+        timeout: float | None,
+        max_bytes: int,
+        owner: str | None = None,
+    ) -> ShellReadResult:
+        if timeout is not None and timeout < 0:
+            raise ServiceError("invalid_request", "timeout must be non-negative")
+        if max_bytes <= 0:
+            raise ServiceError("invalid_request", "max_bytes must be positive")
+        async with self._operation(session_id, owner) as session:
+            limit = min(max_bytes, self.hosts.get(session.host_id).limits.max_shell_bytes)
+            try:
+                return await session.lease.shell_read(timeout=timeout, max_bytes=limit)
+            except TunnelError as exc:
+                raise ServiceError(exc.code, str(exc), retryable=exc.retryable) from exc
+
+    async def attach_shell(
+        self,
+        session_id: str,
+        stdin: ByteStream,
+        *,
+        max_bytes: int,
+        owner: str | None = None,
+    ) -> ShellAttachment:
+        if max_bytes <= 0:
+            raise ServiceError("invalid_request", "max_bytes must be positive")
+        terminal: asyncio.Future[ShellTerminalResult] = asyncio.get_running_loop().create_future()
+
+        async def output() -> AsyncGenerator[bytes, None]:
+            try:
+                async with self._operation(session_id, owner) as session:
+                    limit = min(max_bytes, self.hosts.get(session.host_id).limits.max_shell_bytes)
+                    attachment = await session.lease.attach_shell(stdin, max_bytes=limit)
+                    async for chunk in attachment:
+                        yield chunk
+                    terminal.set_result(await attachment.finish())
+            except TunnelError as exc:
+                error = ServiceError(exc.code, str(exc), retryable=exc.retryable)
+                if not terminal.done():
+                    terminal.set_exception(error)
+                    terminal.exception()
+                raise error from exc
+            except BaseException as exc:
+                if not terminal.done():
+                    terminal.set_exception(exc)
+                    terminal.exception()
+                raise
+
+        return ShellAttachment(output(), terminal)
+
+    @asynccontextmanager
+    async def _operation(self, session_id: str, owner: str | None) -> AsyncGenerator[RuntimeServiceSession, None]:
+        operation = asyncio.current_task()
+        if operation is None:
+            raise RuntimeError("HostBridge operation is not running in an asyncio task")
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if session is None:
+                raise ServiceError("not_found", f"unknown session {session_id}")
+            self._check_owner(session, owner)
+            if session.closing:
+                raise ServiceError("busy", f"session {session_id} is closing", retryable=True)
+            capacity = max(1, session.lease.capabilities.parallel_channels)
+            if session.active_operations >= capacity:
+                raise ServiceError("busy", f"session {session_id} reached its channel capacity", retryable=True)
+            session.active_operations += 1
+            session.operation_tasks.add(operation)
+        try:
+            yield session
+        finally:
+            async with self._lock:
+                session.active_operations -= 1
+                session.operation_tasks.discard(operation)
+                session.last_used_at = time.time()
+
+    async def _exec_on_session(
+        self,
+        session: RuntimeServiceSession,
+        command: str,
+        *,
+        stdin: ByteStream | None,
+        timeout: float,
+        output_limit: int,
+    ) -> ServiceExecResult:
+        limits = self.hosts.get(session.host_id).limits
+        timeout = min(timeout, limits.command_timeout_seconds)
+        output_limit = min(output_limit, limits.max_output_bytes, MAX_EXEC_OUTPUT_BYTES)
+        bounded_stdin = None if stdin is None else self._bounded_chunks(stdin, limits.max_stdin_bytes, "command stdin")
+        try:
+            result: ExecResult = await session.lease.exec(
+                command,
+                stdin=bounded_stdin,
+                timeout=timeout,
+                output_limit=output_limit,
+            )
+        except TunnelError as exc:
+            raise ServiceError(exc.code, str(exc), retryable=exc.retryable) from exc
+        stdout, stderr = self._limit_combined_output(result.stdout, result.stderr, output_limit)
+        return ServiceExecResult(session.session_id, result.exit_code, stdout, stderr, result.timed_out)
+
+    async def _cleanup_remote_tasks(
+        self,
+        session: RuntimeServiceSession,
+        tasks: tuple[ServiceTask, ...],
+    ) -> None:
+        command = build_remote_tasks_cleanup_command((task.task_id, task.remote_pid) for task in tasks)
+        limits = self.hosts.get(session.host_id).limits
+        try:
+            result = await session.lease.exec(
+                command,
+                stdin=None,
+                timeout=min(10, limits.command_timeout_seconds),
+                output_limit=4096,
+            )
+        except TunnelError as exc:
+            raise ServiceError(exc.code, str(exc), retryable=exc.retryable) from exc
+        if result.timed_out or result.exit_code != 0:
+            raise ServiceError("remote_error", "failed to clean up remote tasks", retryable=True)
+
+    def _check_command_policy(self, command: str) -> None:
+        if self._policy is None:
+            return
+        decision = self._policy.evaluate(command)
+        if not decision.allowed:
+            raise ServiceError("denied", decision.reason, details={"matched_pattern": decision.matched_pattern})
+
+    def _ensure_idle_reaper_locked(self) -> None:
+        if self._idle_timeout_seconds <= 0 or self._reaper_task is not None:
+            return
+        self._reaper_task = asyncio.create_task(
+            self._idle_reaper_loop(),
+            name="hostbridge-session-idle-reaper",
+        )
+
+    async def _idle_reaper_loop(self) -> None:
+        while True:
+            await asyncio.sleep(self._reaper_interval_seconds)
+            await self.reap_idle()
+
+    def _emit(self, event: str, **fields: object) -> None:
+        if self._audit is None:
+            return
+        try:
+            audit_event(
+                self._audit,
+                event=event,
+                **{key: value for key, value in fields.items() if value is not None},  # type: ignore[arg-type]
+            )
+        except Exception:
+            logger.exception("audit event could not be written: %s", event)
+
+    @staticmethod
+    async def _close_lease_quietly(lease: HostRuntimeLease) -> None:
+        with suppress(BaseException):
+            await lease.close()
+
+    async def _bounded_chunks(
+        self,
+        chunks: ByteStream,
+        limit: int,
+        label: str,
+    ) -> AsyncGenerator[bytes, None]:
+        transferred = 0
+        async for chunk in iterate_byte_stream(chunks, label=label):
+            transferred += len(chunk)
+            if transferred > limit:
+                raise ServiceError("resource_limit", f"{label} exceeds {limit} bytes")
+            yield chunk
+
+    async def _session_for_request(self, session_id: str, owner: str | None) -> RuntimeServiceSession:
+        async with self._lock:
+            session = self._sessions.get(session_id)
+            if session is None:
+                raise ServiceError("not_found", f"unknown session {session_id}")
+            self._check_owner(session, owner)
+            if session.closing:
+                raise ServiceError("busy", f"session {session_id} is closing", retryable=True)
+            return session
+
+    async def _task_for_request(self, session_id: str, task_id: str, owner: str | None) -> ServiceTask:
+        await self._session_for_request(session_id, owner)
+        async with self._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
+
+    def _release_pending_locked(self, host_id: str) -> None:
+        pending = self._pending_sessions.get(host_id, 0)
+        if pending <= 1:
+            self._pending_sessions.pop(host_id, None)
+        else:
+            self._pending_sessions[host_id] = pending - 1
+
+    @staticmethod
+    def _check_owner(session: RuntimeServiceSession, 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")
+
+    @staticmethod
+    def _limit_combined_output(stdout: bytes, stderr: bytes, limit: int) -> tuple[bytes, bytes]:
+        if len(stdout) + len(stderr) <= limit:
+            return stdout, stderr
+        stdout_keep = min(len(stdout), (limit + 1) // 2)
+        stderr_keep = min(len(stderr), limit - stdout_keep)
+        remaining = limit - stdout_keep - stderr_keep
+        if remaining:
+            extra_stdout = min(len(stdout) - stdout_keep, remaining)
+            stdout_keep += extra_stdout
+            stderr_keep += min(len(stderr) - stderr_keep, remaining - extra_stdout)
+        return stdout[-stdout_keep:] if stdout_keep else b"", stderr[-stderr_keep:] if stderr_keep else b""