Lift take async result outside of a poll coro

Author: igalshilman

python/restate/server_context.py

@@ -246,19 +246,19 @@ def must_take_notification(self, handle):
         return res
 
 
-    async def create_poll_or_cancel_coroutine(self, handle) -> bytes | None:
+    async def create_poll_or_cancel_coroutine(self, handles):
         """Create a coroutine to poll the handle."""
         await self.take_and_send_output()
         while True:
-            if self.vm.is_completed(handle):
+            #if self.vm.is_completed(handle):
                 # Handle is completed
-                return self.must_take_notification(handle)
+            #    return self.must_take_notification(handle)
 
             # Nothing ready yet, let's try to make some progress
-            do_progress_response = self.vm.do_progress([handle])
+            do_progress_response = self.vm.do_progress(handles)
             if isinstance(do_progress_response, DoProgressAnyCompleted):
                 # One of the handles completed, we can continue
-                continue
+                return
             if isinstance(do_progress_response, DoProgressCancelSignalReceived):
                 # Raise cancel signal
                 raise TerminalError("cancelled", 409)
@@ -280,7 +280,8 @@ def create_df(self, handle: int, serde: Serde[T] | None = None) -> ServerDurable
         """Create a durable future."""
 
         async def transform():
-            res = await self.create_poll_or_cancel_coroutine(handle)
+            await self.create_poll_or_cancel_coroutine(handle)
+            res = self.must_take_notification(handle)
             if res is None or serde is None:
                 return res
             return serde.deserialize(res)
@@ -293,13 +294,15 @@ def create_call_df(self, handle: int, invocation_id_handle: int, serde: Serde[T]
         """Create a durable future."""
 
         async def transform():
-            res = await self.create_poll_or_cancel_coroutine(handle)
+            await self.create_poll_or_cancel_coroutine(handle)
+            res = self.must_take_notification(handle)
             if res is None or serde is None:
                 return res
             return serde.deserialize(res)
 
-        def inv_id_factory():
-            return self.create_poll_or_cancel_coroutine(invocation_id_handle)
+        async def inv_id_factory():
+            await self.create_poll_or_cancel_coroutine(invocation_id_handle)
+            return self.must_take_notification(invocation_id_handle)
 
         return ServerCallDurableFuture(handle, transform, invocation_id_handle, inv_id_factory)