Skip to main content

velesdb_memory/
mcp.rs

1//! MCP transport: exposes the memory service as MCP tools over stdio.
2//!
3//! Only **memory semantics** are exposed (`remember / recall / relate / forget
4//! / why`) — never raw database capabilities. See [`crate`] docs for the
5//! license boundary this enforces.
6
7use std::sync::Arc;
8
9use rmcp::handler::server::router::tool::ToolRouter;
10use rmcp::handler::server::wrapper::{Json, Parameters};
11use rmcp::model::{ErrorCode, Implementation, ServerCapabilities, ServerInfo};
12use rmcp::{tool, tool_handler, tool_router, ErrorData, ServerHandler};
13
14use crate::limits::{DEFAULT_WHY_HOPS, MAX_FACT_BYTES, MAX_RECALL_LIMIT, MAX_WHY_HOPS};
15use crate::model::Explanation;
16use crate::service::MemoryService;
17
18/// Default number of memories returned by `recall`.
19const DEFAULT_RECALL_LIMIT: usize = 10;
20
21// The boxed embedder and the shared, runtime-attached extraction backend the
22// server stores — imported for internal use only. The canonical public paths are
23// `velesdb_memory::DynEmbedder` / `velesdb_memory::DynExtractor` (crate root).
24use crate::embedder::DynEmbedder;
25use crate::extract::DynExtractor;
26
27// --- Tool parameter / result DTOs ------------------------------------------
28//
29// The request envelopes and small id-results live in their own module so this
30// file stays focused on the server and tool wiring; output shapes reuse the
31// domain types from `crate::model` directly (no duplicate wire/domain struct).
32mod dto;
33use dto::{
34    ForgetParams, ForgetResult, RecallParams, RecallResult, RecallWhereParams, RelateParams,
35    RelateResult, RememberExtractedParams, RememberExtractedResult, RememberParams, RememberResult,
36    WhyParams,
37};
38
39// --- The server ------------------------------------------------------------
40
41/// MCP server wrapping a [`MemoryService`].
42#[derive(Clone)]
43pub struct McpServer {
44    service: Arc<MemoryService<DynEmbedder>>,
45    /// Optional extraction backend powering `remember_extracted`. `None` unless
46    /// a backend is attached via [`Self::with_extractor`]; the tool then reports
47    /// extraction as unconfigured.
48    extractor: Option<DynExtractor>,
49    tool_router: ToolRouter<McpServer>,
50}
51
52#[tool_router]
53impl McpServer {
54    /// Wrap a memory service as an MCP server.
55    #[must_use]
56    pub fn new(service: MemoryService<DynEmbedder>) -> Self {
57        Self {
58            service: Arc::new(service),
59            extractor: None,
60            tool_router: Self::tool_router(),
61        }
62    }
63
64    /// Attach an extraction backend, enabling the `remember_extracted` tool.
65    /// Without it the tool reports that extraction is not configured.
66    #[must_use]
67    pub fn with_extractor(mut self, extractor: DynExtractor) -> Self {
68        self.extractor = Some(extractor);
69        self
70    }
71
72    #[tool(
73        name = "remember",
74        description = "Store a fact in durable local memory. Optionally link it to existing memories (graph) and tag it with structured metadata like project/author/type/status/date (ColumnStore) for later filtering. Returns the fact's stable id."
75    )]
76    async fn remember(
77        &self,
78        Parameters(params): Parameters<RememberParams>,
79    ) -> Result<Json<RememberResult>, ErrorData> {
80        if params.fact.len() > MAX_FACT_BYTES {
81            return Err(ErrorData::new(
82                ErrorCode::INVALID_PARAMS,
83                format!("fact exceeds maximum size of {MAX_FACT_BYTES} bytes"),
84                None,
85            ));
86        }
87        let service = Arc::clone(&self.service);
88        let RememberParams {
89            fact,
90            links,
91            metadata,
92        } = params;
93        let id =
94            tokio::task::spawn_blocking(move || service.remember(&fact, &links, metadata.as_ref()))
95                .await
96                .map_err(join_error)?
97                .map_err(to_error)?;
98        Ok(Json(RememberResult { id }))
99    }
100
101    #[tool(
102        name = "recall",
103        description = "Recall memories semantically similar to a query (vector), most similar first. Optionally narrow to exact-match metadata via `filter` (ColumnStore), e.g. {\"project\":\"veles\",\"status\":\"resolved\"}."
104    )]
105    async fn recall(
106        &self,
107        Parameters(params): Parameters<RecallParams>,
108    ) -> Result<Json<RecallResult>, ErrorData> {
109        let limit = params
110            .limit
111            .unwrap_or(DEFAULT_RECALL_LIMIT)
112            .min(MAX_RECALL_LIMIT);
113        let service = Arc::clone(&self.service);
114        let RecallParams { query, filter, .. } = params;
115        let memories =
116            tokio::task::spawn_blocking(move || service.recall(&query, limit, filter.as_ref()))
117                .await
118                .map_err(join_error)?
119                .map_err(to_error)?;
120        Ok(Json(RecallResult { memories }))
121    }
122
123    #[tool(
124        name = "recall_where",
125        description = "Fused recall: semantically similar memories (vector) constrained by structured ColumnStore predicates over metadata — ranges and comparisons, not just equality. Each filter is {field, op (eq/ne/lt/le/gt/ge), value}, ANDed. Use for time-windowed or numeric-scoped recall, e.g. facts about a topic with `ts` in a date range. Most similar first."
126    )]
127    async fn recall_where(
128        &self,
129        Parameters(params): Parameters<RecallWhereParams>,
130    ) -> Result<Json<RecallResult>, ErrorData> {
131        let limit = params
132            .limit
133            .unwrap_or(DEFAULT_RECALL_LIMIT)
134            .min(MAX_RECALL_LIMIT);
135        let service = Arc::clone(&self.service);
136        let RecallWhereParams { query, filters, .. } = params;
137        let memories =
138            tokio::task::spawn_blocking(move || service.recall_where(&query, limit, &filters))
139                .await
140                .map_err(join_error)?
141                .map_err(to_error)?;
142        Ok(Json(RecallResult { memories }))
143    }
144
145    #[tool(
146        name = "relate",
147        description = "Create a typed link from one memory to another. Returns the edge id."
148    )]
149    async fn relate(
150        &self,
151        Parameters(params): Parameters<RelateParams>,
152    ) -> Result<Json<RelateResult>, ErrorData> {
153        let service = Arc::clone(&self.service);
154        let RelateParams { from, to, relation } = params;
155        let edge_id = tokio::task::spawn_blocking(move || service.relate(from, to, &relation))
156            .await
157            .map_err(join_error)?
158            .map_err(to_error)?;
159        Ok(Json(RelateResult { edge_id }))
160    }
161
162    #[tool(name = "forget", description = "Delete a memory by id.")]
163    async fn forget(
164        &self,
165        Parameters(params): Parameters<ForgetParams>,
166    ) -> Result<Json<ForgetResult>, ErrorData> {
167        let service = Arc::clone(&self.service);
168        let id = params.id;
169        tokio::task::spawn_blocking(move || service.forget(id))
170            .await
171            .map_err(join_error)?
172            .map_err(to_error)?;
173        Ok(Json(ForgetResult { id }))
174    }
175
176    #[tool(
177        name = "why",
178        description = "Explain a decision: find the best-matching memory (optionally scoped by a metadata `filter`, e.g. the current project) and return the connected subgraph of related memories reachable through typed links — fusing vector, ColumnStore, and graph to surface context a plain similarity search misses."
179    )]
180    async fn why(
181        &self,
182        Parameters(params): Parameters<WhyParams>,
183    ) -> Result<Json<Explanation>, ErrorData> {
184        let max_hops = params
185            .max_hops
186            .unwrap_or(DEFAULT_WHY_HOPS)
187            .min(MAX_WHY_HOPS);
188        let service = Arc::clone(&self.service);
189        let WhyParams {
190            decision, filter, ..
191        } = params;
192        let explanation =
193            tokio::task::spawn_blocking(move || service.why(&decision, max_hops, filter.as_ref()))
194                .await
195                .map_err(join_error)?
196                .map_err(to_error)?;
197        Ok(Json(explanation))
198    }
199
200    #[tool(
201        name = "remember_extracted",
202        description = "Store a passage of raw text by extracting its atomic facts and auto-building the fact↔topic graph, so `why` can later connect them with no manual links. Requires the server to be started with an extraction backend (set VELESDB_MEMORY_EXTRACTOR; build with --features extract). Returns the stored facts' ids."
203    )]
204    async fn remember_extracted(
205        &self,
206        Parameters(params): Parameters<RememberExtractedParams>,
207    ) -> Result<Json<RememberExtractedResult>, ErrorData> {
208        if params.text.len() > MAX_FACT_BYTES {
209            return Err(ErrorData::new(
210                ErrorCode::INVALID_PARAMS,
211                format!("text exceeds maximum size of {MAX_FACT_BYTES} bytes"),
212                None,
213            ));
214        }
215        let Some(extractor) = self.extractor.clone() else {
216            return Err(ErrorData::new(
217                ErrorCode::INTERNAL_ERROR,
218                "extraction backend not configured: start the server with \
219                 VELESDB_MEMORY_EXTRACTOR set (built with --features extract)",
220                None,
221            ));
222        };
223        // Extraction makes a blocking network call (up to the extractor's
224        // timeout), so run it off the async worker pool to keep the stdio loop
225        // responsive to other tool calls and cancellations.
226        let service = Arc::clone(&self.service);
227        let RememberExtractedParams { text, metadata } = params;
228        let ids = tokio::task::spawn_blocking(move || {
229            service.remember_extracted(&text, &extractor, metadata.as_ref())
230        })
231        .await
232        .map_err(join_error)?
233        .map_err(to_error)?;
234        Ok(Json(RememberExtractedResult { ids }))
235    }
236}
237
238/// `#[tool_handler]` generates `call_tool` / `list_tools` from the router;
239/// `get_info` is overridden so the server identifies itself as `velesdb-memory`
240/// (the macro default falls back to rmcp's own identity). Per-tool guidance
241/// lives in each `#[tool(description = …)]`.
242#[tool_handler(router = self.tool_router)]
243impl ServerHandler for McpServer {
244    fn get_info(&self) -> ServerInfo {
245        let mut info = ServerInfo::default();
246        info.server_info = Implementation::new(env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
247        info.capabilities = ServerCapabilities::builder().enable_tools().build();
248        info.instructions = Some(
249            "Local-first memory for AI agents: remember facts, recall them semantically, \
250             relate them, forget them, and ask why a decision was made (connected subgraph)."
251                .to_owned(),
252        );
253        info
254    }
255}
256
257/// Map a `spawn_blocking` join failure (a panicked or cancelled tool task) to an
258/// MCP error. Every tool body runs on the blocking pool, so they all funnel
259/// through this on the (rare) task-failure path.
260///
261/// Takes the error by value so it can be used as `.map_err(join_error)`.
262#[allow(clippy::needless_pass_by_value)]
263fn join_error(join: tokio::task::JoinError) -> ErrorData {
264    ErrorData::new(
265        ErrorCode::INTERNAL_ERROR,
266        format!("memory task failed: {join}"),
267        None,
268    )
269}
270
271/// Map a domain error to an MCP error.
272///
273/// Map a [`MemoryError`](crate::error::MemoryError) onto a JSON-RPC error,
274/// driven by its transport-neutral [`ErrorCategory`](crate::error::ErrorCategory)
275/// so the MCP taxonomy can never drift from the bindings'. Client-input errors
276/// become `invalid_params` (-32602); genuine faults `internal_error` (-32603).
277/// JSON-RPC defines no "not found" code, so a missing id is reported as
278/// `invalid_params` (a bad id is, from the protocol's view, a bad parameter).
279///
280/// Takes the error by value so it can be used as `.map_err(to_error)` at every
281/// call site without a per-site closure.
282#[allow(clippy::needless_pass_by_value)]
283fn to_error(err: crate::error::MemoryError) -> ErrorData {
284    use crate::error::ErrorCategory;
285    let code = match err.category() {
286        ErrorCategory::InvalidInput | ErrorCategory::NotFound => ErrorCode::INVALID_PARAMS,
287        ErrorCategory::Internal => ErrorCode::INTERNAL_ERROR,
288    };
289    ErrorData::new(code, err.to_string(), None)
290}
291
292#[cfg(test)]
293#[path = "mcp/server_tests.rs"]
294mod tests;