aether_core/core/
agent_builder.rs

1use super::agent::{AgentConfig, AutoContinue, RetryConfig};
2use crate::agent_spec::AgentSpec;
3use crate::context::{CompactionConfig, SessionUsageTracker};
4use crate::core::{Agent, AgentDeps, Prompt, PromptCache, Result};
5use crate::events::{AgentEvent, AgentObserver, Command};
6use crate::mcp::McpHandle;
7use llm::parser::ModelProviderParser;
8use llm::{ChatMessage, Context, ModelSettings, SessionUsageEvent, StreamingModelProvider, ToolDefinition};
9use std::sync::Arc;
10use std::time::Duration;
11use tokio::sync::mpsc::{self, Receiver, Sender};
12use tokio::task::JoinHandle;
13
14/// Handle for communicating with a running Agent
15pub struct AgentHandle {
16    handle: JoinHandle<()>,
17}
18
19impl AgentHandle {
20    /// Abort the agent task immediately.
21    pub fn abort(&self) {
22        self.handle.abort();
23    }
24
25    /// Returns `true` if the agent task has finished.
26    pub fn is_finished(&self) -> bool {
27        self.handle.is_finished()
28    }
29
30    /// Wait for the agent task to complete.
31    pub async fn await_completion(self) {
32        let _ = self.handle.await;
33    }
34}
35
36pub struct AgentBuilder {
37    llm: Arc<dyn StreamingModelProvider>,
38    prompts: Vec<Prompt>,
39    tool_definitions: Vec<ToolDefinition>,
40    initial_messages: Vec<ChatMessage>,
41    mcp: Option<McpHandle>,
42    channel_capacity: usize,
43    tool_timeout: Duration,
44    compaction_config: Option<CompactionConfig>,
45    max_auto_continues: u32,
46    retry_config: RetryConfig,
47    context_window: Option<u32>,
48    model_settings: ModelSettings,
49    observers: Vec<Box<dyn AgentObserver>>,
50    session_usage: SessionUsageTracker,
51    session_affinity_key: String,
52}
53
54impl AgentBuilder {
55    pub fn new(llm: Arc<dyn StreamingModelProvider>) -> Self {
56        Self {
57            llm,
58            prompts: Vec::new(),
59            tool_definitions: Vec::new(),
60            initial_messages: Vec::new(),
61            mcp: None,
62            channel_capacity: 1000,
63            tool_timeout: Duration::from_mins(60),
64            compaction_config: Some(CompactionConfig::default()),
65            max_auto_continues: 3,
66            retry_config: RetryConfig::default(),
67            context_window: None,
68            model_settings: ModelSettings::default(),
69            observers: Vec::new(),
70            session_usage: SessionUsageTracker::new("agent"),
71            session_affinity_key: uuid::Uuid::new_v4().to_string(),
72        }
73    }
74
75    /// Create a builder from a resolved `AgentSpec`.
76    ///
77    /// The LLM provider is derived from `spec.model` via `ModelProviderParser`.
78    /// `base_prompts` are prepended before the spec's own prompts.
79    pub async fn from_spec(spec: &AgentSpec, base_prompts: Vec<Prompt>, deps: &AgentDeps) -> Result<Self> {
80        let parser = ModelProviderParser::default().with_provider_connections(spec.provider_connections.clone());
81        let parser = match deps.oauth_credential_store.clone() {
82            Some(store) => parser.with_codex_provider(store),
83            None => parser,
84        };
85        let (provider, _) = parser.parse(&spec.model).await?;
86        let mut builder = Self::new(Arc::from(provider))
87            .context_window(spec.context_window)
88            .model_settings(spec.model_settings.clone())
89            .session_usage(SessionUsageTracker::new(&spec.name));
90
91        if let Some(key) = &deps.session_affinity_key {
92            builder = builder.session_affinity_key(key.clone());
93        }
94        if let Some(observer) = deps.observer(&spec.name) {
95            builder = builder.observer(observer);
96        }
97
98        for prompt in base_prompts {
99            builder = builder.system_prompt(prompt);
100        }
101
102        for prompt in &spec.prompts {
103            builder = builder.system_prompt(prompt.clone());
104        }
105
106        Ok(builder)
107    }
108
109    /// Add a prompt to the system prompt.
110    ///
111    /// Multiple prompts are concatenated with double newlines.
112    pub fn system_prompt(mut self, prompt: Prompt) -> Self {
113        self.prompts.push(prompt);
114        self
115    }
116
117    pub fn tools(mut self, mcp: McpHandle, tools: Vec<ToolDefinition>) -> Self {
118        self.tool_definitions = tools;
119        self.mcp = Some(mcp);
120        self
121    }
122
123    /// Set the timeout for tool execution
124    ///
125    /// If a tool does not return a result within this duration, it will be marked as failed
126    /// and the agent will continue processing.
127    ///
128    /// Default: 60 minutes
129    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
130        self.tool_timeout = timeout;
131        self
132    }
133
134    /// Configure context compaction settings.
135    ///
136    /// By default, agents automatically compact context when token usage exceeds
137    /// 85% of the context window, preventing overflow during long-running tasks.
138    ///
139    /// # Examples
140    /// ```ignore
141    /// // Custom threshold
142    /// agent(llm).compaction(CompactionConfig::with_threshold(0.9))
143    ///
144    /// // Disable compaction entirely
145    /// agent(llm).compaction(CompactionConfig::disabled())
146    ///
147    /// // Full customization
148    /// agent(llm).compaction(
149    ///     CompactionConfig::with_threshold(0.85)
150    ///         .keep_recent_tool_results(3)
151    ///         .min_messages(20)
152    /// )
153    /// ```
154    pub fn compaction(mut self, config: CompactionConfig) -> Self {
155        self.compaction_config = Some(config);
156        self
157    }
158
159    /// Disable context compaction entirely.
160    ///
161    /// Overflow errors from the model will be surfaced directly to callers.
162    pub fn disable_compaction(mut self) -> Self {
163        self.compaction_config = None;
164        self
165    }
166
167    /// Configure the maximum number of auto-continue attempts.
168    ///
169    /// When the LLM stops without making tool calls, the agent may inject a
170    /// continuation prompt and restart the LLM stream for resumable stop
171    /// reasons (for example, token length limits).
172    ///
173    /// This setting limits how many times the agent will attempt to continue
174    /// before giving up and ending the turn with [`TurnEvent::Ended`](crate::events::TurnEvent::Ended).
175    ///
176    /// Default: 3
177    ///
178    /// # Example
179    /// ```ignore
180    /// // Allow up to 5 auto-continue attempts
181    /// agent(llm).max_auto_continues(5)
182    ///
183    /// // Disable auto-continue entirely
184    /// agent(llm).max_auto_continues(0)
185    /// ```
186    pub fn max_auto_continues(mut self, max: u32) -> Self {
187        self.max_auto_continues = max;
188        self
189    }
190
191    /// Configure retry behavior for transient LLM provider failures.
192    pub fn retry(mut self, config: RetryConfig) -> Self {
193        self.retry_config = config;
194        self
195    }
196
197    /// Override the effective model context window in tokens.
198    pub fn context_window(mut self, context_window: Option<u32>) -> Self {
199        self.context_window = context_window;
200        self
201    }
202
203    /// Set the sampling controls (`temperature`, `top_p`, `max_tokens`) applied to
204    /// every model call this agent makes.
205    pub fn model_settings(mut self, model_settings: ModelSettings) -> Self {
206        self.model_settings = model_settings;
207        self
208    }
209
210    pub fn session_affinity_key(mut self, key: impl Into<String>) -> Self {
211        self.session_affinity_key = key.into();
212        self
213    }
214
215    /// Pre-populate the context with conversation history (e.g. from a restored session).
216    ///
217    /// These messages are inserted after the system prompt.
218    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
219        self.initial_messages = messages;
220        self
221    }
222
223    /// Attach an observer of the agent's event stream.
224    pub fn observer(mut self, observer: Box<dyn AgentObserver>) -> Self {
225        self.observers.push(observer);
226        self
227    }
228
229    /// Record usage under `tracker`, which names this agent in usage events.
230    pub fn session_usage(mut self, tracker: SessionUsageTracker) -> Self {
231        self.session_usage = tracker;
232        self
233    }
234
235    /// Continue session totals from the last persisted usage event, e.g. when
236    /// resuming a session.
237    pub fn resume_usage(mut self, last: &SessionUsageEvent) -> Self {
238        self.session_usage.resume_from(last);
239        self
240    }
241
242    pub async fn spawn(self) -> Result<(Sender<Command>, Receiver<AgentEvent>, AgentHandle)> {
243        let mut prompt_cache = PromptCache::new(self.prompts);
244        let system_content = prompt_cache.render().await?;
245        let mut messages = Vec::new();
246
247        if !system_content.is_empty() {
248            messages.push(ChatMessage::system(system_content));
249        }
250
251        messages.extend(self.initial_messages);
252        let (command_tx, command_rx) = mpsc::channel::<Command>(self.channel_capacity);
253        let (message_tx, agent_event_rx) = mpsc::channel::<AgentEvent>(self.channel_capacity);
254        let mut context = Context::new(messages, self.tool_definitions);
255        context.set_model_settings(self.model_settings);
256        context.set_session_affinity_key(Some(self.session_affinity_key));
257
258        let config = AgentConfig {
259            llm: self.llm,
260            context,
261            mcp: self.mcp,
262            tool_timeout: self.tool_timeout,
263            compaction_config: self.compaction_config,
264            auto_continue: AutoContinue::new(self.max_auto_continues),
265            retry_config: self.retry_config,
266            context_window: self.context_window,
267            prompt_cache,
268            observers: self.observers,
269            session_usage: self.session_usage,
270        };
271
272        let agent = Agent::new(config, command_rx, message_tx);
273        let agent_handle = tokio::spawn(agent.run());
274
275        Ok((command_tx, agent_event_rx, AgentHandle { handle: agent_handle }))
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::agent_spec::AgentSpecExposure;
283    use llm::ProviderConnectionOverrides;
284    use mcp_utils::client::ToolFilter;
285
286    #[tokio::test]
287    async fn test_agent_handle_is_finished() {
288        let handle = AgentHandle { handle: tokio::spawn(async {}) };
289        handle.await_completion().await;
290    }
291
292    #[tokio::test]
293    async fn test_agent_handle_abort() {
294        let handle = AgentHandle { handle: tokio::spawn(std::future::pending::<()>()) };
295        assert!(!handle.is_finished());
296        handle.abort();
297        while !handle.is_finished() {
298            tokio::task::yield_now().await;
299        }
300    }
301
302    #[tokio::test]
303    async fn system_prompt_preserves_add_order() {
304        let builder = AgentBuilder::new(Arc::new(llm::testing::FakeLlmProvider::new(vec![])))
305            .system_prompt(Prompt::text("first"))
306            .system_prompt(Prompt::text("second"))
307            .system_prompt(Prompt::text("third"));
308
309        let rendered = Prompt::build_all(&builder.prompts).await.unwrap();
310
311        assert_eq!(rendered, "first\n\nsecond\n\nthird");
312    }
313
314    #[tokio::test]
315    async fn from_spec_applies_context_window_and_model_settings() {
316        let settings = ModelSettings { temperature: Some(0.0), max_tokens: Some(128), ..Default::default() };
317        let spec = AgentSpec {
318            name: "alloy".to_string(),
319            description: "alloy".to_string(),
320            model: "ollama:llama3.2,llamacpp:local".to_string(),
321            reasoning_effort: None,
322            model_settings: settings.clone(),
323            context_window: Some(200_000),
324            prompts: vec![],
325            provider_connections: ProviderConnectionOverrides::default(),
326            mcp_config_sources: Vec::new(),
327            exposure: AgentSpecExposure::both(),
328            tools: ToolFilter::default(),
329        };
330
331        let dependencies = AgentDeps::default().with_session_affinity_key("conversation-123");
332        let builder = AgentBuilder::from_spec(&spec, vec![], &dependencies).await.unwrap();
333
334        assert_eq!(builder.context_window, Some(200_000));
335        assert_eq!(builder.model_settings, settings);
336        assert_eq!(builder.session_affinity_key, "conversation-123");
337    }
338
339    #[tokio::test]
340    async fn from_spec_accepts_alloy_model_specs() {
341        let spec = AgentSpec {
342            name: "alloy".to_string(),
343            description: "alloy".to_string(),
344            model: "ollama:llama3.2,llamacpp:local".to_string(),
345            reasoning_effort: None,
346            model_settings: ModelSettings::default(),
347            context_window: None,
348            prompts: vec![],
349            provider_connections: ProviderConnectionOverrides::default(),
350            mcp_config_sources: Vec::new(),
351            exposure: AgentSpecExposure::both(),
352            tools: ToolFilter::default(),
353        };
354
355        let builder = AgentBuilder::from_spec(&spec, vec![], &AgentDeps::default()).await;
356        assert!(builder.is_ok());
357    }
358}
aether_core/core/agent_builder.rs

aether_core/core/
agent_builder.rs
