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