Fix wait_condition to wake registered wakers on state_mut

crates/sdk-core/tests/integ_tests/workflow_tests/timers.rs

@@ -1,5 +1,6 @@
 use crate::common::{CoreWfStarter, build_fake_sdk, init_core_and_create_wf};
-use std::time::Duration;
+use futures_util::{StreamExt, stream::FuturesUnordered};
+use std::{future::Future, pin::Pin, time::Duration};
 use temporalio_client::WorkflowStartOptions;
 use temporalio_common::{
     prost_dur,
@@ -300,3 +301,43 @@ async fn cancel_before_sent_to_server() {
     worker.register_workflow::<CancelBeforeSentWf>();
     worker.run().await.unwrap();
 }
+
+#[workflow]
+#[derive(Default)]
+struct WaitConditionWakerWf {
+    done: bool,
+}
+
+#[workflow_methods]
+impl WaitConditionWakerWf {
+    #[run(name = DEFAULT_WORKFLOW_TYPE)]
+    async fn run(ctx: &mut WorkflowContext<Self>) -> WorkflowResult<()> {
+        let mut futs: FuturesUnordered<Pin<Box<dyn Future<Output = ()>>>> = FuturesUnordered::new();
+
+        // Future 1: await timer, then set flag via state_mut
+        let ctx1 = ctx.clone();
+        futs.push(Box::pin(async move {
+            ctx1.timer(Duration::from_millis(500)).await;
+            ctx1.state_mut(|s| s.done = true);
+        }));
+
+        // Future 2: wait_condition on the flag (waker-dependent inside FuturesUnordered)
+        let ctx2 = ctx.clone();
+        futs.push(Box::pin(async move {
+            ctx2.wait_condition(|s| s.done).await;
+        }));
+
+        // Drive both to completion
+        while futs.next().await.is_some() {}
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn wait_condition_waker_in_futures_unordered() {
+    let t = canned_histories::single_timer_wf_completes("1");
+    let mock_cfg = MockPollCfg::from_hist_builder(t);
+    let mut worker = build_fake_sdk(mock_cfg);
+    worker.register_workflow::<WaitConditionWakerWf>();
+    worker.run().await.unwrap();
+}

crates/sdk/src/workflow_context.rs

@@ -24,7 +24,7 @@ use std::{
         atomic::{AtomicBool, Ordering},
         mpsc::{Receiver, Sender},
     },
-    task::Poll,
+    task::{Poll, Waker},
     time::{Duration, SystemTime},
 };
 use temporalio_common::{
@@ -127,13 +127,18 @@ pub struct WorkflowContext<W> {
     sync: SyncWorkflowContext<W>,
     /// The workflow instance
     workflow_state: Rc<RefCell<W>>,
+    /// Wakers registered by `wait_condition` futures. Drained and woken on
+    /// every `state_mut` call so that waker-based combinators (e.g.
+    /// `FuturesOrdered`) re-poll the condition after state changes.
+    condition_wakers: Rc<RefCell<Vec<Waker>>>,
 }
 
 impl<W> Clone for WorkflowContext<W> {
     fn clone(&self) -> Self {
         Self {
             sync: self.sync.clone(),
             workflow_state: self.workflow_state.clone(),
+            condition_wakers: self.condition_wakers.clone(),
         }
     }
 }
@@ -791,6 +796,7 @@ impl<W> WorkflowContext<W> {
                 _phantom: PhantomData,
             },
             workflow_state,
+            condition_wakers: Rc::new(RefCell::new(Vec::new())),
         }
     }
 
@@ -803,6 +809,7 @@ impl<W> WorkflowContext<W> {
                 _phantom: PhantomData,
             },
             workflow_state: self.workflow_state.clone(),
+            condition_wakers: self.condition_wakers.clone(),
         }
     }
 
@@ -979,8 +986,16 @@ impl<W> WorkflowContext<W> {
     ///
     /// The borrow is scoped to the closure and cannot escape, preventing
     /// borrows from being held across await points.
+    ///
+    /// After the mutation, all wakers registered by pending `wait_condition`
+    /// futures are woken so that waker-based combinators (e.g.
+    /// `FuturesOrdered`) re-poll them on the next pass.
     pub fn state_mut<R>(&self, f: impl FnOnce(&mut W) -> R) -> R {
-        f(&mut *self.workflow_state.borrow_mut())
+        let result = f(&mut *self.workflow_state.borrow_mut());
+        for waker in self.condition_wakers.borrow_mut().drain(..) {
+            waker.wake();
+        }
+        result
     }
 
     /// Wait for some condition on workflow state to become true, yielding the workflow if not.
@@ -991,10 +1006,11 @@ impl<W> WorkflowContext<W> {
         &'a self,
         mut condition: impl FnMut(&W) -> bool + 'a,
     ) -> impl Future<Output = ()> + 'a {
-        future::poll_fn(move |_cx: &mut Context<'_>| {
+        future::poll_fn(move |cx: &mut Context<'_>| {
             if condition(&*self.workflow_state.borrow()) {
                 Poll::Ready(())
             } else {
+                self.condition_wakers.borrow_mut().push(cx.waker().clone());
                 Poll::Pending
             }
         })