aether_core/core/
agent.rs

1use crate::context::{CompactionConfig, CompactionError, CompactionResult, Compactor, TokenTracker};
2use crate::core::PromptCache;
3use crate::core::prompt_cache_key::derive_prompt_cache_key;
4pub use crate::core::retry_config::RetryConfig;
5use crate::events::{
6    AgentCommand, AgentEvent, AgentObserver, Command, CompactionOutcome, ContextEvent, ContextUsage, LlmCallOutcome,
7    LlmCallPurpose, ModelEvent, ToolEvent, TurnEvent, TurnOutcome, UserCommand,
8};
9use crate::mcp::run_mcp_task::{McpCommand, ToolExecutionEvent};
10use futures::Stream;
11use llm::types::IsoString;
12use llm::{
13    AssistantReasoning, ChatMessage, Context, EncryptedReasoningContent, LlmError, LlmResponse, StopReason,
14    StreamingModelProvider, TokenUsage, ToolCallError, ToolCallRequest, ToolCallResult,
15};
16use std::collections::{HashMap, HashSet, VecDeque};
17use std::pin::Pin;
18use std::sync::Arc;
19use std::time::Duration;
20use tokio::sync::mpsc;
21use tokio::time::sleep;
22use tokio_stream::StreamExt;
23use tokio_stream::StreamMap;
24use tokio_stream::wrappers::ReceiverStream;
25
26/// Internal event type for merging LLM and tool result streams
27#[derive(Debug)]
28enum StreamEvent {
29    LlmRequestStarted { attempt: u32 },
30    Llm(Result<LlmResponse, LlmError>),
31    ToolExecution(ToolExecutionEvent),
32    Command(Command),
33    Compaction(Result<CompactionResult, CompactionError>),
34}
35
36type EventStream = Pin<Box<dyn Stream<Item = StreamEvent> + Send>>;
37
38/// Keys for the merged stream map. Tool-call ids come from the provider, so a
39/// typed key keeps them from colliding with the reserved streams.
40#[derive(Debug, Clone, PartialEq, Eq, Hash)]
41enum StreamKey {
42    User,
43    Llm,
44    Compaction,
45    Tool(String),
46}
47
48pub(crate) struct AgentConfig {
49    pub llm: Arc<dyn StreamingModelProvider>,
50    pub context: Context,
51    pub mcp_command_tx: Option<mpsc::Sender<McpCommand>>,
52    pub tool_timeout: Duration,
53    pub compaction_config: Option<CompactionConfig>,
54    pub auto_continue: AutoContinue,
55    pub retry_config: RetryConfig,
56    pub context_window: Option<u32>,
57    pub prompt_cache: PromptCache,
58    pub observers: Vec<Box<dyn AgentObserver>>,
59}
60
61pub struct Agent {
62    llm: Arc<dyn StreamingModelProvider>,
63    context: Context,
64    mcp_command_tx: Option<mpsc::Sender<McpCommand>>,
65    message_tx: mpsc::Sender<AgentEvent>,
66    observers: Vec<Box<dyn AgentObserver>>,
67    streams: StreamMap<StreamKey, EventStream>,
68    tool_timeout: Duration,
69    token_tracker: TokenTracker,
70    compaction_config: Option<CompactionConfig>,
71    auto_continue: AutoContinue,
72    retry_config: RetryConfig,
73    active_requests: HashMap<String, ToolCallRequest>,
74    queued_user_messages: VecDeque<Vec<llm::ContentBlock>>,
75    pending_turn_messages: Vec<Vec<llm::ContentBlock>>,
76    context_window: Option<u32>,
77    prompt_cache: PromptCache,
78    turn_active: bool,
79    llm_call_active: bool,
80}
81
82impl Agent {
83    pub(crate) fn new(
84        config: AgentConfig,
85        command_rx: mpsc::Receiver<Command>,
86        message_tx: mpsc::Sender<AgentEvent>,
87    ) -> Self {
88        let mut streams: StreamMap<StreamKey, EventStream> = StreamMap::new();
89        streams.insert(StreamKey::User, Box::pin(ReceiverStream::new(command_rx).map(StreamEvent::Command)));
90
91        let context_limit = config.context_window.or_else(|| config.llm.context_window());
92
93        Self {
94            llm: config.llm,
95            context: config.context,
96            mcp_command_tx: config.mcp_command_tx,
97            message_tx,
98            observers: config.observers,
99            streams,
100            tool_timeout: config.tool_timeout,
101            token_tracker: TokenTracker::new(context_limit),
102            compaction_config: config.compaction_config,
103            auto_continue: config.auto_continue,
104            retry_config: config.retry_config,
105            active_requests: HashMap::new(),
106            queued_user_messages: VecDeque::new(),
107            pending_turn_messages: Vec::new(),
108            context_window: config.context_window,
109            prompt_cache: config.prompt_cache,
110            turn_active: false,
111            llm_call_active: false,
112        }
113    }
114
115    pub fn current_model_display_name(&self) -> String {
116        self.llm.display_name()
117    }
118
119    /// Get a reference to the token tracker
120    pub fn token_tracker(&self) -> &TokenTracker {
121        &self.token_tracker
122    }
123
124    pub async fn run(mut self) {
125        let mut state = IterationState::new();
126        self.emit_tool_definitions().await;
127
128        while let Some((_, event)) = self.streams.next().await {
129            match event {
130                StreamEvent::Command(Command::UserCommand(UserCommand::Cancel)) => {
131                    self.on_user_cancel(&mut state).await;
132                }
133
134                StreamEvent::Command(Command::UserCommand(UserCommand::ClearContext)) => {
135                    self.on_user_clear_context(&mut state).await;
136                }
137
138                StreamEvent::Command(Command::UserCommand(UserCommand::Text { content })) => {
139                    if self.is_busy() {
140                        self.queued_user_messages.push_back(content);
141                    } else {
142                        state = IterationState::new();
143                        self.on_user_text(content).await;
144                    }
145                }
146
147                StreamEvent::Command(Command::AgentCommand(AgentCommand::SwitchModel(new_provider))) => {
148                    self.on_switch_model(new_provider).await;
149                }
150
151                StreamEvent::Command(Command::AgentCommand(AgentCommand::UpdateTools(tools))) => {
152                    self.context.set_tools(tools);
153                    self.emit_tool_definitions().await;
154                }
155
156                StreamEvent::Command(Command::AgentCommand(AgentCommand::UpdateMcpInstructions { server, body })) => {
157                    self.on_update_instruction(server, body).await;
158                }
159
160                StreamEvent::Command(Command::AgentCommand(AgentCommand::SetReasoningEffort(effort))) => {
161                    self.context.set_reasoning_effort(effort);
162                }
163
164                StreamEvent::Command(Command::AgentCommand(AgentCommand::ReplaceConversation(messages))) => {
165                    self.on_replace_conversation(messages, &mut state).await;
166                }
167
168                StreamEvent::LlmRequestStarted { attempt } => {
169                    self.begin_chat_call(attempt).await;
170                }
171
172                StreamEvent::Llm(llm_event) => {
173                    self.on_llm_event(llm_event, &mut state).await;
174                }
175
176                StreamEvent::ToolExecution(tool_event) => {
177                    self.on_tool_execution_event(tool_event, &mut state).await;
178                }
179
180                StreamEvent::Compaction(result) => {
181                    self.on_compaction_complete(result).await;
182                }
183            }
184
185            if state.is_complete() {
186                let Some(id) = state.current_message_id.take() else {
187                    continue;
188                };
189                let iteration = std::mem::replace(&mut state, IterationState::new());
190                self.on_iteration_complete(id, iteration).await;
191            }
192        }
193
194        tracing::debug!("Agent task shutting down - input channel closed");
195    }
196
197    async fn on_iteration_complete(&mut self, id: String, iteration: IterationState) {
198        let IterationState {
199            message_content,
200            reasoning_summary_text,
201            encrypted_reasoning,
202            completed_tool_calls,
203            stop_reason,
204            ..
205        } = iteration;
206        let has_tool_calls = !completed_tool_calls.is_empty();
207        let has_content = !message_content.is_empty() || has_tool_calls;
208        let should_auto_continue = self.auto_continue.should_continue(stop_reason.as_ref());
209
210        if has_content {
211            let reasoning = AssistantReasoning::from_parts(reasoning_summary_text.clone(), encrypted_reasoning);
212            self.update_context(&message_content, reasoning, completed_tool_calls);
213
214            self.emit(AgentEvent::text(&id, &message_content, true)).await;
215
216            if !reasoning_summary_text.is_empty() {
217                self.emit(AgentEvent::thought(&id, &reasoning_summary_text, true)).await;
218            }
219        }
220
221        let has_queued_text = !self.queued_user_messages.is_empty();
222        if has_queued_text || has_tool_calls {
223            self.auto_continue.reset();
224            self.start_next_turn().await;
225        } else if should_auto_continue {
226            self.auto_continue.advance();
227            tracing::info!(
228                "LLM stopped with {:?}, auto-continuing (attempt {}/{})",
229                stop_reason,
230                self.auto_continue.count(),
231                self.auto_continue.max()
232            );
233
234            self.emit(AgentEvent::Turn(TurnEvent::AutoContinue {
235                attempt: self.auto_continue.count(),
236                max_attempts: self.auto_continue.max(),
237            }))
238            .await;
239
240            self.inject_continuation_prompt(&message_content, stop_reason.as_ref());
241            self.start_next_turn().await;
242        } else {
243            tracing::debug!("LLM completed turn with stop reason: {:?}", stop_reason);
244            self.auto_continue.reset();
245            self.finish_turn(TurnOutcome::Completed).await;
246        }
247    }
248
249    async fn start_next_turn(&mut self) {
250        self.pending_turn_messages.extend(self.queued_user_messages.drain(..));
251        if self.compaction_needed() {
252            self.begin_compaction().await;
253        } else {
254            self.start_chat_turn().await;
255        }
256    }
257
258    async fn start_chat_turn(&mut self) {
259        self.commit_pending_turn_messages();
260        self.start_llm_stream(None, 0).await;
261    }
262
263    async fn on_user_cancel(&mut self, state: &mut IterationState) {
264        self.abort_in_flight_work().await;
265        self.commit_pending_turn_messages();
266        self.queued_user_messages.clear();
267        *state = IterationState::new();
268        self.finish_turn(TurnOutcome::Cancelled).await;
269    }
270
271    async fn on_user_clear_context(&mut self, state: &mut IterationState) {
272        self.abort_in_flight_work().await;
273        self.queued_user_messages.clear();
274        self.pending_turn_messages.clear();
275        self.context.clear_conversation();
276        self.token_tracker.reset_current_usage();
277        self.auto_continue.reset();
278        *state = IterationState::new();
279
280        self.emit(AgentEvent::Context(ContextEvent::Cleared)).await;
281        self.finish_turn(TurnOutcome::Cancelled).await;
282    }
283
284    async fn on_replace_conversation(&mut self, messages: Vec<ChatMessage>, state: &mut IterationState) {
285        self.abort_in_flight_work().await;
286        self.queued_user_messages.clear();
287        self.pending_turn_messages.clear();
288        self.context.replace_conversation(messages);
289        self.auto_continue.reset();
290        *state = IterationState::new();
291        self.emit(self.context_usage_message()).await;
292        self.finish_turn(TurnOutcome::Cancelled).await;
293    }
294
295    async fn on_user_text(&mut self, content: Vec<llm::ContentBlock>) {
296        self.auto_continue.reset();
297        self.turn_active = true;
298        self.emit(AgentEvent::Turn(TurnEvent::Started { content: content.clone() })).await;
299        self.queued_user_messages.push_back(content);
300        self.start_next_turn().await;
301    }
302
303    async fn on_update_instruction(&mut self, server: String, body: Option<String>) {
304        self.prompt_cache.update_mcp_instruction(server, body);
305        match self.prompt_cache.render().await {
306            Ok(content) => self.context.set_system_content(content),
307            Err(e) => tracing::warn!("Failed to rebuild system prompt after instructions update: {e}"),
308        }
309    }
310
311    async fn on_switch_model(&mut self, new_provider: Box<dyn StreamingModelProvider>) {
312        let previous = self.llm.display_name();
313        let new_context_limit = self.context_window.or_else(|| new_provider.context_window());
314        self.llm = Arc::from(new_provider);
315        self.token_tracker.reset_current_usage();
316        self.token_tracker.set_context_limit(new_context_limit);
317        let new = self.llm.display_name();
318        self.emit(AgentEvent::Model(ModelEvent::Switched { previous, new })).await;
319
320        self.emit(self.context_usage_message()).await;
321    }
322
323    async fn start_llm_stream(&mut self, delay: Option<Duration>, attempt: u32) {
324        self.refresh_prompt_cache_key();
325        self.streams.remove(&StreamKey::Llm);
326        let stream: EventStream = match delay {
327            None => {
328                self.begin_chat_call(attempt).await;
329                Box::pin(self.llm.stream_response(&self.context).map(StreamEvent::Llm))
330            }
331            Some(delay) => {
332                self.emit(AgentEvent::Turn(TurnEvent::RetryScheduled {
333                    purpose: LlmCallPurpose::Chat,
334                    attempt,
335                    max_attempts: self.retry_config.max_attempts,
336                    delay_ms: u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
337                }))
338                .await;
339                let llm = Arc::clone(&self.llm);
340                let context = self.context.clone();
341                Box::pin(async_stream::stream! {
342                    sleep(delay).await;
343                    yield StreamEvent::LlmRequestStarted { attempt };
344                    let mut inner = llm.stream_response(&context);
345                    while let Some(item) = inner.next().await {
346                        yield StreamEvent::Llm(item);
347                    }
348                })
349            }
350        };
351        self.streams.insert(StreamKey::Llm, stream);
352    }
353
354    async fn on_llm_error(&mut self, error: LlmError, state: &mut IterationState) {
355        let will_retry = error.is_retryable() && state.retry_attempt < self.retry_config.max_attempts;
356        let error_message = error.to_string();
357        self.finish_chat_call(LlmCallOutcome::Failed { error: error_message.clone(), will_retry }).await;
358
359        if !will_retry {
360            self.finish_turn(TurnOutcome::Failed { error: error_message }).await;
361            return;
362        }
363
364        state.retry_attempt += 1;
365        let delay = self.retry_config.compute_delay(state.retry_attempt);
366
367        tracing::warn!(
368            attempt = state.retry_attempt,
369            max_attempts = self.retry_config.max_attempts,
370            delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
371            error = %error,
372            "Retrying LLM request after transient failure"
373        );
374
375        // The previous stream may have emitted partial tool-call deltas
376        // before interrupting so we drop them to ensure we rebuild tool state
377        self.active_requests.clear();
378        self.start_llm_stream(Some(delay), state.retry_attempt).await;
379    }
380
381    fn is_busy(&self) -> bool {
382        self.streams.contains_key(&StreamKey::Llm)
383            || self.streams.contains_key(&StreamKey::Compaction)
384            || !self.active_requests.is_empty()
385    }
386
387    async fn abort_in_flight_work(&mut self) {
388        if self.llm_call_active {
389            self.finish_chat_call(LlmCallOutcome::Cancelled).await;
390        }
391        if self.streams.remove(&StreamKey::Compaction).is_some() {
392            self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded {
393                purpose: LlmCallPurpose::Compaction,
394                outcome: LlmCallOutcome::Cancelled,
395            }))
396            .await;
397            self.emit(AgentEvent::Context(ContextEvent::CompactionEnded { outcome: CompactionOutcome::Cancelled }))
398                .await;
399        }
400        self.streams.remove(&StreamKey::Llm);
401        for tool_id in self.active_requests.keys().cloned().collect::<Vec<_>>() {
402            self.streams.remove(&StreamKey::Tool(tool_id));
403        }
404        self.active_requests.clear();
405    }
406
407    /// Inject a continuation prompt when the LLM stops due to a resumable reason.
408    fn inject_continuation_prompt(&mut self, previous_response: &str, stop_reason: Option<&StopReason>) {
409        if !previous_response.is_empty() {
410            self.context.add_message(ChatMessage::Assistant {
411                content: previous_response.to_string(),
412                reasoning: AssistantReasoning::default(),
413                timestamp: IsoString::now(),
414                tool_calls: Vec::new(),
415            });
416        }
417
418        let reason = stop_reason.map_or_else(|| "Unknown".to_string(), |reason| format!("{reason:?}"));
419
420        self.context.add_message(ChatMessage::User {
421            content: vec![llm::ContentBlock::text(format!(
422                "<system-notification>The LLM API stopped with reason '{reason}'. Continue from where you left off and finish your task.</system-notification>"
423            ))],
424            timestamp: IsoString::now(),
425        });
426    }
427
428    async fn on_llm_event(&mut self, result: Result<LlmResponse, LlmError>, state: &mut IterationState) {
429        use LlmResponse::{
430            Done, EncryptedReasoning, Error, Reasoning, Start, Text, ToolRequestArg, ToolRequestComplete,
431            ToolRequestStart, Usage,
432        };
433
434        let response = match result {
435            Ok(response) => response,
436            Err(e) => {
437                self.on_llm_error(e, state).await;
438                return;
439            }
440        };
441
442        match response {
443            Start { message_id } => {
444                state.on_llm_start(message_id);
445            }
446
447            Text { chunk } => {
448                self.handle_llm_text(chunk, state).await;
449            }
450
451            Reasoning { chunk } => {
452                state.reasoning_summary_text.push_str(&chunk);
453                if let Some(id) = state.current_message_id.clone() {
454                    self.emit(AgentEvent::thought(&id, &chunk, false)).await;
455                }
456            }
457
458            EncryptedReasoning { id, content } => {
459                if let Some(model) = self.llm.model() {
460                    state.encrypted_reasoning = Some(EncryptedReasoningContent { id, model, content });
461                }
462            }
463
464            ToolRequestStart { id, name } => {
465                self.handle_tool_request_start(id, name).await;
466            }
467
468            ToolRequestArg { id, chunk } => {
469                self.handle_tool_request_arg(id, chunk).await;
470            }
471
472            ToolRequestComplete { tool_call } => {
473                self.handle_tool_completion(tool_call, state).await;
474            }
475
476            Done { stop_reason } => {
477                state.llm_done = true;
478                state.stop_reason = stop_reason;
479                self.finish_chat_call(LlmCallOutcome::Completed {
480                    stop_reason: state.stop_reason.clone(),
481                    usage: state.call_usage.take(),
482                })
483                .await;
484            }
485
486            Error { message } => {
487                self.finish_chat_call(LlmCallOutcome::Failed { error: message.clone(), will_retry: false }).await;
488                self.finish_turn(TurnOutcome::Failed { error: message }).await;
489            }
490
491            Usage { tokens: sample } => {
492                self.handle_llm_usage(sample, state).await;
493            }
494        }
495    }
496
497    async fn handle_llm_text(&mut self, chunk: String, state: &mut IterationState) {
498        state.message_content.push_str(&chunk);
499
500        if let Some(id) = state.current_message_id.clone() {
501            self.emit(AgentEvent::text(&id, &chunk, false)).await;
502        }
503    }
504
505    async fn handle_tool_request_start(&mut self, id: String, name: String) {
506        let request = ToolCallRequest { id: id.clone(), name, arguments: String::new() };
507        self.active_requests.insert(id, request.clone());
508
509        self.emit(AgentEvent::Tool(ToolEvent::Call { request })).await;
510    }
511
512    async fn handle_tool_request_arg(&mut self, id: String, chunk: String) {
513        let Some(request) = self.active_requests.get_mut(&id) else {
514            return;
515        };
516        request.arguments.push_str(&chunk);
517
518        self.emit(AgentEvent::Tool(ToolEvent::CallUpdate { tool_call_id: id, chunk })).await;
519    }
520
521    async fn handle_tool_completion(&mut self, tool_call: ToolCallRequest, state: &mut IterationState) {
522        state.pending_tool_ids.insert(tool_call.id.clone());
523        debug_assert!(
524            self.active_requests.contains_key(&tool_call.id),
525            "tool call {} should already be in active_requests from handle_tool_request_start",
526            tool_call.id
527        );
528
529        let (tx, rx) = mpsc::channel(100);
530        let stream = ReceiverStream::new(rx).map(StreamEvent::ToolExecution);
531        self.streams.insert(StreamKey::Tool(tool_call.id.clone()), Box::pin(stream));
532
533        let tool_id = tool_call.id.clone();
534        tracing::debug!("Tool execution started: {} ({})", tool_call.name, tool_id);
535        self.emit(AgentEvent::Tool(ToolEvent::ExecutionStarted {
536            tool_id: tool_id.clone(),
537            tool_name: tool_call.name.clone(),
538        }))
539        .await;
540
541        let Some(mcp_command_tx) = self.mcp_command_tx.clone() else {
542            emit_tool_failure(&tx, ToolCallError::from_request(&tool_call, "MCP task is not available")).await;
543            return;
544        };
545
546        let trace_context = self.observers.iter().find_map(|observer| observer.tool_trace_context(&tool_id));
547        let tool_error = ToolCallError::from_request(&tool_call, "Failed to send tool request to MCP task");
548        let command =
549            McpCommand::ExecuteTool { request: tool_call, trace_context, timeout: self.tool_timeout, tx: tx.clone() };
550
551        if mcp_command_tx.send(command).await.is_err() {
552            tracing::warn!("Failed to send tool request to MCP task");
553            emit_tool_failure(&tx, tool_error).await;
554        }
555    }
556
557    async fn handle_llm_usage(&mut self, sample: TokenUsage, state: &mut IterationState) {
558        state.call_usage = Some(sample);
559        self.token_tracker.record_usage(sample);
560        let ratio_pct = self.token_tracker.usage_ratio().map(|r| r * 100.0);
561        let remaining = self.token_tracker.tokens_remaining();
562        tracing::debug!(?sample, ?ratio_pct, ?remaining, "Token usage");
563
564        self.emit(self.context_usage_message()).await;
565    }
566
567    fn context_usage_message(&self) -> AgentEvent {
568        let last = self.token_tracker.last_usage();
569        AgentEvent::Context(ContextEvent::UsageUpdated {
570            usage: ContextUsage {
571                usage_ratio: self.token_tracker.usage_ratio(),
572                context_limit: self.token_tracker.context_limit(),
573                input_tokens: last.input_tokens,
574                output_tokens: last.output_tokens,
575                cache_read_tokens: last.cache_read_tokens,
576                cache_creation_tokens: last.cache_creation_tokens,
577                reasoning_tokens: last.reasoning_tokens,
578                total_input_tokens: self.token_tracker.total_input_tokens(),
579                total_output_tokens: self.token_tracker.total_output_tokens(),
580                total_cache_read_tokens: self.token_tracker.total_cache_read_tokens(),
581                total_cache_creation_tokens: self.token_tracker.total_cache_creation_tokens(),
582                total_reasoning_tokens: self.token_tracker.total_reasoning_tokens(),
583            },
584        })
585    }
586
587    fn compaction_needed(&self) -> bool {
588        self.compaction_config.as_ref().is_some_and(|config| {
589            self.token_tracker.needs_compaction(self.context.estimated_token_count(), config.threshold)
590        })
591    }
592
593    async fn begin_compaction(&mut self) {
594        tracing::info!("Starting context compaction - {} messages", self.context.message_count());
595        self.emit(AgentEvent::Context(ContextEvent::CompactionStarted { message_count: self.context.message_count() }))
596            .await;
597        self.emit(self.llm_call_started(LlmCallPurpose::Compaction, 0)).await;
598
599        let compactor = Compactor::new(self.llm.clone());
600        let context = self.context.clone();
601        let stream: EventStream =
602            Box::pin(futures::stream::once(async move { StreamEvent::Compaction(compactor.compact(context).await) }));
603        self.streams.insert(StreamKey::Compaction, stream);
604    }
605
606    async fn on_compaction_complete(&mut self, result: Result<CompactionResult, CompactionError>) {
607        let outcome = match &result {
608            Ok(result) => LlmCallOutcome::Completed { stop_reason: None, usage: result.usage },
609            Err(e) => LlmCallOutcome::Failed { error: e.to_string(), will_retry: false },
610        };
611        self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Compaction, outcome })).await;
612
613        match result {
614            Ok(result) => {
615                tracing::info!("Context compacted: {} messages removed", result.messages_removed);
616                self.context = self.context.with_compacted_summary(&result.summary);
617                self.token_tracker.reset_current_usage();
618                self.emit(AgentEvent::Context(ContextEvent::CompactionResult {
619                    summary: result.summary,
620                    messages_removed: result.messages_removed,
621                }))
622                .await;
623                self.emit(AgentEvent::Context(ContextEvent::CompactionEnded { outcome: CompactionOutcome::Completed }))
624                    .await;
625            }
626            Err(e) => {
627                tracing::warn!("Context compaction failed: {e}");
628                self.emit(AgentEvent::Context(ContextEvent::CompactionEnded {
629                    outcome: CompactionOutcome::Failed { error: e.to_string() },
630                }))
631                .await;
632            }
633        }
634
635        self.start_chat_turn().await;
636    }
637
638    async fn on_tool_execution_event(&mut self, event: ToolExecutionEvent, state: &mut IterationState) {
639        match event {
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 refresh_prompt_cache_key(&mut self) {
686        let key = derive_prompt_cache_key(self.llm.as_ref(), &self.context);
687        self.context.set_prompt_cache_key(Some(key));
688    }
689
690    fn update_context(
691        &mut self,
692        message_content: &str,
693        reasoning: AssistantReasoning,
694        completed_tools: Vec<Result<ToolCallResult, ToolCallError>>,
695    ) {
696        self.context.push_assistant_turn(message_content, reasoning, completed_tools);
697    }
698
699    fn commit_pending_turn_messages(&mut self) {
700        let content: Vec<_> = self.pending_turn_messages.drain(..).flatten().collect();
701        if !content.is_empty() {
702            self.context.add_message(ChatMessage::User { content, timestamp: IsoString::now() });
703        }
704    }
705
706    async fn emit_tool_definitions(&mut self) {
707        let tools = self.context.tools().clone();
708        if !tools.is_empty() {
709            self.emit(AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools })).await;
710        }
711    }
712
713    async fn emit(&mut self, message: AgentEvent) {
714        for observer in &mut self.observers {
715            observer.on_event(&message);
716        }
717
718        if let Err(e) = self.message_tx.send(message).await {
719            tracing::warn!("Failed to send agent message: {e:?}");
720        }
721    }
722
723    async fn finish_turn(&mut self, outcome: TurnOutcome) {
724        if std::mem::take(&mut self.turn_active) {
725            self.emit(AgentEvent::turn_ended(outcome)).await;
726        }
727    }
728
729    async fn begin_chat_call(&mut self, attempt: u32) {
730        self.llm_call_active = true;
731        self.emit(self.llm_call_started(LlmCallPurpose::Chat, attempt)).await;
732    }
733
734    async fn finish_chat_call(&mut self, outcome: LlmCallOutcome) {
735        if std::mem::take(&mut self.llm_call_active) {
736            self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Chat, outcome })).await;
737        }
738    }
739
740    fn llm_call_started(&self, purpose: LlmCallPurpose, attempt: u32) -> AgentEvent {
741        let model = self.llm.model();
742        AgentEvent::Turn(TurnEvent::LlmCallStarted {
743            purpose,
744            provider: model.as_ref().map(|m| m.provider().to_string()),
745            model: model.as_ref().map(|m| m.model_id().into_owned()),
746            pricing: model.and_then(|m| m.pricing()),
747            display_name: self.llm.display_name(),
748            attempt,
749            max_attempts: self.retry_config.max_attempts,
750        })
751    }
752}
753
754async fn emit_tool_failure(tx: &mpsc::Sender<ToolExecutionEvent>, error: ToolCallError) {
755    let tool_id = error.id.clone();
756    let _ = tx.send(ToolExecutionEvent::Complete { tool_id, result: Err(error), result_meta: None }).await;
757}
758
759pub(crate) struct AutoContinue {
760    max: u32,
761    count: u32,
762}
763
764impl AutoContinue {
765    pub(crate) fn new(max: u32) -> Self {
766        Self { max, count: 0 }
767    }
768
769    fn reset(&mut self) {
770        self.count = 0;
771    }
772
773    fn should_continue(&self, stop_reason: Option<&StopReason>) -> bool {
774        matches!(stop_reason, Some(StopReason::Length)) && self.count < self.max
775    }
776
777    fn advance(&mut self) {
778        self.count += 1;
779    }
780
781    fn count(&self) -> u32 {
782        self.count
783    }
784
785    fn max(&self) -> u32 {
786        self.max
787    }
788}
789
790#[derive(Debug)]
791struct IterationState {
792    current_message_id: Option<String>,
793    message_content: String,
794    reasoning_summary_text: String,
795    encrypted_reasoning: Option<EncryptedReasoningContent>,
796    pending_tool_ids: HashSet<String>,
797    completed_tool_calls: Vec<Result<ToolCallResult, ToolCallError>>,
798    llm_done: bool,
799    stop_reason: Option<StopReason>,
800    retry_attempt: u32,
801    call_usage: Option<TokenUsage>,
802}
803
804impl IterationState {
805    fn new() -> Self {
806        Self {
807            current_message_id: None,
808            message_content: String::new(),
809            reasoning_summary_text: String::new(),
810            encrypted_reasoning: None,
811            pending_tool_ids: HashSet::new(),
812            completed_tool_calls: Vec::new(),
813            llm_done: false,
814            stop_reason: None,
815            retry_attempt: 0,
816            call_usage: None,
817        }
818    }
819
820    fn on_llm_start(&mut self, message_id: String) {
821        self.current_message_id = Some(message_id);
822        self.message_content.clear();
823        self.reasoning_summary_text.clear();
824        self.encrypted_reasoning = None;
825        self.stop_reason = None;
826        self.call_usage = None;
827    }
828
829    fn is_complete(&self) -> bool {
830        self.llm_done && self.pending_tool_ids.is_empty()
831    }
832}
aether_core/core/agent.rs

aether_core/core/
agent.rs
