fix(firmware): pace + recover multi-MB flash writes that wedge the RPC session

src/flipperzero_mcp/firmware/installer.py

@@ -5,10 +5,12 @@
 import asyncio
 import hashlib
 import logging
-from typing import Any, Protocol
+from collections.abc import Awaitable, Callable
+from typing import Any, Literal, Protocol
 
 from flipperzero_mcp.errors import FlipperTimeoutError
 from flipperzero_mcp.firmware.codes import update_code_message
+from flipperzero_mcp.rpc.protobuf_gen import system_pb2
 
 logger = logging.getLogger(__name__)
 
@@ -17,6 +19,25 @@
 # to SD; retry across this many attempts with a settle delay before trusting it.
 _MD5_ATTEMPTS = 5
 _MD5_SETTLE_S = 2.0
+# storage_md5sum hashes the whole file on-device; for a multi-MB blob that is
+# slow (60-120 s for 11 MB) and wedges the session. Above this size, verify by
+# storage_stat size instead (immediate, no wedge) and let the device's own
+# update validation provide cryptographic integrity on apply.
+_MD5_VERIFY_MAX_BYTES = 2 * 1024 * 1024
+# A multi-MB storage_write transfers the data but can wedge the RPC session, so
+# the post-write digest reads back unreadable. The wedge clears on a transport
+# reconnect; reconnect and re-verify (without rewriting) this many times.
+_RESYNC_ATTEMPTS = 2
+# After reconnecting from a large-write wedge, the device needs a moment to be
+# ready to negotiate a new RPC session (past the negotiation cooldown) before
+# the re-verify; settle this long so the fresh session's first call succeeds.
+_POST_WRITE_SETTLE_S = 8.0
+# system_update issued right after the final large write can return a transient
+# UnspecifiedError while the device is still settling; settle this long before
+# (re)trying the update trigger.
+_UPDATE_SETTLE_S = 3.0
+
+_Md5Status = Literal["match", "mismatch", "unreadable"]
 
 
 class FlashError(RuntimeError):
@@ -28,6 +49,7 @@ async def get_device_info(self) -> dict[str, Any]: ...
     async def storage_mkdir(self, path: str) -> bool: ...
     async def storage_write(self, path: str, content: bytes) -> bool: ...
     async def storage_md5sum(self, path: str) -> str | None: ...
+    async def storage_stat(self, path: str) -> dict[str, Any] | None: ...
     async def system_update(self, manifest_path: str) -> int: ...
     async def system_reboot_update(self) -> None: ...
 
@@ -42,15 +64,51 @@ class _Md5Reader(Protocol):
     async def storage_md5sum(self, path: str) -> str | None: ...
 
 
-async def _verify_md5(
+class _StatReader(Protocol):
+    async def storage_stat(self, path: str) -> dict[str, Any] | None: ...
+
+
+_Resync = Callable[[], Awaitable["_RPCLike"]]
+
+
+async def _check_size(rpc: _StatReader, path: str, expected_len: int) -> _Md5Status:
+    """Verify a written file by its on-device size.
+
+    Used for files too large to md5sum reliably: ``storage_stat`` returns the
+    size immediately without hashing, so it neither stalls nor wedges the
+    session. A size match catches the dominant transfer failure (a truncated
+    push); cryptographic integrity is enforced by the device's update
+    validation when the bundle is applied.
+
+    Returns:
+        ``"match"`` if the on-device size equals ``expected_len``, ``"mismatch"``
+        if it differs, or ``"unreadable"`` if the file cannot be stat'd.
+    """
+    stat = await rpc.storage_stat(path)
+    if stat is None:
+        return "unreadable"
+    return "match" if stat.get("size") == expected_len else "mismatch"
+
+
+async def _verify_written(rpc: _RPCLike, dest: str, data: bytes, expected_md5: str) -> _Md5Status:
+    """Verify a freshly written file, choosing the check by size.
+
+    Small files use md5 (cryptographic); files above ``_MD5_VERIFY_MAX_BYTES``
+    use size, because hashing a multi-MB blob on-device is slow and wedge-prone.
+    """
+    if len(data) > _MD5_VERIFY_MAX_BYTES:
+        return await _check_size(rpc, dest, len(data))
+    return await _check_md5(rpc, dest, expected_md5)
+
+
+async def _check_md5(
     rpc: _Md5Reader, path: str, expected: str, *, settle_s: float = _MD5_SETTLE_S
-) -> bool:
-    """Confirm the device-side md5 matches, retrying while the digest is unreadable.
+) -> _Md5Status:
+    """Classify the device-side md5 against ``expected``.
 
-    ``storage_md5sum`` returns ``None`` when the device is still flushing a large
-    write to SD, so a ``None`` is treated as "not ready yet" and retried after a
-    settle delay. A non-``None`` digest that differs is a definitive mismatch and
-    fails immediately.
+    ``storage_md5sum`` returns ``None`` while the device is still flushing a
+    large write to SD, so a ``None`` is treated as "not ready yet" and retried
+    after a settle delay.
 
     Args:
         rpc: Live RPC client.
@@ -59,27 +117,133 @@ async def _verify_md5(
         settle_s: Delay between retries (0 in tests for determinism).
 
     Returns:
-        True if the device digest matches; False on a definitive mismatch or if
-        the digest stays unreadable across all attempts.
+        ``"match"`` if the digest equals ``expected``, ``"mismatch"`` for a
+        definitive non-equal digest (corrupt transfer or wrong data), or
+        ``"unreadable"`` if it stays ``None`` across all attempts.
     """
     for attempt in range(_MD5_ATTEMPTS):
         device_md5 = await rpc.storage_md5sum(path)
         if device_md5 == expected:
-            return True
+            return "match"
         if device_md5 is not None:
-            return False
+            return "mismatch"
         if attempt + 1 < _MD5_ATTEMPTS:
             await asyncio.sleep(settle_s)
-    return False
+    return "unreadable"
+
+
+async def _push_file(
+    rpc: _RPCLike, dest: str, data: bytes, *, resync: _Resync | None
+) -> _RPCLike:
+    """Write one file and verify it on-device, recovering a wedged session.
+
+    A multi-MB ``storage_write`` transfers the data and acks it, but then wedges
+    the RPC session so the immediate digest reads back ``None`` or garbage. A
+    successful write means the data is durably on the SD card, so on a wedge the
+    file is *not* rewritten (a rewrite just re-wedges and, for an 11 MB blob,
+    wastes minutes): instead ``resync`` reconnects, settles past the negotiation
+    cooldown, and re-verifies the existing data. Only a write that itself fails
+    is re-sent, and only a digest that keeps disagreeing after a clean reconnect
+    is a real mismatch.
+
+    Args:
+        rpc: Live RPC client.
+        dest: Device path to write.
+        data: File contents.
+        resync: Reconnect callback returning a fresh RPC client, or ``None`` to
+            fail closed without recovery.
+
+    Returns:
+        The live RPC client, which ``resync`` may have replaced.
 
+    Raises:
+        FlashError: If the digest never matches after the allotted reconnect
+            retries (a persistent mismatch points at a corrupt transfer; a
+            persistent unreadable digest points at an unrecoverable wedge).
+    """
+    expected = hashlib.md5(data, usedforsecurity=False).hexdigest()
+    written = False
+    status: _Md5Status = "unreadable"
+    for attempt in range(_RESYNC_ATTEMPTS + 1):
+        if not written:
+            try:
+                written = await rpc.storage_write(dest, data)
+            except (OSError, RuntimeError, FlipperTimeoutError):
+                written = False
+        if written:
+            try:
+                status = await _verify_written(rpc, dest, data, expected)
+            except (OSError, RuntimeError, FlipperTimeoutError):
+                status = "unreadable"  # write wedged the session; verify after reconnect
+            if status == "match":
+                return rpc
+            if status == "mismatch":
+                written = False  # on-disk data disagrees; re-push on the next pass
+        if resync is None or attempt == _RESYNC_ATTEMPTS:
+            break
+        rpc = await resync()
+        await asyncio.sleep(_POST_WRITE_SETTLE_S)
+    if status == "mismatch":
+        raise FlashError(
+            f"verification failed after writing {dest}; the on-device data does not "
+            "match the bundle even after reconnecting - corrupt or truncated transfer, "
+            "flash aborted"
+        )
+    raise FlashError(
+        f"device RPC session stopped responding after writing {dest}; the update was "
+        "not applied - reconnect to re-establish the session and retry (a transport "
+        "reconnect clears the wedge; power-cycle only if it persists)"
+    )
+
+
+async def _trigger_update(rpc: _RPCLike, manifest: str, *, resync: _Resync | None) -> _RPCLike:
+    """Stage the update via ``system_update``, retrying a transient error.
+
+    Issued right after the final large write, ``system_update`` can return
+    ``UnspecifiedError`` while the device is still settling - a short settle and
+    a clean session clear it (the CLI ``update install`` succeeds the same way).
+    A specific code (target/manifest/integrity mismatch) is a real rejection and
+    fails immediately.
 
-async def install_bundle(rpc: _RPCLike, bundle: _BundleLike, *, pkg_name: str) -> None:
+    Args:
+        rpc: Live RPC client.
+        manifest: Device path to the staged ``update.fuf``.
+        resync: Reconnect callback for a fresh session between retries, or
+            ``None`` to retry on the same session.
+
+    Returns:
+        The live RPC client to reboot with.
+
+    Raises:
+        FlashError: On a definitive rejection, a dropped link, or an
+            ``UnspecifiedError`` that persists across retries.
+    """
+    for attempt in range(_RESYNC_ATTEMPTS + 1):
+        await asyncio.sleep(_UPDATE_SETTLE_S)
+        try:
+            code = await rpc.system_update(manifest)
+        except FlipperTimeoutError as e:
+            raise FlashError(f"device link dropped during update validation: {e}") from e
+        if code == system_pb2.UpdateResponse.OK:
+            return rpc
+        if code != system_pb2.UpdateResponse.UnspecifiedError or attempt == _RESYNC_ATTEMPTS:
+            raise FlashError(f"device rejected update: {update_code_message(code)}")
+        if resync is not None:
+            rpc = await resync()
+    raise FlashError("device rejected update: unspecified update error (persisted after retries)")
+
+
+async def install_bundle(
+    rpc: _RPCLike, bundle: _BundleLike, *, pkg_name: str, resync: _Resync | None = None
+) -> None:
     """Push ``bundle`` to ``/ext/update/<pkg_name>`` and reboot into the updater.
 
     Args:
         rpc: Live RPC client.
         bundle: Resolved local bundle (manifest + files + target).
         pkg_name: Update subfolder name on the device.
+        resync: Reconnect callback returning a fresh RPC client, used to recover
+            a session wedged by a large write. ``None`` fails closed instead.
 
     Raises:
         FlashError: On target mismatch, push/verify failure, or a non-OK update
@@ -105,19 +269,8 @@ async def install_bundle(rpc: _RPCLike, bundle: _BundleLike, *, pkg_name: str) -
         parent = dest.rsplit("/", 1)[0]
         if parent != pkg_dir:
             await rpc.storage_mkdir(parent)
-        if not await rpc.storage_write(dest, data):
-            raise FlashError(f"failed to write {dest}")
-        expected_md5 = hashlib.md5(data, usedforsecurity=False).hexdigest()
-        if not await _verify_md5(rpc, dest, expected_md5):
-            raise FlashError(f"md5 mismatch after writing {dest}")
+        rpc = await _push_file(rpc, dest, data, resync=resync)
 
     manifest = f"{pkg_dir}/{bundle.manifest_name}"
-    try:
-        code = await rpc.system_update(manifest)
-    except FlipperTimeoutError as e:
-        raise FlashError(f"device link dropped during update validation: {e}") from e
-    message = update_code_message(code)
-    if message is not None:
-        raise FlashError(f"device rejected update: {message}")
-
+    rpc = await _trigger_update(rpc, manifest, resync=resync)
     await rpc.system_reboot_update()

src/flipperzero_mcp/rpc/protobuf_rpc.py

@@ -347,6 +347,8 @@ async def start_session_attempt() -> bool:
     async def _send_rpc_message(
         self,
         main_message: Any,  # flipper_pb2.Main
+        *,
+        response_timeout: float = 2.5,
     ) -> Any | None:  # Optional[flipper_pb2.Main]
         """
         Send a protobuf RPC message and receive response.
@@ -357,6 +359,10 @@ async def _send_rpc_message(
 
         Args:
             main_message: Main protobuf message to send
+            response_timeout: Max seconds to wait for the response frame. The
+                2.5 s default suits quick commands; long device-side work (e.g.
+                hashing a multi-MB file) needs a larger value or the response is
+                missed and the call reads as a failure.
 
         Returns:
             Main response message or None
@@ -374,7 +380,7 @@ async def _send_rpc_message(
             await self.transport.send(message)
 
             # Receive one response Main message
-            return await self._receive_main_message(timeout=2.5)
+            return await self._receive_main_message(timeout=response_timeout)
 
         except Exception:
             logger.debug("_send_rpc_message failed", exc_info=True)
@@ -1070,6 +1076,11 @@ async def _storage_rename_internal(self, old_path: str, new_path: str) -> bool:
             logger.debug("_storage_rename_internal failed", exc_info=True)
             return False
 
+    # The device hashes the whole file before replying, so a multi-MB file can
+    # take far longer than a quick command; an 11 MB blob takes 60-120 s and
+    # varies with SD state, so allow generous headroom.
+    _MD5SUM_RESPONSE_TIMEOUT_S = 180.0
+
     async def storage_md5sum(self, path: str) -> str | None:
         """Compute the device-side MD5 of a file via storage_md5sum_request.
 
@@ -1081,7 +1092,10 @@ async def storage_md5sum(self, path: str) -> str | None:
         """
         async with self._io_lock:
             try:
-                return await asyncio.wait_for(self._storage_md5sum_internal(path), timeout=10.0)
+                return await asyncio.wait_for(
+                    self._storage_md5sum_internal(path),
+                    timeout=self._MD5SUM_RESPONSE_TIMEOUT_S + 5.0,
+                )
             except Exception:
                 logger.debug("storage_md5sum(%s) timed out or failed", path, exc_info=True)
                 return None
@@ -1096,7 +1110,9 @@ async def _storage_md5sum_internal(self, path: str) -> str | None:
             req.path = path
             main_request.storage_md5sum_request.CopyFrom(req)
 
-            main_response = await self._send_rpc_message(main_request)
+            main_response = await self._send_rpc_message(
+                main_request, response_timeout=self._MD5SUM_RESPONSE_TIMEOUT_S
+            )
             if (
                 main_response
                 and main_response.command_status == flipper_pb2.CommandStatus.OK
@@ -1108,6 +1124,11 @@ async def _storage_md5sum_internal(self, path: str) -> str | None:
         return None
 
     _WRITE_CHUNK_SIZE = 1024  # Bytes per RPC write frame; conservative for the Flipper serial link.
+    # The device reassembles has_next fragments and acks only the final one, so
+    # an unpaced multi-MB write floods USB CDC (~1.5 MB/s) far faster than the
+    # device drains to SD (~80 KB/s). Pace fragment sends to the device's
+    # sustainable throughput to apply backpressure.
+    _WRITE_THROUGHPUT_BYTES_S = 80_000
 
     async def storage_write(self, path: str, content: bytes) -> bool:
         async with self._io_lock:
@@ -1145,6 +1166,9 @@ async def _storage_write_internal(self, path: str, content: bytes) -> bool:
                     return bool(
                         response and response.command_status == flipper_pb2.CommandStatus.OK
                     )
+                # Backpressure: hold the host to the device's drain rate so the
+                # unacked fragment stream does not overrun and wedge the session.
+                await asyncio.sleep(len(framed) / self._WRITE_THROUGHPUT_BYTES_S)
             return False
         except Exception:
             logger.debug("_storage_write_internal failed", exc_info=True)

src/flipperzero_mcp/tools/firmware.py

@@ -71,6 +71,24 @@ async def _reconnect_and_classify(client: Any, before: Any) -> tuple[Any, bool]:
     )
 
 
+async def _resync_session(client: Any) -> Any:
+    """Reconnect the transport to clear a wedged RPC session.
+
+    Large multi-MB writes can wedge the device's RPC session; a transport
+    teardown + reconnect re-establishes a fresh session that responds again.
+
+    Returns:
+        The fresh RPC client after a successful reconnect.
+
+    Raises:
+        FlashError: If the device cannot be reconnected.
+    """
+    await client.disconnect()
+    if not await client.connect() or client.rpc is None:
+        raise FlashError("could not re-establish the device session after a session wedge")
+    return client.rpc
+
+
 def register_firmware_tools(mcp: FastMCP) -> None:
     """Register the firmware-flash tool."""
 
@@ -124,7 +142,9 @@ async def flipperzero_firmware_install(
 
         try:
             bundle = await _resolve(source, before.target)
-            await install_bundle(rpc, bundle, pkg_name=_PKG_NAME)
+            await install_bundle(
+                rpc, bundle, pkg_name=_PKG_NAME, resync=lambda: _resync_session(client)
+            )
         except (BundleError, FlashError) as e:
             raise ToolError(str(e)) from e