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

aether_core/core/
agent.rs
