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

aether_core/core/
agent_builder.rs
