1use std::collections::{HashMap, HashSet};
4use std::sync::RwLock;
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::{Duration, Instant, SystemTime};
7
8use agent_uri::AgentUri;
9use agent_uri_attestation::{SigningKey, Verifier};
10use async_trait::async_trait;
11
12use crate::expiration::Expirations;
13use crate::{
14 Cursor, Dht, DhtError, DhtKey, DhtStats, Endpoint, MatchMode, MigrationResult, Mutation,
15 MutationKind, MutationProof, NodeId, Page, PeerAddr, Query, Quorum, ReadOptions, Registration,
16 SimulationConfig, WriteOptions, WriteReceipt,
17};
18
19static NEXT_NODE_ID: AtomicU64 = AtomicU64::new(1);
23
24struct Indices {
35 records: HashMap<DhtKey, Vec<Registration>>,
37
38 keys_by_uri: HashMap<String, Vec<DhtKey>>,
40
41 path_by_identity: HashMap<String, String>,
43
44 expirations: Expirations,
47}
48
49impl Indices {
50 fn new() -> Self {
51 Self {
52 records: HashMap::new(),
53 keys_by_uri: HashMap::new(),
54 path_by_identity: HashMap::new(),
55 expirations: Expirations::default(),
56 }
57 }
58
59 fn clear(&mut self) {
60 self.records.clear();
61 self.keys_by_uri.clear();
62 self.path_by_identity.clear();
63 self.expirations.clear();
64 }
65
66 fn forget(&mut self, uri_str: &str) {
74 self.expirations.forget(uri_str);
75 if let Ok(uri) = AgentUri::parse(uri_str) {
76 self.path_by_identity
77 .remove(&SimulatedDht::identity_key(&uri));
78 }
79 let Some(keys) = self.keys_by_uri.remove(uri_str) else {
80 return;
81 };
82 for key in keys {
83 if let Some(registrations) = self.records.get_mut(&key) {
84 registrations.retain(|r| r.agent_uri().canonical() != *uri_str);
85 if registrations.is_empty() {
86 self.records.remove(&key);
87 }
88 }
89 }
90 }
91
92 fn evict_due(&mut self, now: SystemTime) -> usize {
97 let lapsed = self.expirations.take_due(now);
98 for uri_str in &lapsed {
99 self.forget(uri_str);
100 }
101 lapsed.len()
102 }
103
104 fn estimate_memory_usage(&self) -> usize {
105 let key_bytes = self.records.len() * (32 + 64);
112 let registration_bytes = self.keys_by_uri.len() * 500;
113 let uri_index_bytes = self
114 .keys_by_uri
115 .values()
116 .map(|keys| 100 + keys.len() * 32 + 64)
117 .sum::<usize>();
118 let expiry_bytes = self.expirations.queued() * (16 + 100 + 64);
123
124 key_bytes + registration_bytes + uri_index_bytes + expiry_bytes
125 }
126}
127
128pub struct SimulatedDht {
178 indices: RwLock<Indices>,
180
181 verifier: Verifier,
183
184 node_id: NodeId,
186
187 config: SimulationConfig,
189}
190
191impl SimulatedDht {
192 fn next_node_id() -> NodeId {
193 NodeId::from_bytes(NEXT_NODE_ID.fetch_add(1, Ordering::Relaxed).to_be_bytes())
194 }
195
196 fn receipt(&self, quorum: Quorum) -> WriteReceipt {
199 let required = quorum.resolve(self.config.replication_factor);
200 WriteReceipt::new(self.config.replication_factor.get(), required)
201 }
202
203 fn identity_key(uri: &AgentUri) -> String {
204 format!("{}/{}", uri.trust_root(), uri.agent_id())
205 }
206
207 fn ancestor_keys(uri: &AgentUri) -> Vec<DhtKey> {
208 (1..=uri.capability_path().depth())
209 .filter_map(|depth| {
210 DhtKey::derive_at_depth(uri.trust_root(), uri.capability_path(), depth)
211 })
212 .collect()
213 }
214
215 #[must_use]
217 pub fn new(config: SimulationConfig) -> Self {
218 Self {
219 indices: RwLock::new(Indices::new()),
220 verifier: Verifier::new(),
221 node_id: Self::next_node_id(),
222 config,
223 }
224 }
225
226 #[must_use]
228 pub fn with_verifier(config: SimulationConfig, verifier: Verifier) -> Self {
229 Self {
230 indices: RwLock::new(Indices::new()),
231 verifier,
232 node_id: Self::next_node_id(),
233 config,
234 }
235 }
236
237 #[must_use]
239 pub fn with_defaults() -> Self {
240 Self::new(SimulationConfig::default())
241 }
242
243 #[must_use]
245 pub const fn config(&self) -> &SimulationConfig {
246 &self.config
247 }
248
249 pub fn register_batch(
262 &self,
263 registrations: Vec<(Registration, MutationProof)>,
264 ) -> Result<usize, DhtError> {
265 let mut count = 0;
266 for (registration, proof) in registrations {
267 if self.do_register(®istration, &proof).is_ok() {
268 count += 1;
269 }
270 }
271 Ok(count)
272 }
273
274 #[must_use]
280 pub fn stats(&self) -> DhtStats {
281 let indices = self.indices.read().expect("lock poisoned");
282
283 let total_registrations = indices.keys_by_uri.len();
284 let unique_keys = indices.records.len();
285 let unique_trust_roots = indices
289 .keys_by_uri
290 .keys()
291 .filter_map(|uri| AgentUri::parse(uri).ok())
292 .map(|uri| uri.trust_root().as_str().to_string())
293 .collect::<HashSet<_>>()
294 .len();
295
296 let max_registrations_per_key = indices.records.values().map(Vec::len).max().unwrap_or(0);
297
298 let indexed_copies = indices.records.values().map(Vec::len).sum::<usize>();
300 let total_u32 = u32::try_from(indexed_copies).unwrap_or(u32::MAX);
301 let keys_u32 = u32::try_from(unique_keys).unwrap_or(u32::MAX);
302 let avg_registrations_per_key = if keys_u32 > 0 {
303 f64::from(total_u32) / f64::from(keys_u32)
304 } else {
305 0.0
306 };
307
308 let memory_bytes = indices.estimate_memory_usage();
310
311 DhtStats {
312 total_registrations,
313 unique_keys,
314 unique_trust_roots,
315 max_registrations_per_key,
316 avg_registrations_per_key,
317 path_depth_histogram: Vec::new(),
318 memory_bytes,
319 }
320 }
321
322 pub fn simulate_migration(
337 &self,
338 agent_uri: &AgentUri,
339 new_endpoint: Endpoint,
340 agent_key: &SigningKey,
341 ) -> Result<MigrationResult, DhtError> {
342 let uri_str = agent_uri.canonical();
343
344 let current = {
349 let indices = self.indices.read().expect("lock poisoned");
350 let keys = indices
351 .keys_by_uri
352 .get(&uri_str)
353 .ok_or_else(|| DhtError::not_found(&uri_str))?;
354 let key = keys.last().ok_or_else(|| DhtError::not_found(&uri_str))?;
355
356 let registrations = indices
357 .records
358 .get(key)
359 .ok_or_else(|| DhtError::not_found(&uri_str))?;
360
361 registrations
362 .iter()
363 .find(|r| r.agent_uri().canonical() == uri_str)
364 .cloned()
365 .ok_or_else(|| DhtError::not_found(&uri_str))?
366 };
367 let old_endpoints = current.endpoints().to_vec();
368
369 let endpoints = std::slice::from_ref(&new_endpoint);
370 let proof = MutationProof::sign_next(
373 agent_key,
374 ¤t,
375 &Mutation::UpdateEndpoint {
376 endpoints,
377 expires_at: current.expires_at(),
378 },
379 );
380
381 let start = Instant::now();
383 self.apply_mutation(agent_uri, &proof, |record| Mutation::UpdateEndpoint {
384 endpoints,
385 expires_at: record.expires_at(),
386 })?;
387 let update_latency = start.elapsed();
388
389 Ok(MigrationResult::success(
390 uri_str,
391 old_endpoints,
392 vec![new_endpoint],
393 update_latency,
394 ))
395 }
396
397 pub fn clear(&self) {
403 self.indices.write().expect("lock poisoned").clear();
404 }
405
406 pub fn expire_stale(&self) -> usize {
425 self.indices
426 .write()
427 .expect("lock poisoned")
428 .evict_due(SystemTime::now())
429 }
430}
431
432impl SimulatedDht {
433 fn do_register(
436 &self,
437 registration: &Registration,
438 proof: &MutationProof,
439 ) -> Result<(), DhtError> {
440 if registration.endpoints().is_empty() {
441 return Err(DhtError::NoEndpoints);
442 }
443
444 let uri_str = registration.agent_uri().canonical();
445 let keys = Self::ancestor_keys(registration.agent_uri());
446 let exact_key = DhtKey::derive(
447 registration.agent_uri().trust_root(),
448 registration.agent_uri().capability_path(),
449 );
450 let identity = Self::identity_key(registration.agent_uri());
451
452 proof.authorizes_registration(registration)?;
456
457 if self.config.verify_attestations {
458 let token = registration.attestation().ok_or_else(|| {
459 DhtError::invalid_attestation(&uri_str, "registration has no attestation token")
460 })?;
461 let claims = self
462 .verifier
463 .verify_for_capability(
464 token,
465 registration.agent_uri(),
466 registration.agent_uri().capability_path(),
467 )
468 .map_err(|error| DhtError::invalid_attestation(&uri_str, error.to_string()))?;
469
470 let attested = claims
474 .agent_verifying_key()
475 .map_err(|error| DhtError::invalid_attestation(&uri_str, error.to_string()))?;
476 if attested != *registration.agent_key() {
477 return Err(DhtError::AgentKeyMismatch { agent_uri: uri_str });
478 }
479 }
480
481 {
485 let mut indices = self.indices.write().expect("lock poisoned");
486
487 if self.config.auto_expire {
493 indices.evict_due(SystemTime::now());
494 }
495
496 if indices.keys_by_uri.contains_key(&uri_str) {
497 return Err(DhtError::already_registered(&uri_str));
498 }
499 if let Some(existing_path) = indices.path_by_identity.get(&identity)
500 && existing_path != registration.agent_uri().capability_path().as_str()
501 {
502 return Err(DhtError::IdentityCapabilityConflict {
503 identity,
504 existing_path: existing_path.clone(),
505 requested_path: registration
506 .agent_uri()
507 .capability_path()
508 .as_str()
509 .to_string(),
510 });
511 }
512 if let Some(registrations) = indices.records.get(&exact_key)
520 && registrations.len() >= self.config.max_registrations_per_capability
521 {
522 return Err(DhtError::key_capacity_exceeded(
523 format!("{exact_key}"),
524 self.config.max_registrations_per_capability,
525 ));
526 }
527
528 for key in &keys {
532 indices
533 .records
534 .entry(*key)
535 .or_default()
536 .push(registration.clone());
537 }
538
539 indices.expirations.set(&uri_str, registration.expires_at());
540 indices.keys_by_uri.insert(uri_str, keys);
541 indices.path_by_identity.insert(
542 identity,
543 registration
544 .agent_uri()
545 .capability_path()
546 .as_str()
547 .to_string(),
548 );
549 }
550
551 Ok(())
552 }
553
554 fn rewrite_copies(
559 records: &mut HashMap<DhtKey, Vec<Registration>>,
560 keys: &[DhtKey],
561 uri_str: &str,
562 sequence: u64,
563 mut mutate: impl FnMut(&mut Registration),
564 ) -> Result<(), DhtError> {
565 for key in keys {
566 let registrations = records
567 .get_mut(key)
568 .ok_or_else(|| DhtError::not_found(uri_str))?;
569 let registration = registrations
570 .iter_mut()
571 .find(|r| r.agent_uri().canonical() == uri_str)
572 .ok_or_else(|| DhtError::not_found(uri_str))?;
573 mutate(registration);
574 registration.set_sequence(sequence);
575 }
576 Ok(())
577 }
578
579 fn apply_mutation<'m>(
594 &self,
595 agent_uri: &AgentUri,
596 proof: &MutationProof,
597 describe: impl FnOnce(&Registration) -> Mutation<'m>,
598 ) -> Result<(), DhtError> {
599 let uri_str = agent_uri.canonical();
600
601 let mut guard = self.indices.write().expect("lock poisoned");
602 let indices = &mut *guard;
603
604 let keys = indices
605 .keys_by_uri
606 .get(&uri_str)
607 .cloned()
608 .ok_or_else(|| DhtError::not_found(&uri_str))?;
609
610 let current = keys
613 .iter()
614 .filter_map(|key| indices.records.get(key))
615 .flatten()
616 .find(|r| r.agent_uri().canonical() == uri_str)
617 .cloned()
618 .ok_or_else(|| DhtError::not_found(&uri_str))?;
619
620 if current.is_expired() && self.config.auto_expire {
621 return Err(DhtError::expired(&uri_str));
622 }
623
624 let mutation = describe(¤t);
625 proof.authorizes(¤t, &mutation)?;
626
627 match mutation {
628 Mutation::UpdateEndpoint {
632 endpoints,
633 expires_at,
634 } => {
635 Self::rewrite_copies(
636 &mut indices.records,
637 &keys,
638 &uri_str,
639 proof.sequence(),
640 |registration| {
641 registration.update_endpoints(endpoints.to_vec());
642 registration.set_expires_at(expires_at);
643 },
644 )?;
645 indices.expirations.set(&uri_str, expires_at);
646 }
647 Mutation::Refresh { expires_at, .. } => {
652 Self::rewrite_copies(
653 &mut indices.records,
654 &keys,
655 &uri_str,
656 proof.sequence(),
657 |registration| registration.set_expires_at(expires_at),
658 )?;
659 indices.expirations.set(&uri_str, expires_at);
660 }
661 Mutation::Deregister => indices.forget(&uri_str),
662 Mutation::Register { .. } => {
665 return Err(DhtError::Unauthorized {
666 agent_uri: uri_str,
667 operation: MutationKind::Register,
668 });
669 }
670 }
671
672 if self.config.auto_expire {
678 indices.evict_due(SystemTime::now());
679 }
680
681 Ok(())
682 }
683
684 fn do_lookup(
685 &self,
686 query: &Query,
687 options: &ReadOptions,
688 ) -> Result<Page<Registration>, DhtError> {
689 let key = DhtKey::derive(query.trust_root(), query.capability_path());
690
691 let mut matches: Vec<Registration> = {
692 let indices = self.indices.read().expect("lock poisoned");
693 indices
694 .records
695 .get(&key)
696 .map(|registrations| {
697 registrations
698 .iter()
699 .filter(|r| match query.match_mode() {
700 MatchMode::Prefix => true,
704 MatchMode::Exact => {
705 r.agent_uri().capability_path() == query.capability_path()
706 }
707 })
708 .filter(|r| !r.is_expired() || !self.config.auto_expire)
709 .cloned()
710 .collect()
711 })
712 .unwrap_or_default()
713 };
714
715 matches.sort_by_key(|r| r.agent_uri().canonical());
718
719 let offset = options.cursor.map_or(0, |c| c.position() as usize);
720 if offset >= matches.len() {
721 return Ok(Page::complete(Vec::new()));
722 }
723
724 let page_size = self.config.page_size.get();
725 let end = offset.saturating_add(page_size).min(matches.len());
726 let items = matches[offset..end].to_vec();
727
728 if end < matches.len() {
729 let next = u32::try_from(end).map_err(|_| DhtError::QuorumFailed {
730 achieved: 0,
731 required: 1,
732 })?;
733 Ok(Page::partial(items, Cursor::at(next)))
734 } else {
735 Ok(Page::complete(items))
736 }
737 }
738}
739
740#[async_trait]
741impl Dht for SimulatedDht {
742 async fn register(
743 &self,
744 registration: Registration,
745 proof: &MutationProof,
746 options: WriteOptions,
747 ) -> Result<WriteReceipt, DhtError> {
748 self.do_register(®istration, proof)?;
749 Ok(self.receipt(options.quorum))
750 }
751
752 async fn update_endpoint(
753 &self,
754 agent_uri: &AgentUri,
755 new_endpoints: Vec<Endpoint>,
756 proof: &MutationProof,
757 options: WriteOptions,
758 ) -> Result<WriteReceipt, DhtError> {
759 if new_endpoints.is_empty() {
762 return Err(DhtError::NoEndpoints);
763 }
764 self.apply_mutation(agent_uri, proof, |record| Mutation::UpdateEndpoint {
765 endpoints: &new_endpoints,
766 expires_at: record.expires_at(),
767 })?;
768 Ok(self.receipt(options.quorum))
769 }
770
771 async fn refresh(
772 &self,
773 agent_uri: &AgentUri,
774 ttl: Duration,
775 expires_at: SystemTime,
776 proof: &MutationProof,
777 options: WriteOptions,
778 ) -> Result<WriteReceipt, DhtError> {
779 self.apply_mutation(agent_uri, proof, |_| Mutation::Refresh { ttl, expires_at })?;
780 Ok(self.receipt(options.quorum))
781 }
782
783 async fn deregister(
784 &self,
785 agent_uri: &AgentUri,
786 proof: &MutationProof,
787 _options: WriteOptions,
788 ) -> Result<(), DhtError> {
789 self.apply_mutation(agent_uri, proof, |_| Mutation::Deregister)
790 }
791
792 async fn lookup(
793 &self,
794 query: &Query,
795 options: &ReadOptions,
796 ) -> Result<Page<Registration>, DhtError> {
797 self.do_lookup(query, options)
798 }
799
800 fn local_id(&self) -> NodeId {
801 self.node_id.clone()
802 }
803
804 async fn bootstrap(&self, peers: &[PeerAddr]) -> Result<(), DhtError> {
805 if peers.is_empty() {
809 return Err(DhtError::NoPeers);
810 }
811 Ok(())
812 }
813}
814
815#[cfg(test)]
816mod tests {
817 use super::*;
818 use crate::RecordVersion;
819 use crate::registration::system_time_to_millis;
822 use agent_uri::{CapabilityPath, TrustRoot};
823 use agent_uri_attestation::{AcceptAll, Issuer};
824 use futures::executor::block_on;
825 use std::num::NonZeroUsize;
826 use std::time::{Duration, SystemTime};
827
828 fn unverified_dht() -> SimulatedDht {
829 SimulatedDht::new(SimulationConfig::default().with_verify_attestations(false))
830 }
831
832 fn test_uri(suffix: &str) -> AgentUri {
833 AgentUri::parse(&format!(
834 "agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn0{suffix}"
835 ))
836 .unwrap()
837 }
838
839 fn uri_at(path: &str, suffix: &str) -> AgentUri {
842 AgentUri::parse(&format!(
843 "agent://anthropic.com/{path}/llm_01h455vb4pex5vsknk084sn0{suffix}"
844 ))
845 .unwrap()
846 }
847
848 fn dht_with_capacity_for_one() -> SimulatedDht {
851 SimulatedDht::new(
852 SimulationConfig::new()
853 .with_max_registrations_per_capability(1)
854 .with_verify_attestations(false),
855 )
856 }
857
858 fn test_endpoint() -> Endpoint {
859 Endpoint::https("agent.anthropic.com:443")
860 }
861
862 fn registration(uri: &AgentUri, key: &SigningKey) -> Registration {
863 Registration::new(uri.clone(), key.verifying_key(), vec![test_endpoint()])
864 }
865
866 fn register(
872 dht: &SimulatedDht,
873 registration: Registration,
874 key: &SigningKey,
875 ) -> Result<WriteReceipt, DhtError> {
876 let proof = MutationProof::sign_registration(key, ®istration);
877 block_on(dht.register(registration, &proof, WriteOptions::default()))
878 }
879
880 fn record(dht: &SimulatedDht, uri: &AgentUri) -> Registration {
883 stored(dht, uri).expect("record is registered")
884 }
885
886 fn stored(dht: &SimulatedDht, uri: &AgentUri) -> Option<Registration> {
887 lookup(
888 dht,
889 &Query::exact(uri.trust_root().clone(), uri.capability_path().clone()),
890 )
891 .into_iter()
892 .find(|r| r.agent_uri() == uri)
893 }
894
895 fn version_at(dht: &SimulatedDht, uri: &AgentUri, sequence: u64) -> RecordVersion {
900 RecordVersion::new(
901 stored(dht, uri).map_or_else(SystemTime::now, |r| r.registered_at()),
902 sequence,
903 )
904 }
905
906 fn deregister(
907 dht: &SimulatedDht,
908 uri: &AgentUri,
909 key: &SigningKey,
910 sequence: u64,
911 ) -> Result<(), DhtError> {
912 let proof = MutationProof::sign(
913 key,
914 uri,
915 version_at(dht, uri, sequence),
916 &Mutation::Deregister,
917 );
918 block_on(dht.deregister(uri, &proof, WriteOptions::default()))
919 }
920
921 fn expiry_of(dht: &SimulatedDht, uri: &AgentUri) -> SystemTime {
924 stored(dht, uri).map_or_else(SystemTime::now, |r| r.expires_at())
925 }
926
927 fn update_endpoint(
928 dht: &SimulatedDht,
929 uri: &AgentUri,
930 key: &SigningKey,
931 sequence: u64,
932 endpoints: Vec<Endpoint>,
933 ) -> Result<WriteReceipt, DhtError> {
934 let proof = MutationProof::sign(
935 key,
936 uri,
937 version_at(dht, uri, sequence),
938 &Mutation::UpdateEndpoint {
939 endpoints: &endpoints,
940 expires_at: expiry_of(dht, uri),
941 },
942 );
943 block_on(dht.update_endpoint(uri, endpoints, &proof, WriteOptions::default()))
944 }
945
946 fn refresh(
947 dht: &SimulatedDht,
948 uri: &AgentUri,
949 key: &SigningKey,
950 sequence: u64,
951 ttl: Duration,
952 ) -> Result<WriteReceipt, DhtError> {
953 let expires_at = SystemTime::now() + ttl;
957 let proof = MutationProof::sign(
958 key,
959 uri,
960 version_at(dht, uri, sequence),
961 &Mutation::Refresh { ttl, expires_at },
962 );
963 block_on(dht.refresh(uri, ttl, expires_at, &proof, WriteOptions::default()))
964 }
965
966 fn lookup(dht: &SimulatedDht, query: &Query) -> Vec<Registration> {
967 block_on(dht.lookup(query, &ReadOptions::default()))
968 .unwrap()
969 .into_items()
970 }
971
972 fn exact(trust_root: &str, path: &str) -> Query {
973 Query::exact(
974 TrustRoot::parse(trust_root).unwrap(),
975 CapabilityPath::parse(path).unwrap(),
976 )
977 }
978
979 fn prefix(trust_root: &str, path: &str) -> Query {
980 Query::prefix(
981 TrustRoot::parse(trust_root).unwrap(),
982 CapabilityPath::parse(path).unwrap(),
983 )
984 }
985
986 #[test]
987 fn register_and_lookup_exact() {
988 let dht = unverified_dht();
989 let key = SigningKey::generate();
990 let uri = test_uri("2q");
991
992 register(&dht, registration(&uri, &key), &key).unwrap();
993
994 let results = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
995
996 assert_eq!(results.len(), 1);
997 assert_eq!(results[0].agent_uri(), &uri);
998 }
999
1000 #[test]
1001 fn register_empty_endpoints_fails() {
1002 let dht = unverified_dht();
1003 let key = SigningKey::generate();
1004 let registration = Registration::new(test_uri("2q"), key.verifying_key(), vec![]);
1005
1006 let result = register(&dht, registration, &key);
1007
1008 assert!(matches!(result, Err(DhtError::NoEndpoints)));
1009 }
1010
1011 #[test]
1012 fn double_registration_fails() {
1013 let dht = unverified_dht();
1014 let key = SigningKey::generate();
1015 let uri = test_uri("2q");
1016
1017 register(&dht, registration(&uri, &key), &key).unwrap();
1018
1019 let result = register(&dht, registration(&uri, &key), &key);
1020
1021 assert!(matches!(result, Err(DhtError::AlreadyRegistered { .. })));
1022 }
1023
1024 #[test]
1025 fn deregister_removes_agent() {
1026 let dht = unverified_dht();
1027 let key = SigningKey::generate();
1028 let uri = test_uri("2q");
1029
1030 register(&dht, registration(&uri, &key), &key).unwrap();
1031 deregister(&dht, &uri, &key, 1).unwrap();
1032
1033 let results = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
1034
1035 assert!(results.is_empty());
1036 }
1037
1038 #[test]
1039 fn lookup_prefix_finds_descendants() {
1040 let dht = unverified_dht();
1041 let key = SigningKey::generate();
1042
1043 let uri1 =
1044 AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1045 .unwrap();
1046 let uri2 =
1047 AgentUri::parse("agent://anthropic.com/assistant/code/llm_01h455vb4pex5vsknk084sn02r")
1048 .unwrap();
1049
1050 register(&dht, registration(&uri1, &key), &key).unwrap();
1051 register(&dht, registration(&uri2, &key), &key).unwrap();
1052
1053 let results = lookup(&dht, &prefix("anthropic.com", "assistant"));
1054
1055 assert_eq!(results.len(), 2);
1056 }
1057
1058 #[test]
1059 fn lookup_is_scoped_to_one_trust_root() {
1060 let dht = unverified_dht();
1061 let key = SigningKey::generate();
1062
1063 let uri1 =
1064 AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1065 .unwrap();
1066 let uri2 =
1067 AgentUri::parse("agent://openai.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02r")
1068 .unwrap();
1069
1070 register(&dht, registration(&uri1, &key), &key).unwrap();
1071 register(&dht, registration(&uri2, &key), &key).unwrap();
1072
1073 let anthropic = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
1077 let openai = lookup(&dht, &exact("openai.com", "assistant/chat"));
1078
1079 assert_eq!(anthropic.len(), 1);
1080 assert_eq!(
1081 anthropic[0].agent_uri().trust_root().as_str(),
1082 "anthropic.com"
1083 );
1084 assert_eq!(openai.len(), 1);
1085 assert_eq!(openai[0].agent_uri().trust_root().as_str(), "openai.com");
1086 }
1087
1088 #[test]
1089 fn update_endpoint_changes_endpoints() {
1090 let dht = unverified_dht();
1091 let key = SigningKey::generate();
1092 let uri = test_uri("2q");
1093
1094 register(&dht, registration(&uri, &key), &key).unwrap();
1095
1096 let new_endpoint = Endpoint::grpc("agent.anthropic.com:50051");
1097 update_endpoint(&dht, &uri, &key, 1, vec![new_endpoint.clone()]).unwrap();
1098
1099 let results = lookup(&dht, &exact("anthropic.com", "assistant/chat"));
1100
1101 assert_eq!(results.len(), 1);
1102 assert_eq!(results[0].endpoints(), &[new_endpoint]);
1103 }
1104
1105 #[test]
1106 fn stats_reports_correct_counts() {
1107 let dht = unverified_dht();
1108 let key = SigningKey::generate();
1109
1110 let uri1 = test_uri("2q");
1111 let uri2 =
1112 AgentUri::parse("agent://anthropic.com/assistant/code/llm_01h455vb4pex5vsknk084sn02r")
1113 .unwrap();
1114
1115 register(&dht, registration(&uri1, &key), &key).unwrap();
1116 register(&dht, registration(&uri2, &key), &key).unwrap();
1117
1118 let stats = dht.stats();
1119 assert_eq!(stats.total_registrations(), 2);
1120 assert_eq!(stats.unique_trust_roots(), 1);
1121 }
1122
1123 #[test]
1126 fn a_second_agent_at_a_full_capability_path_is_refused() {
1127 let dht = dht_with_capacity_for_one();
1128 let key = SigningKey::generate();
1129
1130 register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1131
1132 let result = register(&dht, registration(&test_uri("2r"), &key), &key);
1133
1134 assert!(matches!(result, Err(DhtError::KeyCapacityExceeded { .. })));
1135 }
1136
1137 #[test]
1141 fn a_full_ancestor_does_not_refuse_a_sibling_path() {
1142 let dht = dht_with_capacity_for_one();
1143 let key = SigningKey::generate();
1144
1145 let occupant = uri_at("assistant/chat", "2q");
1146 let sibling = uri_at("assistant/code", "2r");
1147 register(&dht, registration(&occupant, &key), &key).unwrap();
1148
1149 register(&dht, registration(&sibling, &key), &key)
1150 .expect("a path of its own has capacity of its own");
1151 }
1152
1153 #[test]
1156 fn a_full_ancestor_does_not_refuse_a_path_beneath_it() {
1157 let dht = dht_with_capacity_for_one();
1158 let key = SigningKey::generate();
1159
1160 register(
1161 &dht,
1162 registration(&uri_at("assistant/chat", "2q"), &key),
1163 &key,
1164 )
1165 .unwrap();
1166
1167 register(
1168 &dht,
1169 registration(&uri_at("assistant/chat/vision", "2r"), &key),
1170 &key,
1171 )
1172 .expect("an empty key is not full because its ancestor is");
1173 }
1174
1175 #[test]
1180 fn an_ancestor_key_holds_more_than_one_path_may() {
1181 let dht = dht_with_capacity_for_one();
1182 let key = SigningKey::generate();
1183
1184 register(
1185 &dht,
1186 registration(&uri_at("assistant/chat", "2q"), &key),
1187 &key,
1188 )
1189 .unwrap();
1190 register(
1191 &dht,
1192 registration(&uri_at("assistant/code", "2r"), &key),
1193 &key,
1194 )
1195 .unwrap();
1196
1197 let under_assistant = lookup(&dht, &prefix("anthropic.com", "assistant"));
1198
1199 assert_eq!(under_assistant.len(), 2);
1200 assert!(dht.stats().max_registrations_per_key() > 1);
1201 }
1202
1203 #[test]
1204 fn clear_removes_all() {
1205 let dht = unverified_dht();
1206 let key = SigningKey::generate();
1207
1208 register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1209
1210 dht.clear();
1211
1212 let stats = dht.stats();
1213 assert_eq!(stats.total_registrations(), 0);
1214 }
1215
1216 #[test]
1217 fn register_batch_counts_successes() {
1218 let dht = unverified_dht();
1219 let key = SigningKey::generate();
1220
1221 let registrations = vec![
1222 registration(&test_uri("2q"), &key),
1223 registration(&test_uri("2r"), &key),
1224 Registration::new(test_uri("2s"), key.verifying_key(), vec![]),
1226 ]
1227 .into_iter()
1228 .map(|r| {
1229 let proof = MutationProof::sign_registration(&key, &r);
1230 (r, proof)
1231 })
1232 .collect();
1233
1234 let count = dht.register_batch(registrations).unwrap();
1235
1236 assert_eq!(count, 2);
1237 }
1238
1239 #[test]
1240 fn secure_default_rejects_missing_attestation() {
1241 let dht = SimulatedDht::with_defaults();
1242 let key = SigningKey::generate();
1243
1244 let result = register(&dht, registration(&test_uri("2q"), &key), &key);
1245
1246 assert!(matches!(result, Err(DhtError::InvalidAttestation { .. })));
1247 }
1248
1249 #[test]
1250 fn trusted_attestation_is_verified_before_registration() {
1251 let signing_key = SigningKey::generate();
1252 let issuer = Issuer::new(
1253 "anthropic.com",
1254 signing_key.clone(),
1255 Duration::from_hours(1),
1256 );
1257 let mut verifier = Verifier::new().with_revocation(AcceptAll);
1258 verifier.add_trusted_root("anthropic.com", signing_key.verifying_key());
1259 let dht = SimulatedDht::with_verifier(SimulationConfig::default(), verifier);
1260 let agent_key = SigningKey::generate();
1261 let uri = test_uri("2q");
1262 let token = issuer
1263 .issue(
1264 &uri,
1265 &agent_key.verifying_key(),
1266 vec!["assistant/chat".to_string()],
1267 )
1268 .unwrap();
1269
1270 register(
1271 &dht,
1272 registration(&uri, &agent_key).with_attestation(token),
1273 &agent_key,
1274 )
1275 .unwrap();
1276
1277 let results = lookup(
1278 &dht,
1279 &Query::exact(uri.trust_root().clone(), uri.capability_path().clone()),
1280 );
1281 assert_eq!(results.len(), 1);
1282 }
1283
1284 fn attesting_dht(agent_key: &SigningKey, uri: &AgentUri) -> (SimulatedDht, String) {
1286 let root_key = SigningKey::generate();
1287 let issuer = Issuer::new("anthropic.com", root_key.clone(), Duration::from_hours(1));
1288 let mut verifier = Verifier::new().with_revocation(AcceptAll);
1289 verifier.add_trusted_root("anthropic.com", root_key.verifying_key());
1290 let token = issuer
1291 .issue(
1292 uri,
1293 &agent_key.verifying_key(),
1294 vec!["assistant/chat".to_string()],
1295 )
1296 .unwrap();
1297 (
1298 SimulatedDht::with_verifier(SimulationConfig::default(), verifier),
1299 token,
1300 )
1301 }
1302
1303 #[test]
1304 fn a_lifted_token_does_not_register_the_thief_at_its_own_endpoints() {
1305 let victim = SigningKey::generate();
1310 let attacker = SigningKey::generate();
1311 let uri = test_uri("2q");
1312 let (dht, lifted) = attesting_dht(&victim, &uri);
1313
1314 let hijack = Registration::new(
1315 uri.clone(),
1316 attacker.verifying_key(),
1317 vec![Endpoint::https("attacker.example")],
1318 )
1319 .with_attestation(lifted);
1320
1321 let result = register(&dht, hijack, &attacker);
1322
1323 assert!(matches!(result, Err(DhtError::AgentKeyMismatch { .. })));
1324 }
1325
1326 #[test]
1327 fn a_lifted_token_does_not_register_the_thief_under_the_victims_key_either() {
1328 let victim = SigningKey::generate();
1332 let attacker = SigningKey::generate();
1333 let uri = test_uri("2q");
1334 let (dht, lifted) = attesting_dht(&victim, &uri);
1335
1336 let hijack = Registration::new(
1337 uri.clone(),
1338 victim.verifying_key(),
1339 vec![Endpoint::https("attacker.example")],
1340 )
1341 .with_attestation(lifted);
1342
1343 let result = register(&dht, hijack, &attacker);
1344
1345 assert!(matches!(
1346 result,
1347 Err(DhtError::Unauthorized {
1348 operation: MutationKind::Register,
1349 ..
1350 })
1351 ));
1352 }
1353
1354 #[test]
1355 fn the_attested_agent_registers_normally() {
1356 let agent = SigningKey::generate();
1357 let uri = test_uri("2q");
1358 let (dht, token) = attesting_dht(&agent, &uri);
1359
1360 register(
1361 &dht,
1362 registration(&uri, &agent).with_attestation(token),
1363 &agent,
1364 )
1365 .unwrap();
1366
1367 assert_eq!(
1368 lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
1369 1
1370 );
1371 }
1372
1373 #[test]
1374 fn a_registration_proof_does_not_cover_endpoints_it_did_not_sign() {
1375 let dht = unverified_dht();
1379 let agent = SigningKey::generate();
1380 let uri = test_uri("2q");
1381 let honest = Registration::new(
1382 uri.clone(),
1383 agent.verifying_key(),
1384 vec![Endpoint::https("honest.example")],
1385 );
1386 let proof = MutationProof::sign_registration(&agent, &honest);
1387
1388 let substituted = Registration::new(
1389 uri,
1390 agent.verifying_key(),
1391 vec![Endpoint::https("attacker.example")],
1392 )
1393 .with_registered_at(honest.registered_at());
1394
1395 let result = block_on(dht.register(substituted, &proof, WriteOptions::default()));
1396
1397 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1398 }
1399
1400 #[test]
1401 fn a_registration_proof_does_not_cover_a_different_expiry() {
1402 let dht = unverified_dht();
1403 let agent = SigningKey::generate();
1404 let uri = test_uri("2q");
1405 let honest = registration(&uri, &agent).with_ttl(Duration::from_mins(1));
1406 let proof = MutationProof::sign_registration(&agent, &honest);
1407
1408 let stretched = honest.clone().with_ttl(Duration::from_hours(24));
1409
1410 let result = block_on(dht.register(stretched, &proof, WriteOptions::default()));
1411
1412 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1413 }
1414
1415 #[test]
1416 fn a_registration_proof_does_not_transfer_to_another_agent() {
1417 let dht = unverified_dht();
1418 let agent = SigningKey::generate();
1419 let mine = registration(&test_uri("2q"), &agent);
1420 let proof = MutationProof::sign_registration(&agent, &mine);
1421
1422 let theirs = registration(&test_uri("2r"), &agent).with_registered_at(mine.registered_at());
1423
1424 let result = block_on(dht.register(theirs, &proof, WriteOptions::default()));
1425
1426 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1427 }
1428
1429 #[test]
1430 fn refresh_extends_an_expiring_registration() {
1431 let dht = unverified_dht();
1432 let key = SigningKey::generate();
1433 let uri = test_uri("2q");
1434 register(
1435 &dht,
1436 registration(&uri, &key).with_ttl(Duration::from_millis(50)),
1437 &key,
1438 )
1439 .unwrap();
1440
1441 let before = lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].expires_at();
1442 refresh(&dht, &uri, &key, 1, Duration::from_hours(1)).unwrap();
1443 let after = lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].expires_at();
1444
1445 assert!(
1446 after > before,
1447 "refresh must push expiry out, got {before:?} -> {after:?}"
1448 );
1449 }
1450
1451 #[test]
1452 fn a_refresh_lands_on_the_instant_the_agent_signed() {
1453 let dht = unverified_dht();
1459 let key = SigningKey::generate();
1460 let uri = test_uri("2q");
1461 register(&dht, registration(&uri, &key), &key).unwrap();
1462
1463 let ttl = Duration::from_hours(1);
1464 let chosen = SystemTime::now() + Duration::from_hours(9);
1467 let proof = MutationProof::sign(
1468 &key,
1469 &uri,
1470 version_at(&dht, &uri, 1),
1471 &Mutation::Refresh {
1472 ttl,
1473 expires_at: chosen,
1474 },
1475 );
1476 block_on(dht.refresh(&uri, ttl, chosen, &proof, WriteOptions::default())).unwrap();
1477
1478 let stored = lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].expires_at();
1479 assert_eq!(
1480 system_time_to_millis(stored),
1481 system_time_to_millis(chosen),
1482 "the store wrote an expiry the agent did not sign"
1483 );
1484 }
1485
1486 #[test]
1487 fn a_migration_leaves_the_expiry_where_it_was() {
1488 let dht = unverified_dht();
1493 let key = SigningKey::generate();
1494 let uri = test_uri("2q");
1495 register(&dht, registration(&uri, &key), &key).unwrap();
1496 let before = record(&dht, &uri).expires_at();
1497
1498 update_endpoint(
1499 &dht,
1500 &uri,
1501 &key,
1502 1,
1503 vec![Endpoint::https("eu-west-1.example")],
1504 )
1505 .unwrap();
1506
1507 assert_eq!(
1508 system_time_to_millis(record(&dht, &uri).expires_at()),
1509 system_time_to_millis(before)
1510 );
1511 }
1512
1513 #[test]
1514 fn refresh_updates_every_ancestor_copy() {
1515 let dht = unverified_dht();
1519 let key = SigningKey::generate();
1520 let uri = test_uri("2q");
1521 register(
1522 &dht,
1523 registration(&uri, &key).with_ttl(Duration::from_millis(50)),
1524 &key,
1525 )
1526 .unwrap();
1527
1528 refresh(&dht, &uri, &key, 1, Duration::from_hours(1)).unwrap();
1529
1530 let from_ancestor = lookup(&dht, &prefix("anthropic.com", "assistant"));
1531 assert_eq!(from_ancestor.len(), 1);
1532 assert!(!from_ancestor[0].is_expired());
1533 }
1534
1535 fn lapsed(uri: &AgentUri, key: &SigningKey) -> Registration {
1539 registration(uri, key).with_expires_at(SystemTime::now() - Duration::from_secs(1))
1540 }
1541
1542 fn held(dht: &SimulatedDht) -> usize {
1545 dht.stats().total_registrations()
1546 }
1547
1548 #[test]
1549 fn expire_stale_removes_a_lapsed_record_from_every_index() {
1550 let dht = unverified_dht();
1551 let key = SigningKey::generate();
1552 let uri = test_uri("2q");
1553 register(&dht, lapsed(&uri, &key), &key).unwrap();
1554
1555 assert_eq!(dht.expire_stale(), 1);
1556
1557 assert_eq!(held(&dht), 0, "the record is still indexed by URI");
1558 assert_eq!(dht.stats().unique_keys(), 0, "an ancestor key outlived it");
1559 assert!(lookup(&dht, &prefix("anthropic.com", "assistant")).is_empty());
1560 }
1561
1562 #[test]
1563 fn expire_stale_leaves_a_live_record_alone() {
1564 let dht = unverified_dht();
1565 let key = SigningKey::generate();
1566 let uri = test_uri("2q");
1567 register(&dht, registration(&uri, &key), &key).unwrap();
1568
1569 assert_eq!(dht.expire_stale(), 0);
1570 assert!(stored(&dht, &uri).is_some());
1571 }
1572
1573 #[test]
1574 fn a_write_evicts_what_has_lapsed_without_being_asked() {
1575 let dht = unverified_dht();
1579 let key = SigningKey::generate();
1580 register(&dht, lapsed(&test_uri("2q"), &key), &key).unwrap();
1581
1582 register(&dht, registration(&test_uri("2r"), &key), &key).unwrap();
1583
1584 assert_eq!(held(&dht), 1, "the lapsed record is still stored");
1585 }
1586
1587 #[test]
1588 fn a_lapsed_registration_does_not_hold_its_name() {
1589 let dht = unverified_dht();
1592 let key = SigningKey::generate();
1593 let uri = test_uri("2q");
1594 register(&dht, lapsed(&uri, &key), &key).unwrap();
1595
1596 register(&dht, registration(&uri, &key), &key)
1597 .expect("a lapsed record does not keep its URI registered");
1598
1599 assert_eq!(held(&dht), 1);
1600 assert!(stored(&dht, &uri).is_some_and(|r| !r.is_expired()));
1601 }
1602
1603 #[test]
1604 fn a_lapsed_registration_does_not_hold_its_capacity() {
1605 let dht = dht_with_capacity_for_one();
1608 let key = SigningKey::generate();
1609 register(&dht, lapsed(&test_uri("2q"), &key), &key).unwrap();
1610
1611 register(&dht, registration(&test_uri("2r"), &key), &key)
1612 .expect("a lapsed record does not keep its slot");
1613 }
1614
1615 #[test]
1616 fn an_expiry_moved_earlier_is_evicted_at_the_instant_it_names() {
1617 let dht = unverified_dht();
1621 let key = SigningKey::generate();
1622 let uri = test_uri("2q");
1623 register(&dht, registration(&uri, &key), &key).unwrap();
1624
1625 let ttl = Duration::from_hours(1);
1626 let lapsed_at = SystemTime::now() - Duration::from_secs(1);
1627 let proof = MutationProof::sign(
1628 &key,
1629 &uri,
1630 version_at(&dht, &uri, 1),
1631 &Mutation::Refresh {
1632 ttl,
1633 expires_at: lapsed_at,
1634 },
1635 );
1636 block_on(dht.refresh(&uri, ttl, lapsed_at, &proof, WriteOptions::default())).unwrap();
1637
1638 assert_eq!(held(&dht), 0, "the record lapsed the moment it was written");
1639 }
1640
1641 #[test]
1642 fn a_renewed_record_outlives_the_deadline_it_replaced() {
1643 let dht = unverified_dht();
1648 let key = SigningKey::generate();
1649 let uri = test_uri("2q");
1650 register(
1651 &dht,
1652 registration(&uri, &key).with_ttl(Duration::from_millis(40)),
1653 &key,
1654 )
1655 .unwrap();
1656
1657 refresh(&dht, &uri, &key, 1, Duration::from_hours(1)).unwrap();
1658 std::thread::sleep(Duration::from_millis(60));
1659
1660 assert_eq!(dht.expire_stale(), 0, "the renewed record was evicted");
1661 assert!(stored(&dht, &uri).is_some());
1662 }
1663
1664 #[test]
1665 fn expiry_waits_to_be_asked_when_auto_expire_is_off() {
1666 let dht = SimulatedDht::new(
1670 SimulationConfig::new()
1671 .with_verify_attestations(false)
1672 .with_auto_expire(false),
1673 );
1674 let key = SigningKey::generate();
1675 register(&dht, lapsed(&test_uri("2q"), &key), &key).unwrap();
1676
1677 register(&dht, registration(&test_uri("2r"), &key), &key).unwrap();
1678 assert_eq!(held(&dht), 2, "a write evicted what the caller kept");
1679
1680 assert_eq!(dht.expire_stale(), 1);
1681 assert_eq!(held(&dht), 1);
1682 }
1683
1684 #[test]
1685 fn every_stored_registration_has_a_deadline() {
1686 let dht = unverified_dht();
1691 let key = SigningKey::generate();
1692 for suffix in ["2q", "2r", "2s"] {
1693 register(&dht, registration(&test_uri(suffix), &key), &key).unwrap();
1694 }
1695 refresh(&dht, &test_uri("2q"), &key, 1, Duration::from_hours(2)).unwrap();
1696 deregister(&dht, &test_uri("2s"), &key, 1).unwrap();
1697 update_endpoint(
1698 &dht,
1699 &test_uri("2r"),
1700 &key,
1701 1,
1702 vec![Endpoint::https("eu-west-1.example")],
1703 )
1704 .unwrap();
1705
1706 let indices = dht.indices.read().expect("lock poisoned");
1707 assert_eq!(
1708 indices.expirations.tracked(),
1709 indices.keys_by_uri.len(),
1710 "the schedule and the store disagree about what is registered"
1711 );
1712 }
1713
1714 #[test]
1715 fn refresh_of_unknown_agent_is_not_found() {
1716 let dht = unverified_dht();
1717 let key = SigningKey::generate();
1718 let result = refresh(&dht, &test_uri("2q"), &key, 1, Duration::from_hours(1));
1719 assert!(matches!(result, Err(DhtError::NotFound { .. })));
1720 }
1721
1722 #[test]
1723 fn lookup_pages_and_the_pages_partition_the_results() {
1724 let config = SimulationConfig::new()
1725 .with_verify_attestations(false)
1726 .with_page_size(NonZeroUsize::new(2).unwrap());
1727 let dht = SimulatedDht::new(config);
1728 let key = SigningKey::generate();
1729
1730 for suffix in ["2q", "2r", "2s", "2t", "2v"] {
1731 register(&dht, registration(&test_uri(suffix), &key), &key).unwrap();
1732 }
1733
1734 let query = exact("anthropic.com", "assistant/chat");
1735 let mut options = ReadOptions::default();
1736 let mut seen: Vec<String> = Vec::new();
1737 let mut pages = 0;
1738 loop {
1739 let page = block_on(dht.lookup(&query, &options)).unwrap();
1740 pages += 1;
1741 assert!(page.len() <= 2, "page exceeded the configured page size");
1742 let next = page.next_cursor();
1743 seen.extend(page.into_items().iter().map(|r| r.agent_uri().canonical()));
1744 match next {
1745 Some(cursor) => options = options.with_cursor(cursor),
1746 None => break,
1747 }
1748 }
1749
1750 assert_eq!(pages, 3, "5 results at page size 2 should take 3 pages");
1751 assert_eq!(seen.len(), 5);
1752 let mut deduped = seen.clone();
1753 deduped.sort();
1754 deduped.dedup();
1755 assert_eq!(deduped.len(), 5, "pages must not repeat a registration");
1756 }
1757
1758 #[test]
1759 fn paging_order_is_stable_across_calls() {
1760 let config = SimulationConfig::new()
1763 .with_verify_attestations(false)
1764 .with_page_size(NonZeroUsize::new(2).unwrap());
1765 let dht = SimulatedDht::new(config);
1766 let key = SigningKey::generate();
1767 for suffix in ["2v", "2q", "2s", "2r"] {
1768 register(&dht, registration(&test_uri(suffix), &key), &key).unwrap();
1769 }
1770
1771 let query = exact("anthropic.com", "assistant/chat");
1772 let first = block_on(dht.lookup(&query, &ReadOptions::default())).unwrap();
1773 let again = block_on(dht.lookup(&query, &ReadOptions::default())).unwrap();
1774
1775 assert_eq!(
1776 first
1777 .items()
1778 .iter()
1779 .map(|r| r.agent_uri().canonical())
1780 .collect::<Vec<_>>(),
1781 again
1782 .items()
1783 .iter()
1784 .map(|r| r.agent_uri().canonical())
1785 .collect::<Vec<_>>()
1786 );
1787 }
1788
1789 #[test]
1790 fn cursor_past_the_end_yields_an_exhausted_page() {
1791 let dht = unverified_dht();
1792 let key = SigningKey::generate();
1793 register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1794
1795 let page = block_on(dht.lookup(
1796 &exact("anthropic.com", "assistant/chat"),
1797 &ReadOptions::default().with_cursor(Cursor::at(99)),
1798 ))
1799 .unwrap();
1800
1801 assert!(page.is_empty());
1802 assert!(!page.has_more());
1803 }
1804
1805 #[test]
1806 fn write_receipt_reports_the_resolved_quorum() {
1807 let dht = unverified_dht();
1808 let key = SigningKey::generate();
1809 let receipt = register(&dht, registration(&test_uri("2q"), &key), &key).unwrap();
1810
1811 assert_eq!(receipt.acknowledged(), 1);
1813 assert_eq!(receipt.required(), 1);
1814 assert!(!receipt.exceeded_quorum());
1815 }
1816
1817 #[test]
1818 fn instances_have_distinct_identities() {
1819 let a = unverified_dht();
1820 let b = unverified_dht();
1821 assert_ne!(a.local_id(), b.local_id());
1822 assert!(!a.local_id().is_empty());
1823 }
1824
1825 #[test]
1826 fn bootstrap_without_peers_fails_the_way_a_real_backend_would() {
1827 let dht = unverified_dht();
1828 assert!(matches!(
1829 block_on(dht.bootstrap(&[])),
1830 Err(DhtError::NoPeers)
1831 ));
1832 assert!(block_on(dht.bootstrap(&[PeerAddr::new("/ip4/127.0.0.1/tcp/4001")])).is_ok());
1833 }
1834
1835 #[test]
1836 fn the_simulator_is_usable_as_a_trait_object() {
1837 let key = SigningKey::generate();
1840 let dht: Box<dyn Dht> = Box::new(unverified_dht());
1841 let record = registration(&test_uri("2q"), &key);
1842 let proof = MutationProof::sign_registration(&key, &record);
1843 block_on(dht.register(record, &proof, WriteOptions::default())).unwrap();
1844
1845 let page = block_on(dht.lookup(
1846 &exact("anthropic.com", "assistant/chat"),
1847 &ReadOptions::default(),
1848 ))
1849 .unwrap();
1850 assert_eq!(page.len(), 1);
1851 }
1852
1853 #[test]
1854 fn agent_id_cannot_be_reused_under_a_different_capability_path() {
1855 let dht = unverified_dht();
1856 let key = SigningKey::generate();
1857 let first =
1858 AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1859 .unwrap();
1860 let reclassified =
1861 AgentUri::parse("agent://anthropic.com/assistant/code/llm_01h455vb4pex5vsknk084sn02q")
1862 .unwrap();
1863 register(&dht, registration(&first, &key), &key).unwrap();
1864
1865 let result = register(&dht, registration(&reclassified, &key), &key);
1866
1867 assert!(matches!(
1868 result,
1869 Err(DhtError::IdentityCapabilityConflict { .. })
1870 ));
1871
1872 let dht = unverified_dht();
1875 let first =
1876 AgentUri::parse("agent://anthropic.com/assistant/chat/llm_01h455vb4pex5vsknk084sn02q")
1877 .unwrap();
1878 let ancestor_reuse =
1879 AgentUri::parse("agent://anthropic.com/assistant/llm_01h455vb4pex5vsknk084sn02q")
1880 .unwrap();
1881 register(&dht, registration(&first, &key), &key).unwrap();
1882
1883 let result = register(&dht, registration(&ancestor_reuse, &key), &key);
1884
1885 assert!(matches!(
1886 result,
1887 Err(DhtError::IdentityCapabilityConflict { .. })
1888 ));
1889 }
1890
1891 #[test]
1892 fn an_attacker_cannot_repoint_an_agents_endpoints() {
1893 let dht = unverified_dht();
1896 let agent = SigningKey::generate();
1897 let attacker = SigningKey::generate();
1898 let uri = test_uri("2q");
1899 register(&dht, registration(&uri, &agent), &agent).unwrap();
1900
1901 let result = update_endpoint(
1902 &dht,
1903 &uri,
1904 &attacker,
1905 1,
1906 vec![Endpoint::https("attacker.example")],
1907 );
1908
1909 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1910 assert_eq!(
1911 lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].endpoints(),
1912 &[test_endpoint()]
1913 );
1914 }
1915
1916 #[test]
1917 fn an_attacker_cannot_evict_an_agent() {
1918 let dht = unverified_dht();
1919 let agent = SigningKey::generate();
1920 let attacker = SigningKey::generate();
1921 let uri = test_uri("2q");
1922 register(&dht, registration(&uri, &agent), &agent).unwrap();
1923
1924 let result = deregister(&dht, &uri, &attacker, 1);
1925
1926 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1927 assert_eq!(
1928 lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
1929 1
1930 );
1931 }
1932
1933 #[test]
1934 fn an_attacker_cannot_refresh_an_agent() {
1935 let dht = unverified_dht();
1938 let agent = SigningKey::generate();
1939 let attacker = SigningKey::generate();
1940 let uri = test_uri("2q");
1941 register(&dht, registration(&uri, &agent), &agent).unwrap();
1942
1943 let result = refresh(&dht, &uri, &attacker, 1, Duration::from_hours(9));
1944
1945 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1946 }
1947
1948 #[test]
1949 fn a_proof_does_not_authorize_endpoints_it_did_not_sign() {
1950 let dht = unverified_dht();
1954 let agent = SigningKey::generate();
1955 let uri = test_uri("2q");
1956 register(&dht, registration(&uri, &agent), &agent).unwrap();
1957
1958 let honest = vec![Endpoint::https("honest.example")];
1959 let held = record(&dht, &uri);
1960 let proof = MutationProof::sign_next(
1961 &agent,
1962 &held,
1963 &Mutation::UpdateEndpoint {
1964 endpoints: &honest,
1965 expires_at: held.expires_at(),
1966 },
1967 );
1968
1969 let result = block_on(dht.update_endpoint(
1970 &uri,
1971 vec![Endpoint::https("attacker.example")],
1972 &proof,
1973 WriteOptions::default(),
1974 ));
1975
1976 assert!(matches!(result, Err(DhtError::Unauthorized { .. })));
1977 }
1978
1979 #[test]
1980 fn a_successful_mutation_advances_the_record_sequence() {
1981 let dht = unverified_dht();
1982 let key = SigningKey::generate();
1983 let uri = test_uri("2q");
1984 register(&dht, registration(&uri, &key), &key).unwrap();
1985
1986 assert_eq!(
1987 lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].sequence(),
1988 0
1989 );
1990
1991 refresh(&dht, &uri, &key, 4, Duration::from_hours(1)).unwrap();
1992
1993 assert_eq!(
1996 lookup(&dht, &exact("anthropic.com", "assistant/chat"))[0].sequence(),
1997 4
1998 );
1999 }
2000
2001 #[test]
2002 fn every_ancestor_copy_carries_the_new_sequence() {
2003 let dht = unverified_dht();
2006 let key = SigningKey::generate();
2007 let uri = test_uri("2q");
2008 register(&dht, registration(&uri, &key), &key).unwrap();
2009
2010 refresh(&dht, &uri, &key, 3, Duration::from_hours(1)).unwrap();
2011
2012 assert_eq!(
2013 lookup(&dht, &prefix("anthropic.com", "assistant"))[0].sequence(),
2014 3
2015 );
2016 }
2017
2018 #[test]
2019 fn a_replayed_proof_is_refused_the_second_time() {
2020 let dht = unverified_dht();
2021 let key = SigningKey::generate();
2022 let uri = test_uri("2q");
2023 register(&dht, registration(&uri, &key), &key).unwrap();
2024
2025 let endpoints = vec![Endpoint::https("first.example")];
2026 let held = record(&dht, &uri);
2027 let proof = MutationProof::sign_next(
2028 &key,
2029 &held,
2030 &Mutation::UpdateEndpoint {
2031 endpoints: &endpoints,
2032 expires_at: held.expires_at(),
2033 },
2034 );
2035
2036 block_on(dht.update_endpoint(&uri, endpoints.clone(), &proof, WriteOptions::default()))
2037 .unwrap();
2038 let replay =
2039 block_on(dht.update_endpoint(&uri, endpoints, &proof, WriteOptions::default()));
2040
2041 assert!(matches!(
2042 replay,
2043 Err(DhtError::StaleSequence {
2044 presented: 1,
2045 current: 1,
2046 ..
2047 })
2048 ));
2049 }
2050
2051 #[test]
2052 fn a_captured_proof_does_not_reach_a_later_record_instance() {
2053 let dht = unverified_dht();
2057 let key = SigningKey::generate();
2058 let uri = test_uri("2q");
2059 register(&dht, registration(&uri, &key), &key).unwrap();
2060
2061 let captured = MutationProof::sign_next(&key, &record(&dht, &uri), &Mutation::Deregister);
2062 block_on(dht.deregister(&uri, &captured, WriteOptions::default())).unwrap();
2063
2064 register(
2067 &dht,
2068 registration(&uri, &key).with_registered_at(SystemTime::now() + Duration::from_secs(1)),
2069 &key,
2070 )
2071 .unwrap();
2072 let replay = block_on(dht.deregister(&uri, &captured, WriteOptions::default()));
2073
2074 assert!(matches!(replay, Err(DhtError::Unauthorized { .. })));
2075 assert_eq!(
2076 lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
2077 1
2078 );
2079 }
2080
2081 #[test]
2082 fn continuing_the_sequence_defeats_a_captured_proof_without_relying_on_the_clock() {
2083 let dht = unverified_dht();
2088 let key = SigningKey::generate();
2089 let uri = test_uri("2q");
2090 let opened_at = SystemTime::now();
2091 register(
2092 &dht,
2093 registration(&uri, &key).with_registered_at(opened_at),
2094 &key,
2095 )
2096 .unwrap();
2097
2098 let captured = MutationProof::sign_next(&key, &record(&dht, &uri), &Mutation::Deregister);
2099 block_on(dht.deregister(&uri, &captured, WriteOptions::default())).unwrap();
2100
2101 register(
2103 &dht,
2104 registration(&uri, &key)
2105 .with_registered_at(opened_at)
2106 .with_sequence(captured.sequence()),
2107 &key,
2108 )
2109 .unwrap();
2110 let replay = block_on(dht.deregister(&uri, &captured, WriteOptions::default()));
2111
2112 assert!(matches!(replay, Err(DhtError::StaleSequence { .. })));
2113 assert_eq!(
2114 lookup(&dht, &exact("anthropic.com", "assistant/chat")).len(),
2115 1
2116 );
2117 }
2118
2119 #[test]
2120 fn a_write_at_the_records_own_sequence_does_not_advance_it() {
2121 let dht = unverified_dht();
2122 let key = SigningKey::generate();
2123 let uri = test_uri("2q");
2124 register(&dht, registration(&uri, &key), &key).unwrap();
2125
2126 let result = refresh(&dht, &uri, &key, 0, Duration::from_hours(1));
2128
2129 assert!(matches!(
2130 result,
2131 Err(DhtError::StaleSequence {
2132 presented: 0,
2133 current: 0,
2134 ..
2135 })
2136 ));
2137 }
2138
2139 #[test]
2140 fn migration_is_authorized_by_the_agents_own_key() {
2141 let dht = unverified_dht();
2142 let agent = SigningKey::generate();
2143 let attacker = SigningKey::generate();
2144 let uri = test_uri("2q");
2145 register(&dht, registration(&uri, &agent), &agent).unwrap();
2146
2147 let moved = Endpoint::https("eu-west-1.agent.anthropic.com");
2148 assert!(
2149 dht.simulate_migration(&uri, moved.clone(), &attacker)
2150 .is_err()
2151 );
2152
2153 let result = dht.simulate_migration(&uri, moved.clone(), &agent).unwrap();
2154
2155 assert_eq!(result.new_endpoints(), &[moved]);
2156 }
2157
2158 #[test]
2174 fn concurrent_registration_and_migration_do_not_deadlock() {
2175 use std::sync::Arc;
2176 use std::sync::mpsc;
2177 use std::thread;
2178
2179 const WORKERS: usize = 4;
2180 const ROUNDS: usize = 60;
2181 const ALPHABET: &[u8] = b"0123456789abcdefghjkmnpqrstvwxyz";
2184
2185 fn agent(index: usize) -> AgentUri {
2186 test_uri(&format!(
2187 "{}{}",
2188 ALPHABET[index / ALPHABET.len()] as char,
2189 ALPHABET[index % ALPHABET.len()] as char,
2190 ))
2191 }
2192
2193 let dht = Arc::new(unverified_dht());
2194
2195 let migrators: Vec<(AgentUri, SigningKey)> = (0..WORKERS)
2201 .map(|index| {
2202 let (uri, key) = (agent(index), SigningKey::generate());
2203 register(&dht, registration(&uri, &key), &key).expect("a fresh agent registers");
2204 (uri, key)
2205 })
2206 .collect();
2207 let churners: Vec<(AgentUri, SigningKey)> = (WORKERS..2 * WORKERS)
2208 .map(|index| (agent(index), SigningKey::generate()))
2209 .collect();
2210
2211 let mut handles = Vec::new();
2212 for (uri, key) in migrators {
2213 let dht = Arc::clone(&dht);
2214 handles.push(thread::spawn(move || {
2215 for round in 0..ROUNDS {
2216 let endpoint = Endpoint::https(format!("eu-west-{round}.anthropic.com"));
2217 dht.simulate_migration(&uri, endpoint, &key)
2218 .expect("a registered agent migrates");
2219 }
2220 }));
2221 }
2222 for (uri, key) in churners {
2223 let dht = Arc::clone(&dht);
2224 handles.push(thread::spawn(move || {
2225 for _ in 0..ROUNDS {
2226 register(&dht, registration(&uri, &key), &key)
2227 .expect("a deregistered agent registers again");
2228 deregister(&dht, &uri, &key, 1).expect("its own key removes it");
2229 }
2230 }));
2231 }
2232
2233 let (done, finished) = mpsc::channel();
2238 thread::spawn(move || {
2239 let outcomes: Vec<_> = handles.into_iter().map(thread::JoinHandle::join).collect();
2240 let _ = done.send(outcomes);
2241 });
2242
2243 let outcomes = finished
2244 .recv_timeout(Duration::from_mins(1))
2245 .expect("registration and migration deadlocked against each other");
2246 for outcome in outcomes {
2247 outcome.expect("no worker panicked");
2248 }
2249
2250 let registered = lookup(&dht, &prefix("anthropic.com", "assistant"));
2253 assert_eq!(registered.len(), WORKERS);
2254 assert_eq!(dht.stats().total_registrations, WORKERS);
2255 }
2256}