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