Skip to main content

agenterra_rmcp/handler/server/
router.rs

1use 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}