agent_client_protocol/jsonrpc/
handlers.rs

1use crate::jsonrpc::{HandleDispatchFrom, Handled, IntoHandled, JsonRpcResponse};
2
3use crate::role::{HasPeer, Role, handle_incoming_dispatch};
4use crate::{ConnectionTo, Dispatch, JsonRpcNotification, JsonRpcRequest, UntypedMessage};
5// Types re-exported from crate root
6use super::Responder;
7use std::marker::PhantomData;
8use std::ops::AsyncFnMut;
9
10/// Null handler that accepts no messages.
11#[derive(Debug)]
12pub struct NullHandler;
13
14impl Default for NullHandler {
15    fn default() -> Self {
16        Self
17    }
18}
19
20impl<Counterpart: Role> HandleDispatchFrom<Counterpart> for NullHandler {
21    fn describe_chain(&self) -> impl std::fmt::Debug {
22        "(null)"
23    }
24
25    async fn handle_dispatch_from(
26        &mut self,
27        message: Dispatch,
28        _cx: ConnectionTo<Counterpart>,
29    ) -> Result<Handled<Dispatch>, crate::Error> {
30        Ok(Handled::No {
31            message,
32            retry: false,
33        })
34    }
35}
36
37/// Handler for typed request messages
38pub struct RequestHandler<
39    Counterpart: Role,
40    Peer: Role,
41    Req: JsonRpcRequest = UntypedMessage,
42    F = (),
43    ToFut = (),
44> {
45    counterpart: Counterpart,
46    peer: Peer,
47    handler: F,
48    to_future_hack: ToFut,
49    phantom: PhantomData<fn(Req)>,
50}
51
52impl<Counterpart: Role, Peer: Role, Req: JsonRpcRequest, F, ToFut>
53    RequestHandler<Counterpart, Peer, Req, F, ToFut>
54{
55    /// Creates a new request handler
56    pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self {
57        Self {
58            counterpart,
59            peer,
60            handler,
61            to_future_hack,
62            phantom: PhantomData,
63        }
64    }
65}
66
67impl<Counterpart: Role, Peer: Role, Req, F, T, ToFut> HandleDispatchFrom<Counterpart>
68    for RequestHandler<Counterpart, Peer, Req, F, ToFut>
69where
70    Counterpart: HasPeer<Peer>,
71    Req: JsonRpcRequest,
72    F: AsyncFnMut(
73            Req,
74            Responder<Req::Response>,
75            ConnectionTo<Counterpart>,
76        ) -> Result<T, crate::Error>
77        + Send,
78    T: crate::IntoHandled<(Req, Responder<Req::Response>)>,
79    ToFut: Fn(
80            &mut F,
81            Req,
82            Responder<Req::Response>,
83            ConnectionTo<Counterpart>,
84        ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
85        + Send
86        + Sync,
87{
88    fn describe_chain(&self) -> impl std::fmt::Debug {
89        std::any::type_name::<Req>()
90    }
91
92    async fn handle_dispatch_from(
93        &mut self,
94        dispatch: Dispatch,
95        connection: ConnectionTo<Counterpart>,
96    ) -> Result<Handled<Dispatch>, crate::Error> {
97        handle_incoming_dispatch(
98            self.counterpart.clone(),
99            self.peer.clone(),
100            dispatch,
101            connection,
102            async |dispatch, connection| {
103                match dispatch {
104                    Dispatch::Request(message, responder) => {
105                        tracing::debug!(
106                            request_type = std::any::type_name::<Req>(),
107                            message = ?message,
108                            "RequestHandler::handle_request"
109                        );
110                        if Req::matches_method(&message.method) {
111                            match Req::parse_message(&message.method, &message.params) {
112                                Ok(req) => {
113                                    tracing::trace!(
114                                        ?req,
115                                        "RequestHandler::handle_request: parse completed"
116                                    );
117                                    let typed_responder = responder.cast();
118                                    let result = (self.to_future_hack)(
119                                        &mut self.handler,
120                                        req,
121                                        typed_responder,
122                                        connection,
123                                    )
124                                    .await?;
125                                    match result.into_handled() {
126                                        Handled::Yes => Ok(Handled::Yes),
127                                        Handled::No {
128                                            message: (request, responder),
129                                            retry,
130                                        } => {
131                                            // Handler returned the request back, convert to untyped
132                                            let untyped = request.to_untyped_message()?;
133                                            Ok(Handled::No {
134                                                message: Dispatch::Request(
135                                                    untyped,
136                                                    responder.erase_to_json(),
137                                                ),
138                                                retry,
139                                            })
140                                        }
141                                    }
142                                }
143                                Err(err) => {
144                                    tracing::trace!(
145                                        ?err,
146                                        "RequestHandler::handle_request: parse errored"
147                                    );
148                                    Err(err)
149                                }
150                            }
151                        } else {
152                            tracing::trace!("RequestHandler::handle_request: method doesn't match");
153                            Ok(Handled::No {
154                                message: Dispatch::Request(message, responder),
155                                retry: false,
156                            })
157                        }
158                    }
159
160                    Dispatch::Notification(..) | Dispatch::Response(..) => Ok(Handled::No {
161                        message: dispatch,
162                        retry: false,
163                    }),
164                }
165            },
166        )
167        .await
168    }
169}
170
171/// Handler for typed notification messages
172pub struct NotificationHandler<
173    Counterpart: Role,
174    Peer: Role,
175    Notif: JsonRpcNotification = UntypedMessage,
176    F = (),
177    ToFut = (),
178> {
179    counterpart: Counterpart,
180    peer: Peer,
181    handler: F,
182    to_future_hack: ToFut,
183    phantom: PhantomData<fn(Notif)>,
184}
185
186impl<Counterpart: Role, Peer: Role, Notif: JsonRpcNotification, F, ToFut>
187    NotificationHandler<Counterpart, Peer, Notif, F, ToFut>
188{
189    /// Creates a new notification handler
190    pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self {
191        Self {
192            counterpart,
193            peer,
194            handler,
195            to_future_hack,
196            phantom: PhantomData,
197        }
198    }
199}
200
201impl<Counterpart: Role, Peer: Role, Notif, F, T, ToFut> HandleDispatchFrom<Counterpart>
202    for NotificationHandler<Counterpart, Peer, Notif, F, ToFut>
203where
204    Counterpart: HasPeer<Peer>,
205    Notif: JsonRpcNotification,
206    F: AsyncFnMut(Notif, ConnectionTo<Counterpart>) -> Result<T, crate::Error> + Send,
207    T: crate::IntoHandled<(Notif, ConnectionTo<Counterpart>)>,
208    ToFut: Fn(
209            &mut F,
210            Notif,
211            ConnectionTo<Counterpart>,
212        ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
213        + Send
214        + Sync,
215{
216    fn describe_chain(&self) -> impl std::fmt::Debug {
217        std::any::type_name::<Notif>()
218    }
219
220    async fn handle_dispatch_from(
221        &mut self,
222        dispatch: Dispatch,
223        connection: ConnectionTo<Counterpart>,
224    ) -> Result<Handled<Dispatch>, crate::Error> {
225        handle_incoming_dispatch(
226            self.counterpart.clone(),
227            self.peer.clone(),
228            dispatch,
229            connection,
230            async |dispatch, connection| {
231                match dispatch {
232                    Dispatch::Notification(message) => {
233                        tracing::debug!(
234                            request_type = std::any::type_name::<Notif>(),
235                            message = ?message,
236                            "NotificationHandler::handle_dispatch"
237                        );
238                        if Notif::matches_method(&message.method) {
239                            match Notif::parse_message(&message.method, &message.params) {
240                                Ok(notif) => {
241                                    tracing::trace!(
242                                        ?notif,
243                                        "NotificationHandler::handle_notification: parse completed"
244                                    );
245                                    let result =
246                                        (self.to_future_hack)(&mut self.handler, notif, connection)
247                                            .await?;
248                                    match result.into_handled() {
249                                        Handled::Yes => Ok(Handled::Yes),
250                                        Handled::No {
251                                            message: (notification, _cx),
252                                            retry,
253                                        } => {
254                                            // Handler returned the notification back, convert to untyped
255                                            let untyped = notification.to_untyped_message()?;
256                                            Ok(Handled::No {
257                                                message: Dispatch::Notification(untyped),
258                                                retry,
259                                            })
260                                        }
261                                    }
262                                }
263                                Err(err) => {
264                                    tracing::trace!(
265                                        ?err,
266                                        "NotificationHandler::handle_notification: parse errored"
267                                    );
268                                    Err(err)
269                                }
270                            }
271                        } else {
272                            tracing::trace!(
273                                "NotificationHandler::handle_notification: method doesn't match"
274                            );
275                            Ok(Handled::No {
276                                message: Dispatch::Notification(message),
277                                retry: false,
278                            })
279                        }
280                    }
281
282                    Dispatch::Request(..) | Dispatch::Response(..) => Ok(Handled::No {
283                        message: dispatch,
284                        retry: false,
285                    }),
286                }
287            },
288        )
289        .await
290    }
291}
292
293/// Handler for typed requests, notifications, and matching responses.
294pub struct MessageHandler<
295    Counterpart: Role,
296    Peer: Role,
297    Req: JsonRpcRequest = UntypedMessage,
298    Notif: JsonRpcNotification = UntypedMessage,
299    F = (),
300    ToFut = (),
301> {
302    counterpart: Counterpart,
303    peer: Peer,
304    handler: F,
305    to_future_hack: ToFut,
306    phantom: PhantomData<fn(Dispatch<Req, Notif>)>,
307}
308
309impl<Counterpart: Role, Peer: Role, Req: JsonRpcRequest, Notif: JsonRpcNotification, F, ToFut>
310    MessageHandler<Counterpart, Peer, Req, Notif, F, ToFut>
311{
312    /// Creates a new message handler
313    pub fn new(counterpart: Counterpart, peer: Peer, handler: F, to_future_hack: ToFut) -> Self {
314        Self {
315            counterpart,
316            peer,
317            handler,
318            to_future_hack,
319            phantom: PhantomData,
320        }
321    }
322}
323
324impl<Counterpart: Role, Peer: Role, Req: JsonRpcRequest, Notif: JsonRpcNotification, F, T, ToFut>
325    HandleDispatchFrom<Counterpart> for MessageHandler<Counterpart, Peer, Req, Notif, F, ToFut>
326where
327    Counterpart: HasPeer<Peer>,
328    F: AsyncFnMut(Dispatch<Req, Notif>, ConnectionTo<Counterpart>) -> Result<T, crate::Error>
329        + Send,
330    T: IntoHandled<Dispatch<Req, Notif>>,
331    ToFut: Fn(
332            &mut F,
333            Dispatch<Req, Notif>,
334            ConnectionTo<Counterpart>,
335        ) -> crate::BoxFuture<'_, Result<T, crate::Error>>
336        + Send
337        + Sync,
338{
339    fn describe_chain(&self) -> impl std::fmt::Debug {
340        format!(
341            "({}, {})",
342            std::any::type_name::<Req>(),
343            std::any::type_name::<Notif>()
344        )
345    }
346
347    async fn handle_dispatch_from(
348        &mut self,
349        dispatch: Dispatch,
350        connection: ConnectionTo<Counterpart>,
351    ) -> Result<Handled<Dispatch>, crate::Error> {
352        handle_incoming_dispatch(
353            self.counterpart.clone(),
354            self.peer.clone(),
355            dispatch,
356            connection,
357            async |dispatch, connection| match dispatch.into_typed_dispatch::<Req, Notif>()? {
358                Ok(typed_dispatch) => {
359                    let result =
360                        (self.to_future_hack)(&mut self.handler, typed_dispatch, connection)
361                            .await?;
362                    match result.into_handled() {
363                        Handled::Yes => Ok(Handled::Yes),
364                        Handled::No {
365                            message: Dispatch::Request(request, responder),
366                            retry,
367                        } => {
368                            let untyped = request.to_untyped_message()?;
369                            Ok(Handled::No {
370                                message: Dispatch::Request(untyped, responder.erase_to_json()),
371                                retry,
372                            })
373                        }
374                        Handled::No {
375                            message: Dispatch::Notification(notification),
376                            retry,
377                        } => {
378                            let untyped = notification.to_untyped_message()?;
379                            Ok(Handled::No {
380                                message: Dispatch::Notification(untyped),
381                                retry,
382                            })
383                        }
384                        Handled::No {
385                            message: Dispatch::Response(result, responder),
386                            retry,
387                        } => {
388                            let method = responder.method();
389                            let untyped_result = match result {
390                                Ok(response) => response.into_json(method).map(Ok),
391                                Err(err) => Ok(Err(err)),
392                            }?;
393                            Ok(Handled::No {
394                                message: Dispatch::Response(
395                                    untyped_result,
396                                    responder.erase_to_json(),
397                                ),
398                                retry,
399                            })
400                        }
401                    }
402                }
403
404                Err(dispatch) => Ok(Handled::No {
405                    message: dispatch,
406                    retry: false,
407                }),
408            },
409        )
410        .await
411    }
412}
413
414/// Wraps a handler with an optional name for tracing/debugging.
415pub struct NamedHandler<H> {
416    name: Option<String>,
417    handler: H,
418}
419
420impl<H> NamedHandler<H> {
421    /// Creates a new named handler
422    pub fn new(name: Option<String>, handler: H) -> Self {
423        Self { name, handler }
424    }
425}
426
427impl<Counterpart: Role, H: HandleDispatchFrom<Counterpart>> HandleDispatchFrom<Counterpart>
428    for NamedHandler<H>
429{
430    fn describe_chain(&self) -> impl std::fmt::Debug {
431        format!(
432            "NamedHandler({:?}, {:?})",
433            self.name,
434            self.handler.describe_chain()
435        )
436    }
437
438    async fn handle_dispatch_from(
439        &mut self,
440        message: Dispatch,
441        connection: ConnectionTo<Counterpart>,
442    ) -> Result<Handled<Dispatch>, crate::Error> {
443        if let Some(name) = &self.name {
444            crate::util::instrumented_with_connection_name(
445                name.clone(),
446                self.handler.handle_dispatch_from(message, connection),
447            )
448            .await
449        } else {
450            self.handler.handle_dispatch_from(message, connection).await
451        }
452    }
453}
454
455/// Chains two handlers together, trying the first handler and falling back to the second
456pub struct ChainedHandler<H1, H2> {
457    handler1: H1,
458    handler2: H2,
459}
460
461impl<H1, H2> ChainedHandler<H1, H2> {
462    /// Creates a new chain handler
463    pub fn new(handler1: H1, handler2: H2) -> Self {
464        Self { handler1, handler2 }
465    }
466}
467
468impl<Counterpart: Role, H1, H2> HandleDispatchFrom<Counterpart> for ChainedHandler<H1, H2>
469where
470    H1: HandleDispatchFrom<Counterpart>,
471    H2: HandleDispatchFrom<Counterpart>,
472{
473    fn describe_chain(&self) -> impl std::fmt::Debug {
474        format!(
475            "{:?}, {:?}",
476            self.handler1.describe_chain(),
477            self.handler2.describe_chain()
478        )
479    }
480
481    async fn handle_dispatch_from(
482        &mut self,
483        message: Dispatch,
484        connection: ConnectionTo<Counterpart>,
485    ) -> Result<Handled<Dispatch>, crate::Error> {
486        match self
487            .handler1
488            .handle_dispatch_from(message, connection.clone())
489            .await?
490        {
491            Handled::Yes => Ok(Handled::Yes),
492            Handled::No {
493                message,
494                retry: retry1,
495            } => match self
496                .handler2
497                .handle_dispatch_from(message, connection)
498                .await?
499            {
500                Handled::Yes => Ok(Handled::Yes),
501                Handled::No {
502                    message,
503                    retry: retry2,
504                } => Ok(Handled::No {
505                    message,
506                    retry: retry1 | retry2,
507                }),
508            },
509        }
510    }
511}
agent_client_protocol/jsonrpc/handlers.rs

agent_client_protocol/jsonrpc/
handlers.rs
