enact_core/workflow/

story_loop.rs

//! Story Loop - Iterate over planned stories with per-item verification
//!
//! Implements Antfarm's story loop pattern:
//! - Iterate over planned stories
//! - Optional fresh-session per story
//! - Per-story verification
//! - Independent retry counters

use crate::workflow::contract::{ContractParser, ParsedOutput, StepStatus};
use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Story definition
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Story {
    /// Story identifier
    pub id: String,
    /// Story title
    pub title: String,
    /// Story description
    pub description: String,
    /// Acceptance criteria
    #[serde(default)]
    pub acceptance_criteria: Vec<String>,
    /// Test criteria
    #[serde(default)]
    pub test_criteria: Vec<String>,
    /// Dependencies (story IDs that must be completed first)
    #[serde(default)]
    pub depends_on: Vec<String>,
    /// Estimated effort
    #[serde(default)]
    pub effort: Option<String>,
    /// Whether this story is completed
    #[serde(default)]
    pub completed: bool,
    /// Number of retry attempts
    #[serde(default)]
    pub retry_count: u32,
    /// Verification result
    #[serde(default)]
    pub verified: Option<bool>,
    /// Feedback from verifier
    #[serde(default)]
    pub verify_feedback: Option<String>,
}

/// Story loop configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoryLoopConfig {
    /// Field containing the stories array (e.g., "plan.STORIES_JSON")
    pub over: String,
    /// Completion condition
    #[serde(default)]
    pub completion: CompletionCondition,
    /// Whether to use fresh session for each story
    #[serde(default)]
    pub fresh_session: bool,
    /// Whether to verify each story
    #[serde(default)]
    pub verify_each: bool,
    /// Verification step ID (if verify_each is true)
    #[serde(default)]
    pub verify_step: Option<String>,
}

/// Completion conditions for story loop
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum CompletionCondition {
    /// All stories must be done
    #[default]
    AllDone,
    /// At least one story done
    AtLeastOne,
    /// Specific number of stories
    Count(usize),
    /// Percentage of stories
    Percentage(f32),
}

/// Story loop state
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoryLoopState {
    /// All stories
    pub stories: Vec<Story>,
    /// Current story index
    pub current_index: usize,
    /// Completed story IDs
    pub completed_ids: Vec<String>,
    /// Stories pending verification
    pub pending_verification: Vec<String>,
    /// Loop iteration count (safety)
    pub iteration_count: usize,
    /// Maximum iterations (prevent infinite loops)
    pub max_iterations: usize,
}

impl StoryLoopState {
    /// Create new story loop state
    pub fn new(stories: Vec<Story>) -> Self {
        Self {
            max_iterations: stories.len() * 3, // Allow 3 attempts per story
            stories,
            current_index: 0,
            completed_ids: vec![],
            pending_verification: vec![],
            iteration_count: 0,
        }
    }

    /// Get current story
    pub fn current_story(&self) -> Option<&Story> {
        self.stories.get(self.current_index)
    }

    /// Get current story (mutable)
    pub fn current_story_mut(&mut self) -> Option<&mut Story> {
        self.stories.get_mut(self.current_index)
    }

    /// Mark current story as completed
    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(())
    }

    /// Mark current story for retry with feedback
    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(())
    }

    /// Move to next story
    pub fn next_story(&mut self) -> bool {
        self.iteration_count += 1;

        // Find next incomplete story
        for i in (self.current_index + 1)..self.stories.len() {
            if !self.stories[i].completed {
                self.current_index = i;
                return true;
            }
        }

        // Check if we need to retry any stories
        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
    }

    /// Check if completion condition is met
    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)
                }
            }
        }
    }

    /// Check if we've exceeded max iterations
    pub fn should_stop(&self) -> bool {
        self.iteration_count >= self.max_iterations
    }

    /// Get completion percentage
    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
        }
    }

    /// Serialize stories to JSON
    pub fn stories_json(&self) -> Result<String> {
        serde_json::to_string(&self.stories).context("Failed to serialize stories")
    }

    /// Get context variables for current iteration
    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
    }
}

/// Story loop executor
pub struct StoryLoopExecutor;

impl StoryLoopExecutor {
    /// Execute a story loop
    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 {
            // Check stopping conditions
            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,
                });
            }

            // Check if there's a current story
            if state.current_story().is_none() {
                return Ok(StoryLoopResult {
                    success: false,
                    state,
                    reason: Some("No more stories to process".to_string()),
                });
            }

            // Execute the step
            let output = step_fn(state.clone()).await?;

            // Handle the result
            match output.status {
                StepStatus::Done => {
                    state.mark_completed()?;

                    // If verification is enabled, add to pending
                    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());

                    // Get story ID before mutable borrow
                    let current_story_id = state.current_story().map(|s| s.id.clone());

                    state.mark_retry(&feedback)?;

                    // Check if max retries reached
                    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)),
                    });
                }
            }
        }
    }
}

/// Story loop execution result
#[derive(Debug, Clone)]
pub struct StoryLoopResult {
    /// Whether the loop completed successfully
    pub success: bool,
    /// Final state
    pub state: StoryLoopState,
    /// Reason for failure (if any)
    pub reason: Option<String>,
}

/// Parse stories from JSON 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);

        // First story
        assert_eq!(state.current_story().unwrap().id, "story-1");

        // Mark first as complete, move to second
        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);

        // Not complete with 0/2 done
        assert!(!state.is_complete(&CompletionCondition::AllDone));
        assert!(!state.is_complete(&CompletionCondition::Count(1)));

        // Complete one
        state.mark_completed().unwrap();
        assert!(state.is_complete(&CompletionCondition::Count(1)));
        assert!(!state.is_complete(&CompletionCondition::AllDone));

        // Complete second
        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");
    }
}