[Rust Frontend] Start current wave for a stale DP FirstRequest (#46833)

Signed-off-by: Blas Rodriguez Irizar <rodrigblas@gmail.com>

rust/src/engine-core-client/src/coordinator/handle.rs

@@ -20,6 +20,27 @@ pub(crate) struct CoordinatorStateSnapshot {
     pub engines_running: bool,
 }
 
+impl CoordinatorStateSnapshot {
+    /// Resume the engines for a `FirstRequest` and return the wave to broadcast
+    /// and the engine to exclude from the wakeup.
+    ///
+    /// The request may have been stamped with a `request_wave` older than
+    /// `current_wave` if a `WaveComplete` advanced it after the command was
+    /// enqueued. Such a request still needs serving, so the current wave is
+    /// broadcast to every engine (`exclude = None`); the wave is never rewound.
+    /// A non-stale request excludes the engine that already received it. Mirrors
+    /// the Python coordinator's front-end path.
+    pub(crate) fn start_wave_for_first_request(
+        &mut self,
+        request_wave: u32,
+        target_engine_index: u32,
+    ) -> (u32, Option<u32>) {
+        self.engines_running = true;
+        let exclude = (request_wave >= self.current_wave).then_some(target_engine_index);
+        (self.current_wave, exclude)
+    }
+}
+
 /// Shared in-process coordinator state.
 pub(crate) type CoordinatorState = Mutex<CoordinatorStateSnapshot>;
 

rust/src/engine-core-client/src/coordinator/inproc.rs

@@ -27,9 +27,10 @@ use crate::protocol::{
 struct StartDpWaveMessage {
     /// DP wave number that all engines should start processing.
     wave: u32,
-    /// Engine index that already received the triggering request and should not
-    /// receive an extra wakeup notification.
-    exclude_engine_index: u32,
+    /// Engine index that already received the triggering request and so does not
+    /// need an extra wakeup. `None` wakes every engine (used when the triggering
+    /// request was for a stale wave).
+    exclude_engine_index: Option<u32>,
 }
 
 /// Background half of the in-process coordinator.
@@ -57,7 +58,11 @@ impl InProcCoordinatorRunner {
     }
 
     /// Broadcast Python-compatible `START_DP_WAVE` to all connected engines.
-    async fn broadcast_start_wave(&mut self, wave: u32, exclude_engine_index: u32) -> Result<()> {
+    async fn broadcast_start_wave(
+        &mut self,
+        wave: u32,
+        exclude_engine_index: Option<u32>,
+    ) -> Result<()> {
         let payload = encode_msgpack(&StartDpWaveMessage {
             wave,
             exclude_engine_index,
@@ -86,13 +91,17 @@ impl InProcCoordinatorRunner {
                         engine_id: target_engine_id.to_vec(),
                     }
                 })?;
-                self.state.lock().current_wave = wave;
+                let (current_wave, exclude) = {
+                    let mut state = self.state.lock();
+                    state.start_wave_for_first_request(wave, target_engine_index)
+                };
                 debug!(
-                    wave,
-                    exclude_engine_index = target_engine_index,
+                    current_wave,
+                    request_wave = wave,
+                    ?exclude,
                     "starting DP wave after first request while engines were paused"
                 );
-                self.broadcast_start_wave(wave, target_engine_index).await?;
+                self.broadcast_start_wave(current_wave, exclude).await?;
             }
         }
         Ok(())
@@ -150,7 +159,7 @@ impl InProcCoordinatorRunner {
                             exclude_engine_index = engine_index,
                             "starting DP wave after stale-wave notification from engine"
                         );
-                        self.broadcast_start_wave(wave, engine_index).await?;
+                        self.broadcast_start_wave(wave, Some(engine_index)).await?;
                     }
                 }
             },
@@ -202,3 +211,48 @@ impl InProcCoordinatorRunner {
         inner.close_registries(Arc::new(error));
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::coordinator::handle::CoordinatorStateSnapshot;
+
+    /// A `FirstRequest` for the current wave starts that wave and excludes the
+    /// engine that already received the triggering request.
+    #[test]
+    fn first_request_for_current_wave_excludes_target() {
+        let mut state = CoordinatorStateSnapshot {
+            current_wave: 3,
+            engines_running: false,
+        };
+
+        let (wave, exclude) = state.start_wave_for_first_request(3, 2);
+
+        assert_eq!(wave, 3);
+        assert_eq!(exclude, Some(2));
+        assert!(state.engines_running);
+        assert_eq!(state.current_wave, 3);
+    }
+
+    /// A `FirstRequest` whose wave was superseded by a racing `WaveComplete`
+    /// (`request_wave < current_wave`) must still start the request's wave: it
+    /// broadcasts the current wave and wakes every engine (`exclude = None`)
+    /// rather than rewinding the wave or dropping the request.
+    #[test]
+    fn stale_first_request_starts_current_wave_for_all_engines() {
+        let mut state = CoordinatorStateSnapshot {
+            current_wave: 4,
+            engines_running: false,
+        };
+
+        // Request stamped with wave 3 while the coordinator already advanced to 4.
+        let (wave, exclude) = state.start_wave_for_first_request(3, 2);
+
+        assert_eq!(
+            wave, 4,
+            "must broadcast the current wave, not the stale one"
+        );
+        assert_eq!(exclude, None, "a stale request must wake every engine");
+        assert!(state.engines_running);
+        assert_eq!(state.current_wave, 4, "wave must not be rewound");
+    }
+}