Skip to main content

antigravity_codes/
handlers.rs

1//! Answering the requests the harness sends *back* to the client.
2//!
3//! A turn does not just stream output. Depending on how the session was
4//! configured, the harness will stop and wait for the client on four
5//! occasions, and the turn stays blocked until each is answered:
6//!
7//! | Request | Raised when | Answered with |
8//! |---|---|---|
9//! | [`ToolCall`] | the model calls a tool declared via [`HarnessOptions::tool`](crate::HarnessOptions::tool) | [`ToolResponse`] |
10//! | [`CallHookRequest`] | a lifecycle hook registered via [`HarnessOptions::hook`](crate::HarnessOptions::hook) fires | [`CallHookResponse`] |
11//! | [`PolicyDecisionRequest`] | a dynamic [`PolicyRule`](crate::protocol::PolicyRule) needs adjudicating | [`PolicyDecisionResponse`] |
12//! | [`UserQuestionsRequest`] | the agent asks the user something | [`UserQuestionsResponse`] |
13//!
14//! None of these arrive unless the corresponding feature was configured, so
15//! the empty [`Handlers`] is a perfectly good default for a plain chat
16//! session. When one *does* arrive with no handler registered, the defaults
17//! below keep the turn moving rather than deadlocking it.
18
19use 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/// The callbacks a [`Client`](crate::Client) uses to answer the harness.
38#[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    /// No handlers. Fine for a session with no client-side tools or hooks.
61    pub fn new() -> Self {
62        Self::default()
63    }
64
65    /// Registers the implementation of a client-side tool.
66    ///
67    /// `name` must match the [`Tool`](crate::protocol::Tool) declared on
68    /// [`HarnessOptions::tool`](crate::HarnessOptions::tool); the harness
69    /// dispatches by name.
70    ///
71    /// ```
72    /// use antigravity_codes::handlers::Handlers;
73    /// use antigravity_codes::protocol::ToolResponse;
74    ///
75    /// let handlers = Handlers::new().tool("clock", |call| async move {
76    ///     ToolResponse::ok(call.id.unwrap_or_default(), r#"{"now":"2026-08-08T00:00:00Z"}"#)
77    /// });
78    /// assert!(handlers.handles_tool("clock"));
79    /// ```
80    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    /// Handles every lifecycle hook the session subscribed to.
91    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    /// Adjudicates dynamic policy rules.
101    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    /// Answers questions the agent puts to the user.
111    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    /// Approves or refuses a tool the harness wants confirmed before running.
121    ///
122    /// The handler receives the whole [`StepUpdate`], because the action being
123    /// confirmed — the command line, the file path, the diff — is on the step,
124    /// not on the (empty) request itself.
125    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    /// Whether a tool of this name has an implementation registered.
135    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            // Reported to the model as a tool failure rather than raised to the
145            // caller: the harness only asks for tools the session declared, so
146            // this is a client-side wiring bug, and failing the one call lets
147            // the agent recover or explain itself instead of wedging the turn.
148            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            // `NO_MATCH` defers to whatever static rules the harness has, which
172            // is the safe reading of "the client expressed no opinion".
173            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            // Cancelling is the only answer that cannot be wrong: it tells the
188            // agent nobody is there, and unblocks the turn.
189            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            // Refusing by default. The harness only asks when the session was
200            // configured to require confirmation, so silently approving would
201            // undo the very control the caller asked for.
202            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}