aether_core/core/
agent.rs

1use crate::context::{
2    CompactionConfig, CompactionError, CompactionResult, Compactor, SessionUsageTracker, TokenTracker,
3};
4use crate::core::PromptCache;
5use crate::core::prompt_cache_key::derive_prompt_cache_key;
6use crate::core::queued_input::QueuedInput;
7pub use crate::core::retry_config::RetryConfig;
8use crate::core::tool_execution::{ToolAbortPolicy, ToolExecutionUpdate, ToolExecutions};
9use crate::events::{
10    AgentCommand, AgentEvent, AgentObserver, Command, CompactionId, CompactionOutcome, ContextEvent, LlmCallOutcome,
11    ModelEvent, StreamState, TaskOutcome, TaskOutcomeState, ToolEvent, TraceContext, TurnEvent, TurnOutcome,
12    UserCommand,
13};
14use crate::mcp::McpCatalog;
15use crate::mcp::tool_bridge::call_tool;
16use futures::{FutureExt, Stream};
17use llm::{
18    AssistantReasoning, ChatMessage, Context, EncryptedReasoningContent, LlmCallPurpose, LlmError, LlmModel,
19    LlmResponse, MessageId, ModelIdentity, StopReason, StreamingModelProvider, TokenUsage, ToolCallError,
20    ToolCallRequest, ToolCallResult,
21};
22use mcp_utils::client::{ToolCallEvent, ToolCallOptions};
23use std::collections::VecDeque;
24use std::panic::{AssertUnwindSafe, catch_unwind};
25use std::pin::Pin;
26use std::sync::Arc;
27use std::time::Duration;
28use tokio::sync::mpsc;
29use tokio::time::sleep;
30use tokio_stream::StreamExt;
31use tokio_stream::StreamMap;
32use tokio_stream::wrappers::ReceiverStream;
33use utils::panic::panic_message;
34
35/// Internal event type for merging LLM and tool result streams
36#[derive(Debug)]
37#[allow(clippy::large_enum_variant)]
38enum StreamEvent {
39    LlmRequestStarted { attempt: u32 },
40    Llm(Result<LlmResponse, LlmError>),
41    ToolExecution(ToolCallEvent),
42    Command(Command),
43    InputClosed,
44    McpChanged,
45    Compaction(Result<CompactionResult, CompactionError>),
46}
47
48type EventStream = Pin<Box<dyn Stream<Item = StreamEvent> + Send>>;
49
50/// Keys for the merged stream map. Tool-call IDs come from providers, so the
51/// typed key keeps them from colliding with reserved streams.
52#[derive(Debug, Clone, PartialEq, Eq, Hash)]
53enum StreamKey {
54    Input,
55    Mcp,
56    Llm,
57    Compaction,
58    Tool(String),
59}
60
61pub(crate) struct AgentConfig {
62    pub llm: Arc<dyn StreamingModelProvider>,
63    pub context: Context,
64    pub mcp: McpCatalog,
65    pub tool_timeout: Duration,
66    pub compaction_config: Option<CompactionConfig>,
67    pub auto_continue: AutoContinue,
68    pub retry_config: RetryConfig,
69    pub context_window: Option<u32>,
70    pub prompt_cache: PromptCache,
71    pub observers: Vec<Box<dyn AgentObserver>>,
72    pub session_usage: SessionUsageTracker,
73}
74
75pub struct Agent {
76    llm: Arc<dyn StreamingModelProvider>,
77    context: Context,
78    mcp: McpCatalog,
79    message_tx: mpsc::Sender<AgentEvent>,
80    observers: Observers,
81    streams: StreamMap<StreamKey, EventStream>,
82    tool_timeout: Duration,
83    token_tracker: TokenTracker,
84    compaction_config: Option<CompactionConfig>,
85    auto_continue: AutoContinue,
86    retry_config: RetryConfig,
87    tool_executions: ToolExecutions,
88    pending_inputs: VecDeque<QueuedInput>,
89    queued_inputs: VecDeque<QueuedInput>,
90    context_window: Option<u32>,
91    prompt_cache: PromptCache,
92    turn_active: bool,
93    llm_call_active: bool,
94    active_compaction: Option<CompactionId>,
95    active_model: Option<LlmModel>,
96    session_usage: SessionUsageTracker,
97}
98
99impl Agent {
100    pub(crate) fn new(
101        config: AgentConfig,
102        command_rx: mpsc::Receiver<Command>,
103        message_tx: mpsc::Sender<AgentEvent>,
104    ) -> Self {
105        let mut streams: StreamMap<StreamKey, EventStream> = StreamMap::new();
106        let input_stream = ReceiverStream::new(command_rx)
107            .map(StreamEvent::Command)
108            .chain(futures::stream::once(async { StreamEvent::InputClosed }));
109        streams.insert(StreamKey::Input, Box::pin(input_stream));
110
111        streams.insert(StreamKey::Mcp, Box::pin(config.mcp.changes().map(|()| StreamEvent::McpChanged)));
112
113        let context_limit = config.context_window.or_else(|| config.llm.context_window());
114
115        Self {
116            llm: config.llm,
117            context: config.context,
118            mcp: config.mcp,
119            message_tx,
120            observers: Observers(config.observers),
121            streams,
122            tool_timeout: config.tool_timeout,
123            token_tracker: TokenTracker::new(context_limit),
124            compaction_config: config.compaction_config,
125            auto_continue: config.auto_continue,
126            retry_config: config.retry_config,
127            tool_executions: ToolExecutions::default(),
128            pending_inputs: VecDeque::new(),
129            queued_inputs: VecDeque::new(),
130            context_window: config.context_window,
131            prompt_cache: config.prompt_cache,
132            turn_active: false,
133            llm_call_active: false,
134            active_compaction: None,
135            active_model: None,
136            session_usage: config.session_usage,
137        }
138    }
139
140    pub fn current_model_display_name(&self) -> String {
141        self.llm.display_name()
142    }
143
144    /// Get a reference to the token tracker
145    pub fn token_tracker(&self) -> &TokenTracker {
146        &self.token_tracker
147    }
148
149    pub async fn run(mut self) {
150        let Err(panic) = AssertUnwindSafe(self.event_loop()).catch_unwind().await else {
151            return;
152        };
153
154        let error = format!("Agent panicked: {}", panic_message(&*panic));
155        tracing::error!("{error}");
156
157        self.end_in_flight_llm_calls(
158            LlmCallOutcome::failed(error.clone(), false),
159            CompactionOutcome::Failed { error: error.clone() },
160        )
161        .await;
162
163        if std::mem::take(&mut self.turn_active) {
164            self.emit(AgentEvent::turn_ended(TurnOutcome::failed(error))).await;
165        }
166    }
167
168    async fn event_loop(&mut self) {
169        let mut state = IterationState::default();
170        let mut input_closed = false;
171        self.sync_mcp().await;
172
173        while let Some((stream_key, event)) = self.streams.next().await {
174            match event {
175                StreamEvent::Command(Command::UserCommand(UserCommand::Cancel)) => {
176                    self.on_user_cancel(&mut state).await;
177                }
178
179                StreamEvent::Command(Command::UserCommand(UserCommand::ClearContext)) => {
180                    self.on_user_clear_context(&mut state).await;
181                }
182
183                StreamEvent::Command(Command::UserCommand(UserCommand::Text { message_id, content })) => {
184                    self.receive_input(QueuedInput::User { message_id, content }, &mut state).await;
185                }
186
187                StreamEvent::Command(Command::AgentCommand(AgentCommand::SwitchModel(new_provider))) => {
188                    self.on_switch_model(new_provider).await;
189                }
190
191                StreamEvent::McpChanged => {
192                    self.sync_mcp().await;
193                }
194
195                StreamEvent::Command(Command::AgentCommand(AgentCommand::SetReasoningEffort(effort))) => {
196                    self.context.set_reasoning_effort(effort.unwrap_or_default());
197                }
198
199                StreamEvent::Command(Command::AgentCommand(AgentCommand::ReplaceConversation(messages))) => {
200                    self.on_replace_conversation(messages, &mut state).await;
201                }
202
203                StreamEvent::InputClosed => {
204                    input_closed = true;
205                }
206
207                StreamEvent::LlmRequestStarted { attempt } => {
208                    self.begin_chat_call(attempt).await;
209                }
210
211                StreamEvent::Llm(llm_event) => {
212                    self.on_llm_event(llm_event, &mut state).await;
213                }
214
215                StreamEvent::ToolExecution(tool_event) => {
216                    let StreamKey::Tool(tool_id) = stream_key else {
217                        unreachable!("tool events must come from a tool stream")
218                    };
219                    self.on_tool_execution_event(tool_id, tool_event, &mut state).await;
220                }
221
222                StreamEvent::Compaction(result) => {
223                    self.on_compaction_complete(result).await;
224                }
225            }
226
227            if state.is_complete(self.tool_executions.has_foreground())
228                && let Some(id) = state.current_message_id.take()
229            {
230                let iteration = std::mem::take(&mut state);
231                self.on_iteration_complete(id, iteration).await;
232            }
233
234            if input_closed && !self.turn_active && !self.is_busy() {
235                if self.tool_executions.is_empty() {
236                    break;
237                }
238                self.abort_in_flight_work(ToolAbortPolicy::CancelAll).await;
239            }
240        }
241
242        tracing::debug!("Agent task shutting down - input channel closed");
243    }
244
245    async fn on_iteration_complete(&mut self, id: MessageId, iteration: IterationState) {
246        let IterationState {
247            message_content,
248            reasoning_summary_text,
249            encrypted_reasoning,
250            completed_tool_calls,
251            stop_reason,
252            ..
253        } = iteration;
254        let has_tool_calls = !completed_tool_calls.is_empty();
255        let has_content = !message_content.is_empty() || !reasoning_summary_text.is_empty() || has_tool_calls;
256        let should_auto_continue = self.auto_continue.should_continue(stop_reason.as_ref());
257
258        if has_content {
259            let reasoning = AssistantReasoning::from_parts(reasoning_summary_text.clone(), encrypted_reasoning);
260            self.context.push_assistant_turn(id.clone(), &message_content, reasoning, completed_tool_calls);
261
262            self.emit(AgentEvent::text(&id, &message_content, StreamState::Complete)).await;
263
264            if !reasoning_summary_text.is_empty() {
265                self.emit(AgentEvent::thought(&id, &reasoning_summary_text, StreamState::Complete)).await;
266            }
267        }
268
269        let has_queued_input = !self.queued_inputs.is_empty();
270        if has_queued_input || has_tool_calls {
271            self.auto_continue.reset();
272            self.start_next_turn().await;
273        } else if should_auto_continue {
274            self.auto_continue.advance();
275            tracing::info!(
276                "LLM stopped with {:?}, auto-continuing (attempt {}/{})",
277                stop_reason,
278                self.auto_continue.count,
279                self.auto_continue.max
280            );
281
282            self.inject_continuation_prompt(stop_reason.as_ref()).await;
283            self.start_next_turn().await;
284        } else {
285            tracing::debug!("LLM completed turn with stop reason: {:?}", stop_reason);
286            self.auto_continue.reset();
287            self.finish_turn(TurnOutcome::Completed).await;
288        }
289    }
290
291    async fn start_next_turn(&mut self) {
292        debug_assert!(self.pending_inputs.is_empty());
293        self.pending_inputs.append(&mut self.queued_inputs);
294        if self.compaction_needed() {
295            self.begin_compaction().await;
296        } else {
297            self.start_chat_turn().await;
298        }
299    }
300
301    async fn start_chat_turn(&mut self) {
302        self.commit_pending_inputs().await;
303        self.start_llm_stream(None, 0).await;
304    }
305
306    async fn on_user_cancel(&mut self, state: &mut IterationState) {
307        self.abort_in_flight_work(ToolAbortPolicy::PreserveBackgroundAcknowledgements).await;
308        *state = IterationState::default();
309        self.finish_turn(TurnOutcome::Cancelled).await;
310    }
311
312    async fn discard_in_flight_work(&mut self, state: &mut IterationState) {
313        self.abort_in_flight_work(ToolAbortPolicy::CancelAll).await;
314        let mut inputs = std::mem::take(&mut self.pending_inputs);
315        inputs.append(&mut self.queued_inputs);
316        self.discard_inputs(inputs, false).await;
317        self.auto_continue.reset();
318        *state = IterationState::default();
319    }
320
321    async fn on_user_clear_context(&mut self, state: &mut IterationState) {
322        self.discard_in_flight_work(state).await;
323        self.context.clear_conversation();
324        self.token_tracker.reset_current_usage();
325        self.emit(AgentEvent::Context(ContextEvent::Cleared)).await;
326        self.finish_turn(TurnOutcome::Cancelled).await;
327    }
328
329    async fn on_replace_conversation(&mut self, messages: Vec<ChatMessage>, state: &mut IterationState) {
330        self.discard_in_flight_work(state).await;
331        self.context.replace_conversation(messages);
332        self.emit(self.context_usage_message()).await;
333        self.finish_turn(TurnOutcome::Cancelled).await;
334    }
335
336    async fn begin_turn(&mut self, input: QueuedInput, state: &mut IterationState) {
337        *state = IterationState::default();
338        self.auto_continue.reset();
339        self.turn_active = true;
340        let content = input.content_blocks();
341        self.emit(AgentEvent::Turn(TurnEvent::Started { content })).await;
342        self.queued_inputs.push_back(input);
343        self.start_next_turn().await;
344    }
345
346    async fn receive_input(&mut self, input: QueuedInput, state: &mut IterationState) {
347        if self.turn_active {
348            self.queued_inputs.push_back(input);
349        } else {
350            self.begin_turn(input, state).await;
351        }
352    }
353
354    async fn sync_mcp(&mut self) {
355        let (instructions, tools) = (self.mcp.instructions(), self.mcp.tool_definitions());
356
357        if self.prompt_cache.set_mcp_instructions(instructions) {
358            match self.prompt_cache.render().await {
359                Ok(content) => self.context.set_system_content(content),
360                Err(e) => tracing::warn!("Failed to rebuild system prompt after instructions update: {e}"),
361            }
362        }
363
364        if tools != *self.context.tools() {
365            self.context.set_tools(tools.clone());
366            self.emit(AgentEvent::Tool(ToolEvent::DefinitionsUpdated { tools })).await;
367        }
368    }
369
370    async fn on_switch_model(&mut self, new_provider: Box<dyn StreamingModelProvider>) {
371        let previous = self.llm.display_name();
372        let new_context_limit = self.context_window.or_else(|| new_provider.context_window());
373        self.llm = Arc::from(new_provider);
374        self.token_tracker.reset_current_usage();
375        self.token_tracker.set_context_limit(new_context_limit);
376        let new = self.llm.display_name();
377        self.emit(AgentEvent::Model(ModelEvent::Switched { previous, new })).await;
378
379        self.emit(self.context_usage_message()).await;
380    }
381
382    async fn start_llm_stream(&mut self, delay: Option<Duration>, attempt: u32) {
383        self.refresh_prompt_cache_key();
384        self.streams.remove(&StreamKey::Llm);
385        let stream: EventStream = match delay {
386            None => {
387                self.begin_chat_call(attempt).await;
388                Box::pin(self.llm.stream_response(&self.context).map(StreamEvent::Llm))
389            }
390            Some(delay) => {
391                self.emit(AgentEvent::Turn(TurnEvent::RetryScheduled {
392                    purpose: LlmCallPurpose::Chat,
393                    attempt,
394                    max_attempts: self.retry_config.max_attempts,
395                    delay_ms: u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
396                }))
397                .await;
398                let llm = Arc::clone(&self.llm);
399                let context = self.context.clone();
400                Box::pin(async_stream::stream! {
401                    sleep(delay).await;
402                    yield StreamEvent::LlmRequestStarted { attempt };
403                    let mut inner = llm.stream_response(&context);
404                    while let Some(item) = inner.next().await {
405                        yield StreamEvent::Llm(item);
406                    }
407                })
408            }
409        };
410        self.streams.insert(StreamKey::Llm, stream);
411    }
412
413    async fn on_llm_error(&mut self, error: LlmError, state: &mut IterationState) {
414        let will_retry = error.is_retryable() && state.retry_attempt < self.retry_config.max_attempts;
415        let outcome = LlmCallOutcome::from_llm_error(&error, will_retry);
416        let error_message = error.to_string();
417        self.finish_chat_call(outcome).await;
418
419        if !will_retry {
420            self.finish_turn(TurnOutcome::failed(error_message)).await;
421            return;
422        }
423
424        state.retry_attempt += 1;
425        let delay = self.retry_config.compute_delay(state.retry_attempt);
426
427        tracing::warn!(
428            attempt = state.retry_attempt,
429            max_attempts = self.retry_config.max_attempts,
430            delay_ms = u64::try_from(delay.as_millis()).unwrap_or(u64::MAX),
431            error = %error,
432            "Retrying LLM request after transient failure"
433        );
434
435        self.tool_executions.retire_foreground();
436        self.start_llm_stream(Some(delay), state.retry_attempt).await;
437    }
438
439    fn is_busy(&self) -> bool {
440        self.streams.contains_key(&StreamKey::Llm)
441            || self.active_compaction.is_some()
442            || self.tool_executions.has_foreground()
443    }
444
445    async fn abort_in_flight_work(&mut self, tool_policy: ToolAbortPolicy) {
446        self.end_in_flight_llm_calls(LlmCallOutcome::Cancelled, CompactionOutcome::Cancelled).await;
447        self.streams.remove(&StreamKey::Llm);
448        for tool_id in self.tool_executions.abort(&tool_policy) {
449            self.streams.remove(&StreamKey::Tool(tool_id));
450        }
451    }
452
453    async fn end_in_flight_llm_calls(&mut self, call_outcome: LlmCallOutcome, compaction_outcome: CompactionOutcome) {
454        self.finish_chat_call(call_outcome.clone()).await;
455        if let Some(compaction_id) = self.active_compaction.take() {
456            self.streams.remove(&StreamKey::Compaction);
457            self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded {
458                purpose: LlmCallPurpose::Compaction,
459                outcome: call_outcome,
460            }))
461            .await;
462            self.emit(AgentEvent::Context(ContextEvent::CompactionEnded {
463                compaction_id,
464                outcome: compaction_outcome,
465            }))
466            .await;
467        }
468    }
469
470    /// Inject a continuation prompt when the LLM stops due to a resumable reason.
471    async fn inject_continuation_prompt(&mut self, stop_reason: Option<&StopReason>) {
472        let reason = stop_reason.map_or_else(|| "Unknown".to_string(), |reason| format!("{reason:?}"));
473        let message_id = MessageId::new();
474        let content = vec![llm::ContentBlock::text(format!(
475            "<system-notification>The LLM API stopped with reason '{reason}'. Continue from where you left off and finish your task.</system-notification>"
476        ))];
477        self.context.add_message(ChatMessage::user_with_id(message_id.clone(), content.clone()));
478        self.emit(AgentEvent::Turn(TurnEvent::AutoContinue {
479            attempt: self.auto_continue.count,
480            max_attempts: self.auto_continue.max,
481            message_id,
482            content,
483        }))
484        .await;
485    }
486
487    async fn on_llm_event(&mut self, result: Result<LlmResponse, LlmError>, state: &mut IterationState) {
488        use LlmResponse::{
489            Done, EncryptedReasoning, Error, Reasoning, Start, Text, ToolRequestArg, ToolRequestComplete,
490            ToolRequestStart, Usage,
491        };
492
493        let response = match result {
494            Ok(response) => response,
495            Err(e) => {
496                self.on_llm_error(e, state).await;
497                return;
498            }
499        };
500
501        match response {
502            Start => state.on_llm_start(MessageId::new()),
503
504            Text { chunk } => {
505                self.handle_llm_text(chunk, state).await;
506            }
507
508            Reasoning { chunk } => {
509                state.reasoning_summary_text.push_str(&chunk);
510                if let Some(id) = state.current_message_id.clone() {
511                    self.emit(AgentEvent::thought(&id, &chunk, StreamState::Partial)).await;
512                }
513            }
514
515            EncryptedReasoning { id, content } => {
516                if let Some(model) = self.active_model.clone() {
517                    state.encrypted_reasoning = Some(EncryptedReasoningContent { id, model, content });
518                }
519            }
520
521            ToolRequestStart { id, name } => {
522                self.emit(AgentEvent::Tool(ToolEvent::InputStarted { id, name })).await;
523            }
524
525            ToolRequestArg { id, chunk } => {
526                self.emit(AgentEvent::Tool(ToolEvent::InputDelta { id, chunk })).await;
527            }
528
529            ToolRequestComplete { tool_call } => {
530                self.handle_tool_completion(tool_call, state).await;
531            }
532
533            Done { stop_reason } => {
534                state.llm_done = true;
535                state.stop_reason = stop_reason;
536                self.finish_chat_call(LlmCallOutcome::Completed {
537                    stop_reason: state.stop_reason.clone(),
538                    usage: state.call_usage.take(),
539                })
540                .await;
541            }
542
543            Error { message } => {
544                self.finish_chat_call(LlmCallOutcome::failed(message.clone(), false)).await;
545                self.finish_turn(TurnOutcome::failed(message)).await;
546            }
547
548            Usage { tokens: sample } => {
549                self.handle_llm_usage(sample, state).await;
550            }
551        }
552    }
553
554    async fn handle_llm_text(&mut self, chunk: String, state: &mut IterationState) {
555        state.message_content.push_str(&chunk);
556
557        if let Some(id) = state.current_message_id.clone() {
558            self.emit(AgentEvent::text(&id, &chunk, StreamState::Partial)).await;
559        }
560    }
561
562    async fn handle_tool_completion(&mut self, tool_call: ToolCallRequest, state: &mut IterationState) {
563        let cancel = self.tool_executions.start(tool_call.clone());
564
565        let tool_id = tool_call.id.clone();
566        tracing::debug!("Tool execution started: {} ({})", tool_call.name, tool_id);
567        self.emit(AgentEvent::Tool(ToolEvent::Call { request: tool_call.clone() })).await;
568
569        let trace_context = self.observers.tool_trace_context(&tool_id);
570        let options = ToolCallOptions {
571            timeout: Some(self.tool_timeout),
572            meta: trace_context.as_ref().map(TraceContext::to_meta),
573            cancel,
574        };
575        match call_tool(&self.mcp, &tool_call, options) {
576            Ok(call) => {
577                self.streams.insert(StreamKey::Tool(tool_id), Box::pin(call.map(StreamEvent::ToolExecution)));
578            }
579            Err(rejection) => {
580                let update = self.tool_executions.reject(&tool_id, &rejection);
581                self.apply_tool_update(tool_id, update, state).await;
582            }
583        }
584    }
585
586    async fn handle_llm_usage(&mut self, sample: TokenUsage, state: &mut IterationState) {
587        state.call_usage = Some(sample);
588        self.token_tracker.record_usage(sample);
589        let ratio_pct = self.token_tracker.usage_ratio().map(|r| r * 100.0);
590        let remaining = self.token_tracker.tokens_remaining();
591        tracing::debug!(?sample, ?ratio_pct, ?remaining, "Token usage");
592
593        self.emit(self.context_usage_message()).await;
594        self.emit_session_usage(LlmCallPurpose::Chat, sample).await;
595    }
596
597    async fn emit_session_usage(&mut self, purpose: LlmCallPurpose, tokens: TokenUsage) {
598        let model = ModelIdentity::of(self.active_model.as_ref());
599        let event = self.session_usage.record(purpose, model, tokens);
600        self.emit(AgentEvent::SessionUsage(event)).await;
601    }
602
603    fn context_usage_message(&self) -> AgentEvent {
604        AgentEvent::Context(ContextEvent::UsageUpdated { usage: self.token_tracker.snapshot().clone() })
605    }
606
607    fn compaction_needed(&self) -> bool {
608        self.compaction_config.as_ref().is_some_and(|config| {
609            self.token_tracker.needs_compaction(self.context.estimated_token_count(), config.threshold)
610        })
611    }
612
613    async fn begin_compaction(&mut self) {
614        tracing::info!("Starting context compaction - {} messages", self.context.message_count());
615        let compaction_id = CompactionId::new();
616        self.active_compaction = Some(compaction_id.clone());
617        self.emit(AgentEvent::Context(ContextEvent::CompactionStarted {
618            compaction_id,
619            message_count: self.context.message_count(),
620        }))
621        .await;
622        let started = self.begin_llm_call(LlmCallPurpose::Compaction, 0);
623        self.emit(started).await;
624
625        let compactor = Compactor::new(self.llm.clone());
626        let context = self.context.clone();
627        let stream: EventStream =
628            Box::pin(futures::stream::once(async move { StreamEvent::Compaction(compactor.compact(context).await) }));
629        self.streams.insert(StreamKey::Compaction, stream);
630    }
631
632    async fn on_compaction_complete(&mut self, result: Result<CompactionResult, CompactionError>) {
633        let compaction_id = self.active_compaction.take().expect("completed compaction has an identity");
634        if let Ok(result) = &result
635            && let Some(usage) = result.usage
636        {
637            self.emit_session_usage(LlmCallPurpose::Compaction, usage).await;
638        }
639        let outcome = match &result {
640            Ok(result) => LlmCallOutcome::Completed { stop_reason: None, usage: result.usage },
641            Err(e) => LlmCallOutcome::failed(e.to_string(), false),
642        };
643        self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Compaction, outcome })).await;
644
645        match result {
646            Ok(result) => {
647                tracing::info!("Context compacted: {} messages removed", result.messages_removed);
648                let message_id = MessageId::new();
649                self.context = self.context.with_compacted_summary(message_id.clone(), &result.summary);
650                self.token_tracker.reset_current_usage();
651                self.emit(AgentEvent::Context(ContextEvent::CompactionResult {
652                    compaction_id: compaction_id.clone(),
653                    message_id,
654                    summary: result.summary,
655                    messages_removed: result.messages_removed,
656                }))
657                .await;
658                self.emit(AgentEvent::Context(ContextEvent::CompactionEnded {
659                    compaction_id,
660                    outcome: CompactionOutcome::Completed,
661                }))
662                .await;
663            }
664            Err(e) => {
665                tracing::warn!("Context compaction failed: {e}");
666                self.emit(AgentEvent::Context(ContextEvent::CompactionEnded {
667                    compaction_id,
668                    outcome: CompactionOutcome::Failed { error: e.to_string() },
669                }))
670                .await;
671            }
672        }
673
674        self.start_chat_turn().await;
675    }
676
677    async fn on_tool_execution_event(&mut self, tool_id: String, event: ToolCallEvent, state: &mut IterationState) {
678        let update = self.tool_executions.on_event(&tool_id, event);
679        self.apply_tool_update(tool_id, update, state).await;
680    }
681
682    async fn apply_tool_update(&mut self, tool_id: String, update: ToolExecutionUpdate, state: &mut IterationState) {
683        match update {
684            ToolExecutionUpdate::Event(event) => {
685                if let ToolEvent::SubAgentProgress { payload, .. } = &event
686                    && let AgentEvent::SessionUsage(child) = &payload.event
687                {
688                    let folded = self.session_usage.record_child(&payload.task_id, child.clone());
689                    self.emit(AgentEvent::SessionUsage(folded)).await;
690                }
691                self.emit(AgentEvent::Tool(event)).await;
692            }
693            ToolExecutionUpdate::Completed { result, event } => {
694                self.streams.remove(&StreamKey::Tool(tool_id));
695                state.completed_tool_calls.push(result);
696                self.emit(AgentEvent::Tool(event)).await;
697            }
698            ToolExecutionUpdate::TaskCreated { result, event } => {
699                state.completed_tool_calls.push(Ok(result));
700                self.emit(AgentEvent::Tool(event)).await;
701            }
702            ToolExecutionUpdate::TaskFinished(outcome) => {
703                self.streams.remove(&StreamKey::Tool(tool_id));
704                if matches!(outcome.state, TaskOutcomeState::Cancelled) {
705                    self.record_task_outcome(outcome).await;
706                } else {
707                    self.receive_input(QueuedInput::TaskOutcome(Box::new(outcome)), state).await;
708                }
709            }
710            ToolExecutionUpdate::Retired => {
711                self.streams.remove(&StreamKey::Tool(tool_id));
712            }
713            ToolExecutionUpdate::Ignored => {
714                tracing::debug!(%tool_id, "Ignoring unexpected tool execution event");
715            }
716        }
717    }
718
719    async fn record_task_outcome(&mut self, outcome: TaskOutcome) {
720        self.context.add_message(outcome.context_message());
721        self.emit(AgentEvent::Tool(outcome.into())).await;
722    }
723
724    fn refresh_prompt_cache_key(&mut self) {
725        let key = derive_prompt_cache_key(self.llm.as_ref(), &self.context);
726        self.context.set_prompt_cache_key(Some(key));
727    }
728
729    async fn commit_pending_inputs(&mut self) {
730        let inputs = std::mem::take(&mut self.pending_inputs);
731        self.commit_inputs(inputs).await;
732    }
733
734    async fn discard_inputs(&mut self, inputs: VecDeque<QueuedInput>, keep_task_outcomes: bool) {
735        for input in inputs {
736            match input {
737                QueuedInput::User { message_id, .. } => {
738                    self.emit(AgentEvent::Turn(TurnEvent::UserMessageDiscarded { message_id })).await;
739                }
740                QueuedInput::TaskOutcome(outcome) if keep_task_outcomes => self.record_task_outcome(*outcome).await,
741                QueuedInput::TaskOutcome(_) => {}
742            }
743        }
744    }
745
746    async fn commit_inputs(&mut self, inputs: VecDeque<QueuedInput>) {
747        for input in inputs {
748            match input {
749                QueuedInput::User { message_id, content } => {
750                    self.context.add_message(ChatMessage::user_with_id(message_id.clone(), content));
751                    self.emit(AgentEvent::Turn(TurnEvent::UserMessageInserted { message_id })).await;
752                }
753                QueuedInput::TaskOutcome(outcome) => self.record_task_outcome(*outcome).await,
754            }
755        }
756    }
757
758    async fn emit(&mut self, message: AgentEvent) {
759        self.observers.on_event(&message);
760
761        if let Err(e) = self.message_tx.send(message).await {
762            tracing::warn!("Failed to send agent message: {e:?}");
763        }
764    }
765
766    async fn finish_turn(&mut self, outcome: TurnOutcome) {
767        self.commit_pending_inputs().await;
768        let queued = std::mem::take(&mut self.queued_inputs);
769        self.discard_inputs(queued, true).await;
770        if std::mem::take(&mut self.turn_active) {
771            self.emit(AgentEvent::turn_ended(outcome)).await;
772        }
773    }
774
775    async fn begin_chat_call(&mut self, attempt: u32) {
776        self.llm_call_active = true;
777        let started = self.begin_llm_call(LlmCallPurpose::Chat, attempt);
778        if let Some(system_prompt) = self.context.system_content() {
779            self.observers.on_system_prompt(system_prompt);
780        }
781        self.emit(started).await;
782    }
783
784    async fn finish_chat_call(&mut self, outcome: LlmCallOutcome) {
785        if std::mem::take(&mut self.llm_call_active) {
786            self.emit(AgentEvent::Turn(TurnEvent::LlmCallEnded { purpose: LlmCallPurpose::Chat, outcome })).await;
787        }
788    }
789
790    fn begin_llm_call(&mut self, purpose: LlmCallPurpose, attempt: u32) -> AgentEvent {
791        self.active_model = self.llm.model();
792        AgentEvent::Turn(TurnEvent::LlmCallStarted {
793            purpose,
794            model: ModelIdentity::of(self.active_model.as_ref()),
795            display_name: self.llm.display_name(),
796            attempt,
797            max_attempts: self.retry_config.max_attempts,
798        })
799    }
800}
801
802pub(crate) struct AutoContinue {
803    max: u32,
804    count: u32,
805}
806
807impl AutoContinue {
808    pub(crate) fn new(max: u32) -> Self {
809        Self { max, count: 0 }
810    }
811
812    fn reset(&mut self) {
813        self.count = 0;
814    }
815
816    fn should_continue(&self, stop_reason: Option<&StopReason>) -> bool {
817        matches!(stop_reason, Some(StopReason::Length)) && self.count < self.max
818    }
819
820    fn advance(&mut self) {
821        self.count += 1;
822    }
823}
824
825#[derive(Debug, Default)]
826struct IterationState {
827    current_message_id: Option<MessageId>,
828    message_content: String,
829    reasoning_summary_text: String,
830    encrypted_reasoning: Option<EncryptedReasoningContent>,
831    completed_tool_calls: Vec<Result<ToolCallResult, ToolCallError>>,
832    llm_done: bool,
833    stop_reason: Option<StopReason>,
834    retry_attempt: u32,
835    call_usage: Option<TokenUsage>,
836}
837
838impl IterationState {
839    fn on_llm_start(&mut self, message_id: MessageId) {
840        self.current_message_id = Some(message_id);
841        self.message_content.clear();
842        self.reasoning_summary_text.clear();
843        self.encrypted_reasoning = None;
844        self.stop_reason = None;
845        self.call_usage = None;
846    }
847
848    fn is_complete(&self, has_foreground_tools: bool) -> bool {
849        self.llm_done && !has_foreground_tools
850    }
851}
852
853struct Observers(Vec<Box<dyn AgentObserver>>);
854
855impl Observers {
856    fn on_event(&mut self, event: &AgentEvent) {
857        self.notify(|observer| observer.on_event(event));
858    }
859
860    fn on_system_prompt(&mut self, prompt: &str) {
861        self.notify(|observer| observer.on_system_prompt(prompt));
862    }
863
864    fn tool_trace_context(&mut self, tool_id: &str) -> Option<TraceContext> {
865        let mut trace_context = None;
866        self.notify(|observer| trace_context = trace_context.take().or_else(|| observer.tool_trace_context(tool_id)));
867        trace_context
868    }
869
870    fn notify(&mut self, mut callback: impl FnMut(&mut dyn AgentObserver)) {
871        self.0.retain_mut(|observer| {
872            catch_unwind(AssertUnwindSafe(|| callback(observer.as_mut())))
873                .inspect_err(|panic| {
874                    tracing::error!("Removing agent observer that panicked: {}", panic_message(&**panic));
875                })
876                .is_ok()
877        });
878    }
879}
aether_core/core/agent.rs

aether_core/core/
agent.rs
