aether_core/core/
agent.rs

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

aether_core/core/
agent.rs
