aether_core/core/
agent.rs

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

aether_core/core/
agent.rs
