From b46736062a63973e6426aedf9d1f48821ce68914 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 27 Oct 2026 19:54:39 +0000
Subject: [PATCH] Retry failed and cancelled async predictions instead of
 caching errors

---
 great_ai/deploy/great_ai.py     |  2 +-
 great_ai/external/async_lru.py  |  4 ++--
 tests/test_cache_regressions.py | 27 +++++++++++++++++++++++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py
index 607cd7a..c529ac4 100644
--- a/great_ai/deploy/great_ai.py
+++ b/great_ai/deploy/great_ai.py
@@ -265,7 +265,7 @@ class GreatAI(Generic[T, V]):
                 result = func(*args, **kwargs)
                 return cast(T, t.finalise(output=result))
 
-        @alru_cache(maxsize=get_context().prediction_cache_size)
+        @alru_cache(maxsize=get_context().prediction_cache_size, cache_exceptions=False)
         async def func_in_tracing_context_async(
             *args: Any,
             do_not_persist_traces: bool = False,
diff --git a/great_ai/external/async_lru.py b/great_ai/external/async_lru.py
index c209d5a..52c39f2 100644
--- a/great_ai/external/async_lru.py
+++ b/great_ai/external/async_lru.py
@@ -148,9 +148,9 @@ def alru_cache(
                     _cache_hit(wrapped, key)
                     return await asyncio.shield(fut)
 
-                exc = fut._exception
+                exc = None if fut.cancelled() else fut.exception()
 
-                if exc is None or cache_exceptions:
+                if not fut.cancelled() and (exc is None or cache_exceptions):
                     _cache_hit(wrapped, key)
                     return fut.result()
 
diff --git a/tests/test_cache_regressions.py b/tests/test_cache_regressions.py
index 4edfa5c..004dbbe 100644
--- a/tests/test_cache_regressions.py
+++ b/tests/test_cache_regressions.py
@@ -1,5 +1,8 @@
+import asyncio
 from typing import Dict, Tuple
 
+import pytest
+
 from great_ai import GreatAI
 
 
@@ -15,3 +18,27 @@ def test_nested_tuple_inputs_can_be_cached(isolated_context: None) -> None:
     assert predict(({"x": 7},)).output == 7
     assert predict(({"x": 7},)).output == 7
     assert calls == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("cancelled", [False, True])
+async def test_async_predictions_retry_after_failure(
+    isolated_context: None, cancelled: bool
+) -> None:
+    calls = 0
+
+    @GreatAI.create
+    async def predict(value: int) -> int:
+        nonlocal calls
+        calls += 1
+        if calls == 1:
+            if cancelled:
+                raise asyncio.CancelledError()
+            raise ValueError("temporary failure")
+        return value
+
+    with pytest.raises(asyncio.CancelledError if cancelled else ValueError):
+        await predict(7)
+    assert (await predict(7)).output == 7
+    assert (await predict(7)).output == 7
+    assert calls == 2