Restore original 404 retry logic (with counter) and add separate 406 handling

Agent-Logs-Url: https://github.com/Azure/azure-sdk-for-python/sessions/dbd04bba-1b0d-4c00-92c0-5ee3d41d6653

Co-authored-by: PallabPaul <20746357+PallabPaul@users.noreply.github.com>

Author: Copilot

sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/_operations/_patch.py

@@ -55,6 +55,7 @@ def __init__(
         self._latest_response: JSON = {}
 
         self._retry_not_found = retry_not_found
+        self._not_found_count = 0
 
     def initialize(self, client, initial_response, deserialization_callback):  # pylint: disable=unused-argument
         self._evaluate_response(initial_response)
@@ -101,17 +102,22 @@ def run(self) -> None:
                 try:
                     response = self._operation()
                     self._evaluate_response(response)
+                except ResourceNotFoundError as not_found_exception:
+                    # We'll allow some instances of resource not found to account for replication
+                    # delay if session stickiness is lost.
+
+                    self._not_found_count += 1
+
+                    not_retryable = not self._retry_not_found or self._give_up_not_found_error(not_found_exception)
+
+                    if not_retryable or self._not_found_count >= 3:
+                        raise
                 except HttpResponseError as http_exception:
-                    # 404: node doesn't know about the transaction yet (replication lag)
-                    # 406: node knows the transaction but it hasn't committed yet
-                    # Both are transient during transaction-status polling — treat as
+                    # 406: node knows the transaction but it hasn't committed yet.
+                    # This is transient during transaction-status polling — treat as
                     # pending so the poller retries.
-                    if http_exception.status_code in (404, 406) and self._retry_not_found:
-                        if (
-                            isinstance(http_exception, ResourceNotFoundError)
-                            and self._give_up_not_found_error(http_exception)
-                        ):
-                            raise
+                    if http_exception.status_code == 406 and self._retry_not_found:
+                        pass  # stay in polling loop
                     else:
                         raise
                 if not self.finished():

sdk/confidentialledger/azure-confidentialledger/azure/confidentialledger/aio/_operations/_patch.py

@@ -56,17 +56,22 @@ async def run(self) -> None:
                 try:
                     response = await self._operation()
                     self._evaluate_response(response)
+                except ResourceNotFoundError as not_found_exception:
+                    # We'll allow some instances of resource not found to account for replication
+                    # delay if session stickiness is lost.
+
+                    self._not_found_count += 1
+
+                    not_retryable = not self._retry_not_found or self._give_up_not_found_error(not_found_exception)
+
+                    if not_retryable or self._not_found_count >= 3:
+                        raise
                 except HttpResponseError as http_exception:
-                    # 404: node doesn't know about the transaction yet (replication lag)
-                    # 406: node knows the transaction but it hasn't committed yet
-                    # Both are transient during transaction-status polling — treat as
+                    # 406: node knows the transaction but it hasn't committed yet.
+                    # This is transient during transaction-status polling — treat as
                     # pending so the poller retries.
-                    if http_exception.status_code in (404, 406) and self._retry_not_found:
-                        if (
-                            isinstance(http_exception, ResourceNotFoundError)
-                            and self._give_up_not_found_error(http_exception)
-                        ):
-                            raise
+                    if http_exception.status_code == 406 and self._retry_not_found:
+                        pass  # stay in polling loop
                     else:
                         raise
                 if not self.finished():

sdk/confidentialledger/azure-confidentialledger/tests/test_polling.py

@@ -38,17 +38,17 @@ def _make_http_response_error(status_code, message="Error"):
 
 class TestStatePollingMethodTransient404:
     """Verify that 404 responses during polling are treated as transient when
-    retry_not_found is True.
+    retry_not_found is True, up to 3 retries.
     """
 
     def test_404_retried_then_committed(self):
-        """404 should be retried and eventually succeed when the transaction commits."""
+        """404 should be retried (up to 3 times) and eventually succeed."""
         call_count = 0
 
         def operation():
             nonlocal call_count
             call_count += 1
-            if call_count <= 3:
+            if call_count <= 2:
                 raise _make_resource_not_found()
             return {"state": "Committed"}
 
@@ -57,7 +57,7 @@ def operation():
         poller.run()
 
         assert poller.status() == "finished"
-        assert call_count == 4
+        assert call_count == 3
 
     def test_404_raises_when_retry_not_found_false(self):
         """404 should raise immediately when retry_not_found is False."""
@@ -89,23 +89,19 @@ def operation():
 
         assert poller.status() == "failed"
 
-    def test_many_404s_retried_without_limit(self):
-        """More than 3 consecutive 404s should still be retried (no hard limit)."""
-        call_count = 0
+    def test_404_raises_after_3_retries(self):
+        """After 3 consecutive 404s the poller should give up and raise."""
 
         def operation():
-            nonlocal call_count
-            call_count += 1
-            if call_count <= 10:
-                raise _make_resource_not_found()
-            return {"state": "Committed"}
+            raise _make_resource_not_found()
 
         poller = StatePollingMethod(operation, "Committed", 0, retry_not_found=True)
         poller._status = "polling"
-        poller.run()
 
-        assert poller.status() == "finished"
-        assert call_count == 11
+        with pytest.raises(ResourceNotFoundError):
+            poller.run()
+
+        assert poller.status() == "failed"
 
 
 class TestStatePollingMethodTransient406:
@@ -237,7 +233,7 @@ async def test_404_retried_then_committed(self):
         async def operation():
             nonlocal call_count
             call_count += 1
-            if call_count <= 3:
+            if call_count <= 2:
                 raise _make_resource_not_found()
             return {"state": "Committed"}
 
@@ -246,7 +242,7 @@ async def operation():
         await poller.run()
 
         assert poller.status() == "finished"
-        assert call_count == 4
+        assert call_count == 3
 
     @pytest.mark.asyncio
     async def test_404_raises_when_retry_not_found_false(self):
@@ -275,22 +271,19 @@ async def operation():
         assert poller.status() == "failed"
 
     @pytest.mark.asyncio
-    async def test_many_404s_retried_without_limit(self):
-        call_count = 0
+    async def test_404_raises_after_3_retries(self):
+        """After 3 consecutive 404s the async poller should give up and raise."""
 
         async def operation():
-            nonlocal call_count
-            call_count += 1
-            if call_count <= 10:
-                raise _make_resource_not_found()
-            return {"state": "Committed"}
+            raise _make_resource_not_found()
 
         poller = AsyncStatePollingMethod(operation, "Committed", 0, retry_not_found=True)
         poller._status = "polling"
-        await poller.run()
 
-        assert poller.status() == "finished"
-        assert call_count == 11
+        with pytest.raises(ResourceNotFoundError):
+            await poller.run()
+
+        assert poller.status() == "failed"
 
 
 class TestAsyncStatePollingMethodTransient406: