1use rmcp::handler::server::router::tool::ToolRouter;
2use rmcp::handler::server::wrapper::Parameters;
3use rmcp::model::{CallToolResult, ErrorData, ServerInfo};
4use rmcp::transport::io::stdio;
5use rmcp::{serve_server, tool, tool_handler, tool_router};
6use schemars::JsonSchema;
7use serde::Deserialize;
8use sui_crypto::ed25519::Ed25519PrivateKey;
9use sui_sdk_types::Address;
10
11use crate::{api, config, keystore, models};
12
13#[derive(JsonSchema, Deserialize)]
14#[allow(dead_code)]
15struct SearchParams {
16 q: String,
17 #[serde(default)]
18 owner: Option<String>,
19 #[serde(default)]
20 limit: Option<u32>,
21}
22
23#[derive(JsonSchema, Deserialize)]
24#[allow(dead_code)]
25struct ChunkParams {
26 q: String,
27 #[serde(default)]
28 owner: Option<String>,
29 #[serde(default)]
30 limit: Option<u32>,
31 #[serde(default)]
32 expand: Option<bool>,
33 #[serde(default)]
34 article_ids: Option<Vec<String>>,
35}
36
37#[derive(JsonSchema, Deserialize)]
38#[allow(dead_code)]
39struct ExpandParams {
40 chunk_ids: Vec<u64>,
42}
43
44#[allow(dead_code)]
45pub struct ZingMcpServer {
46 rpc_url: String,
47 api_base_url: String,
48 keypair: Ed25519PrivateKey,
49 sender: Address,
50 platform_usdc_address: Address,
51 tool_router: ToolRouter<Self>,
52}
53
54#[tool_router(router = tool_router)]
55impl ZingMcpServer {
56 pub async fn new(api_override: Option<String>) -> anyhow::Result<Self> {
57 let cfg = config::load_config()?;
58 let rpc_url = cfg.rpc_url;
59 let api_base_url = api_override.unwrap_or(cfg.api_base_url);
60 let sender = cfg.active_address;
61 let platform_usdc_address = cfg.platform_usdc_address;
62
63 let keypair = keystore::load_keypair(&cfg.keystore_path, &sender)?;
64
65 Ok(Self {
66 rpc_url,
67 api_base_url,
68 keypair,
69 sender,
70 platform_usdc_address,
71 tool_router: Self::tool_router(),
72 })
73 }
74
75 pub async fn serve(self) -> anyhow::Result<()> {
76 tracing::info!("Starting Zing MCP server on stdio");
77 let running = serve_server(self, stdio()).await?;
78 running.waiting().await?;
79 Ok(())
80 }
81
82 #[tool(
83 description = "Search the Zing decentralized knowledge base. Provide short keyword queries (2-4 words preferred). Returns articles with relevance scores, excerpts, tags, and budget info. If you already have specific article_ids from previous results, use zing_chunks with article_ids to retrieve their content directly instead of searching again. Default limit is 20."
84 )]
85 async fn zing_search(
86 &self,
87 Parameters(params): Parameters<SearchParams>,
88 ) -> Result<CallToolResult, ErrorData> {
89 tracing::info!("MCP zing_search q={}", params.q);
90
91 let wiki = "global".to_string();
92 let owner_param = params.owner.as_deref();
93 let limit = params.limit.unwrap_or(20).min(50);
94
95 let response = api::search(
96 &self.rpc_url,
97 &self.api_base_url,
98 &self.keypair,
99 &self.sender,
100 &self.platform_usdc_address,
101 ¶ms.q,
102 &wiki,
103 owner_param,
104 limit,
105 )
106 .await
107 .map_err(|e| {
108 tracing::error!("search failed: {e}");
109 ErrorData::internal_error(e.to_string(), None)
110 })?;
111
112 let agent_results: Vec<models::AgentSearchResult> = response
113 .results
114 .iter()
115 .map(|r| {
116 let excerpt = r.best_match.as_ref().map(|m| m.excerpt.clone());
117 let heading_path = r
118 .best_match
119 .as_ref()
120 .map(|m| m.heading_path.clone())
121 .unwrap_or_default();
122 models::AgentSearchResult {
123 article_id: r.article_id.clone(),
124 title: r.title.clone().unwrap_or_else(|| "Untitled".into()),
125 excerpt,
126 heading_path,
127 score: r.signals.relevance_score,
128 article_token_count: r.signals.article_token_count,
129 recency_days: r.signals.recency_days,
130 tags: r.tags.clone(),
131 }
132 })
133 .collect();
134
135 let agent_response = models::AgentSearchResponse {
136 results: agent_results,
137 budget: models::AgentBudget {
138 paid_usdc: response.budget.paid_usdc,
139 consumed_usdc: response.budget.consumed_usdc,
140 remaining_usdc: response.budget.remaining_usdc,
141 },
142 };
143
144 Ok(CallToolResult::structured(
145 serde_json::to_value(&agent_response)
146 .map_err(|e| ErrorData::internal_error(e.to_string(), None))?,
147 ))
148 }
149
150 #[tool(
151 description = "Retrieve raw text segments from search results with per-chunk pricing. Provide short keyword queries. Returns chunks with text, scores, content_type, and truncation metadata. Set expand=true (no extra cost) to return full text instead of excerpts. Use article_ids to filter to specific articles. When truncation metadata is present, call zing_expand_chunks with those chunk_ids to retrieve full text. Default limit is 20."
152 )]
153 async fn zing_chunks(
154 &self,
155 Parameters(params): Parameters<ChunkParams>,
156 ) -> Result<CallToolResult, ErrorData> {
157 tracing::info!("MCP zing_chunks q={} expand={:?}", params.q, params.expand);
158
159 let wiki = "global".to_string();
160 let owner_param = params.owner.as_deref();
161 let limit = params.limit.unwrap_or(20).min(50);
162
163 let response = api::chunks(
164 &self.rpc_url,
165 &self.api_base_url,
166 &self.keypair,
167 &self.sender,
168 &self.platform_usdc_address,
169 ¶ms.q,
170 &wiki,
171 owner_param,
172 limit,
173 params.expand,
174 params.article_ids.clone(),
175 )
176 .await
177 .map_err(|e| {
178 tracing::error!("chunks failed: {e}");
179 ErrorData::internal_error(e.to_string(), None)
180 })?;
181
182 let agent_chunks: Vec<models::AgentChunkResult> = response
183 .chunks
184 .iter()
185 .map(|c| models::AgentChunkResult {
186 chunk_id: c.chunk_id,
187 article_id: c.article_id.clone(),
188 title: c.title.clone(),
189 text: c.text.clone(),
190 score: c.scores.blended,
191 chunk_token_count: c.chunk_token_count,
192 heading_path: c.heading_path.clone(),
193 content_type: c.content_type.clone(),
194 language: c.language.clone(),
195 truncated: c.truncated.clone(),
196 })
197 .collect();
198
199 let agent_response = models::AgentChunksResponse {
200 chunks: agent_chunks,
201 budget: models::AgentBudget {
202 paid_usdc: response.budget.paid_usdc,
203 consumed_usdc: response.budget.consumed_usdc,
204 remaining_usdc: response.budget.remaining_usdc,
205 },
206 };
207
208 Ok(CallToolResult::structured(
209 serde_json::to_value(&agent_response)
210 .map_err(|e| ErrorData::internal_error(e.to_string(), None))?,
211 ))
212 }
213
214 #[tool(
215 description = "Expand truncated chunks to retrieve full untruncated text. Pass chunk_ids from chunks results that have non-null truncated fields. Max 20 chunk IDs per call."
216 )]
217 async fn zing_expand_chunks(
218 &self,
219 Parameters(params): Parameters<ExpandParams>,
220 ) -> Result<CallToolResult, ErrorData> {
221 tracing::info!("MCP zing_expand_chunks chunk_ids={:?}", params.chunk_ids);
222
223 let chunk_ids_i64: Vec<i64> = params.chunk_ids.iter().map(|&id| id as i64).collect();
224
225 let response = api::expand_chunks(
226 &self.rpc_url,
227 &self.api_base_url,
228 &self.keypair,
229 &self.sender,
230 &self.platform_usdc_address,
231 &chunk_ids_i64,
232 )
233 .await
234 .map_err(|e| {
235 tracing::error!("expand_chunks failed: {e}");
236 ErrorData::internal_error(e.to_string(), None)
237 })?;
238
239 let agent_chunks: Vec<models::AgentExpandedChunk> = response
240 .chunks
241 .iter()
242 .map(|c| models::AgentExpandedChunk {
243 chunk_id: c.chunk_id,
244 article_id: c.article_id.clone(),
245 heading_path: c.heading_path.clone(),
246 chunk_text: c.chunk_text.clone(),
247 content_type: c.content_type.clone(),
248 token_count: c.token_count,
249 truncated: c.truncated.clone(),
250 })
251 .collect();
252
253 let agent_response = models::AgentExpandResponse {
254 chunks: agent_chunks,
255 budget: models::AgentBudget {
256 paid_usdc: response.budget.paid_usdc,
257 consumed_usdc: response.budget.consumed_usdc,
258 remaining_usdc: response.budget.remaining_usdc,
259 },
260 };
261
262 Ok(CallToolResult::structured(
263 serde_json::to_value(&agent_response)
264 .map_err(|e| ErrorData::internal_error(e.to_string(), None))?,
265 ))
266 }
267}
268
269#[tool_handler(router = self.tool_router)]
270impl rmcp::ServerHandler for ZingMcpServer {
271 fn get_info(&self) -> ServerInfo {
272 let capabilities = rmcp::model::ServerCapabilities::builder()
273 .enable_tools()
274 .build();
275 rmcp::model::InitializeResult::new(capabilities).with_instructions(
276 "Search the Zing decentralized knowledge base. \
277 Use zing_search to find articles, zing_chunks to retrieve semantic chunks, \
278 and zing_expand_chunks to get full untruncated text for truncated chunks.",
279 )
280 }
281}