Skip to main content

oino_agent/
config.rs

1//! Loop configuration and hooks (`AgentLoopConfig`).
2//!
3//! Hooks are traits so that `oino-core` can implement them on its session type. Every hook
4//! returns a boxed future; an `Err(HookError)` from a loop-level hook ends the run the way a
5//! thrown error ends the run (the [`crate::Agent`] turns it into a failure message). Errors
6//! from tool hooks become error tool results instead.
7
8use std::sync::Arc;
9
10use futures::future::BoxFuture;
11use futures::stream::BoxStream;
12use oino_ai::{
13    AgentMessage, AssistantMessage, AssistantMessageEvent, Context, JsonObject, JsonValue, Message,
14    Model, StreamOptions, ThinkingLevel, ToolCall, ToolResultMessage, Usage, UserContent,
15};
16use serde::{Deserialize, Serialize};
17use tokio_util::sync::CancellationToken;
18
19use crate::tool::{AgentTool, AgentToolResult};
20
21/// Error raised by a loop hook; its text becomes the run's `errorMessage`.
22#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
23#[error("{0}")]
24pub struct HookError(pub String);
25
26impl HookError {
27    /// Error with the given message.
28    #[must_use]
29    pub fn new(message: impl Into<String>) -> Self {
30        Self(message.into())
31    }
32}
33
34/// Messages and executable tools of a run.
35#[derive(Clone, Default)]
36pub struct AgentContext {
37    /// Transcript (all roles).
38    pub messages: Vec<AgentMessage>,
39    /// Tools the runtime can execute.
40    pub tools: Vec<Arc<dyn AgentTool>>,
41}
42
43impl std::fmt::Debug for AgentContext {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("AgentContext")
46            .field("messages", &self.messages.len())
47            .field(
48                "tools",
49                &self
50                    .tools
51                    .iter()
52                    .map(|t| t.meta().name.clone())
53                    .collect::<Vec<_>>(),
54            )
55            .finish()
56    }
57}
58
59impl AgentContext {
60    /// The tool with `name`, if executable.
61    #[must_use]
62    pub fn tool(&self, name: &str) -> Option<&Arc<dyn AgentTool>> {
63        self.tools.iter().find(|tool| tool.meta().name == name)
64    }
65}
66
67/// One provider request.
68#[derive(Clone, Debug)]
69pub struct StreamRequest {
70    /// Model to call.
71    pub model: Model,
72    /// Normalised transcript.
73    pub context: Context,
74    /// Stream options with `reasoning` set from the thinking level.
75    pub options: StreamOptions,
76    /// Credential resolved for this request.
77    pub api_key: Option<String>,
78    /// Abort signal.
79    pub cancel: CancellationToken,
80}
81
82/// Provider dispatch. The stream must end with `Done` or `Error`.
83pub trait StreamFn: Send + Sync {
84    /// Starts one request.
85    fn stream(&self, request: StreamRequest) -> BoxStream<'static, AssistantMessageEvent>;
86}
87
88/// `AgentMessage[]` → `Message[]` at the request boundary.
89pub type ConvertToLlm = Arc<dyn Fn(&[AgentMessage]) -> Vec<Message> + Send + Sync>;
90
91/// Rewrites the transcript before conversion.
92pub trait TransformContext: Send + Sync {
93    /// Returns the messages to convert.
94    fn transform<'a>(
95        &'a self,
96        messages: Vec<AgentMessage>,
97        cancel: &'a CancellationToken,
98    ) -> BoxFuture<'a, Result<Vec<AgentMessage>, HookError>>;
99}
100
101/// Resolves a short-lived credential per request.
102pub trait GetApiKey: Send + Sync {
103    /// Credential for `provider`, if any.
104    fn api_key<'a>(&'a self, provider: &'a str)
105    -> BoxFuture<'a, Result<Option<String>, HookError>>;
106}
107
108/// The before tool call context record.
109#[derive(Clone, Copy, Debug)]
110pub struct BeforeToolCallContext<'a> {
111    /// Assistant message that issued the call.
112    pub assistant_message: &'a AssistantMessage,
113    /// The call as issued.
114    pub tool_call: &'a ToolCall,
115    /// Validated arguments.
116    pub args: &'a JsonObject,
117    /// Current run context.
118    pub context: &'a AgentContext,
119}
120
121/// The before tool call result record.
122#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
123#[serde(rename_all = "camelCase")]
124pub struct BeforeToolCallResult {
125    /// Block the call.
126    #[serde(default)]
127    pub block: bool,
128    /// Text of the error result when blocked.
129    #[serde(default, skip_serializing_if = "Option::is_none")]
130    pub reason: Option<String>,
131    /// Mark the blocked result as terminating.
132    #[serde(default)]
133    pub terminate: bool,
134    /// Replacement arguments for the call (handlers mutate `args` in place).
135    #[serde(default, skip_serializing_if = "Option::is_none")]
136    pub input: Option<JsonObject>,
137}
138
139/// B e f o r e   T o o l   C a l l.
140pub trait BeforeToolCall: Send + Sync {
141    /// Inspects (and may block) a validated call.
142    fn before<'a>(
143        &'a self,
144        cx: BeforeToolCallContext<'a>,
145        cancel: &'a CancellationToken,
146    ) -> BoxFuture<'a, Result<Option<BeforeToolCallResult>, HookError>>;
147}
148
149/// The after tool call context record.
150#[derive(Clone, Copy, Debug)]
151pub struct AfterToolCallContext<'a> {
152    /// Assistant message that issued the call.
153    pub assistant_message: &'a AssistantMessage,
154    /// The call as issued.
155    pub tool_call: &'a ToolCall,
156    /// Validated arguments.
157    pub args: &'a JsonObject,
158    /// Result produced by the tool.
159    pub result: &'a AgentToolResult,
160    /// Error flag of the result.
161    pub is_error: bool,
162    /// Current run context.
163    pub context: &'a AgentContext,
164}
165
166/// Field-wise overrides.
167#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
168#[serde(rename_all = "camelCase")]
169pub struct AfterToolCallResult {
170    /// Replacement content (drops `structuredContent` unless also replaced).
171    #[serde(default, skip_serializing_if = "Option::is_none")]
172    pub content: Option<Vec<UserContent>>,
173    /// Replacement details.
174    #[serde(default, skip_serializing_if = "Option::is_none")]
175    pub details: Option<JsonValue>,
176    /// Replacement structured content.
177    #[serde(default, skip_serializing_if = "Option::is_none")]
178    pub structured_content: Option<JsonValue>,
179    /// Replacement error flag.
180    #[serde(default, skip_serializing_if = "Option::is_none")]
181    pub is_error: Option<bool>,
182    /// Replacement usage.
183    #[serde(default, skip_serializing_if = "Option::is_none")]
184    pub usage: Option<Usage>,
185    /// Replacement terminate flag.
186    #[serde(default, skip_serializing_if = "Option::is_none")]
187    pub terminate: Option<bool>,
188}
189
190/// A f t e r   T o o l   C a l l.
191pub trait AfterToolCall: Send + Sync {
192    /// Post-processes a result.
193    fn after<'a>(
194        &'a self,
195        cx: AfterToolCallContext<'a>,
196        cancel: &'a CancellationToken,
197    ) -> BoxFuture<'a, Result<Option<AfterToolCallResult>, HookError>>;
198}
199
200/// The turn context passed to hooks.
201#[derive(Clone, Copy, Debug)]
202pub struct AgentTurnContext<'a> {
203    /// Assistant message of the turn.
204    pub message: &'a AssistantMessage,
205    /// Tool results of the turn.
206    pub tool_results: &'a [ToolResultMessage],
207    /// Context after the turn.
208    pub context: &'a AgentContext,
209    /// Messages added by the run so far.
210    pub new_messages: &'a [AgentMessage],
211}
212
213/// The agent turn decision record.
214#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
215#[serde(tag = "action", rename_all = "lowercase")]
216pub enum TurnDecision {
217    /// Request one more (context-only) turn unless work is already scheduled.
218    Continue,
219    /// End the run after this turn.
220    End,
221}
222
223/// Decides after each turn, before `turn_end`.
224pub trait FinishTurn: Send + Sync {
225    /// Decision for this turn.
226    fn finish<'a>(
227        &'a self,
228        turn: AgentTurnContext<'a>,
229        cancel: &'a CancellationToken,
230    ) -> BoxFuture<'a, Result<Option<TurnDecision>, HookError>>;
231}
232
233/// The agent loop turn update record.
234#[derive(Clone, Debug, Default)]
235pub struct AgentLoopTurnUpdate {
236    /// Replacement context.
237    pub context: Option<AgentContext>,
238    /// Messages to add before the next request.
239    pub messages: Option<Vec<AgentMessage>>,
240    /// Replacement model.
241    pub model: Option<Model>,
242    /// Replacement thinking level.
243    pub thinking_level: Option<ThinkingLevel>,
244}
245
246/// Runs after `turn_end` when the run continues (e.g. compaction).
247pub trait PrepareNextTurn: Send + Sync {
248    /// Optional update for the next turn.
249    fn prepare<'a>(
250        &'a self,
251        turn: AgentTurnContext<'a>,
252        cancel: &'a CancellationToken,
253    ) -> BoxFuture<'a, Result<Option<AgentLoopTurnUpdate>, HookError>>;
254}
255
256/// The prepare request context record.
257#[derive(Clone, Copy, Debug)]
258pub struct PrepareRequestContext<'a> {
259    /// Context about to be sent.
260    pub context: &'a AgentContext,
261    /// Model about to be called.
262    pub model: &'a Model,
263    /// Thinking level about to be requested.
264    pub thinking_level: ThinkingLevel,
265}
266
267/// The agent request update record.
268#[derive(Clone, Debug, Default)]
269pub struct AgentRequestUpdate {
270    /// Replacement context.
271    pub context: Option<AgentContext>,
272    /// Replacement model.
273    pub model: Option<Model>,
274    /// Replacement thinking level.
275    pub thinking_level: Option<ThinkingLevel>,
276}
277
278/// Runs before every request, including the first (virtual routing).
279pub trait PrepareRequest: Send + Sync {
280    /// Optional update for this request.
281    fn prepare<'a>(
282        &'a self,
283        request: PrepareRequestContext<'a>,
284        cancel: &'a CancellationToken,
285    ) -> BoxFuture<'a, Result<Option<AgentRequestUpdate>, HookError>>;
286}
287
288/// Reads queued steering / follow-up messages.
289pub trait QueueSource: Send + Sync {
290    /// Drains the next messages (possibly none).
291    fn drain(&self) -> BoxFuture<'_, Vec<AgentMessage>>;
292}
293
294/// The loop's tool execution mode.
295#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
296#[serde(rename_all = "lowercase")]
297pub enum ToolExecution {
298    /// Prepare sequentially, execute concurrently.
299    #[default]
300    Parallel,
301    /// One call after another.
302    Sequential,
303}
304
305/// Configuration of one loop run.
306///
307/// Provider-level callbacks (`onPayload`, `onResponse`, `onProviderStreamEvent`) belong to the
308/// [`StreamFn`] implementation in Oino, which owns the provider hooks.
309#[derive(Clone)]
310pub struct AgentLoopConfig {
311    /// Model of the first request.
312    pub model: Model,
313    /// Thinking level of the first request (`Off` sends no `reasoning`).
314    pub thinking_level: ThinkingLevel,
315    /// Base stream options (`reasoning` is overwritten per request).
316    pub stream_options: StreamOptions,
317    /// Provider dispatch.
318    pub stream: Arc<dyn StreamFn>,
319    /// Request-boundary conversion.
320    pub convert_to_llm: ConvertToLlm,
321    /// Context rewrite before conversion.
322    pub transform_context: Option<Arc<dyn TransformContext>>,
323    /// Per-request credential.
324    pub get_api_key: Option<Arc<dyn GetApiKey>>,
325    /// Static credential fallback.
326    pub api_key: Option<String>,
327    /// Tool gate.
328    pub before_tool_call: Option<Arc<dyn BeforeToolCall>>,
329    /// Tool result post-processing.
330    pub after_tool_call: Option<Arc<dyn AfterToolCall>>,
331    /// Turn decision.
332    pub finish_turn: Option<Arc<dyn FinishTurn>>,
333    /// Per-request update.
334    pub prepare_request: Option<Arc<dyn PrepareRequest>>,
335    /// Between-turn update.
336    pub prepare_next_turn: Option<Arc<dyn PrepareNextTurn>>,
337    /// Steering queue.
338    pub get_steering_messages: Option<Arc<dyn QueueSource>>,
339    /// Follow-up queue.
340    pub get_follow_up_messages: Option<Arc<dyn QueueSource>>,
341    /// Batch scheduling.
342    pub tool_execution: ToolExecution,
343}
344
345impl AgentLoopConfig {
346    /// A config with no hooks and the default conversion (system/user/assistant/toolResult).
347    #[must_use]
348    pub fn new(model: Model, stream: Arc<dyn StreamFn>) -> Self {
349        Self {
350            model,
351            thinking_level: ThinkingLevel::Off,
352            stream_options: StreamOptions::default(),
353            stream,
354            convert_to_llm: Arc::new(default_convert_to_llm),
355            transform_context: None,
356            get_api_key: None,
357            api_key: None,
358            before_tool_call: None,
359            after_tool_call: None,
360            finish_turn: None,
361            prepare_request: None,
362            prepare_next_turn: None,
363            get_steering_messages: None,
364            get_follow_up_messages: None,
365            tool_execution: ToolExecution::Parallel,
366        }
367    }
368}
369
370/// The default conversion: keeps the four LLM roles.
371#[must_use]
372pub fn default_convert_to_llm(messages: &[AgentMessage]) -> Vec<Message> {
373    messages.iter().filter_map(AgentMessage::as_llm).collect()
374}