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
19use crate::{
20 ConnectionTo, Dispatch, HandleDispatchFrom, Handled, JsonRpcNotification, JsonRpcRequest,
21 JsonRpcResponse, Responder, ResponseRouter, UntypedMessage,
22 role::{HasPeer, Role, handle_incoming_dispatch},
23};
24
25fn preserve_retry(
26 state: Result<Handled<Dispatch>, crate::Error>,
27 prior_retry: bool,
28) -> Result<Handled<Dispatch>, crate::Error> {
29 state.map(|state| match state {
30 Handled::Yes => Handled::Yes,
31 Handled::No { message, retry } => Handled::No {
32 message,
33 retry: prior_retry | retry,
34 },
35 })
36}
37
38/// Role-agnostic helper for pattern-matching on untyped JSON-RPC messages.
39///
40/// Use this when you already have an unwrapped message and just need to parse it,
41/// such as inside a [`MatchDispatchFrom`] callback or when processing messages
42/// that don't need peer transforms.
43///
44/// For connection handlers where you need proper peer-aware transforms,
45/// use [`MatchDispatchFrom`] instead.
46///
47/// # Example
48///
49/// ```
50/// # use agent_client_protocol::Dispatch;
51/// # use agent_client_protocol::schema::v1::{AgentCapabilities, InitializeRequest, InitializeResponse};
52/// # use agent_client_protocol::util::MatchDispatch;
53/// # async fn example(message: Dispatch) -> Result<(), agent_client_protocol::Error> {
54/// MatchDispatch::new(message)
55/// .if_request(|req: InitializeRequest, responder: agent_client_protocol::Responder<InitializeResponse>| async move {
56/// let response = InitializeResponse::new(req.protocol_version)
57/// .agent_capabilities(AgentCapabilities::new());
58/// responder.respond(response)
59/// })
60/// .await
61/// .otherwise(|message| async move {
62/// match message {
63/// Dispatch::Request(_, responder) => {
64/// responder.respond_with_error(agent_client_protocol::util::internal_error("unknown method"))
65/// }
66/// Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(()),
67/// }
68/// })
69/// .await
70/// # }
71/// ```
72#[must_use]
73#[derive(Debug)]
74pub struct MatchDispatch {
75 state: Result<Handled<Dispatch>, crate::Error>,
76}
77
78impl MatchDispatch {
79 /// Create a new pattern matcher for the given message.
80 pub fn new(message: Dispatch) -> Self {
81 Self {
82 state: Ok(Handled::No {
83 message,
84 retry: false,
85 }),
86 }
87 }
88
89 /// Try to handle the message as a request of type `Req`.
90 ///
91 /// If the message can be parsed as `Req`, the handler `op` is called with the parsed
92 /// request and a typed request context. If parsing fails or the message was already
93 /// handled by a previous call, this has no effect.
94 pub async fn if_request<Req: JsonRpcRequest, H>(
95 mut self,
96 op: impl AsyncFnOnce(Req, Responder<Req::Response>) -> Result<H, crate::Error>,
97 ) -> Self
98 where
99 H: crate::IntoHandled<(Req, Responder<Req::Response>)>,
100 {
101 if let Ok(Handled::No {
102 message: dispatch,
103 retry,
104 }) = self.state
105 {
106 self.state = match dispatch {
107 Dispatch::Request(untyped_request, untyped_responder) => {
108 if Req::matches_method(untyped_request.method()) {
109 match Req::parse_message(untyped_request.method(), untyped_request.params())
110 {
111 Ok(typed_request) => {
112 let typed_responder = untyped_responder.cast();
113 match op(typed_request, typed_responder).await {
114 Ok(result) => match result.into_handled() {
115 Handled::Yes => Ok(Handled::Yes),
116 Handled::No {
117 message: (request, responder),
118 retry: request_retry,
119 } => match request.to_untyped_message() {
120 Ok(untyped) => Ok(Handled::No {
121 message: Dispatch::Request(
122 untyped,
123 responder.erase_to_json(),
124 ),
125 retry: retry | request_retry,
126 }),
127 Err(err) => Err(err),
128 },
129 },
130 Err(err) => Err(err),
131 }
132 }
133 Err(err) => Err(err),
134 }
135 } else {
136 Ok(Handled::No {
137 message: Dispatch::Request(untyped_request, untyped_responder),
138 retry,
139 })
140 }
141 }
142 Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(Handled::No {
143 message: dispatch,
144 retry,
145 }),
146 };
147 }
148 self
149 }
150
151 /// Try to handle the message as a notification of type `N`.
152 ///
153 /// If the message can be parsed as `N`, the handler `op` is called with the parsed
154 /// notification. If parsing fails or the message was already handled, this has no effect.
155 pub async fn if_notification<N: JsonRpcNotification, H>(
156 mut self,
157 op: impl AsyncFnOnce(N) -> Result<H, crate::Error>,
158 ) -> Self
159 where
160 H: crate::IntoHandled<N>,
161 {
162 if let Ok(Handled::No {
163 message: dispatch,
164 retry,
165 }) = self.state
166 {
167 self.state = match dispatch {
168 Dispatch::Notification(untyped_notification) => {
169 if N::matches_method(untyped_notification.method()) {
170 match N::parse_message(
171 untyped_notification.method(),
172 untyped_notification.params(),
173 ) {
174 Ok(typed_notification) => match op(typed_notification).await {
175 Ok(result) => match result.into_handled() {
176 Handled::Yes => Ok(Handled::Yes),
177 Handled::No {
178 message: notification,
179 retry: notification_retry,
180 } => match notification.to_untyped_message() {
181 Ok(untyped) => Ok(Handled::No {
182 message: Dispatch::Notification(untyped),
183 retry: retry | notification_retry,
184 }),
185 Err(err) => Err(err),
186 },
187 },
188 Err(err) => Err(err),
189 },
190 Err(err) => Err(err),
191 }
192 } else {
193 Ok(Handled::No {
194 message: Dispatch::Notification(untyped_notification),
195 retry,
196 })
197 }
198 }
199 Dispatch::Request(_, _) | Dispatch::Response(_, _) => Ok(Handled::No {
200 message: dispatch,
201 retry,
202 }),
203 };
204 }
205 self
206 }
207
208 /// Try to handle the message as a typed `Dispatch<R, N>`.
209 ///
210 /// This attempts to parse the message as either request type `R` or notification type `N`,
211 /// providing a typed `Dispatch` to the handler if successful.
212 pub async fn if_dispatch<R: JsonRpcRequest, N: JsonRpcNotification, H>(
213 mut self,
214 op: impl AsyncFnOnce(Dispatch<R, N>) -> Result<H, crate::Error>,
215 ) -> Self
216 where
217 H: crate::IntoHandled<Dispatch<R, N>>,
218 {
219 if let Ok(Handled::No {
220 message: dispatch,
221 retry,
222 }) = self.state
223 {
224 self.state = match dispatch.into_typed_dispatch::<R, N>() {
225 Ok(Ok(typed_dispatch)) => match op(typed_dispatch).await {
226 Ok(result) => match result.into_handled() {
227 Handled::Yes => Ok(Handled::Yes),
228 Handled::No {
229 message: typed_dispatch,
230 retry: message_retry,
231 } => {
232 let untyped = match typed_dispatch {
233 Dispatch::Request(request, responder) => {
234 match request.to_untyped_message() {
235 Ok(untyped) => {
236 Dispatch::Request(untyped, responder.erase_to_json())
237 }
238 Err(err) => return Self { state: Err(err) },
239 }
240 }
241 Dispatch::Notification(notification) => {
242 match notification.to_untyped_message() {
243 Ok(untyped) => Dispatch::Notification(untyped),
244 Err(err) => return Self { state: Err(err) },
245 }
246 }
247 Dispatch::Response(result, router) => {
248 let method = router.method();
249 let untyped_result = match result {
250 Ok(response) => match response.into_json(method) {
251 Ok(json) => Ok(json),
252 Err(err) => return Self { state: Err(err) },
253 },
254 Err(err) => Err(err),
255 };
256 Dispatch::Response(untyped_result, router.erase_to_json())
257 }
258 };
259 Ok(Handled::No {
260 message: untyped,
261 retry: retry | message_retry,
262 })
263 }
264 },
265 Err(err) => Err(err),
266 },
267 Ok(Err(dispatch)) => Ok(Handled::No {
268 message: dispatch,
269 retry,
270 }),
271 Err(err) => Err(err),
272 };
273 }
274 self
275 }
276
277 /// Try to handle the message as a response to a request of type `Req`.
278 ///
279 /// If the message is a `Response` variant and the method matches `Req`, the handler
280 /// is called with the result (which may be `Ok` or `Err`) and a typed response context.
281 /// Use this when you need to handle both success and error responses.
282 ///
283 /// For handling only successful responses, see [`if_ok_response_to`](Self::if_ok_response_to).
284 pub async fn if_response_to<Req: JsonRpcRequest, H>(
285 mut self,
286 op: impl AsyncFnOnce(
287 Result<Req::Response, crate::Error>,
288 ResponseRouter<Req::Response>,
289 ) -> Result<H, crate::Error>,
290 ) -> Self
291 where
292 H: crate::IntoHandled<(
293 Result<Req::Response, crate::Error>,
294 ResponseRouter<Req::Response>,
295 )>,
296 {
297 if let Ok(Handled::No {
298 message: dispatch,
299 retry,
300 }) = self.state
301 {
302 self.state = match dispatch {
303 Dispatch::Response(result, router) => {
304 // Check if the request type matches this method
305 if Req::matches_method(router.method()) {
306 // Method matches, parse the response
307 let typed_router: ResponseRouter<Req::Response> = router.cast();
308 let typed_result = match result {
309 Ok(value) => Req::Response::from_value(typed_router.method(), value),
310 Err(err) => Err(err),
311 };
312
313 match op(typed_result, typed_router).await {
314 Ok(handler_result) => match handler_result.into_handled() {
315 Handled::Yes => Ok(Handled::Yes),
316 Handled::No {
317 message: (result, router),
318 retry: response_retry,
319 } => {
320 // Convert typed result back to untyped
321 let untyped_result = match result {
322 Ok(response) => response.into_json(router.method()),
323 Err(err) => Err(err),
324 };
325 Ok(Handled::No {
326 message: Dispatch::Response(
327 untyped_result,
328 router.erase_to_json(),
329 ),
330 retry: retry | response_retry,
331 })
332 }
333 },
334 Err(err) => Err(err),
335 }
336 } else {
337 // Method doesn't match, return unhandled
338 Ok(Handled::No {
339 message: Dispatch::Response(result, router),
340 retry,
341 })
342 }
343 }
344 Dispatch::Request(_, _) | Dispatch::Notification(_) => Ok(Handled::No {
345 message: dispatch,
346 retry,
347 }),
348 };
349 }
350 self
351 }
352
353 /// Try to handle the message as a successful response to a request of type `Req`.
354 ///
355 /// If the message is a `Response` variant with an `Ok` result and the method matches `Req`,
356 /// the handler is called with the parsed response and a typed response context.
357 /// Error responses are passed through without calling the handler.
358 ///
359 /// This is a convenience wrapper around [`if_response_to`](Self::if_response_to) for the
360 /// common case where you only care about successful responses.
361 pub async fn if_ok_response_to<Req: JsonRpcRequest, H>(
362 self,
363 op: impl AsyncFnOnce(Req::Response, ResponseRouter<Req::Response>) -> Result<H, crate::Error>,
364 ) -> Self
365 where
366 H: crate::IntoHandled<(Req::Response, ResponseRouter<Req::Response>)>,
367 {
368 self.if_response_to::<Req, _>(async move |result, router| match result {
369 Ok(response) => {
370 let handler_result = op(response, router).await?;
371 match handler_result.into_handled() {
372 Handled::Yes => Ok(Handled::Yes),
373 Handled::No {
374 message: (resp, router),
375 retry,
376 } => Ok(Handled::No {
377 message: (Ok(resp), router),
378 retry,
379 }),
380 }
381 }
382 Err(err) => Ok(Handled::No {
383 message: (Err(err), router),
384 retry: false,
385 }),
386 })
387 .await
388 }
389
390 /// Complete matching, returning `Handled::No` if no match was found.
391 pub fn done(self) -> Result<Handled<Dispatch>, crate::Error> {
392 self.state
393 }
394
395 /// Handle messages that didn't match any previous handler.
396 pub async fn otherwise(
397 self,
398 op: impl AsyncFnOnce(Dispatch) -> Result<(), crate::Error>,
399 ) -> Result<(), crate::Error> {
400 match self.state {
401 Ok(Handled::Yes) => Ok(()),
402 Ok(Handled::No { message, retry: _ }) => op(message).await,
403 Err(err) => Err(err),
404 }
405 }
406
407 /// Handle messages that didn't match any previous handler.
408 pub fn otherwise_ignore(self) -> Result<(), crate::Error> {
409 match self.state {
410 Ok(_) => Ok(()),
411 Err(err) => Err(err),
412 }
413 }
414}
415
416/// Role-aware helper for pattern-matching on untyped JSON-RPC requests.
417///
418/// **Prefer this over implementing [`HandleDispatchFrom`] directly.** This provides
419/// a more ergonomic API for matching on message types in connection handlers.
420///
421/// Use this when you need peer-aware transforms (e.g., unwrapping proxy envelopes)
422/// before parsing messages. For simple parsing without peer awareness (e.g., inside
423/// a callback), use [`MatchDispatch`] instead.
424///
425/// This wraps [`MatchDispatch`] and applies peer-specific message transformations
426/// via `remote_style().handle_incoming_dispatch()` before delegating to `MatchDispatch`
427/// for the actual parsing.
428///
429/// [`HandleDispatchFrom`]: crate::HandleDispatchFrom
430///
431/// # Example
432///
433/// ```
434/// # use agent_client_protocol::Dispatch;
435/// # use agent_client_protocol::schema::v1::{AgentCapabilities, InitializeRequest, InitializeResponse, PromptRequest, PromptResponse, StopReason};
436/// # use agent_client_protocol::util::MatchDispatchFrom;
437/// # async fn example(message: Dispatch, cx: &agent_client_protocol::ConnectionTo<agent_client_protocol::Client>) -> Result<(), agent_client_protocol::Error> {
438/// MatchDispatchFrom::new(message, cx)
439/// .if_request(|req: InitializeRequest, responder: agent_client_protocol::Responder<InitializeResponse>| async move {
440/// // Handle initialization
441/// let response = InitializeResponse::new(req.protocol_version)
442/// .agent_capabilities(AgentCapabilities::new());
443/// responder.respond(response)
444/// })
445/// .await
446/// .if_request(|_req: PromptRequest, responder: agent_client_protocol::Responder<PromptResponse>| async move {
447/// // Handle prompts
448/// responder.respond(PromptResponse::new(StopReason::EndTurn))
449/// })
450/// .await
451/// .otherwise(|message| async move {
452/// // Fallback for unrecognized messages
453/// match message {
454/// Dispatch::Request(_, responder) => responder.respond_with_error(agent_client_protocol::util::internal_error("unknown method")),
455/// Dispatch::Notification(_) | Dispatch::Response(_, _) => Ok(()),
456/// }
457/// })
458/// .await
459/// # }
460/// ```
461#[must_use]
462#[derive(Debug)]
463pub struct MatchDispatchFrom<Counterpart: Role> {
464 state: Result<Handled<Dispatch>, crate::Error>,
465 connection: ConnectionTo<Counterpart>,
466}
467
468impl<Counterpart: Role> MatchDispatchFrom<Counterpart> {
469 /// Create a new pattern matcher for the given untyped request message.
470 pub fn new(message: Dispatch, cx: &ConnectionTo<Counterpart>) -> Self {
471 Self {
472 state: Ok(Handled::No {
473 message,
474 retry: false,
475 }),
476 connection: cx.clone(),
477 }
478 }
479
480 /// Try to handle the message as a request of type `Req`.
481 ///
482 /// If the message can be parsed as `Req`, the handler `op` is called with the parsed
483 /// request and a typed request context. If parsing fails or the message was already
484 /// handled by a previous `handle_if`, this call has no effect.
485 ///
486 /// The handler can return either `()` (which becomes `Handled::Yes`) or an explicit
487 /// `Handled` value to control whether the message should be passed to the next handler.
488 ///
489 /// Returns `self` to allow chaining multiple `handle_if` calls.
490 pub async fn if_request<Req: JsonRpcRequest, H>(
491 self,
492 op: impl AsyncFnOnce(Req, Responder<Req::Response>) -> Result<H, crate::Error>,
493 ) -> Self
494 where
495 Counterpart: HasPeer<Counterpart>,
496 H: crate::IntoHandled<(Req, Responder<Req::Response>)>,
497 {
498 let counterpart = self.connection.counterpart();
499 self.if_request_from(counterpart, op).await
500 }
501
502 /// Try to handle the message as a request of type `Req` from a specific peer.
503 ///
504 /// This is similar to [`if_request`](Self::if_request), but first applies peer-specific
505 /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes when receiving
506 /// from an agent via a proxy).
507 ///
508 /// # Parameters
509 ///
510 /// * `peer` - The peer the message is expected to come from
511 /// * `op` - The handler to call if the message matches
512 pub async fn if_request_from<Peer: Role, Req: JsonRpcRequest, H>(
513 mut self,
514 peer: Peer,
515 op: impl AsyncFnOnce(Req, Responder<Req::Response>) -> Result<H, crate::Error>,
516 ) -> Self
517 where
518 Counterpart: HasPeer<Peer>,
519 H: crate::IntoHandled<(Req, Responder<Req::Response>)>,
520 {
521 if let Ok(Handled::No { message, retry }) = self.state {
522 let state = handle_incoming_dispatch(
523 self.connection.counterpart(),
524 peer,
525 message,
526 self.connection.clone(),
527 async |dispatch, _connection| {
528 // Delegate to MatchDispatch for parsing
529 MatchDispatch::new(dispatch).if_request(op).await.done()
530 },
531 )
532 .await;
533 self.state = preserve_retry(state, retry);
534 }
535 self
536 }
537
538 /// Try to handle the message as a notification of type `N`.
539 ///
540 /// If the message can be parsed as `N`, the handler `op` is called with the parsed
541 /// notification and connection context. If parsing fails or the message was already
542 /// handled by a previous `handle_if`, this call has no effect.
543 ///
544 /// The handler can return either `()` (which becomes `Handled::Yes`) or an explicit
545 /// `Handled` value to control whether the message should be passed to the next handler.
546 ///
547 /// Returns `self` to allow chaining multiple `handle_if` calls.
548 pub async fn if_notification<N: JsonRpcNotification, H>(
549 self,
550 op: impl AsyncFnOnce(N) -> Result<H, crate::Error>,
551 ) -> Self
552 where
553 Counterpart: HasPeer<Counterpart>,
554 H: crate::IntoHandled<N>,
555 {
556 let counterpart = self.connection.counterpart();
557 self.if_notification_from(counterpart, op).await
558 }
559
560 /// Try to handle the message as a notification of type `N` from a specific peer.
561 ///
562 /// This is similar to [`if_notification`](Self::if_notification), but first applies peer-specific
563 /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes when receiving
564 /// from an agent via a proxy).
565 ///
566 /// # Parameters
567 ///
568 /// * `peer` - The peer the message is expected to come from
569 /// * `op` - The handler to call if the message matches
570 pub async fn if_notification_from<Peer: Role, N: JsonRpcNotification, H>(
571 mut self,
572 peer: Peer,
573 op: impl AsyncFnOnce(N) -> Result<H, crate::Error>,
574 ) -> Self
575 where
576 Counterpart: HasPeer<Peer>,
577 H: crate::IntoHandled<N>,
578 {
579 if let Ok(Handled::No { message, retry }) = self.state {
580 let state = handle_incoming_dispatch(
581 self.connection.counterpart(),
582 peer,
583 message,
584 self.connection.clone(),
585 async |dispatch, _connection| {
586 // Delegate to MatchDispatch for parsing
587 MatchDispatch::new(dispatch)
588 .if_notification(op)
589 .await
590 .done()
591 },
592 )
593 .await;
594 self.state = preserve_retry(state, retry);
595 }
596 self
597 }
598
599 /// Try to handle the message as a typed `Dispatch<Req, N>` from a specific peer.
600 ///
601 /// This is similar to [`MatchDispatch::if_dispatch`], but first applies peer-specific
602 /// message transformation (e.g., unwrapping `SuccessorMessage` envelopes).
603 ///
604 /// # Parameters
605 ///
606 /// * `peer` - The peer the message is expected to come from
607 /// * `op` - The handler to call if the message matches
608 pub async fn if_dispatch_from<Peer: Role, Req: JsonRpcRequest, N: JsonRpcNotification, H>(
609 mut self,
610 peer: Peer,
611 op: impl AsyncFnOnce(Dispatch<Req, N>) -> Result<H, crate::Error>,
612 ) -> Self
613 where
614 Counterpart: HasPeer<Peer>,
615 H: crate::IntoHandled<Dispatch<Req, N>>,
616 {
617 if let Ok(Handled::No { message, retry }) = self.state {
618 let state = handle_incoming_dispatch(
619 self.connection.counterpart(),
620 peer,
621 message,
622 self.connection.clone(),
623 async |dispatch, _connection| {
624 // Delegate to MatchDispatch for parsing
625 MatchDispatch::new(dispatch).if_dispatch(op).await.done()
626 },
627 )
628 .await;
629 self.state = preserve_retry(state, retry);
630 }
631 self
632 }
633
634 /// Try to handle the message as a response to a request of type `Req`.
635 ///
636 /// If the message is a `Response` variant and the method matches `Req`, the handler
637 /// is called with the result (which may be `Ok` or `Err`) and a typed response context.
638 ///
639 /// Unlike requests and notifications, responses don't need peer-specific transforms
640 /// (they don't have the `SuccessorMessage` envelope structure), so this method
641 /// delegates directly to [`MatchDispatch::if_response_to`].
642 pub async fn if_response_to<Req: JsonRpcRequest, H>(
643 mut self,
644 op: impl AsyncFnOnce(
645 Result<Req::Response, crate::Error>,
646 ResponseRouter<Req::Response>,
647 ) -> Result<H, crate::Error>,
648 ) -> Self
649 where
650 H: crate::IntoHandled<(
651 Result<Req::Response, crate::Error>,
652 ResponseRouter<Req::Response>,
653 )>,
654 {
655 if let Ok(Handled::No { message, retry }) = self.state {
656 let state = MatchDispatch::new(message)
657 .if_response_to::<Req, H>(op)
658 .await
659 .done();
660 self.state = preserve_retry(state, retry);
661 }
662 self
663 }
664
665 /// Try to handle the message as a successful response to a request of type `Req`.
666 ///
667 /// If the message is a `Response` variant with an `Ok` result and the method matches `Req`,
668 /// the handler is called with the parsed response and a typed response context.
669 /// Error responses are passed through without calling the handler.
670 ///
671 /// This is a convenience wrapper around [`if_response_to`](Self::if_response_to).
672 pub async fn if_ok_response_to<Req: JsonRpcRequest, H>(
673 self,
674 op: impl AsyncFnOnce(Req::Response, ResponseRouter<Req::Response>) -> Result<H, crate::Error>,
675 ) -> Self
676 where
677 Counterpart: HasPeer<Counterpart>,
678 H: crate::IntoHandled<(Req::Response, ResponseRouter<Req::Response>)>,
679 {
680 let counterpart = self.connection.counterpart();
681 self.if_ok_response_to_from::<Req, Counterpart, H>(counterpart, op)
682 .await
683 }
684
685 /// Try to handle the message as a response to a request of type `Req` from a specific peer.
686 ///
687 /// If the message is a `Response` variant, the method matches `Req`, and the `role_id`
688 /// matches the expected peer, the handler is called with the result and a typed response context.
689 ///
690 /// This is used to filter responses by the peer they came from, which is important
691 /// in proxy scenarios where responses might arrive from multiple peers.
692 pub async fn if_response_to_from<Req: JsonRpcRequest, Peer: Role, H>(
693 mut self,
694 peer: Peer,
695 op: impl AsyncFnOnce(
696 Result<Req::Response, crate::Error>,
697 ResponseRouter<Req::Response>,
698 ) -> Result<H, crate::Error>,
699 ) -> Self
700 where
701 Counterpart: HasPeer<Peer>,
702 H: crate::IntoHandled<(
703 Result<Req::Response, crate::Error>,
704 ResponseRouter<Req::Response>,
705 )>,
706 {
707 if let Ok(Handled::No { message, retry }) = self.state {
708 let state = handle_incoming_dispatch(
709 self.connection.counterpart(),
710 peer,
711 message,
712 self.connection.clone(),
713 async |dispatch, _connection| {
714 // Delegate to MatchDispatch for parsing
715 MatchDispatch::new(dispatch)
716 .if_response_to::<Req, H>(op)
717 .await
718 .done()
719 },
720 )
721 .await;
722 self.state = preserve_retry(state, retry);
723 }
724 self
725 }
726
727 /// Try to handle the message as a successful response to a request of type `Req` from a specific peer.
728 ///
729 /// This is a convenience wrapper around [`if_response_to_from`](Self::if_response_to_from)
730 /// for the common case where you only care about successful responses.
731 pub async fn if_ok_response_to_from<Req: JsonRpcRequest, Peer: Role, H>(
732 self,
733 peer: Peer,
734 op: impl AsyncFnOnce(Req::Response, ResponseRouter<Req::Response>) -> Result<H, crate::Error>,
735 ) -> Self
736 where
737 Counterpart: HasPeer<Peer>,
738 H: crate::IntoHandled<(Req::Response, ResponseRouter<Req::Response>)>,
739 {
740 self.if_response_to_from::<Req, _, _>(peer, async move |result, router| match result {
741 Ok(response) => {
742 let handler_result = op(response, router).await?;
743 match handler_result.into_handled() {
744 Handled::Yes => Ok(Handled::Yes),
745 Handled::No {
746 message: (resp, router),
747 retry,
748 } => Ok(Handled::No {
749 message: (Ok(resp), router),
750 retry,
751 }),
752 }
753 }
754 Err(err) => Ok(Handled::No {
755 message: (Err(err), router),
756 retry: false,
757 }),
758 })
759 .await
760 }
761
762 /// Complete matching, returning `Handled::No` if no match was found.
763 pub fn done(self) -> Result<Handled<Dispatch>, crate::Error> {
764 match self.state {
765 Ok(Handled::Yes) => Ok(Handled::Yes),
766 Ok(Handled::No { message, retry }) => Ok(Handled::No { message, retry }),
767 Err(err) => Err(err),
768 }
769 }
770
771 /// Handle messages that didn't match any previous `handle_if` call.
772 ///
773 /// This is the fallback handler that receives the original untyped message if none
774 /// of the typed handlers matched. You must call this method to complete the pattern
775 /// matching chain and get the final result.
776 pub async fn otherwise(
777 self,
778 op: impl AsyncFnOnce(Dispatch) -> Result<(), crate::Error>,
779 ) -> Result<(), crate::Error> {
780 match self.state {
781 Ok(Handled::Yes) => Ok(()),
782 Ok(Handled::No { message, retry: _ }) => op(message).await,
783 Err(err) => Err(err),
784 }
785 }
786
787 /// Handle messages that didn't match any previous `handle_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_delegate(
793 self,
794 mut handler: impl HandleDispatchFrom<Counterpart>,
795 ) -> Result<Handled<Dispatch>, crate::Error> {
796 match self.state? {
797 Handled::Yes => Ok(Handled::Yes),
798 Handled::No {
799 message,
800 retry: outer_retry,
801 } => match handler
802 .handle_dispatch_from(message, self.connection.clone())
803 .await?
804 {
805 Handled::Yes => Ok(Handled::Yes),
806 Handled::No {
807 message,
808 retry: inner_retry,
809 } => Ok(Handled::No {
810 message,
811 retry: inner_retry | outer_retry,
812 }),
813 },
814 }
815 }
816}
817
818/// Builder for pattern-matching on untyped JSON-RPC notifications.
819///
820/// Similar to [`MatchDispatch`] but specifically for notifications (fire-and-forget messages with no response).
821///
822/// # Pattern
823///
824/// The typical pattern is:
825/// 1. Create a `TypeNotification` from an untyped message
826/// 2. Chain `.handle_if()` calls for each type you want to try
827/// 3. End with `.otherwise()` for messages that don't match any type
828///
829/// # Example
830///
831/// ```
832/// # use agent_client_protocol::UntypedMessage;
833/// # use agent_client_protocol::schema::v1::SessionNotification;
834/// # use agent_client_protocol::util::TypeNotification;
835/// # async fn example(message: UntypedMessage) -> Result<(), agent_client_protocol::Error> {
836/// TypeNotification::new(message)
837/// .handle_if(|notif: SessionNotification| async move {
838/// // Handle session notifications
839/// println!("Session update: {:?}", notif);
840/// Ok(())
841/// })
842/// .await
843/// .otherwise(|untyped: UntypedMessage| async move {
844/// // Fallback for unrecognized notifications
845/// println!("Unknown notification: {}", untyped.method);
846/// Ok(())
847/// })
848/// .await
849/// # }
850/// ```
851///
852/// Since notifications don't expect responses, handlers only receive the parsed
853/// notification (not a request context).
854#[must_use]
855#[derive(Debug)]
856pub struct TypeNotification {
857 state: Option<TypeNotificationState>,
858}
859
860#[derive(Debug)]
861enum TypeNotificationState {
862 Unhandled(String, serde_json::Value),
863 Handled(Result<(), crate::Error>),
864}
865
866impl TypeNotification {
867 /// Create a new pattern matcher for the given untyped notification message.
868 pub fn new(request: UntypedMessage) -> Self {
869 let UntypedMessage { method, params } = request;
870 Self {
871 state: Some(TypeNotificationState::Unhandled(method, params)),
872 }
873 }
874
875 /// Try to handle the message as type `N`.
876 ///
877 /// If the message can be parsed as `N`, the handler `op` is called with the parsed
878 /// notification. If parsing fails, the malformed matching notification is consumed
879 /// and logged without invoking the handler or sending a response. If the message was
880 /// already handled by a previous `handle_if`, this call has no effect.
881 ///
882 /// Returns `self` to allow chaining multiple `handle_if` calls.
883 pub async fn handle_if<N: JsonRpcNotification>(
884 mut self,
885 op: impl AsyncFnOnce(N) -> Result<(), crate::Error>,
886 ) -> Self {
887 self.state = Some(match self.state.take().expect("valid state") {
888 TypeNotificationState::Unhandled(method, params) => {
889 if N::matches_method(&method) {
890 match N::parse_message(&method, ¶ms) {
891 Ok(request) => TypeNotificationState::Handled(op(request).await),
892 Err(error) => {
893 tracing::warn!(
894 ?error,
895 method,
896 "Invalid notification params; ignoring notification without replying"
897 );
898 TypeNotificationState::Handled(Ok(()))
899 }
900 }
901 } else {
902 TypeNotificationState::Unhandled(method, params)
903 }
904 }
905
906 TypeNotificationState::Handled(err) => TypeNotificationState::Handled(err),
907 });
908 self
909 }
910
911 /// Handle messages that didn't match any previous `handle_if` call.
912 ///
913 /// This is the fallback handler that receives the original untyped message if none
914 /// of the typed handlers matched. You must call this method to complete the pattern
915 /// matching chain and get the final result.
916 pub async fn otherwise(
917 mut self,
918 op: impl AsyncFnOnce(UntypedMessage) -> Result<(), crate::Error>,
919 ) -> Result<(), crate::Error> {
920 match self.state.take().expect("valid state") {
921 TypeNotificationState::Unhandled(method, params) => {
922 match UntypedMessage::new(&method, params) {
923 Ok(m) => op(m).await,
924 Err(error) => {
925 tracing::warn!(
926 ?error,
927 method,
928 "Invalid untyped notification; ignoring notification without replying"
929 );
930 Ok(())
931 }
932 }
933 }
934 TypeNotificationState::Handled(r) => r,
935 }
936 }
937}
938
939#[cfg(test)]
940mod type_notification_tests {
941 use std::{cell::Cell, rc::Rc};
942
943 use serde::{Deserialize, Serialize};
944
945 use super::{TypeNotification, UntypedMessage};
946
947 #[derive(Clone, Debug, Serialize, Deserialize)]
948 struct TestNotification {
949 required: String,
950 }
951
952 impl crate::JsonRpcMessage for TestNotification {
953 fn matches_method(method: &str) -> bool {
954 method == "test/notification"
955 }
956
957 fn method(&self) -> &'static str {
958 "test/notification"
959 }
960
961 fn to_untyped_message(&self) -> Result<UntypedMessage, crate::Error> {
962 UntypedMessage::new(self.method(), self)
963 }
964
965 fn parse_message(method: &str, params: &impl Serialize) -> Result<Self, crate::Error> {
966 if !Self::matches_method(method) {
967 return Err(crate::Error::method_not_found());
968 }
969 crate::util::json_cast_params(params)
970 }
971 }
972
973 impl crate::JsonRpcNotification for TestNotification {}
974
975 #[test]
976 fn invalid_matching_notification_is_consumed_without_fallback() {
977 futures::executor::block_on(async {
978 let handler_called = Rc::new(Cell::new(false));
979 let fallback_called = Rc::new(Cell::new(false));
980 let handler_called_in_callback = Rc::clone(&handler_called);
981 let fallback_called_in_callback = Rc::clone(&fallback_called);
982 let message =
983 UntypedMessage::new("test/notification", serde_json::json!({ "wrong": "field" }))
984 .expect("test notification should serialize");
985
986 TypeNotification::new(message)
987 .handle_if(move |_: TestNotification| async move {
988 handler_called_in_callback.set(true);
989 Ok(())
990 })
991 .await
992 .otherwise(move |_| async move {
993 fallback_called_in_callback.set(true);
994 Ok(())
995 })
996 .await
997 .expect("invalid notification should be ignored");
998
999 assert!(!handler_called.get());
1000 assert!(!fallback_called.get());
1001 });
1002 }
1003}