Skip to main content

mentedb_server/
grpc.rs

1//! gRPC service implementations for MenteDB.
2//!
3//! Provides a CognitionService with bidirectional streaming that wraps the
4//! CognitionStream engine, and a MemoryService with standard CRUD operations.
5
6use std::collections::HashMap;
7use std::pin::Pin;
8use std::sync::Arc;
9
10use mentedb_cognitive::stream::{CognitionStream, StreamAlert};
11use mentedb_core::edge::EdgeType;
12use mentedb_core::memory::{AttributeValue, MemoryType};
13use mentedb_core::{MemoryEdge, MemoryNode};
14use tokio::sync::mpsc;
15use tokio_stream::{Stream, StreamExt, wrappers::ReceiverStream};
16use tonic::{Request, Response, Status, Streaming};
17use tracing::{error, info};
18use uuid::Uuid;
19
20use crate::auth;
21use crate::state::AppState;
22
23pub mod pb {
24    tonic::include_proto!("mentedb");
25}
26
27use mentedb_core::types::{AgentId, MemoryId, SpaceId, UserId};
28use pb::cognition_service_server::CognitionService;
29use pb::memory_service_server::MemoryService;
30
31// ---------------------------------------------------------------------------
32// CognitionService
33// ---------------------------------------------------------------------------
34
35/// gRPC service implementation for cognitive memory operations.
36pub struct CognitionServiceImpl {
37    #[allow(dead_code)]
38    pub state: Arc<AppState>,
39}
40
41type CognitionStream_ =
42    Pin<Box<dyn Stream<Item = Result<pb::CognitionAlert, Status>> + Send + 'static>>;
43
44#[tonic::async_trait]
45impl CognitionService for CognitionServiceImpl {
46    type StreamCognitionStream = CognitionStream_;
47
48    async fn stream_cognition(
49        &self,
50        request: Request<Streaming<pb::CognitionRequest>>,
51    ) -> Result<Response<Self::StreamCognitionStream>, Status> {
52        authenticate_grpc_streaming(&self.state, request.metadata())?;
53        info!("gRPC cognition stream opened");
54
55        let stream_engine = Arc::new(CognitionStream::new(1000));
56        let known_facts: Arc<std::sync::Mutex<Vec<(MemoryId, String)>>> =
57            Arc::new(std::sync::Mutex::new(Vec::new()));
58
59        let (tx, rx) = mpsc::channel(128);
60        let mut inbound = request.into_inner();
61
62        let engine = stream_engine.clone();
63        let facts = known_facts.clone();
64
65        tokio::spawn(async move {
66            while let Some(result) = inbound.next().await {
67                let req = match result {
68                    Ok(r) => r,
69                    Err(e) => {
70                        error!("cognition stream receive error: {e}");
71                        break;
72                    }
73                };
74
75                let payload = match req.payload {
76                    Some(p) => p,
77                    None => continue,
78                };
79
80                match payload {
81                    pb::cognition_request::Payload::Token(t) => {
82                        engine.feed_token(&t.text);
83
84                        let current_facts = facts.lock().unwrap().clone();
85                        let alerts = engine.check_alerts(&current_facts);
86                        for alert in alerts {
87                            let proto_alert = stream_alert_to_proto(alert);
88                            if tx.send(Ok(proto_alert)).await.is_err() {
89                                return;
90                            }
91                        }
92                    }
93                    pb::cognition_request::Payload::KnownFact(kf) => {
94                        let mid = kf
95                            .memory_id
96                            .parse::<MemoryId>()
97                            .unwrap_or_else(|_| MemoryId::nil());
98                        facts.lock().unwrap().push((mid, kf.content));
99                    }
100                    pb::cognition_request::Payload::EndOfTurn(_) => {
101                        let current_facts = facts.lock().unwrap().clone();
102                        let alerts = engine.check_alerts(&current_facts);
103                        for alert in alerts {
104                            let proto_alert = stream_alert_to_proto(alert);
105                            if tx.send(Ok(proto_alert)).await.is_err() {
106                                return;
107                            }
108                        }
109                    }
110                    pb::cognition_request::Payload::Flush(_) => {
111                        let text = engine.drain_buffer();
112                        let alert = pb::CognitionAlert {
113                            alert: Some(pb::cognition_alert::Alert::BufferFlushed(
114                                pb::BufferFlushedAlert {
115                                    accumulated_text: text,
116                                },
117                            )),
118                        };
119                        if tx.send(Ok(alert)).await.is_err() {
120                            return;
121                        }
122                    }
123                }
124            }
125            info!("gRPC cognition stream closed");
126        });
127
128        let output = ReceiverStream::new(rx);
129        Ok(Response::new(Box::pin(output)))
130    }
131}
132
133fn stream_alert_to_proto(alert: StreamAlert) -> pb::CognitionAlert {
134    let inner = match alert {
135        StreamAlert::Contradiction {
136            memory_id,
137            ai_said,
138            stored,
139        } => pb::cognition_alert::Alert::Contradiction(pb::ContradictionAlert {
140            memory_id: memory_id.to_string(),
141            ai_said,
142            stored,
143        }),
144        StreamAlert::Forgotten { memory_id, summary } => {
145            pb::cognition_alert::Alert::Forgotten(pb::ForgottenAlert {
146                memory_id: memory_id.to_string(),
147                summary,
148            })
149        }
150        StreamAlert::Correction {
151            memory_id,
152            old,
153            new,
154        } => pb::cognition_alert::Alert::Correction(pb::CorrectionAlert {
155            memory_id: memory_id.to_string(),
156            old_content: old,
157            new_content: new,
158        }),
159        StreamAlert::Reinforcement { memory_id } => {
160            pb::cognition_alert::Alert::Reinforcement(pb::ReinforcementAlert {
161                memory_id: memory_id.to_string(),
162            })
163        }
164    };
165    pb::CognitionAlert { alert: Some(inner) }
166}
167
168// ---------------------------------------------------------------------------
169// MemoryService
170// ---------------------------------------------------------------------------
171
172/// gRPC service implementation for basic memory CRUD operations.
173pub struct MemoryServiceImpl {
174    /// Shared application state containing the database handle.
175    pub state: Arc<AppState>,
176}
177
178#[tonic::async_trait]
179impl MemoryService for MemoryServiceImpl {
180    async fn store(
181        &self,
182        request: Request<pb::StoreRequest>,
183    ) -> Result<Response<pb::StoreResponse>, Status> {
184        let caller = authenticate_grpc_request(&self.state, &request)?;
185        let req = request.into_inner();
186
187        let agent_id = AgentId(parse_uuid_field(&req.agent_id, "agent_id")?);
188        if let Some(ref ta) = caller {
189            let tu: AgentId = ta
190                .parse()
191                .map_err(|_| Status::internal("bad token agent_id"))?;
192            if tu != agent_id {
193                return Err(Status::permission_denied("agent_id mismatch"));
194            }
195        }
196        let memory_type = parse_memory_type_str(&req.memory_type)?;
197        let space_id = if req.space_id.is_empty() {
198            SpaceId::nil()
199        } else {
200            SpaceId(parse_uuid_field(&req.space_id, "space_id")?)
201        };
202
203        let salience = if req.salience == 0.0 {
204            0.5
205        } else {
206            req.salience
207        };
208        let confidence = if req.confidence == 0.0 {
209            1.0
210        } else {
211            req.confidence
212        };
213
214        let now = std::time::SystemTime::now()
215            .duration_since(std::time::UNIX_EPOCH)
216            .unwrap_or_default()
217            .as_micros() as u64;
218
219        let id = MemoryId::new();
220
221        let attributes = req
222            .attributes
223            .into_iter()
224            .map(|(k, v)| (k, AttributeValue::String(v)))
225            .collect::<HashMap<String, AttributeValue>>();
226
227        let node = MemoryNode {
228            id,
229            agent_id,
230            user_id: UserId::nil(),
231            memory_type,
232            embedding: req.embedding,
233            content: req.content,
234            created_at: now,
235            accessed_at: now,
236            access_count: 0,
237            salience,
238            confidence,
239            space_id,
240            attributes,
241            tags: req.tags,
242            valid_from: req.valid_from,
243            valid_until: req.valid_until,
244        };
245
246        let db = &*self.state.db;
247        db.store(node).map_err(|e| {
248            error!("gRPC store failed: {e}");
249            Status::internal(format!("store failed: {e}"))
250        })?;
251
252        Ok(Response::new(pb::StoreResponse {
253            id: id.to_string(),
254            status: "stored".into(),
255        }))
256    }
257
258    async fn recall(
259        &self,
260        request: Request<pb::RecallRequest>,
261    ) -> Result<Response<pb::RecallResponse>, Status> {
262        authenticate_grpc_request(&self.state, &request)?;
263        let req = request.into_inner();
264
265        let db = &*self.state.db;
266        let window = db.recall(&req.query).map_err(|e| {
267            error!("gRPC recall failed: {e}");
268            Status::internal(format!("recall failed: {e}"))
269        })?;
270
271        let memory_count: usize = window.blocks.iter().map(|b| b.memories.len()).sum();
272
273        Ok(Response::new(pb::RecallResponse {
274            context: window.format,
275            total_tokens: window.total_tokens as u64,
276            memory_count: memory_count as u64,
277        }))
278    }
279
280    async fn search(
281        &self,
282        request: Request<pb::SearchRequest>,
283    ) -> Result<Response<pb::SearchResponse>, Status> {
284        authenticate_grpc_request(&self.state, &request)?;
285        let req = request.into_inner();
286        let k = if req.k == 0 { 10 } else { req.k as usize };
287
288        if req.embedding.is_empty() {
289            return Err(Status::invalid_argument("missing embedding vector"));
290        }
291
292        let db = &*self.state.db;
293        let results = db.recall_similar(&req.embedding, k).map_err(|e| {
294            error!("gRPC search failed: {e}");
295            Status::internal(format!("search failed: {e}"))
296        })?;
297
298        let items: Vec<pb::SearchResult> = results
299            .iter()
300            .map(|(id, score)| pb::SearchResult {
301                id: id.to_string(),
302                score: *score,
303            })
304            .collect();
305
306        Ok(Response::new(pb::SearchResponse { results: items }))
307    }
308
309    async fn forget(
310        &self,
311        request: Request<pb::ForgetRequest>,
312    ) -> Result<Response<pb::ForgetResponse>, Status> {
313        authenticate_grpc_request(&self.state, &request)?;
314        let req = request.into_inner();
315        let id = MemoryId(parse_uuid_field(&req.id, "id")?);
316
317        let db = &*self.state.db;
318        db.forget(id).map_err(|e| {
319            error!("gRPC forget failed: {e}");
320            Status::internal(format!("forget failed: {e}"))
321        })?;
322
323        Ok(Response::new(pb::ForgetResponse {
324            status: "deleted".into(),
325        }))
326    }
327
328    async fn relate(
329        &self,
330        request: Request<pb::RelateRequest>,
331    ) -> Result<Response<pb::RelateResponse>, Status> {
332        authenticate_grpc_request(&self.state, &request)?;
333        let req = request.into_inner();
334        let source = MemoryId(parse_uuid_field(&req.source, "source")?);
335        let target = MemoryId(parse_uuid_field(&req.target, "target")?);
336        let edge_type = parse_edge_type_str(&req.edge_type)?;
337        let weight = if req.weight == 0.0 { 1.0 } else { req.weight };
338
339        let now = std::time::SystemTime::now()
340            .duration_since(std::time::UNIX_EPOCH)
341            .unwrap_or_default()
342            .as_micros() as u64;
343
344        let edge = MemoryEdge {
345            source,
346            target,
347            edge_type,
348            weight,
349            created_at: now,
350            valid_from: req.valid_from,
351            valid_until: req.valid_until,
352            label: None,
353        };
354
355        let db = &*self.state.db;
356        db.relate(edge).map_err(|e| {
357            error!("gRPC relate failed: {e}");
358            Status::internal(format!("relate failed: {e}"))
359        })?;
360
361        Ok(Response::new(pb::RelateResponse {
362            status: "created".into(),
363        }))
364    }
365}
366
367#[allow(clippy::result_large_err)]
368fn authenticate_grpc_request<T>(
369    state: &AppState,
370    request: &Request<T>,
371) -> Result<Option<String>, Status> {
372    let secret = match &state.jwt_secret {
373        Some(s) => s,
374        None => return Ok(None),
375    };
376    let token = request
377        .metadata()
378        .get("authorization")
379        .and_then(|v| v.to_str().ok())
380        .and_then(|v| v.strip_prefix("Bearer "))
381        .ok_or_else(|| Status::unauthenticated("missing or invalid authorization metadata"))?;
382    auth::validate_token(secret, token)
383        .map(|c| Some(c.agent_id))
384        .map_err(|e| Status::unauthenticated(format!("invalid token: {e}")))
385}
386#[allow(clippy::result_large_err)]
387fn authenticate_grpc_streaming(
388    state: &AppState,
389    metadata: &tonic::metadata::MetadataMap,
390) -> Result<Option<String>, Status> {
391    let secret = match &state.jwt_secret {
392        Some(s) => s,
393        None => return Ok(None),
394    };
395    let token = metadata
396        .get("authorization")
397        .and_then(|v| v.to_str().ok())
398        .and_then(|v| v.strip_prefix("Bearer "))
399        .ok_or_else(|| Status::unauthenticated("missing or invalid authorization metadata"))?;
400    auth::validate_token(secret, token)
401        .map(|c| Some(c.agent_id))
402        .map_err(|e| Status::unauthenticated(format!("invalid token: {e}")))
403}
404
405// ---------------------------------------------------------------------------
406// Helpers
407// ---------------------------------------------------------------------------
408
409#[allow(clippy::result_large_err)]
410fn parse_uuid_field(s: &str, field: &str) -> Result<Uuid, Status> {
411    Uuid::parse_str(s)
412        .map_err(|_| Status::invalid_argument(format!("invalid UUID for '{field}': {s}")))
413}
414
415#[allow(clippy::result_large_err)]
416fn parse_memory_type_str(s: &str) -> Result<MemoryType, Status> {
417    match s.to_lowercase().as_str() {
418        "episodic" => Ok(MemoryType::Episodic),
419        "semantic" => Ok(MemoryType::Semantic),
420        "procedural" => Ok(MemoryType::Procedural),
421        "antipattern" | "anti_pattern" => Ok(MemoryType::AntiPattern),
422        "reasoning" => Ok(MemoryType::Reasoning),
423        "correction" => Ok(MemoryType::Correction),
424        _ => Err(Status::invalid_argument(format!(
425            "unknown memory_type: {s}"
426        ))),
427    }
428}
429
430#[allow(clippy::result_large_err)]
431fn parse_edge_type_str(s: &str) -> Result<EdgeType, Status> {
432    match s.to_lowercase().as_str() {
433        "caused" => Ok(EdgeType::Caused),
434        "before" => Ok(EdgeType::Before),
435        "related" => Ok(EdgeType::Related),
436        "contradicts" => Ok(EdgeType::Contradicts),
437        "supports" => Ok(EdgeType::Supports),
438        "supersedes" => Ok(EdgeType::Supersedes),
439        "derived" => Ok(EdgeType::Derived),
440        "partof" | "part_of" => Ok(EdgeType::PartOf),
441        _ => Err(Status::invalid_argument(format!("unknown edge_type: {s}"))),
442    }
443}
444
445// ---------------------------------------------------------------------------
446// Tests
447// ---------------------------------------------------------------------------
448
449#[cfg(test)]
450mod tests {
451    use super::*;
452    use mentedb::MenteDb;
453    use std::time::Instant;
454    use tempfile::TempDir;
455
456    fn make_test_state() -> (Arc<AppState>, TempDir) {
457        let tmp = TempDir::new().unwrap();
458        let db = MenteDb::open(tmp.path()).unwrap();
459        let state = Arc::new(AppState {
460            db: Arc::new(db),
461            spaces: Arc::new(tokio::sync::RwLock::new(mentedb_core::SpaceManager::new())),
462            jwt_secret: None,
463            admin_key: None,
464            start_time: Instant::now(),
465            extraction_config: None,
466            auto_extract: false,
467            extraction_tx: None,
468        });
469        (state, tmp)
470    }
471
472    #[tokio::test]
473    async fn test_grpc_memory_store_and_recall() {
474        let (state, _tmp) = make_test_state();
475        let svc = MemoryServiceImpl {
476            state: state.clone(),
477        };
478
479        let agent_id = AgentId::new().to_string();
480        let store_req = Request::new(pb::StoreRequest {
481            agent_id: agent_id.clone(),
482            memory_type: "episodic".into(),
483            content: "The user prefers dark mode".into(),
484            embedding: vec![],
485            tags: vec!["preference".into()],
486            attributes: HashMap::new(),
487            space_id: String::new(),
488            salience: 0.8,
489            confidence: 1.0,
490            valid_from: None,
491            valid_until: None,
492        });
493
494        let resp = svc.store(store_req).await.unwrap();
495        let inner = resp.into_inner();
496        assert_eq!(inner.status, "stored");
497        assert!(!inner.id.is_empty());
498
499        let recall_req = Request::new(pb::RecallRequest {
500            query: "RECALL memories LIMIT 100".into(),
501        });
502        let resp = svc.recall(recall_req).await.unwrap();
503        let inner = resp.into_inner();
504        // Recall succeeded without error; count may be 0 if the engine has not
505        // indexed the memory yet, so we just verify the response is valid.
506        let _ = inner.memory_count;
507    }
508
509    #[tokio::test]
510    async fn test_grpc_memory_forget() {
511        let (state, _tmp) = make_test_state();
512        let svc = MemoryServiceImpl {
513            state: state.clone(),
514        };
515
516        let agent_id = AgentId::new().to_string();
517        let store_resp = svc
518            .store(Request::new(pb::StoreRequest {
519                agent_id,
520                memory_type: "semantic".into(),
521                content: "Temporary memory".into(),
522                embedding: vec![],
523                tags: vec![],
524                attributes: HashMap::new(),
525                space_id: String::new(),
526                salience: 0.5,
527                confidence: 1.0,
528                valid_from: None,
529                valid_until: None,
530            }))
531            .await
532            .unwrap();
533        let stored_id = store_resp.into_inner().id;
534
535        let forget_resp = svc
536            .forget(Request::new(pb::ForgetRequest {
537                id: stored_id.clone(),
538            }))
539            .await
540            .unwrap();
541        assert_eq!(forget_resp.into_inner().status, "deleted");
542    }
543
544    #[tokio::test]
545    async fn test_grpc_memory_relate() {
546        let (state, _tmp) = make_test_state();
547        let svc = MemoryServiceImpl {
548            state: state.clone(),
549        };
550
551        let agent_id = AgentId::new().to_string();
552        let id1 = svc
553            .store(Request::new(pb::StoreRequest {
554                agent_id: agent_id.clone(),
555                memory_type: "episodic".into(),
556                content: "Event A".into(),
557                embedding: vec![],
558                tags: vec![],
559                attributes: HashMap::new(),
560                space_id: String::new(),
561                salience: 0.5,
562                confidence: 1.0,
563                valid_from: None,
564                valid_until: None,
565            }))
566            .await
567            .unwrap()
568            .into_inner()
569            .id;
570
571        let id2 = svc
572            .store(Request::new(pb::StoreRequest {
573                agent_id,
574                memory_type: "episodic".into(),
575                content: "Event B".into(),
576                embedding: vec![],
577                tags: vec![],
578                attributes: HashMap::new(),
579                space_id: String::new(),
580                salience: 0.5,
581                confidence: 1.0,
582                valid_from: None,
583                valid_until: None,
584            }))
585            .await
586            .unwrap()
587            .into_inner()
588            .id;
589
590        let relate_resp = svc
591            .relate(Request::new(pb::RelateRequest {
592                source: id1,
593                target: id2,
594                edge_type: "caused".into(),
595                weight: 0.9,
596                valid_from: None,
597                valid_until: None,
598            }))
599            .await
600            .unwrap();
601        assert_eq!(relate_resp.into_inner().status, "created");
602    }
603
604    #[tokio::test]
605    async fn test_grpc_cognition_stream() {
606        use pb::cognition_service_server::CognitionServiceServer;
607        use tokio::net::TcpListener;
608        use tonic::transport::{Channel, Server};
609
610        let (state, _tmp) = make_test_state();
611
612        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
613        let addr = listener.local_addr().unwrap();
614
615        let svc = CognitionServiceServer::new(CognitionServiceImpl {
616            state: state.clone(),
617        });
618
619        tokio::spawn(async move {
620            let incoming = tokio_stream::wrappers::TcpListenerStream::new(listener);
621            Server::builder()
622                .add_service(svc)
623                .serve_with_incoming(incoming)
624                .await
625                .unwrap();
626        });
627
628        // Small delay to let the server start
629        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
630
631        let channel = Channel::from_shared(format!("http://{addr}"))
632            .unwrap()
633            .connect()
634            .await
635            .unwrap();
636
637        let mut client = pb::cognition_service_client::CognitionServiceClient::new(channel);
638
639        let (tx, rx) = mpsc::channel(32);
640        let outbound = ReceiverStream::new(rx);
641
642        let response = client.stream_cognition(outbound).await.unwrap();
643        let mut inbound = response.into_inner();
644
645        // Register a known fact
646        let mid = MemoryId::new();
647        tx.send(pb::CognitionRequest {
648            payload: Some(pb::cognition_request::Payload::KnownFact(pb::KnownFact {
649                memory_id: mid.to_string(),
650                content: "The system uses PostgreSQL for storage".into(),
651            })),
652        })
653        .await
654        .unwrap();
655
656        // Feed tokens that contradict the known fact
657        tx.send(pb::CognitionRequest {
658            payload: Some(pb::cognition_request::Payload::Token(pb::TokenPayload {
659                text: "The system does not use PostgreSQL, actually it uses MySQL".into(),
660            })),
661        })
662        .await
663        .unwrap();
664
665        // We should get a contradiction alert
666        if let Some(Ok(alert)) = inbound.next().await {
667            match alert.alert {
668                Some(pb::cognition_alert::Alert::Contradiction(c)) => {
669                    assert_eq!(c.memory_id, mid.to_string());
670                    assert!(!c.ai_said.is_empty());
671                    assert!(!c.stored.is_empty());
672                }
673                other => panic!("expected contradiction alert, got: {:?}", other),
674            }
675        } else {
676            panic!("expected an alert from the cognition stream");
677        }
678
679        // Test flush
680        tx.send(pb::CognitionRequest {
681            payload: Some(pb::cognition_request::Payload::Flush(pb::FlushBuffer {})),
682        })
683        .await
684        .unwrap();
685
686        if let Some(Ok(alert)) = inbound.next().await {
687            match alert.alert {
688                Some(pb::cognition_alert::Alert::BufferFlushed(f)) => {
689                    assert!(!f.accumulated_text.is_empty());
690                }
691                other => panic!("expected buffer_flushed alert, got: {:?}", other),
692            }
693        }
694
695        drop(tx);
696    }
697
698    #[tokio::test]
699    async fn test_grpc_invalid_uuid() {
700        let (state, _tmp) = make_test_state();
701        let svc = MemoryServiceImpl { state };
702
703        let result = svc
704            .store(Request::new(pb::StoreRequest {
705                agent_id: "not-a-uuid".into(),
706                memory_type: "episodic".into(),
707                content: "test".into(),
708                embedding: vec![],
709                tags: vec![],
710                attributes: HashMap::new(),
711                space_id: String::new(),
712                salience: 0.5,
713                confidence: 1.0,
714                valid_from: None,
715                valid_until: None,
716            }))
717            .await;
718
719        assert!(result.is_err());
720        let status = result.unwrap_err();
721        assert_eq!(status.code(), tonic::Code::InvalidArgument);
722    }
723
724    #[tokio::test]
725    async fn test_grpc_invalid_memory_type() {
726        let (state, _tmp) = make_test_state();
727        let svc = MemoryServiceImpl { state };
728
729        let result = svc
730            .store(Request::new(pb::StoreRequest {
731                agent_id: AgentId::new().to_string(),
732                memory_type: "nonexistent".into(),
733                content: "test".into(),
734                embedding: vec![],
735                tags: vec![],
736                attributes: HashMap::new(),
737                space_id: String::new(),
738                salience: 0.5,
739                confidence: 1.0,
740                valid_from: None,
741                valid_until: None,
742            }))
743            .await;
744
745        assert!(result.is_err());
746        let status = result.unwrap_err();
747        assert_eq!(status.code(), tonic::Code::InvalidArgument);
748    }
749
750    #[tokio::test]
751    async fn test_grpc_search_missing_embedding() {
752        let (state, _tmp) = make_test_state();
753        let svc = MemoryServiceImpl { state };
754
755        let result = svc
756            .search(Request::new(pb::SearchRequest {
757                embedding: vec![],
758                k: 10,
759            }))
760            .await;
761
762        assert!(result.is_err());
763        assert_eq!(result.unwrap_err().code(), tonic::Code::InvalidArgument);
764    }
765
766    fn make_auth_test_state(secret: &str) -> (Arc<AppState>, TempDir) {
767        let tmp = TempDir::new().unwrap();
768        let db = MenteDb::open(tmp.path()).unwrap();
769        (
770            Arc::new(AppState {
771                db: Arc::new(db),
772                spaces: Arc::new(tokio::sync::RwLock::new(mentedb_core::SpaceManager::new())),
773                jwt_secret: Some(secret.into()),
774                admin_key: Some("ak".into()),
775                start_time: Instant::now(),
776                extraction_config: None,
777                auto_extract: false,
778                extraction_tx: None,
779            }),
780            tmp,
781        )
782    }
783    #[tokio::test]
784    async fn test_grpc_auth_required_when_secret_set() {
785        let (s, _t) = make_auth_test_state("s");
786        let svc = MemoryServiceImpl { state: s };
787        let r = svc
788            .recall(Request::new(pb::RecallRequest {
789                query: "RECALL memories LIMIT 10".into(),
790            }))
791            .await;
792        assert!(r.is_err());
793        assert_eq!(r.unwrap_err().code(), tonic::Code::Unauthenticated);
794    }
795    #[tokio::test]
796    async fn test_grpc_auth_succeeds_with_valid_token() {
797        let (s, _t) = make_auth_test_state("s");
798        let svc = MemoryServiceImpl { state: s };
799        let a = AgentId::new();
800        let tok = crate::auth::create_token("s", &a.to_string(), false, 1);
801        let mut r = Request::new(pb::StoreRequest {
802            agent_id: a.to_string(),
803            memory_type: "episodic".into(),
804            content: "t".into(),
805            embedding: vec![],
806            tags: vec![],
807            attributes: HashMap::new(),
808            space_id: String::new(),
809            salience: 0.5,
810            confidence: 1.0,
811            valid_from: None,
812            valid_until: None,
813        });
814        r.metadata_mut()
815            .insert("authorization", format!("Bearer {tok}").parse().unwrap());
816        assert_eq!(svc.store(r).await.unwrap().into_inner().status, "stored");
817    }
818    #[tokio::test]
819    async fn test_grpc_auth_rejects_wrong_agent_id() {
820        let (s, _t) = make_auth_test_state("s");
821        let svc = MemoryServiceImpl { state: s };
822        let ta = AgentId::new();
823        let ra = AgentId::new();
824        let tok = crate::auth::create_token("s", &ta.to_string(), false, 1);
825        let mut r = Request::new(pb::StoreRequest {
826            agent_id: ra.to_string(),
827            memory_type: "episodic".into(),
828            content: "t".into(),
829            embedding: vec![],
830            tags: vec![],
831            attributes: HashMap::new(),
832            space_id: String::new(),
833            salience: 0.5,
834            confidence: 1.0,
835            valid_from: None,
836            valid_until: None,
837        });
838        r.metadata_mut()
839            .insert("authorization", format!("Bearer {tok}").parse().unwrap());
840        assert_eq!(
841            svc.store(r).await.unwrap_err().code(),
842            tonic::Code::PermissionDenied
843        );
844    }
845    #[tokio::test]
846    async fn test_grpc_auth_invalid_token() {
847        let (s, _t) = make_auth_test_state("s");
848        let svc = MemoryServiceImpl { state: s };
849        let mut r = Request::new(pb::RecallRequest {
850            query: "RECALL memories LIMIT 10".into(),
851        });
852        r.metadata_mut()
853            .insert("authorization", "Bearer bad.tok".parse().unwrap());
854        assert_eq!(
855            svc.recall(r).await.unwrap_err().code(),
856            tonic::Code::Unauthenticated
857        );
858    }
859}