agent_client_protocol/jsonrpc/
handlers.rs

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

agent_client_protocol/jsonrpc/
handlers.rs
