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

agent_client_protocol/jsonrpc/
handlers.rs
