agent_client_protocol/util/typed.rs
1//! Utilities for pattern matching on untyped JSON-RPC messages.
2//!
3//! When handling [`UntypedMessage`]s, you can use [`MatchDispatch`] for simple parsing
4//! or [`MatchDispatchFrom`] when you need peer-aware transforms (e.g., unwrapping
5//! proxy envelopes).
6//!
7//! # When to use which
8//!
9//! - **[`MatchDispatchFrom`]**: Preferred over implementing [`HandleDispatchFrom`] directly.
10//! Use this in connection handlers when you need to match on message types with
11//! proper peer-aware transforms (e.g., unwrapping `SuccessorMessage` envelopes).
12//!
13//! - **[`MatchDispatch`]**: Use this when you already have an unwrapped message and
14//! just need to parse it, such as inside a [`MatchDispatchFrom`] callback or when
15//! processing messages that don't need peer transforms.
16//!
17//! [`HandleDispatchFrom`]: crate::HandleDispatchFrom
18//! [`UntypedMessage`]: crate::UntypedMessage
19
20use crate::{
21 ConnectionTo, Dispatch, HandleDispatchFrom, Handled, JsonRpcNotification, JsonRpcRequest,
22 JsonRpcResponse, Responder, ResponseRouter,
23 role::{HasPeer, Role, handle_incoming_dispatch},
24};
25
26fn preserve_retry(
27 state: Result<Handled<Dispatch>, crate::Error>,
28 prior_retry: bool,
29) -> Result<Handled<Dispatch>, crate::Error> {
30 state.map(|state| match state {
31 Handled::Yes => Handled::Yes,
32 Handled::No { message, retry } => Handled::No {
33 message,
34 retry: prior_retry | retry,
35 },
36 })
37}
38
39/// Role-agnostic helper for pattern-matching on untyped JSON-RPC messages.
40///
41/// Use this when you already have an unwrapped message and just need to parse it,
42/// such as inside a [`MatchDispatchFrom`] callback or when processing messages
43/// that don't need peer transforms.
44///
45/// For connection handlers where you need proper peer-aware transforms,
46/// use [`MatchDispatchFrom`] instead.
47///
48/// # Example
49///
50/// ```
51/// # use agent_client_protocol::Dispatch;
52/// # use agent_client_protocol::v1::schema::{AgentCapabilities, InitializeRequest, InitializeResponse};
53/// # use agent_client_protocol::util::MatchDispatch;
54/// # async fn example(message: Dispatch) -> Result<(), agent_client_protocol::Error> {
55/// MatchDispatch::new(message)
56/// .if_request(|req: InitializeRequest, responder: agent_client_protocol::Responder<InitializeResponse>| async move {
57/// let response = InitializeResponse::new(req.protocol_version)
58/// .agent_capabilities(AgentCapabilities::new());
59/// responder.respond(response)
60/// })
61/// .await
62/// .otherwise(|message| async move {
63/// match message {
64/// Dispatch::Request(_, responder) => {
65/// responder.respond_with_error(agent_client_protocol::Error::method_not_found())
66/// }
67/// Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(()),
68/// }
69/// })
70/// .await
71/// # }
72/// ```
73#[must_use]
74#[derive(Debug)]
75pub struct MatchDispatch {
76 state: Result<Handled<Dispatch>, crate::Error>,
77}
78
79impl MatchDispatch {
80 /// Create a new pattern matcher for the given message.
81 pub fn new(message: Dispatch) -> Self {
82 Self {
83 state: Ok(Handled::No {
84 message,
85 retry: false,
86 }),
87 }
88 }
89
90 /// Try to handle the message as a request of type `Req`.
91 ///
92 /// If the message is a request for one of `Req`'s methods, it is parsed and
93 /// `op` is called with the request and a typed responder. If the method
94 /// matches but the params do not parse, matching stops with the parse
95 /// error. Other messages, and messages already handled by a previous call,
96 /// pass through unchanged.
97 pub async fn if_request<Req: JsonRpcRequest, H>(
98 mut self,
99 op: impl AsyncFnOnce(Req, Responder<Req::Response>) -> Result<H, crate::Error>,
100 ) -> Self
101 where
102 H: crate::IntoHandled<(Req, Responder<Req::Response>)>,
103 {
104 if let Ok(Handled::No {
105 message: dispatch,
106 retry,
107 }) = self.state
108 {
109 self.state = match dispatch {
110 Dispatch::Request(untyped_request, untyped_responder) => {
111 if Req::matches_method(untyped_request.method()) {
112 match Req::parse_message(untyped_request.method(), untyped_request.params())
113 {
114 Ok(typed_request) => {
115 let typed_responder = untyped_responder.cast();
116 match op(typed_request, typed_responder).await {
117 Ok(result) => match result.into_handled() {
118 Handled::Yes => Ok(Handled::Yes),
119 Handled::No {
120 message: (request, responder),
121 retry: request_retry,
122 } => match request.to_untyped_message() {
123 Ok(untyped) => Ok(Handled::No {
124 message: Dispatch::Request(
125 untyped,
126 responder.erase_to_json(),
127 ),
128 retry: retry | request_retry,
129 }),
130 Err(err) => Err(err),
131 },
132 },
133 Err(err) => Err(err),
134 }
135 }
136 Err(err) => Err(err),
137 }
138 } else {
139 Ok(Handled::No {
140 message: Dispatch::Request(untyped_request, untyped_responder),
141 retry,
142 })
143 }
144 }
145 Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(Handled::No {
146 message: dispatch,
147 retry,
148 }),
149 };
150 }
151 self
152 }
153
154 /// Try to handle the message as a notification of type `N`.
155 ///
156 /// If the message is a notification for one of `N`'s methods, it is parsed
157 /// and `op` is called with the notification. If the method matches but the
158 /// params do not parse, matching stops with the parse error. Other
159 /// messages, and messages already handled by a previous call, pass through
160 /// unchanged.
161 pub async fn if_notification<N: JsonRpcNotification, H>(
162 mut self,
163 op: impl AsyncFnOnce(N) -> Result<H, crate::Error>,
164 ) -> Self
165 where
166 H: crate::IntoHandled<N>,
167 {
168 if let Ok(Handled::No {
169 message: dispatch,
170 retry,
171 }) = self.state
172 {
173 self.state = match dispatch {
174 Dispatch::Notification(untyped_notification) => {
175 if N::matches_method(untyped_notification.method()) {
176 match N::parse_message(
177 untyped_notification.method(),
178 untyped_notification.params(),
179 ) {
180 Ok(typed_notification) => match op(typed_notification).await {
181 Ok(result) => match result.into_handled() {
182 Handled::Yes => Ok(Handled::Yes),
183 Handled::No {
184 message: notification,
185 retry: notification_retry,
186 } => match notification.to_untyped_message() {
187 Ok(untyped) => Ok(Handled::No {
188 message: Dispatch::Notification(untyped),
189 retry: retry | notification_retry,
190 }),
191 Err(err) => Err(err),
192 },
193 },
194 Err(err) => Err(err),
195 },
196 Err(err) => Err(err),
197 }
198 } else {
199 Ok(Handled::No {
200 message: Dispatch::Notification(untyped_notification),
201 retry,
202 })
203 }
204 }
205 Dispatch::Request(_, _) | Dispatch::Response(_, _) => Ok(Handled::No {
206 message: dispatch,
207 retry,
208 }),
209 };
210 }
211 self
212 }
213
214 /// Try to handle the message as a typed `Dispatch<R, N>`.
215 ///
216 /// This attempts to parse the message as either request type `R` or notification type `N`,
217 /// providing a typed `Dispatch` to the handler if successful.
218 pub async fn if_dispatch<R: JsonRpcRequest, N: JsonRpcNotification, H>(
219 mut self,
220 op: impl AsyncFnOnce(Dispatch<R, N>) -> Result<H, crate::Error>,
221 ) -> Self
222 where
223 H: crate::IntoHandled<Dispatch<R, N>>,
224 {
225 if let Ok(Handled::No {
226 message: dispatch,
227 retry,
228 }) = self.state
229 {
230 self.state = match dispatch.into_typed_dispatch::<R, N>() {
231 Ok(Ok(typed_dispatch)) => match op(typed_dispatch).await {
232 Ok(result) => match result.into_handled() {
233 Handled::Yes => Ok(Handled::Yes),
234 Handled::No {
235 message: typed_dispatch,
236 retry: message_retry,
237 } => {
238 let untyped = match typed_dispatch {
239 Dispatch::Request(request, responder) => {
240 match request.to_untyped_message() {
241 Ok(untyped) => {
242 Dispatch::Request(untyped, responder.erase_to_json())
243 }
244 Err(err) => return Self { state: Err(err) },
245 }
246 }
247 Dispatch::Notification(notification) => {
248 match notification.to_untyped_message() {
249 Ok(untyped) => Dispatch::Notification(untyped),
250 Err(err) => return Self { state: Err(err) },
251 }
252 }
253 Dispatch::Response(result, router) => {
254 let method = router.method();
255 let untyped_result = match result {
256 Ok(response) => match response.into_json(method) {
257 Ok(json) => Ok(json),
258 Err(err) => return Self { state: Err(err) },
259 },
260 Err(err) => Err(err),
261 };
262 Dispatch::Response(untyped_result, router.erase_to_json())
263 }
264 };
265 Ok(Handled::No {
266 message: untyped,
267 retry: retry | message_retry,
268 })
269 }
270 },
271 Err(err) => Err(err),
272 },
273 Ok(Err(dispatch)) => Ok(Handled::No {
274 message: dispatch,
275 retry,
276 }),
277 Err(err) => Err(err),
278 };
279 }
280 self
281 }
282
283 /// Try to handle the message as a response to a request of type `Req`.
284 ///
285 /// If the message is a `Response` variant and the method matches `Req`, the handler
286 /// is called with the result (which may be `Ok` or `Err`) and a typed response context.
287 /// Use this when you need to handle both success and error responses.
288 ///
289 /// For handling only successful responses, see [`if_ok_response_to`](Self::if_ok_response_to).
290 pub async fn if_response_to<Req: JsonRpcRequest, H>(
291 mut self,
292 op: impl AsyncFnOnce(
293 Result<Req::Response, crate::Error>,
294 ResponseRouter<Req::Response>,
295 ) -> Result<H, crate::Error>,
296 ) -> Self
297 where
298 H: crate::IntoHandled<(
299 Result<Req::Response, crate::Error>,
300 ResponseRouter<Req::Response>,
301 )>,
302 {
303 if let Ok(Handled::No {
304 message: dispatch,
305 retry,
306 }) = self.state
307 {
308 self.state = match dispatch {
309 Dispatch::Response(result, router) => {
310 // Check if the request type matches this method
311 if Req::matches_method(router.method()) {
312 // Method matches, parse the response
313 let typed_router: ResponseRouter<Req::Response> = router.cast();
314 let typed_result = match result {
315 Ok(value) => Req::Response::from_json(typed_router.method(), value),
316 Err(err) => Err(err),
317 };
318
319 match op(typed_result, typed_router).await {
320 Ok(handler_result) => match handler_result.into_handled() {
321 Handled::Yes => Ok(Handled::Yes),
322 Handled::No {
323 message: (result, router),
324 retry: response_retry,
325 } => {
326 // Convert typed result back to untyped
327 let untyped_result = match result {
328 Ok(response) => response.into_json(router.method()),
329 Err(err) => Err(err),
330 };
331 Ok(Handled::No {
332 message: Dispatch::Response(
333 untyped_result,
334 router.erase_to_json(),
335 ),
336 retry: retry | response_retry,
337 })
338 }
339 },
340 Err(err) => Err(err),
341 }
342 } else {
343 // Method doesn't match, return unhandled
344 Ok(Handled::No {
345 message: Dispatch::Response(result, router),
346 retry,
347 })
348 }
349 }
350 Dispatch::Request(_, _) | Dispatch::Notification(_) => Ok(Handled::No {
351 message: dispatch,
352 retry,
353 }),
354 };
355 }
356 self
357 }
358
359 /// Try to handle the message as a successful response to a request of type `Req`.
360 ///
361 /// If the message is a `Response` variant with an `Ok` result and the method matches `Req`,
362 /// the handler is called with the parsed response and a typed response context.
363 /// Error responses are passed through without calling the handler.
364 ///
365 /// This is a convenience wrapper around [`if_response_to`](Self::if_response_to) for the
366 /// common case where you only care about successful responses.
367 pub async fn if_ok_response_to<Req: JsonRpcRequest, H>(
368 self,
369 op: impl AsyncFnOnce(Req::Response, ResponseRouter<Req::Response>) -> Result<H, crate::Error>,
370 ) -> Self
371 where
372 H: crate::IntoHandled<(Req::Response, ResponseRouter<Req::Response>)>,
373 {
374 self.if_response_to::<Req, _>(async move |result, router| match result {
375 Ok(response) => {
376 let handler_result = op(response, router).await?;
377 match handler_result.into_handled() {
378 Handled::Yes => Ok(Handled::Yes),
379 Handled::No {
380 message: (resp, router),
381 retry,
382 } => Ok(Handled::No {
383 message: (Ok(resp), router),
384 retry,
385 }),
386 }
387 }
388 Err(err) => Ok(Handled::No {
389 message: (Err(err), router),
390 retry: false,
391 }),
392 })
393 .await
394 }
395
396 /// Complete matching, returning `Handled::No` if no match was found.
397 pub fn done(self) -> Result<Handled<Dispatch>, crate::Error> {
398 self.state
399 }
400
401 /// Handle messages that didn't match any previous handler.
402 pub async fn otherwise(
403 self,
404 op: impl AsyncFnOnce(Dispatch) -> Result<(), crate::Error>,
405 ) -> Result<(), crate::Error> {
406 match self.state {
407 Ok(Handled::Yes) => Ok(()),
408 Ok(Handled::No { message, retry: _ }) => op(message).await,
409 Err(err) => Err(err),
410 }
411 }
412
413 /// Finish matching, answering an unmatched request with a method not
414 /// found error (`-32601`) and ignoring an unmatched notification or
415 /// response.
416 pub fn otherwise_ignore(self) -> Result<(), crate::Error> {
417 match self.state {
418 Ok(Handled::No {
419 message: Dispatch::Request(_, responder),
420 retry: _,
421 }) => responder.respond_with_error(crate::Error::method_not_found()),
422 Ok(_) => Ok(()),
423 Err(err) => Err(err),
424 }
425 }
426}
427
428/// Role-aware helper for pattern-matching on untyped JSON-RPC requests.
429///
430/// **Prefer this over implementing [`HandleDispatchFrom`] directly.** This provides
431/// a more ergonomic API for matching on message types in connection handlers.
432///
433/// Use this when you need peer-aware transforms (e.g., unwrapping proxy envelopes)
434/// before parsing messages. For simple parsing without peer awareness (e.g., inside
435/// a callback), use [`MatchDispatch`] instead.
436///
437/// This wraps [`MatchDispatch`] and applies peer-specific message transformations
438/// via `remote_style().handle_incoming_dispatch()` before delegating to `MatchDispatch`
439/// for the actual parsing.
440///
441/// [`HandleDispatchFrom`]: crate::HandleDispatchFrom
442///
443/// # Example
444///
445/// ```
446/// # use agent_client_protocol::Dispatch;
447/// # use agent_client_protocol::v1::schema::{AgentCapabilities, InitializeRequest, InitializeResponse, PromptRequest, PromptResponse, StopReason};
448/// # use agent_client_protocol::util::MatchDispatchFrom;
449/// # async fn example(message: Dispatch, cx: &agent_client_protocol::ConnectionTo<agent_client_protocol::Client>) -> Result<(), agent_client_protocol::Error> {
450/// MatchDispatchFrom::new(message, cx)
451/// .if_request(|req: InitializeRequest, responder: agent_client_protocol::Responder<InitializeResponse>| async move {
452/// // Handle initialization
453/// let response = InitializeResponse::new(req.protocol_version)
454/// .agent_capabilities(AgentCapabilities::new());
455/// responder.respond(response)
456/// })
457/// .await
458/// .if_request(|_req: PromptRequest, responder: agent_client_protocol::Responder<PromptResponse>| async move {
459/// // Handle prompts
460/// responder.respond(PromptResponse::new(StopReason::EndTurn))
461/// })
462/// .await
463/// .otherwise(|message| async move {
464/// // Fallback for unrecognized messages
465/// match message {
466/// Dispatch::Request(_, responder) => responder.respond_with_error(agent_client_protocol::Error::method_not_found()),
467/// Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(()),
468/// }
469/// })
470/// .await
471/// # }
472/// ```
473#[must_use]
474#[derive(Debug)]
475pub struct MatchDispatchFrom<Counterpart: Role> {
476 state: Result<Handled<Dispatch>, crate::Error>,
477 connection: ConnectionTo<Counterpart>,
478}
479
480impl<Counterpart: Role> MatchDispatchFrom<Counterpart> {
481 /// Create a new pattern matcher for the given untyped request message.
482 pub fn new(message: Dispatch, cx: &ConnectionTo<Counterpart>) -> Self {
483 Self {
484 state: Ok(Handled::No {
485 message,
486 retry: false,
487 }),
488 connection: cx.clone(),
489 }
490 }
491
492 /// Try to handle the message as a request of type `Req`.
493 ///
494 /// If the message is a request for one of `Req`'s methods, it is parsed and
495 /// `op` is called with the request and a typed responder. If the method
496 /// matches but the params do not parse, matching stops with the parse
497 /// error. Other messages, and messages already handled by a previous
498 /// `if_*` call, pass through unchanged.
499 ///
500 /// The handler can return either `()` (which becomes `Handled::Yes`) or an explicit
501 /// `Handled` value to control whether the message should be passed to the next handler.
502 ///
503 /// Returns `self` to allow chaining multiple `if_*` calls.
504 pub async fn if_request<Req: JsonRpcRequest, H>(
505 self,
506 op: impl AsyncFnOnce(Req, Responder<Req::Response>) -> Result<H, crate::Error>,
507 ) -> Self
508 where
509 Counterpart: HasPeer<Counterpart>,
510 H: crate::IntoHandled<(Req, Responder<Req::Response>)>,
511 {
512 let counterpart = self.connection.counterpart();
513 self.if_request_from(counterpart, op).await
514 }
515
516 /// Try to handle the message as a request of type `Req` from a specific peer.
517 ///
518 /// This is similar to [`if_request`](Self::if_request), but first applies peer-specific
519 /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes when receiving
520 /// from an agent via a proxy).
521 ///
522 /// # Parameters
523 ///
524 /// * `peer` - The peer the message is expected to come from
525 /// * `op` - The handler to call if the message matches
526 pub async fn if_request_from<Peer: Role, Req: JsonRpcRequest, H>(
527 mut self,
528 peer: Peer,
529 op: impl AsyncFnOnce(Req, Responder<Req::Response>) -> Result<H, crate::Error>,
530 ) -> Self
531 where
532 Counterpart: HasPeer<Peer>,
533 H: crate::IntoHandled<(Req, Responder<Req::Response>)>,
534 {
535 if let Ok(Handled::No { message, retry }) = self.state {
536 let state = handle_incoming_dispatch(
537 self.connection.counterpart(),
538 peer,
539 message,
540 self.connection.clone(),
541 async |dispatch, _connection| {
542 // Delegate to MatchDispatch for parsing
543 MatchDispatch::new(dispatch).if_request(op).await.done()
544 },
545 )
546 .await;
547 self.state = preserve_retry(state, retry);
548 }
549 self
550 }
551
552 /// Try to handle the message as a notification of type `N`.
553 ///
554 /// If the message is a notification for one of `N`'s methods, it is parsed
555 /// and `op` is called with the notification. If the method matches but the
556 /// params do not parse, matching stops with the parse error. Other
557 /// messages, and messages already handled by a previous `if_*` call, pass
558 /// through unchanged.
559 ///
560 /// The handler can return either `()` (which becomes `Handled::Yes`) or an explicit
561 /// `Handled` value to control whether the message should be passed to the next handler.
562 ///
563 /// Returns `self` to allow chaining multiple `if_*` calls.
564 pub async fn if_notification<N: JsonRpcNotification, H>(
565 self,
566 op: impl AsyncFnOnce(N) -> Result<H, crate::Error>,
567 ) -> Self
568 where
569 Counterpart: HasPeer<Counterpart>,
570 H: crate::IntoHandled<N>,
571 {
572 let counterpart = self.connection.counterpart();
573 self.if_notification_from(counterpart, op).await
574 }
575
576 /// Try to handle the message as a notification of type `N` from a specific peer.
577 ///
578 /// This is similar to [`if_notification`](Self::if_notification), but first applies peer-specific
579 /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes when receiving
580 /// from an agent via a proxy).
581 ///
582 /// # Parameters
583 ///
584 /// * `peer` - The peer the message is expected to come from
585 /// * `op` - The handler to call if the message matches
586 pub async fn if_notification_from<Peer: Role, N: JsonRpcNotification, H>(
587 mut self,
588 peer: Peer,
589 op: impl AsyncFnOnce(N) -> Result<H, crate::Error>,
590 ) -> Self
591 where
592 Counterpart: HasPeer<Peer>,
593 H: crate::IntoHandled<N>,
594 {
595 if let Ok(Handled::No { message, retry }) = self.state {
596 let state = handle_incoming_dispatch(
597 self.connection.counterpart(),
598 peer,
599 message,
600 self.connection.clone(),
601 async |dispatch, _connection| {
602 // Delegate to MatchDispatch for parsing
603 MatchDispatch::new(dispatch)
604 .if_notification(op)
605 .await
606 .done()
607 },
608 )
609 .await;
610 self.state = preserve_retry(state, retry);
611 }
612 self
613 }
614
615 /// Try to handle the message as a typed `Dispatch<Req, N>` from a specific peer.
616 ///
617 /// This is similar to [`MatchDispatch::if_dispatch`], but first applies peer-specific
618 /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes).
619 ///
620 /// # Parameters
621 ///
622 /// * `peer` - The peer the message is expected to come from
623 /// * `op` - The handler to call if the message matches
624 pub async fn if_dispatch_from<Peer: Role, Req: JsonRpcRequest, N: JsonRpcNotification, H>(
625 mut self,
626 peer: Peer,
627 op: impl AsyncFnOnce(Dispatch<Req, N>) -> Result<H, crate::Error>,
628 ) -> Self
629 where
630 Counterpart: HasPeer<Peer>,
631 H: crate::IntoHandled<Dispatch<Req, N>>,
632 {
633 if let Ok(Handled::No { message, retry }) = self.state {
634 let state = handle_incoming_dispatch(
635 self.connection.counterpart(),
636 peer,
637 message,
638 self.connection.clone(),
639 async |dispatch, _connection| {
640 // Delegate to MatchDispatch for parsing
641 MatchDispatch::new(dispatch).if_dispatch(op).await.done()
642 },
643 )
644 .await;
645 self.state = preserve_retry(state, retry);
646 }
647 self
648 }
649
650 /// Try to handle the message as a response to a request of type `Req`.
651 ///
652 /// If the message is a `Response` variant and the method matches `Req`, the handler
653 /// is called with the result (which may be `Ok` or `Err`) and a typed response context.
654 ///
655 /// Unlike requests and notifications, responses don't need peer-specific transforms
656 /// (they don't have the `SuccessorMessage` envelope structure), so this method
657 /// delegates directly to [`MatchDispatch::if_response_to`].
658 pub async fn if_response_to<Req: JsonRpcRequest, H>(
659 mut self,
660 op: impl AsyncFnOnce(
661 Result<Req::Response, crate::Error>,
662 ResponseRouter<Req::Response>,
663 ) -> Result<H, crate::Error>,
664 ) -> Self
665 where
666 H: crate::IntoHandled<(
667 Result<Req::Response, crate::Error>,
668 ResponseRouter<Req::Response>,
669 )>,
670 {
671 if let Ok(Handled::No { message, retry }) = self.state {
672 let state = MatchDispatch::new(message)
673 .if_response_to::<Req, H>(op)
674 .await
675 .done();
676 self.state = preserve_retry(state, retry);
677 }
678 self
679 }
680
681 /// Try to handle the message as a successful response to a request of type `Req`.
682 ///
683 /// If the message is a `Response` variant with an `Ok` result and the method matches `Req`,
684 /// the handler is called with the parsed response and a typed response context.
685 /// Error responses are passed through without calling the handler.
686 ///
687 /// This is a convenience wrapper around [`if_response_to`](Self::if_response_to).
688 pub async fn if_ok_response_to<Req: JsonRpcRequest, H>(
689 self,
690 op: impl AsyncFnOnce(Req::Response, ResponseRouter<Req::Response>) -> Result<H, crate::Error>,
691 ) -> Self
692 where
693 Counterpart: HasPeer<Counterpart>,
694 H: crate::IntoHandled<(Req::Response, ResponseRouter<Req::Response>)>,
695 {
696 let counterpart = self.connection.counterpart();
697 self.if_ok_response_to_from::<Req, Counterpart, H>(counterpart, op)
698 .await
699 }
700
701 /// Try to handle the message as a response to a request of type `Req` from a specific peer.
702 ///
703 /// If the message is a `Response` variant, the method matches `Req`, and the `role_id`
704 /// matches the expected peer, the handler is called with the result and a typed response context.
705 ///
706 /// This is used to filter responses by the peer they came from, which is important
707 /// in proxy scenarios where responses might arrive from multiple peers.
708 pub async fn if_response_to_from<Req: JsonRpcRequest, Peer: Role, H>(
709 mut self,
710 peer: Peer,
711 op: impl AsyncFnOnce(
712 Result<Req::Response, crate::Error>,
713 ResponseRouter<Req::Response>,
714 ) -> Result<H, crate::Error>,
715 ) -> Self
716 where
717 Counterpart: HasPeer<Peer>,
718 H: crate::IntoHandled<(
719 Result<Req::Response, crate::Error>,
720 ResponseRouter<Req::Response>,
721 )>,
722 {
723 if let Ok(Handled::No { message, retry }) = self.state {
724 let state = handle_incoming_dispatch(
725 self.connection.counterpart(),
726 peer,
727 message,
728 self.connection.clone(),
729 async |dispatch, _connection| {
730 // Delegate to MatchDispatch for parsing
731 MatchDispatch::new(dispatch)
732 .if_response_to::<Req, H>(op)
733 .await
734 .done()
735 },
736 )
737 .await;
738 self.state = preserve_retry(state, retry);
739 }
740 self
741 }
742
743 /// Try to handle the message as a successful response to a request of type `Req` from a specific peer.
744 ///
745 /// This is a convenience wrapper around [`if_response_to_from`](Self::if_response_to_from)
746 /// for the common case where you only care about successful responses.
747 pub async fn if_ok_response_to_from<Req: JsonRpcRequest, Peer: Role, H>(
748 self,
749 peer: Peer,
750 op: impl AsyncFnOnce(Req::Response, ResponseRouter<Req::Response>) -> Result<H, crate::Error>,
751 ) -> Self
752 where
753 Counterpart: HasPeer<Peer>,
754 H: crate::IntoHandled<(Req::Response, ResponseRouter<Req::Response>)>,
755 {
756 self.if_response_to_from::<Req, _, _>(peer, async move |result, router| match result {
757 Ok(response) => {
758 let handler_result = op(response, router).await?;
759 match handler_result.into_handled() {
760 Handled::Yes => Ok(Handled::Yes),
761 Handled::No {
762 message: (resp, router),
763 retry,
764 } => Ok(Handled::No {
765 message: (Ok(resp), router),
766 retry,
767 }),
768 }
769 }
770 Err(err) => Ok(Handled::No {
771 message: (Err(err), router),
772 retry: false,
773 }),
774 })
775 .await
776 }
777
778 /// Complete matching, returning `Handled::No` if no match was found.
779 pub fn done(self) -> Result<Handled<Dispatch>, crate::Error> {
780 match self.state {
781 Ok(Handled::Yes) => Ok(Handled::Yes),
782 Ok(Handled::No { message, retry }) => Ok(Handled::No { message, retry }),
783 Err(err) => Err(err),
784 }
785 }
786
787 /// Handle messages that didn't match any previous `if_*` call.
788 ///
789 /// This is the fallback handler that receives the original untyped message if none
790 /// of the typed handlers matched. You must call this method to complete the pattern
791 /// matching chain and get the final result.
792 pub async fn otherwise(
793 self,
794 op: impl AsyncFnOnce(Dispatch) -> Result<(), crate::Error>,
795 ) -> Result<(), crate::Error> {
796 match self.state {
797 Ok(Handled::Yes) => Ok(()),
798 Ok(Handled::No { message, retry: _ }) => op(message).await,
799 Err(err) => Err(err),
800 }
801 }
802
803 /// Handle messages that didn't match any previous `if_*` call.
804 ///
805 /// This is the fallback handler that receives the original untyped message if none
806 /// of the typed handlers matched. You must call this method to complete the pattern
807 /// matching chain and get the final result.
808 pub async fn otherwise_delegate(
809 self,
810 mut handler: impl HandleDispatchFrom<Counterpart>,
811 ) -> Result<Handled<Dispatch>, crate::Error> {
812 match self.state? {
813 Handled::Yes => Ok(Handled::Yes),
814 Handled::No {
815 message,
816 retry: outer_retry,
817 } => match handler
818 .handle_dispatch_from(message, self.connection.clone())
819 .await?
820 {
821 Handled::Yes => Ok(Handled::Yes),
822 Handled::No {
823 message,
824 retry: inner_retry,
825 } => Ok(Handled::No {
826 message,
827 retry: inner_retry | outer_retry,
828 }),
829 },
830 }
831 }
832}