@@ -266,6 +266,7 @@ async def upload(
266266 first_timeout_s : float = 40.0 ,
267267 subsequent_timeout_s : float | None = None ,
268268 use_sha : bool = True ,
269+ attempts : int = 3 ,
269270 ) -> AsyncIterator [int ]:
270271 """Iteratively upload an `image` to `slot`, yielding the offset.
271272
@@ -288,6 +289,7 @@ async def upload(
288289 Zephyr's SMP server will fail with `MGMT_ERR.EINVAL` if the
289290 MTU is too small to include both the SHA256 and the first 32-bytes
290291 of the image. Increase the MTU or set `use_sha=False` in this case.
292+ attempts: Maximum number of attempts per chunk.
291293
292294 Yields:
293295 the offset of the image upload
@@ -324,20 +326,33 @@ async def upload(
324326 assert_never (response ) # pragma: no cover
325327
326328 # send chunks until the SMP server reports that the offset is at the end of the image
329+ current_attempt = 0
327330 while response .off != len (image ):
328- response = await self .request (
329- self ._maximize_upload_packet (
330- ImageUploadWrite (
331- off = response .off ,
332- data = b"" ,
333- len = len (image ) if response .off == 0 else None ,
334- image = slot if response .off == 0 else None ,
335- upgrade = upgrade if response .off == 0 else None ,
331+ if current_attempt >= attempts :
332+ raise SMPUploadError (f"Timed out { attempts } times" )
333+
334+ try :
335+ response = await self .request (
336+ self ._maximize_upload_packet (
337+ ImageUploadWrite (
338+ off = response .off ,
339+ data = b"" ,
340+ len = len (image ) if response .off == 0 else None ,
341+ image = slot if response .off == 0 else None ,
342+ upgrade = upgrade if response .off == 0 else None ,
343+ ),
344+ image ,
336345 ),
337- image ,
338- ),
339- timeout_s = subsequent_timeout_s ,
340- )
346+ timeout_s = subsequent_timeout_s ,
347+ )
348+ current_attempt = 0
349+ except TimeoutError :
350+ current_attempt += 1
351+ continue
352+ except SMPBadSequence :
353+ current_attempt += 1
354+ continue
355+
341356 if error (response ):
342357 raise SMPUploadError (response )
343358 elif success (response ):
0 commit comments