1use 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
31pub 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(¤t_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(¤t_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
168pub struct MemoryServiceImpl {
174 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#[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#[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 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 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 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 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 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 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}