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

aether_core/core/
agent.rs
