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

rpi_agent/
hooks.rs
