mod common;
use std::sync::Arc;
use std::time::Duration;
use serde_json::json;
use tokio::sync::broadcast;
use butterfly_bot::brain::manager::BrainManager;
use butterfly_bot::domains::agent::AIAgent;
use butterfly_bot::interfaces::providers::{LlmResponse, ToolCall};
use butterfly_bot::services::agent::{AgentService, UiEvent};
use common::{DummyTool, QueueLlmProvider};
async fn recv_event(rx: &mut broadcast::Receiver<UiEvent>) -> UiEvent {
tokio::time::timeout(Duration::from_secs(2), rx.recv())
.await
.expect("timed out waiting for ui event")
.expect("ui event channel closed")
}
#[tokio::test]
async fn execution_trace_redacts_tool_event_payload() {
let llm = Arc::new(QueueLlmProvider::new(vec![
LlmResponse {
text: "Action: call trace_tool for secret work".to_string(),
tool_calls: vec![ToolCall {
name: "trace_tool".to_string(),
arguments: json!({
"api_key": "sk-abc123",
"authorization": "Bearer raw-token",
"note": "github_pat_abcdef12345"
}),
}],
},
LlmResponse {
text: "Summary: complete".to_string(),
tool_calls: Vec::new(),
},
]));
let brain = Arc::new(BrainManager::new(json!({})));
let (tx, mut rx) = broadcast::channel(16);
let agent = AIAgent {
name: "trace-agent".to_string(),
instructions: "trace".to_string(),
specialization: "ops".to_string(),
};
let service = AgentService::new(llm, agent, None, None, None, None, brain, Some(tx));
let tool = Arc::new(DummyTool::new("dummy_tool"));
assert!(service.tool_registry.register_tool(tool).await);
assert!(
service
.tool_registry
.assign_tool_to_agent(service.agent_name(), "dummy_tool")
.await
);
let _ = service
.generate_response("u1", "run trace", "", None)
.await
.unwrap();
let event = recv_event(&mut rx).await;
assert_eq!(event.event_type, "tool");
assert_eq!(event.user_id, "u1");
assert_eq!(event.tool, "trace_tool");
assert_eq!(event.status, "not_found");
assert!(event.timestamp > 0);
let args = event.payload.get("args").expect("args payload");
assert_eq!(
args.get("api_key").and_then(|v| v.as_str()),
Some("[REDACTED]")
);
assert_eq!(
args.get("authorization").and_then(|v| v.as_str()),
Some("[REDACTED]")
);
assert_eq!(
args.get("note").and_then(|v| v.as_str()),
Some("github_pat_[REDACTED]")
);
assert_eq!(
event.payload.get("message").and_then(|v| v.as_str()),
Some("Tool not found")
);
}
#[tokio::test]
async fn execution_trace_emits_events_in_call_order() {
let llm = Arc::new(QueueLlmProvider::new(vec![
LlmResponse {
text: "Action: call trace_tool first".to_string(),
tool_calls: vec![ToolCall {
name: "trace_tool".to_string(),
arguments: json!({"step": 1}),
}],
},
LlmResponse {
text: "Action: call trace_tool second".to_string(),
tool_calls: vec![ToolCall {
name: "trace_tool".to_string(),
arguments: json!({"step": 2}),
}],
},
LlmResponse {
text: "Summary: done".to_string(),
tool_calls: Vec::new(),
},
]));
let brain = Arc::new(BrainManager::new(json!({})));
let (tx, mut rx) = broadcast::channel(16);
let agent = AIAgent {
name: "trace-agent".to_string(),
instructions: "trace".to_string(),
specialization: "ops".to_string(),
};
let service = AgentService::new(llm, agent, None, None, None, None, brain, Some(tx));
let tool = Arc::new(DummyTool::new("dummy_tool"));
assert!(service.tool_registry.register_tool(tool).await);
assert!(
service
.tool_registry
.assign_tool_to_agent(service.agent_name(), "dummy_tool")
.await
);
let _ = service
.generate_response("u2", "run two steps", "", None)
.await
.unwrap();
let first = recv_event(&mut rx).await;
let second = recv_event(&mut rx).await;
assert_eq!(first.event_type, "tool");
assert_eq!(second.event_type, "tool");
assert_eq!(first.user_id, "u2");
assert_eq!(second.user_id, "u2");
assert_eq!(first.tool, "trace_tool");
assert_eq!(second.tool, "trace_tool");
assert_eq!(first.status, "not_found");
assert_eq!(second.status, "not_found");
let first_step = first
.payload
.get("args")
.and_then(|v| v.get("step"))
.and_then(|v| v.as_i64());
let second_step = second
.payload
.get("args")
.and_then(|v| v.get("step"))
.and_then(|v| v.as_i64());
assert_eq!(first_step, Some(1));
assert_eq!(second_step, Some(2));
assert!(second.timestamp >= first.timestamp);
}