Skip to main content

yapi_agent/
hooks.rs

1//! Points where the session steers the loop.
2
3use futures_util::future::BoxFuture;
4use serde_json::Value;
5use yapi_ai::registry::Auth;
6use yapi_types::event::{AgentEvent, ToolResult};
7use yapi_types::message::{AssistantMessage, Message, ToolCall, ToolResultMessage, Usage};
8use yapi_types::model::Model;
9
10/// A tool call about to run.
11#[derive(Debug)]
12pub struct BeforeToolCall<'a> {
13    /// The response that made the call.
14    pub assistant_message: &'a AssistantMessage,
15    /// The call as the model made it.
16    pub tool_call: &'a ToolCall,
17    /// Validated arguments, which the hook may change. The tool runs with
18    /// them as they are when the hook returns.
19    pub args: &'a mut Value,
20    /// The transcript so far.
21    pub messages: &'a [Message],
22    /// The tool call that made this call, for calls a tool made.
23    pub parent_tool_call_id: Option<&'a str>,
24}
25
26/// Blocks a tool call.
27#[derive(Clone, Debug, Default, PartialEq)]
28pub struct Block {
29    /// Shown to the model, or "Tool execution was blocked" when absent.
30    pub reason: Option<String>,
31    /// Ends the run after this batch.
32    pub terminate: bool,
33}
34
35/// A tool call that ran.
36#[derive(Debug)]
37pub struct AfterToolCall<'a> {
38    /// The response that made the call.
39    pub assistant_message: &'a AssistantMessage,
40    /// The call as the model made it.
41    pub tool_call: &'a ToolCall,
42    /// Validated arguments.
43    pub args: &'a Value,
44    /// What the tool returned.
45    pub result: &'a ToolResult,
46    /// Whether the result is an error.
47    pub is_error: bool,
48    /// The tool call that made this call, for calls a tool made.
49    pub parent_tool_call_id: Option<&'a str>,
50}
51
52/// Replacement fields for a tool result. `None` keeps the original.
53#[derive(Clone, Debug, Default, PartialEq)]
54pub struct ResultPatch {
55    /// New content. It drops the structured content, which may no longer
56    /// match, unless `structured_content` replaces that too.
57    pub content: Option<Vec<yapi_types::message::ContentBlock>>,
58    /// New details.
59    pub details: Option<Value>,
60    /// New structured content.
61    pub structured_content: Option<Value>,
62    /// New error flag.
63    pub is_error: Option<bool>,
64    /// New usage.
65    pub usage: Option<Usage>,
66    /// New terminate flag.
67    pub terminate: Option<bool>,
68}
69
70/// What the session decides as a turn ends (pi's `AgentTurnDecision`).
71#[derive(Clone, Debug, Default, PartialEq)]
72pub struct TurnDecision {
73    /// Makes the loop send one more provider request when nothing else would.
74    pub continue_run: bool,
75}
76
77/// Hooks into the agent loop. Every method has a no-op default.
78///
79/// The loop awaits each hook before it continues, so a hook sees a consistent
80/// transcript. Hooks never call back into the loop.
81pub trait AgentHooks: Send + Sync {
82    /// Receives every loop event, in order.
83    fn on_event<'a>(&'a self, _event: &'a AgentEvent) -> BoxFuture<'a, ()> {
84        Box::pin(async {})
85    }
86
87    /// Rewrites the transcript before the loop converts it for the provider.
88    fn transform_context(&self, messages: Vec<Message>) -> BoxFuture<'_, Vec<Message>> {
89        Box::pin(async { messages })
90    }
91
92    /// Converts session messages to the four roles providers accept.
93    fn convert_to_llm(&self, messages: Vec<Message>) -> Vec<Message> {
94        messages
95    }
96
97    /// Credentials for a request to `model`, resolved fresh for each request.
98    fn auth<'a>(&'a self, _model: &'a Model) -> BoxFuture<'a, Auth> {
99        Box::pin(async { Auth::default() })
100    }
101
102    /// May block a validated call.
103    fn before_tool_call<'a>(&'a self, _call: BeforeToolCall<'a>) -> BoxFuture<'a, Option<Block>> {
104        Box::pin(async { None })
105    }
106
107    /// May patch a result.
108    fn after_tool_call<'a>(
109        &'a self,
110        _call: AfterToolCall<'a>,
111    ) -> BoxFuture<'a, Option<ResultPatch>> {
112        Box::pin(async { None })
113    }
114
115    /// A turn finished: the loop emitted its assistant message and tool
116    /// results, and `turn_end` is next (pi's `finishTurn`). The loop ignores
117    /// the decision after a failed or aborted response.
118    fn finish_turn<'a>(
119        &'a self,
120        _message: &'a AssistantMessage,
121        _tool_results: &'a [ToolResultMessage],
122    ) -> BoxFuture<'a, TurnDecision> {
123        Box::pin(async { TurnDecision::default() })
124    }
125
126    /// The message to record in place of `message` once it is complete, before
127    /// the loop emits its `message_end` (pi's `message_end` replacement). Must
128    /// keep the message's role.
129    fn finish_message(&self, message: Message) -> BoxFuture<'_, Message> {
130        Box::pin(async { message })
131    }
132
133    /// Completes a tool result message before it is emitted, for example
134    /// with the calls the tool made.
135    fn complete_tool_result(&self, _message: &mut ToolResultMessage) {}
136
137    /// The transcript the next turn continues from, read before each turn
138    /// after the first (pi's `prepareNextTurn` context). `None` keeps the
139    /// loop's own. An error ends the run with it, as pi's agent ends a run
140    /// whose hook throws.
141    fn next_turn_messages(&self) -> BoxFuture<'_, Result<Option<Vec<Message>>, String>> {
142        Box::pin(async { Ok(None) })
143    }
144
145    /// The tools for the next request, read before each turn. `None` keeps
146    /// the current ones.
147    fn current_tools(&self) -> BoxFuture<'_, Option<Vec<std::sync::Arc<dyn crate::Tool>>>> {
148        Box::pin(async { None })
149    }
150
151    /// Messages the user queued to steer the current run. Reading drains them.
152    fn steering_messages(&self) -> BoxFuture<'_, Vec<Message>> {
153        Box::pin(async { Vec::new() })
154    }
155
156    /// Messages to send once the run would otherwise stop. Reading drains them.
157    fn follow_up_messages(&self) -> BoxFuture<'_, Vec<Message>> {
158        Box::pin(async { Vec::new() })
159    }
160}