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

aether_core/core/
agent.rs
