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

aether_core/core/
agent.rs
