fix: Retry on 404 in run details to handle read replica lag (#440)

Author: nahua-aignx

src/aignostics/platform/resources/runs.py

@@ -13,6 +13,7 @@
 from typing import Any, cast
 
 from aignx.codegen.api.public_api import PublicApi
+from aignx.codegen.exceptions import NotFoundException
 from aignx.codegen.exceptions import ServiceException
 from aignx.codegen.models import (
     CustomMetadataUpdateRequest,
@@ -42,6 +43,7 @@
     Retrying,
     retry_if_exception_type,
     stop_after_attempt,
+    stop_after_delay,
     wait_exponential_jitter,
 )
 from urllib3.exceptions import IncompleteRead, PoolError, ProtocolError, ProxyError
@@ -137,7 +139,8 @@ def for_run_id(cls, run_id: str, cache_token: bool = True) -> "Run":
     def details(self, nocache: bool = False, hide_platform_queue_position: bool = False) -> RunData:
         """Retrieves the current status of the application run.
 
-        Retries on network and server errors.
+        Retries on network and server errors. Additionally retries on
+        NotFoundException for up to 5 seconds to handle read replica lag.
 
         Args:
             nocache (bool): If True, skip reading from cache and fetch fresh data from the API.
@@ -149,24 +152,37 @@ def details(self, nocache: bool = False, hide_platform_queue_position: bool = Fa
             RunData: The run data.
 
         Raises:
+            NotFoundException: If the run is not found after retries.
             Exception: If the API request fails.
         """
 
         @cached_operation(ttl=settings().run_cache_ttl, use_token=True)
         def details_with_retry(run_id: str) -> RunData:
+            def _fetch() -> RunData:
+                return Retrying(
+                    retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
+                    stop=stop_after_attempt(settings().run_retry_attempts),
+                    wait=wait_exponential_jitter(
+                        initial=settings().run_retry_wait_min, max=settings().run_retry_wait_max
+                    ),
+                    before_sleep=_log_retry_attempt,
+                    reraise=True,
+                )(
+                    lambda: self._api.get_run_v1_runs_run_id_get(
+                        run_id,
+                        _request_timeout=settings().run_timeout,
+                        _headers={"User-Agent": user_agent()},
+                    )
+                )
+
+            # NOTE(nahua): Outer retry handles NotFoundException (read replica lag)
             return Retrying(
-                retry=retry_if_exception_type(exception_types=RETRYABLE_EXCEPTIONS),
-                stop=stop_after_attempt(settings().run_retry_attempts),
-                wait=wait_exponential_jitter(initial=settings().run_retry_wait_min, max=settings().run_retry_wait_max),
+                retry=retry_if_exception_type(exception_types=(NotFoundException,)),
+                stop=stop_after_delay(5),
+                wait=wait_exponential_jitter(initial=0.5, max=3),
                 before_sleep=_log_retry_attempt,
                 reraise=True,
-            )(
-                lambda: self._api.get_run_v1_runs_run_id_get(
-                    run_id,
-                    _request_timeout=settings().run_timeout,
-                    _headers={"User-Agent": user_agent()},
-                )
-            )
+            )(_fetch)
 
         run_data: RunData = details_with_retry(self.run_id, nocache=nocache)  # type: ignore[call-arg]
         if hide_platform_queue_position:

tests/aignostics/platform/resources/runs_test.py

@@ -600,3 +600,71 @@ def test_run_details_can_hide_platform_queue_position(
     result = app_run.details(hide_platform_queue_position=hide_platform_queue_position)
     assert result.num_preceding_items_org == run_data.num_preceding_items_org
     assert result.num_preceding_items_platform == expected_platform_queue_position
+
+
+@pytest.mark.unit
+def test_run_details_retries_on_not_found_then_succeeds(app_run, mock_api) -> None:
+    """Test that Run.details retries on NotFoundException and succeeds when the run becomes available.
+
+    This verifies the outer retry logic that handles read replica lag by retrying
+    NotFoundException until the run is found.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_api: Mock ExternalsApi instance.
+    """
+    from aignx.codegen.exceptions import NotFoundException
+
+    run_data = RunReadResponse.model_construct(run_id="test-run-id")
+    mock_api.get_run_v1_runs_run_id_get.side_effect = [
+        NotFoundException(),
+        NotFoundException(),
+        run_data,
+    ]
+
+    result = app_run.details()
+
+    assert result.run_id == "test-run-id"
+    assert mock_api.get_run_v1_runs_run_id_get.call_count == 3
+
+
+@pytest.mark.unit
+def test_run_details_raises_not_found_after_timeout(app_run, mock_api) -> None:
+    """Test that Run.details re-raises NotFoundException after the retry timeout expires.
+
+    This verifies that the outer retry gives up after the configured delay and
+    surfaces the NotFoundException to the caller.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_api: Mock ExternalsApi instance.
+    """
+    from aignx.codegen.exceptions import NotFoundException
+
+    mock_api.get_run_v1_runs_run_id_get.side_effect = NotFoundException()
+
+    with pytest.raises(NotFoundException):
+        app_run.details()
+
+    assert mock_api.get_run_v1_runs_run_id_get.call_count > 1
+
+
+@pytest.mark.unit
+def test_run_details_does_not_retry_other_exceptions(app_run, mock_api) -> None:
+    """Test that the outer retry does not catch non-NotFoundException errors.
+
+    This verifies that exceptions like ForbiddenException pass straight through
+    the outer retry without being retried.
+
+    Args:
+        app_run: Run instance with mock API.
+        mock_api: Mock ExternalsApi instance.
+    """
+    from aignx.codegen.exceptions import ForbiddenException
+
+    mock_api.get_run_v1_runs_run_id_get.side_effect = ForbiddenException()
+
+    with pytest.raises(ForbiddenException):
+        app_run.details()
+
+    assert mock_api.get_run_v1_runs_run_id_get.call_count == 1