1use std::collections::HashMap;
20
21use futures_util::future::BoxFuture;
22
23use crate::protocol::{
24 CallHookRequest, CallHookResponse, EmptyResult, PolicyDecisionRequest, PolicyDecisionResponse,
25 PolicyEvaluationOutcome, StepUpdate, ToolCall, ToolResponse, UserQuestionsRequest,
26 UserQuestionsResponse,
27};
28
29type ToolFn = Box<dyn Fn(ToolCall) -> BoxFuture<'static, ToolResponse> + Send + Sync>;
30type HookFn = Box<dyn Fn(CallHookRequest) -> BoxFuture<'static, CallHookResponse> + Send + Sync>;
31type PolicyFn =
32 Box<dyn Fn(PolicyDecisionRequest) -> BoxFuture<'static, PolicyDecisionResponse> + Send + Sync>;
33type QuestionFn =
34 Box<dyn Fn(UserQuestionsRequest) -> BoxFuture<'static, UserQuestionsResponse> + Send + Sync>;
35type ConfirmFn = Box<dyn Fn(StepUpdate) -> BoxFuture<'static, bool> + Send + Sync>;
36
37#[derive(Default)]
39pub struct Handlers {
40 tools: HashMap<String, ToolFn>,
41 hook: Option<HookFn>,
42 policy: Option<PolicyFn>,
43 questions: Option<QuestionFn>,
44 confirm: Option<ConfirmFn>,
45}
46
47impl std::fmt::Debug for Handlers {
48 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49 f.debug_struct("Handlers")
50 .field("tools", &self.tools.keys().collect::<Vec<_>>())
51 .field("hook", &self.hook.is_some())
52 .field("policy", &self.policy.is_some())
53 .field("questions", &self.questions.is_some())
54 .field("confirm", &self.confirm.is_some())
55 .finish()
56 }
57}
58
59impl Handlers {
60 pub fn new() -> Self {
62 Self::default()
63 }
64
65 pub fn tool<F, Fut>(mut self, name: impl Into<String>, f: F) -> Self
81 where
82 F: Fn(ToolCall) -> Fut + Send + Sync + 'static,
83 Fut: std::future::Future<Output = ToolResponse> + Send + 'static,
84 {
85 self.tools
86 .insert(name.into(), Box::new(move |call| Box::pin(f(call))));
87 self
88 }
89
90 pub fn on_hook<F, Fut>(mut self, f: F) -> Self
92 where
93 F: Fn(CallHookRequest) -> Fut + Send + Sync + 'static,
94 Fut: std::future::Future<Output = CallHookResponse> + Send + 'static,
95 {
96 self.hook = Some(Box::new(move |request| Box::pin(f(request))));
97 self
98 }
99
100 pub fn on_policy<F, Fut>(mut self, f: F) -> Self
102 where
103 F: Fn(PolicyDecisionRequest) -> Fut + Send + Sync + 'static,
104 Fut: std::future::Future<Output = PolicyDecisionResponse> + Send + 'static,
105 {
106 self.policy = Some(Box::new(move |request| Box::pin(f(request))));
107 self
108 }
109
110 pub fn on_questions<F, Fut>(mut self, f: F) -> Self
112 where
113 F: Fn(UserQuestionsRequest) -> Fut + Send + Sync + 'static,
114 Fut: std::future::Future<Output = UserQuestionsResponse> + Send + 'static,
115 {
116 self.questions = Some(Box::new(move |request| Box::pin(f(request))));
117 self
118 }
119
120 pub fn on_tool_confirmation<F, Fut>(mut self, f: F) -> Self
126 where
127 F: Fn(StepUpdate) -> Fut + Send + Sync + 'static,
128 Fut: std::future::Future<Output = bool> + Send + 'static,
129 {
130 self.confirm = Some(Box::new(move |step| Box::pin(f(step))));
131 self
132 }
133
134 pub fn handles_tool(&self, name: &str) -> bool {
136 self.tools.contains_key(name)
137 }
138
139 pub(crate) async fn call_tool(&self, call: ToolCall) -> ToolResponse {
140 let id = call.id.clone().unwrap_or_default();
141 let name = call.name.clone().unwrap_or_default();
142 match self.tools.get(&name) {
143 Some(f) => f(call).await,
144 None => ToolResponse::error(id, format!("no handler registered for tool `{name}`")),
149 }
150 }
151
152 pub(crate) async fn call_hook(&self, request: CallHookRequest) -> CallHookResponse {
153 let request_id = request.request_id.clone();
154 match &self.hook {
155 Some(f) => f(request).await,
156 None => CallHookResponse {
157 request_id,
158 empty_result: Some(EmptyResult {}),
159 ..Default::default()
160 },
161 }
162 }
163
164 pub(crate) async fn call_policy(
165 &self,
166 request: PolicyDecisionRequest,
167 ) -> PolicyDecisionResponse {
168 let request_id = request.request_id.clone();
169 match &self.policy {
170 Some(f) => f(request).await,
171 None => PolicyDecisionResponse {
174 request_id,
175 outcome: Some(PolicyEvaluationOutcome::NoMatch),
176 ..Default::default()
177 },
178 }
179 }
180
181 pub(crate) async fn call_questions(
182 &self,
183 request: UserQuestionsRequest,
184 ) -> UserQuestionsResponse {
185 match &self.questions {
186 Some(f) => f(request).await,
187 None => UserQuestionsResponse {
190 cancelled: Some(true),
191 ..Default::default()
192 },
193 }
194 }
195
196 pub(crate) async fn call_confirm(&self, step: StepUpdate) -> bool {
197 match &self.confirm {
198 Some(f) => f(step).await,
199 None => false,
203 }
204 }
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210
211 #[tokio::test]
212 async fn an_unregistered_tool_fails_that_call_only() {
213 let handlers = Handlers::new();
214 let response = handlers
215 .call_tool(ToolCall {
216 id: Some("call-1".into()),
217 name: Some("missing".into()),
218 ..Default::default()
219 })
220 .await;
221 assert_eq!(response.id.as_deref(), Some("call-1"));
222 assert!(response
223 .error_message
224 .unwrap()
225 .contains("no handler registered"));
226 }
227
228 #[tokio::test]
229 async fn a_registered_tool_is_dispatched_by_name() {
230 let handlers = Handlers::new().tool("echo", |call| async move {
231 ToolResponse::ok(call.id.unwrap_or_default(), "42")
232 });
233 let response = handlers
234 .call_tool(ToolCall {
235 id: Some("c".into()),
236 name: Some("echo".into()),
237 ..Default::default()
238 })
239 .await;
240 assert_eq!(response.response_json.as_deref(), Some("42"));
241 }
242
243 #[tokio::test]
244 async fn defaults_keep_the_turn_moving() {
245 let handlers = Handlers::new();
246 let hook = handlers
247 .call_hook(CallHookRequest {
248 request_id: Some("r".into()),
249 ..Default::default()
250 })
251 .await;
252 assert_eq!(hook.request_id.as_deref(), Some("r"));
253 assert!(hook.empty_result.is_some());
254
255 let policy = handlers
256 .call_policy(PolicyDecisionRequest {
257 request_id: Some("p".into()),
258 ..Default::default()
259 })
260 .await;
261 assert_eq!(policy.outcome, Some(PolicyEvaluationOutcome::NoMatch));
262
263 let questions = handlers
264 .call_questions(UserQuestionsRequest::default())
265 .await;
266 assert_eq!(questions.cancelled, Some(true));
267 }
268}