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

aether_core/core/
agent.rs
