1use std::io::{BufRead, BufReader, Read, Write};
47use std::path::Path;
48
49use serde::{Deserialize, Serialize};
50use serde_json::{json, Map, Value};
51
52use crate::error::{MnemoError, Result};
53use crate::memory::{Memory, MemoryType, Metric, Scope};
54use crate::result_cache::{CachePutOpts, SemanticCachePutOpts, DEFAULT_SEMANTIC_THRESHOLD};
55use crate::store::{Mnemo, RecallRequest};
56
57const PROTOCOL_VERSION: &str = "2025-06-18";
59
60const PASSPHRASE_ENV: &str = "MNEMO_PASSPHRASE";
62
63#[doc(hidden)]
68pub fn __fuzz_parse_request_line(line: &str) -> std::result::Result<(), String> {
69 serde_json::from_str::<Request>(line)
70 .map(|_| ())
71 .map_err(|e| e.to_string())
72}
73
74#[derive(Debug, Deserialize)]
79struct Request {
80 #[serde(default)]
81 jsonrpc: String,
82 #[serde(default)]
83 id: Option<Value>,
84 method: String,
85 #[serde(default)]
86 params: Value,
87}
88
89#[derive(Debug, Serialize)]
91struct Response {
92 jsonrpc: &'static str,
93 id: Value,
94 #[serde(skip_serializing_if = "Option::is_none")]
95 result: Option<Value>,
96 #[serde(skip_serializing_if = "Option::is_none")]
97 error: Option<RpcError>,
98}
99
100#[derive(Debug, Serialize)]
102struct RpcError {
103 code: i32,
104 message: String,
105 #[serde(skip_serializing_if = "Option::is_none")]
106 data: Option<Value>,
107}
108
109impl RpcError {
110 fn new(code: i32, message: impl Into<String>) -> Self {
111 Self {
112 code,
113 message: message.into(),
114 data: None,
115 }
116 }
117}
118
119const PARSE_ERROR: i32 = -32700;
121const INVALID_REQUEST: i32 = -32600;
122const METHOD_NOT_FOUND: i32 = -32601;
123const INVALID_PARAMS: i32 = -32602;
124const INTERNAL_ERROR: i32 = -32603;
125
126pub fn serve_stdio(path: &Path) -> Result<()> {
139 serve_stdio_with_embedder(path, None)
140}
141
142pub fn serve_stdio_with_embedder(
152 path: &Path,
153 embed_spec: Option<(String, String, bool, std::time::Duration)>,
154) -> Result<()> {
155 let stdin = std::io::stdin();
156 let stdout = std::io::stdout();
157 serve_with_streams(path, stdin.lock(), stdout.lock(), embed_spec)
158}
159
160pub fn serve_with_streams<R: Read, W: Write>(
164 path: &Path,
165 reader: R,
166 mut writer: W,
167 embed_spec: Option<(String, String, bool, std::time::Duration)>,
168) -> Result<()> {
169 let passphrase = std::env::var(PASSPHRASE_ENV).map_err(|_| {
170 MnemoError::Invalid(format!(
171 "{PASSPHRASE_ENV} is required to serve — set it in the shell before starting"
172 ))
173 })?;
174
175 eprintln!(
176 "mnemo mcp: opening {} (protocol {PROTOCOL_VERSION})",
177 path.display()
178 );
179 let mut db = Mnemo::open(path, &passphrase)?;
180 if let Some((cmd, name, normalize, timeout)) = embed_spec {
181 let identity = crate::EmbedderIdentity {
182 name,
183 dimensions: db.dimensions(),
184 normalize,
185 };
186 let embedder = crate::EmbedCmd::new(&cmd, identity)?.with_timeout(timeout);
187 db.set_embedder(embedder)?;
188 eprintln!(
189 "mnemo mcp: embedder ready (text params accepted, timeout {}s)",
190 timeout.as_secs()
191 );
192 }
193 eprintln!("mnemo mcp: ready ({} live memories)", db.len());
194
195 let reader = BufReader::new(reader);
196 for line in reader.lines() {
197 let line = line.map_err(MnemoError::Io)?;
198 let trimmed = line.trim();
199 if trimmed.is_empty() {
200 continue;
201 }
202
203 let (id, out) = match serde_json::from_str::<Request>(trimmed) {
204 Ok(req) => {
205 let id = req.id.clone();
206 let is_notification = id.is_none();
207 let result = handle(&mut db, &req);
208 if is_notification {
209 continue;
211 }
212 (id.unwrap_or(Value::Null), result)
213 }
214 Err(e) => (
215 Value::Null,
216 Err(RpcError::new(PARSE_ERROR, format!("parse error: {e}"))),
217 ),
218 };
219
220 let resp = match out {
221 Ok(value) => Response {
222 jsonrpc: "2.0",
223 id,
224 result: Some(value),
225 error: None,
226 },
227 Err(err) => Response {
228 jsonrpc: "2.0",
229 id,
230 result: None,
231 error: Some(err),
232 },
233 };
234
235 let line =
236 serde_json::to_string(&resp).map_err(|e| MnemoError::Serialize(e.to_string()))?;
237 writeln!(writer, "{line}").map_err(MnemoError::Io)?;
238 writer.flush().map_err(MnemoError::Io)?;
239 }
240 if let Err(e) = db.flush() {
252 eprintln!("mnemo mcp: flush on shutdown failed: {e}");
253 }
254 eprintln!("mnemo mcp: stdin closed, shutting down");
255 Ok(())
256}
257
258fn handle(db: &mut Mnemo, req: &Request) -> std::result::Result<Value, RpcError> {
263 if req.jsonrpc != "2.0" {
264 return Err(RpcError::new(
265 INVALID_REQUEST,
266 format!("expected jsonrpc \"2.0\", got {:?}", req.jsonrpc),
267 ));
268 }
269 match req.method.as_str() {
270 "initialize" => Ok(initialize_result(&req.params)),
271 "notifications/initialized" | "notifications/cancelled" => Ok(Value::Null),
272 "ping" => Ok(json!({})),
273 "tools/list" => Ok(tools_list()),
274 "tools/call" => tools_call(db, &req.params),
275 other => Err(RpcError::new(
276 METHOD_NOT_FOUND,
277 format!("method '{other}' not implemented"),
278 )),
279 }
280}
281
282fn initialize_result(params: &Value) -> Value {
285 let client_version = params
286 .get("protocolVersion")
287 .and_then(|v| v.as_str())
288 .unwrap_or(PROTOCOL_VERSION)
289 .to_string();
290 json!({
291 "protocolVersion": client_version,
292 "capabilities": {
293 "tools": { "listChanged": false },
294 },
295 "serverInfo": {
296 "name": "mnemo",
297 "version": env!("CARGO_PKG_VERSION"),
298 },
299 })
300}
301
302fn tools_list() -> Value {
308 json!({
309 "tools": [
310 {
311 "name": "about",
312 "description": "Read-only briefing for this .mnemo file: manifest scaffold plus any memories tagged as onboarding. Run this first to learn how the file expects to be used.",
313 "inputSchema": {
314 "type": "object",
315 "properties": {},
316 "additionalProperties": false
317 }
318 },
319 {
320 "name": "remember",
321 "description": "Insert or overwrite a memory. Pass EITHER `vector` (explicit embedding) OR `text` (embedded by the server if it was started with --embed-cmd / --embed-name). Persists to disk before returning. Returns `{id}`.",
322 "inputSchema": {
323 "type": "object",
324 "required": ["content", "memory_type"],
325 "properties": {
326 "content": { "type": "string" },
327 "memory_type": { "type": "string", "enum": ["episodic", "semantic", "procedural", "working"] },
328 "vector": { "type": "array", "items": { "type": "number" }, "description": "Explicit embedding; mutually exclusive with `text`." },
329 "text": { "type": "string", "description": "Content to embed via the server's registered embedder; mutually exclusive with `vector`. Requires the server to have been started with --embed-cmd." },
330 "importance": { "type": "number", "minimum": 0, "maximum": 1 },
331 "agent_id": { "type": "string" },
332 "session_id": { "type": "string" },
333 "ttl_secs": { "type": "integer" },
334 "shared": { "type": "boolean" },
335 "metadata": { "type": "object" }
336 },
337 "additionalProperties": false
338 }
339 },
340 {
341 "name": "recall",
342 "description": "Score every live memory (or the ANN-narrowed candidate set) against a query and return the top-k. Pass EITHER `query` (explicit vector) OR `text` (embedded by the server; requires --embed-cmd at serve start). Does not track access on this handle.",
343 "inputSchema": {
344 "type": "object",
345 "properties": {
346 "query": { "type": "array", "items": { "type": "number" }, "description": "Explicit query vector; mutually exclusive with `text`." },
347 "text": { "type": "string", "description": "Query text; embedded by the server's registered embedder. Mutually exclusive with `query`." },
348 "top_k": { "type": "integer", "default": 10 },
349 "types": { "type": "array", "items": { "type": "string" } },
350 "agent": { "type": "string" },
351 "metric": { "type": "string", "enum": ["cosine", "l2", "dot"], "default": "cosine" }
352 },
353 "additionalProperties": false
354 }
355 },
356 {
357 "name": "forget",
358 "description": "Soft-delete a memory by ULID. Space is reclaimed on the next `mnemo compact`. Persists to disk before returning.",
359 "inputSchema": {
360 "type": "object",
361 "required": ["id"],
362 "properties": {
363 "id": { "type": "string", "description": "ULID of the memory to delete" }
364 },
365 "additionalProperties": false
366 }
367 },
368 {
369 "name": "list",
370 "description": "Return every live memory as JSON. Sorted by creation time, oldest first. Vectors are omitted by default to keep payloads small.",
371 "inputSchema": {
372 "type": "object",
373 "properties": {
374 "include_vectors": { "type": "boolean", "default": false }
375 },
376 "additionalProperties": false
377 }
378 },
379 {
380 "name": "snapshot_list",
381 "description": "List every restorable snapshot (one per committed flush). Each entry gives the transaction ID, creation timestamp, and memory count.",
382 "inputSchema": {
383 "type": "object",
384 "properties": {},
385 "additionalProperties": false
386 }
387 },
388 {
389 "name": "stats",
390 "description": "Summary statistics: memory count, deleted count, agents present, snapshot count, page cache occupancy.",
391 "inputSchema": {
392 "type": "object",
393 "properties": {},
394 "additionalProperties": false
395 }
396 },
397 {
398 "name": "cache_put",
399 "description": "Insert or overwrite an exact-key cache entry. Value is a UTF-8 string; use `content_type` to label the payload. Persists to disk before returning. Namespace and key are caller-chosen; the key is SHA-256 hashed by the engine.",
400 "inputSchema": {
401 "type": "object",
402 "required": ["namespace", "key", "value"],
403 "properties": {
404 "namespace": { "type": "string" },
405 "key": { "type": "string" },
406 "value": { "type": "string" },
407 "content_type": { "type": "string", "default": "text" },
408 "ttl_secs": { "type": "integer", "minimum": 0 }
409 },
410 "additionalProperties": false
411 }
412 },
413 {
414 "name": "cache_get",
415 "description": "Look up an exact-key cache entry. Returns `{value, content_type, hit}` on a hit, `{hit: false}` on miss (missing or TTL-expired). Bumps access stats catalog-only.",
416 "inputSchema": {
417 "type": "object",
418 "required": ["namespace", "key"],
419 "properties": {
420 "namespace": { "type": "string" },
421 "key": { "type": "string" }
422 },
423 "additionalProperties": false
424 }
425 },
426 {
427 "name": "cache_stats",
428 "description": "Cache statistics for a namespace (omit for the whole cache). Returns entries, bytes, hits, misses, hit_rate, evictions.",
429 "inputSchema": {
430 "type": "object",
431 "properties": {
432 "namespace": { "type": "string" }
433 },
434 "additionalProperties": false
435 }
436 },
437 {
438 "name": "cache_put_semantic",
439 "description": "Insert an entry with an embedding vector for semantic recall (Phase 10.2). Pass EITHER `vector` + `model` (explicit) OR `text` (embedded by the server; `model` defaults to the server embedder's identity name). Later cache_get_semantic calls under a different `model` won't see this entry (a different-model hit would be a semantics bug).",
440 "inputSchema": {
441 "type": "object",
442 "required": ["namespace", "key", "value"],
443 "properties": {
444 "namespace": { "type": "string" },
445 "key": { "type": "string" },
446 "vector": { "type": "array", "items": { "type": "number" }, "description": "Explicit embedding; mutually exclusive with `text`. Requires `model`." },
447 "text": { "type": "string", "description": "Content to embed via the server's registered embedder; mutually exclusive with `vector`." },
448 "value": { "type": "string" },
449 "model": { "type": "string", "description": "Required with `vector`; optional with `text` (defaults to the embedder's identity name)." },
450 "content_type": { "type": "string", "default": "text" },
451 "ttl_secs": { "type": "integer", "minimum": 0 }
452 },
453 "additionalProperties": false
454 }
455 },
456 {
457 "name": "cache_get_semantic",
458 "description": "Semantic recall over the cache: top-1 cosine over live vectored entries in the namespace whose `model` matches. Pass EITHER `query` + `model` OR `text` (which embeds via the server and defaults the model to the embedder's identity). Returns `{hit: true, value, similarity, ...}` iff similarity >= threshold, else `{hit: false}`. Threshold defaults to 0.97.",
459 "inputSchema": {
460 "type": "object",
461 "required": ["namespace"],
462 "properties": {
463 "namespace": { "type": "string" },
464 "query": { "type": "array", "items": { "type": "number" }, "description": "Explicit query vector; mutually exclusive with `text`. Requires `model`." },
465 "text": { "type": "string", "description": "Query text; embedded by the server. Mutually exclusive with `query`." },
466 "model": { "type": "string", "description": "Required with `query`; optional with `text`." },
467 "threshold": { "type": "number", "minimum": -1.0, "maximum": 1.0 }
468 },
469 "additionalProperties": false
470 }
471 }
472 ]
473 })
474}
475
476fn tools_call(db: &mut Mnemo, params: &Value) -> std::result::Result<Value, RpcError> {
481 let name = params
482 .get("name")
483 .and_then(|v| v.as_str())
484 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "missing tool name"))?;
485 let empty = json!({});
486 let args = params.get("arguments").unwrap_or(&empty);
487
488 let text = match name {
489 "about" => tool_about(db)?,
490 "remember" => tool_remember(db, args)?,
491 "recall" => tool_recall(db, args)?,
492 "forget" => tool_forget(db, args)?,
493 "list" => tool_list(db, args)?,
494 "snapshot_list" => tool_snapshot_list(db)?,
495 "stats" => tool_stats(db)?,
496 "cache_put" => tool_cache_put(db, args)?,
497 "cache_get" => tool_cache_get(db, args)?,
498 "cache_stats" => tool_cache_stats(db, args)?,
499 "cache_put_semantic" => tool_cache_put_semantic(db, args)?,
500 "cache_get_semantic" => tool_cache_get_semantic(db, args)?,
501 other => {
502 return Err(RpcError::new(
503 METHOD_NOT_FOUND,
504 format!("unknown tool '{other}'"),
505 ))
506 }
507 };
508
509 Ok(json!({
510 "content": [ { "type": "text", "text": text } ],
511 "isError": false,
512 }))
513}
514
515fn embed_text_via_server(
526 db: &mut Mnemo,
527 tool: &str,
528 text: &str,
529) -> std::result::Result<Vec<f32>, RpcError> {
530 let embedder = db.embedder().ok_or_else(|| {
531 RpcError::new(
532 INVALID_PARAMS,
533 format!(
534 "{tool}: 'text' requires an embedder on the server. \
535 Start with: mnemo serve --mcp <file> --embed-cmd '<cmd>' --embed-name <name>"
536 ),
537 )
538 })?;
539 let mut vs = embedder.embed(&[text]).map_err(mnemo_err)?;
540 vs.pop().ok_or_else(|| {
541 RpcError::new(
542 INTERNAL_ERROR,
543 format!("{tool}: embedder returned no vector for the input text"),
544 )
545 })
546}
547
548fn mnemo_err(e: MnemoError) -> RpcError {
549 RpcError::new(INTERNAL_ERROR, e.to_string())
550}
551
552fn tool_about(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
553 let mems = db.about().map_err(mnemo_err)?;
554 let arr: Vec<Value> = mems.iter().map(memory_to_json).collect();
555 Ok(serde_json::to_string_pretty(&arr).unwrap_or_else(|_| "[]".into()))
556}
557
558fn tool_remember(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
559 let content = args
560 .get("content")
561 .and_then(|v| v.as_str())
562 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "remember: missing 'content'"))?;
563 let mtype_str = args
564 .get("memory_type")
565 .and_then(|v| v.as_str())
566 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "remember: missing 'memory_type'"))?;
567 let mtype = MemoryType::parse(mtype_str).ok_or_else(|| {
568 RpcError::new(
569 INVALID_PARAMS,
570 format!("remember: unknown memory_type '{mtype_str}'"),
571 )
572 })?;
573 let vector = match (args.get("vector"), args.get("text")) {
576 (Some(_), Some(_)) => {
577 return Err(RpcError::new(
578 INVALID_PARAMS,
579 "remember: pass EITHER 'vector' or 'text', not both",
580 ));
581 }
582 (Some(v), None) => v
583 .as_array()
584 .ok_or_else(|| {
585 RpcError::new(INVALID_PARAMS, "remember: 'vector' must be a JSON array")
586 })?
587 .iter()
588 .map(|v| v.as_f64().map(|f| f as f32))
589 .collect::<Option<Vec<f32>>>()
590 .ok_or_else(|| {
591 RpcError::new(
592 INVALID_PARAMS,
593 "remember: 'vector' must be an array of numbers",
594 )
595 })?,
596 (None, Some(t)) => {
597 let text = t.as_str().ok_or_else(|| {
598 RpcError::new(INVALID_PARAMS, "remember: 'text' must be a string")
599 })?;
600 embed_text_via_server(db, "remember", text)?
601 }
602 (None, None) => {
603 return Err(RpcError::new(
604 INVALID_PARAMS,
605 "remember: missing 'vector' or 'text'",
606 ))
607 }
608 };
609
610 let mut mem = Memory::new(content, mtype, vector);
611 if let Some(agent) = args.get("agent_id").and_then(|v| v.as_str()) {
612 mem = mem.with_agent(agent);
613 }
614 if let Some(sess) = args.get("session_id").and_then(|v| v.as_str()) {
615 mem = mem.with_session(sess);
616 }
617 if let Some(imp) = args.get("importance").and_then(|v| v.as_f64()) {
618 mem = mem.with_importance(imp as f32);
619 }
620 if let Some(ttl) = args.get("ttl_secs").and_then(|v| v.as_i64()) {
621 mem = mem.with_ttl(ttl);
622 }
623 if let Some(shared) = args.get("shared").and_then(|v| v.as_bool()) {
624 if shared {
625 mem = mem.with_scope(Scope::Shared);
626 }
627 }
628 if let Some(meta) = args.get("metadata").and_then(|v| v.as_object()) {
629 for (k, v) in meta {
630 mem.metadata.insert(k.clone(), v.clone());
631 }
632 }
633
634 let id = db.remember(mem).map_err(mnemo_err)?;
635 db.flush().map_err(mnemo_err)?;
636 Ok(json!({ "id": id.to_string() }).to_string())
637}
638
639fn tool_recall(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
640 let query = match (args.get("query"), args.get("text")) {
642 (Some(_), Some(_)) => {
643 return Err(RpcError::new(
644 INVALID_PARAMS,
645 "recall: pass EITHER 'query' or 'text', not both",
646 ))
647 }
648 (Some(v), None) => v
649 .as_array()
650 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "recall: 'query' must be a JSON array"))?
651 .iter()
652 .map(|v| v.as_f64().map(|f| f as f32))
653 .collect::<Option<Vec<f32>>>()
654 .ok_or_else(|| {
655 RpcError::new(
656 INVALID_PARAMS,
657 "recall: 'query' must be an array of numbers",
658 )
659 })?,
660 (None, Some(t)) => {
661 let text = t
662 .as_str()
663 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "recall: 'text' must be a string"))?;
664 embed_text_via_server(db, "recall", text)?
665 }
666 (None, None) => {
667 return Err(RpcError::new(
668 INVALID_PARAMS,
669 "recall: missing 'query' or 'text'",
670 ))
671 }
672 };
673 let top_k = args.get("top_k").and_then(|v| v.as_u64()).unwrap_or(10) as usize;
674 let metric = match args
675 .get("metric")
676 .and_then(|v| v.as_str())
677 .unwrap_or("cosine")
678 {
679 "cosine" => Metric::Cosine,
680 "l2" => Metric::L2,
681 "dot" => Metric::Dot,
682 other => {
683 return Err(RpcError::new(
684 INVALID_PARAMS,
685 format!("recall: unknown metric '{other}'"),
686 ))
687 }
688 };
689 let mut req = RecallRequest::new(query)
702 .top_k(top_k)
703 .metric(metric)
704 .track_access(true);
705 if let Some(types) = args.get("types").and_then(|v| v.as_array()) {
706 let parsed: Vec<MemoryType> = types
707 .iter()
708 .filter_map(|t| t.as_str().and_then(MemoryType::parse))
709 .collect();
710 if !parsed.is_empty() {
711 req = req.types(parsed);
712 }
713 }
714 if let Some(agent) = args.get("agent").and_then(|v| v.as_str()) {
715 req = req.agent(agent.to_string());
716 }
717
718 let hits = db.recall(&req).map_err(mnemo_err)?;
719 let out: Vec<Value> = hits
720 .into_iter()
721 .map(|h| {
722 let mut obj = memory_to_json(&h.memory);
723 if let Some(map) = obj.as_object_mut() {
724 map.insert("score".into(), json!(h.score));
725 map.insert("similarity".into(), json!(h.similarity));
726 }
727 obj
728 })
729 .collect();
730 Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
731}
732
733fn tool_forget(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
734 let id_str = args
735 .get("id")
736 .and_then(|v| v.as_str())
737 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "forget: missing 'id'"))?;
738 let ulid = crate::Ulid::from_string(id_str)
739 .map_err(|_| RpcError::new(INVALID_PARAMS, format!("forget: invalid ULID '{id_str}'")))?;
740 db.delete(&ulid).map_err(mnemo_err)?;
741 db.flush().map_err(mnemo_err)?;
742 Ok(json!({ "deleted": id_str }).to_string())
743}
744
745fn tool_list(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
746 let include_vectors = args
747 .get("include_vectors")
748 .and_then(|v| v.as_bool())
749 .unwrap_or(false);
750 let mems = db.memories().map_err(mnemo_err)?;
751 let out: Vec<Value> = mems
752 .iter()
753 .map(|m| {
754 let mut v = memory_to_json(m);
755 if !include_vectors {
756 if let Some(o) = v.as_object_mut() {
757 o.remove("vector");
758 }
759 }
760 v
761 })
762 .collect();
763 Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
764}
765
766fn tool_snapshot_list(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
767 let snaps = db.snapshots();
768 let out: Vec<Value> = snaps
769 .into_iter()
770 .map(|s| {
771 json!({
772 "txn_id": s.txn_id,
773 "created_at": s.created_at,
774 "memory_count": s.memory_count,
775 })
776 })
777 .collect();
778 Ok(serde_json::to_string_pretty(&out).unwrap_or_else(|_| "[]".into()))
779}
780
781fn tool_stats(db: &mut Mnemo) -> std::result::Result<String, RpcError> {
782 let s = db.stats().map_err(mnemo_err)?;
783 let (cache_pages, cache_capacity) = db.page_cache_stats();
784 let snapshot_count = db.snapshots().len();
785 Ok(json!({
788 "memories": s.memories,
789 "deleted": s.deleted,
790 "dimensions": s.dimensions,
791 "file_bytes": s.file_bytes,
792 "agents": s.agents,
793 "encrypted": s.encrypted,
794 "created_at": s.created_at,
795 "wal_pages": s.wal_pages,
796 "index": s.index,
797 "snapshot_count": snapshot_count,
798 "cache_pages": cache_pages,
799 "cache_capacity": cache_capacity,
800 })
801 .to_string())
802}
803
804fn tool_cache_put(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
807 let ns = args
808 .get("namespace")
809 .and_then(|v| v.as_str())
810 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'namespace'"))?;
811 let key = args
812 .get("key")
813 .and_then(|v| v.as_str())
814 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'key'"))?;
815 let value = args
816 .get("value")
817 .and_then(|v| v.as_str())
818 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put: missing 'value'"))?;
819 let content_type = args
820 .get("content_type")
821 .and_then(|v| v.as_str())
822 .unwrap_or("text")
823 .to_string();
824 let ttl_secs = args.get("ttl_secs").and_then(|v| v.as_u64());
825 let opts = CachePutOpts {
826 content_type,
827 ttl_secs,
828 cost_hint_ms: None,
829 };
830 db.cache_put(ns, key, value.as_bytes(), opts)
831 .map_err(mnemo_err)?;
832 db.flush().map_err(mnemo_err)?;
833 Ok(json!({ "cached": true, "namespace": ns, "key": key, "bytes": value.len() }).to_string())
834}
835
836fn tool_cache_get(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
837 let ns = args
838 .get("namespace")
839 .and_then(|v| v.as_str())
840 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get: missing 'namespace'"))?;
841 let key = args
842 .get("key")
843 .and_then(|v| v.as_str())
844 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get: missing 'key'"))?;
845 match db.cache_get(ns, key).map_err(mnemo_err)? {
846 Some(hit) => Ok(json!({
847 "hit": true,
848 "value": String::from_utf8_lossy(&hit.value),
849 "content_type": hit.content_type,
850 "created_at": hit.created_at,
851 "accessed_at": hit.accessed_at,
852 "access_count": hit.access_count,
853 })
854 .to_string()),
855 None => Ok(json!({ "hit": false }).to_string()),
856 }
857}
858
859fn tool_cache_stats(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
860 let ns = args.get("namespace").and_then(|v| v.as_str());
861 let s = db.cache_stats(ns);
862 Ok(json!({
863 "namespace": ns,
864 "entries": s.entries,
865 "bytes": s.bytes,
866 "hits": s.hits,
867 "misses": s.misses,
868 "hit_rate": s.hit_rate,
869 "evictions": s.evictions,
870 })
871 .to_string())
872}
873
874fn tool_cache_put_semantic(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
875 let ns = args
876 .get("namespace")
877 .and_then(|v| v.as_str())
878 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'namespace'"))?;
879 let key = args
880 .get("key")
881 .and_then(|v| v.as_str())
882 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'key'"))?;
883 let value = args
884 .get("value")
885 .and_then(|v| v.as_str())
886 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'value'"))?;
887 let (vector, model) = match (args.get("vector"), args.get("text")) {
890 (Some(_), Some(_)) => {
891 return Err(RpcError::new(
892 INVALID_PARAMS,
893 "cache_put_semantic: pass EITHER 'vector' or 'text', not both",
894 ));
895 }
896 (Some(v), None) => {
897 let vec = v
898 .as_array()
899 .ok_or_else(|| {
900 RpcError::new(
901 INVALID_PARAMS,
902 "cache_put_semantic: 'vector' must be a JSON array",
903 )
904 })?
905 .iter()
906 .map(|v| v.as_f64().map(|f| f as f32))
907 .collect::<Option<Vec<f32>>>()
908 .ok_or_else(|| {
909 RpcError::new(
910 INVALID_PARAMS,
911 "cache_put_semantic: 'vector' must be an array of numbers",
912 )
913 })?;
914 let model = args
915 .get("model")
916 .and_then(|v| v.as_str())
917 .ok_or_else(|| {
918 RpcError::new(INVALID_PARAMS, "cache_put_semantic: missing 'model'")
919 })?
920 .to_string();
921 (vec, model)
922 }
923 (None, Some(t)) => {
924 let text = t.as_str().ok_or_else(|| {
925 RpcError::new(
926 INVALID_PARAMS,
927 "cache_put_semantic: 'text' must be a string",
928 )
929 })?;
930 let vec = embed_text_via_server(db, "cache_put_semantic", text)?;
931 let model = match args.get("model").and_then(|v| v.as_str()) {
933 Some(m) => m.to_string(),
934 None => {
935 db.embedder()
936 .expect("embed_text_via_server would have errored otherwise")
937 .identity()
938 .name
939 }
940 };
941 (vec, model)
942 }
943 (None, None) => {
944 return Err(RpcError::new(
945 INVALID_PARAMS,
946 "cache_put_semantic: missing 'vector' or 'text'",
947 ))
948 }
949 };
950 let content_type = args
951 .get("content_type")
952 .and_then(|v| v.as_str())
953 .unwrap_or("text")
954 .to_string();
955 let ttl_secs = args.get("ttl_secs").and_then(|v| v.as_u64());
956 let opts = SemanticCachePutOpts {
957 content_type,
958 ttl_secs,
959 cost_hint_ms: None,
960 model,
961 };
962 db.cache_put_semantic(ns, key, vector, value.as_bytes(), opts)
963 .map_err(mnemo_err)?;
964 db.flush().map_err(mnemo_err)?;
965 Ok(json!({ "cached": true, "namespace": ns, "key": key, "bytes": value.len() }).to_string())
966}
967
968fn tool_cache_get_semantic(db: &mut Mnemo, args: &Value) -> std::result::Result<String, RpcError> {
969 let ns = args
970 .get("namespace")
971 .and_then(|v| v.as_str())
972 .ok_or_else(|| RpcError::new(INVALID_PARAMS, "cache_get_semantic: missing 'namespace'"))?;
973 let (query, model_owned) = match (args.get("query"), args.get("text")) {
976 (Some(_), Some(_)) => {
977 return Err(RpcError::new(
978 INVALID_PARAMS,
979 "cache_get_semantic: pass EITHER 'query' or 'text', not both",
980 ));
981 }
982 (Some(q), None) => {
983 let vec = q
984 .as_array()
985 .ok_or_else(|| {
986 RpcError::new(
987 INVALID_PARAMS,
988 "cache_get_semantic: 'query' must be a JSON array",
989 )
990 })?
991 .iter()
992 .map(|v| v.as_f64().map(|f| f as f32))
993 .collect::<Option<Vec<f32>>>()
994 .ok_or_else(|| {
995 RpcError::new(
996 INVALID_PARAMS,
997 "cache_get_semantic: 'query' must be an array of numbers",
998 )
999 })?;
1000 let model = args
1001 .get("model")
1002 .and_then(|v| v.as_str())
1003 .ok_or_else(|| {
1004 RpcError::new(INVALID_PARAMS, "cache_get_semantic: missing 'model'")
1005 })?
1006 .to_string();
1007 (vec, model)
1008 }
1009 (None, Some(t)) => {
1010 let text = t.as_str().ok_or_else(|| {
1011 RpcError::new(
1012 INVALID_PARAMS,
1013 "cache_get_semantic: 'text' must be a string",
1014 )
1015 })?;
1016 let vec = embed_text_via_server(db, "cache_get_semantic", text)?;
1017 let model = match args.get("model").and_then(|v| v.as_str()) {
1018 Some(m) => m.to_string(),
1019 None => {
1020 db.embedder()
1021 .expect("embed_text_via_server would have errored otherwise")
1022 .identity()
1023 .name
1024 }
1025 };
1026 (vec, model)
1027 }
1028 (None, None) => {
1029 return Err(RpcError::new(
1030 INVALID_PARAMS,
1031 "cache_get_semantic: missing 'query' or 'text'",
1032 ))
1033 }
1034 };
1035 let model = model_owned.as_str();
1036 let threshold = args
1037 .get("threshold")
1038 .and_then(|v| v.as_f64())
1039 .map(|f| f as f32)
1040 .unwrap_or(DEFAULT_SEMANTIC_THRESHOLD);
1041 match db
1042 .cache_get_semantic(ns, &query, threshold, model)
1043 .map_err(mnemo_err)?
1044 {
1045 Some((hit, similarity)) => Ok(json!({
1046 "hit": true,
1047 "similarity": similarity,
1048 "value": String::from_utf8_lossy(&hit.value),
1049 "content_type": hit.content_type,
1050 "created_at": hit.created_at,
1051 "accessed_at": hit.accessed_at,
1052 "access_count": hit.access_count,
1053 })
1054 .to_string()),
1055 None => Ok(json!({ "hit": false }).to_string()),
1056 }
1057}
1058
1059fn memory_to_json(m: &Memory) -> Value {
1064 serde_json::to_value(m).unwrap_or_else(|_| {
1065 let mut map = Map::new();
1066 map.insert("id".into(), json!(m.id.to_string()));
1067 Value::Object(map)
1068 })
1069}