aether_core/core/
agent.rs

1use crate::context::{CompactionConfig, Compactor, TokenTracker};
2use crate::core::PromptCache;
3pub use crate::core::retry_config::RetryConfig;
4use crate::events::{
5    AgentCommand, AgentEvent, AgentObserver, Command, ContextEvent, ContextUsage, LlmCallOutcome, LlmCallPurpose,
6    ModelEvent, ToolEvent, TurnEvent, TurnOutcome, UserCommand,
7};
8use crate::mcp::run_mcp_task::{McpCommand, ToolExecutionEvent};
9use futures::Stream;
10use llm::types::IsoString;
11use llm::{
12    AssistantReasoning, ChatMessage, Context, EncryptedReasoningContent, LlmError, LlmResponse, StopReason,
13    StreamingModelProvider, TokenUsage, ToolCallError, ToolCallRequest, ToolCallResult,
14};
15use std::collections::{HashMap, HashSet, VecDeque};
16use std::pin::Pin;
17use std::sync::Arc;
18use std::time::Duration;
19use tokio::sync::mpsc;
20use tokio::time::sleep;
21use tokio_stream::StreamExt;
22use tokio_stream::StreamMap;
23use tokio_stream::wrappers::ReceiverStream;
24
25/// Internal event type for merging LLM and tool result streams
26#[derive(Debug)]
27enum StreamEvent {
28    LlmRequestStarted { attempt: u32 },
29    Llm(Result<LlmResponse, LlmError>),
30    ToolExecution(ToolExecutionEvent),
31    Command(Command),
32}
33
34type EventStream = Pin<Box<dyn Stream<Item = StreamEvent> + Send>>;
35
36const USER_STREAM_KEY: &str = "user";
37const LLM_STREAM_KEY: &str = "llm";
38
39pub(crate) struct AgentConfig {
40    pub llm: Arc<dyn StreamingModelProvider>,
41    pub context: Context,
42    pub mcp_command_tx: Option<mpsc::Sender<McpCommand>>,
43    pub tool_timeout: Duration,
44    pub compaction_config: Option<CompactionConfig>,
45    pub auto_continue: AutoContinue,
46    pub retry_config: RetryConfig,
47    pub context_window: Option<u32>,
48    pub prompt_cache: PromptCache,
49    pub observers: Vec<Box<dyn AgentObserver>>,
50}
51
52pub struct Agent {
53    llm: Arc<dyn StreamingModelProvider>,
54    context: Context,
55    mcp_command_tx: Option<mpsc::Sender<McpCommand>>,
56    message_tx: mpsc::Sender<AgentEvent>,
57    observers: Vec<Box<dyn AgentObserver>>,
58    streams: StreamMap<String, EventStream>,
59    tool_timeout: Duration,
60    token_tracker: TokenTracker,
61    compaction_config: Option<CompactionConfig>,
62    auto_continue: AutoContinue,
63    retry_config: RetryConfig,
64    active_requests: HashMap<String, ToolCallRequest>,
65    queued_user_messages: VecDeque<Vec<llm::ContentBlock>>,
66    context_window: Option<u32>,
67    prompt_cache: PromptCache,
68    turn_active: bool,
69    llm_call_active: bool,
70}
71
72impl Agent {
73    pub(crate) fn new(
74        config: AgentConfig,
75        command_rx: mpsc::Receiver<Command>,
76        message_tx: mpsc::Sender<AgentEvent>,
77    ) -> Self {
78        let mut streams: StreamMap<String, EventStream> = StreamMap::new();
79        streams
80            .insert(USER_STREAM_KEY.to_string(), Box::pin(ReceiverStream::new(command_rx).map(StreamEvent::Command)));
81
82        let context_limit = config.context_window.or_else(|| config.llm.context_window());
83
84        Self {
85            llm: config.llm,
86            context: config.context,
87            mcp_command_tx: config.mcp_command_tx,
88            message_tx,
89            observers: config.observers,
90            streams,
91            tool_timeout: config.tool_timeout,
92            token_tracker: TokenTracker::new(context_limit),
93            compaction_config: config.compaction_config,
94            auto_continue: config.auto_continue,
95            retry_config: config.retry_config,
96            active_requests: HashMap::new(),
97            queued_user_messages: VecDeque::new(),
98            context_window: config.context_window,
99            prompt_cache: config.prompt_cache,
100            turn_active: false,
101            llm_call_active: false,
102        }
103    }
104
105    pub fn current_model_display_name(&self) -> String {
106        self.llm.display_name()
107    }
108
109    /// Get a reference to the token tracker
110    pub fn token_tracker(&self) -> &TokenTracker {
111        &self.token_tracker
112    }
113
114    pub async fn run(mut self) {
115        let mut state = IterationState::new();
116        self.emit_tool_definitions().await;
117
118        while let Some((_, event)) = self.streams.next().await {
119            match event {
120                StreamEvent::Command(Command::UserCommand(UserCommand::Cancel)) => {
121                    self.on_user_cancel(&mut state).await;
122                }
123
124                StreamEvent::Command(Command::UserCommand(UserCommand::ClearContext)) => {
125                    self.on_user_clear_context(&mut state).await;
126                }
127
128                StreamEvent::Command(Command::UserCommand(UserCommand::Text { content })) => {
129                    if self.is_busy() {
130                        self.queued_user_messages.push_back(content);
131                    } else {
132                        state = IterationState::new();
133                        self.on_user_text(content).await;
134                    }
135                }
136
137                StreamEvent::Command(Command::AgentCommand(AgentCommand::SwitchModel(new_provider))) => {
138                    self.on_switch_model(new_provider).await;
139                }
140
141                StreamEvent::Command(Command::AgentCommand(AgentCommand::UpdateTools(tools))) => {
142                    self.context.set_tools(tools);
143                    self.emit_tool_definitions().await;
144                }
145
146                StreamEvent::Command(Command::AgentCommand(AgentCommand::UpdateMcpInstructions { server, body })) => {
147                    self.on_update_instruction(server, body).await;
148                }
149
150                StreamEvent::Command(Command::AgentCommand(AgentCommand::SetReasoningEffort(effort))) => {
151                    self.context.set_reasoning_effort(effort);
152                }
153
154                StreamEvent::Command(Command::AgentCommand(AgentCommand::ReplaceConversation(messages))) => {
155                    self.on_replace_conversation(messages, &mut state).await;
156                }
157
158                StreamEvent::LlmRequestStarted { attempt } => {
159                    self.begin_chat_call(attempt).await;
160                }
161
162                StreamEvent::Llm(llm_event) => {
163                    self.on_llm_event(llm_event, &mut state).await;
164                }
165
166                StreamEvent::ToolExecution(tool_event) => {
167                    self.on_tool_execution_event(tool_event, &mut state).await;
168                }
169            }
170
171            if state.is_complete() {
172                let Some(id) = state.current_message_id.take() else {
173                    continue;
174                };
175                let iteration = std::mem::replace(&mut state, IterationState::new());
176                self.on_iteration_complete(id, iteration).await;
177            }
178        }
179
180        tracing::debug!("Agent task shutting down - input channel closed");
181    }
182
183    async fn on_iteration_complete(&mut self, id: String, iteration: IterationState) {
184        let IterationState {
185            message_content,
186            reasoning_summary_text,
187            encrypted_reasoning,
188            completed_tool_calls,
189            stop_reason,
190            ..
191        } = iteration;
192        let has_tool_calls = !completed_tool_calls.is_empty();
193        let has_content = !message_content.is_empty() || has_tool_calls;
194        let should_auto_continue = self.auto_continue.should_continue(stop_reason.as_ref());
195
196        if has_content {
197            let reasoning = AssistantReasoning::from_parts(reasoning_summary_text.clone(), encrypted_reasoning);
198            self.update_context(&message_content, reasoning, completed_tool_calls);
199
200            self.emit(AgentEvent::text(&id, &message_content, true)).await;
201
202            if !reasoning_summary_text.is_empty() {
203                self.emit(AgentEvent::thought(&id, &reasoning_summary_text, true)).await;
204            }
205        }
206
207        let has_queued_text = !self.queued_user_messages.is_empty();
208        if has_queued_text {
209            let content: Vec<_> = self.queued_user_messages.drain(..).flatten().collect();
210            self.context.add_message(ChatMessage::User { content, timestamp: IsoString::now() });
211        }
212
213        if has_queued_text || has_tool_calls {
214            self.auto_continue.reset();
215            self.start_next_turn().await;
216        } else if should_auto_continue {
217            self.auto_continue.advance();
218            tracing::info!(
219                "LLM stopped with {:?}, auto-continuing (attempt {}/{})",
220                stop_reason,
221                self.auto_continue.count(),
222                self.auto_continue.max()
223            );
224
225            self.emit(AgentEvent::Turn(TurnEvent::AutoContinue {
226                attempt: self.auto_continue.count(),
227                max_attempts: self.auto_continue.max(),
228            }))
229            .await;
230
231            self.inject_continuation_prompt(&message_content, stop_reason.as_ref());
232            self.start_next_turn().await;
233        } else {
234            tracing::debug!("LLM completed turn with stop reason: {:?}", stop_reason);
235            self.auto_continue.reset();
236            self.finish_turn(TurnOutcome::Completed).await;
237        }
238    }
239
240    async fn start_next_turn(&mut self) {
241        self.maybe_preflight_compact().await;
242        self.start_llm_stream(None, 0).await;
243    }
244
245    async fn on_user_cancel(&mut self, state: &mut IterationState) {
246        self.abort_in_flight_work().await;
247        *state = IterationState::new();
248        self.finish_turn(TurnOutcome::Cancelled).await;
249    }
250
251    async fn on_user_clear_context(&mut self, state: &mut IterationState) {
252        self.abort_in_flight_work().await;
253        self.context.clear_conversation();
254        self.token_tracker.reset_current_usage();
255        self.auto_continue.reset();
256        *state = IterationState::new();
257
258        self.emit(AgentEvent::Context(ContextEvent::Cleared)).await;
259        self.finish_turn(TurnOutcome::Cancelled).await;
260    }
261
262    async fn on_replace_conversation(&mut self, messages: Vec<ChatMessage>, state: &mut IterationState) {
263        self.abort_in_flight_work().await;
264        self.context.replace_conversation(messages);
265        self.auto_continue.reset();
266        *state = IterationState::new();
267        self.emit(self.context_usage_message()).await;
268        self.finish_turn(TurnOutcome::Cancelled).await;
269    }
270
271    async fn on_user_text(&mut self, content: Vec<llm::ContentBlock>) {
272        self.context.add_message(ChatMessage::User { content: content.clone(), timestamp: IsoString::now() });
273        self.auto_continue.reset();
274        self.turn_active = true;
275        self.emit(AgentEvent::Turn(TurnEvent::Started { content })).await;
276        self.start_llm_stream(None, 0).await;
277    }
278
279    async fn on_update_instruction(&mut self, server: String, body: Option<String>) {
280        self.prompt_cache.update_mcp_instruction(server, body);
281        match self.prompt_cache.render().await {
282            Ok(content) => self.context.set_system_content(content),
283            Err(e) => tracing::warn!("Failed to rebuild system prompt after instructions update: {e}"),
284        }
285    }
286
287    async fn on_switch_model(&mut self, new_provider: Box<dyn StreamingModelProvider>) {
288        let previous = self.llm.display_name();
289        let new_context_limit = self.context_window.or_else(|| new_provider.context_window());
290        self.llm = Arc::from(new_provider);
291        self.token_tracker.reset_current_usage();
292        self.token_tracker.set_context_limit(new_context_limit);
293        let new = self.llm.display_name();
294        self.emit(AgentEvent::Model(ModelEvent::Switched { previous, new })).await;
295
296        self.emit(self.context_usage_message()).await;
297    }
298
299    async fn start_llm_stream(&mut self, delay: Option<Duration>, attempt: u32) {
300        self.streams.remove(LLM_STREAM_KEY);
301        let stream: EventStream = match delay {
302            None => {
303                self.begin_chat_call(attempt).await;
304                Box::pin(self.llm.stream_response(&self.context).map(StreamEvent::Llm))
305            }
306            Some(delay) => {
307                self.emit(AgentEvent::Turn(TurnEvent::RetryScheduled {
308                    purpose: LlmCallPurpose::Chat,
309                    attempt,
310                    max_attempts: self.retry_config.max_attempts,
311                    delay_ms: u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
312                }))
313                .await;
314                let llm = Arc::clone(&self.llm);
315                let context = self.context.clone();
316                Box::pin(async_stream::stream! {
317                    sleep(delay).await;
318                    yield StreamEvent::LlmRequestStarted { attempt };
319                    let mut inner = llm.stream_response(&context);
320                    while let Some(item) = inner.next().await {
321                        yield StreamEvent::Llm(item);
322                    }
323                })
324            }
325        };
326        self.streams.insert(LLM_STREAM_KEY.to_string(), stream);
327    }
328
329    async fn on_llm_error(&mut self, error: LlmError, state: &mut IterationState) {
330        let will_retry = error.is_retryable() && state.retry_attempt < self.retry_config.max_attempts;
331        let error_message = error.to_string();
332        self.finish_chat_call(LlmCallOutcome::Failed { error: error_message.clone(), will_retry }).await;
333
334        if !will_retry {
335            self.finish_turn(TurnOutcome::Failed { error: error_message }).await;
336            return;
337        }
338
339        state.retry_attempt += 1;
340        let delay = self.retry_config.compute_delay(state.retry_attempt);
341
342        tracing::warn!(
343            attempt = state.retry_attempt,
344            max_attempts = self.retry_config.max_attempts,
345            delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
346            error = %error,
347            "Retrying LLM request after transient failure"
348        );
349
350        // The previous stream may have emitted partial tool-call deltas
351        // before interrupting so we drop them to ensure we rebuild tool state
352        self.active_requests.clear();
353        self.start_llm_stream(Some(delay), state.retry_attempt).await;
354    }
355
356    fn is_busy(&self) -> bool {
357        self.streams.contains_key(LLM_STREAM_KEY) || !self.active_requests.is_empty()
358    }
359
360    async fn abort_in_flight_work(&mut self) {
361        if self.llm_call_active {
362            self.finish_chat_call(LlmCallOutcome::Cancelled).await;
363        }
364        self.streams.remove(LLM_STREAM_KEY);
365        for stream_key in self.active_requests.keys().cloned().collect::<Vec<_>>() {
366            self.streams.remove(&stream_key);
367        }
368        self.active_requests.clear();
369        self.queued_user_messages.clear();
370    }
371
372    /// Inject a continuation prompt when the LLM stops due to a resumable reason.
373    fn inject_continuation_prompt(&mut self, previous_response: &str, stop_reason: Option<&StopReason>) {
374        if !previous_response.is_empty() {
375            self.context.add_message(ChatMessage::Assistant {
376                content: previous_response.to_string(),
377                reasoning: AssistantReasoning::default(),
378                timestamp: IsoString::now(),
379                tool_calls: Vec::new(),
380            });
381        }
382
383        let reason = stop_reason.map_or_else(|| "Unknown".to_string(), |reason| format!("{reason:?}"));
384
385        self.context.add_message(ChatMessage::User {
386            content: vec![llm::ContentBlock::text(format!(
387                "<system-notification>The LLM API stopped with reason '{reason}'. Continue from where you left off and finish your task.</system-notification>"
388            ))],
389            timestamp: IsoString::now(),
390        });
391    }
392
393    async fn on_llm_event(&mut self, result: Result<LlmResponse, LlmError>, state: &mut IterationState) {
394        use LlmResponse::{
395            Done, EncryptedReasoning, Error, Reasoning, Start, Text, ToolRequestArg, ToolRequestComplete,
396            ToolRequestStart, Usage,
397        };
398
399        let response = match result {
400            Ok(response) => response,
401            Err(e) => {
402                self.on_llm_error(e, state).await;
403                return;
404            }
405        };
406
407        match response {
408            Start { message_id } => {
409                state.on_llm_start(message_id);
410            }
411
412            Text { chunk } => {
413                self.handle_llm_text(chunk, state).await;
414            }
415
416            Reasoning { chunk } => {
417                state.reasoning_summary_text.push_str(&chunk);
418                if let Some(id) = state.current_message_id.clone() {
419                    self.emit(AgentEvent::thought(&id, &chunk, false)).await;
420                }
421            }
422
423            EncryptedReasoning { id, content } => {
424                if let Some(model) = self.llm.model() {
425                    state.encrypted_reasoning = Some(EncryptedReasoningContent { id, model, content });
426                }
427            }
428
429            ToolRequestStart { id, name } => {
430                self.handle_tool_request_start(id, name).await;
431            }
432
433            ToolRequestArg { id, chunk } => {
434                self.handle_tool_request_arg(id, chunk).await;
435            }
436
437            ToolRequestComplete { tool_call } => {
438                self.handle_tool_completion(tool_call, state).await;
439            }
440
441            Done { stop_reason } => {
442                state.llm_done = true;
443                state.stop_reason = stop_reason;
444                self.finish_chat_call(LlmCallOutcome::Completed {
445                    stop_reason: state.stop_reason.clone(),
446                    usage: state.call_usage.take(),
447                })
448                .await;
449            }
450
451            Error { message } => {
452                self.finish_chat_call(LlmCallOutcome::Failed { error: message.clone(), will_retry: false }).await;
453                self.finish_turn(TurnOutcome::Failed { error: message }).await;
454            }
455
456            Usage { tokens: sample } => {
457                self.handle_llm_usage(sample, state).await;
458            }
459        }
460    }
461
462    async fn handle_llm_text(&mut self, chunk: String, state: &mut IterationState) {
463        state.message_content.push_str(&chunk);
464
465        if let Some(id) = state.current_message_id.clone() {
466            self.emit(AgentEvent::text(&id, &chunk, false)).await;
467        }
468    }
469
470    async fn handle_tool_request_start(&mut self, id: String, name: String) {
471        let request = ToolCallRequest { id: id.clone(), name, arguments: String::new() };
472        self.active_requests.insert(id, request.clone());
473
474        self.emit(AgentEvent::Tool(ToolEvent::Call { request })).await;
475    }
476
477    async fn handle_tool_request_arg(&mut self, id: String, chunk: String) {
478        let Some(request) = self.active_requests.get_mut(&id) else {
479            return;
480        };
481        request.arguments.push_str(&chunk);
482
483        self.emit(AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id: id, chunk })).await;
484    }
485
486    async fn handle_tool_completion(&mut self, tool_call: ToolCallRequest, state: &mut IterationState) {
487        state.pending_tool_ids.insert(tool_call.id.clone());
488        debug_assert!(
489            self.active_requests.contains_key(&tool_call.id),
490            "tool call {} should already be in active_requests from handle_tool_request_start",
491            tool_call.id
492        );
493
494        let (tx, rx) = mpsc::channel(100);
495        let stream = ReceiverStream::new(rx).map(StreamEvent::ToolExecution);
496        let stream_key = tool_call.id.clone();
497        self.streams.insert(stream_key, Box::pin(stream));
498
499        if let Some(ref mcp_command_tx) = self.mcp_command_tx {
500            let mcp_future =
501                mcp_command_tx.send(McpCommand::ExecuteTool { request: tool_call, timeout: self.tool_timeout, tx });
502            if let Err(e) = mcp_future.await {
503                tracing::warn!("Failed to send tool request to MCP task: {:?}", e);
504            }
505        }
506    }
507
508    async fn handle_llm_usage(&mut self, sample: TokenUsage, state: &mut IterationState) {
509        state.call_usage = Some(sample);
510        self.token_tracker.record_usage(sample);
511        let ratio_pct = self.token_tracker.usage_ratio().map(|r| r * 100.0);
512        let remaining = self.token_tracker.tokens_remaining();
513        tracing::debug!(?sample, ?ratio_pct, ?remaining, "Token usage");
514
515        self.emit(self.context_usage_message()).await;
516
517        self.maybe_compact_context().await;
518    }
519
520    fn context_usage_message(&self) -> AgentEvent {
521        let last = self.token_tracker.last_usage();
522        AgentEvent::Context(ContextEvent::UsageUpdated {
523            usage: ContextUsage {
524                usage_ratio: self.token_tracker.usage_ratio(),
525                context_limit: self.token_tracker.context_limit(),
526                input_tokens: last.input_tokens,
527                output_tokens: last.output_tokens,
528                cache_read_tokens: last.cache_read_tokens,
529                cache_creation_tokens: last.cache_creation_tokens,
530                reasoning_tokens: last.reasoning_tokens,
531                total_input_tokens: self.token_tracker.total_input_tokens(),
532                total_output_tokens: self.token_tracker.total_output_tokens(),
533                total_cache_read_tokens: self.token_tracker.total_cache_read_tokens(),
534                total_cache_creation_tokens: self.token_tracker.total_cache_creation_tokens(),
535                total_reasoning_tokens: self.token_tracker.total_reasoning_tokens(),
536            },
537        })
538    }
539
540    /// Pre-flight check: estimate context size and compact proactively if it would
541    /// overflow before the LLM even sees it. This catches the case where large tool
542    /// results push context past the limit before usage-based compaction can fire.
543    async fn maybe_preflight_compact(&mut self) {
544        let Some(context_limit) = self.token_tracker.context_limit() else {
545            return;
546        };
547        let Some(config) = self.compaction_config.as_ref() else {
548            return;
549        };
550        let estimated = self.context.estimated_token_count();
551        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
552        let threshold = (f64::from(context_limit) * config.threshold).ceil() as u32;
553        if estimated >= threshold {
554            tracing::info!(
555                "Pre-flight compaction triggered: estimated {estimated} tokens >= {:.1}% of {context_limit} limit",
556                config.threshold * 100.0
557            );
558            if let CompactionOutcome::Failed(e) = self.compact_context().await {
559                tracing::warn!("Pre-flight compaction failed: {e}");
560            }
561        }
562    }
563
564    /// Check if compaction is needed and perform it if so.
565    async fn maybe_compact_context(&mut self) {
566        if !self.compaction_config.as_ref().is_some_and(|config| self.token_tracker.should_compact(config.threshold)) {
567            return;
568        }
569
570        if let CompactionOutcome::Failed(error_message) = self.compact_context().await {
571            tracing::warn!("Context compaction failed: {}", error_message);
572        }
573    }
574
575    async fn compact_context(&mut self) -> CompactionOutcome {
576        let Some(ref _config) = self.compaction_config else {
577            tracing::warn!("Context compaction requested but compaction is disabled");
578            return CompactionOutcome::SkippedDisabled;
579        };
580
581        match self.token_tracker.usage_ratio() {
582            Some(usage_ratio) => {
583                tracing::info!(
584                    "Starting context compaction - {} messages, {:.1}% of context limit",
585                    self.context.message_count(),
586                    usage_ratio * 100.0
587                );
588            }
589            None => {
590                tracing::info!(
591                    "Starting context compaction - {} messages (context limit unknown)",
592                    self.context.message_count(),
593                );
594            }
595        }
596
597        self.emit(AgentEvent::Context(ContextEvent::CompactionStarted { message_count: self.context.message_count() }))
598            .await;
599
600        let compactor = Compactor::new(self.llm.clone());
601
602        self.emit(self.llm_call_started(LlmCallPurpose::Compaction, 0)).await;
603        match compactor.compact(&self.context).await {
604            Ok(result) => {
605                tracing::info!("Context compacted: {} messages removed", result.messages_removed);
606                self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded {
607                    purpose: LlmCallPurpose::Compaction,
608                    outcome: LlmCallOutcome::Completed { stop_reason: None, usage: result.usage },
609                }))
610                .await;
611
612                self.context = result.context;
613                self.token_tracker.reset_current_usage();
614
615                self.emit(AgentEvent::Context(ContextEvent::CompactionResult {
616                    summary: result.summary,
617                    messages_removed: result.messages_removed,
618                }))
619                .await;
620                CompactionOutcome::Compacted
621            }
622            Err(e) => {
623                self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded {
624                    purpose: LlmCallPurpose::Compaction,
625                    outcome: LlmCallOutcome::Failed { error: e.to_string(), will_retry: false },
626                }))
627                .await;
628                CompactionOutcome::Failed(e.to_string())
629            }
630        }
631    }
632
633    async fn on_tool_execution_event(&mut self, event: ToolExecutionEvent, state: &mut IterationState) {
634        match event {
635            ToolExecutionEvent::Started { tool_id, tool_name } => {
636                tracing::debug!("Tool execution started: {} ({})", tool_name, tool_id);
637                self.emit(AgentEvent::Tool(ToolEvent::ExecutionStarted { tool_id, tool_name })).await;
638            }
639
640            ToolExecutionEvent::Progress { tool_id, progress } => {
641                tracing::debug!(
642                    "Tool progress for {}: {}/{}",
643                    tool_id,
644                    progress.progress,
645                    progress.total.unwrap_or(0.0)
646                );
647
648                if let Some(request) = self.active_requests.get(&tool_id).cloned() {
649                    self.emit(AgentEvent::Tool(ToolEvent::Progress {
650                        request,
651                        progress: progress.progress,
652                        total: progress.total,
653                        message: progress.message.clone(),
654                    }))
655                    .await;
656                }
657            }
658
659            ToolExecutionEvent::Complete { tool_id: _, result, result_meta } => match result {
660                Ok(tool_result) => {
661                    tracing::debug!("Tool result received: {} -> {}", tool_result.name, tool_result.result.len());
662
663                    if state.pending_tool_ids.remove(&tool_result.id) {
664                        self.active_requests.remove(&tool_result.id);
665                        state.completed_tool_calls.push(Ok(tool_result.clone()));
666
667                        self.emit(AgentEvent::Tool(ToolEvent::Result { result: tool_result, result_meta })).await;
668                    } else {
669                        tracing::debug!("Ignoring stale tool result for id: {}", tool_result.id);
670                    }
671                }
672
673                Err(tool_error) => {
674                    if state.pending_tool_ids.remove(&tool_error.id) {
675                        self.active_requests.remove(&tool_error.id);
676                        state.completed_tool_calls.push(Err(tool_error.clone()));
677
678                        self.emit(AgentEvent::Tool(ToolEvent::Error { error: tool_error })).await;
679                    }
680                }
681            },
682        }
683    }
684
685    fn update_context(
686        &mut self,
687        message_content: &str,
688        reasoning: AssistantReasoning,
689        completed_tools: Vec<Result<ToolCallResult, ToolCallError>>,
690    ) {
691        self.context.push_assistant_turn(message_content, reasoning, completed_tools);
692    }
693
694    async fn emit_tool_definitions(&mut self) {
695        let tools = self.context.tools().clone();
696        if !tools.is_empty() {
697            self.emit(AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools })).await;
698        }
699    }
700
701    async fn emit(&mut self, message: AgentEvent) {
702        for observer in &mut self.observers {
703            observer.on_event(&message);
704        }
705
706        if let Err(e) = self.message_tx.send(message).await {
707            tracing::warn!("Failed to send agent message: {e:?}");
708        }
709    }
710
711    async fn finish_turn(&mut self, outcome: TurnOutcome) {
712        if std::mem::take(&mut self.turn_active) {
713            self.emit(AgentEvent::turn_ended(outcome)).await;
714        }
715    }
716
717    async fn begin_chat_call(&mut self, attempt: u32) {
718        self.llm_call_active = true;
719        self.emit(self.llm_call_started(LlmCallPurpose::Chat, attempt)).await;
720    }
721
722    async fn finish_chat_call(&mut self, outcome: LlmCallOutcome) {
723        if std::mem::take(&mut self.llm_call_active) {
724            self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Chat, outcome })).await;
725        }
726    }
727
728    fn llm_call_started(&self, purpose: LlmCallPurpose, attempt: u32) -> AgentEvent {
729        let model = self.llm.model();
730        AgentEvent::Turn(TurnEvent::LlmCallStarted {
731            purpose,
732            provider: model.as_ref().map(|m| m.provider().to_string()),
733            model: model.map(|m| m.model_id().into_owned()),
734            display_name: self.llm.display_name(),
735            attempt,
736            max_attempts: self.retry_config.max_attempts,
737        })
738    }
739}
740
741#[derive(Debug, Clone, PartialEq, Eq)]
742enum CompactionOutcome {
743    Compacted,
744    SkippedDisabled,
745    Failed(String),
746}
747
748pub(crate) struct AutoContinue {
749    max: u32,
750    count: u32,
751}
752
753impl AutoContinue {
754    pub(crate) fn new(max: u32) -> Self {
755        Self { max, count: 0 }
756    }
757
758    fn reset(&mut self) {
759        self.count = 0;
760    }
761
762    fn should_continue(&self, stop_reason: Option<&StopReason>) -> bool {
763        matches!(stop_reason, Some(StopReason::Length)) && self.count < self.max
764    }
765
766    fn advance(&mut self) {
767        self.count += 1;
768    }
769
770    fn count(&self) -> u32 {
771        self.count
772    }
773
774    fn max(&self) -> u32 {
775        self.max
776    }
777}
778
779#[derive(Debug)]
780struct IterationState {
781    current_message_id: Option<String>,
782    message_content: String,
783    reasoning_summary_text: String,
784    encrypted_reasoning: Option<EncryptedReasoningContent>,
785    pending_tool_ids: HashSet<String>,
786    completed_tool_calls: Vec<Result<ToolCallResult, ToolCallError>>,
787    llm_done: bool,
788    stop_reason: Option<StopReason>,
789    retry_attempt: u32,
790    call_usage: Option<TokenUsage>,
791}
792
793impl IterationState {
794    fn new() -> Self {
795        Self {
796            current_message_id: None,
797            message_content: String::new(),
798            reasoning_summary_text: String::new(),
799            encrypted_reasoning: None,
800            pending_tool_ids: HashSet::new(),
801            completed_tool_calls: Vec::new(),
802            llm_done: false,
803            stop_reason: None,
804            retry_attempt: 0,
805            call_usage: None,
806        }
807    }
808
809    fn on_llm_start(&mut self, message_id: String) {
810        self.current_message_id = Some(message_id);
811        self.message_content.clear();
812        self.reasoning_summary_text.clear();
813        self.encrypted_reasoning = None;
814        self.stop_reason = None;
815        self.call_usage = None;
816    }
817
818    fn is_complete(&self) -> bool {
819        self.llm_done && self.pending_tool_ids.is_empty()
820    }
821}
822
823#[cfg(test)]
824mod tests {
825    use crate::core::{AgentBuilder, Prompt};
826
827    use super::*;
828    use llm::{ContentBlock, testing::FakeLlmProvider};
829    use tokio::sync::mpsc;
830
831    #[tokio::test]
832    async fn replace_conversation_preserves_system_prompt_for_next_request() {
833        let llm = FakeLlmProvider::with_single_response(vec![LlmResponse::start("msg"), LlmResponse::done()]);
834
835        let captured_contexts = llm.captured_contexts();
836        let (tx, mut rx, handle) =
837            AgentBuilder::new(Arc::new(llm)).system_prompt(Prompt::text("original system")).spawn().await.unwrap();
838
839        tx.send(Command::AgentCommand(AgentCommand::ReplaceConversation(vec![
840            ChatMessage::User { content: vec![ContentBlock::text("old user")], timestamp: IsoString::now() },
841            ChatMessage::Assistant {
842                content: "old assistant".to_string(),
843                reasoning: AssistantReasoning::default(),
844                timestamp: IsoString::now(),
845                tool_calls: vec![],
846            },
847        ])))
848        .await
849        .unwrap();
850
851        tx.send(Command::UserCommand(UserCommand::Text { content: vec![ContentBlock::text("new user")] }))
852            .await
853            .unwrap();
854
855        while let Some(message) = rx.recv().await {
856            if matches!(message, AgentEvent::Turn(TurnEvent::Ended { .. })) {
857                break;
858            }
859        }
860
861        let contexts = captured_contexts.lock().unwrap();
862        let messages = contexts.last().expect("provider should receive a context").messages();
863        assert!(matches!(messages[0], ChatMessage::System { ref content, .. } if content == "original system"));
864        assert!(
865            matches!(messages[1], ChatMessage::User { ref content, .. } if content == &vec![llm::ContentBlock::text("old user")])
866        );
867        assert!(matches!(messages[2], ChatMessage::Assistant { ref content, .. } if content == "old assistant"));
868        assert!(
869            matches!(messages[3], ChatMessage::User { ref content, .. } if content == &vec![llm::ContentBlock::text("new user")])
870        );
871        handle.abort();
872    }
873
874    #[tokio::test]
875    async fn replace_conversation_preserves_token_usage() {
876        let llm = FakeLlmProvider::new(vec![vec![
877            LlmResponse::start("msg"),
878            LlmResponse::usage(800, 10),
879            LlmResponse::done(),
880        ]])
881        .with_context_window(Some(1000));
882        let (tx, mut rx, handle) = AgentBuilder::new(Arc::new(llm)).spawn().await.unwrap();
883
884        tx.send(Command::UserCommand(UserCommand::Text { content: vec![llm::ContentBlock::text("first user")] }))
885            .await
886            .unwrap();
887
888        while let Some(message) = rx.recv().await {
889            if matches!(message, AgentEvent::Turn(TurnEvent::Ended { .. })) {
890                break;
891            }
892        }
893
894        tx.send(Command::AgentCommand(AgentCommand::ReplaceConversation(vec![ChatMessage::User {
895            content: vec![ContentBlock::text("replacement user")],
896            timestamp: IsoString::now(),
897        }])))
898        .await
899        .unwrap();
900
901        let Some(AgentEvent::Context(ContextEvent::UsageUpdated { usage })) = rx.recv().await else {
902            panic!("expected context usage update after conversation replacement");
903        };
904
905        assert_eq!(usage.input_tokens, 800);
906        assert_eq!(usage.usage_ratio, Some(0.8));
907        handle.abort();
908    }
909
910    #[tokio::test]
911    async fn test_preflight_compaction_uses_configured_threshold() {
912        let llm = Arc::new(
913            FakeLlmProvider::with_single_response(vec![
914                LlmResponse::start("summary"),
915                LlmResponse::text("summary"),
916                LlmResponse::done(),
917            ])
918            .with_context_window(Some(100)),
919        );
920        let context = Context::new(
921            vec![ChatMessage::User {
922                content: vec![llm::ContentBlock::text("x".repeat(344))],
923                timestamp: IsoString::now(),
924            }],
925            vec![],
926        );
927        let (user_tx, user_rx) = mpsc::channel(1);
928        let (message_tx, _message_rx) = mpsc::channel(8);
929        drop(user_tx);
930
931        let mut agent = Agent::new(
932            AgentConfig {
933                llm,
934                context,
935                mcp_command_tx: None,
936                tool_timeout: Duration::from_secs(1),
937                compaction_config: Some(CompactionConfig::with_threshold(0.85)),
938                auto_continue: AutoContinue::new(0),
939                retry_config: RetryConfig::disabled(),
940                context_window: None,
941                prompt_cache: PromptCache::new(vec![]),
942                observers: Vec::new(),
943            },
944            user_rx,
945            message_tx,
946        );
947
948        agent.maybe_preflight_compact().await;
949
950        assert!(
951            matches!(
952                agent.context.messages().as_slice(),
953                [ChatMessage::Summary { content, .. }] if content == "summary"
954            ),
955            "expected context to be compacted, got {:?}",
956            agent.context.messages()
957        );
958    }
959}
aether_core/core/agent.rs

aether_core/core/
agent.rs
