@neoline/hostbridge 2.0.4 → 2.0.5

package/src/server_control_mcp/mock_ssh_session.py

@@ -56,7 +56,7 @@ class HostBridgeSessionFactory:
         last_error: Exception | None = None
         for _ in range(2):
             try:
-                self._client.close_session(session_id, owner="mock-ssh")
+                self._client.release_mock_ssh_session(session_id)
                 return
             except HostBridgeError as exc:
                 if exc.code == "not_found":
@@ -64,5 +64,6 @@ class HostBridgeSessionFactory:
                 last_error = exc
             except Exception as exc:
                 last_error = exc
-        assert last_error is not None
+        if last_error is None:
+            raise RuntimeError("backend close failed without an error")
         raise last_error

package/src/server_control_mcp/mock_ssh_sftp.py

@@ -32,6 +32,8 @@ class HostBridgeSFTPHandle:
     atime: int | None = None
     mtime: int | None = None
     closed: bool = False
+    reserved_bytes: int = 0
+    slot_reserved: bool = False
 
 
 @dataclass(frozen=True, slots=True)
@@ -49,7 +51,15 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         session_factory: HostBridgeSessionFactory,
         timeout: int,
         max_bytes: int,
+        max_open_handles: int,
+        max_temp_bytes: int,
     ) -> None:
+        if not isinstance(max_bytes, int) or isinstance(max_bytes, bool) or max_bytes <= 0:
+            raise ValueError("max_bytes must be a positive integer")
+        if not isinstance(max_open_handles, int) or isinstance(max_open_handles, bool) or max_open_handles <= 0:
+            raise ValueError("max_open_handles must be a positive integer")
+        if not isinstance(max_temp_bytes, int) or isinstance(max_temp_bytes, bool) or max_temp_bytes <= 0:
+            raise ValueError("max_temp_bytes must be a positive integer")
         super().__init__(chan)
         self._client = client
         self._session_factory = session_factory
@@ -58,6 +68,11 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         self._operation_lock = asyncio.Lock()
         self._timeout = timeout
         self._max_bytes = max_bytes
+        self._max_open_handles = max_open_handles
+        self._max_temp_bytes = max_temp_bytes
+        self._resource_lock = threading.Lock()
+        self._open_handle_count = 0
+        self._retained_bytes = 0
         self._handles: list[HostBridgeSFTPHandle] = []
 
     async def _run_blocking(self, function, /, *args, **kwargs):
@@ -92,6 +107,9 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         return await self._run_blocking(self._open_sync, path, pflags, attrs)
 
     def _open_sync(self, path: bytes, pflags: int, attrs: asyncssh.SFTPAttrs) -> HostBridgeSFTPHandle:
+        self._reserve_handle_slot()
+        registered = False
+        temporary_path: Path | None = None
         remote_path = self._path(path)
         readable = bool(pflags & asyncssh.FXF_READ)
         writable = bool(pflags & asyncssh.FXF_WRITE)
@@ -99,26 +117,35 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         truncate = bool(pflags & asyncssh.FXF_TRUNC)
         exclusive = bool(pflags & asyncssh.FXF_EXCL)
         append = bool(pflags & asyncssh.FXF_APPEND)
-        if not readable and not writable:
-            raise asyncssh.SFTPFailure("file must be opened for reading or writing")
-        if (truncate or append) and not writable:
-            raise asyncssh.SFTPPermissionDenied("write flags require write access")
-        if create:
-            create_mode = 0o644 if attrs.permissions is None else attrs.permissions
-            self._create_remote(remote_path, mode=create_mode, exclusive=exclusive)
-        remote_attrs = self._stat(remote_path, follow=True)
-        if remote_attrs is not None and remote_attrs.type == asyncssh.FILEXFER_TYPE_DIRECTORY:
-            raise asyncssh.SFTPFileIsADirectory(remote_path)
-        mode = attrs.permissions if attrs.permissions is not None else 0o644
-        if remote_attrs is not None and remote_attrs.permissions is not None and attrs.permissions is None:
-            mode = remote_attrs.permissions
-        temporary_path = self._new_temporary_path()
         try:
+            if not readable and not writable:
+                raise asyncssh.SFTPFailure("file must be opened for reading or writing")
+            if (truncate or append) and not writable:
+                raise asyncssh.SFTPPermissionDenied("write flags require write access")
+            if create:
+                create_mode = 0o644 if attrs.permissions is None else attrs.permissions
+                self._create_remote(remote_path, mode=create_mode, exclusive=exclusive)
+            remote_attrs = self._stat(remote_path, follow=True)
+            if remote_attrs is not None and remote_attrs.type == asyncssh.FILEXFER_TYPE_DIRECTORY:
+                raise asyncssh.SFTPFileIsADirectory(remote_path)
+            mode = attrs.permissions if attrs.permissions is not None else 0o644
+            if remote_attrs is not None and remote_attrs.permissions is not None and attrs.permissions is None:
+                mode = remote_attrs.permissions
+            temporary_path = self._new_temporary_path()
             if truncate:
                 self._truncate_remote(remote_path)
                 temporary_path.touch(mode=0o600)
             else:
-                self._download_to(remote_path, temporary_path)
+                available = self._available_temp_bytes()
+                declared_size = 0 if remote_attrs.size is None else int(remote_attrs.size)
+                if declared_size < 0 or declared_size > available:
+                    raise asyncssh.SFTPFailure(
+                        f"mock-ssh SFTP temporary storage limit exceeded ({self._max_temp_bytes} bytes)"
+                    )
+                if available == 0:
+                    temporary_path.touch(mode=0o600)
+                else:
+                    self._download_to(remote_path, temporary_path, max_bytes=available)
             size = temporary_path.stat().st_size
             handle = HostBridgeSFTPHandle(
                 remote_path,
@@ -129,11 +156,16 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
                 mode=mode,
                 size=size,
                 dirty=False,
+                slot_reserved=True,
             )
-            self._handles.append(handle)
+            self._register_handle(handle)
+            registered = True
             return handle
         except BaseException:
-            temporary_path.unlink(missing_ok=True)
+            if temporary_path is not None:
+                temporary_path.unlink(missing_ok=True)
+            if not registered:
+                self._release_handle_slot()
             raise
 
     async def read(self, file_obj: object, offset: int, size: int) -> bytes:
@@ -163,9 +195,15 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         final_size = max(handle.size, write_offset + len(payload))
         if final_size > self._max_bytes:
             raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
-        with handle.temporary_path.open("r+b") as stream:
-            stream.seek(write_offset)
-            stream.write(payload)
+        growth = final_size - handle.size
+        self._reserve_growth(handle, growth)
+        try:
+            with handle.temporary_path.open("r+b") as stream:
+                stream.seek(write_offset)
+                stream.write(payload)
+        except BaseException:
+            self._release_growth(handle, growth)
+            raise
         handle.size = final_size
         handle.dirty = True
         return len(payload)
@@ -185,8 +223,7 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         finally:
             handle.closed = True
             handle.temporary_path.unlink(missing_ok=True)
-            with suppress(ValueError):
-                self._handles.remove(handle)
+            self._release_handle(handle)
 
     async def stat(self, path: bytes) -> asyncssh.SFTPAttrs:
         return await self._run_blocking(self._stat, self._path(path), follow=True)
@@ -218,8 +255,16 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
         if attrs.size is not None:
             if attrs.size < 0 or attrs.size > self._max_bytes:
                 raise asyncssh.SFTPFailure(f"file exceeds mock-ssh limit ({self._max_bytes} bytes)")
-            with handle.temporary_path.open("r+b") as stream:
-                stream.truncate(attrs.size)
+            growth = max(0, attrs.size - handle.size)
+            shrink = max(0, handle.size - attrs.size)
+            self._reserve_growth(handle, growth)
+            try:
+                with handle.temporary_path.open("r+b") as stream:
+                    stream.truncate(attrs.size)
+            except BaseException:
+                self._release_growth(handle, growth)
+                raise
+            self._release_growth(handle, shrink)
             handle.size = attrs.size
             handle.dirty = True
         handle.atime = attrs.atime if attrs.atime is not None else handle.atime
@@ -291,7 +336,7 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
             mtime = int(attrs.mtime if attrs.mtime is not None else current.mtime or time.time())
             self._checked_run(f"python3 -c {shlex.quote(script)} {shlex.quote(remote_path)} {atime} {mtime}")
 
-    def _download_to(self, remote_path: str, temporary_path: Path) -> None:
+    def _download_to(self, remote_path: str, temporary_path: Path, *, max_bytes: int) -> None:
         try:
             result = self._client.download_file(
                 self._backend_session_id(),
@@ -299,12 +344,12 @@ class HostBridgeSFTPServer(asyncssh.SFTPServer):
                 temporary_path,
                 timeout=self._timeout,
                 owner="mock-ssh",
-                max_bytes=self._max_bytes,
+                max_bytes=min(self._max_bytes, max_bytes),
             )
-            if result.bytes_transferred > self._max_bytes:
+            if result.bytes_transferred > min(self._max_bytes, max_bytes):
                 raise HostBridgeError(
                     "resource_limit",
-                    f"file exceeds mock-ssh limit ({self._max_bytes} bytes)",
+                    f"file exceeds mock-ssh SFTP temporary storage limit ({max_bytes} bytes)",
                 )
         except HostBridgeError as exc:
             raise self._sftp_error(exc) from exc
@@ -474,6 +519,59 @@ print(json.dumps(items))
                 self._session_id = self._session_factory.open()
             return self._session_id
 
+    def _reserve_handle_slot(self) -> None:
+        with self._resource_lock:
+            if self._open_handle_count >= self._max_open_handles:
+                raise asyncssh.SFTPFailure(f"mock-ssh SFTP open handle limit reached ({self._max_open_handles})")
+            self._open_handle_count += 1
+
+    def _release_handle_slot(self) -> None:
+        with self._resource_lock:
+            self._open_handle_count -= 1
+
+    def _available_temp_bytes(self) -> int:
+        with self._resource_lock:
+            return self._max_temp_bytes - self._retained_bytes
+
+    def _register_handle(self, handle: HostBridgeSFTPHandle) -> None:
+        with self._resource_lock:
+            if self._retained_bytes + handle.size > self._max_temp_bytes:
+                raise asyncssh.SFTPFailure(
+                    f"mock-ssh SFTP temporary storage limit exceeded ({self._max_temp_bytes} bytes)"
+                )
+            handle.reserved_bytes = handle.size
+            self._retained_bytes += handle.size
+            self._handles.append(handle)
+
+    def _reserve_growth(self, handle: HostBridgeSFTPHandle, growth: int) -> None:
+        if growth <= 0:
+            return
+        with self._resource_lock:
+            if self._retained_bytes + growth > self._max_temp_bytes:
+                raise asyncssh.SFTPFailure(
+                    f"mock-ssh SFTP temporary storage limit exceeded ({self._max_temp_bytes} bytes)"
+                )
+            self._retained_bytes += growth
+            handle.reserved_bytes += growth
+
+    def _release_growth(self, handle: HostBridgeSFTPHandle, released: int) -> None:
+        if released <= 0:
+            return
+        with self._resource_lock:
+            actual = min(released, handle.reserved_bytes)
+            handle.reserved_bytes -= actual
+            self._retained_bytes -= actual
+
+    def _release_handle(self, handle: HostBridgeSFTPHandle) -> None:
+        with self._resource_lock:
+            with suppress(ValueError):
+                self._handles.remove(handle)
+            self._retained_bytes -= handle.reserved_bytes
+            handle.reserved_bytes = 0
+            if handle.slot_reserved:
+                self._open_handle_count -= 1
+                handle.slot_reserved = False
+
     @staticmethod
     def _path(path: bytes) -> str:
         text = path.decode("utf-8", "surrogateescape")