aether_core/core/
agent_builder.rs

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

aether_core/core/
agent_builder.rs
