Skip to main content

aster/agents/subagent_scheduler/
config.rs

1//! SubAgent 调度器配置
2//!
3//! 定义调度器的各种配置选项
4
5use serde::{Deserialize, Serialize};
6use std::time::Duration;
7
8use crate::agents::context::ContextInheritanceConfig;
9
10/// 调度器配置
11#[derive(Debug, Clone, Serialize, Deserialize)]
12#[serde(rename_all = "camelCase")]
13pub struct SchedulerConfig {
14    /// 最大并发数
15    pub max_concurrency: usize,
16    /// 默认任务超时时间
17    pub default_timeout: Duration,
18    /// 是否在失败时重试
19    pub retry_on_failure: bool,
20    /// 首次错误时停止
21    pub stop_on_first_error: bool,
22    /// 最大重试次数
23    pub max_retries: usize,
24    /// 重试延迟
25    pub retry_delay: Duration,
26    /// 上下文继承配置
27    pub context_inheritance: ContextInheritanceConfig,
28    /// 是否自动生成摘要
29    pub auto_summarize: bool,
30    /// 摘要最大 token 数
31    pub summary_max_tokens: usize,
32    /// 默认模型
33    pub default_model: Option<String>,
34    /// 是否启用进度回调
35    pub enable_progress_callback: bool,
36}
37
38impl Default for SchedulerConfig {
39    fn default() -> Self {
40        Self {
41            max_concurrency: 5,
42            default_timeout: Duration::from_secs(300), // 5 分钟
43            retry_on_failure: true,
44            stop_on_first_error: false,
45            max_retries: 3,
46            retry_delay: Duration::from_secs(1),
47            context_inheritance: ContextInheritanceConfig::default(),
48            auto_summarize: true,
49            summary_max_tokens: 2000,
50            default_model: None,
51            enable_progress_callback: true,
52        }
53    }
54}
55
56impl SchedulerConfig {
57    /// 创建高并发配置(适合研究任务)
58    pub fn high_concurrency() -> Self {
59        Self {
60            max_concurrency: 10,
61            default_timeout: Duration::from_secs(600),
62            ..Default::default()
63        }
64    }
65
66    /// 创建低并发配置(适合编码任务)
67    pub fn low_concurrency() -> Self {
68        Self {
69            max_concurrency: 2,
70            stop_on_first_error: true,
71            ..Default::default()
72        }
73    }
74
75    /// 创建串行配置
76    pub fn sequential() -> Self {
77        Self {
78            max_concurrency: 1,
79            stop_on_first_error: true,
80            ..Default::default()
81        }
82    }
83
84    /// 设置最大并发数
85    pub fn with_max_concurrency(mut self, max: usize) -> Self {
86        self.max_concurrency = max;
87        self
88    }
89
90    /// 设置默认超时
91    pub fn with_timeout(mut self, timeout: Duration) -> Self {
92        self.default_timeout = timeout;
93        self
94    }
95
96    /// 设置重试配置
97    pub fn with_retry(mut self, enabled: bool, max_retries: usize) -> Self {
98        self.retry_on_failure = enabled;
99        self.max_retries = max_retries;
100        self
101    }
102
103    /// 设置首次错误停止
104    pub fn with_stop_on_first_error(mut self, stop: bool) -> Self {
105        self.stop_on_first_error = stop;
106        self
107    }
108
109    /// 设置上下文继承配置
110    pub fn with_context_inheritance(mut self, config: ContextInheritanceConfig) -> Self {
111        self.context_inheritance = config;
112        self
113    }
114
115    /// 设置默认模型
116    pub fn with_default_model(mut self, model: impl Into<String>) -> Self {
117        self.default_model = Some(model.into());
118        self
119    }
120}
121
122#[cfg(test)]
123mod tests {
124    use super::*;
125
126    #[test]
127    fn test_default_config() {
128        let config = SchedulerConfig::default();
129        assert_eq!(config.max_concurrency, 5);
130        assert!(config.retry_on_failure);
131        assert!(!config.stop_on_first_error);
132    }
133
134    #[test]
135    fn test_high_concurrency_config() {
136        let config = SchedulerConfig::high_concurrency();
137        assert_eq!(config.max_concurrency, 10);
138    }
139
140    #[test]
141    fn test_sequential_config() {
142        let config = SchedulerConfig::sequential();
143        assert_eq!(config.max_concurrency, 1);
144        assert!(config.stop_on_first_error);
145    }
146
147    #[test]
148    fn test_config_builder() {
149        let config = SchedulerConfig::default()
150            .with_max_concurrency(8)
151            .with_stop_on_first_error(true)
152            .with_default_model("sonnet");
153
154        assert_eq!(config.max_concurrency, 8);
155        assert!(config.stop_on_first_error);
156        assert_eq!(config.default_model, Some("sonnet".to_string()));
157    }
158}