oino_agent/
agent_loop.rs

1//! The agent loop.
2//!
3//! A run emits `agent_start`, then turns: each turn adds queued messages, requests one
4//! assistant response, executes its tool calls and ends with `turn_end`. The run continues
5//! while tool calls or steering messages remain, then polls follow-ups, then honours an
6//! explicit `finish_turn → Continue` with one context-only turn. There is no turn cap.
7
8use std::sync::Arc;
9use std::sync::atomic::Ordering;
10use std::time::Instant;
11
12use futures::StreamExt;
13use futures::future::{BoxFuture, join_all};
14use oino_ai::transcript::{SystemState, tool_state_changes};
15use oino_ai::{
16    AgentMessage, AssistantContent, AssistantMessage, AssistantMessageEvent, Context, JsonObject,
17    StopReason, SystemMessage, ThinkingLevel, ToolCall, ToolDecl, ToolRef, ToolResultMessage,
18    UserTextField,
19};
20use tokio_util::sync::CancellationToken;
21
22use crate::config::{
23    AfterToolCall, AfterToolCallContext, AgentContext, AgentLoopConfig, AgentTurnContext,
24    BeforeToolCall, BeforeToolCallContext, HookError, PrepareRequestContext, StreamRequest,
25    ToolExecution, TurnDecision,
26};
27use crate::event::{AgentEvent, EventSink};
28use crate::tool::{AgentTool, AgentToolResult, ExecutionMode, ToolInvocation, UpdateSink};
29use crate::validate::validate_tool_arguments;
30
31/// The loop cannot start or a hook failed.
32#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
33#[non_exhaustive]
34pub enum AgentLoopError {
35    /// `agent_loop_continue` without messages.
36    #[error("Cannot continue: no messages in context")]
37    NoMessages,
38    /// `agent_loop_continue` after an assistant message.
39    #[error("Cannot continue from message role: assistant")]
40    ContinueFromAssistant,
41    /// A loop hook failed; the run ends like a thrown error.
42    #[error("{0}")]
43    Hook(#[from] HookError),
44}
45
46/// Current wall-clock time in epoch milliseconds.
47#[must_use]
48pub fn now_ms() -> i64 {
49    std::time::SystemTime::now()
50        .duration_since(std::time::UNIX_EPOCH)
51        .map_or(0, |d| i64::try_from(d.as_millis()).unwrap_or(i64::MAX))
52}
53
54/// Starts a run with `prompts`. Returns the messages the run added.
55///
56/// # Errors
57///
58/// [`AgentLoopError::Hook`] when a loop hook fails (events emitted so far stand).
59pub async fn agent_loop(
60    prompts: Vec<AgentMessage>,
61    context: AgentContext,
62    config: &AgentLoopConfig,
63    sink: &dyn EventSink,
64    cancel: CancellationToken,
65) -> Result<Vec<AgentMessage>, AgentLoopError> {
66    let initial = declare_tool_changes(&context, prompts);
67    let mut new_messages = initial.clone();
68    let mut current = context;
69    current.messages.extend(initial.iter().cloned());
70    sink.emit(AgentEvent::AgentStart).await;
71    sink.emit(AgentEvent::TurnStart).await;
72    for message in initial {
73        sink.emit(AgentEvent::MessageStart {
74            message: message.clone(),
75        })
76        .await;
77        sink.emit(AgentEvent::MessageEnd { message }).await;
78    }
79    run_loop(current, &mut new_messages, config, sink, &cancel).await?;
80    Ok(new_messages)
81}
82
83/// Continues from the current transcript.
84///
85/// # Errors
86///
87/// [`AgentLoopError::NoMessages`] / [`AgentLoopError::ContinueFromAssistant`] before anything
88/// is emitted; [`AgentLoopError::Hook`] when a hook fails.
89pub async fn agent_loop_continue(
90    context: AgentContext,
91    config: &AgentLoopConfig,
92    sink: &dyn EventSink,
93    cancel: CancellationToken,
94) -> Result<Vec<AgentMessage>, AgentLoopError> {
95    match context.messages.last() {
96        None => return Err(AgentLoopError::NoMessages),
97        Some(AgentMessage::Assistant(_)) => return Err(AgentLoopError::ContinueFromAssistant),
98        Some(_) => {}
99    }
100    let mut new_messages = Vec::new();
101    sink.emit(AgentEvent::AgentStart).await;
102    sink.emit(AgentEvent::TurnStart).await;
103    run_loop(context, &mut new_messages, config, sink, &cancel).await?;
104    Ok(new_messages)
105}
106
107struct CompletedTurn {
108    message: AssistantMessage,
109    tool_results: Vec<ToolResultMessage>,
110}
111
112#[allow(
113    clippy::too_many_lines,
114    reason = "the loop is ported one-to-one; splitting obscures the control flow"
115)]
116async fn run_loop(
117    mut current: AgentContext,
118    new_messages: &mut Vec<AgentMessage>,
119    initial_config: &AgentLoopConfig,
120    sink: &dyn EventSink,
121    cancel: &CancellationToken,
122) -> Result<(), AgentLoopError> {
123    let mut model = initial_config.model.clone();
124    let mut thinking = initial_config.thinking_level;
125    let config = initial_config;
126    let mut last_turn: Option<CompletedTurn> = None;
127    let mut explicit_continuation = false;
128    let mut pending = poll(config.get_steering_messages.as_ref()).await;
129
130    loop {
131        let mut has_more_tool_calls = true;
132        while has_more_tool_calls || !pending.is_empty() {
133            let mut prepared = Vec::new();
134            if let Some(turn) = &last_turn {
135                if let Some(hook) = &config.prepare_next_turn {
136                    let cx = AgentTurnContext {
137                        message: &turn.message,
138                        tool_results: &turn.tool_results,
139                        context: &current,
140                        new_messages,
141                    };
142                    if let Some(update) = hook.prepare(cx, cancel).await? {
143                        if let Some(context) = update.context {
144                            current = context;
145                        }
146                        prepared = update.messages.unwrap_or_default();
147                        model = update.model.unwrap_or(model);
148                        thinking = update.thinking_level.unwrap_or(thinking);
149                    }
150                }
151                if pending.is_empty() {
152                    pending = poll(config.get_steering_messages.as_ref()).await;
153                }
154                sink.emit(AgentEvent::TurnStart).await;
155            }
156
157            prepared.append(&mut pending);
158            for message in declare_tool_changes(&current, prepared) {
159                sink.emit(AgentEvent::MessageStart {
160                    message: message.clone(),
161                })
162                .await;
163                sink.emit(AgentEvent::MessageEnd {
164                    message: message.clone(),
165                })
166                .await;
167                current.messages.push(message.clone());
168                new_messages.push(message);
169            }
170
171            if let Some(hook) = &config.prepare_request {
172                let cx = PrepareRequestContext {
173                    context: &current,
174                    model: &model,
175                    thinking_level: thinking,
176                };
177                if let Some(update) = hook.prepare(cx, cancel).await? {
178                    if let Some(context) = update.context {
179                        current = context;
180                    }
181                    model = update.model.unwrap_or(model);
182                    thinking = update.thinking_level.unwrap_or(thinking);
183                }
184            }
185
186            let message =
187                stream_assistant_response(&mut current, config, &model, thinking, sink, cancel)
188                    .await?;
189            new_messages.push(AgentMessage::Assistant(message.clone()));
190
191            if matches!(message.stop_reason, StopReason::Error | StopReason::Aborted) {
192                if let Some(hook) = &config.finish_turn {
193                    let cx = AgentTurnContext {
194                        message: &message,
195                        tool_results: &[],
196                        context: &current,
197                        new_messages,
198                    };
199                    hook.finish(cx, cancel).await?;
200                }
201                sink.emit(AgentEvent::TurnEnd {
202                    message: AgentMessage::Assistant(message),
203                    tool_results: Vec::new(),
204                })
205                .await;
206                sink.emit(AgentEvent::AgentEnd {
207                    messages: new_messages.clone(),
208                })
209                .await;
210                return Ok(());
211            }
212
213            let tool_calls = tool_calls_of(&message);
214            let mut tool_results = Vec::new();
215            has_more_tool_calls = false;
216            if !tool_calls.is_empty() {
217                let batch = if message.stop_reason == StopReason::Length {
218                    fail_truncated_tool_calls(&tool_calls, sink).await
219                } else {
220                    execute_tool_calls(&current, &message, tool_calls, config, sink, cancel).await
221                };
222                has_more_tool_calls = !batch.terminate;
223                tool_results = batch.messages;
224                for result in &tool_results {
225                    current
226                        .messages
227                        .push(AgentMessage::ToolResult(result.clone()));
228                    new_messages.push(AgentMessage::ToolResult(result.clone()));
229                }
230            }
231
232            let decision = match &config.finish_turn {
233                Some(hook) => {
234                    let cx = AgentTurnContext {
235                        message: &message,
236                        tool_results: &tool_results,
237                        context: &current,
238                        new_messages,
239                    };
240                    hook.finish(cx, cancel).await?
241                }
242                None => None,
243            };
244            sink.emit(AgentEvent::TurnEnd {
245                message: AgentMessage::Assistant(message.clone()),
246                tool_results: tool_results.clone(),
247            })
248            .await;
249            last_turn = Some(CompletedTurn {
250                message,
251                tool_results,
252            });
253            if decision == Some(TurnDecision::End) {
254                sink.emit(AgentEvent::AgentEnd {
255                    messages: new_messages.clone(),
256                })
257                .await;
258                return Ok(());
259            }
260            explicit_continuation = decision == Some(TurnDecision::Continue);
261            pending = poll(config.get_steering_messages.as_ref()).await;
262            if has_more_tool_calls || !pending.is_empty() {
263                explicit_continuation = false;
264            }
265        }
266
267        let follow_ups = poll(config.get_follow_up_messages.as_ref()).await;
268        if !follow_ups.is_empty() {
269            explicit_continuation = false;
270            pending = follow_ups;
271            continue;
272        }
273        if explicit_continuation {
274            explicit_continuation = false;
275            continue;
276        }
277        break;
278    }
279    sink.emit(AgentEvent::AgentEnd {
280        messages: new_messages.clone(),
281    })
282    .await;
283    Ok(())
284}
285
286async fn poll(source: Option<&Arc<dyn crate::config::QueueSource>>) -> Vec<AgentMessage> {
287    match source {
288        Some(source) => source.drain().await,
289        None => Vec::new(),
290    }
291}
292
293fn tool_calls_of(message: &AssistantMessage) -> Vec<ToolCall> {
294    message
295        .content
296        .iter()
297        .filter_map(|block| match block {
298            AssistantContent::ToolCall(call) => Some(call.clone()),
299            _ => None,
300        })
301        .collect()
302}
303
304/// Announces loadout changes with a system message before the next
305/// request, or rewrites the tool fields of a pending system message.
306#[must_use]
307pub fn declare_tool_changes(
308    context: &AgentContext,
309    mut pending: Vec<AgentMessage>,
310) -> Vec<AgentMessage> {
311    let system_index = pending
312        .iter()
313        .rposition(|message| matches!(message, AgentMessage::System(_)));
314    // The pending system message counts with its tool fields cleared, i.e. not at all.
315    let previous = SystemState::fold(
316        context
317            .messages
318            .iter()
319            .chain(
320                pending
321                    .iter()
322                    .enumerate()
323                    .filter(|(index, _)| Some(*index) != system_index)
324                    .map(|(_, message)| message),
325            )
326            .filter_map(|message| match message {
327                AgentMessage::System(system) => Some(system),
328                _ => None,
329            }),
330    );
331    let current: Vec<ToolDecl> = context
332        .tools
333        .iter()
334        .map(|tool| tool.meta().declaration())
335        .collect();
336    let (added, removed) = tool_state_changes(&previous.tools, &current);
337    let unchanged = added.is_empty() && removed.is_empty();
338    if let Some(index) = system_index {
339        if let AgentMessage::System(system) = &pending[index] {
340            let declares_nothing = system.tools_added.as_ref().is_none_or(Vec::is_empty)
341                && system.tools_removed.as_ref().is_none_or(Vec::is_empty);
342            if unchanged && declares_nothing {
343                return pending;
344            }
345            pending[index] = AgentMessage::System(with_tool_changes(system, added, removed));
346        }
347        return pending;
348    }
349    if unchanged {
350        return pending;
351    }
352    let update = SystemMessage {
353        content: UserTextField::Text(String::new()),
354        sections: None,
355        tools_added: None,
356        tools_removed: None,
357        timestamp: now_ms(),
358    };
359    let update = AgentMessage::System(with_tool_changes(&update, added, removed));
360    let insert_at = pending
361        .iter()
362        .position(|message| !matches!(message, AgentMessage::System(_)))
363        .unwrap_or(pending.len());
364    pending.insert(insert_at, update);
365    pending
366}
367
368fn with_tool_changes(
369    message: &SystemMessage,
370    added: Vec<ToolDecl>,
371    removed: Vec<ToolRef>,
372) -> SystemMessage {
373    SystemMessage {
374        tools_added: (!added.is_empty()).then_some(added),
375        tools_removed: (!removed.is_empty()).then_some(removed),
376        ..message.clone()
377    }
378}
379
380async fn stream_assistant_response(
381    current: &mut AgentContext,
382    config: &AgentLoopConfig,
383    model: &oino_ai::Model,
384    thinking: ThinkingLevel,
385    sink: &dyn EventSink,
386    cancel: &CancellationToken,
387) -> Result<AssistantMessage, AgentLoopError> {
388    let messages = match &config.transform_context {
389        Some(hook) => hook.transform(current.messages.clone(), cancel).await?,
390        None => current.messages.clone(),
391    };
392    let llm_messages = (config.convert_to_llm)(&messages);
393    let api_key = match &config.get_api_key {
394        Some(hook) => hook.api_key(&model.provider).await?,
395        None => None,
396    }
397    .filter(|key| !key.is_empty())
398    .or_else(|| config.api_key.clone());
399    let reasoning = (thinking != ThinkingLevel::Off).then_some(thinking);
400    let mut options = config.stream_options.clone();
401    options.reasoning = reasoning;
402    let mut events = config.stream.stream(StreamRequest {
403        model: model.clone(),
404        context: Context {
405            messages: llm_messages,
406        },
407        options,
408        api_key,
409        cancel: cancel.clone(),
410    });
411
412    let mut started = false;
413    let mut last_partial: Option<AssistantMessage> = None;
414    while let Some(event) = events.next().await {
415        match &event {
416            AssistantMessageEvent::Start { partial } => {
417                started = true;
418                last_partial = Some(partial.clone());
419                sink.emit(AgentEvent::MessageStart {
420                    message: AgentMessage::Assistant(partial.clone()),
421                })
422                .await;
423            }
424            AssistantMessageEvent::Done { message, .. }
425            | AssistantMessageEvent::Error { error: message, .. } => {
426                let mut final_message = message.clone();
427                final_message.thinking_level = Some(thinking);
428                return Ok(finish_message(current, final_message, started, sink).await);
429            }
430            other => {
431                if started {
432                    let partial = partial_of(other).cloned();
433                    if let Some(partial) = partial {
434                        last_partial = Some(partial.clone());
435                        sink.emit(AgentEvent::MessageUpdate {
436                            message: AgentMessage::Assistant(partial),
437                            assistant_message_event: event,
438                        })
439                        .await;
440                    }
441                }
442            }
443        }
444    }
445    // The stream ended without a terminal event: finish with an error message.
446    let mut final_message = last_partial.unwrap_or_else(|| {
447        AssistantMessage::empty(&model.api, &model.provider, &model.id, now_ms())
448    });
449    final_message.stop_reason = if cancel.is_cancelled() {
450        StopReason::Aborted
451    } else {
452        StopReason::Error
453    };
454    final_message
455        .error_message
456        .get_or_insert_with(|| "Stream ended without a final event".to_owned());
457    final_message.thinking_level = Some(thinking);
458    Ok(finish_message(current, final_message, started, sink).await)
459}
460
461fn partial_of(event: &AssistantMessageEvent) -> Option<&AssistantMessage> {
462    match event {
463        AssistantMessageEvent::Start { partial }
464        | AssistantMessageEvent::TextStart { partial, .. }
465        | AssistantMessageEvent::TextDelta { partial, .. }
466        | AssistantMessageEvent::TextEnd { partial, .. }
467        | AssistantMessageEvent::ThinkingStart { partial, .. }
468        | AssistantMessageEvent::ThinkingDelta { partial, .. }
469        | AssistantMessageEvent::ThinkingEnd { partial, .. }
470        | AssistantMessageEvent::ToolcallStart { partial, .. }
471        | AssistantMessageEvent::ToolcallDelta { partial, .. }
472        | AssistantMessageEvent::ToolcallEnd { partial, .. } => Some(partial),
473        AssistantMessageEvent::Done { .. } | AssistantMessageEvent::Error { .. } => None,
474    }
475}
476
477async fn finish_message(
478    current: &mut AgentContext,
479    message: AssistantMessage,
480    started: bool,
481    sink: &dyn EventSink,
482) -> AssistantMessage {
483    current
484        .messages
485        .push(AgentMessage::Assistant(message.clone()));
486    if !started {
487        sink.emit(AgentEvent::MessageStart {
488            message: AgentMessage::Assistant(message.clone()),
489        })
490        .await;
491    }
492    sink.emit(AgentEvent::MessageEnd {
493        message: AgentMessage::Assistant(message.clone()),
494    })
495    .await;
496    message
497}
498
499/// A settled tool call.
500#[derive(Clone, Debug, PartialEq)]
501pub struct ToolCallOutcome {
502    /// The call as issued.
503    pub tool_call: ToolCall,
504    /// Final result.
505    pub result: AgentToolResult,
506    /// Error flag.
507    pub is_error: bool,
508    /// Execution time (absent when the tool never ran).
509    pub duration_ms: Option<u64>,
510}
511
512struct ToolBatch {
513    messages: Vec<ToolResultMessage>,
514    terminate: bool,
515}
516
517async fn fail_truncated_tool_calls(tool_calls: &[ToolCall], sink: &dyn EventSink) -> ToolBatch {
518    let mut messages = Vec::new();
519    for call in tool_calls {
520        emit_start(call, sink).await;
521        let outcome = ToolCallOutcome {
522            tool_call: call.clone(),
523            result: AgentToolResult::error(format!(
524                "Tool call \"{}\" was not executed: the response hit the output token limit, so its arguments may be truncated. Re-issue the tool call with complete arguments.",
525                call.name
526            )),
527            is_error: true,
528            duration_ms: None,
529        };
530        emit_end(&outcome, sink).await;
531        let message = tool_result_message(&outcome);
532        emit_result_message(&message, sink).await;
533        messages.push(message);
534    }
535    ToolBatch {
536        messages,
537        terminate: false,
538    }
539}
540
541async fn execute_tool_calls(
542    current: &AgentContext,
543    assistant: &AssistantMessage,
544    tool_calls: Vec<ToolCall>,
545    config: &AgentLoopConfig,
546    sink: &dyn EventSink,
547    cancel: &CancellationToken,
548) -> ToolBatch {
549    let sequential = config.tool_execution == ToolExecution::Sequential
550        || tool_calls.iter().any(|call| {
551            current
552                .tool(&call.name)
553                .is_some_and(|tool| tool.meta().execution_mode == ExecutionMode::Sequential)
554        });
555    if sequential {
556        execute_sequential(current, assistant, tool_calls, config, sink, cancel).await
557    } else {
558        execute_parallel(current, assistant, tool_calls, config, sink, cancel).await
559    }
560}
561
562async fn execute_sequential(
563    current: &AgentContext,
564    assistant: &AssistantMessage,
565    tool_calls: Vec<ToolCall>,
566    config: &AgentLoopConfig,
567    sink: &dyn EventSink,
568    cancel: &CancellationToken,
569) -> ToolBatch {
570    let mut outcomes = Vec::new();
571    let mut messages = Vec::new();
572    for call in tool_calls {
573        emit_start(&call, sink).await;
574        let hooks = Hooks::from_config(config);
575        let outcome = match prepare_tool_call(
576            current,
577            assistant,
578            &call,
579            &current.tools,
580            &hooks,
581            cancel,
582        )
583        .await
584        {
585            Preparation::Immediate { result } => ToolCallOutcome {
586                tool_call: call.clone(),
587                result,
588                is_error: true,
589                duration_ms: None,
590            },
591            Preparation::Prepared { tool, args } => {
592                let executed =
593                    execute_prepared(&tool, &call, args.clone(), None, cancel, Some(sink), None)
594                        .await;
595                finalize(current, assistant, &call, &args, executed, &hooks, cancel).await
596            }
597        };
598        emit_end(&outcome, sink).await;
599        let message = tool_result_message(&outcome);
600        emit_result_message(&message, sink).await;
601        outcomes.push(outcome);
602        messages.push(message);
603        if cancel.is_cancelled() {
604            break;
605        }
606    }
607    ToolBatch {
608        messages,
609        terminate: should_terminate(&outcomes),
610    }
611}
612
613async fn execute_parallel(
614    current: &AgentContext,
615    assistant: &AssistantMessage,
616    tool_calls: Vec<ToolCall>,
617    config: &AgentLoopConfig,
618    sink: &dyn EventSink,
619    cancel: &CancellationToken,
620) -> ToolBatch {
621    let hooks = Hooks::from_config(config);
622    let mut slots: Vec<BoxFuture<'_, ToolCallOutcome>> = Vec::new();
623    for call in tool_calls {
624        emit_start(&call, sink).await;
625        match prepare_tool_call(current, assistant, &call, &current.tools, &hooks, cancel).await {
626            Preparation::Immediate { result } => {
627                let outcome = ToolCallOutcome {
628                    tool_call: call,
629                    result,
630                    is_error: true,
631                    duration_ms: None,
632                };
633                emit_end(&outcome, sink).await;
634                slots.push(Box::pin(std::future::ready(outcome)));
635            }
636            Preparation::Prepared { tool, args } => {
637                let hooks = hooks.clone();
638                slots.push(Box::pin(async move {
639                    if cancel.is_cancelled() {
640                        let outcome = ToolCallOutcome {
641                            tool_call: call,
642                            result: AgentToolResult::error("Operation aborted"),
643                            is_error: true,
644                            duration_ms: None,
645                        };
646                        emit_end(&outcome, sink).await;
647                        return outcome;
648                    }
649                    let executed = execute_prepared(
650                        &tool,
651                        &call,
652                        args.clone(),
653                        None,
654                        cancel,
655                        Some(sink),
656                        None,
657                    )
658                    .await;
659                    let outcome =
660                        finalize(current, assistant, &call, &args, executed, &hooks, cancel).await;
661                    emit_end(&outcome, sink).await;
662                    outcome
663                }));
664            }
665        }
666        if cancel.is_cancelled() {
667            break;
668        }
669    }
670    let outcomes = join_all(slots).await;
671    let mut messages = Vec::new();
672    for outcome in &outcomes {
673        let message = tool_result_message(outcome);
674        emit_result_message(&message, sink).await;
675        messages.push(message);
676    }
677    ToolBatch {
678        messages,
679        terminate: should_terminate(&outcomes),
680    }
681}
682
683fn should_terminate(outcomes: &[ToolCallOutcome]) -> bool {
684    !outcomes.is_empty() && outcomes.iter().all(|outcome| outcome.result.terminate)
685}
686
687#[derive(Clone)]
688struct Hooks {
689    before: Option<Arc<dyn BeforeToolCall>>,
690    after: Option<Arc<dyn AfterToolCall>>,
691}
692
693impl Hooks {
694    fn from_config(config: &AgentLoopConfig) -> Self {
695        Self {
696            before: config.before_tool_call.clone(),
697            after: config.after_tool_call.clone(),
698        }
699    }
700}
701
702enum Preparation {
703    Immediate {
704        result: AgentToolResult,
705    },
706    Prepared {
707        tool: Arc<dyn AgentTool>,
708        args: JsonObject,
709    },
710}
711
712async fn prepare_tool_call(
713    current: &AgentContext,
714    assistant: &AssistantMessage,
715    call: &ToolCall,
716    tools: &[Arc<dyn AgentTool>],
717    hooks: &Hooks,
718    cancel: &CancellationToken,
719) -> Preparation {
720    let Some(tool) = tools.iter().find(|tool| tool.meta().name == call.name) else {
721        return Preparation::Immediate {
722            result: AgentToolResult::error(format!("Tool {} not found", call.name)),
723        };
724    };
725    let mut prepared_call = call.clone();
726    if call.arguments_error.is_none() {
727        prepared_call.arguments = tool.prepare_arguments(call.arguments.clone());
728    }
729    let mut args = match validate_tool_arguments(&tool.meta().parameters, &prepared_call) {
730        Ok(args) => args,
731        Err(message) => {
732            return Preparation::Immediate {
733                result: AgentToolResult::error(message),
734            };
735        }
736    };
737    if let Some(before) = &hooks.before {
738        let cx = BeforeToolCallContext {
739            assistant_message: assistant,
740            tool_call: call,
741            args: &args,
742            context: current,
743        };
744        let decision = match before.before(cx, cancel).await {
745            Ok(decision) => decision,
746            Err(error) => {
747                return Preparation::Immediate {
748                    result: AgentToolResult::error(error.0),
749                };
750            }
751        };
752        if cancel.is_cancelled() {
753            return Preparation::Immediate {
754                result: AgentToolResult::error("Operation aborted"),
755            };
756        }
757        let decision = decision.map(|mut decision| {
758            if let Some(input) = decision.input.take() {
759                args = input;
760            }
761            decision
762        });
763        if let Some(decision) = decision.filter(|d| d.block) {
764            let mut result = AgentToolResult::error(
765                decision
766                    .reason
767                    .filter(|reason| !reason.is_empty())
768                    .unwrap_or_else(|| "Tool execution was blocked".to_owned()),
769            );
770            result.terminate = decision.terminate;
771            return Preparation::Immediate { result };
772        }
773    }
774    if cancel.is_cancelled() {
775        return Preparation::Immediate {
776            result: AgentToolResult::error("Operation aborted"),
777        };
778    }
779    Preparation::Prepared {
780        tool: Arc::clone(tool),
781        args,
782    }
783}
784
785struct Executed {
786    result: AgentToolResult,
787    is_error: bool,
788    duration_ms: u64,
789}
790
791/// Where partial results of a running tool go.
792type UpdateCallback<'a> = &'a (dyn Fn(AgentToolResult) -> BoxFuture<'a, ()> + Send + Sync);
793
794async fn execute_prepared(
795    tool: &Arc<dyn AgentTool>,
796    call: &ToolCall,
797    args: JsonObject,
798    parent_tool_call_id: Option<String>,
799    cancel: &CancellationToken,
800    sink: Option<&dyn EventSink>,
801    on_update: Option<UpdateCallback<'_>>,
802) -> Executed {
803    let (updates, mut receiver, open) = UpdateSink::channel();
804    let invocation = ToolInvocation {
805        tool_call_id: call.id.clone(),
806        args,
807        parent_tool_call_id,
808    };
809    let started = Instant::now();
810    let mut running = tool.execute(invocation, updates, cancel.clone());
811    let forward = |partial: AgentToolResult| async move {
812        if let Some(sink) = sink {
813            sink.emit(AgentEvent::ToolExecutionUpdate {
814                tool_call_id: call.id.clone(),
815                tool_name: call.name.clone(),
816                args: call.arguments.clone(),
817                partial_result: partial.clone(),
818            })
819            .await;
820        }
821        if let Some(on_update) = on_update {
822            on_update(partial).await;
823        }
824    };
825    let outcome = loop {
826        tokio::select! {
827            biased;
828            result = &mut running => break result,
829            Some(partial) = receiver.recv() => forward(partial).await,
830        }
831    };
832    let duration_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
833    open.store(false, Ordering::Release);
834    while let Ok(partial) = receiver.try_recv() {
835        forward(partial).await;
836    }
837    match outcome {
838        Ok(result) => Executed {
839            is_error: result.is_error,
840            result,
841            duration_ms,
842        },
843        Err(error) => Executed {
844            result: AgentToolResult::error(error.0),
845            is_error: true,
846            duration_ms,
847        },
848    }
849}
850
851async fn finalize(
852    current: &AgentContext,
853    assistant: &AssistantMessage,
854    call: &ToolCall,
855    args: &JsonObject,
856    executed: Executed,
857    hooks: &Hooks,
858    cancel: &CancellationToken,
859) -> ToolCallOutcome {
860    let mut result = executed.result;
861    let mut is_error = executed.is_error;
862    if let Some(after) = &hooks.after {
863        let cx = AfterToolCallContext {
864            assistant_message: assistant,
865            tool_call: call,
866            args,
867            result: &result,
868            is_error,
869            context: current,
870        };
871        match after.after(cx, cancel).await {
872            Ok(Some(update)) => {
873                let structured = match (&update.structured_content, &update.content) {
874                    (Some(structured), _) => Some(structured.clone()),
875                    (None, Some(_)) => None,
876                    (None, None) => result.structured_content.clone(),
877                };
878                if let Some(content) = update.content {
879                    result.content = content;
880                }
881                if let Some(details) = update.details {
882                    result.details = Some(details);
883                }
884                if let Some(usage) = update.usage {
885                    result.usage = Some(usage);
886                }
887                if let Some(terminate) = update.terminate {
888                    result.terminate = terminate;
889                }
890                result.structured_content = structured;
891                is_error = update.is_error.unwrap_or(is_error);
892            }
893            Ok(None) => {}
894            Err(error) => {
895                result = AgentToolResult::error(error.0);
896                is_error = true;
897            }
898        }
899    }
900    ToolCallOutcome {
901        tool_call: call.clone(),
902        result,
903        is_error,
904        duration_ms: Some(executed.duration_ms),
905    }
906}
907
908async fn emit_start(call: &ToolCall, sink: &dyn EventSink) {
909    sink.emit(AgentEvent::ToolExecutionStart {
910        tool_call_id: call.id.clone(),
911        tool_name: call.name.clone(),
912        args: call.arguments.clone(),
913    })
914    .await;
915}
916
917async fn emit_end(outcome: &ToolCallOutcome, sink: &dyn EventSink) {
918    sink.emit(AgentEvent::ToolExecutionEnd {
919        tool_call_id: outcome.tool_call.id.clone(),
920        tool_name: outcome.tool_call.name.clone(),
921        result: outcome.result.clone(),
922        is_error: outcome.is_error,
923        duration_ms: outcome.duration_ms,
924    })
925    .await;
926}
927
928/// C r e a t e   T o o l   R e s u l t   M e s s a g e.
929#[must_use]
930pub fn tool_result_message(outcome: &ToolCallOutcome) -> ToolResultMessage {
931    ToolResultMessage {
932        tool_call_id: outcome.tool_call.id.clone(),
933        tool_name: outcome.tool_call.name.clone(),
934        content: outcome.result.content.clone(),
935        details: outcome.result.details.clone(),
936        usage: outcome.result.usage,
937        nested_calls: None,
938        is_error: outcome.is_error,
939        timestamp: now_ms(),
940        duration_ms: outcome.duration_ms,
941    }
942}
943
944async fn emit_result_message(message: &ToolResultMessage, sink: &dyn EventSink) {
945    sink.emit(AgentEvent::MessageStart {
946        message: AgentMessage::ToolResult(message.clone()),
947    })
948    .await;
949    sink.emit(AgentEvent::MessageEnd {
950        message: AgentMessage::ToolResult(message.clone()),
951    })
952    .await;
953}
954
955/// Options of [`run_tool_call`].
956pub struct RunToolCallOptions<'a> {
957    /// Assistant message the call belongs to.
958    pub assistant_message: &'a AssistantMessage,
959    /// Context the hooks see.
960    pub context: &'a AgentContext,
961    /// Tools to resolve the call against.
962    pub tools: &'a [Arc<dyn AgentTool>],
963    /// Tool gate.
964    pub before_tool_call: Option<Arc<dyn BeforeToolCall>>,
965    /// Result post-processing.
966    pub after_tool_call: Option<Arc<dyn AfterToolCall>>,
967    /// Abort signal.
968    pub cancel: CancellationToken,
969    /// Calling tool for nested calls (`ctx.executeTool`).
970    pub parent_tool_call_id: Option<String>,
971    /// Partial-result callback.
972    pub on_update: Option<UpdateCallback<'a>>,
973}
974
975/// One call through preparation, validation, hooks and execution. Emits no
976/// events and never fails (errors become `is_error` outcomes).
977pub async fn run_tool_call(call: ToolCall, options: RunToolCallOptions<'_>) -> ToolCallOutcome {
978    let hooks = Hooks {
979        before: options.before_tool_call.clone(),
980        after: options.after_tool_call.clone(),
981    };
982    match prepare_tool_call(
983        options.context,
984        options.assistant_message,
985        &call,
986        options.tools,
987        &hooks,
988        &options.cancel,
989    )
990    .await
991    {
992        Preparation::Immediate { result } => ToolCallOutcome {
993            tool_call: call,
994            result,
995            is_error: true,
996            duration_ms: None,
997        },
998        Preparation::Prepared { tool, args } => {
999            let executed = execute_prepared(
1000                &tool,
1001                &call,
1002                args.clone(),
1003                options.parent_tool_call_id.clone(),
1004                &options.cancel,
1005                None,
1006                options.on_update,
1007            )
1008            .await;
1009            finalize(
1010                options.context,
1011                options.assistant_message,
1012                &call,
1013                &args,
1014                executed,
1015                &hooks,
1016                &options.cancel,
1017            )
1018            .await
1019        }
1020    }
1021}
oino_agent/agent_loop.rs

oino_agent/
agent_loop.rs
