1use std::sync::Arc;
18
19use tonic::{Request, Response, Status};
20use uuid::Uuid;
21
22use mnemo_core::model::acl::Permission;
23use mnemo_core::model::delegation::{Delegation, DelegationScope};
24use mnemo_core::model::memory::{MemoryType, Scope, SourceType};
25use mnemo_core::query::MnemoEngine;
26use mnemo_core::query::branch::BranchRequest as CoreBranchRequest;
27use mnemo_core::query::checkpoint::CheckpointRequest as CoreCheckpointRequest;
28use mnemo_core::query::forget::{
29 ForgetRequest as CoreForgetRequest, ForgetStrategy,
30 ForgetSubjectRequest as CoreForgetSubjectRequest,
31};
32use mnemo_core::query::merge::{MergeRequest as CoreMergeRequest, MergeStrategy};
33use mnemo_core::query::recall::RecallRequest as CoreRecallRequest;
34use mnemo_core::query::remember::RememberRequest as CoreRememberRequest;
35use mnemo_core::query::replay::ReplayRequest as CoreReplayRequest;
36use mnemo_core::query::share::ShareRequest as CoreShareRequest;
37
38pub mod proto {
43 tonic::include_proto!("mnemo.v1");
44}
45
46use proto::mnemo_service_server::{MnemoService, MnemoServiceServer};
47use proto::{
48 BranchRequest as ProtoBranchRequest, BranchResponse as ProtoBranchResponse,
49 CheckpointRequest as ProtoCheckpointRequest, CheckpointResponse as ProtoCheckpointResponse,
50 DelegateRequest as ProtoDelegateRequest, DelegateResponse as ProtoDelegateResponse,
51 ForgetError as ProtoForgetError, ForgetRequest as ProtoForgetRequest,
52 ForgetResponse as ProtoForgetResponse, ForgetSubjectRequest as ProtoForgetSubjectRequest,
53 ForgetSubjectResponse as ProtoForgetSubjectResponse, HealthRequest, HealthResponse,
54 MergeRequest as ProtoMergeRequest, MergeResponse as ProtoMergeResponse,
55 RecallRequest as ProtoRecallRequest, RecallResponse as ProtoRecallResponse,
56 RememberRequest as ProtoRememberRequest, RememberResponse as ProtoRememberResponse,
57 ReplayMemory as ProtoReplayMemory, ReplayRequest as ProtoReplayRequest,
58 ReplayResponse as ProtoReplayResponse, ScoredMemory as ProtoScoredMemory,
59 ShareRequest as ProtoShareRequest, ShareResponse as ProtoShareResponse,
60 VerifyRequest as ProtoVerifyRequest, VerifyResponse as ProtoVerifyResponse,
61};
62
63#[derive(Clone)]
69pub struct MnemoGrpcServer {
70 engine: Arc<MnemoEngine>,
71}
72
73impl MnemoGrpcServer {
74 pub fn new(engine: Arc<MnemoEngine>) -> Self {
76 Self { engine }
77 }
78}
79
80#[tonic::async_trait]
81impl MnemoService for MnemoGrpcServer {
82 async fn remember(
85 &self,
86 request: Request<ProtoRememberRequest>,
87 ) -> Result<Response<ProtoRememberResponse>, Status> {
88 let req = request.into_inner();
89
90 let memory_type = match req.memory_type {
91 Some(ref s) => match s.parse::<MemoryType>() {
92 Ok(mt) => Some(mt),
93 Err(_) => {
94 return Err(Status::invalid_argument(format!(
95 "invalid memory_type '{}': expected one of: episodic, semantic, procedural, working",
96 s
97 )));
98 }
99 },
100 None => None,
101 };
102
103 let scope = match req.scope {
104 Some(ref s) => match s.parse::<Scope>() {
105 Ok(sc) => Some(sc),
106 Err(_) => {
107 return Err(Status::invalid_argument(format!(
108 "invalid scope '{}': expected one of: private, shared, public, global",
109 s
110 )));
111 }
112 },
113 None => None,
114 };
115
116 let source_type = match req.source_type {
117 Some(ref s) => match s.parse::<SourceType>() {
118 Ok(st) => Some(st),
119 Err(_) => {
120 return Err(Status::invalid_argument(format!(
121 "invalid source_type '{}': expected one of: agent, human, system, user_input, tool_output, model_response, retrieval, consolidation, import",
122 s
123 )));
124 }
125 },
126 None => None,
127 };
128
129 let metadata: Option<serde_json::Value> = match req.metadata {
130 Some(ref s) => match serde_json::from_str(s) {
131 Ok(v) => Some(v),
132 Err(e) => {
133 return Err(Status::invalid_argument(format!(
134 "invalid metadata JSON: {}",
135 e
136 )));
137 }
138 },
139 None => None,
140 };
141
142 let tags = if req.tags.is_empty() {
143 None
144 } else {
145 Some(req.tags)
146 };
147
148 let related_to = if req.related_to.is_empty() {
149 None
150 } else {
151 Some(req.related_to)
152 };
153
154 let core_req = CoreRememberRequest {
155 content: req.content,
156 agent_id: req.agent_id,
157 memory_type,
158 scope,
159 importance: req.importance,
160 tags,
161 metadata,
162 source_type,
163 source_id: req.source_id,
164 org_id: req.org_id,
165 thread_id: req.thread_id,
166 ttl_seconds: req.ttl_seconds,
167 related_to,
168 decay_rate: req.decay_rate,
169 created_by: req.created_by,
170 };
171
172 let result = self
173 .engine
174 .remember(core_req)
175 .await
176 .map_err(core_error_to_status)?;
177
178 Ok(Response::new(ProtoRememberResponse {
179 id: result.id.to_string(),
180 content_hash: result.content_hash,
181 }))
182 }
183
184 async fn recall(
187 &self,
188 request: Request<ProtoRecallRequest>,
189 ) -> Result<Response<ProtoRecallResponse>, Status> {
190 let req = request.into_inner();
191
192 let memory_type = match req.memory_type {
193 Some(ref s) => match s.parse::<MemoryType>() {
194 Ok(mt) => Some(mt),
195 Err(_) => {
196 return Err(Status::invalid_argument(format!(
197 "invalid memory_type '{}': expected one of: episodic, semantic, procedural, working",
198 s
199 )));
200 }
201 },
202 None => None,
203 };
204
205 let scope = match req.scope {
206 Some(ref s) => match s.parse::<Scope>() {
207 Ok(sc) => Some(sc),
208 Err(_) => {
209 return Err(Status::invalid_argument(format!(
210 "invalid scope '{}': expected one of: private, shared, public, global",
211 s
212 )));
213 }
214 },
215 None => None,
216 };
217
218 let tags = if req.tags.is_empty() {
219 None
220 } else {
221 Some(req.tags)
222 };
223
224 let hybrid_weights = if req.hybrid_weights.is_empty() {
225 None
226 } else {
227 Some(req.hybrid_weights)
228 };
229
230 let core_req = CoreRecallRequest {
231 query: req.query,
232 agent_id: req.agent_id,
233 limit: req.limit.map(|l| l as usize),
234 memory_type,
235 memory_types: None,
236 scope,
237 min_importance: req.min_importance,
238 tags,
239 org_id: req.org_id,
240 strategy: req.strategy,
241 temporal_range: None,
242 recency_half_life_hours: None,
243 hybrid_weights,
244 rrf_k: req.rrf_k,
245 as_of: req.as_of,
246 explain: req.explain,
247 with_provenance: None,
248 };
249
250 let result = self
251 .engine
252 .recall(core_req)
253 .await
254 .map_err(core_error_to_status)?;
255
256 let memories: Vec<ProtoScoredMemory> = result
257 .memories
258 .into_iter()
259 .map(|m| ProtoScoredMemory {
260 id: m.id.to_string(),
261 content: m.content,
262 memory_type: format!("{:?}", m.memory_type),
263 importance: m.importance,
264 score: m.score,
265 created_at: m.created_at,
266 agent_id: m.agent_id,
267 scope: format!("{:?}", m.scope),
268 tags: m.tags,
269 metadata: m.metadata.to_string(),
270 access_count: m.access_count,
271 updated_at: m.updated_at,
272 score_breakdown: m.score_breakdown.map(|b| proto::ScoreBreakdown {
273 vector: b.vector,
274 bm25: b.bm25,
275 graph: b.graph,
276 recency: b.recency,
277 rrf_rank: b.rrf_rank,
278 }),
279 })
280 .collect();
281
282 let total = result.total as u32;
283
284 Ok(Response::new(ProtoRecallResponse { memories, total }))
285 }
286
287 async fn forget(
290 &self,
291 request: Request<ProtoForgetRequest>,
292 ) -> Result<Response<ProtoForgetResponse>, Status> {
293 let req = request.into_inner();
294
295 let memory_ids: Vec<Uuid> = req
296 .memory_ids
297 .iter()
298 .map(|s| {
299 Uuid::parse_str(s)
300 .map_err(|e| Status::invalid_argument(format!("invalid UUID '{s}': {e}")))
301 })
302 .collect::<Result<Vec<_>, _>>()?;
303
304 let strategy = match req.strategy {
305 Some(ref s) => {
306 let st = match s.as_str() {
307 "soft_delete" => ForgetStrategy::SoftDelete,
308 "hard_delete" => ForgetStrategy::HardDelete,
309 "decay" => ForgetStrategy::Decay,
310 "consolidate" => ForgetStrategy::Consolidate,
311 "archive" => ForgetStrategy::Archive,
312 "redact" => ForgetStrategy::Redact,
313 _ => {
314 return Err(Status::invalid_argument(format!(
315 "invalid forget strategy '{}': expected one of: soft_delete, hard_delete, decay, consolidate, archive, redact",
316 s
317 )));
318 }
319 };
320 Some(st)
321 }
322 None => None,
323 };
324
325 let core_req = CoreForgetRequest {
326 memory_ids,
327 agent_id: req.agent_id,
328 strategy,
329 criteria: None,
330 };
331
332 let result = self
333 .engine
334 .forget(core_req)
335 .await
336 .map_err(core_error_to_status)?;
337
338 let forgotten: Vec<String> = result.forgotten.iter().map(|id| id.to_string()).collect();
339
340 let errors: Vec<ProtoForgetError> = result
341 .errors
342 .into_iter()
343 .map(|e| ProtoForgetError {
344 id: e.id.to_string(),
345 error: e.error,
346 })
347 .collect();
348
349 Ok(Response::new(ProtoForgetResponse { forgotten, errors }))
350 }
351
352 async fn health(
355 &self,
356 _request: Request<HealthRequest>,
357 ) -> Result<Response<HealthResponse>, Status> {
358 Ok(Response::new(HealthResponse {
359 status: "ok".to_string(),
360 version: env!("CARGO_PKG_VERSION").to_string(),
361 }))
362 }
363
364 async fn share(
367 &self,
368 request: Request<ProtoShareRequest>,
369 ) -> Result<Response<ProtoShareResponse>, Status> {
370 let req = request.into_inner();
371 let memory_id = Uuid::parse_str(&req.memory_id)
372 .map_err(|e| Status::invalid_argument(format!("invalid UUID: {e}")))?;
373 let permission = match req.permission {
374 Some(ref s) => match s.parse::<Permission>() {
375 Ok(p) => Some(p),
376 Err(_) => {
377 return Err(Status::invalid_argument(format!(
378 "invalid permission '{}': expected one of: read, write, delete, share, delegate, admin",
379 s
380 )));
381 }
382 },
383 None => None,
384 };
385 let target_agent_ids = if req.target_agent_ids.is_empty() {
386 None
387 } else {
388 Some(req.target_agent_ids)
389 };
390
391 let core_req = CoreShareRequest {
392 memory_id,
393 agent_id: req.agent_id,
394 target_agent_id: req.target_agent_id,
395 target_agent_ids,
396 permission,
397 expires_in_hours: req.expires_in_hours,
398 };
399 let result = self
400 .engine
401 .share(core_req)
402 .await
403 .map_err(core_error_to_status)?;
404
405 Ok(Response::new(ProtoShareResponse {
406 acl_id: result.acl_id.to_string(),
407 acl_ids: result.acl_ids.iter().map(|id| id.to_string()).collect(),
408 memory_id: result.memory_id.to_string(),
409 shared_with: result.shared_with,
410 shared_with_all: result.shared_with_all,
411 permission: result.permission.to_string(),
412 }))
413 }
414
415 async fn checkpoint(
418 &self,
419 request: Request<ProtoCheckpointRequest>,
420 ) -> Result<Response<ProtoCheckpointResponse>, Status> {
421 let req = request.into_inner();
422 let state_snapshot: serde_json::Value = serde_json::from_str(&req.state_snapshot)
423 .map_err(|e| Status::invalid_argument(format!("invalid JSON state_snapshot: {e}")))?;
424 let metadata: Option<serde_json::Value> = match req.metadata {
425 Some(ref s) => match serde_json::from_str(s) {
426 Ok(v) => Some(v),
427 Err(e) => {
428 return Err(Status::invalid_argument(format!(
429 "invalid metadata JSON: {}",
430 e
431 )));
432 }
433 },
434 None => None,
435 };
436
437 let core_req = CoreCheckpointRequest {
438 thread_id: req.thread_id,
439 agent_id: req.agent_id,
440 branch_name: req.branch_name,
441 state_snapshot,
442 label: req.label,
443 metadata,
444 };
445 let result = self
446 .engine
447 .checkpoint(core_req)
448 .await
449 .map_err(core_error_to_status)?;
450
451 Ok(Response::new(ProtoCheckpointResponse {
452 checkpoint_id: result.id.to_string(),
453 parent_id: result.parent_id.map(|id| id.to_string()),
454 branch_name: result.branch_name,
455 }))
456 }
457
458 async fn branch(
461 &self,
462 request: Request<ProtoBranchRequest>,
463 ) -> Result<Response<ProtoBranchResponse>, Status> {
464 let req = request.into_inner();
465 let source_checkpoint_id = match req.source_checkpoint_id {
466 Some(ref s) => match Uuid::parse_str(s) {
467 Ok(id) => Some(id),
468 Err(e) => {
469 return Err(Status::invalid_argument(format!(
470 "invalid source_checkpoint_id '{}': {}",
471 s, e
472 )));
473 }
474 },
475 None => None,
476 };
477
478 let core_req = CoreBranchRequest {
479 thread_id: req.thread_id,
480 agent_id: req.agent_id,
481 new_branch_name: req.new_branch_name,
482 source_checkpoint_id,
483 source_branch: req.source_branch,
484 };
485 let result = self
486 .engine
487 .branch(core_req)
488 .await
489 .map_err(core_error_to_status)?;
490
491 Ok(Response::new(ProtoBranchResponse {
492 checkpoint_id: result.checkpoint_id.to_string(),
493 branch_name: result.branch_name,
494 source_checkpoint_id: result.source_checkpoint_id.to_string(),
495 }))
496 }
497
498 async fn merge(
501 &self,
502 request: Request<ProtoMergeRequest>,
503 ) -> Result<Response<ProtoMergeResponse>, Status> {
504 let req = request.into_inner();
505 let strategy = match req.strategy {
506 Some(ref s) => {
507 let st = match s.as_str() {
508 "full_merge" => MergeStrategy::FullMerge,
509 "cherry_pick" => MergeStrategy::CherryPick,
510 "squash" => MergeStrategy::Squash,
511 _ => {
512 return Err(Status::invalid_argument(format!(
513 "invalid merge strategy '{}': expected one of: full_merge, cherry_pick, squash",
514 s
515 )));
516 }
517 };
518 Some(st)
519 }
520 None => None,
521 };
522 let cherry_pick_ids = if req.cherry_pick_ids.is_empty() {
523 None
524 } else {
525 let ids: Result<Vec<Uuid>, _> = req
526 .cherry_pick_ids
527 .iter()
528 .map(|s| Uuid::parse_str(s))
529 .collect();
530 Some(ids.map_err(|e| Status::invalid_argument(format!("invalid UUID: {e}")))?)
531 };
532
533 let core_req = CoreMergeRequest {
534 thread_id: req.thread_id,
535 agent_id: req.agent_id,
536 source_branch: req.source_branch,
537 target_branch: req.target_branch,
538 strategy,
539 cherry_pick_ids,
540 };
541 let result = self
542 .engine
543 .merge(core_req)
544 .await
545 .map_err(core_error_to_status)?;
546
547 Ok(Response::new(ProtoMergeResponse {
548 checkpoint_id: result.checkpoint_id.to_string(),
549 target_branch: result.target_branch,
550 merged_memory_count: result.merged_memory_count as u32,
551 }))
552 }
553
554 async fn replay(
557 &self,
558 request: Request<ProtoReplayRequest>,
559 ) -> Result<Response<ProtoReplayResponse>, Status> {
560 let req = request.into_inner();
561 let checkpoint_id = match req.checkpoint_id {
562 Some(ref s) => match Uuid::parse_str(s) {
563 Ok(id) => Some(id),
564 Err(e) => {
565 return Err(Status::invalid_argument(format!(
566 "invalid checkpoint_id '{}': {}",
567 s, e
568 )));
569 }
570 },
571 None => None,
572 };
573
574 let core_req = CoreReplayRequest {
575 thread_id: req.thread_id,
576 agent_id: req.agent_id,
577 checkpoint_id,
578 branch_name: req.branch_name,
579 as_of: req.as_of,
580 };
581 let result = self
582 .engine
583 .replay(core_req)
584 .await
585 .map_err(core_error_to_status)?;
586
587 let checkpoint_json =
588 serde_json::to_string(&result.checkpoint).unwrap_or_else(|_| "{}".to_string());
589 let memories: Vec<ProtoReplayMemory> = result
590 .memories
591 .iter()
592 .map(|m| ProtoReplayMemory {
593 id: m.id.to_string(),
594 content: m.content.clone(),
595 memory_type: format!("{:?}", m.memory_type),
596 created_at: m.created_at.clone(),
597 })
598 .collect();
599
600 let (chain_valid, chain_total, chain_verified) =
601 if let Some(ref cv) = result.chain_verification {
602 (
603 Some(cv.valid),
604 Some(cv.total_records as u32),
605 Some(cv.verified_records as u32),
606 )
607 } else {
608 (None, None, None)
609 };
610
611 Ok(Response::new(ProtoReplayResponse {
612 checkpoint_json,
613 memories,
614 event_count: result.events.len() as u32,
615 chain_valid,
616 chain_total,
617 chain_verified,
618 }))
619 }
620
621 async fn delegate(
624 &self,
625 request: Request<ProtoDelegateRequest>,
626 ) -> Result<Response<ProtoDelegateResponse>, Status> {
627 let req = request.into_inner();
628 let permission: Permission = req
629 .permission
630 .parse()
631 .map_err(|e: mnemo_core::error::Error| Status::invalid_argument(e.to_string()))?;
632
633 let scope = if !req.memory_ids.is_empty() {
634 let ids: Vec<Uuid> = req
635 .memory_ids
636 .iter()
637 .map(|s| {
638 Uuid::parse_str(s)
639 .map_err(|e| Status::invalid_argument(format!("invalid UUID: {e}")))
640 })
641 .collect::<Result<Vec<_>, _>>()?;
642 DelegationScope::ByMemoryId(ids)
643 } else if !req.tags.is_empty() {
644 DelegationScope::ByTag(req.tags)
645 } else {
646 DelegationScope::AllMemories
647 };
648
649 let now = chrono::Utc::now();
650 let expires_at = req
651 .expires_in_hours
652 .map(|h| (now + chrono::Duration::seconds((h * 3600.0) as i64)).to_rfc3339());
653
654 let delegation = Delegation {
655 id: Uuid::now_v7(),
656 delegator_id: req.delegator_id,
657 delegate_id: req.delegate_id,
658 permission,
659 scope,
660 max_depth: req.max_depth.unwrap_or(0),
661 current_depth: 0,
662 parent_delegation_id: None,
663 created_at: now.to_rfc3339(),
664 expires_at,
665 revoked_at: None,
666 };
667
668 self.engine
669 .storage
670 .insert_delegation(&delegation)
671 .await
672 .map_err(core_error_to_status)?;
673
674 Ok(Response::new(ProtoDelegateResponse {
675 delegation_id: delegation.id.to_string(),
676 }))
677 }
678
679 async fn verify(
682 &self,
683 request: Request<ProtoVerifyRequest>,
684 ) -> Result<Response<ProtoVerifyResponse>, Status> {
685 let req = request.into_inner();
686 let result = self
687 .engine
688 .verify_integrity(req.agent_id, req.thread_id.as_deref())
689 .await
690 .map_err(core_error_to_status)?;
691
692 Ok(Response::new(ProtoVerifyResponse {
693 valid: result.valid,
694 total_records: result.total_records as u32,
695 verified_records: result.verified_records as u32,
696 first_broken_at: result.first_broken_at.map(|id| id.to_string()),
697 error_message: result.error_message,
698 }))
699 }
700
701 async fn forget_subject(
704 &self,
705 request: Request<ProtoForgetSubjectRequest>,
706 ) -> Result<Response<ProtoForgetSubjectResponse>, Status> {
707 let req = request.into_inner();
708
709 let strategy = match req.strategy.as_deref().unwrap_or("redact") {
710 "redact" => ForgetStrategy::Redact,
711 "hard_delete" => ForgetStrategy::HardDelete,
712 "soft_delete" => ForgetStrategy::SoftDelete,
713 other => {
714 return Err(Status::invalid_argument(format!(
715 "invalid forget_subject strategy '{}': expected one of: redact, hard_delete, soft_delete",
716 other
717 )));
718 }
719 };
720
721 let core_req = CoreForgetSubjectRequest {
722 subject_id: req.subject_id,
723 agent_id: req.agent_id,
724 strategy,
725 };
726
727 let result = self
728 .engine
729 .forget_subject(core_req)
730 .await
731 .map_err(core_error_to_status)?;
732
733 let errors: Vec<ProtoForgetError> = result
734 .errors
735 .into_iter()
736 .map(|e| ProtoForgetError {
737 id: e.id.to_string(),
738 error: e.error,
739 })
740 .collect();
741
742 let strategy_str = match result.strategy {
743 ForgetStrategy::SoftDelete => "soft_delete",
744 ForgetStrategy::HardDelete => "hard_delete",
745 ForgetStrategy::Decay => "decay",
746 ForgetStrategy::Consolidate => "consolidate",
747 ForgetStrategy::Archive => "archive",
748 ForgetStrategy::Redact => "redact",
749 }
750 .to_string();
751
752 Ok(Response::new(ProtoForgetSubjectResponse {
753 subject_id: result.subject_id,
754 strategy: strategy_str,
755 matched: result.matched as u32,
756 forgotten: result.forgotten.iter().map(|id| id.to_string()).collect(),
757 cascaded_events: result.cascaded_events as u32,
758 errors,
759 }))
760 }
761}
762
763pub fn router(engine: Arc<MnemoEngine>) -> tonic::transport::server::Router {
787 let svc = MnemoGrpcServer::new(engine);
788 tonic::transport::Server::builder().add_service(MnemoServiceServer::new(svc))
789}
790
791fn core_error_to_status(err: mnemo_core::error::Error) -> Status {
797 use mnemo_core::error::Error;
798
799 match err {
800 Error::Validation(msg) => Status::invalid_argument(msg),
801 Error::PermissionDenied(msg) => Status::permission_denied(msg),
802 Error::NotFound(msg) => Status::not_found(msg),
803 other => Status::internal(other.to_string()),
804 }
805}
806
807#[cfg(test)]
812mod tests {
813 use super::*;
814
815 #[test]
816 fn core_error_maps_correctly() {
817 let validation =
818 core_error_to_status(mnemo_core::error::Error::Validation("bad input".into()));
819 assert_eq!(validation.code(), tonic::Code::InvalidArgument);
820
821 let perm = core_error_to_status(mnemo_core::error::Error::PermissionDenied(
822 "forbidden".into(),
823 ));
824 assert_eq!(perm.code(), tonic::Code::PermissionDenied);
825
826 let not_found = core_error_to_status(mnemo_core::error::Error::NotFound("missing".into()));
827 assert_eq!(not_found.code(), tonic::Code::NotFound);
828 }
829}