From 9650a3647c958de849efa35ac05e372138905800 Mon Sep 17 00:00:00 2001
From: Michael Zimmermann <sigmaepsilon92@gmail.com>
Date: Wed, 8 Jul 2026 14:46:09 +0200
Subject: [PATCH] smpclient: upload: retry on timeouts

This is important for unrealiable transports like UDP over OpenThread.

Fixes: https://github.com/intercreate/smpclient/issues/56
---
 src/smpclient/__init__.py | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/src/smpclient/__init__.py b/src/smpclient/__init__.py
index 9db4a55..583fe6e 100644
--- a/src/smpclient/__init__.py
+++ b/src/smpclient/__init__.py
@@ -266,6 +266,7 @@ async def upload(
         first_timeout_s: float = 40.0,
         subsequent_timeout_s: float | None = None,
         use_sha: bool = True,
+        attempts: int = 3,
     ) -> AsyncIterator[int]:
         """Iteratively upload an `image` to `slot`, yielding the offset.
 
@@ -288,6 +289,7 @@ async def upload(
                 Zephyr's SMP server will fail with `MGMT_ERR.EINVAL` if the
                 MTU is too small to include both the SHA256 and the first 32-bytes
                 of the image.  Increase the MTU or set `use_sha=False` in this case.
+            attempts: Maximum number of attempts per chunk.
 
         Yields:
             the offset of the image upload
@@ -324,20 +326,33 @@ async def upload(
             assert_never(response)  # pragma: no cover
 
         # send chunks until the SMP server reports that the offset is at the end of the image
+        current_attempt = 0
         while response.off != len(image):
-            response = await self.request(
-                self._maximize_upload_packet(
-                    ImageUploadWrite(
-                        off=response.off,
-                        data=b"",
-                        len=len(image) if response.off == 0 else None,
-                        image=slot if response.off == 0 else None,
-                        upgrade=upgrade if response.off == 0 else None,
+            if current_attempt >= attempts:
+                raise SMPUploadError(f"Timed out {attempts} times")
+
+            try:
+                response = await self.request(
+                    self._maximize_upload_packet(
+                        ImageUploadWrite(
+                            off=response.off,
+                            data=b"",
+                            len=len(image) if response.off == 0 else None,
+                            image=slot if response.off == 0 else None,
+                            upgrade=upgrade if response.off == 0 else None,
+                        ),
+                        image,
                     ),
-                    image,
-                ),
-                timeout_s=subsequent_timeout_s,
-            )
+                    timeout_s=subsequent_timeout_s,
+                )
+                current_attempt = 0
+            except TimeoutError:
+                current_attempt += 1
+                continue
+            except SMPBadSequence:
+                current_attempt += 1
+                continue
+
             if error(response):
                 raise SMPUploadError(response)
             elif success(response):