rpi_agent/
hooks.rs

1//! Mirrors `packages/agent/src/types.ts::AgentLoopConfig` — the config bundle
2//! the low-level loop consumes. TS `AgentLoopConfig extends SimpleStreamOptions`
3//! and adds `model` + a set of async hooks. Rust models it as a struct of
4//! `Option<Arc<dyn Fn>>>` hooks (required `convert_to_llm` is non-optional),
5//! the `model`, and a `SimpleStreamOptions`-derived subset the loop forwards
6//! to `StreamFn`.
7//!
8//! Hook signatures use `Arc<dyn Fn(...) -> BoxFuture<'static, T> + Send + Sync>`
9//! so an `Agent` can store closures that capture `Arc<Agent>` / shared state
10//! without lifetime gymnastics. Every hook mirrors the TS contract: must not
11//! throw — return a safe fallback instead.
12
13use crate::message::AgentMessage;
14use crate::types::{
15    AfterToolCallContext, AfterToolCallResult, AgentLoopTurnUpdate, BeforeToolCallContext,
16    BeforeToolCallResult, ShouldStopAfterTurnContext, ToolExecutionMode,
17};
18use futures::future::BoxFuture;
19use rpi_ai::provider::{CacheRetention, SimpleStreamOptions};
20use rpi_ai::types::{Message, ThinkingLevel};
21use rpi_ai::Model;
22use std::sync::Arc;
23use std::time::Duration;
24use tokio_util::sync::CancellationToken;
25
26/// `(messages: AgentMessage[]) -> Message[]` — the required
27/// LLM-boundary converter. Filters/transforms custom messages into the
28/// provider-facing `Message` union. Must not panic.
29pub type ConvertToLlm =
30    Arc<dyn Fn(Vec<AgentMessage>) -> BoxFuture<'static, Vec<Message>> + Send + Sync>;
31
32/// `(messages, signal) -> AgentMessage[]` — optional pre-convert transform at
33/// the `AgentMessage` level (context-window pruning, external injection).
34pub type TransformContext = Arc<
35    dyn Fn(Vec<AgentMessage>, CancellationToken) -> BoxFuture<'static, Vec<AgentMessage>>
36        + Send
37        + Sync,
38>;
39
40/// `(provider: &str) -> Option<String>` — dynamic API-key resolver per turn.
41pub type GetApiKey = Arc<dyn Fn(&str) -> BoxFuture<'static, Option<String>> + Send + Sync>;
42
43/// `(context) -> bool` — return true to stop after the current turn (before
44/// steering/follow-up drain).
45pub type ShouldStopAfterTurn =
46    Arc<dyn Fn(ShouldStopAfterTurnContext<'_>) -> BoxFuture<'static, bool> + Send + Sync>;
47
48/// `(context) -> Option<AgentLoopTurnUpdate>` — replacement context/model/thinking
49/// for the next turn, if any.
50pub type PrepareNextTurn = Arc<
51    dyn Fn(ShouldStopAfterTurnContext<'_>) -> BoxFuture<'static, Option<AgentLoopTurnUpdate>>
52        + Send
53        + Sync,
54>;
55
56/// Hook invoked after a tool-result batch is appended and before the next
57/// assistant request. Implementations may replace the context (for example by
58/// inserting a compaction boundary).
59pub type AfterToolResults = Arc<
60    dyn Fn(ShouldStopAfterTurnContext<'_>) -> BoxFuture<'static, Option<AgentLoopTurnUpdate>>
61        + Send
62        + Sync,
63>;
64
65/// `() -> Vec<AgentMessage>` — messages to inject mid-run after a tool batch.
66pub type GetSteeringMessages = Arc<dyn Fn() -> BoxFuture<'static, Vec<AgentMessage>> + Send + Sync>;
67
68/// `() -> Vec<AgentMessage>` — messages to inject after the loop would stop.
69pub type GetFollowUpMessages = Arc<dyn Fn() -> BoxFuture<'static, Vec<AgentMessage>> + Send + Sync>;
70
71/// `(context, signal) -> Option<BeforeToolCallResult>` — can block a tool call
72/// before it runs.
73pub type BeforeToolCall = Arc<
74    dyn Fn(
75            BeforeToolCallContext<'_>,
76            CancellationToken,
77        ) -> BoxFuture<'static, Option<BeforeToolCallResult>>
78        + Send
79        + Sync,
80>;
81
82/// `(context, signal) -> Option<AfterToolCallResult>` — overrides fields of an
83/// executed tool result before `tool_execution_end` / `MessageEnd`.
84pub type AfterToolCall = Arc<
85    dyn Fn(
86            AfterToolCallContext<'_>,
87            CancellationToken,
88        ) -> BoxFuture<'static, Option<AfterToolCallResult>>
89        + Send
90        + Sync,
91>;
92
93/// The config bundle handed to `run_agent_loop` / `run_agent_loop_continue`.
94/// Mirrors TS `AgentLoopConfig`. Hook fields are `Option` except `convert_to_llm`
95/// (required). The provider-options subset (`api_key`, `timeout`, `cache_retention`,
96/// `session_id`, `max_retries`, `max_retry_delay`, `signal`) is kept as plain
97/// fields the loop stuffs into a `SimpleStreamOptions` per turn.
98#[derive(Clone)]
99pub struct AgentLoopConfig {
100    pub model: Model,
101
102    /// Required: `AgentMessage[]` → `Message[]`. Never `None` at call time; the
103    /// `AgentBuilder` installs `default_convert_to_llm` (drop custom roles) when
104    /// the caller doesn't supply one.
105    pub convert_to_llm: ConvertToLlm,
106
107    pub transform_context: Option<TransformContext>,
108    pub get_api_key: Option<GetApiKey>,
109    pub should_stop_after_turn: Option<ShouldStopAfterTurn>,
110    pub prepare_next_turn: Option<PrepareNextTurn>,
111    pub after_tool_results: Option<AfterToolResults>,
112    pub get_steering_messages: Option<GetSteeringMessages>,
113    pub get_follow_up_messages: Option<GetFollowUpMessages>,
114    pub before_tool_call: Option<BeforeToolCall>,
115    pub after_tool_call: Option<AfterToolCall>,
116
117    /// Per-batch execution mode. Default `Parallel`.
118    pub tool_execution: ToolExecutionMode,
119
120    // ---- provider-options subset (forwarded into SimpleStreamOptions) ----
121    pub thinking_level: ThinkingLevel,
122    pub api_key: Option<String>,
123    pub timeout: Option<Duration>,
124    pub max_retries: Option<u32>,
125    pub max_retry_delay: Option<Duration>,
126    pub cache_retention: CacheRetention,
127    pub session_id: Option<String>,
128    /// Cancellation token for the whole run. The loop clones a child per tool.
129    pub signal: CancellationToken,
130}
131
132impl std::fmt::Debug for AgentLoopConfig {
133    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
134        f.debug_struct("AgentLoopConfig")
135            .field("model", &self.model)
136            .field("tool_execution", &self.tool_execution)
137            .field("thinking_level", &self.thinking_level)
138            .field("cache_retention", &self.cache_retention)
139            .field("session_id", &self.session_id)
140            .field("transform_context", &self.transform_context.is_some())
141            .field("get_api_key", &self.get_api_key.is_some())
142            .field(
143                "should_stop_after_turn",
144                &self.should_stop_after_turn.is_some(),
145            )
146            .field("prepare_next_turn", &self.prepare_next_turn.is_some())
147            .field("after_tool_results", &self.after_tool_results.is_some())
148            .field(
149                "get_steering_messages",
150                &self.get_steering_messages.is_some(),
151            )
152            .field(
153                "get_follow_up_messages",
154                &self.get_follow_up_messages.is_some(),
155            )
156            .field("before_tool_call", &self.before_tool_call.is_some())
157            .field("after_tool_call", &self.after_tool_call.is_some())
158            .finish()
159    }
160}
161
162/// Build a `SimpleStreamOptions` from this config's provider-options subset,
163/// overriding `signal` with the run's token. Mirrors the TS spread
164/// `{ ...config, apiKey: resolvedApiKey, signal }` passed to `streamFunction`.
165impl AgentLoopConfig {
166    pub fn to_stream_options(&self, api_key: Option<String>) -> SimpleStreamOptions {
167        let mut opts = SimpleStreamOptions {
168            api_key,
169            timeout: self.timeout,
170            max_retries: self.max_retries,
171            max_retry_delay: self.max_retry_delay,
172            headers: None,
173            metadata: None,
174            cache_retention: self.cache_retention,
175            session_id: self.session_id.clone(),
176            signal: self.signal.clone(),
177            ..SimpleStreamOptions::default()
178        };
179        // Forward the run's thinking level as `reasoning` (the provider-level
180        // knob) so a real provider like anthropic can map it to adaptive vs
181        // budget-based thinking. TS `AgentLoopConfig extends SimpleStreamOptions`
182        // and so carries `reasoning` through directly.
183        match self.thinking_level {
184            ThinkingLevel::Off => opts.reasoning = None,
185            other => opts.reasoning = Some(other),
186        }
187        opts
188    }
189}
190
191/// The default `convert_to_llm`: keep only `user`/`assistant`/`toolResult`
192/// messages, drop `custom`. Mirrors TS `defaultConvertToLlm`.
193pub fn default_convert_to_llm(messages: Vec<AgentMessage>) -> Vec<Message> {
194    messages
195        .into_iter()
196        .filter_map(|m| match m {
197            AgentMessage::User(u) => Some(Message::User(u)),
198            AgentMessage::Assistant(a) => Some(Message::Assistant(a)),
199            AgentMessage::ToolResult(t) => Some(Message::ToolResult(t)),
200            AgentMessage::Custom(_) => None,
201        })
202        .collect()
203}
204
205/// Wrap a `default_convert_to_llm` into the `ConvertToLlm` Arc shape expected
206/// by `AgentLoopConfig`.
207pub fn default_convert_to_llm_fn() -> ConvertToLlm {
208    Arc::new(|messages: Vec<AgentMessage>| {
209        let out = default_convert_to_llm(messages);
210        Box::pin(async move { out })
211    })
212}
rpi_agent/hooks.rs

rpi_agent/
hooks.rs
