@neoline/hostbridge 2.0.2 → 2.0.4

package/src/server_control_mcp/services.py

@@ -3,19 +3,28 @@
 from __future__ import annotations
 
 import hashlib
+import logging
 import re
-import shlex
 import threading
 import time
 import uuid
 from collections.abc import Callable, Iterable, Iterator
+from contextlib import contextmanager, suppress
 from dataclasses import dataclass, field
 
 from .hosts import HostConfig, 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 .transports.base import ExecResult, ShellReadResult, TransferResult, Transport, TransportBusy, TransportError
 
-TransportFactory = Callable[[HostConfig], Transport]
+TransportFactory = Callable[[HostConfig, float], Transport]
+logger = logging.getLogger(__name__)
+MAX_EXEC_OUTPUT_BYTES = 3 * 1024 * 1024 - 64 * 1024
 
 
 class ServiceError(RuntimeError):
@@ -42,6 +51,7 @@ class ServiceSession:
     created_at: float = field(default_factory=time.time)
     last_used_at: float = field(default_factory=time.time)
     active_operations: int = 0
+    closing: bool = False
 
     def describe(self) -> dict[str, object]:
         return {
@@ -85,31 +95,69 @@ class HostBridgeServices:
         hosts: HostRegistry,
         *,
         transport_factory: TransportFactory,
+        default_connect_timeout: float = 30,
         policy: PolicyConfig | None = None,
         audit_sink: AuditSink | None = None,
+        max_tasks_per_session: int = 32,
     ) -> None:
+        if max_tasks_per_session <= 0:
+            raise ValueError("max_tasks_per_session must be positive")
         self.hosts = hosts
         self._transport_factory = transport_factory
+        self._default_connect_timeout = default_connect_timeout
         self._policy = policy
         self._audit = audit_sink
+        self._max_tasks_per_session = max_tasks_per_session
         self._sessions: dict[str, ServiceSession] = {}
+        self._pending_sessions: dict[str, int] = {}
         self._tasks: dict[str, ServiceTask] = {}
+        self._launching_tasks: dict[str, ServiceTask] = {}
         self._sessions_lock = threading.RLock()
+        self._operation_local = threading.local()
+        self._closing = False
         self._reaper_stop = threading.Event()
         self._reaper_thread: threading.Thread | None = None
 
-    def open_session(self, host_id: str, *, owner: str | None = None) -> ServiceSession:
+    def open_session(
+        self,
+        host_id: str,
+        *,
+        owner: str | None = None,
+        connect_timeout: float | None = None,
+    ) -> ServiceSession:
         try:
             host = self.hosts.get(host_id)
         except KeyError as exc:
             raise ServiceError("not_found", str(exc)) from exc
+        with self._sessions_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
+        effective_timeout = self._default_connect_timeout if connect_timeout is None else connect_timeout
+        if effective_timeout <= 0:
+            self._release_pending_session(host.id)
+            raise ServiceError("invalid_request", "connect_timeout must be positive")
         try:
-            transport = self._transport_factory(host)
-        except Exception as exc:
+            transport = self._transport_factory(host, effective_timeout)
+        except BaseException as exc:
+            self._release_pending_session(host.id)
+            if not isinstance(exc, Exception):
+                raise
             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._release_pending_session_locked(host.id)
+            accepted = not self._closing
+            if accepted:
+                self._sessions[session.session_id] = session
+        if not accepted:
+            with suppress(Exception):
+                transport.close()
+            raise ServiceError("connection_lost", "HostBridge services are shutting down", retryable=True)
         self._emit("session_opened", session_id=session.session_id, host_id=host.id, owner_agent_id=owner, outcome="ok")
         return session
 
@@ -133,25 +181,24 @@ class HostBridgeServices:
             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()
+        with self._operation(session_id, owner) as session:
+            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)
+            stdin = self._bounded_chunks(stdin, limits.max_stdin_bytes, "command stdin")
+            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
+        if result.timed_out:
+            self.close_session(session_id, force=True)
         self._emit(
             "command_run",
             session_id=session.session_id,
@@ -163,38 +210,84 @@ class HostBridgeServices:
         return ServiceExecResult(
             session.session_id,
             result.exit_code,
-            result.stdout,
-            result.stderr,
+            *self._limit_combined_output(result.stdout, result.stderr, output_limit),
             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:
+            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
+            tasks = [
+                task
+                for task in (*self._tasks.values(), *self._launching_tasks.values())
+                if task.session_id == session_id
+            ]
+        try:
+            self._cleanup_tasks(session, tasks)
+        except Exception as exc:
+            if not force:
+                with self._sessions_lock:
+                    session.closing = False
+                raise ServiceError("connection_lost", f"failed to close session {session_id}", retryable=True) from exc
+            logger.exception("failed to clean up tasks while force-closing HostBridge session %s", session_id)
+        try:
+            session.transport.close()
+        except Exception as exc:
             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",
-            )
+                session.closing = False
+            raise ServiceError("connection_lost", f"failed to close session {session_id}", retryable=True) from exc
+        with self._sessions_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=owner,
+            outcome="forced" if force else "user",
+        )
 
-    def close_all(self) -> dict[str, object]:
+    def close_all(self, *, owner: str | None = None, force: bool = False) -> dict[str, object]:
         with self._sessions_lock:
-            sessions = list(self._sessions.values())
-            self._sessions.clear()
-            self._tasks.clear()
+            if force:
+                self._closing = True
+            candidates = [session for session in self._sessions.values() if force or session.owner == owner]
+            busy = [session.session_id for session in candidates if session.active_operations and not force]
+            sessions = [session for session in candidates if force or not session.active_operations]
+        closed: list[str] = []
+        failed: list[str] = []
         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)}
+            try:
+                self.close_session(session.session_id, owner=owner, force=force)
+            except ServiceError:
+                failed.append(session.session_id)
+                logger.exception("failed to close HostBridge session %s", session.session_id)
+            else:
+                closed.append(session.session_id)
+        return {"closed": closed, "count": len(closed), "busy": busy, "failed": failed}
+
+    def _release_pending_session(self, host_id: str) -> None:
+        with self._sessions_lock:
+            self._release_pending_session_locked(host_id)
+
+    def _release_pending_session_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
 
     def reap_idle(self, timeout: float) -> int:
         if timeout <= 0:
@@ -206,17 +299,18 @@ class HostBridgeServices:
                 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
-                }
+        closed = 0
         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)
+            try:
+                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, outcome="expired"
+                )
+        return closed
 
     def start_idle_reaper(self, *, timeout: float, interval: float = 30) -> None:
         if timeout <= 0 or interval <= 0:
@@ -242,30 +336,47 @@ class HostBridgeServices:
     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
+        with self._operation(session_id, owner) as session:
+            task_id = uuid.uuid4().hex
+            job_dir = remote_job_dir(task_id)
+            script = build_remote_task_command(command, task_id)
+            launching_task = ServiceTask(task_id, session_id, session.host_id, command, job_dir, None, owner)
+            with self._sessions_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_task
+            try:
+                result = self.exec(session_id, script, timeout=10, output_limit=4096, owner=owner)
+                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,
+                    job_dir,
+                    int(pid_match.group(1)),
+                    owner,
+                )
+                with self._sessions_lock:
+                    self._launching_tasks.pop(task_id, None)
+                    self._tasks[task_id] = task
+            except BaseException:
+                with self._sessions_lock:
+                    needs_cleanup = task_id in self._launching_tasks and session_id in self._sessions
+                if needs_cleanup:
+                    try:
+                        self._cleanup_tasks(session, [launching_task])
+                    except Exception:
+                        logger.exception("failed to clean up task %s after launch failure", task_id)
+                    else:
+                        with self._sessions_lock:
+                            self._launching_tasks.pop(task_id, None)
+                raise
         self._emit(
             "task_started",
             session_id=session_id,
@@ -285,17 +396,18 @@ class HostBridgeServices:
         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)
+        with self._operation(session_id, owner):
+            task = self._get_task(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 = 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():
@@ -316,16 +428,27 @@ class HostBridgeServices:
         }
 
     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"
+        with self._operation(session_id, owner):
+            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"{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 = 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"
         )
-        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}
 
@@ -340,23 +463,21 @@ class HostBridgeServices:
         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()
+        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")
+            bounded_chunks = self._bounded_chunks(chunks, limit, "upload")
+            assert bounded_chunks is not None
+            try:
+                result = session.transport.upload(
+                    bounded_chunks,
+                    remote_path,
+                    mode=mode,
+                    expected_size=expected_size,
+                )
+            except TransportError as exc:
+                raise ServiceError("remote_error", str(exc), retryable=True) from exc
         if expected_size is not None and result.bytes_transferred != expected_size:
             raise ServiceError(
                 "integrity_error",
@@ -382,33 +503,29 @@ class HostBridgeServices:
         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}",
-                )
+            with self._operation(session_id, owner) as session:
+                limit = self.hosts.get(session.host_id).limits.max_transfer_bytes
+                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)
+                        if transferred > limit:
+                            raise ServiceError("resource_limit", f"download exceeds {limit} bytes")
+                        yield chunk
+                except TransportError as exc:
+                    raise ServiceError("remote_error", str(exc), retryable=True) from exc
+                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}",
+                    )
 
         return stream()
 
@@ -423,8 +540,11 @@ class HostBridgeServices:
         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")
+        session = self._get_session(session_id)
+        self._check_owner(session, 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 self.upload(
             session_id,
             iter([data]),
@@ -446,21 +566,24 @@ class HostBridgeServices:
         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")
+            session = self._get_session(session_id)
+            limit = min(max_bytes, self.hosts.get(session.host_id).limits.max_text_bytes)
+            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
 
     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()
+        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:
+                session.transport.shell_write(data)
+            except TransportError as exc:
+                raise ServiceError("remote_error", str(exc), retryable=True) from exc
         return len(data)
 
     def shell_read(
@@ -471,14 +594,12 @@ class HostBridgeServices:
         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
+        with self._operation(session_id, owner) as session:
+            max_bytes = min(max_bytes, self.hosts.get(session.host_id).limits.max_shell_bytes)
+            try:
+                return session.transport.shell_read(timeout=timeout, max_bytes=max_bytes)
+            except TransportError as exc:
+                raise ServiceError("remote_error", str(exc), retryable=True) from exc
 
     def _get_session(self, session_id: str) -> ServiceSession:
         with self._sessions_lock:
@@ -487,6 +608,76 @@ class HostBridgeServices:
             raise ServiceError("not_found", f"unknown session {session_id}")
         return session
 
+    @staticmethod
+    def _cleanup_tasks(session: ServiceSession, tasks: list[ServiceTask]) -> None:
+        if not tasks:
+            return
+        command = build_remote_tasks_cleanup_command((task.task_id, task.remote_pid) for task in tasks)
+        result = session.transport.exec(command, stdin=None, timeout=15, output_limit=4096)
+        if result.timed_out or result.exit_code != 0:
+            raise TransportError(f"failed to clean up tasks for session {session.session_id}")
+
+    @contextmanager
+    def _operation(self, session_id: str, owner: str | None) -> Iterator[ServiceSession]:
+        leased_sessions = getattr(self._operation_local, "session_ids", None)
+        if leased_sessions is None:
+            leased_sessions = set()
+            self._operation_local.session_ids = leased_sessions
+        if session_id in leased_sessions:
+            leased_session = self._get_session(session_id)
+            self._check_owner(leased_session, owner)
+            yield leased_session
+            return
+        with self._sessions_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.transport.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
+            leased_sessions.add(session_id)
+        try:
+            yield session
+        finally:
+            with self._sessions_lock:
+                session.active_operations -= 1
+                session.last_used_at = time.time()
+                leased_sessions.remove(session_id)
+
+    @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""
+
+    @staticmethod
+    def _bounded_chunks(chunks: Iterable[bytes] | None, limit: int, label: str) -> Iterator[bytes] | None:
+        if chunks is None:
+            return None
+
+        def bounded() -> Iterator[bytes]:
+            transferred = 0
+            for chunk in chunks:
+                if not isinstance(chunk, bytes):
+                    raise TypeError(f"{label} chunks must be bytes")
+                transferred += len(chunk)
+                if transferred > limit:
+                    raise ServiceError("resource_limit", f"{label} exceeds {limit} bytes")
+                yield chunk
+
+        return bounded()
+
     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)
@@ -505,4 +696,11 @@ class HostBridgeServices:
 
     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]
+            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)