forked from boos/BoosAPI
chore: sync current workspace changes
This commit is contained in:
@@ -0,0 +1,251 @@
|
||||
use open_agent_sdk::types::*;
|
||||
use open_agent_sdk::{Agent, AgentOptions};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_creation() {
|
||||
let agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
assert!(agent.get_messages().is_empty());
|
||||
assert!(!agent.session_id().is_empty());
|
||||
assert!(agent.tool_names().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_honors_explicit_session_id() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
session_id: Some("caller-session".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.session_id(), "caller-session");
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_custom_model() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
model: Some("claude-haiku-4-5-20251001".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.model(), "claude-haiku-4-5-20251001");
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_custom_cwd() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
cwd: Some("/tmp".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should use /tmp as working directory
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_allowed_tools() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
allowed_tools: Some(vec!["Read".to_string(), "Glob".to_string()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.tool_names(), ["Glob", "Read"]);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_rejects_unknown_allowed_tool() {
|
||||
let result = Agent::new(AgentOptions {
|
||||
allowed_tools: Some(vec!["WebFetch".to_string()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(matches!(result, Err(error) if error.contains("unknown or unavailable tool")));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_custom_tools() {
|
||||
use async_trait::async_trait;
|
||||
|
||||
struct TestTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for TestTool {
|
||||
fn name(&self) -> &str {
|
||||
"TestTool"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"A test tool"
|
||||
}
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
async fn call(
|
||||
&self,
|
||||
_input: serde_json::Value,
|
||||
_ctx: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
Ok(ToolResult::text("test result"))
|
||||
}
|
||||
}
|
||||
|
||||
let agent = Agent::new(AgentOptions {
|
||||
custom_tools: vec![Arc::new(TestTool)],
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_clear() {
|
||||
let mut agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
// Manually add a message for testing
|
||||
agent.messages.push(Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "test".to_string(),
|
||||
}],
|
||||
});
|
||||
|
||||
assert!(!agent.get_messages().is_empty());
|
||||
agent.clear();
|
||||
assert!(agent.get_messages().is_empty());
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_set_model() {
|
||||
let mut agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
agent.set_model("claude-opus-4-6-20250514");
|
||||
assert_eq!(agent.model(), "claude-opus-4-6-20250514");
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_max_turns() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
max_turns: Some(5),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.max_turns, 5);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_max_budget() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
max_budget_usd: Some(1.0),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(agent.max_budget_usd, Some(1.0));
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_thinking() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
thinking: Some(ThinkingConfig::enabled(10000)),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(agent.thinking.is_some());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_system_prompt() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
system_prompt: Some("You are a code reviewer.".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
agent.system_prompt.as_deref(),
|
||||
Some("You are a code reviewer.")
|
||||
);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_append_system_prompt() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
append_system_prompt: Some("Always respond in JSON.".to_string()),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(
|
||||
agent.append_system_prompt.as_deref(),
|
||||
Some("Always respond in JSON.")
|
||||
);
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_cost_tracker() {
|
||||
let agent = Agent::new(AgentOptions::default()).await.unwrap();
|
||||
|
||||
let cost = agent.cost_tracker().total_cost().await;
|
||||
assert_eq!(cost, 0.0);
|
||||
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_permission_mode() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
permission_mode: Some(PermissionMode::AcceptEdits),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should be created successfully with the permission mode
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_agent_with_disallowed_tools() {
|
||||
let agent = Agent::new(AgentOptions {
|
||||
disallowed_tools: Some(vec!["Bash".to_string()]),
|
||||
..Default::default()
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Agent should not have Bash tool
|
||||
assert!(!agent.session_id().is_empty());
|
||||
agent.close().await;
|
||||
}
|
||||
@@ -0,0 +1,197 @@
|
||||
use open_agent_sdk::costtracker::CostTracker;
|
||||
use open_agent_sdk::types::Usage;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_new() {
|
||||
let tracker = CostTracker::new();
|
||||
assert_eq!(tracker.total_cost().await, 0.0);
|
||||
assert_eq!(tracker.total_tokens().await, 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_add_usage() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
let usage = Usage {
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
};
|
||||
|
||||
tracker
|
||||
.add_usage("claude-sonnet-4-6-20250514", &usage)
|
||||
.await;
|
||||
|
||||
assert_eq!(tracker.total_tokens().await, 1500);
|
||||
assert!(tracker.total_cost().await > 0.0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_multiple_models() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-sonnet-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-opus-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 500,
|
||||
output_tokens: 200,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(tracker.total_tokens().await, 2200);
|
||||
|
||||
let sonnet_usage = tracker.get_model_usage("claude-sonnet-4-6-20250514").await;
|
||||
assert!(sonnet_usage.is_some());
|
||||
assert_eq!(sonnet_usage.unwrap().input_tokens, 1000);
|
||||
|
||||
let opus_usage = tracker.get_model_usage("claude-opus-4-6-20250514").await;
|
||||
assert!(opus_usage.is_some());
|
||||
assert_eq!(opus_usage.unwrap().input_tokens, 500);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_accumulates() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
let usage = Usage {
|
||||
input_tokens: 1000,
|
||||
output_tokens: 500,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
tracker
|
||||
.add_usage("claude-sonnet-4-6-20250514", &usage)
|
||||
.await;
|
||||
tracker
|
||||
.add_usage("claude-sonnet-4-6-20250514", &usage)
|
||||
.await;
|
||||
|
||||
assert_eq!(tracker.total_tokens().await, 3000);
|
||||
|
||||
let model_usage = tracker
|
||||
.get_model_usage("claude-sonnet-4-6-20250514")
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(model_usage.input_tokens, 2000);
|
||||
assert_eq!(model_usage.output_tokens, 1000);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_durations() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker.add_api_duration(100).await;
|
||||
tracker.add_api_duration(200).await;
|
||||
tracker.add_tool_duration(50).await;
|
||||
|
||||
let summary = tracker.summary().await;
|
||||
assert_eq!(summary.api_duration_ms, 300);
|
||||
assert_eq!(summary.tool_duration_ms, 50);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_code_changes() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker.add_code_changes(10, 5).await;
|
||||
tracker.add_code_changes(20, 3).await;
|
||||
|
||||
let summary = tracker.summary().await;
|
||||
assert_eq!(summary.lines_added, 30);
|
||||
assert_eq!(summary.lines_removed, 8);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_tracker_web_searches() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker.add_web_search().await;
|
||||
tracker.add_web_search().await;
|
||||
|
||||
let summary = tracker.summary().await;
|
||||
assert_eq!(summary.web_searches, 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_calculation_sonnet() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
// 1M input + 1M output for sonnet = $3 + $15 = $18
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-sonnet-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 1_000_000,
|
||||
output_tokens: 1_000_000,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let cost = tracker.total_cost().await;
|
||||
assert!((cost - 18.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_cost_calculation_opus() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
// 1M input + 1M output for opus = $15 + $75 = $90
|
||||
tracker
|
||||
.add_usage(
|
||||
"claude-opus-4-6-20250514",
|
||||
&Usage {
|
||||
input_tokens: 1_000_000,
|
||||
output_tokens: 1_000_000,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let cost = tracker.total_cost().await;
|
||||
assert!((cost - 90.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_get_all_usage() {
|
||||
let tracker = CostTracker::new();
|
||||
|
||||
tracker
|
||||
.add_usage(
|
||||
"model-a",
|
||||
&Usage {
|
||||
input_tokens: 100,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
tracker
|
||||
.add_usage(
|
||||
"model-b",
|
||||
&Usage {
|
||||
input_tokens: 200,
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.await;
|
||||
|
||||
let all = tracker.get_all_usage().await;
|
||||
assert_eq!(all.len(), 2);
|
||||
assert!(all.contains_key("model-a"));
|
||||
assert!(all.contains_key("model-b"));
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
use open_agent_sdk::hooks::*;
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_config_default() {
|
||||
let config = HookConfig::default();
|
||||
assert!(config.pre_tool_use.is_empty());
|
||||
assert!(config.post_tool_use.is_empty());
|
||||
assert!(config.post_sampling.is_empty());
|
||||
assert!(config.stop.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pre_tool_use_hook_allow() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "*".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: false,
|
||||
message: None,
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = config.run_pre_tool_use("Bash", &json!({})).await;
|
||||
assert!(result.is_none()); // Not blocked
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_pre_tool_use_hook_block() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "Bash".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: true,
|
||||
message: Some("Bash is blocked".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = config.run_pre_tool_use("Bash", &json!({})).await;
|
||||
assert!(result.is_some());
|
||||
assert!(result.unwrap().blocked);
|
||||
|
||||
// Different tool should not be blocked
|
||||
let result = config.run_pre_tool_use("Read", &json!({})).await;
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_matcher_wildcard() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "mcp__*".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: true,
|
||||
message: Some("MCP blocked".to_string()),
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
let result = config
|
||||
.run_pre_tool_use("mcp__server__tool", &json!({}))
|
||||
.await;
|
||||
assert!(result.is_some());
|
||||
|
||||
let result = config.run_pre_tool_use("Read", &json!({})).await;
|
||||
assert!(result.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_matcher_pipe_or() {
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "Bash|Edit|Write".to_string(),
|
||||
handler: Arc::new(|_input| {
|
||||
Box::pin(async move {
|
||||
HookOutput {
|
||||
blocked: true,
|
||||
message: None,
|
||||
..Default::default()
|
||||
}
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert!(config.run_pre_tool_use("Bash", &json!({})).await.is_some());
|
||||
assert!(config.run_pre_tool_use("Edit", &json!({})).await.is_some());
|
||||
assert!(config.run_pre_tool_use("Write", &json!({})).await.is_some());
|
||||
assert!(config.run_pre_tool_use("Read", &json!({})).await.is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_post_tool_use_hook() {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
let called = Arc::new(AtomicBool::new(false));
|
||||
let called_clone = called.clone();
|
||||
|
||||
let config = HookConfig {
|
||||
post_tool_use: vec![HookRule {
|
||||
matcher: "Bash".to_string(),
|
||||
handler: Arc::new(move |_input| {
|
||||
let called = called_clone.clone();
|
||||
Box::pin(async move {
|
||||
called.store(true, Ordering::SeqCst);
|
||||
HookOutput::default()
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
config.run_post_tool_use("Bash", &json!({}), "output").await;
|
||||
assert!(called.load(Ordering::SeqCst));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_hook_input_fields() {
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
let verified = Arc::new(AtomicBool::new(false));
|
||||
let verified_clone = verified.clone();
|
||||
|
||||
let config = HookConfig {
|
||||
pre_tool_use: vec![HookRule {
|
||||
matcher: "*".to_string(),
|
||||
handler: Arc::new(move |input| {
|
||||
let verified = verified_clone.clone();
|
||||
Box::pin(async move {
|
||||
assert_eq!(input.event, HookEvent::PreToolUse);
|
||||
assert_eq!(input.tool_name.as_deref(), Some("Bash"));
|
||||
assert!(input.tool_input.is_some());
|
||||
verified.store(true, Ordering::SeqCst);
|
||||
HookOutput::default()
|
||||
})
|
||||
}),
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
config
|
||||
.run_pre_tool_use("Bash", &json!({"command": "ls"}))
|
||||
.await;
|
||||
assert!(verified.load(Ordering::SeqCst));
|
||||
}
|
||||
@@ -0,0 +1,163 @@
|
||||
use open_agent_sdk::permissions::*;
|
||||
use open_agent_sdk::types::PermissionMode;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_bypass_mode_allows_everything() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::BypassPermissions,
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({"command": "rm -rf /"})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
assert_eq!(
|
||||
check_permission(&config, "Write", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_rules() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
deny_rules: vec![Rule {
|
||||
tool_name: "Bash".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({})),
|
||||
PermissionResult::Deny("Tool 'Bash' denied by rule".to_string())
|
||||
);
|
||||
// Other tools should not be denied
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allow_rules() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
allow_rules: vec![Rule {
|
||||
tool_name: "Read".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_wildcard_deny() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
deny_rules: vec![Rule {
|
||||
tool_name: "mcp__*".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
allow_rules: vec![Rule {
|
||||
tool_name: "*".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "mcp__server__tool", &json!({})),
|
||||
PermissionResult::Deny("Tool 'mcp__server__tool' denied by rule".to_string())
|
||||
);
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_allowed_tools_whitelist() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
allowed_tools: Some(vec!["Read".to_string(), "Glob".to_string()]),
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
check_permission(&config, "Read", &json!({})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
match check_permission(&config, "Bash", &json!({})) {
|
||||
PermissionResult::Deny(_) => {}
|
||||
_ => panic!("Expected deny"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_deny_with_pattern() {
|
||||
let config = PermissionConfig {
|
||||
mode: PermissionMode::Default,
|
||||
deny_rules: vec![Rule {
|
||||
tool_name: "Bash".to_string(),
|
||||
pattern: Some("rm".to_string()),
|
||||
}],
|
||||
allow_rules: vec![Rule {
|
||||
tool_name: "*".to_string(),
|
||||
pattern: None,
|
||||
}],
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
// Deny when pattern matches
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({"command": "rm file.txt"})),
|
||||
PermissionResult::Deny("Tool 'Bash' denied by rule".to_string())
|
||||
);
|
||||
|
||||
// Allow when pattern doesn't match
|
||||
assert_eq!(
|
||||
check_permission(&config, "Bash", &json!({"command": "ls"})),
|
||||
PermissionResult::Allow
|
||||
);
|
||||
}
|
||||
|
||||
// --- Filesystem Validator Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_is_sensitive_path() {
|
||||
assert!(FilesystemValidator::is_sensitive_path("/home/user/.env"));
|
||||
assert!(FilesystemValidator::is_sensitive_path(
|
||||
"/home/user/.env.local"
|
||||
));
|
||||
assert!(FilesystemValidator::is_sensitive_path(
|
||||
"/home/user/.ssh/id_rsa"
|
||||
));
|
||||
assert!(FilesystemValidator::is_sensitive_path(
|
||||
"/app/credentials.json"
|
||||
));
|
||||
assert!(!FilesystemValidator::is_sensitive_path("/app/main.rs"));
|
||||
assert!(!FilesystemValidator::is_sensitive_path(
|
||||
"/app/src/config.rs"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_filesystem_validator() {
|
||||
let validator = FilesystemValidator::new("/tmp/project");
|
||||
|
||||
// Paths within working dir should be allowed
|
||||
assert!(validator
|
||||
.validate_path("/tmp/project/src/main.rs", false)
|
||||
.is_ok());
|
||||
assert!(validator
|
||||
.validate_path("/tmp/project/src/main.rs", true)
|
||||
.is_ok());
|
||||
}
|
||||
@@ -0,0 +1,695 @@
|
||||
use open_agent_sdk::tools::ToolRegistry;
|
||||
use open_agent_sdk::types::{Tool, ToolUseContext};
|
||||
use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
fn create_test_context(dir: &str) -> ToolUseContext {
|
||||
ToolUseContext {
|
||||
working_dir: dir.to_string(),
|
||||
abort_signal: tokio_util::sync::CancellationToken::new(),
|
||||
read_file_state: Arc::new(RwLock::new(std::collections::HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Registry Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_default_registry() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
assert!(registry.len() > 0);
|
||||
|
||||
// Check core tools exist
|
||||
assert!(registry.get("Bash").is_some());
|
||||
assert!(registry.get("Read").is_some());
|
||||
assert!(registry.get("Write").is_some());
|
||||
assert!(registry.get("Edit").is_some());
|
||||
assert!(registry.get("Glob").is_some());
|
||||
assert!(registry.get("Grep").is_some());
|
||||
assert!(registry.get("WebFetch").is_some());
|
||||
assert!(registry.get("WebSearch").is_some());
|
||||
assert!(registry.get("AskUserQuestion").is_some());
|
||||
assert!(registry.get("TaskCreate").is_some());
|
||||
assert!(registry.get("TaskGet").is_some());
|
||||
assert!(registry.get("TaskList").is_some());
|
||||
assert!(registry.get("TaskUpdate").is_some());
|
||||
assert!(registry.get("ToolSearch").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_safe_builtin_registry_is_exact_and_rejects_broad_tools() {
|
||||
let registry = ToolRegistry::safe_builtin_registry(&["Read", "Grep"]).unwrap();
|
||||
let mut names = registry.names();
|
||||
names.sort();
|
||||
assert_eq!(names, ["Grep", "Read"]);
|
||||
assert!(ToolRegistry::safe_builtin_registry(&["WebFetch"]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_register_custom() {
|
||||
use async_trait::async_trait;
|
||||
use open_agent_sdk::types::{ToolError, ToolInputSchema, ToolResult};
|
||||
use serde_json::Value;
|
||||
|
||||
struct CustomTool;
|
||||
|
||||
#[async_trait]
|
||||
impl Tool for CustomTool {
|
||||
fn name(&self) -> &str {
|
||||
"CustomTool"
|
||||
}
|
||||
fn description(&self) -> &str {
|
||||
"A test tool"
|
||||
}
|
||||
fn input_schema(&self) -> ToolInputSchema {
|
||||
ToolInputSchema::default()
|
||||
}
|
||||
async fn call(
|
||||
&self,
|
||||
_input: Value,
|
||||
_ctx: &ToolUseContext,
|
||||
) -> Result<ToolResult, ToolError> {
|
||||
Ok(ToolResult::text("custom result"))
|
||||
}
|
||||
}
|
||||
|
||||
let mut registry = ToolRegistry::new();
|
||||
assert!(registry.is_empty());
|
||||
|
||||
registry.register(Arc::new(CustomTool));
|
||||
assert_eq!(registry.len(), 1);
|
||||
assert!(registry.get("CustomTool").is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_filter() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read_only = registry.filter(|t| t.is_read_only(&json!({})));
|
||||
assert!(read_only.len() > 0);
|
||||
|
||||
// Read, Glob, Grep should be read-only
|
||||
let names: Vec<&str> = read_only.iter().map(|t| t.name()).collect();
|
||||
assert!(names.contains(&"Read"));
|
||||
assert!(names.contains(&"Glob"));
|
||||
assert!(names.contains(&"Grep"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_retain() {
|
||||
let mut registry = ToolRegistry::default_registry();
|
||||
let initial_count = registry.len();
|
||||
|
||||
registry.retain(&["Read", "Glob", "Grep"]);
|
||||
assert_eq!(registry.len(), 3);
|
||||
assert!(registry.len() < initial_count);
|
||||
|
||||
assert!(registry.get("Read").is_some());
|
||||
assert!(registry.get("Bash").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_registry_remove() {
|
||||
let mut registry = ToolRegistry::default_registry();
|
||||
assert!(registry.get("Bash").is_some());
|
||||
|
||||
registry.remove(&["Bash"]);
|
||||
assert!(registry.get("Bash").is_none());
|
||||
}
|
||||
|
||||
// --- Bash Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_echo() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = bash
|
||||
.call(json!({"command": "echo 'hello world'"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("hello world"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_exit_code() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = bash.call(json!({"command": "false"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("Exit code:"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_destructive_detection() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = bash
|
||||
.call(json!({"command": "rm -rf /"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("destructive"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_bash_is_read_only() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let bash = registry.get("Bash").unwrap();
|
||||
|
||||
assert!(bash.is_read_only(&json!({"command": "ls -la"})));
|
||||
assert!(bash.is_read_only(&json!({"command": "git status"})));
|
||||
assert!(!bash.is_read_only(&json!({"command": "rm file.txt"})));
|
||||
assert!(!bash.is_read_only(&json!({"command": "cargo build"})));
|
||||
}
|
||||
|
||||
// --- File Read Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("test.txt");
|
||||
std::fs::write(&file_path, "line 1\nline 2\nline 3\n").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = read
|
||||
.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("line 1"));
|
||||
assert!(text.contains("line 2"));
|
||||
assert!(text.contains("line 3"));
|
||||
// Should have line numbers
|
||||
assert!(text.contains("1\t"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file_not_found() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = read
|
||||
.call(
|
||||
json!({"file_path": "/tmp/nonexistent_file_12345.txt"}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("not found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_file_with_offset_and_limit() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("test.txt");
|
||||
let content: String = (1..=100).map(|i| format!("line {}\n", i)).collect();
|
||||
std::fs::write(&file_path, &content).unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = read
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"offset": 10,
|
||||
"limit": 5
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("line 11"));
|
||||
assert!(text.contains("line 15"));
|
||||
assert!(!text.contains("line 16"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_directory_error() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = read.call(json!({"file_path": "/tmp"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("directory"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_read_is_read_only() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let read = registry.get("Read").unwrap();
|
||||
assert!(read.is_read_only(&json!({})));
|
||||
assert!(read.is_concurrency_safe(&json!({})));
|
||||
}
|
||||
|
||||
// --- File Write Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_new_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("new_file.txt");
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let write = registry.get("Write").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = write
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"content": "hello world"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Created"));
|
||||
assert_eq!(std::fs::read_to_string(&file_path).unwrap(), "hello world");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_write_creates_directories() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("a/b/c/file.txt");
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let write = registry.get("Write").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = write
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"content": "nested"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(file_path.exists());
|
||||
}
|
||||
|
||||
// --- File Edit Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_file() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_test.txt");
|
||||
std::fs::write(&file_path, "hello world").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
// Must read first for staleness tracking
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "hello",
|
||||
"new_string": "goodbye"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&file_path).unwrap(),
|
||||
"goodbye world"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_string_not_found() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_test2.txt");
|
||||
std::fs::write(&file_path, "hello world").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "xyz",
|
||||
"new_string": "abc"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("not found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_multiple_occurrences() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_multi.txt");
|
||||
std::fs::write(&file_path, "foo bar foo baz foo").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Without replace_all, should fail for multiple occurrences
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "qux"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("3 times"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_replace_all() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_all.txt");
|
||||
std::fs::write(&file_path, "foo bar foo baz foo").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let read = registry.get("Read").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
read.call(json!({"file_path": file_path.to_str().unwrap()}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "foo",
|
||||
"new_string": "qux",
|
||||
"replace_all": true
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(
|
||||
std::fs::read_to_string(&file_path).unwrap(),
|
||||
"qux bar qux baz qux"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_edit_same_string_error() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let file_path = dir.path().join("edit_same.txt");
|
||||
std::fs::write(&file_path, "hello world").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let edit = registry.get("Edit").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = edit
|
||||
.call(
|
||||
json!({
|
||||
"file_path": file_path.to_str().unwrap(),
|
||||
"old_string": "hello",
|
||||
"new_string": "hello"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("different"));
|
||||
}
|
||||
|
||||
// --- Glob Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_find_files() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("a.rs"), "").unwrap();
|
||||
std::fs::write(dir.path().join("b.rs"), "").unwrap();
|
||||
std::fs::write(dir.path().join("c.txt"), "").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob.call(json!({"pattern": "*.rs"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("a.rs"));
|
||||
assert!(text.contains("b.rs"));
|
||||
assert!(!text.contains("c.txt"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_no_matches() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob.call(json!({"pattern": "*.xyz"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("No files found"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_dot_lists_current_directory_without_recursing() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(dir.path().join("root.rs"), "").unwrap();
|
||||
std::fs::create_dir(dir.path().join("nested")).unwrap();
|
||||
std::fs::write(dir.path().join("nested").join("child.rs"), "").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob.call(json!({"pattern": "."}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("root.rs"));
|
||||
assert!(!text.contains("child.rs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_skips_heavy_directories_before_descending() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::create_dir(dir.path().join("node_modules")).unwrap();
|
||||
std::fs::write(dir.path().join("node_modules").join("dep.rs"), "").unwrap();
|
||||
std::fs::create_dir(dir.path().join("src")).unwrap();
|
||||
std::fs::write(dir.path().join("src").join("lib.rs"), "").unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = glob
|
||||
.call(json!({"pattern": "**/*.rs"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("lib.rs"));
|
||||
assert!(!text.contains("dep.rs"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_glob_returns_when_cancelled() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
ctx.abort_signal.cancel();
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let glob = registry.get("Glob").unwrap();
|
||||
|
||||
let result = glob.call(json!({"pattern": "**/*"}), &ctx).await.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("No files found"));
|
||||
}
|
||||
|
||||
// --- Grep Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_grep_search() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
std::fs::write(
|
||||
dir.path().join("test.txt"),
|
||||
"hello world\nfoo bar\nhello again\n",
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let grep = registry.get("Grep").unwrap();
|
||||
let ctx = create_test_context(dir.path().to_str().unwrap());
|
||||
|
||||
let result = grep
|
||||
.call(
|
||||
json!({
|
||||
"pattern": "hello",
|
||||
"path": dir.path().to_str().unwrap(),
|
||||
"output_mode": "content"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert!(!result.is_error);
|
||||
let text = result.get_text();
|
||||
assert!(text.contains("hello world") || text.contains("hello"));
|
||||
}
|
||||
|
||||
// --- Task Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_task_lifecycle() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
// Create a task
|
||||
let create = registry.get("TaskCreate").unwrap();
|
||||
let result = create
|
||||
.call(
|
||||
json!({
|
||||
"subject": "Test task",
|
||||
"description": "A test task"
|
||||
}),
|
||||
&ctx,
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Created task"));
|
||||
|
||||
// List tasks
|
||||
let list = registry.get("TaskList").unwrap();
|
||||
let result = list.call(json!({}), &ctx).await.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Test task"));
|
||||
|
||||
// Get task
|
||||
let get = registry.get("TaskGet").unwrap();
|
||||
let result = get.call(json!({"id": "task_1"}), &ctx).await.unwrap();
|
||||
assert!(!result.is_error);
|
||||
assert!(result.get_text().contains("Test task"));
|
||||
|
||||
// Update task
|
||||
let update = registry.get("TaskUpdate").unwrap();
|
||||
let result = update
|
||||
.call(json!({"id": "task_1", "status": "completed"}), &ctx)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!result.is_error);
|
||||
|
||||
// Verify update
|
||||
let result = get.call(json!({"id": "task_1"}), &ctx).await.unwrap();
|
||||
assert!(result.get_text().to_lowercase().contains("completed"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_task_not_found() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let get = registry.get("TaskGet").unwrap();
|
||||
let result = get.call(json!({"id": "nonexistent"}), &ctx).await.unwrap();
|
||||
assert!(result.is_error);
|
||||
assert!(result.get_text().contains("not found"));
|
||||
}
|
||||
|
||||
// --- WebFetch Tool Tests ---
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_webfetch_missing_url() {
|
||||
let registry = ToolRegistry::default_registry();
|
||||
let webfetch = registry.get("WebFetch").unwrap();
|
||||
let ctx = create_test_context("/tmp");
|
||||
|
||||
let result = webfetch.call(json!({}), &ctx).await;
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
// --- Diff Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_unified_diff() {
|
||||
let old = "line 1\nline 2\nline 3\n";
|
||||
let new = "line 1\nline 2 modified\nline 3\n";
|
||||
let diff = open_agent_sdk::tools::diff::unified_diff(old, new, "test.txt");
|
||||
|
||||
assert!(diff.contains("--- a/test.txt"));
|
||||
assert!(diff.contains("+++ b/test.txt"));
|
||||
assert!(diff.contains("-line 2"));
|
||||
assert!(diff.contains("+line 2 modified"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_count_changes() {
|
||||
let old = "a\nb\nc\n";
|
||||
let new = "a\nx\nc\nd\n";
|
||||
let (added, removed) = open_agent_sdk::tools::diff::count_changes(old, new);
|
||||
assert!(added >= 1);
|
||||
assert!(removed >= 1);
|
||||
}
|
||||
@@ -0,0 +1,210 @@
|
||||
use open_agent_sdk::types::*;
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn test_tool_result_text() {
|
||||
let result = ToolResult::text("hello world");
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(result.get_text(), "hello world");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_result_error() {
|
||||
let result = ToolResult::error("something went wrong");
|
||||
assert!(result.is_error);
|
||||
assert_eq!(result.get_text(), "something went wrong");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_result_image() {
|
||||
let result = ToolResult::image("base64data".to_string(), "image/png".to_string());
|
||||
assert!(!result.is_error);
|
||||
assert_eq!(result.content.len(), 1);
|
||||
match &result.content[0] {
|
||||
ToolResultContent::Image { source } => {
|
||||
assert_eq!(source.data, "base64data");
|
||||
assert_eq!(source.media_type, "image/png");
|
||||
}
|
||||
_ => panic!("Expected image content"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_tool_input_schema_default() {
|
||||
let schema = ToolInputSchema::default();
|
||||
assert_eq!(schema.schema_type, "object");
|
||||
assert!(schema.properties.is_empty());
|
||||
assert!(schema.required.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_permission_mode_default() {
|
||||
let mode = PermissionMode::default();
|
||||
assert_eq!(mode, PermissionMode::BypassPermissions);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_permission_mode_serialization() {
|
||||
let mode = PermissionMode::AcceptEdits;
|
||||
let json = serde_json::to_string(&mode).unwrap();
|
||||
assert_eq!(json, "\"acceptEdits\"");
|
||||
|
||||
let deserialized: PermissionMode = serde_json::from_str(&json).unwrap();
|
||||
assert_eq!(deserialized, PermissionMode::AcceptEdits);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_creation() {
|
||||
let msg = Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
};
|
||||
assert_eq!(msg.role, MessageRole::User);
|
||||
assert_eq!(msg.content.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_text_from_message() {
|
||||
let msg = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "Hello ".to_string(),
|
||||
},
|
||||
ContentBlock::Text {
|
||||
text: "World".to_string(),
|
||||
},
|
||||
],
|
||||
};
|
||||
assert_eq!(extract_text(&msg), "Hello World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_text_skips_non_text() {
|
||||
let msg = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "Hello".to_string(),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "123".to_string(),
|
||||
name: "Read".to_string(),
|
||||
input: json!({}),
|
||||
},
|
||||
ContentBlock::Text {
|
||||
text: " World".to_string(),
|
||||
},
|
||||
],
|
||||
};
|
||||
assert_eq!(extract_text(&msg), "Hello World");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_tool_uses() {
|
||||
let msg = Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "Let me read that".to_string(),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "tu_1".to_string(),
|
||||
name: "Read".to_string(),
|
||||
input: json!({"file_path": "/tmp/test.txt"}),
|
||||
},
|
||||
ContentBlock::ToolUse {
|
||||
id: "tu_2".to_string(),
|
||||
name: "Glob".to_string(),
|
||||
input: json!({"pattern": "*.rs"}),
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
let tool_uses = extract_tool_uses(&msg);
|
||||
assert_eq!(tool_uses.len(), 2);
|
||||
assert_eq!(tool_uses[0].1, "Read");
|
||||
assert_eq!(tool_uses[1].1, "Glob");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_usage_default() {
|
||||
let usage = Usage::default();
|
||||
assert_eq!(usage.input_tokens, 0);
|
||||
assert_eq!(usage.output_tokens, 0);
|
||||
assert_eq!(usage.cache_creation_input_tokens, 0);
|
||||
assert_eq!(usage.cache_read_input_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thinking_config() {
|
||||
let config = ThinkingConfig::enabled(10000);
|
||||
assert_eq!(config.thinking_type, "enabled");
|
||||
assert_eq!(config.budget_tokens, Some(10000));
|
||||
|
||||
let disabled = ThinkingConfig::disabled();
|
||||
assert_eq!(disabled.thinking_type, "disabled");
|
||||
assert_eq!(disabled.budget_tokens, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_message_role_serialization() {
|
||||
let user: MessageRole = serde_json::from_str("\"user\"").unwrap();
|
||||
assert_eq!(user, MessageRole::User);
|
||||
|
||||
let assistant: MessageRole = serde_json::from_str("\"assistant\"").unwrap();
|
||||
assert_eq!(assistant, MessageRole::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_block_text_serialization() {
|
||||
let block = ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
};
|
||||
let json = serde_json::to_value(&block).unwrap();
|
||||
assert_eq!(json["type"], "text");
|
||||
assert_eq!(json["text"], "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_content_block_tool_use_serialization() {
|
||||
let block = ContentBlock::ToolUse {
|
||||
id: "tu_1".to_string(),
|
||||
name: "Read".to_string(),
|
||||
input: json!({"file_path": "/test.txt"}),
|
||||
};
|
||||
let json = serde_json::to_value(&block).unwrap();
|
||||
assert_eq!(json["type"], "tool_use");
|
||||
assert_eq!(json["name"], "Read");
|
||||
assert_eq!(json["input"]["file_path"], "/test.txt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sdk_message_serialization() {
|
||||
let msg = SDKMessage::System {
|
||||
message: "Agent started".to_string(),
|
||||
};
|
||||
let json = serde_json::to_value(&msg).unwrap();
|
||||
assert_eq!(json["type"], "system");
|
||||
assert_eq!(json["message"], "Agent started");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_cache_control() {
|
||||
let cc = CacheControl::ephemeral();
|
||||
assert_eq!(cc.control_type, "ephemeral");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mcp_server_config_stdio() {
|
||||
let config = McpServerConfig::Stdio {
|
||||
command: "npx".to_string(),
|
||||
args: vec!["server".to_string()],
|
||||
env: std::collections::HashMap::new(),
|
||||
};
|
||||
let json = serde_json::to_value(&config).unwrap();
|
||||
assert_eq!(json["transport"], "stdio");
|
||||
assert_eq!(json["command"], "npx");
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
use open_agent_sdk::types::*;
|
||||
use open_agent_sdk::utils::{compact, messages, retry, tokens};
|
||||
use serde_json::json;
|
||||
|
||||
// --- Token Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_estimate_tokens() {
|
||||
assert_eq!(tokens::estimate_tokens(""), 0);
|
||||
assert_eq!(tokens::estimate_tokens("1234"), 1);
|
||||
assert_eq!(tokens::estimate_tokens("12345678"), 2);
|
||||
assert!(tokens::estimate_tokens("hello world this is a test") > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_messages_tokens() {
|
||||
let msgs = vec![
|
||||
Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "hello".to_string(),
|
||||
}],
|
||||
},
|
||||
Message {
|
||||
role: MessageRole::Assistant,
|
||||
content: vec![ContentBlock::Text {
|
||||
text: "world".to_string(),
|
||||
}],
|
||||
},
|
||||
];
|
||||
let tokens = tokens::estimate_messages_tokens(&msgs);
|
||||
assert!(tokens > 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_context_window_size() {
|
||||
assert_eq!(
|
||||
tokens::get_context_window_size("claude-sonnet-4-6-20250514"),
|
||||
200_000
|
||||
);
|
||||
assert_eq!(
|
||||
tokens::get_context_window_size("claude-opus-4-6-20250514"),
|
||||
200_000
|
||||
);
|
||||
assert_eq!(
|
||||
tokens::get_context_window_size("claude-haiku-4-5-20251001"),
|
||||
200_000
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_auto_compact_threshold() {
|
||||
let threshold = tokens::get_auto_compact_threshold("claude-sonnet-4-6-20250514");
|
||||
assert_eq!(threshold, 200_000 - 13_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_cost() {
|
||||
let usage = Usage {
|
||||
input_tokens: 1_000_000,
|
||||
output_tokens: 1_000_000,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
};
|
||||
|
||||
let cost = tokens::estimate_cost("claude-sonnet-4-6-20250514", &usage);
|
||||
// Sonnet: $3/M input + $15/M output = $18
|
||||
assert!((cost - 18.0).abs() < 0.01);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_estimate_cost_with_cache() {
|
||||
let usage = Usage {
|
||||
input_tokens: 0,
|
||||
output_tokens: 0,
|
||||
cache_creation_input_tokens: 1_000_000,
|
||||
cache_read_input_tokens: 1_000_000,
|
||||
};
|
||||
|
||||
let cost = tokens::estimate_cost("claude-sonnet-4-6-20250514", &usage);
|
||||
// Sonnet cache: $3.75/M write + $0.3/M read = $4.05
|
||||
assert!((cost - 4.05).abs() < 0.01);
|
||||
}
|
||||
|
||||
// --- Message Utility Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_create_user_message() {
|
||||
let msg = messages::create_user_message("hello");
|
||||
assert_eq!(msg.role, MessageRole::User);
|
||||
assert_eq!(messages::extract_text(&msg), "hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_create_assistant_message() {
|
||||
let msg = messages::create_assistant_message("hi there");
|
||||
assert_eq!(msg.role, MessageRole::Assistant);
|
||||
assert_eq!(messages::extract_text(&msg), "hi there");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_messages_empty() {
|
||||
let normalized = messages::normalize_messages(&[]);
|
||||
assert!(normalized.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_messages_starts_with_user() {
|
||||
let msgs = vec![messages::create_assistant_message("I'm an assistant")];
|
||||
let normalized = messages::normalize_messages(&msgs);
|
||||
assert_eq!(normalized[0].role, MessageRole::User);
|
||||
assert_eq!(normalized.len(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_normalize_messages_merges_consecutive() {
|
||||
let msgs = vec![
|
||||
messages::create_user_message("hello"),
|
||||
messages::create_user_message("world"),
|
||||
messages::create_assistant_message("hi"),
|
||||
];
|
||||
let normalized = messages::normalize_messages(&msgs);
|
||||
assert_eq!(normalized.len(), 2); // Merged two user messages
|
||||
assert_eq!(normalized[0].role, MessageRole::User);
|
||||
assert_eq!(normalized[1].role, MessageRole::Assistant);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_strip_images() {
|
||||
let msgs = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![
|
||||
ContentBlock::Text {
|
||||
text: "look at this".to_string(),
|
||||
},
|
||||
ContentBlock::Image {
|
||||
source: ImageContentSource {
|
||||
source_type: "base64".to_string(),
|
||||
media_type: "image/png".to_string(),
|
||||
data: "abcdef".to_string(),
|
||||
},
|
||||
},
|
||||
],
|
||||
}];
|
||||
|
||||
let stripped = messages::strip_images(&msgs);
|
||||
assert_eq!(stripped[0].content.len(), 1);
|
||||
match &stripped[0].content[0] {
|
||||
ContentBlock::Text { text } => assert_eq!(text, "look at this"),
|
||||
_ => panic!("Expected text block"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_truncate_text() {
|
||||
assert_eq!(messages::truncate_text("hello", 10), "hello");
|
||||
assert_eq!(
|
||||
messages::truncate_text("hello world", 5),
|
||||
"hello... (truncated)"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Compact Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_should_auto_compact_false() {
|
||||
let msgs = vec![messages::create_user_message("short message")];
|
||||
assert!(!compact::should_auto_compact(
|
||||
&msgs,
|
||||
"claude-sonnet-4-6-20250514"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_micro_compact_preserves_short_results() {
|
||||
let msgs = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::ToolResult {
|
||||
tool_use_id: "tu_1".to_string(),
|
||||
content: vec![ToolResultContentBlock::Text {
|
||||
text: "short result".to_string(),
|
||||
}],
|
||||
is_error: false,
|
||||
}],
|
||||
}];
|
||||
|
||||
let compacted = compact::micro_compact_messages(&msgs);
|
||||
if let ContentBlock::ToolResult { content, .. } = &compacted[0].content[0] {
|
||||
if let ToolResultContentBlock::Text { text } = &content[0] {
|
||||
assert_eq!(text, "short result");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_micro_compact_truncates_long_results() {
|
||||
let long_text = "x".repeat(100_000);
|
||||
let msgs = vec![Message {
|
||||
role: MessageRole::User,
|
||||
content: vec![ContentBlock::ToolResult {
|
||||
tool_use_id: "tu_1".to_string(),
|
||||
content: vec![ToolResultContentBlock::Text {
|
||||
text: long_text.clone(),
|
||||
}],
|
||||
is_error: false,
|
||||
}],
|
||||
}];
|
||||
|
||||
let compacted = compact::micro_compact_messages(&msgs);
|
||||
if let ContentBlock::ToolResult { content, .. } = &compacted[0].content[0] {
|
||||
if let ToolResultContentBlock::Text { text } = &content[0] {
|
||||
assert!(text.len() < long_text.len());
|
||||
assert!(text.contains("truncated"));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// --- Retry Tests ---
|
||||
|
||||
#[test]
|
||||
fn test_retry_config_default() {
|
||||
let config = retry::RetryConfig::default();
|
||||
assert_eq!(config.max_retries, 3);
|
||||
assert_eq!(config.base_delay_ms, 2000);
|
||||
assert_eq!(config.max_delay_ms, 30_000);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_retryable() {
|
||||
use open_agent_sdk::api::ApiError;
|
||||
|
||||
assert!(retry::is_retryable(&ApiError::RateLimitError));
|
||||
assert!(retry::is_retryable(&ApiError::Timeout));
|
||||
assert!(retry::is_retryable(&ApiError::NetworkError(
|
||||
"conn reset".to_string()
|
||||
)));
|
||||
assert!(retry::is_retryable(&ApiError::HttpError {
|
||||
status: 500,
|
||||
message: "internal error".to_string(),
|
||||
}));
|
||||
assert!(!retry::is_retryable(&ApiError::AuthError(
|
||||
"bad key".to_string()
|
||||
)));
|
||||
assert!(!retry::is_retryable(&ApiError::PromptTooLong(
|
||||
"too long".to_string()
|
||||
)));
|
||||
assert!(!retry::is_retryable(&ApiError::PartialStream(
|
||||
"visible output already emitted".to_string()
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_auth_error() {
|
||||
use open_agent_sdk::api::ApiError;
|
||||
|
||||
assert!(retry::is_auth_error(&ApiError::AuthError(
|
||||
"bad key".to_string()
|
||||
)));
|
||||
assert!(!retry::is_auth_error(&ApiError::RateLimitError));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_is_prompt_too_long() {
|
||||
use open_agent_sdk::api::ApiError;
|
||||
|
||||
assert!(retry::is_prompt_too_long(&ApiError::PromptTooLong(
|
||||
"too long".to_string()
|
||||
)));
|
||||
assert!(!retry::is_prompt_too_long(&ApiError::RateLimitError));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_retry_delay() {
|
||||
let config = retry::RetryConfig::default();
|
||||
let delay = retry::get_retry_delay(&config, 0);
|
||||
// Should be between 1s and 3s (base 2s with jitter)
|
||||
assert!(delay.as_millis() >= 1000);
|
||||
assert!(delay.as_millis() <= 4000);
|
||||
}
|
||||
Reference in New Issue
Block a user