ralph_workflow/reducer/state/agent_chain/
transitions.rs

1// State transition methods for AgentChainState.
2//
3// These methods implement the fallback chain progression: advancing models,
4// switching agents, and starting retry cycles with backoff.
5
6use std::sync::Arc;
7
8use super::backoff::calculate_backoff_delay_ms;
9use super::{AgentChainState, AgentDrain, AgentRole, DrainMode, RateLimitContinuationPrompt};
10
11impl AgentChainState {
12    #[must_use]
13    pub fn advance_to_next_model(&self) -> Self {
14        let start_agent_index = self.current_agent_index;
15
16        // When models are configured, we try each model for the current agent once.
17        // If the models list is exhausted, advance to the next agent/retry cycle
18        // instead of looping models indefinitely.
19        let mut next = match self.models_per_agent.get(self.current_agent_index) {
20            Some(models) if !models.is_empty() => {
21                if self.current_model_index + 1 < models.len() {
22                    // Simple model advance - only increment model index
23                    Self {
24                        agents: Arc::clone(&self.agents),
25                        current_agent_index: self.current_agent_index,
26                        models_per_agent: Arc::clone(&self.models_per_agent),
27                        current_model_index: self.current_model_index + 1,
28                        retry_cycle: self.retry_cycle,
29                        max_cycles: self.max_cycles,
30                        retry_delay_ms: self.retry_delay_ms,
31                        backoff_multiplier: self.backoff_multiplier,
32                        max_backoff_ms: self.max_backoff_ms,
33                        backoff_pending_ms: self.backoff_pending_ms,
34                        current_role: self.current_role,
35                        current_drain: self.current_drain,
36                        current_mode: self.current_mode,
37                        rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
38                        last_session_id: self.last_session_id.clone(),
39                    }
40                } else {
41                    self.switch_to_next_agent()
42                }
43            }
44            _ => self.switch_to_next_agent(),
45        };
46
47        if next.current_agent_index != start_agent_index {
48            next.last_session_id = None;
49        }
50
51        next
52    }
53
54    #[must_use]
55    pub fn switch_to_next_agent(&self) -> Self {
56        if self.current_agent_index + 1 < self.agents.len() {
57            // Advance to next agent
58            Self {
59                agents: Arc::clone(&self.agents),
60                current_agent_index: self.current_agent_index + 1,
61                models_per_agent: Arc::clone(&self.models_per_agent),
62                current_model_index: 0,
63                retry_cycle: self.retry_cycle,
64                max_cycles: self.max_cycles,
65                retry_delay_ms: self.retry_delay_ms,
66                backoff_multiplier: self.backoff_multiplier,
67                max_backoff_ms: self.max_backoff_ms,
68                backoff_pending_ms: None,
69                current_role: self.current_role,
70                current_drain: self.current_drain,
71                current_mode: self.current_mode,
72                rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
73                last_session_id: self.last_session_id.clone(),
74            }
75        } else {
76            // Wrap around to first agent and increment retry cycle
77            let new_retry_cycle = self.retry_cycle + 1;
78            let new_backoff_pending_ms = if new_retry_cycle >= self.max_cycles {
79                None
80            } else {
81                // Create temporary state to calculate backoff
82                let temp = Self {
83                    agents: Arc::clone(&self.agents),
84                    current_agent_index: 0,
85                    models_per_agent: Arc::clone(&self.models_per_agent),
86                    current_model_index: 0,
87                    retry_cycle: new_retry_cycle,
88                    max_cycles: self.max_cycles,
89                    retry_delay_ms: self.retry_delay_ms,
90                    backoff_multiplier: self.backoff_multiplier,
91                    max_backoff_ms: self.max_backoff_ms,
92                    backoff_pending_ms: None,
93                    current_role: self.current_role,
94                    current_drain: self.current_drain,
95                    current_mode: self.current_mode,
96                    rate_limit_continuation_prompt: None,
97                    last_session_id: None,
98                };
99                Some(temp.calculate_backoff_delay_ms_for_retry_cycle())
100            };
101
102            Self {
103                agents: Arc::clone(&self.agents),
104                current_agent_index: 0,
105                models_per_agent: Arc::clone(&self.models_per_agent),
106                current_model_index: 0,
107                retry_cycle: new_retry_cycle,
108                max_cycles: self.max_cycles,
109                retry_delay_ms: self.retry_delay_ms,
110                backoff_multiplier: self.backoff_multiplier,
111                max_backoff_ms: self.max_backoff_ms,
112                backoff_pending_ms: new_backoff_pending_ms,
113                current_role: self.current_role,
114                current_drain: self.current_drain,
115                current_mode: self.current_mode,
116                rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
117                last_session_id: self.last_session_id.clone(),
118            }
119        }
120    }
121
122    /// Switch to a specific agent by name.
123    ///
124    /// If `to_agent` is unknown, falls back to `switch_to_next_agent()` to keep the
125    /// reducer deterministic.
126    #[must_use]
127    pub fn switch_to_agent_named(&self, to_agent: &str) -> Self {
128        let Some(target_index) = self.agents.iter().position(|a| a == to_agent) else {
129            return self.switch_to_next_agent();
130        };
131
132        if target_index == self.current_agent_index {
133            // Same agent - just reset model index
134            return Self {
135                agents: Arc::clone(&self.agents),
136                current_agent_index: self.current_agent_index,
137                models_per_agent: Arc::clone(&self.models_per_agent),
138                current_model_index: 0,
139                retry_cycle: self.retry_cycle,
140                max_cycles: self.max_cycles,
141                retry_delay_ms: self.retry_delay_ms,
142                backoff_multiplier: self.backoff_multiplier,
143                max_backoff_ms: self.max_backoff_ms,
144                backoff_pending_ms: None,
145                current_role: self.current_role,
146                current_drain: self.current_drain,
147                current_mode: self.current_mode,
148                rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
149                last_session_id: self.last_session_id.clone(),
150            };
151        }
152
153        if target_index <= self.current_agent_index {
154            // Treat switching to an earlier agent as starting a new retry cycle.
155            let new_retry_cycle = self.retry_cycle + 1;
156            let new_backoff_pending_ms = if new_retry_cycle >= self.max_cycles && target_index == 0
157            {
158                None
159            } else {
160                // Create temporary state to calculate backoff
161                let temp = Self {
162                    agents: Arc::clone(&self.agents),
163                    current_agent_index: target_index,
164                    models_per_agent: Arc::clone(&self.models_per_agent),
165                    current_model_index: 0,
166                    retry_cycle: new_retry_cycle,
167                    max_cycles: self.max_cycles,
168                    retry_delay_ms: self.retry_delay_ms,
169                    backoff_multiplier: self.backoff_multiplier,
170                    max_backoff_ms: self.max_backoff_ms,
171                    backoff_pending_ms: None,
172                    current_role: self.current_role,
173                    current_drain: self.current_drain,
174                    current_mode: self.current_mode,
175                    rate_limit_continuation_prompt: None,
176                    last_session_id: None,
177                };
178                Some(temp.calculate_backoff_delay_ms_for_retry_cycle())
179            };
180
181            Self {
182                agents: Arc::clone(&self.agents),
183                current_agent_index: target_index,
184                models_per_agent: Arc::clone(&self.models_per_agent),
185                current_model_index: 0,
186                retry_cycle: new_retry_cycle,
187                max_cycles: self.max_cycles,
188                retry_delay_ms: self.retry_delay_ms,
189                backoff_multiplier: self.backoff_multiplier,
190                max_backoff_ms: self.max_backoff_ms,
191                backoff_pending_ms: new_backoff_pending_ms,
192                current_role: self.current_role,
193                current_drain: self.current_drain,
194                current_mode: self.current_mode,
195                rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
196                last_session_id: self.last_session_id.clone(),
197            }
198        } else {
199            // Advancing to later agent
200            Self {
201                agents: Arc::clone(&self.agents),
202                current_agent_index: target_index,
203                models_per_agent: Arc::clone(&self.models_per_agent),
204                current_model_index: 0,
205                retry_cycle: self.retry_cycle,
206                max_cycles: self.max_cycles,
207                retry_delay_ms: self.retry_delay_ms,
208                backoff_multiplier: self.backoff_multiplier,
209                max_backoff_ms: self.max_backoff_ms,
210                backoff_pending_ms: None,
211                current_role: self.current_role,
212                current_drain: self.current_drain,
213                current_mode: self.current_mode,
214                rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
215                last_session_id: self.last_session_id.clone(),
216            }
217        }
218    }
219
220    /// Switch to next agent after rate limit, preserving prompt for continuation.
221    ///
222    /// This is used when an agent hits a 429 rate limit error. Instead of
223    /// retrying with the same agent (which would likely hit rate limits again),
224    /// we switch to the next agent and preserve the prompt so the new agent
225    /// can continue the same work.
226    #[must_use]
227    pub fn switch_to_next_agent_with_prompt(&self, prompt: Option<String>) -> Self {
228        let base = self.switch_to_next_agent();
229        // Back-compat: older callers didn't track role. Preserve prompt only.
230        Self {
231            agents: base.agents,
232            current_agent_index: base.current_agent_index,
233            models_per_agent: base.models_per_agent,
234            current_model_index: base.current_model_index,
235            retry_cycle: base.retry_cycle,
236            max_cycles: base.max_cycles,
237            retry_delay_ms: base.retry_delay_ms,
238            backoff_multiplier: base.backoff_multiplier,
239            max_backoff_ms: base.max_backoff_ms,
240            backoff_pending_ms: base.backoff_pending_ms,
241            current_role: base.current_role,
242            current_drain: base.current_drain,
243            current_mode: base.current_mode,
244            rate_limit_continuation_prompt: prompt.map(|p| RateLimitContinuationPrompt {
245                drain: base.current_drain,
246                role: base.current_role,
247                prompt: p,
248            }),
249            last_session_id: base.last_session_id,
250        }
251    }
252
253    /// Switch to next agent after rate limit, preserving prompt for continuation (role-scoped).
254    #[must_use]
255    pub fn switch_to_next_agent_with_prompt_for_role(
256        &self,
257        role: AgentRole,
258        prompt: Option<String>,
259    ) -> Self {
260        let base = self.switch_to_next_agent();
261        Self {
262            agents: base.agents,
263            current_agent_index: base.current_agent_index,
264            models_per_agent: base.models_per_agent,
265            current_model_index: base.current_model_index,
266            retry_cycle: base.retry_cycle,
267            max_cycles: base.max_cycles,
268            retry_delay_ms: base.retry_delay_ms,
269            backoff_multiplier: base.backoff_multiplier,
270            max_backoff_ms: base.max_backoff_ms,
271            backoff_pending_ms: base.backoff_pending_ms,
272            current_role: base.current_role,
273            current_drain: base.current_drain,
274            current_mode: base.current_mode,
275            rate_limit_continuation_prompt: prompt.map(|p| RateLimitContinuationPrompt {
276                drain: base.current_drain,
277                role,
278                prompt: p,
279            }),
280            last_session_id: base.last_session_id,
281        }
282    }
283
284    /// Clear continuation prompt after successful execution.
285    ///
286    /// Called when an agent successfully completes its task, clearing any
287    /// saved prompt context from previous rate-limited agents.
288    #[must_use]
289    pub fn clear_continuation_prompt(&self) -> Self {
290        Self {
291            agents: Arc::clone(&self.agents),
292            current_agent_index: self.current_agent_index,
293            models_per_agent: Arc::clone(&self.models_per_agent),
294            current_model_index: self.current_model_index,
295            retry_cycle: self.retry_cycle,
296            max_cycles: self.max_cycles,
297            retry_delay_ms: self.retry_delay_ms,
298            backoff_multiplier: self.backoff_multiplier,
299            max_backoff_ms: self.max_backoff_ms,
300            backoff_pending_ms: self.backoff_pending_ms,
301            current_role: self.current_role,
302            current_drain: self.current_drain,
303            current_mode: self.current_mode,
304            rate_limit_continuation_prompt: None,
305            last_session_id: self.last_session_id.clone(),
306        }
307    }
308
309    #[must_use]
310    pub fn reset_for_drain(&self, drain: AgentDrain) -> Self {
311        Self {
312            agents: Arc::clone(&self.agents),
313            current_agent_index: 0,
314            models_per_agent: Arc::clone(&self.models_per_agent),
315            current_model_index: 0,
316            retry_cycle: 0,
317            max_cycles: self.max_cycles,
318            retry_delay_ms: self.retry_delay_ms,
319            backoff_multiplier: self.backoff_multiplier,
320            max_backoff_ms: self.max_backoff_ms,
321            backoff_pending_ms: None,
322            current_role: drain.role(),
323            current_drain: drain,
324            current_mode: DrainMode::Normal,
325            rate_limit_continuation_prompt: None,
326            last_session_id: None,
327        }
328    }
329
330    #[must_use]
331    pub fn reset_for_role(&self, role: AgentRole) -> Self {
332        self.reset_for_drain(match role {
333            AgentRole::Developer => AgentDrain::Development,
334            AgentRole::Reviewer => AgentDrain::Review,
335            AgentRole::Commit => AgentDrain::Commit,
336            AgentRole::Analysis => AgentDrain::Analysis,
337        })
338    }
339
340    #[must_use]
341    pub fn reset(&self) -> Self {
342        Self {
343            agents: Arc::clone(&self.agents),
344            current_agent_index: 0,
345            models_per_agent: Arc::clone(&self.models_per_agent),
346            current_model_index: 0,
347            retry_cycle: self.retry_cycle,
348            max_cycles: self.max_cycles,
349            retry_delay_ms: self.retry_delay_ms,
350            backoff_multiplier: self.backoff_multiplier,
351            max_backoff_ms: self.max_backoff_ms,
352            backoff_pending_ms: None,
353            current_role: self.current_role,
354            current_drain: self.current_drain,
355            current_mode: DrainMode::Normal,
356            rate_limit_continuation_prompt: None,
357            last_session_id: None,
358        }
359    }
360
361    /// Store session ID from agent response for potential reuse.
362    #[must_use]
363    pub fn with_session_id(&self, session_id: Option<String>) -> Self {
364        Self {
365            agents: Arc::clone(&self.agents),
366            current_agent_index: self.current_agent_index,
367            models_per_agent: Arc::clone(&self.models_per_agent),
368            current_model_index: self.current_model_index,
369            retry_cycle: self.retry_cycle,
370            max_cycles: self.max_cycles,
371            retry_delay_ms: self.retry_delay_ms,
372            backoff_multiplier: self.backoff_multiplier,
373            max_backoff_ms: self.max_backoff_ms,
374            backoff_pending_ms: self.backoff_pending_ms,
375            current_role: self.current_role,
376            current_drain: self.current_drain,
377            current_mode: self.current_mode,
378            rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
379            last_session_id: session_id,
380        }
381    }
382
383    /// Clear session ID (e.g., when switching agents or starting new work).
384    #[must_use]
385    pub fn clear_session_id(&self) -> Self {
386        Self {
387            agents: Arc::clone(&self.agents),
388            current_agent_index: self.current_agent_index,
389            models_per_agent: Arc::clone(&self.models_per_agent),
390            current_model_index: self.current_model_index,
391            retry_cycle: self.retry_cycle,
392            max_cycles: self.max_cycles,
393            retry_delay_ms: self.retry_delay_ms,
394            backoff_multiplier: self.backoff_multiplier,
395            max_backoff_ms: self.max_backoff_ms,
396            backoff_pending_ms: self.backoff_pending_ms,
397            current_role: self.current_role,
398            current_drain: self.current_drain,
399            current_mode: self.current_mode,
400            rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
401            last_session_id: None,
402        }
403    }
404
405    #[must_use]
406    pub fn start_retry_cycle(&self) -> Self {
407        let new_retry_cycle = self.retry_cycle + 1;
408        let new_backoff_pending_ms = if new_retry_cycle >= self.max_cycles {
409            None
410        } else {
411            // Create temporary state to calculate backoff
412            let temp = Self {
413                agents: Arc::clone(&self.agents),
414                current_agent_index: 0,
415                models_per_agent: Arc::clone(&self.models_per_agent),
416                current_model_index: 0,
417                retry_cycle: new_retry_cycle,
418                max_cycles: self.max_cycles,
419                retry_delay_ms: self.retry_delay_ms,
420                backoff_multiplier: self.backoff_multiplier,
421                max_backoff_ms: self.max_backoff_ms,
422                backoff_pending_ms: None,
423                current_role: self.current_role,
424                current_drain: self.current_drain,
425                current_mode: self.current_mode,
426                rate_limit_continuation_prompt: None,
427                last_session_id: None,
428            };
429            Some(temp.calculate_backoff_delay_ms_for_retry_cycle())
430        };
431
432        Self {
433            agents: Arc::clone(&self.agents),
434            current_agent_index: 0,
435            models_per_agent: Arc::clone(&self.models_per_agent),
436            current_model_index: 0,
437            retry_cycle: new_retry_cycle,
438            max_cycles: self.max_cycles,
439            retry_delay_ms: self.retry_delay_ms,
440            backoff_multiplier: self.backoff_multiplier,
441            max_backoff_ms: self.max_backoff_ms,
442            backoff_pending_ms: new_backoff_pending_ms,
443            current_role: self.current_role,
444            current_drain: self.current_drain,
445            current_mode: self.current_mode,
446            rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
447            last_session_id: self.last_session_id.clone(),
448        }
449    }
450
451    #[must_use]
452    pub fn clear_backoff_pending(&self) -> Self {
453        Self {
454            agents: Arc::clone(&self.agents),
455            current_agent_index: self.current_agent_index,
456            models_per_agent: Arc::clone(&self.models_per_agent),
457            current_model_index: self.current_model_index,
458            retry_cycle: self.retry_cycle,
459            max_cycles: self.max_cycles,
460            retry_delay_ms: self.retry_delay_ms,
461            backoff_multiplier: self.backoff_multiplier,
462            max_backoff_ms: self.max_backoff_ms,
463            backoff_pending_ms: None,
464            current_role: self.current_role,
465            current_drain: self.current_drain,
466            current_mode: self.current_mode,
467            rate_limit_continuation_prompt: self.rate_limit_continuation_prompt.clone(),
468            last_session_id: self.last_session_id.clone(),
469        }
470    }
471
472    pub(super) fn calculate_backoff_delay_ms_for_retry_cycle(&self) -> u64 {
473        // The first retry cycle should use the base delay.
474        let cycle_index = self.retry_cycle.saturating_sub(1);
475        calculate_backoff_delay_ms(
476            self.retry_delay_ms,
477            self.backoff_multiplier,
478            self.max_backoff_ms,
479            cycle_index,
480        )
481    }
482}
483
484#[cfg(test)]
485mod backoff_semantics_tests {
486    use super::*;
487
488    #[test]
489    fn test_switch_to_agent_named_preserves_backoff_when_retry_cycle_hits_max_but_state_is_not_exhausted(
490    ) {
491        let mut state = AgentChainState::initial().with_agents(
492            vec!["a".to_string(), "b".to_string(), "c".to_string()],
493            vec![vec![], vec![], vec![]],
494            AgentRole::Developer,
495        );
496        state.max_cycles = 2;
497        state.retry_cycle = 1;
498        state.current_agent_index = 2;
499
500        let next = state.switch_to_agent_named("b");
501
502        assert_eq!(next.current_agent_index, 1);
503        assert_eq!(next.retry_cycle, 2);
504        assert!(
505            next.backoff_pending_ms.is_some(),
506            "backoff should remain pending unless the state is fully exhausted"
507        );
508    }
509}
ralph_workflow/reducer/state/agent_chain/transitions.rs

ralph_workflow/reducer/state/agent_chain/
transitions.rs
