aether_core/core/
agent.rs

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

aether_core/core/
agent.rs
