oino_agent/
agent.rs

1//! Stateful agent (`Agent`).
2//!
3//! [`Agent`] owns the transcript, reduces loop events into [`AgentState`], awaits listeners in
4//! subscription order, and exposes the steering / follow-up queues. It is cheap to clone; all
5//! clones share one agent.
6
7use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
8use std::sync::{Arc, Mutex, MutexGuard, PoisonError, Weak};
9
10use futures::future::BoxFuture;
11use indexmap::IndexSet;
12use oino_ai::transcript::{SystemState, initial_system_message, system_message_text};
13use oino_ai::{
14    AgentMessage, AssistantContent, AssistantMessage, ImageContent, Model, StopReason,
15    StreamOptions, ThinkingLevel, UserContent, UserContentField, UserMessage,
16};
17use serde::{Deserialize, Serialize};
18use tokio::sync::watch;
19use tokio_util::sync::CancellationToken;
20
21use crate::agent_loop::{AgentLoopError, agent_loop, agent_loop_continue, now_ms};
22use crate::config::{
23    AfterToolCall, AgentContext, AgentLoopConfig, BeforeToolCall, ConvertToLlm, FinishTurn,
24    GetApiKey, PrepareNextTurn, PrepareRequest, QueueSource, StreamFn, ToolExecution,
25    TransformContext, default_convert_to_llm,
26};
27use crate::event::{AgentEvent, EventSink};
28use crate::tool::AgentTool;
29
30/// How queued messages are drained.
31#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
32pub enum QueueMode {
33    /// Every queued message at once.
34    #[serde(rename = "all")]
35    All,
36    /// One message per turn (default).
37    #[default]
38    #[serde(rename = "one-at-a-time")]
39    OneAtATime,
40}
41
42/// Operation refused by the agent.
43#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
44#[non_exhaustive]
45pub enum AgentError {
46    /// `prompt` while a run is active.
47    #[error(
48        "Agent is already processing a prompt. Use steer() or followUp() to queue messages, or wait for completion."
49    )]
50    PromptWhileBusy,
51    /// `continue_run` while a run is active.
52    #[error("Agent is already processing. Wait for completion before continuing.")]
53    ContinueWhileBusy,
54    /// `reset` while a run is active.
55    #[error("Agent is already processing. Wait for completion before resetting.")]
56    ResetWhileBusy,
57    /// `continue_run` on an empty (or system-only) transcript.
58    #[error("No messages to continue from")]
59    NoMessages,
60    /// `continue_run` after an assistant message with empty queues.
61    #[error("Cannot continue from message role: assistant")]
62    CannotContinue,
63}
64
65/// Snapshot of the agent state.
66#[derive(Clone)]
67pub struct AgentState {
68    /// Model of the next request.
69    pub model: Model,
70    /// Thinking level of the next request.
71    pub thinking_level: ThinkingLevel,
72    /// Executable tools.
73    pub tools: Vec<Arc<dyn AgentTool>>,
74    /// Transcript.
75    pub messages: Vec<AgentMessage>,
76    /// A run is active.
77    pub is_streaming: bool,
78    /// The message currently streaming.
79    pub streaming_message: Option<AgentMessage>,
80    /// Tool calls started but not settled, in start order.
81    pub pending_tool_calls: IndexSet<String>,
82    /// Error of the last failed turn.
83    pub error_message: Option<String>,
84}
85
86impl std::fmt::Debug for AgentState {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("AgentState")
89            .field("model", &self.model.key())
90            .field("thinking_level", &self.thinking_level)
91            .field("messages", &self.messages.len())
92            .field("is_streaming", &self.is_streaming)
93            .field("error_message", &self.error_message)
94            .finish_non_exhaustive()
95    }
96}
97
98impl AgentState {
99    /// The replayed prompt text of the transcript.
100    #[must_use]
101    pub fn system_prompt(&self) -> String {
102        system_prompt_of(&self.messages)
103    }
104}
105
106/// Replayed system prompt text of `messages`.
107#[must_use]
108pub fn system_prompt_of(messages: &[AgentMessage]) -> String {
109    SystemState::fold(messages.iter().filter_map(|m| match m {
110        AgentMessage::System(system) => Some(system),
111        _ => None,
112    }))
113    .into_message()
114    .map(|message| system_message_text(&message))
115    .unwrap_or_default()
116}
117
118/// The default model id.
119#[must_use]
120pub fn unknown_model() -> Model {
121    Model {
122        id: "unknown".into(),
123        name: "unknown".into(),
124        api: "unknown".into(),
125        provider: "unknown".into(),
126        base_url: String::new(),
127        input: Vec::new(),
128        input_limits: None,
129        cost: oino_ai::ModelCost::default(),
130        headers: None,
131        kind: None,
132        reasoning: false,
133        thinking_level_map: None,
134        prompt_cache: None,
135        context_window: 0,
136        max_tokens: 0,
137        sampling_params: None,
138        sampling_params_by_thinking_level: None,
139        compat: None,
140    }
141}
142
143/// Hooks forwarded to every run.
144#[derive(Clone, Default)]
145pub struct AgentHooks {
146    /// Context rewrite.
147    pub transform_context: Option<Arc<dyn TransformContext>>,
148    /// Per-request credential.
149    pub get_api_key: Option<Arc<dyn GetApiKey>>,
150    /// Tool gate.
151    pub before_tool_call: Option<Arc<dyn BeforeToolCall>>,
152    /// Tool result post-processing.
153    pub after_tool_call: Option<Arc<dyn AfterToolCall>>,
154    /// Turn decision.
155    pub finish_turn: Option<Arc<dyn FinishTurn>>,
156    /// Per-request update.
157    pub prepare_request: Option<Arc<dyn PrepareRequest>>,
158    /// Between-turn update.
159    pub prepare_next_turn: Option<Arc<dyn PrepareNextTurn>>,
160}
161
162/// Construction options.
163#[derive(Clone)]
164pub struct AgentOptions {
165    /// Initial system prompt (becomes the leading system message).
166    pub system_prompt: String,
167    /// Initial model.
168    pub model: Model,
169    /// Initial thinking level.
170    pub thinking_level: ThinkingLevel,
171    /// Initial tools.
172    pub tools: Vec<Arc<dyn AgentTool>>,
173    /// Initial transcript.
174    pub messages: Vec<AgentMessage>,
175    /// Provider dispatch.
176    pub stream: Arc<dyn StreamFn>,
177    /// Request-boundary conversion.
178    pub convert_to_llm: ConvertToLlm,
179    /// Hooks.
180    pub hooks: AgentHooks,
181    /// Steering drain mode.
182    pub steering_mode: QueueMode,
183    /// Follow-up drain mode.
184    pub follow_up_mode: QueueMode,
185    /// Base stream options (`sessionId`, `transport`, budgets, retry caps).
186    pub stream_options: StreamOptions,
187    /// Batch scheduling.
188    pub tool_execution: ToolExecution,
189}
190
191impl AgentOptions {
192    /// Options with defaults for `stream`.
193    #[must_use]
194    pub fn new(stream: Arc<dyn StreamFn>) -> Self {
195        Self {
196            system_prompt: String::new(),
197            model: unknown_model(),
198            thinking_level: ThinkingLevel::Off,
199            tools: Vec::new(),
200            messages: Vec::new(),
201            stream,
202            convert_to_llm: Arc::new(default_convert_to_llm),
203            hooks: AgentHooks::default(),
204            steering_mode: QueueMode::OneAtATime,
205            follow_up_mode: QueueMode::OneAtATime,
206            stream_options: StreamOptions {
207                transport: Some(oino_ai::Transport::Auto),
208                ..StreamOptions::default()
209            },
210            tool_execution: ToolExecution::Parallel,
211        }
212    }
213}
214
215/// Listener of agent events (the `subscribe` callback). Awaited in subscription order.
216pub trait AgentListener: Send + Sync {
217    /// Handles one event of the active run.
218    fn on_event<'a>(
219        &'a self,
220        event: &'a AgentEvent,
221        cancel: &'a CancellationToken,
222    ) -> BoxFuture<'a, ()>;
223}
224
225/// Unsubscribes on drop.
226#[must_use = "dropping the subscription unsubscribes"]
227pub struct Subscription {
228    inner: Weak<Inner>,
229    id: u64,
230}
231
232impl Drop for Subscription {
233    fn drop(&mut self) {
234        if let Some(inner) = self.inner.upgrade() {
235            lock(&inner.listeners).retain(|(id, _)| *id != self.id);
236        }
237    }
238}
239
240impl std::fmt::Debug for Subscription {
241    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
242        f.debug_struct("Subscription")
243            .field("id", &self.id)
244            .finish_non_exhaustive()
245    }
246}
247
248#[derive(Debug)]
249struct PendingQueue {
250    messages: Vec<AgentMessage>,
251    mode: QueueMode,
252}
253
254impl PendingQueue {
255    const fn new(mode: QueueMode) -> Self {
256        Self {
257            messages: Vec::new(),
258            mode,
259        }
260    }
261
262    fn peek(&self) -> Vec<AgentMessage> {
263        match self.mode {
264            QueueMode::All => self.messages.clone(),
265            QueueMode::OneAtATime => self.messages.first().cloned().into_iter().collect(),
266        }
267    }
268
269    fn drain(&mut self) -> Vec<AgentMessage> {
270        let drained = self.peek();
271        self.messages.drain(..drained.len());
272        drained
273    }
274}
275
276struct ActiveRun {
277    cancel: CancellationToken,
278}
279
280struct Settings {
281    stream: Arc<dyn StreamFn>,
282    convert_to_llm: ConvertToLlm,
283    hooks: AgentHooks,
284    stream_options: StreamOptions,
285    tool_execution: ToolExecution,
286}
287
288struct Inner {
289    state: Mutex<AgentState>,
290    listeners: Mutex<Vec<(u64, Arc<dyn AgentListener>)>>,
291    steering: Mutex<PendingQueue>,
292    follow_up: Mutex<PendingQueue>,
293    settings: Mutex<Settings>,
294    active: Mutex<Option<ActiveRun>>,
295    idle: watch::Sender<bool>,
296    next_listener: AtomicU64,
297}
298
299fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
300    mutex.lock().unwrap_or_else(PoisonError::into_inner)
301}
302
303/// Stateful wrapper around the agent loop.
304#[derive(Clone)]
305pub struct Agent {
306    inner: Arc<Inner>,
307}
308
309impl std::fmt::Debug for Agent {
310    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
311        f.debug_struct("Agent")
312            .field("state", &*lock(&self.inner.state))
313            .finish_non_exhaustive()
314    }
315}
316
317impl Agent {
318    /// Creates an agent; a leading system message is added for the prompt and tools unless the
319    /// transcript already starts with one.
320    #[must_use]
321    pub fn new(options: AgentOptions) -> Self {
322        let mut messages = options.messages;
323        let declarations: Vec<_> = options
324            .tools
325            .iter()
326            .map(|t| t.meta().declaration())
327            .collect();
328        if !matches!(messages.first(), Some(AgentMessage::System(_)))
329            && let Some(initial) = initial_system_message(&options.system_prompt, &declarations)
330        {
331            messages.insert(0, AgentMessage::System(initial));
332        }
333        let (idle, _) = watch::channel(true);
334        Self {
335            inner: Arc::new(Inner {
336                state: Mutex::new(AgentState {
337                    model: options.model,
338                    thinking_level: options.thinking_level,
339                    tools: options.tools,
340                    messages,
341                    is_streaming: false,
342                    streaming_message: None,
343                    pending_tool_calls: IndexSet::new(),
344                    error_message: None,
345                }),
346                listeners: Mutex::new(Vec::new()),
347                steering: Mutex::new(PendingQueue::new(options.steering_mode)),
348                follow_up: Mutex::new(PendingQueue::new(options.follow_up_mode)),
349                settings: Mutex::new(Settings {
350                    stream: options.stream,
351                    convert_to_llm: options.convert_to_llm,
352                    hooks: options.hooks,
353                    stream_options: options.stream_options,
354                    tool_execution: options.tool_execution,
355                }),
356                active: Mutex::new(None),
357                idle,
358                next_listener: AtomicU64::new(1),
359            }),
360        }
361    }
362
363    /// Registers a listener; dropping the returned [`Subscription`] removes it.
364    pub fn subscribe(&self, listener: Arc<dyn AgentListener>) -> Subscription {
365        let id = self.inner.next_listener.fetch_add(1, Ordering::Relaxed);
366        lock(&self.inner.listeners).push((id, listener));
367        Subscription {
368            inner: Arc::downgrade(&self.inner),
369            id,
370        }
371    }
372
373    /// Snapshot of the state.
374    #[must_use]
375    pub fn state(&self) -> AgentState {
376        lock(&self.inner.state).clone()
377    }
378
379    /// Runs `f` on the state under the lock (avoids cloning the transcript).
380    pub fn with_state<R>(&self, f: impl FnOnce(&AgentState) -> R) -> R {
381        f(&lock(&self.inner.state))
382    }
383
384    /// Copy of the transcript.
385    #[must_use]
386    pub fn messages(&self) -> Vec<AgentMessage> {
387        lock(&self.inner.state).messages.clone()
388    }
389
390    /// Replaces the transcript (copy-on-set).
391    pub fn set_messages(&self, messages: Vec<AgentMessage>) {
392        lock(&self.inner.state).messages = messages;
393    }
394
395    /// Appends one message to the transcript.
396    pub fn append_message(&self, message: AgentMessage) {
397        lock(&self.inner.state).messages.push(message);
398    }
399
400    /// Replaces the last transcript message (in-place `message_end` replacement).
401    pub fn replace_last_message(&self, message: AgentMessage) {
402        if let Some(last) = lock(&self.inner.state).messages.last_mut() {
403            *last = message;
404        }
405    }
406
407    /// Replaces the executable tools.
408    pub fn set_tools(&self, tools: Vec<Arc<dyn AgentTool>>) {
409        lock(&self.inner.state).tools = tools;
410    }
411
412    /// Current model.
413    #[must_use]
414    pub fn model(&self) -> Model {
415        lock(&self.inner.state).model.clone()
416    }
417
418    /// Sets the model of the next request.
419    pub fn set_model(&self, model: Model) {
420        lock(&self.inner.state).model = model;
421    }
422
423    /// Current thinking level.
424    #[must_use]
425    pub fn thinking_level(&self) -> ThinkingLevel {
426        lock(&self.inner.state).thinking_level
427    }
428
429    /// Sets the thinking level of the next request.
430    pub fn set_thinking_level(&self, level: ThinkingLevel) {
431        lock(&self.inner.state).thinking_level = level;
432    }
433
434    /// S t a t e . s y s t e m   P r o m p t.
435    #[must_use]
436    pub fn system_prompt(&self) -> String {
437        system_prompt_of(&lock(&self.inner.state).messages)
438    }
439
440    /// Whether a run is active.
441    #[must_use]
442    pub fn is_streaming(&self) -> bool {
443        lock(&self.inner.state).is_streaming
444    }
445
446    /// Replaces the hooks.
447    pub fn set_hooks(&self, hooks: AgentHooks) {
448        lock(&self.inner.settings).hooks = hooks;
449    }
450
451    /// Updates the base stream options (`sessionId`, `transport`, budgets, …).
452    pub fn update_stream_options(&self, f: impl FnOnce(&mut StreamOptions)) {
453        f(&mut lock(&self.inner.settings).stream_options);
454    }
455
456    /// Sets the batch scheduling.
457    pub fn set_tool_execution(&self, mode: ToolExecution) {
458        lock(&self.inner.settings).tool_execution = mode;
459    }
460
461    /// Steering drain mode.
462    #[must_use]
463    pub fn steering_mode(&self) -> QueueMode {
464        lock(&self.inner.steering).mode
465    }
466
467    /// Sets the steering drain mode.
468    pub fn set_steering_mode(&self, mode: QueueMode) {
469        lock(&self.inner.steering).mode = mode;
470    }
471
472    /// Follow-up drain mode.
473    #[must_use]
474    pub fn follow_up_mode(&self) -> QueueMode {
475        lock(&self.inner.follow_up).mode
476    }
477
478    /// Sets the follow-up drain mode.
479    pub fn set_follow_up_mode(&self, mode: QueueMode) {
480        lock(&self.inner.follow_up).mode = mode;
481    }
482
483    /// Queues a message injected after the current turn.
484    pub fn steer(&self, message: AgentMessage) {
485        lock(&self.inner.steering).messages.push(message);
486    }
487
488    /// Queues a message delivered when the agent would otherwise stop.
489    pub fn follow_up(&self, message: AgentMessage) {
490        lock(&self.inner.follow_up).messages.push(message);
491    }
492
493    /// Edits the steering (`steering = true`) or follow-up queue in place, atomically with
494    /// respect to draining (single-item queue edits).
495    pub fn update_queue<R>(
496        &self,
497        steering: bool,
498        edit: impl FnOnce(&mut Vec<AgentMessage>) -> R,
499    ) -> R {
500        let queue = if steering {
501            &self.inner.steering
502        } else {
503            &self.inner.follow_up
504        };
505        edit(&mut lock(queue).messages)
506    }
507
508    /// Clears the steering queue.
509    pub fn clear_steering_queue(&self) {
510        lock(&self.inner.steering).messages.clear();
511    }
512
513    /// Clears the follow-up queue.
514    pub fn clear_follow_up_queue(&self) {
515        lock(&self.inner.follow_up).messages.clear();
516    }
517
518    /// Clears both queues.
519    pub fn clear_all_queues(&self) {
520        self.clear_steering_queue();
521        self.clear_follow_up_queue();
522    }
523
524    /// Whether either queue has messages.
525    #[must_use]
526    pub fn has_queued_messages(&self) -> bool {
527        !lock(&self.inner.steering).messages.is_empty()
528            || !lock(&self.inner.follow_up).messages.is_empty()
529    }
530
531    /// Messages the next turn would take, without consuming them.
532    #[must_use]
533    pub fn peek_queued_messages(&self) -> Vec<AgentMessage> {
534        let steering = lock(&self.inner.steering).peek();
535        if steering.is_empty() {
536            lock(&self.inner.follow_up).peek()
537        } else {
538            steering
539        }
540    }
541
542    /// Every queued steering message (in order).
543    #[must_use]
544    pub fn steering_queue(&self) -> Vec<AgentMessage> {
545        lock(&self.inner.steering).messages.clone()
546    }
547
548    /// Every queued follow-up message (in order).
549    #[must_use]
550    pub fn follow_up_queue(&self) -> Vec<AgentMessage> {
551        lock(&self.inner.follow_up).messages.clone()
552    }
553
554    /// Abort signal of the active run.
555    #[must_use]
556    pub fn signal(&self) -> Option<CancellationToken> {
557        lock(&self.inner.active)
558            .as_ref()
559            .map(|run| run.cancel.clone())
560    }
561
562    /// Aborts the active run, if any.
563    pub fn abort(&self) {
564        if let Some(run) = lock(&self.inner.active).as_ref() {
565            run.cancel.cancel();
566        }
567    }
568
569    /// Clears a run orphaned by a crashed driver task (X-PANIC): a panic's unwind dropped
570    /// the run future before `finish_run` ran, leaving the run "active" and the session
571    /// forever busy. Discards streaming state and pending tool calls; the transcript is
572    /// kept — callers rebuild from the session log. No-op when no run is active.
573    pub fn recover_orphaned_run(&self) {
574        if lock(&self.inner.active).take().is_some() {
575            {
576                let mut state = lock(&self.inner.state);
577                state.is_streaming = false;
578                state.streaming_message = None;
579                state.pending_tool_calls.clear();
580            }
581            self.inner.idle.send_replace(true);
582        }
583    }
584
585    /// Resolves once no run is active and every `agent_end` listener has returned.
586    pub async fn wait_for_idle(&self) {
587        let mut idle = self.inner.idle.subscribe();
588        // `wait_for` only fails when the sender is dropped, which cannot happen while `self`
589        // keeps `inner` alive.
590        let _ = idle.wait_for(|idle| *idle).await;
591    }
592
593    /// Clears the transcript (keeping the replayed system baseline) and both queues.
594    ///
595    /// # Errors
596    ///
597    /// [`AgentError::ResetWhileBusy`] while a run is active.
598    pub fn reset(&self) -> Result<(), AgentError> {
599        if lock(&self.inner.active).is_some() {
600            return Err(AgentError::ResetWhileBusy);
601        }
602        {
603            let mut state = lock(&self.inner.state);
604            let baseline = SystemState::fold(state.messages.iter().filter_map(|m| match m {
605                AgentMessage::System(system) => Some(system),
606                _ => None,
607            }))
608            .into_message();
609            state.messages = baseline.map(AgentMessage::System).into_iter().collect();
610            state.is_streaming = false;
611            state.streaming_message = None;
612            state.pending_tool_calls.clear();
613            state.error_message = None;
614        }
615        self.clear_all_queues();
616        Ok(())
617    }
618
619    /// A user message with text and images.
620    ///
621    /// # Errors
622    ///
623    /// [`AgentError::PromptWhileBusy`] while a run is active.
624    pub async fn prompt_text(
625        &self,
626        text: impl Into<String>,
627        images: Vec<ImageContent>,
628    ) -> Result<(), AgentError> {
629        let mut content = vec![UserContent::text(text)];
630        content.extend(images.into_iter().map(UserContent::Image));
631        self.prompt(vec![AgentMessage::User(UserMessage {
632            content: UserContentField::Blocks(content),
633            timestamp: now_ms(),
634        })])
635        .await
636    }
637
638    /// Starts a run with `messages`; resolves when the run is idle.
639    ///
640    /// # Errors
641    ///
642    /// [`AgentError::PromptWhileBusy`] while a run is active.
643    pub async fn prompt(&self, messages: Vec<AgentMessage>) -> Result<(), AgentError> {
644        let cancel = self.begin_run().ok_or(AgentError::PromptWhileBusy)?;
645        self.run(RunKind::Prompt(messages, false), cancel).await;
646        Ok(())
647    }
648
649    /// Continues from the transcript; after an assistant message the queues
650    /// are drained instead.
651    ///
652    /// # Errors
653    ///
654    /// [`AgentError::ContinueWhileBusy`], [`AgentError::NoMessages`],
655    /// [`AgentError::CannotContinue`].
656    pub async fn continue_run(&self) -> Result<(), AgentError> {
657        if lock(&self.inner.active).is_some() {
658            return Err(AgentError::ContinueWhileBusy);
659        }
660        let (last_is_assistant, empty) = {
661            let state = lock(&self.inner.state);
662            (
663                matches!(state.messages.last(), Some(AgentMessage::Assistant(_))),
664                state
665                    .messages
666                    .iter()
667                    .all(|m| matches!(m, AgentMessage::System(_))),
668            )
669        };
670        if empty {
671            return Err(AgentError::NoMessages);
672        }
673        let kind = if last_is_assistant {
674            let steering = lock(&self.inner.steering).drain();
675            if steering.is_empty() {
676                let follow_ups = lock(&self.inner.follow_up).drain();
677                if follow_ups.is_empty() {
678                    return Err(AgentError::CannotContinue);
679                }
680                RunKind::Prompt(follow_ups, false)
681            } else {
682                RunKind::Prompt(steering, true)
683            }
684        } else {
685            RunKind::Continue
686        };
687        let cancel = self.begin_run().ok_or(AgentError::ContinueWhileBusy)?;
688        self.run(kind, cancel).await;
689        Ok(())
690    }
691
692    fn begin_run(&self) -> Option<CancellationToken> {
693        let mut active = lock(&self.inner.active);
694        if active.is_some() {
695            return None;
696        }
697        let cancel = CancellationToken::new();
698        *active = Some(ActiveRun {
699            cancel: cancel.clone(),
700        });
701        drop(active);
702        self.inner.idle.send_replace(false);
703        let mut state = lock(&self.inner.state);
704        state.is_streaming = true;
705        state.streaming_message = None;
706        state.error_message = None;
707        Some(cancel)
708    }
709
710    async fn run(&self, kind: RunKind, cancel: CancellationToken) {
711        let (context, config) = self.loop_inputs(matches!(kind, RunKind::Prompt(_, true)));
712        let sink = AgentSink {
713            inner: Arc::clone(&self.inner),
714            cancel: cancel.clone(),
715        };
716        let outcome = match kind {
717            RunKind::Prompt(messages, _) => {
718                agent_loop(messages, context, &config, &sink, cancel.clone()).await
719            }
720            RunKind::Continue => agent_loop_continue(context, &config, &sink, cancel.clone()).await,
721        };
722        if let Err(error) = outcome {
723            self.handle_run_failure(&error, cancel.is_cancelled(), &sink)
724                .await;
725        }
726        self.finish_run();
727    }
728
729    fn loop_inputs(&self, skip_initial_steering_poll: bool) -> (AgentContext, AgentLoopConfig) {
730        let (context, model, thinking_level) = {
731            let state = lock(&self.inner.state);
732            (
733                AgentContext {
734                    messages: state.messages.clone(),
735                    tools: state.tools.clone(),
736                },
737                state.model.clone(),
738                state.thinking_level,
739            )
740        };
741        let settings = lock(&self.inner.settings);
742        let hooks = settings.hooks.clone();
743        let config = AgentLoopConfig {
744            model,
745            thinking_level,
746            stream_options: settings.stream_options.clone(),
747            stream: Arc::clone(&settings.stream),
748            convert_to_llm: Arc::clone(&settings.convert_to_llm),
749            transform_context: hooks.transform_context,
750            get_api_key: hooks.get_api_key,
751            api_key: None,
752            before_tool_call: hooks.before_tool_call,
753            after_tool_call: hooks.after_tool_call,
754            finish_turn: hooks.finish_turn,
755            prepare_request: hooks.prepare_request,
756            prepare_next_turn: hooks.prepare_next_turn,
757            get_steering_messages: Some(Arc::new(QueueDrain {
758                inner: Arc::downgrade(&self.inner),
759                steering: true,
760                skip_next: AtomicBool::new(skip_initial_steering_poll),
761            })),
762            get_follow_up_messages: Some(Arc::new(QueueDrain {
763                inner: Arc::downgrade(&self.inner),
764                steering: false,
765                skip_next: AtomicBool::new(false),
766            })),
767            tool_execution: settings.tool_execution,
768        };
769        (context, config)
770    }
771
772    async fn handle_run_failure(&self, error: &AgentLoopError, aborted: bool, sink: &AgentSink) {
773        let model = self.model();
774        let mut failure = AssistantMessage::empty(&model.api, &model.provider, &model.id, now_ms());
775        failure.content = vec![AssistantContent::text("")];
776        failure.stop_reason = if aborted {
777            StopReason::Aborted
778        } else {
779            StopReason::Error
780        };
781        failure.error_message = Some(error.to_string());
782        let message = AgentMessage::Assistant(failure);
783        sink.emit(AgentEvent::MessageStart {
784            message: message.clone(),
785        })
786        .await;
787        sink.emit(AgentEvent::MessageEnd {
788            message: message.clone(),
789        })
790        .await;
791        sink.emit(AgentEvent::TurnEnd {
792            message: message.clone(),
793            tool_results: Vec::new(),
794        })
795        .await;
796        sink.emit(AgentEvent::AgentEnd {
797            messages: vec![message],
798        })
799        .await;
800    }
801
802    fn finish_run(&self) {
803        {
804            let mut state = lock(&self.inner.state);
805            state.is_streaming = false;
806            state.streaming_message = None;
807            state.pending_tool_calls.clear();
808        }
809        *lock(&self.inner.active) = None;
810        self.inner.idle.send_replace(true);
811    }
812}
813
814enum RunKind {
815    /// Prompt messages; `true` skips the initial steering poll (they came from steering).
816    Prompt(Vec<AgentMessage>, bool),
817    Continue,
818}
819
820struct QueueDrain {
821    inner: Weak<Inner>,
822    steering: bool,
823    skip_next: AtomicBool,
824}
825
826impl QueueSource for QueueDrain {
827    fn drain(&self) -> BoxFuture<'_, Vec<AgentMessage>> {
828        Box::pin(async move {
829            if self.skip_next.swap(false, Ordering::AcqRel) {
830                return Vec::new();
831            }
832            let Some(inner) = self.inner.upgrade() else {
833                return Vec::new();
834            };
835            if self.steering {
836                lock(&inner.steering).drain()
837            } else {
838                lock(&inner.follow_up).drain()
839            }
840        })
841    }
842}
843
844struct AgentSink {
845    inner: Arc<Inner>,
846    cancel: CancellationToken,
847}
848
849impl AgentSink {
850    fn reduce(&self, event: &AgentEvent) {
851        let mut state = lock(&self.inner.state);
852        match event {
853            AgentEvent::MessageStart { message } | AgentEvent::MessageUpdate { message, .. } => {
854                state.streaming_message = Some(message.clone());
855            }
856            AgentEvent::MessageEnd { message } => {
857                state.streaming_message = None;
858                state.messages.push(message.clone());
859            }
860            AgentEvent::ToolExecutionStart { tool_call_id, .. } => {
861                state.pending_tool_calls.insert(tool_call_id.clone());
862            }
863            AgentEvent::ToolExecutionEnd { tool_call_id, .. } => {
864                state.pending_tool_calls.shift_remove(tool_call_id);
865            }
866            AgentEvent::TurnEnd {
867                message: AgentMessage::Assistant(message),
868                ..
869            } => {
870                if let Some(error) = &message.error_message {
871                    state.error_message = Some(error.clone());
872                }
873            }
874            AgentEvent::AgentEnd { .. } => state.streaming_message = None,
875            _ => {}
876        }
877    }
878}
879
880impl EventSink for AgentSink {
881    fn emit(&self, event: AgentEvent) -> BoxFuture<'_, ()> {
882        Box::pin(async move {
883            self.reduce(&event);
884            let listeners: Vec<Arc<dyn AgentListener>> = lock(&self.inner.listeners)
885                .iter()
886                .map(|(_, listener)| Arc::clone(listener))
887                .collect();
888            for listener in listeners {
889                listener.on_event(&event, &self.cancel).await;
890            }
891        })
892    }
893}
oino_agent/agent.rs

oino_agent/
agent.rs
