[OTA-R] Fix retries on Query Image failure due to invalid session

src/app/clusters/ota-requestor/DefaultOTARequestor.cpp

@@ -648,6 +648,9 @@ IdleStateReason DefaultOTARequestor::MapErrorToIdleStateReason(CHIP_ERROR error)
 
 void DefaultOTARequestor::RecordNewUpdateState(OTAUpdateStateEnum newState, OTAChangeReasonEnum reason, CHIP_ERROR error)
 {
+    bool handleIdleStateEnter       = true;
+    IdleStateReason idleStateReason = IdleStateReason::kUnknown;
+
     // Set server UpdateState attribute
     OtaRequestorServerSetUpdateState(newState);
 
@@ -670,17 +673,25 @@ void DefaultOTARequestor::RecordNewUpdateState(OTAUpdateStateEnum newState, OTAC
 
     if ((newState == OTAUpdateStateEnum::kIdle) && (mCurrentUpdateState != OTAUpdateStateEnum::kIdle))
     {
-        IdleStateReason idleStateReason = MapErrorToIdleStateReason(error);
-
-        // Inform the driver that the core logic has entered the Idle state
-        mOtaRequestorDriver->HandleIdleStateEnter(idleStateReason);
+        idleStateReason      = MapErrorToIdleStateReason(error);
+        handleIdleStateEnter = true;
     }
     else if ((mCurrentUpdateState == OTAUpdateStateEnum::kIdle) && (newState != OTAUpdateStateEnum::kIdle))
     {
-        mOtaRequestorDriver->HandleIdleStateExit();
+        handleIdleStateEnter = false;
     }
 
+    // Update the new state before handling the state transition
     mCurrentUpdateState = newState;
+
+    if (handleIdleStateEnter)
+    {
+        mOtaRequestorDriver->HandleIdleStateEnter(idleStateReason);
+    }
+    else
+    {
+        mOtaRequestorDriver->HandleIdleStateExit();
+    }
 }
 
 void DefaultOTARequestor::RecordErrorUpdateState(CHIP_ERROR error, OTAChangeReasonEnum reason)

src/app/clusters/ota-requestor/DefaultOTARequestorDriver.cpp

@@ -45,6 +45,7 @@ namespace {
 using namespace app::Clusters::OtaSoftwareUpdateRequestor;
 using namespace app::Clusters::OtaSoftwareUpdateRequestor::Structs;
 
+constexpr uint8_t kMaxInvalidSessionRetries        = 3;  // Max # of query image retries to perform on invalid session error
 constexpr uint32_t kDelayQueryUponCommissioningSec = 30; // Delay before sending the initial image query after commissioning
 constexpr uint32_t kImmediateStartDelaySec         = 1;  // Delay before sending a query in response to UrgentUpdateAvailable
 constexpr System::Clock::Seconds32 kDefaultDelayedActionTime = System::Clock::Seconds32(120);
@@ -58,9 +59,10 @@ DefaultOTARequestorDriver * ToDriver(void * context)
 
 void DefaultOTARequestorDriver::Init(OTARequestorInterface * requestor, OTAImageProcessorInterface * processor)
 {
-    mRequestor          = requestor;
-    mImageProcessor     = processor;
-    mProviderRetryCount = 0;
+    mRequestor                = requestor;
+    mImageProcessor           = processor;
+    mProviderRetryCount       = 0;
+    mInvalidSessionRetryCount = 0;
 
     if (mImageProcessor->IsFirstImageRun())
     {
@@ -127,6 +129,11 @@ void DefaultOTARequestorDriver::HandleIdleStateExit()
 
 void DefaultOTARequestorDriver::HandleIdleStateEnter(IdleStateReason reason)
 {
+    if (reason != IdleStateReason::kInvalidSession)
+    {
+        mInvalidSessionRetryCount = 0;
+    }
+
     switch (reason)
     {
     case IdleStateReason::kUnknown:
@@ -138,8 +145,17 @@ void DefaultOTARequestorDriver::HandleIdleStateEnter(IdleStateReason reason)
         StartSelectedTimer(SelectedTimer::kPeriodicQueryTimer);
         break;
     case IdleStateReason::kInvalidSession:
-        // An invalid session is detected which may be temporary so try to query the same provider again
-        SendQueryImage();
+        if (mInvalidSessionRetryCount < kMaxInvalidSessionRetries)
+        {
+            // An invalid session is detected which may be temporary so try to query the same provider again
+            SendQueryImage();
+            mInvalidSessionRetryCount++;
+        }
+        else
+        {
+            mInvalidSessionRetryCount = 0;
+            StartSelectedTimer(SelectedTimer::kPeriodicQueryTimer);
+        }
         break;
     }
 }

src/app/clusters/ota-requestor/DefaultOTARequestorDriver.h

@@ -109,7 +109,10 @@ class DefaultOTARequestorDriver : public OTARequestorDriver
     uint16_t maxDownloadBlockSize  = 1024;
     // Maximum number of times to retry a BUSY OTA provider before moving to the next available one
     static constexpr uint8_t kMaxBusyProviderRetryCount = 3;
-    uint8_t mProviderRetryCount; // Track retry count for the current provider
+    // Track retry count for the current provider
+    uint8_t mProviderRetryCount;
+    // Track query image retry count on invalid session error
+    uint8_t mInvalidSessionRetryCount;
 };
 
 } // namespace DeviceLayer