use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use crate::agent::{Agent, AgentFactory};
use crate::consensus::{ConsensusConfig, ConsensusEngine};
use crate::error::{MagiError, ProviderError};
use crate::provider::{CompletionConfig, LlmProvider};
use crate::reporting::{MagiReport, ReportConfig, ReportFormatter};
use crate::schema::{AgentName, AgentOutput, Mode};
use crate::validate::{ValidationLimits, Validator};
use tokio::task::AbortHandle;
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct MagiConfig {
pub timeout: Duration,
pub max_input_len: usize,
pub completion: CompletionConfig,
}
impl Default for MagiConfig {
fn default() -> Self {
Self {
timeout: Duration::from_secs(300),
max_input_len: 1_048_576,
completion: CompletionConfig::default(),
}
}
}
pub struct MagiBuilder {
default_provider: Arc<dyn LlmProvider>,
agent_providers: BTreeMap<AgentName, Arc<dyn LlmProvider>>,
custom_prompts: BTreeMap<AgentName, String>,
prompts_dir: Option<PathBuf>,
config: MagiConfig,
validation_limits: ValidationLimits,
consensus_config: ConsensusConfig,
report_config: ReportConfig,
}
impl MagiBuilder {
pub fn new(default_provider: Arc<dyn LlmProvider>) -> Self {
Self {
default_provider,
agent_providers: BTreeMap::new(),
custom_prompts: BTreeMap::new(),
prompts_dir: None,
config: MagiConfig::default(),
validation_limits: ValidationLimits::default(),
consensus_config: ConsensusConfig::default(),
report_config: ReportConfig::default(),
}
}
pub fn with_provider(mut self, name: AgentName, provider: Arc<dyn LlmProvider>) -> Self {
self.agent_providers.insert(name, provider);
self
}
pub fn with_custom_prompt(mut self, name: AgentName, prompt: String) -> Self {
self.custom_prompts.insert(name, prompt);
self
}
pub fn with_prompts_dir(mut self, dir: PathBuf) -> Self {
self.prompts_dir = Some(dir);
self
}
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.config.timeout = timeout;
self
}
pub fn with_max_input_len(mut self, max: usize) -> Self {
self.config.max_input_len = max;
self
}
pub fn with_completion_config(mut self, config: CompletionConfig) -> Self {
self.config.completion = config;
self
}
pub fn with_validation_limits(mut self, limits: ValidationLimits) -> Self {
self.validation_limits = limits;
self
}
pub fn with_consensus_config(mut self, config: ConsensusConfig) -> Self {
self.consensus_config = config;
self
}
pub fn with_report_config(mut self, config: ReportConfig) -> Self {
self.report_config = config;
self
}
pub fn build(self) -> Result<Magi, MagiError> {
let mut factory = AgentFactory::new(self.default_provider);
for (name, provider) in self.agent_providers {
factory = factory.with_provider(name, provider);
}
for (name, prompt) in self.custom_prompts {
factory = factory.with_custom_prompt(name, prompt);
}
if let Some(dir) = self.prompts_dir {
factory = factory.from_directory(&dir)?;
}
Ok(Magi {
config: self.config,
agent_factory: factory,
validator: Validator::with_limits(self.validation_limits),
consensus_engine: ConsensusEngine::new(self.consensus_config),
formatter: ReportFormatter::with_config(self.report_config),
})
}
}
struct AbortGuard(Vec<AbortHandle>);
impl Drop for AbortGuard {
fn drop(&mut self) {
for handle in &self.0 {
handle.abort();
}
}
}
pub struct Magi {
config: MagiConfig,
agent_factory: AgentFactory,
validator: Validator,
consensus_engine: ConsensusEngine,
formatter: ReportFormatter,
}
impl Magi {
pub fn new(provider: Arc<dyn LlmProvider>) -> Self {
MagiBuilder::new(provider).build().expect(
"Magi::new uses all defaults and cannot fail; \
this is an internal invariant violation",
)
}
pub fn builder(provider: Arc<dyn LlmProvider>) -> MagiBuilder {
MagiBuilder::new(provider)
}
pub async fn analyze(&self, mode: &Mode, content: &str) -> Result<MagiReport, MagiError> {
if content.len() > self.config.max_input_len {
return Err(MagiError::InputTooLarge {
size: content.len(),
max: self.config.max_input_len,
});
}
let agents = self.agent_factory.create_agents(*mode);
let prompt = build_prompt(mode, content);
let agent_results = self.launch_agents(agents, &prompt).await;
let (successful, failed_agents) = self.process_results(agent_results)?;
let consensus = self.consensus_engine.determine(&successful)?;
let banner = self.formatter.format_banner(&successful, &consensus);
let report = self.formatter.format_report(&successful, &consensus);
let degraded = successful.len() < 3;
Ok(MagiReport {
agents: successful,
consensus,
banner,
report,
degraded,
failed_agents,
})
}
async fn launch_agents(
&self,
agents: Vec<Agent>,
prompt: &str,
) -> Vec<(AgentName, Result<String, MagiError>)> {
let timeout = self.config.timeout;
let completion = self.config.completion.clone();
let mut handles = Vec::new();
let mut abort_handles = Vec::new();
for agent in agents {
let name = agent.name();
let user_prompt = prompt.to_string();
let config = completion.clone();
let handle = tokio::spawn(async move {
let result =
tokio::time::timeout(timeout, agent.execute(&user_prompt, &config)).await;
match result {
Ok(Ok(response)) => Ok(response),
Ok(Err(provider_err)) => Err(MagiError::Provider(provider_err)),
Err(_elapsed) => Err(MagiError::Provider(ProviderError::Timeout {
message: format!("agent timed out after {timeout:?}"),
})),
}
});
abort_handles.push(handle.abort_handle());
handles.push((name, handle));
}
let _guard = AbortGuard(abort_handles);
let mut results = Vec::new();
for (name, handle) in handles {
match handle.await {
Ok(result) => results.push((name, result)),
Err(join_err) => results.push((
name,
Err(MagiError::Provider(ProviderError::Process {
exit_code: None,
stderr: format!("agent task panicked: {join_err}"),
})),
)),
}
}
results
}
fn process_results(
&self,
results: Vec<(AgentName, Result<String, MagiError>)>,
) -> Result<(Vec<AgentOutput>, BTreeMap<AgentName, String>), MagiError> {
let mut successful = Vec::new();
let mut failed_agents = BTreeMap::new();
for (name, result) in results {
match result {
Ok(raw) => match parse_agent_response(&raw) {
Ok(output) => match self.validator.validate(&output) {
Ok(()) => successful.push(output),
Err(e) => {
failed_agents.insert(name, format!("validation: {e}"));
}
},
Err(e) => {
failed_agents.insert(name, format!("parse: {e}"));
}
},
Err(e) => {
failed_agents.insert(name, e.to_string());
}
}
}
let min_agents = self.consensus_engine.min_agents();
if successful.len() < min_agents {
return Err(MagiError::InsufficientAgents {
succeeded: successful.len(),
required: min_agents,
});
}
Ok((successful, failed_agents))
}
}
fn build_prompt(mode: &Mode, content: &str) -> String {
format!("MODE: {mode}\nCONTEXT:\n{content}")
}
fn parse_agent_response(raw: &str) -> Result<AgentOutput, MagiError> {
let trimmed = raw.trim();
let stripped = if trimmed.starts_with("```") {
let without_opening = if let Some(rest) = trimmed.strip_prefix("```json") {
rest
} else {
trimmed.strip_prefix("```").unwrap_or(trimmed)
};
without_opening
.strip_suffix("```")
.unwrap_or(without_opening)
.trim()
} else {
trimmed
};
if let Ok(output) = serde_json::from_str::<AgentOutput>(stripped) {
return Ok(output);
}
for (start, _) in stripped.match_indices('{') {
let candidate = &stripped[start..];
if let Ok(output) = serde_json::from_str::<AgentOutput>(candidate) {
return Ok(output);
}
}
Err(MagiError::Deserialization(
"no valid JSON object found in agent response".to_string(),
))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::schema::*;
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;
fn mock_agent_json(agent: &str, verdict: &str, confidence: f64) -> String {
format!(
r#"{{
"agent": "{agent}",
"verdict": "{verdict}",
"confidence": {confidence},
"summary": "Summary from {agent}",
"reasoning": "Reasoning from {agent}",
"findings": [],
"recommendation": "Recommendation from {agent}"
}}"#
)
}
struct MockProvider {
name: String,
model: String,
responses: Vec<Result<String, ProviderError>>,
call_count: AtomicUsize,
}
impl MockProvider {
fn success(name: &str, model: &str, responses: Vec<String>) -> Self {
Self {
name: name.to_string(),
model: model.to_string(),
responses: responses.into_iter().map(Ok).collect(),
call_count: AtomicUsize::new(0),
}
}
fn mixed(name: &str, model: &str, responses: Vec<Result<String, ProviderError>>) -> Self {
Self {
name: name.to_string(),
model: model.to_string(),
responses,
call_count: AtomicUsize::new(0),
}
}
fn calls(&self) -> usize {
self.call_count.load(Ordering::SeqCst)
}
}
#[async_trait::async_trait]
impl LlmProvider for MockProvider {
async fn complete(
&self,
_system_prompt: &str,
_user_prompt: &str,
_config: &CompletionConfig,
) -> Result<String, ProviderError> {
let idx = self.call_count.fetch_add(1, Ordering::SeqCst);
let idx = idx % self.responses.len();
self.responses[idx].clone()
}
fn name(&self) -> &str {
&self.name
}
fn model(&self) -> &str {
&self.model
}
}
#[tokio::test]
async fn test_analyze_unanimous_approve_returns_complete_report() {
let responses = vec![
mock_agent_json("melchior", "approve", 0.9),
mock_agent_json("balthasar", "approve", 0.85),
mock_agent_json("caspar", "approve", 0.95),
];
let provider = Arc::new(MockProvider::success("mock", "test-model", responses));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "fn main() {}").await;
let report = result.expect("analyze should succeed");
assert_eq!(report.agents.len(), 3);
assert!(!report.degraded);
assert!(report.failed_agents.is_empty());
assert_eq!(report.consensus.consensus_verdict, Verdict::Approve);
assert!(!report.banner.is_empty());
assert!(!report.report.is_empty());
}
#[tokio::test]
async fn test_analyze_one_agent_timeout_degrades_gracefully() {
let responses = vec![
Ok(mock_agent_json("melchior", "approve", 0.9)),
Ok(mock_agent_json("balthasar", "approve", 0.85)),
Err(ProviderError::Timeout {
message: "exceeded timeout".to_string(),
}),
];
let provider = Arc::new(MockProvider::mixed("mock", "test-model", responses));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "fn main() {}").await;
let report = result.expect("analyze should succeed with degradation");
assert!(report.degraded);
assert_eq!(report.failed_agents.len(), 1);
assert_eq!(report.agents.len(), 2);
}
#[tokio::test]
async fn test_analyze_one_agent_bad_json_degrades_gracefully() {
let responses = vec![
Ok(mock_agent_json("melchior", "approve", 0.9)),
Ok(mock_agent_json("balthasar", "approve", 0.85)),
Ok("not valid json at all".to_string()),
];
let provider = Arc::new(MockProvider::mixed("mock", "test-model", responses));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "fn main() {}").await;
let report = result.expect("analyze should succeed with degradation");
assert!(report.degraded);
}
#[tokio::test]
async fn test_analyze_two_agents_fail_returns_insufficient_agents() {
let responses = vec![
Ok(mock_agent_json("melchior", "approve", 0.9)),
Err(ProviderError::Timeout {
message: "timeout".to_string(),
}),
Err(ProviderError::Network {
message: "connection refused".to_string(),
}),
];
let provider = Arc::new(MockProvider::mixed("mock", "test-model", responses));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "fn main() {}").await;
match result {
Err(MagiError::InsufficientAgents {
succeeded,
required,
}) => {
assert_eq!(succeeded, 1);
assert_eq!(required, 2);
}
other => panic!("Expected InsufficientAgents, got: {other:?}"),
}
}
#[tokio::test]
async fn test_analyze_all_agents_fail_returns_insufficient_agents() {
let responses = vec![
Err(ProviderError::Timeout {
message: "timeout".to_string(),
}),
Err(ProviderError::Network {
message: "network".to_string(),
}),
Err(ProviderError::Auth {
message: "auth".to_string(),
}),
];
let provider = Arc::new(MockProvider::mixed("mock", "test-model", responses));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "fn main() {}").await;
match result {
Err(MagiError::InsufficientAgents {
succeeded,
required,
}) => {
assert_eq!(succeeded, 0);
assert_eq!(required, 2);
}
other => panic!("Expected InsufficientAgents, got: {other:?}"),
}
}
#[tokio::test]
async fn test_analyze_plain_text_response_treated_as_failure() {
let responses = vec![
Ok(mock_agent_json("melchior", "approve", 0.9)),
Ok(mock_agent_json("balthasar", "approve", 0.85)),
Ok("I think the code is good".to_string()),
];
let provider = Arc::new(MockProvider::mixed("mock", "test-model", responses));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "fn main() {}").await;
let report = result.expect("should succeed with degradation");
assert!(report.degraded);
assert_eq!(report.agents.len(), 2);
}
#[tokio::test]
async fn test_magi_new_creates_with_defaults() {
let responses = vec![
mock_agent_json("melchior", "approve", 0.9),
mock_agent_json("balthasar", "approve", 0.85),
mock_agent_json("caspar", "approve", 0.95),
];
let provider = Arc::new(MockProvider::success(
"test-provider",
"test-model",
responses,
));
let magi = Magi::new(provider as Arc<dyn LlmProvider>);
let result = magi.analyze(&Mode::CodeReview, "test content").await;
let report = result.expect("should succeed");
assert_eq!(report.agents.len(), 3);
}
#[tokio::test]
async fn test_builder_with_mixed_providers_and_custom_config() {
let default_responses = vec![
mock_agent_json("melchior", "approve", 0.9),
mock_agent_json("balthasar", "approve", 0.85),
];
let caspar_responses = vec![mock_agent_json("caspar", "reject", 0.8)];
let default_provider = Arc::new(MockProvider::success(
"default-provider",
"model-a",
default_responses,
));
let caspar_provider = Arc::new(MockProvider::success(
"caspar-provider",
"model-b",
caspar_responses,
));
let magi = MagiBuilder::new(default_provider.clone() as Arc<dyn LlmProvider>)
.with_provider(
AgentName::Caspar,
caspar_provider.clone() as Arc<dyn LlmProvider>,
)
.with_timeout(Duration::from_secs(60))
.build()
.expect("build should succeed");
let result = magi.analyze(&Mode::CodeReview, "test content").await;
let report = result.expect("should succeed");
assert_eq!(report.agents.len(), 3);
assert!(caspar_provider.calls() > 0);
}
#[tokio::test]
async fn test_analyze_input_too_large_rejects_without_launching_agents() {
let responses = vec![mock_agent_json("melchior", "approve", 0.9)];
let provider = Arc::new(MockProvider::success("mock", "test-model", responses));
let magi = MagiBuilder::new(provider.clone() as Arc<dyn LlmProvider>)
.with_max_input_len(100)
.build()
.expect("build should succeed");
let content = "x".repeat(200);
let result = magi.analyze(&Mode::CodeReview, &content).await;
match result {
Err(MagiError::InputTooLarge { size, max }) => {
assert_eq!(size, 200);
assert_eq!(max, 100);
}
other => panic!("Expected InputTooLarge, got: {other:?}"),
}
assert_eq!(provider.calls(), 0, "No agents should have been launched");
}
#[test]
fn test_magi_config_default_values() {
let config = MagiConfig::default();
assert_eq!(config.timeout, Duration::from_secs(300));
assert_eq!(config.max_input_len, 1_048_576);
}
#[test]
fn test_build_prompt_formats_mode_and_content() {
let result = build_prompt(&Mode::CodeReview, "fn main() {}");
assert_eq!(result, "MODE: code-review\nCONTEXT:\nfn main() {}");
}
#[test]
fn test_parse_agent_response_strips_code_fences() {
let json = mock_agent_json("melchior", "approve", 0.9);
let raw = format!("```json\n{json}\n```");
let result = parse_agent_response(&raw);
let output = result.expect("should parse successfully");
assert_eq!(output.agent, AgentName::Melchior);
assert_eq!(output.verdict, Verdict::Approve);
}
#[test]
fn test_parse_agent_response_extracts_json_from_preamble() {
let json = mock_agent_json("melchior", "approve", 0.9);
let raw = format!("Here is my analysis:\n{json}");
let result = parse_agent_response(&raw);
assert!(result.is_ok(), "should find JSON in preamble text");
}
#[test]
fn test_parse_agent_response_fails_on_invalid_input() {
let result = parse_agent_response("no json here");
assert!(result.is_err(), "should fail on invalid input");
}
#[test]
fn test_magi_builder_build_returns_result() {
let responses = vec![mock_agent_json("melchior", "approve", 0.9)];
let provider =
Arc::new(MockProvider::success("mock", "model", responses)) as Arc<dyn LlmProvider>;
let magi = MagiBuilder::new(provider).build();
assert!(magi.is_ok());
}
}