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

aether_core/core/
agent_builder.rs
