agenterra_rmcp/handler/server/
router.rs1use std::sync::Arc;
2
3use tool::{IntoToolRoute, ToolRoute};
4
5use super::ServerHandler;
6use crate::{
7 RoleServer, Service,
8 model::{ClientRequest, ListToolsResult, ServerResult},
9 service::NotificationContext,
10};
11
12pub mod tool;
13
14pub struct Router<S> {
15 pub tool_router: tool::ToolRouter<S>,
16 pub service: Arc<S>,
17}
18
19impl<S> Router<S>
20where
21 S: ServerHandler,
22{
23 pub fn new(service: S) -> Self {
24 Self {
25 tool_router: tool::ToolRouter::new(),
26 service: Arc::new(service),
27 }
28 }
29
30 pub fn with_tool<R, A>(mut self, route: R) -> Self
31 where
32 R: IntoToolRoute<S, A>,
33 {
34 self.tool_router.add_route(route.into_tool_route());
35 self
36 }
37
38 pub fn with_tools(mut self, routes: impl IntoIterator<Item = ToolRoute<S>>) -> Self {
39 for route in routes {
40 self.tool_router.add_route(route);
41 }
42 self
43 }
44}
45
46impl<S> Service<RoleServer> for Router<S>
47where
48 S: ServerHandler,
49{
50 async fn handle_notification(
51 &self,
52 notification: <RoleServer as crate::service::ServiceRole>::PeerNot,
53 context: NotificationContext<RoleServer>,
54 ) -> Result<(), crate::Error> {
55 self.service
56 .handle_notification(notification, context)
57 .await
58 }
59 async fn handle_request(
60 &self,
61 request: <RoleServer as crate::service::ServiceRole>::PeerReq,
62 context: crate::service::RequestContext<RoleServer>,
63 ) -> Result<<RoleServer as crate::service::ServiceRole>::Resp, crate::Error> {
64 match request {
65 ClientRequest::CallToolRequest(request) => {
66 if self.tool_router.has_route(request.params.name.as_ref())
67 || !self.tool_router.transparent_when_not_found
68 {
69 let tool_call_context = crate::handler::server::tool::ToolCallContext::new(
70 self.service.as_ref(),
71 request.params,
72 context,
73 );
74 let result = self.tool_router.call(tool_call_context).await?;
75 Ok(ServerResult::CallToolResult(result))
76 } else {
77 self.service
78 .handle_request(ClientRequest::CallToolRequest(request), context)
79 .await
80 }
81 }
82 ClientRequest::ListToolsRequest(_) => {
83 let tools = self.tool_router.list_all();
84 Ok(ServerResult::ListToolsResult(ListToolsResult {
85 tools,
86 next_cursor: None,
87 }))
88 }
89 rest => self.service.handle_request(rest, context).await,
90 }
91 }
92
93 fn get_info(&self) -> <RoleServer as crate::service::ServiceRole>::Info {
94 self.service.get_info()
95 }
96}