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}