aether_core/core/
agent_builder.rs

1use super::agent::{AgentConfig, AutoContinue, RetryConfig};
2use crate::agent_spec::AgentSpec;
3use crate::context::CompactionConfig;
4use crate::core::{Agent, AgentDeps, Prompt, PromptCache, Result};
5use crate::events::{AgentEvent, AgentObserver, Command};
6use crate::mcp::run_mcp_task::McpCommand;
7use llm::parser::ModelProviderParser;
8use llm::types::IsoString;
9use llm::{ChatMessage, Context, ModelSettings, 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_tx: Option<Sender<McpCommand>>,
43    channel_capacity: usize,
44    tool_timeout: Duration,
45    compaction_config: Option<CompactionConfig>,
46    max_auto_continues: u32,
47    retry_config: RetryConfig,
48    prompt_cache_key: Option<String>,
49    context_window: Option<u32>,
50    model_settings: ModelSettings,
51    observers: Vec<Box<dyn AgentObserver>>,
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_tx: None,
62            channel_capacity: 1000,
63            tool_timeout: Duration::from_mins(20),
64            compaction_config: Some(CompactionConfig::default()),
65            max_auto_continues: 3,
66            retry_config: RetryConfig::default(),
67            prompt_cache_key: None,
68            context_window: None,
69            model_settings: ModelSettings::default(),
70            observers: Vec::new(),
71        }
72    }
73
74    /// Create a builder from a resolved `AgentSpec`.
75    ///
76    /// The LLM provider is derived from `spec.model` via `ModelProviderParser`.
77    /// `base_prompts` are prepended before the spec's own prompts.
78    pub async fn from_spec(spec: &AgentSpec, base_prompts: Vec<Prompt>, deps: &AgentDeps) -> Result<Self> {
79        let parser = ModelProviderParser::default().with_provider_connections(spec.provider_connections.clone());
80        let parser = match deps.oauth_credential_store.clone() {
81            Some(store) => parser.with_codex_provider(store),
82            None => parser,
83        };
84        let (provider, _) = parser.parse(&spec.model).await?;
85        let mut builder = Self::new(Arc::from(provider))
86            .context_window(spec.context_window)
87            .model_settings(spec.model_settings.clone());
88
89        if let Some(observer) = deps.observer() {
90            builder = builder.observer(observer);
91        }
92
93        for prompt in base_prompts {
94            builder = builder.system_prompt(prompt);
95        }
96
97        for prompt in &spec.prompts {
98            builder = builder.system_prompt(prompt.clone());
99        }
100
101        Ok(builder)
102    }
103
104    /// Add a prompt to the system prompt.
105    ///
106    /// Multiple prompts are concatenated with double newlines.
107    pub fn system_prompt(mut self, prompt: Prompt) -> Self {
108        self.prompts.push(prompt);
109        self
110    }
111
112    pub fn tools(mut self, tx: Sender<McpCommand>, tools: Vec<ToolDefinition>) -> Self {
113        self.tool_definitions = tools;
114        self.mcp_tx = Some(tx);
115        self
116    }
117
118    /// Set the timeout for tool execution
119    ///
120    /// If a tool does not return a result within this duration, it will be marked as failed
121    /// and the agent will continue processing.
122    ///
123    /// Default: 20 minutes
124    pub fn tool_timeout(mut self, timeout: Duration) -> Self {
125        self.tool_timeout = timeout;
126        self
127    }
128
129    /// Configure context compaction settings.
130    ///
131    /// By default, agents automatically compact context when token usage exceeds
132    /// 85% of the context window, preventing overflow during long-running tasks.
133    ///
134    /// # Examples
135    /// ```ignore
136    /// // Custom threshold
137    /// agent(llm).compaction(CompactionConfig::with_threshold(0.9))
138    ///
139    /// // Disable compaction entirely
140    /// agent(llm).compaction(CompactionConfig::disabled())
141    ///
142    /// // Full customization
143    /// agent(llm).compaction(
144    ///     CompactionConfig::with_threshold(0.85)
145    ///         .keep_recent_tool_results(3)
146    ///         .min_messages(20)
147    /// )
148    /// ```
149    pub fn compaction(mut self, config: CompactionConfig) -> Self {
150        self.compaction_config = Some(config);
151        self
152    }
153
154    /// Disable context compaction entirely.
155    ///
156    /// Overflow errors from the model will be surfaced directly to callers.
157    pub fn disable_compaction(mut self) -> Self {
158        self.compaction_config = None;
159        self
160    }
161
162    /// Configure the maximum number of auto-continue attempts.
163    ///
164    /// When the LLM stops without making tool calls, the agent may inject a
165    /// continuation prompt and restart the LLM stream for resumable stop
166    /// reasons (for example, token length limits).
167    ///
168    /// This setting limits how many times the agent will attempt to continue
169    /// before giving up and ending the turn with [`TurnEvent::Ended`](crate::events::TurnEvent::Ended).
170    ///
171    /// Default: 3
172    ///
173    /// # Example
174    /// ```ignore
175    /// // Allow up to 5 auto-continue attempts
176    /// agent(llm).max_auto_continues(5)
177    ///
178    /// // Disable auto-continue entirely
179    /// agent(llm).max_auto_continues(0)
180    /// ```
181    pub fn max_auto_continues(mut self, max: u32) -> Self {
182        self.max_auto_continues = max;
183        self
184    }
185
186    /// Configure retry behavior for transient LLM provider failures.
187    pub fn retry(mut self, config: RetryConfig) -> Self {
188        self.retry_config = config;
189        self
190    }
191
192    /// Set a prompt cache key for LLM provider request routing.
193    ///
194    /// This is typically a session ID (UUID) that remains stable across all
195    /// turns within a conversation, improving prompt cache hit rates.
196    pub fn prompt_cache_key(mut self, key: String) -> Self {
197        self.prompt_cache_key = Some(key);
198        self
199    }
200
201    /// Override the effective model context window in tokens.
202    pub fn context_window(mut self, context_window: Option<u32>) -> Self {
203        self.context_window = context_window;
204        self
205    }
206
207    /// Set the sampling controls (`temperature`, `top_p`, `max_tokens`) applied to
208    /// every model call this agent makes.
209    pub fn model_settings(mut self, model_settings: ModelSettings) -> Self {
210        self.model_settings = model_settings;
211        self
212    }
213
214    /// Pre-populate the context with conversation history (e.g. from a restored session).
215    ///
216    /// These messages are inserted after the system prompt.
217    pub fn messages(mut self, messages: Vec<ChatMessage>) -> Self {
218        self.initial_messages = messages;
219        self
220    }
221
222    /// Attach an observer of the agent's event stream.
223    pub fn observer(mut self, observer: Box<dyn AgentObserver>) -> Self {
224        self.observers.push(observer);
225        self
226    }
227
228    pub async fn spawn(self) -> Result<(Sender<Command>, Receiver<AgentEvent>, AgentHandle)> {
229        let mut prompt_cache = PromptCache::new(self.prompts);
230        let system_content = prompt_cache.render().await?;
231        let mut messages = Vec::new();
232
233        if !system_content.is_empty() {
234            messages.push(ChatMessage::System { content: system_content, timestamp: IsoString::now() });
235        }
236
237        messages.extend(self.initial_messages);
238        let (command_tx, command_rx) = mpsc::channel::<Command>(self.channel_capacity);
239        let (message_tx, agent_event_rx) = mpsc::channel::<AgentEvent>(self.channel_capacity);
240        let mut context = Context::new(messages, self.tool_definitions);
241        context.set_prompt_cache_key(self.prompt_cache_key);
242        context.set_model_settings(self.model_settings);
243
244        let config = AgentConfig {
245            llm: self.llm,
246            context,
247            mcp_command_tx: self.mcp_tx,
248            tool_timeout: self.tool_timeout,
249            compaction_config: self.compaction_config,
250            auto_continue: AutoContinue::new(self.max_auto_continues),
251            retry_config: self.retry_config,
252            context_window: self.context_window,
253            prompt_cache,
254            observers: self.observers,
255        };
256
257        let agent = Agent::new(config, command_rx, message_tx);
258        let agent_handle = tokio::spawn(agent.run());
259
260        Ok((command_tx, agent_event_rx, AgentHandle { handle: agent_handle }))
261    }
262}
263
264#[cfg(test)]
265mod tests {
266    use super::*;
267    use crate::agent_spec::{AgentSpecExposure, ToolFilter};
268    use crate::events::{AgentCommand, ContextEvent, UserCommand};
269    use llm::testing::FakeLlmProvider;
270    use llm::{LlmResponse, ProviderConnectionOverrides};
271
272    #[tokio::test]
273    async fn test_agent_handle_is_finished() {
274        let handle = AgentHandle { handle: tokio::spawn(async {}) };
275        handle.await_completion().await;
276    }
277
278    #[tokio::test]
279    async fn test_agent_handle_abort() {
280        let handle = AgentHandle {
281            handle: tokio::spawn(async {
282                tokio::time::sleep(Duration::from_mins(1)).await;
283            }),
284        };
285        assert!(!handle.is_finished());
286        handle.abort();
287        // Give the runtime a moment to process the abort
288        tokio::time::sleep(Duration::from_millis(10)).await;
289        assert!(handle.is_finished());
290    }
291
292    #[tokio::test]
293    async fn context_window_override_supplies_unknown_provider_limit() {
294        let llm = Arc::new(FakeLlmProvider::with_single_response(vec![
295            LlmResponse::start("msg"),
296            LlmResponse::usage(100_000, 10),
297            LlmResponse::done(),
298        ]));
299
300        let (tx, mut rx, handle) = AgentBuilder::new(llm).context_window(Some(200_000)).spawn().await.unwrap();
301        tx.send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text("hello")] }))
302            .await
303            .unwrap();
304
305        let update = next_context_usage(&mut rx).await;
306        assert_eq!(update.context_limit, Some(200_000));
307        assert_eq!(update.usage_ratio, Some(0.5));
308        assert_eq!(update.input_tokens, 100_000);
309        handle.abort();
310    }
311
312    #[tokio::test]
313    async fn context_window_override_beats_provider_limit() {
314        let llm = Arc::new(
315            FakeLlmProvider::with_single_response(vec![
316                LlmResponse::start("msg"),
317                LlmResponse::usage(100_000, 10),
318                LlmResponse::done(),
319            ])
320            .with_context_window(Some(128_000)),
321        );
322
323        let (tx, mut rx, handle) = AgentBuilder::new(llm).context_window(Some(200_000)).spawn().await.unwrap();
324        tx.send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text("hello")] }))
325            .await
326            .unwrap();
327
328        let update = next_context_usage(&mut rx).await;
329        assert_eq!(update.context_limit, Some(200_000));
330        assert_eq!(update.usage_ratio, Some(0.5));
331        handle.abort();
332    }
333
334    #[tokio::test]
335    async fn context_window_override_survives_model_switch() {
336        let llm = Arc::new(FakeLlmProvider::new(vec![]).with_context_window(Some(128_000)));
337        let (tx, mut rx, handle) = AgentBuilder::new(llm).context_window(Some(200_000)).spawn().await.unwrap();
338
339        tx.send(Command::AgentCommand(AgentCommand::SwitchModel(Box::new(
340            FakeLlmProvider::new(vec![]).with_display_name("new fake").with_context_window(Some(32_000)),
341        ))))
342        .await
343        .unwrap();
344
345        let update = next_context_usage(&mut rx).await;
346        assert_eq!(update.context_limit, Some(200_000));
347        assert_eq!(update.usage_ratio, Some(0.0));
348        handle.abort();
349    }
350
351    async fn next_context_usage(rx: &mut Receiver<AgentEvent>) -> ContextUsageUpdate {
352        loop {
353            if let AgentEvent::Context(ContextEvent::UsageUpdated { usage }) =
354                rx.recv().await.expect("agent should emit context usage")
355            {
356                return ContextUsageUpdate {
357                    usage_ratio: usage.usage_ratio,
358                    context_limit: usage.context_limit,
359                    input_tokens: usage.input_tokens,
360                };
361            }
362        }
363    }
364
365    struct ContextUsageUpdate {
366        usage_ratio: Option<f64>,
367        context_limit: Option<u32>,
368        input_tokens: u32,
369    }
370
371    #[tokio::test]
372    async fn system_prompt_preserves_add_order() {
373        let builder = AgentBuilder::new(Arc::new(llm::testing::FakeLlmProvider::new(vec![])))
374            .system_prompt(Prompt::text("first"))
375            .system_prompt(Prompt::text("second"))
376            .system_prompt(Prompt::text("third"));
377
378        let rendered = Prompt::build_all(&builder.prompts).await.unwrap();
379
380        assert_eq!(rendered, "first\n\nsecond\n\nthird");
381    }
382
383    #[tokio::test]
384    async fn from_spec_applies_context_window_and_model_settings() {
385        let settings = ModelSettings { temperature: Some(0.0), max_tokens: Some(128), ..Default::default() };
386        let spec = AgentSpec {
387            name: "alloy".to_string(),
388            description: "alloy".to_string(),
389            model: "ollama:llama3.2,llamacpp:local".to_string(),
390            reasoning_effort: None,
391            model_settings: settings.clone(),
392            context_window: Some(200_000),
393            prompts: vec![],
394            provider_connections: ProviderConnectionOverrides::default(),
395            mcp_config_sources: Vec::new(),
396            exposure: AgentSpecExposure::both(),
397            tools: ToolFilter::default(),
398        };
399
400        let builder = AgentBuilder::from_spec(&spec, vec![], &AgentDeps::default()).await.unwrap();
401
402        assert_eq!(builder.context_window, Some(200_000));
403        assert_eq!(builder.model_settings, settings);
404    }
405
406    #[tokio::test]
407    async fn spawn_applies_model_settings_to_context() {
408        let fake = FakeLlmProvider::with_single_response(vec![
409            LlmResponse::start("msg"),
410            LlmResponse::usage(10, 10),
411            LlmResponse::done(),
412        ]);
413
414        let captured = fake.captured_contexts();
415        let settings = ModelSettings { temperature: Some(0.0), max_tokens: Some(64), ..Default::default() };
416
417        let (tx, mut rx, handle) =
418            AgentBuilder::new(Arc::new(fake)).model_settings(settings.clone()).spawn().await.unwrap();
419
420        tx.send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text("hello")] }))
421            .await
422            .unwrap();
423
424        let _ = next_context_usage(&mut rx).await;
425        let contexts = captured.lock().unwrap();
426
427        assert_eq!(contexts[0].model_settings(), &settings);
428        handle.abort();
429    }
430
431    #[tokio::test]
432    async fn from_spec_accepts_alloy_model_specs() {
433        let spec = AgentSpec {
434            name: "alloy".to_string(),
435            description: "alloy".to_string(),
436            model: "ollama:llama3.2,llamacpp:local".to_string(),
437            reasoning_effort: None,
438            model_settings: ModelSettings::default(),
439            context_window: None,
440            prompts: vec![],
441            provider_connections: ProviderConnectionOverrides::default(),
442            mcp_config_sources: Vec::new(),
443            exposure: AgentSpecExposure::both(),
444            tools: ToolFilter::default(),
445        };
446
447        let builder = AgentBuilder::from_spec(&spec, vec![], &AgentDeps::default()).await;
448        assert!(builder.is_ok());
449    }
450}
aether_core/core/agent_builder.rs

aether_core/core/
agent_builder.rs
