fix: isolate retry state in parallel batches

pocketflow/__init__.py

@@ -1,4 +1,4 @@
-import asyncio, warnings, copy, time
+import asyncio, warnings, copy, time, contextvars
 
 class BaseNode:
     def __init__(self): self.params,self.successors={},{}
@@ -57,6 +57,12 @@ def _run(self,shared):
         return self.post(shared,pr,None)
 
 class AsyncNode(Node):
+    def __init__(self,max_retries=1,wait=0):
+        super().__init__(max_retries,wait); self._cur_retry=contextvars.ContextVar("cur_retry",default=0)
+    @property
+    def cur_retry(self): return self._cur_retry.get()
+    @cur_retry.setter
+    def cur_retry(self,value): self._cur_retry.set(value)
     async def prep_async(self,shared): pass
     async def exec_async(self,prep_res): pass
     async def exec_fallback_async(self,prep_res,exc): raise exc
@@ -97,4 +103,4 @@ class AsyncParallelBatchFlow(AsyncFlow,BatchFlow):
     async def _run_async(self,shared): 
         pr=await self.prep_async(shared) or []
         await asyncio.gather(*(self._orch_async(shared,{**self.params,**bp}) for bp in pr))
-        return await self.post_async(shared,pr,None)
+        return await self.post_async(shared,pr,None)

tests/test_async_parallel_batch_node.py

@@ -139,5 +139,37 @@ async def exec_async(self, item):
         self.assertLess(execution_order.index(1), execution_order.index(0))
         self.assertLess(execution_order.index(3), execution_order.index(2))
 
+    def test_retry_state_is_isolated_per_item(self):
+        """Each parallel item should observe and control its own retry attempt."""
+        class RetryProcessor(AsyncParallelBatchNode):
+            def __init__(self):
+                super().__init__(max_retries=2)
+                self.second_attempt_started = asyncio.Event()
+                self.attempts = {'a': [], 'b': []}
+
+            async def prep_async(self, shared_storage):
+                return ['a', 'b']
+
+            async def exec_async(self, item):
+                self.attempts[item].append(self.cur_retry)
+                if self.cur_retry == 0:
+                    if item == 'a':
+                        await self.second_attempt_started.wait()
+                    raise RuntimeError(f"retry {item}")
+
+                if item == 'b':
+                    self.second_attempt_started.set()
+                return f"{item}-ok"
+
+            async def post_async(self, shared_storage, prep_result, exec_result):
+                shared_storage['results'] = exec_result
+
+        shared_storage = {}
+        processor = RetryProcessor()
+        self.loop.run_until_complete(processor.run_async(shared_storage))
+
+        self.assertEqual(shared_storage['results'], ['a-ok', 'b-ok'])
+        self.assertEqual(processor.attempts, {'a': [0, 1], 'b': [0, 1]})
+
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()