vtcode 0.140.2

use anstyle::Color;
use serde_json::Value;
use std::fmt::Write;
use vtcode_core::tools::registry::ToolExecutionError;
use vtcode_core::utils::style_helpers::ColorPalette;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ToolDisplayStatus {
    Success,
    Failure,
    Warning,
}

impl ToolDisplayStatus {
    pub(crate) fn from_command_output(output: &Value, command_success: bool) -> Self {
        if !command_success {
            Self::Failure
        } else if output.get("warning").is_some_and(has_warning_payload) {
            Self::Warning
        } else {
            Self::Success
        }
    }

    pub(crate) fn color(self, palette: ColorPalette) -> Color {
        match self {
            Self::Success => palette.success,
            Self::Failure => palette.error,
            Self::Warning => palette.warning,
        }
    }
}

fn has_warning_payload(warning: &Value) -> bool {
    match warning {
        Value::Bool(enabled) => *enabled,
        Value::String(message) => !message.trim().is_empty(),
        Value::Array(items) => !items.is_empty(),
        Value::Object(fields) => !fields.is_empty(),
        Value::Null => false,
        // Numeric warning fields are commonly counts. Treat zero like an
        // empty payload so a harmless `warning: 0` does not turn a successful
        // tool call yellow.
        Value::Number(value) => value.as_f64().is_some_and(|number| number != 0.0),
    }
}

/// Result of a tool execution
#[derive(Debug)]
pub(crate) enum ToolExecutionStatus {
    /// Tool completed
    Success {
        /// Tool output
        output: Value,
        /// Standard output if available
        stdout: Option<String>,
        /// List of modified files
        modified_files: Vec<String>,
        /// Whether the command was successful
        command_success: bool,
    },
    /// Tool execution failed
    Failure {
        /// Error that occurred
        error: ToolExecutionError,
    },
    /// Tool execution timed out
    Timeout {
        /// Timeout error
        error: ToolExecutionError,
    },
    /// Tool execution was cancelled
    Cancelled,
    // TODO: Progress variant planned for streaming tool progress updates
}

impl ToolExecutionStatus {
    pub(crate) fn display_status(&self) -> ToolDisplayStatus {
        match self {
            Self::Success { output, command_success, .. } => {
                ToolDisplayStatus::from_command_output(output, *command_success)
            }
            Self::Failure { .. } | Self::Timeout { .. } => ToolDisplayStatus::Failure,
            Self::Cancelled => ToolDisplayStatus::Warning,
        }
    }

    pub(crate) fn error(&self) -> Option<&ToolExecutionError> {
        match self {
            ToolExecutionStatus::Failure { error } | ToolExecutionStatus::Timeout { error } => Some(error),
            ToolExecutionStatus::Success { .. } | ToolExecutionStatus::Cancelled => None,
        }
    }
}

/// Outcome produced by a tool pipeline run - returns a success/failure wrapper along with stdout and modified files
pub(crate) struct ToolPipelineOutcome {
    pub status: ToolExecutionStatus,
    pub command_success: bool,
    /// When set, the interaction loop should switch the active primary agent to
    /// this name after the tool completes. Used by the plan-mode "switch to
    /// build/auto agent" decision at the plan-ready gate.
    pub pending_primary_agent: Option<String>,
    /// When true, finish the current turn after recording this tool result.
    /// Used when a headless policy requires user confirmation that cannot be
    /// collected through an inline prompt.
    pub stop_after_tool: bool,
}

impl ToolPipelineOutcome {
    pub(crate) fn from_status(status: ToolExecutionStatus) -> Self {
        let command_success = match &status {
            ToolExecutionStatus::Success { command_success, .. } => *command_success,
            _ => false,
        };
        ToolPipelineOutcome {
            status,
            command_success,
            pending_primary_agent: None,
            stop_after_tool: false,
        }
    }

    pub(crate) fn modified_files(&self) -> &[String] {
        match &self.status {
            ToolExecutionStatus::Success { modified_files, .. } => modified_files,
            _ => &[],
        }
    }

    pub(crate) fn modified_files_mut(&mut self) -> Option<&mut Vec<String>> {
        match &mut self.status {
            ToolExecutionStatus::Success { modified_files, .. } => Some(modified_files),
            _ => None,
        }
    }

    pub(crate) fn set_command_success(&mut self, command_success: bool) {
        self.command_success = command_success;
        if let ToolExecutionStatus::Success { command_success: status_success, .. } = &mut self.status {
            *status_success = command_success;
        }
    }
}

// ---------------------------------------------------------------------------
// Batch outcome tracking
// ---------------------------------------------------------------------------

/// Summary of a batch of tool executions (one LLM turn may request multiple
/// tool calls). Tracks per-tool outcomes so the agent loop can detect partial
/// failures and inject recovery context.
#[derive(Debug)]
pub(crate) struct ToolBatchOutcome {
    /// Per-tool results, ordered to match the original batch request order.
    pub entries: Vec<ToolBatchEntry>,
}

/// A single entry inside a [`ToolBatchOutcome`].
#[derive(Debug)]
pub(crate) struct ToolBatchEntry {
    /// High-level result category.
    pub result: ToolBatchResult,
}

/// Simplified result for aggregation — avoids carrying the full
/// [`ToolExecutionStatus`] payload.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum ToolBatchResult {
    Success,
    Failure,
    Timeout,
    Cancelled,
}

impl From<&ToolExecutionStatus> for ToolBatchResult {
    fn from(status: &ToolExecutionStatus) -> Self {
        match status {
            ToolExecutionStatus::Success { .. } => Self::Success,
            ToolExecutionStatus::Failure { .. } => Self::Failure,
            ToolExecutionStatus::Timeout { .. } => Self::Timeout,
            ToolExecutionStatus::Cancelled => Self::Cancelled,
        }
    }
}

/// Aggregate statistics for a [`ToolBatchOutcome`].
#[derive(Debug, Clone, Copy)]
pub(crate) struct ToolBatchStats {
    pub total: usize,
    pub succeeded: usize,
    pub failed: usize,
    pub timed_out: usize,
    pub cancelled: usize,
}

impl ToolBatchOutcome {
    /// Create a new empty batch.
    pub(crate) fn new() -> Self {
        Self { entries: Vec::new() }
    }

    /// Record the result of one tool call.
    pub(crate) fn record(&mut self, status: &ToolExecutionStatus) {
        self.entries.push(ToolBatchEntry { result: ToolBatchResult::from(status) });
    }

    /// Compute aggregate statistics.
    pub(crate) fn stats(&self) -> ToolBatchStats {
        let mut s = ToolBatchStats {
            total: self.entries.len(),
            succeeded: 0,
            failed: 0,
            timed_out: 0,
            cancelled: 0,
        };
        for entry in &self.entries {
            match entry.result {
                ToolBatchResult::Success => s.succeeded += 1,
                ToolBatchResult::Failure => s.failed += 1,
                ToolBatchResult::Timeout => s.timed_out += 1,
                ToolBatchResult::Cancelled => s.cancelled += 1,
            }
        }
        s
    }

    /// Returns `true` when **some** tools succeeded but **some** failed.
    pub(crate) fn is_partial_success(&self) -> bool {
        let s = self.stats();
        // Compare against the total so every non-success outcome (including
        // cancellation and future result variants) is handled consistently.
        s.succeeded > 0 && s.succeeded < s.total
    }

    /// Build a compact one-line summary suitable for structured logging.
    pub(crate) fn summary_line(&self) -> String {
        let s = self.stats();
        if s.total == 0 {
            return "no tools executed".to_string();
        }
        if s.succeeded == s.total {
            return format!("all {} tools succeeded", s.total);
        }
        let mut buf = String::with_capacity(64);
        let mut first = true;
        if s.succeeded > 0 {
            if !first {
                buf.push_str(", ");
            }
            let _ = write!(buf, "{} succeeded", s.succeeded);
            first = false;
        }
        if s.failed > 0 {
            if !first {
                buf.push_str(", ");
            }
            let _ = write!(buf, "{} failed", s.failed);
            first = false;
        }
        if s.timed_out > 0 {
            if !first {
                buf.push_str(", ");
            }
            let _ = write!(buf, "{} timed out", s.timed_out);
            first = false;
        }
        if s.cancelled > 0 {
            if !first {
                buf.push_str(", ");
            }
            let _ = write!(buf, "{} cancelled", s.cancelled);
        }
        let _ = write!(buf, " ({}/{} tools)", s.total, s.total);
        buf
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty_batch_stats() {
        let batch = ToolBatchOutcome::new();
        let s = batch.stats();
        assert_eq!(s.total, 0);
        assert!(!batch.is_partial_success());
        assert_eq!(batch.summary_line(), "no tools executed");
    }

    #[test]
    fn tool_pipeline_outcomes_continue_by_default() {
        let outcome = ToolPipelineOutcome::from_status(ToolExecutionStatus::Cancelled);

        assert!(!outcome.stop_after_tool);
        assert!(outcome.pending_primary_agent.is_none());
    }

    #[test]
    fn all_success_batch() {
        let mut batch = ToolBatchOutcome::new();
        let success = ToolExecutionStatus::Success {
            output: serde_json::json!("ok"),
            stdout: None,
            modified_files: vec![],
            command_success: true,
        };
        batch.record(&success);
        batch.record(&success);

        let s = batch.stats();
        assert_eq!(s.total, 2);
        assert_eq!(s.succeeded, 2);
        assert!(!batch.is_partial_success());
        assert!(batch.summary_line().contains("all 2 tools succeeded"));
    }

    #[test]
    fn partial_failure_batch() {
        let mut batch = ToolBatchOutcome::new();
        let success = ToolExecutionStatus::Success {
            output: serde_json::json!("ok"),
            stdout: None,
            modified_files: vec![],
            command_success: true,
        };
        let failure = ToolExecutionStatus::Failure {
            error: ToolExecutionError::new(
                "tool".to_string(),
                vtcode_core::tools::registry::ToolErrorType::PermissionDenied,
                "permission denied".to_string(),
            ),
        };
        batch.record(&success);
        batch.record(&failure);

        assert!(batch.is_partial_success());
        let s = batch.stats();
        assert_eq!(s.succeeded, 1);
        assert_eq!(s.failed, 1);
        assert!(batch.summary_line().contains("1 succeeded"));
        assert!(batch.summary_line().contains("1 failed"));
    }

    #[test]
    fn all_failure_not_partial_success() {
        let mut batch = ToolBatchOutcome::new();
        let failure = ToolExecutionStatus::Failure {
            error: ToolExecutionError::new(
                "tool".to_string(),
                vtcode_core::tools::registry::ToolErrorType::ExecutionError,
                "boom".to_string(),
            ),
        };
        batch.record(&failure);
        assert!(!batch.is_partial_success());
    }

    #[test]
    fn timeout_entry_tracked() {
        let mut batch = ToolBatchOutcome::new();
        let timeout = ToolExecutionStatus::Timeout {
            error: ToolExecutionError::new(
                "slow_tool".to_string(),
                vtcode_core::tools::registry::ToolErrorType::Timeout,
                "timed out after 30s".to_string(),
            ),
        };
        batch.record(&timeout);
        let s = batch.stats();
        assert_eq!(s.timed_out, 1);
    }

    #[test]
    fn cancellation_with_success_is_partial_success() {
        let mut batch = ToolBatchOutcome::new();
        let success = ToolExecutionStatus::Success {
            output: serde_json::json!("ok"),
            stdout: None,
            modified_files: vec![],
            command_success: true,
        };
        batch.record(&success);
        batch.record(&ToolExecutionStatus::Cancelled);

        assert!(batch.is_partial_success());
        assert!(batch.summary_line().contains("1 cancelled"));
    }

    #[test]
    fn warning_status_requires_an_enabled_or_non_empty_payload() {
        assert_eq!(
            ToolDisplayStatus::from_command_output(&serde_json::json!({"warning": false}), true),
            ToolDisplayStatus::Success
        );
        assert_eq!(
            ToolDisplayStatus::from_command_output(&serde_json::json!({"warning": "  "}), true),
            ToolDisplayStatus::Success
        );
        assert_eq!(
            ToolDisplayStatus::from_command_output(&serde_json::json!({"warning": "no results"}), true),
            ToolDisplayStatus::Warning
        );
    }
}