use crate::workflow::contract::{ContractParser, ParsedOutput, StepStatus};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Story {
pub id: String,
pub title: String,
pub description: String,
#[serde(default)]
pub acceptance_criteria: Vec<String>,
#[serde(default)]
pub test_criteria: Vec<String>,
#[serde(default)]
pub depends_on: Vec<String>,
#[serde(default)]
pub effort: Option<String>,
#[serde(default)]
pub completed: bool,
#[serde(default)]
pub retry_count: u32,
#[serde(default)]
pub verified: Option<bool>,
#[serde(default)]
pub verify_feedback: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoryLoopConfig {
pub over: String,
#[serde(default)]
pub completion: CompletionCondition,
#[serde(default)]
pub fresh_session: bool,
#[serde(default)]
pub verify_each: bool,
#[serde(default)]
pub verify_step: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum CompletionCondition {
#[default]
AllDone,
AtLeastOne,
Count(usize),
Percentage(f32),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoryLoopState {
pub stories: Vec<Story>,
pub current_index: usize,
pub completed_ids: Vec<String>,
pub pending_verification: Vec<String>,
pub iteration_count: usize,
pub max_iterations: usize,
}
impl StoryLoopState {
pub fn new(stories: Vec<Story>) -> Self {
Self {
max_iterations: stories.len() * 3, stories,
current_index: 0,
completed_ids: vec![],
pending_verification: vec![],
iteration_count: 0,
}
}
pub fn current_story(&self) -> Option<&Story> {
self.stories.get(self.current_index)
}
pub fn current_story_mut(&mut self) -> Option<&mut Story> {
self.stories.get_mut(self.current_index)
}
pub fn mark_completed(&mut self) -> Result<()> {
if let Some(story) = self.current_story() {
let id = story.id.clone();
if let Some(story) = self.stories.iter_mut().find(|s| s.id == id) {
story.completed = true;
if !self.completed_ids.contains(&id) {
self.completed_ids.push(id);
}
}
}
Ok(())
}
pub fn mark_retry(&mut self, feedback: &str) -> Result<()> {
if let Some(story) = self.current_story_mut() {
story.retry_count += 1;
story.verify_feedback = Some(feedback.to_string());
}
Ok(())
}
pub fn next_story(&mut self) -> bool {
self.iteration_count += 1;
for i in (self.current_index + 1)..self.stories.len() {
if !self.stories[i].completed {
self.current_index = i;
return true;
}
}
for i in 0..self.stories.len() {
if !self.stories[i].completed && self.stories[i].retry_count < 2 {
self.current_index = i;
return true;
}
}
false
}
pub fn is_complete(&self, condition: &CompletionCondition) -> bool {
let completed = self.completed_ids.len();
let total = self.stories.len();
match condition {
CompletionCondition::AllDone => completed >= total,
CompletionCondition::AtLeastOne => completed >= 1,
CompletionCondition::Count(n) => completed >= *n,
CompletionCondition::Percentage(p) => {
if total == 0 {
true
} else {
(completed as f32 / total as f32) >= (*p / 100.0)
}
}
}
}
pub fn should_stop(&self) -> bool {
self.iteration_count >= self.max_iterations
}
pub fn completion_percentage(&self) -> f32 {
if self.stories.is_empty() {
100.0
} else {
(self.completed_ids.len() as f32 / self.stories.len() as f32) * 100.0
}
}
pub fn stories_json(&self) -> Result<String> {
serde_json::to_string(&self.stories).context("Failed to serialize stories")
}
pub fn context_variables(&self) -> HashMap<String, String> {
let mut vars = HashMap::new();
vars.insert("stories_count".to_string(), self.stories.len().to_string());
vars.insert(
"completed_count".to_string(),
self.completed_ids.len().to_string(),
);
vars.insert(
"remaining_count".to_string(),
(self.stories.len() - self.completed_ids.len()).to_string(),
);
vars.insert(
"completion_percentage".to_string(),
format!("{:.1}", self.completion_percentage()),
);
if let Some(story) = self.current_story() {
vars.insert("current_story_id".to_string(), story.id.clone());
vars.insert("current_story_title".to_string(), story.title.clone());
vars.insert(
"current_story".to_string(),
serde_json::to_string(story).unwrap_or_default(),
);
vars.insert(
"current_story_description".to_string(),
story.description.clone(),
);
if let Some(feedback) = &story.verify_feedback {
vars.insert("verify_feedback".to_string(), feedback.clone());
}
}
vars.insert(
"completed_stories".to_string(),
serde_json::to_string(&self.completed_ids).unwrap_or_default(),
);
vars.insert(
"stories_json".to_string(),
self.stories_json().unwrap_or_default(),
);
vars
}
}
pub struct StoryLoopExecutor;
impl StoryLoopExecutor {
pub async fn execute_loop<F, Fut>(
config: &StoryLoopConfig,
stories: Vec<Story>,
mut step_fn: F,
) -> Result<StoryLoopResult>
where
F: FnMut(StoryLoopState) -> Fut,
Fut: std::future::Future<Output = Result<ParsedOutput>>,
{
let mut state = StoryLoopState::new(stories);
loop {
if state.should_stop() {
return Ok(StoryLoopResult {
success: false,
state,
reason: Some("Max iterations exceeded".to_string()),
});
}
if state.is_complete(&config.completion) {
return Ok(StoryLoopResult {
success: true,
state,
reason: None,
});
}
if state.current_story().is_none() {
return Ok(StoryLoopResult {
success: false,
state,
reason: Some("No more stories to process".to_string()),
});
}
let output = step_fn(state.clone()).await?;
match output.status {
StepStatus::Done => {
state.mark_completed()?;
if config.verify_each {
if let Some(story) = state.current_story() {
state.pending_verification.push(story.id.clone());
}
}
state.next_story();
}
StepStatus::Retry => {
let feedback = ContractParser::get_feedback(&output.raw_output, "ISSUES")
.unwrap_or_else(|| "Retry requested".to_string());
let current_story_id = state.current_story().map(|s| s.id.clone());
state.mark_retry(&feedback)?;
if let Some(story_id) = current_story_id {
if let Some(story) = state.stories.iter().find(|s| s.id == story_id) {
if story.retry_count >= 2 {
return Ok(StoryLoopResult {
success: false,
state,
reason: Some(format!(
"Max retries reached for story: {}",
story_id
)),
});
}
}
}
}
StepStatus::Blocked => {
let blocked_story_id = state
.current_story()
.map(|s| s.id.clone())
.unwrap_or_default();
return Ok(StoryLoopResult {
success: false,
state,
reason: Some(format!("Story blocked: {}", blocked_story_id)),
});
}
}
}
}
}
#[derive(Debug, Clone)]
pub struct StoryLoopResult {
pub success: bool,
pub state: StoryLoopState,
pub reason: Option<String>,
}
pub fn parse_stories(json_str: &str) -> Result<Vec<Story>> {
serde_json::from_str(json_str).context("Failed to parse stories JSON")
}
#[cfg(test)]
mod tests {
use super::*;
fn create_test_stories() -> Vec<Story> {
vec![
Story {
id: "story-1".to_string(),
title: "First Story".to_string(),
description: "Implement feature A".to_string(),
acceptance_criteria: vec!["Feature A works".to_string()],
test_criteria: vec!["Test A passes".to_string()],
depends_on: vec![],
effort: Some("small".to_string()),
completed: false,
retry_count: 0,
verified: None,
verify_feedback: None,
},
Story {
id: "story-2".to_string(),
title: "Second Story".to_string(),
description: "Implement feature B".to_string(),
acceptance_criteria: vec!["Feature B works".to_string()],
test_criteria: vec!["Test B passes".to_string()],
depends_on: vec!["story-1".to_string()],
effort: Some("medium".to_string()),
completed: false,
retry_count: 0,
verified: None,
verify_feedback: None,
},
]
}
#[test]
fn test_story_loop_state() {
let stories = create_test_stories();
let state = StoryLoopState::new(stories);
assert_eq!(state.stories.len(), 2);
assert_eq!(state.current_index, 0);
assert!(state.current_story().is_some());
assert_eq!(state.current_story().unwrap().id, "story-1");
}
#[test]
fn test_mark_completed() {
let stories = create_test_stories();
let mut state = StoryLoopState::new(stories);
state.mark_completed().unwrap();
assert_eq!(state.completed_ids.len(), 1);
assert!(state.stories[0].completed);
}
#[test]
fn test_next_story() {
let stories = create_test_stories();
let mut state = StoryLoopState::new(stories);
assert_eq!(state.current_story().unwrap().id, "story-1");
state.mark_completed().unwrap();
assert!(state.next_story());
assert_eq!(state.current_story().unwrap().id, "story-2");
}
#[test]
fn test_completion_conditions() {
let stories = create_test_stories();
let mut state = StoryLoopState::new(stories);
assert!(!state.is_complete(&CompletionCondition::AllDone));
assert!(!state.is_complete(&CompletionCondition::Count(1)));
state.mark_completed().unwrap();
assert!(state.is_complete(&CompletionCondition::Count(1)));
assert!(!state.is_complete(&CompletionCondition::AllDone));
state.next_story();
state.mark_completed().unwrap();
assert!(state.is_complete(&CompletionCondition::AllDone));
}
#[test]
fn test_context_variables() {
let stories = create_test_stories();
let state = StoryLoopState::new(stories);
let vars = state.context_variables();
assert_eq!(vars.get("stories_count").unwrap(), "2");
assert_eq!(vars.get("completed_count").unwrap(), "0");
assert!(vars.contains_key("current_story"));
assert!(vars.contains_key("stories_json"));
}
#[test]
fn test_parse_stories() {
let json = r#"[
{
"id": "story-1",
"title": "Test Story",
"description": "A test story",
"acceptance_criteria": ["It works"],
"completed": false
}
]"#;
let stories = parse_stories(json).unwrap();
assert_eq!(stories.len(), 1);
assert_eq!(stories[0].id, "story-1");
}
}