fix(replay): route step-action buttons through selected timeline + dispatch carries at_step

Cargo.toml

@@ -14,7 +14,7 @@ members = [
 resolver = "2"
 
 [workspace.package]
-version = "0.14.7"
+version = "0.14.8"
 edition = "2024"
 authors = ["Rewind Contributors"]
 license = "MIT"

crates/rewind-web/src/dispatcher.rs

@@ -21,6 +21,8 @@
 //!   "job_id": "<uuid>",
 //!   "session_id": "<uuid>",
 //!   "replay_context_id": "<uuid>",
+//!   "replay_context_timeline_id": "<uuid>",
+//!   "at_step": <u32>,
 //!   "base_url": "<rewind-server-base>"
 //! }
 //! ```
@@ -74,12 +76,25 @@ type HmacSha256 = Hmac<Sha256>;
 /// the replay context targets — the SDK's `attach_replay_context`
 /// uses it to set `_timeline_id` so any cache-miss live recordings
 /// land in the fork rather than the original timeline.
+///
+/// **`at_step`** (added 2026-04-29): the original fork-point of the
+/// replay-context's timeline (i.e. the step number the user clicked
+/// "Run replay" at in the dashboard). Distinct from the replay
+/// context's `from_step` which the `/api/sessions/{sid}/replay-jobs`
+/// handler hardcodes to 0 (see runners.rs `create_replay_job` —
+/// the agent re-runs the loop from scratch so the replay-lookup
+/// ordinal cursor must start at recorded step #1). Runners use
+/// `at_step` to know which conversation turn the user wanted to
+/// start from — needed for multi-turn replay where edits to
+/// step #N's user message in turn 2+ should drive the agent at
+/// iteration N (companion ray-agent change).
 #[derive(Debug, serde::Serialize)]
 struct DispatchBody<'a> {
     pub job_id: &'a str,
     pub session_id: &'a str,
     pub replay_context_id: &'a str,
     pub replay_context_timeline_id: &'a str,
+    pub at_step: u32,
     pub base_url: &'a str,
 }
 
@@ -159,11 +174,16 @@ impl Dispatcher {
     /// **Review #154 F2:** caller passes `replay_context_timeline_id`
     /// so the dispatch payload can carry it and the SDK's
     /// `attach_replay_context` can set `_timeline_id` correctly.
+    ///
+    /// `at_step` (added 2026-04-29): the fork-point of the timeline
+    /// (= the step number the user clicked Run replay at). Forwarded
+    /// in the dispatch body for runner-side multi-turn replay.
     pub async fn dispatch(
         &self,
         runner: &Runner,
         job: &ReplayJob,
         replay_context_timeline_id: &str,
+        at_step: u32,
     ) -> DispatchOutcome {
         let webhook_url = match runner.webhook_url.as_deref() {
             Some(u) => u,
@@ -208,6 +228,7 @@ impl Dispatcher {
             session_id: &job.session_id,
             replay_context_id,
             replay_context_timeline_id,
+            at_step,
             base_url: &self.base_url,
         };
         let body_bytes = match serde_json::to_vec(&body) {

crates/rewind-web/src/runners.rs

@@ -558,7 +558,7 @@ async fn create_replay_job(
     // 1. Resolve session, runner, and replay-context (creating fork+
     //    context if shape A). All store work in a single lock scope
     //    so we don't race with deletions between checks.
-    let (job, runner, fork_timeline_id) = {
+    let (job, runner, fork_timeline_id, at_step) = {
         let store = state
             .store
             .lock()
@@ -568,7 +568,7 @@ async fn create_replay_job(
             .map_err(internal)?
             .ok_or_else(|| not_found("session"))?;
 
-        let (runner_id, replay_context_id, fork_timeline_id) = match req {
+        let (runner_id, replay_context_id, fork_timeline_id, at_step) = match req {
             CreateReplayJobRequest::CreateAndDispatch(a) => {
                 // Validate runner.
                 let runner = store
@@ -631,7 +631,7 @@ async fn create_replay_job(
                         .set_replay_context_strict_match(&ctx_id, true)
                         .map_err(internal)?;
                 }
-                (a.runner_id, ctx_id, Some(fork.id))
+                (a.runner_id, ctx_id, Some(fork.id), a.at_step)
             }
             CreateReplayJobRequest::ReuseContext(b) => {
                 let runner = store
@@ -691,7 +691,21 @@ async fn create_replay_job(
                         b.replay_context_id
                     )));
                 }
-                (b.runner_id, b.replay_context_id, Some(ctx.timeline_id))
+                // Resolve at_step from the context's fork timeline. The
+                // dispatch payload carries the user-clicked step number
+                // so the runner can drive multi-turn replay; for
+                // ReuseContext that step is baked into the timeline's
+                // fork_at_step. Root timelines (parent_timeline_id =
+                // None) shouldn't reach here because contexts always
+                // target a fork; defensive fallback to 1.
+                let at_step = store
+                    .get_timelines(&session.id)
+                    .map_err(internal)?
+                    .iter()
+                    .find(|t| t.id == ctx.timeline_id)
+                    .and_then(|t| t.fork_at_step)
+                    .unwrap_or(1);
+                (b.runner_id, b.replay_context_id, Some(ctx.timeline_id), at_step)
             }
         };
 
@@ -741,7 +755,7 @@ async fn create_replay_job(
                 None,
             )
             .map_err(internal)?;
-        (job, runner, fork_timeline_id)
+        (job, runner, fork_timeline_id, at_step)
     };
 
     // 3. Fire the dispatcher in the background. The handler returns
@@ -762,9 +776,10 @@ async fn create_replay_job(
     let timeline_for_payload = fork_timeline_id
         .clone()
         .unwrap_or_default();
+    let at_step_for_payload = at_step;
     tokio::spawn(async move {
         let outcome = dispatcher
-            .dispatch(&runner_clone, &job_clone, &timeline_for_payload)
+            .dispatch(&runner_clone, &job_clone, &timeline_for_payload, at_step_for_payload)
             .await;
         let after = {
             let store = match store_arc.lock() {

crates/rewind-web/tests/replay_jobs_tests.rs

@@ -130,6 +130,34 @@ async fn spawn_runner_stub_accepting() -> (String, tokio::sync::mpsc::Receiver<a
     (format!("http://{addr}/wh"), rx)
 }
 
+/// Variant of `spawn_runner_stub_accepting` that ALSO captures the
+/// request body bytes. Used by tests that assert on dispatch payload
+/// fields (e.g. `at_step`).
+async fn spawn_runner_stub_capturing_body() -> (
+    String,
+    tokio::sync::mpsc::Receiver<(axum::http::HeaderMap, axum::body::Bytes)>,
+) {
+    let (tx, rx) = tokio::sync::mpsc::channel(8);
+    let app = axum::Router::new().route(
+        "/wh",
+        axum::routing::post(
+            move |headers: axum::http::HeaderMap, body: axum::body::Bytes| {
+                let tx = tx.clone();
+                async move {
+                    let _ = tx.send((headers, body)).await;
+                    StatusCode::ACCEPTED
+                }
+            },
+        ),
+    );
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let addr = listener.local_addr().unwrap();
+    tokio::spawn(async move {
+        let _ = axum::serve(listener, app).await;
+    });
+    (format!("http://{addr}/wh"), rx)
+}
+
 async fn spawn_runner_stub_returning_500() -> String {
     let app = axum::Router::new().route(
         "/wh",
@@ -565,6 +593,94 @@ async fn shape_a_strict_match_omitted_defaults_to_warn_only() {
     assert!(!ctx.strict_match);
 }
 
+// ── 2026-04-29: dispatch payload carries `at_step` ──
+
+#[tokio::test]
+async fn dispatch_payload_carries_at_step() {
+    // Regression for the dev1 bug where the runner can't drive
+    // multi-turn replay because the dispatch payload doesn't tell it
+    // which conversation turn the user clicked on. The dispatcher now
+    // forwards the timeline's `fork_at_step` (= the `at_step` the
+    // user requested) so a future runner-side change can use it to
+    // reconstruct conversation history and replay from the right
+    // iteration. This test pins the wire-format guarantee end-to-end:
+    // POST shape A with at_step=4 → captured runner body has
+    // `"at_step": 4`.
+    let (api, _callbacks, store, _tmp) = setup();
+    let (webhook_url, mut rx) = spawn_runner_stub_capturing_body().await;
+    let (runner_id, _raw) = register_runner(&store, &webhook_url);
+    // Need >= 4 steps so a fork at_step=4 is valid.
+    let (session_id, root_timeline_id) = seed_session_with_n_steps(&store, 4);
+
+    let (status, _body) = json_post(
+        api,
+        &format!("/api/sessions/{session_id}/replay-jobs"),
+        json!({
+            "runner_id": runner_id,
+            "source_timeline_id": root_timeline_id,
+            "at_step": 4,
+        }),
+    )
+    .await;
+    assert_eq!(status, StatusCode::ACCEPTED);
+
+    let (_headers, body_bytes) =
+        tokio::time::timeout(std::time::Duration::from_secs(3), rx.recv())
+            .await
+            .expect("dispatcher must call the runner within 3s")
+            .expect("stub channel closed");
+    let body_json: Value = serde_json::from_slice(&body_bytes)
+        .expect("dispatch body must be valid JSON");
+    assert_eq!(
+        body_json["at_step"].as_u64(),
+        Some(4),
+        "dispatch payload must carry the user-clicked at_step (got: {body_json})"
+    );
+    // Sanity: the other fields the SDK already reads must still be present.
+    assert!(body_json["job_id"].is_string());
+    assert!(body_json["session_id"].is_string());
+    assert!(body_json["replay_context_id"].is_string());
+    assert!(body_json["replay_context_timeline_id"].is_string());
+    assert!(body_json["base_url"].is_string());
+}
+
+#[tokio::test]
+async fn dispatch_payload_at_step_for_reuse_context_reads_fork_at_step() {
+    // When dispatching against a pre-existing replay context (shape B),
+    // there's no `at_step` on the request body — but the dispatch
+    // payload still needs `at_step` so the runner can drive multi-turn
+    // replay. The handler reads it from the context's fork timeline's
+    // `fork_at_step` field. For root timelines (no parent) the
+    // fallback is 1.
+    let (api, _callbacks, store, _tmp) = setup();
+    let (webhook_url, mut rx) = spawn_runner_stub_capturing_body().await;
+    let (runner_id, _raw) = register_runner(&store, &webhook_url);
+    // seed_session_and_context creates a context against the ROOT
+    // timeline (parent_timeline_id=None, fork_at_step=None). The
+    // handler should default at_step to 1.
+    let (session_id, _root, ctx_id) = seed_session_and_context(&store);
+
+    let (status, _body) = json_post(
+        api,
+        &format!("/api/sessions/{session_id}/replay-jobs"),
+        json!({"runner_id": runner_id, "replay_context_id": ctx_id}),
+    )
+    .await;
+    assert_eq!(status, StatusCode::ACCEPTED);
+
+    let (_headers, body_bytes) =
+        tokio::time::timeout(std::time::Duration::from_secs(3), rx.recv())
+            .await
+            .expect("dispatcher must call the runner within 3s")
+            .expect("stub channel closed");
+    let body_json: Value = serde_json::from_slice(&body_bytes).unwrap();
+    assert_eq!(
+        body_json["at_step"].as_u64(),
+        Some(1),
+        "ReuseContext path: at_step falls back to 1 when timeline has no fork_at_step (got: {body_json})"
+    );
+}
+
 // ── Review #154 F3: Shape B cursor + TTL validation ──
 
 #[tokio::test]
@@ -622,6 +738,13 @@ async fn shape_b_rejects_expired_context() {
 // Helper: session with one recorded step so Shape A's fork(at_step=1)
 // has something to fork from.
 fn seed_session_with_step(store: &Arc<Mutex<Store>>) -> (String, String) {
+    seed_session_with_n_steps(store, 1)
+}
+
+/// Seed a session + root timeline + N owned LLM-call steps. Returns
+/// `(session_id, timeline_id)`. Used by tests that need a session
+/// with enough steps to fork at a meaningful at_step value.
+fn seed_session_with_n_steps(store: &Arc<Mutex<Store>>, n: u32) -> (String, String) {
     use rewind_store::{Session, SessionSource, SessionStatus, Step, StepStatus, Timeline};
     let s = store.lock().unwrap();
     let session = Session {
@@ -641,11 +764,13 @@ fn seed_session_with_step(store: &Arc<Mutex<Store>>) -> (String, String) {
     let timeline = Timeline::new_root(&session.id);
     s.create_session(&session).unwrap();
     s.create_timeline(&timeline).unwrap();
-    let mut step = Step::new_llm_call(&timeline.id, &session.id, 1, "stub-model");
-    step.status = StepStatus::Success;
-    step.duration_ms = 10;
-    s.create_step(&step).unwrap();
-    s.update_session_stats(&session.id, 1, 0).unwrap();
+    for i in 1..=n {
+        let mut step = Step::new_llm_call(&timeline.id, &session.id, i, "stub-model");
+        step.status = StepStatus::Success;
+        step.duration_ms = 10;
+        s.create_step(&step).unwrap();
+    }
+    s.update_session_stats(&session.id, n, 0).unwrap();
     (session.id, timeline.id)
 }