From 8de675f3a5eac2aa82eadefbfae909be3e620cf1 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Tue, 27 Oct 2026 16:52:00 +0000
Subject: [PATCH] Use fresh worker caches for async batch predictions

---
 great_ai/deploy/great_ai.py          |  6 +++++-
 tests/test_prediction_regressions.py | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/great_ai/deploy/great_ai.py b/great_ai/deploy/great_ai.py
index ac14590..607cd7a 100644
--- a/great_ai/deploy/great_ai.py
+++ b/great_ai/deploy/great_ai.py
@@ -205,7 +205,11 @@ class GreatAI(Generic[T, V]):
                 evaluations run part of the CI.
         """
 
-        wrapped_function = self._wrapped_func
+        # Futures in a warmed async cache cannot be serialized into worker
+        # processes. Build a fresh cache around the original decorated function.
+        wrapped_function = freeze_arguments(
+            self._get_cached_traced_function(self._wrapped_func.__wrapped__)
+        )
 
         def inner(value: Any) -> T:
             return (
diff --git a/tests/test_prediction_regressions.py b/tests/test_prediction_regressions.py
index 38e64bd..da84fa5 100644
--- a/tests/test_prediction_regressions.py
+++ b/tests/test_prediction_regressions.py
@@ -58,6 +58,20 @@ def test_http_preserves_nested_models(isolated_context: None, is_async: bool) ->
         assert response.json()["output"] == 5
 
 
+@pytest.mark.asyncio
+async def test_batch_after_async_prediction(isolated_context: None) -> None:
+    @GreatAI.create
+    async def predict(value: int) -> int:
+        return value + 1
+
+    assert (await predict(1)).output == 2
+    assert [trace.output for trace in predict.process_batch([1, 2], concurrency=1)] == [
+        2,
+        3,
+    ]
+    assert (await predict(1)).output == 2
+
+
 def test_flattening_preserves_reserved_trace_fields() -> None:
     trace = Trace(
         trace_id="actual-id",