Skip to main content

mcp_compressor_core/server/
registration.rs

1//! Frontend MCP server registration for compressed wrapper tools.
2
3use std::sync::Arc;
4
5use rmcp::handler::server::ServerHandler;
6use rmcp::model::{
7    Annotated, CallToolRequestParams, CallToolResult, Content, ErrorCode, GetPromptRequestParams,
8    GetPromptResult, InitializeResult, ListPromptsResult, ListResourcesResult, ListToolsResult,
9    PaginatedRequestParams, Prompt,
10    RawResource, ReadResourceRequestParams, ReadResourceResult, Resource, ResourceContents,
11    ServerCapabilities, Tool,
12};
13use rmcp::service::RequestContext;
14use rmcp::{ErrorData as McpError, RoleServer};
15use serde_json::{Map, Value};
16
17use crate::server::CompressedServer;
18
19/// Dynamic frontend MCP service that exposes compressed wrapper tools and
20/// delegates their calls to [`CompressedServer`].
21#[derive(Debug)]
22pub struct FrontendServer {
23    compressed: Arc<CompressedServer>,
24}
25
26impl FrontendServer {
27    pub fn new(compressed: CompressedServer) -> Self {
28        Self {
29            compressed: Arc::new(compressed),
30        }
31    }
32
33    pub fn from_arc(compressed: Arc<CompressedServer>) -> Self {
34        Self { compressed }
35    }
36}
37
38impl ServerHandler for FrontendServer {
39    fn get_info(&self) -> InitializeResult {
40        InitializeResult::new(
41            ServerCapabilities::builder()
42                .enable_tools()
43                .enable_resources()
44                .enable_prompts()
45                .build(),
46        )
47        .with_instructions("Compressed MCP frontend server")
48    }
49
50    async fn list_tools(
51        &self,
52        _request: Option<PaginatedRequestParams>,
53        _context: RequestContext<RoleServer>,
54    ) -> Result<ListToolsResult, McpError> {
55        let tools = self
56            .compressed
57            .list_frontend_tools()
58            .await
59            .map_err(mcp_error)?
60            .into_iter()
61            .map(convert_tool)
62            .collect();
63        Ok(ListToolsResult::with_all_items(tools))
64    }
65
66    async fn call_tool(
67        &self,
68        request: CallToolRequestParams,
69        _context: RequestContext<RoleServer>,
70    ) -> Result<CallToolResult, McpError> {
71        let wrapper_name = request.name.to_string();
72        let arguments = request.arguments.unwrap_or_default();
73        let output = if wrapper_name.ends_with("get_tool_schema") {
74            let tool_name = required_string(&arguments, "tool_name")?;
75            self.compressed
76                .get_tool_schema(&wrapper_name, &tool_name)
77                .await
78        } else if wrapper_name.ends_with("invoke_tool") {
79            let tool_name = required_string(&arguments, "tool_name")?;
80            let tool_input = arguments
81                .get("tool_input")
82                .cloned()
83                .unwrap_or_else(|| Value::Object(Map::new()));
84            self.compressed
85                .invoke_tool(&wrapper_name, &tool_name, tool_input)
86                .await
87        } else if wrapper_name.ends_with("list_tools") {
88            self.compressed.list_backend_tools(&wrapper_name).await
89        } else {
90            Err(crate::Error::ToolNotFound(wrapper_name))
91        }
92        .map_err(mcp_error)?;
93
94        Ok(CallToolResult::success(vec![Content::text(output)]))
95    }
96
97    async fn list_resources(
98        &self,
99        _request: Option<PaginatedRequestParams>,
100        _context: RequestContext<RoleServer>,
101    ) -> Result<ListResourcesResult, McpError> {
102        let resources = self
103            .compressed
104            .list_resources()
105            .await
106            .map_err(mcp_error)?
107            .into_iter()
108            .map(convert_resource)
109            .collect();
110        Ok(ListResourcesResult::with_all_items(resources))
111    }
112
113    async fn read_resource(
114        &self,
115        request: ReadResourceRequestParams,
116        _context: RequestContext<RoleServer>,
117    ) -> Result<ReadResourceResult, McpError> {
118        let text = self
119            .compressed
120            .read_resource(&request.uri)
121            .await
122            .map_err(mcp_error)?;
123        Ok(ReadResourceResult::new(vec![
124            ResourceContents::text(text, request.uri),
125        ]))
126    }
127
128    async fn list_prompts(
129        &self,
130        _request: Option<PaginatedRequestParams>,
131        _context: RequestContext<RoleServer>,
132    ) -> Result<ListPromptsResult, McpError> {
133        let prompts = self
134            .compressed
135            .list_prompts()
136            .await
137            .map_err(mcp_error)?
138            .into_iter()
139            .map(|name| Prompt::new(name, Option::<String>::None, None))
140            .collect();
141        Ok(ListPromptsResult::with_all_items(prompts))
142    }
143
144    async fn get_prompt(
145        &self,
146        request: GetPromptRequestParams,
147        _context: RequestContext<RoleServer>,
148    ) -> Result<GetPromptResult, McpError> {
149        self.compressed
150            .get_prompt(&request.name, request.arguments)
151            .await
152            .map_err(mcp_error)
153    }
154
155    fn get_tool(&self, _name: &str) -> Option<Tool> {
156        None
157    }
158}
159
160fn convert_tool(tool: crate::compression::engine::Tool) -> Tool {
161    let input_schema = match tool.input_schema {
162        Value::Object(map) => map,
163        _ => Map::new(),
164    };
165    Tool::new(
166        tool.name,
167        tool.description.unwrap_or_default(),
168        Arc::new(input_schema),
169    )
170}
171
172fn convert_resource(uri: String) -> Resource {
173    Annotated::new(RawResource {
174        name: uri.clone(),
175        uri,
176        title: None,
177        description: None,
178        mime_type: None,
179        icons: None,
180        size: None,
181        meta: None,
182    }, None)
183}
184
185fn required_string(arguments: &Map<String, Value>, name: &str) -> Result<String, McpError> {
186    arguments
187        .get(name)
188        .and_then(Value::as_str)
189        .map(str::to_string)
190        .ok_or_else(|| McpError::new(ErrorCode::INVALID_PARAMS, format!("missing {name}"), None))
191}
192
193fn mcp_error(error: crate::Error) -> McpError {
194    McpError::new(ErrorCode::INTERNAL_ERROR, error.to_string(), None)
195}