1  //! Agent Circles: Lightweight coordination layer for agent collectives.
  2  //!
  3  //! This module provides a slimmed-down Circle interface designed specifically
  4  //! for agent-to-agent coordination, enabling:
  5  //!
  6  //! - **Resource pooling**: Agents join circles to share inference capacity
  7  //! - **Easy onboarding**: Minimal ceremony for new agents to join
  8  //! - **Serialization-first**: All types designed for wire transmission
  9  //!
 10  //! ## Architecture
 11  //!
 12  //! ```text
 13  //! ┌───────────────────────────────────────────────────────────────┐
 14  //! │                     Agent Circle                              │
 15  //! ├───────────────────────────────────────────────────────────────┤
 16  //! │  Members: [Agent A, Agent B, Agent C]                        │
 17  //! │  Capabilities: [inference, embedding, image_gen]             │
 18  //! │  Epoch: 3 (increments on membership change)                  │
 19  //! └───────────────────────────────────────────────────────────────┘
 20  //!           │                    │                    │
 21  //!      ┌────▼────┐          ┌────▼────┐          ┌────▼────┐
 22  //!      │ Agent A │          │ Agent B │          │ Agent C │
 23  //!      │ llama7b │          │ mistral │          │ embed   │
 24  //!      └─────────┘          └─────────┘          └─────────┘
 25  //! ```
 26  //!
 27  //! ## Joining a Circle
 28  //!
 29  //! ```ignore
 30  //! use abzu_inference::agent_circles::{AgentCircle, AgentProfile, Capability};
 31  //!
 32  //! // Create your agent profile
 33  //! let profile = AgentProfile::new("my-agent")
 34  //!     .with_capability(Capability::Inference { model: "llama-7b".into() })
 35  //!     .with_pubkey(my_pubkey);
 36  //!
 37  //! // Join via invite link
 38  //! let invite = "abzu://circle/abc123/join?token=xyz";
 39  //! circle.join(profile, invite).await?;
 40  //! ```
 41  
 42  use serde::{Deserialize, Serialize};
 43  use std::collections::HashSet;
 44  
 45  /// Capability that an agent offers to the circle.
 46  #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
 47  pub enum Capability {
 48      /// LLM inference with specific model
 49      Inference {
 50          model: String,
 51          max_context: Option<usize>,
 52      },
 53      
 54      /// Text embedding generation
 55      Embedding {
 56          model: String,
 57          dimensions: usize,
 58      },
 59      
 60      /// Image generation
 61      ImageGeneration {
 62          model: String,
 63      },
 64      
 65      /// Speech-to-text transcription
 66      Transcription {
 67          model: String,
 68      },
 69      
 70      /// Text-to-speech synthesis
 71      Speech {
 72          model: String,
 73      },
 74      
 75      /// Code execution sandbox
 76      CodeExecution {
 77          runtime: String,
 78      },
 79      
 80      /// Web browsing
 81      WebBrowsing,
 82      
 83      /// File storage (with capacity in bytes)
 84      Storage {
 85          capacity_bytes: u64,
 86      },
 87      
 88      /// Custom capability
 89      Custom {
 90          name: String,
 91          metadata: String,
 92      },
 93  }
 94  
 95  impl Capability {
 96      /// Create an inference capability
 97      pub fn inference(model: impl Into<String>) -> Self {
 98          Self::Inference {
 99              model: model.into(),
100              max_context: None,
101          }
102      }
103      
104      /// Create an embedding capability
105      pub fn embedding(model: impl Into<String>, dimensions: usize) -> Self {
106          Self::Embedding {
107              model: model.into(),
108              dimensions,
109          }
110      }
111  }
112  
113  /// Agent status within a circle.
114  #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
115  pub enum AgentStatus {
116      /// Available for work
117      Available,
118      /// Currently processing a request
119      Busy,
120      /// Temporarily offline
121      Offline,
122      /// Leaving the circle
123      Departing,
124  }
125  
126  /// Profile of an agent participating in a circle.
127  #[derive(Debug, Clone, Serialize, Deserialize)]
128  pub struct AgentProfile {
129      /// Unique agent identifier (derived from pubkey or chosen)
130      pub id: String,
131      
132      /// Agent's Ed25519 public key (32 bytes)
133      pub pubkey: [u8; 32],
134      
135      /// Human-readable name (optional)
136      pub name: Option<String>,
137      
138      /// Capabilities this agent offers
139      pub capabilities: Vec<Capability>,
140      
141      /// Current status
142      pub status: AgentStatus,
143      
144      /// When the agent joined (Unix ms)
145      pub joined_at: u64,
146      
147      /// Last heartbeat (Unix ms)
148      pub last_seen: u64,
149      
150      /// Reputation score (0.0 - 1.0, based on successful tasks)
151      pub reputation: f32,
152  }
153  
154  impl AgentProfile {
155      /// Create a new agent profile with minimal info
156      pub fn new(id: impl Into<String>) -> Self {
157          let now = std::time::SystemTime::now()
158              .duration_since(std::time::UNIX_EPOCH)
159              .unwrap()
160              .as_millis() as u64;
161          
162          Self {
163              id: id.into(),
164              pubkey: [0u8; 32], // Must be set via with_pubkey
165              name: None,
166              capabilities: vec![],
167              status: AgentStatus::Available,
168              joined_at: now,
169              last_seen: now,
170              reputation: 0.5, // Neutral starting reputation
171          }
172      }
173      
174      /// Set the agent's public key
175      pub fn with_pubkey(mut self, pubkey: [u8; 32]) -> Self {
176          self.pubkey = pubkey;
177          self
178      }
179      
180      /// Set the agent's display name
181      pub fn with_name(mut self, name: impl Into<String>) -> Self {
182          self.name = Some(name.into());
183          self
184      }
185      
186      /// Add a capability
187      pub fn with_capability(mut self, cap: Capability) -> Self {
188          self.capabilities.push(cap);
189          self
190      }
191      
192      /// Set initial reputation
193      pub fn with_reputation(mut self, rep: f32) -> Self {
194          self.reputation = rep.clamp(0.0, 1.0);
195          self
196      }
197      
198      /// Check if this agent has a specific capability
199      pub fn has_capability(&self, cap: &Capability) -> bool {
200          self.capabilities.contains(cap)
201      }
202      
203      /// Check if this agent can handle inference for a model
204      pub fn can_infer(&self, model: &str) -> bool {
205          self.capabilities.iter().any(|c| {
206              matches!(c, Capability::Inference { model: m, .. } if m == model)
207          })
208      }
209      
210      /// Update last_seen to now
211      pub fn heartbeat(&mut self) {
212          self.last_seen = std::time::SystemTime::now()
213              .duration_since(std::time::UNIX_EPOCH)
214              .unwrap()
215              .as_millis() as u64;
216      }
217      
218      /// Check if agent is stale (no heartbeat for duration_ms)
219      pub fn is_stale(&self, duration_ms: u64) -> bool {
220          let now = std::time::SystemTime::now()
221              .duration_since(std::time::UNIX_EPOCH)
222              .unwrap()
223              .as_millis() as u64;
224          
225          now.saturating_sub(self.last_seen) > duration_ms
226      }
227  }
228  
229  /// A lightweight agent circle for coordination.
230  ///
231  /// Unlike full Circle (which handles messaging, encryption, trust), 
232  /// AgentCircle focuses purely on resource coordination.
233  #[derive(Debug, Clone, Serialize, Deserialize)]
234  pub struct AgentCircle {
235      /// Circle ID (32 bytes)
236      pub id: [u8; 32],
237      
238      /// Circle name
239      pub name: String,
240      
241      /// Description of the circle's purpose
242      pub description: Option<String>,
243      
244      /// Member agents
245      pub members: Vec<AgentProfile>,
246      
247      /// Current epoch (increments on membership change)
248      pub epoch: u64,
249      
250      /// Aggregated capabilities available in this circle
251      pub capabilities: HashSet<String>,
252      
253      /// Creation timestamp (Unix ms)
254      pub created_at: u64,
255      
256      /// Maximum members allowed (0 = unlimited)
257      pub max_members: usize,
258      
259      /// Whether new agents can join freely or need invite
260      pub open_membership: bool,
261  }
262  
263  impl AgentCircle {
264      /// Create a new agent circle
265      pub fn new(name: impl Into<String>, founder: AgentProfile) -> Self {
266          let mut hasher = blake3::Hasher::new();
267          hasher.update(b"abzu:agent-circle:");
268          hasher.update(&founder.pubkey);
269          hasher.update(&founder.joined_at.to_le_bytes());
270          let nonce: [u8; 16] = rand::random();
271          hasher.update(&nonce);
272          let id = *hasher.finalize().as_bytes();
273          
274          let mut capabilities = HashSet::new();
275          for cap in &founder.capabilities {
276              capabilities.insert(capability_name(cap));
277          }
278          
279          Self {
280              id,
281              name: name.into(),
282              description: None,
283              members: vec![founder],
284              epoch: 1,
285              capabilities,
286              created_at: std::time::SystemTime::now()
287                  .duration_since(std::time::UNIX_EPOCH)
288                  .unwrap()
289                  .as_millis() as u64,
290              max_members: 0,
291              open_membership: false,
292          }
293      }
294      
295      /// Set description
296      pub fn with_description(mut self, desc: impl Into<String>) -> Self {
297          self.description = Some(desc.into());
298          self
299      }
300      
301      /// Set max members
302      pub fn with_max_members(mut self, max: usize) -> Self {
303          self.max_members = max;
304          self
305      }
306      
307      /// Enable open membership
308      pub fn with_open_membership(mut self) -> Self {
309          self.open_membership = true;
310          self
311      }
312      
313      /// Add a member to the circle
314      pub fn add_member(&mut self, agent: AgentProfile) -> Result<(), String> {
315          // Check capacity
316          if self.max_members > 0 && self.members.len() >= self.max_members {
317              return Err("Circle is full".into());
318          }
319          
320          // Check if already a member
321          if self.members.iter().any(|m| m.pubkey == agent.pubkey) {
322              return Err("Agent is already a member".into());
323          }
324          
325          // Add capabilities
326          for cap in &agent.capabilities {
327              self.capabilities.insert(capability_name(cap));
328          }
329          
330          self.members.push(agent);
331          self.epoch += 1;
332          
333          Ok(())
334      }
335      
336      /// Remove a member from the circle
337      pub fn remove_member(&mut self, pubkey: &[u8; 32]) -> bool {
338          let before = self.members.len();
339          self.members.retain(|m| &m.pubkey != pubkey);
340          
341          if self.members.len() < before {
342              // Recalculate capabilities
343              self.capabilities.clear();
344              for member in &self.members {
345                  for cap in &member.capabilities {
346                      self.capabilities.insert(capability_name(cap));
347                  }
348              }
349              self.epoch += 1;
350              true
351          } else {
352              false
353          }
354      }
355      
356      /// Get available members (status == Available)
357      pub fn available_members(&self) -> Vec<&AgentProfile> {
358          self.members
359              .iter()
360              .filter(|m| m.status == AgentStatus::Available)
361              .collect()
362      }
363      
364      /// Find members that can handle a specific capability
365      pub fn find_capable(&self, cap: &Capability) -> Vec<&AgentProfile> {
366          self.members
367              .iter()
368              .filter(|m| m.status == AgentStatus::Available && m.has_capability(cap))
369              .collect()
370      }
371      
372      /// Find members that can handle inference for a model
373      pub fn find_inference_capable(&self, model: &str) -> Vec<&AgentProfile> {
374          self.members
375              .iter()
376              .filter(|m| m.status == AgentStatus::Available && m.can_infer(model))
377              .collect()
378      }
379      
380      /// Prune stale members (no heartbeat for duration)
381      pub fn prune_stale(&mut self, stale_duration_ms: u64) -> usize {
382          let before = self.members.len();
383          self.members.retain(|m| !m.is_stale(stale_duration_ms));
384          let pruned = before - self.members.len();
385          
386          if pruned > 0 {
387              // Recalculate capabilities
388              self.capabilities.clear();
389              for member in &self.members {
390                  for cap in &member.capabilities {
391                      self.capabilities.insert(capability_name(cap));
392                  }
393              }
394              self.epoch += 1;
395          }
396          
397          pruned
398      }
399      
400      /// Generate a compact invite link
401      pub fn invite_link(&self) -> String {
402          let id_hex = hex::encode(&self.id[..8]);
403          format!("abzu://circle/{}/join", id_hex)
404      }
405      
406      /// Serialize to compact bytes for transmission
407      pub fn to_bytes(&self) -> Result<Vec<u8>, String> {
408          serde_json::to_vec(self).map_err(|e| e.to_string())
409      }
410      
411      /// Deserialize from bytes
412      pub fn from_bytes(bytes: &[u8]) -> Result<Self, String> {
413          serde_json::from_slice(bytes).map_err(|e| e.to_string())
414      }
415  }
416  
417  /// Extract a capability name for aggregation
418  fn capability_name(cap: &Capability) -> String {
419      match cap {
420          Capability::Inference { model, .. } => format!("inference:{}", model),
421          Capability::Embedding { model, .. } => format!("embedding:{}", model),
422          Capability::ImageGeneration { model } => format!("image:{}", model),
423          Capability::Transcription { model } => format!("transcription:{}", model),
424          Capability::Speech { model } => format!("speech:{}", model),
425          Capability::CodeExecution { runtime } => format!("code:{}", runtime),
426          Capability::WebBrowsing => "web".into(),
427          Capability::Storage { .. } => "storage".into(),
428          Capability::Custom { name, .. } => format!("custom:{}", name),
429      }
430  }
431  
432  /// Invite token for joining a circle.
433  #[derive(Debug, Clone, Serialize, Deserialize)]
434  pub struct CircleInvite {
435      /// Circle ID
436      pub circle_id: [u8; 32],
437      
438      /// Circle name (for display)
439      pub circle_name: String,
440      
441      /// Who created the invite
442      pub inviter_pubkey: [u8; 32],
443      
444      /// Expiration timestamp (Unix ms, 0 = never)
445      pub expires_at: u64,
446      
447      /// Maximum uses (0 = unlimited)
448      pub max_uses: u32,
449      
450      /// Signature from inviter
451      pub signature: Vec<u8>,
452  }
453  
454  impl CircleInvite {
455      /// Create a new invite
456      pub fn new(
457          circle_id: [u8; 32],
458          circle_name: String,
459          inviter_pubkey: [u8; 32],
460      ) -> Self {
461          Self {
462              circle_id,
463              circle_name,
464              inviter_pubkey,
465              expires_at: 0,
466              max_uses: 0,
467              signature: vec![],
468          }
469      }
470      
471      /// Set expiration (duration from now)
472      pub fn expires_in(mut self, duration: std::time::Duration) -> Self {
473          let now = std::time::SystemTime::now()
474              .duration_since(std::time::UNIX_EPOCH)
475              .unwrap()
476              .as_millis() as u64;
477          self.expires_at = now + duration.as_millis() as u64;
478          self
479      }
480      
481      /// Set max uses
482      pub fn with_max_uses(mut self, uses: u32) -> Self {
483          self.max_uses = uses;
484          self
485      }
486      
487      /// Check if invite is still valid
488      pub fn is_valid(&self) -> bool {
489          if self.expires_at == 0 {
490              return true;
491          }
492          
493          let now = std::time::SystemTime::now()
494              .duration_since(std::time::UNIX_EPOCH)
495              .unwrap()
496              .as_millis() as u64;
497          
498          now < self.expires_at
499      }
500      
501      /// Serialize to URL-safe string
502      pub fn to_token(&self) -> String {
503          let bytes = serde_json::to_vec(self).unwrap_or_default();
504          base64::Engine::encode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, &bytes)
505      }
506      
507      /// Deserialize from URL-safe string
508      pub fn from_token(token: &str) -> Result<Self, String> {
509          let bytes = base64::Engine::decode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, token)
510              .map_err(|e| e.to_string())?;
511          serde_json::from_slice(&bytes).map_err(|e| e.to_string())
512      }
513  }
514  
515  #[cfg(test)]
516  mod tests {
517      use super::*;
518  
519      #[test]
520      fn test_agent_profile_builder() {
521          let profile = AgentProfile::new("test-agent")
522              .with_name("Test Agent")
523              .with_pubkey([1u8; 32])
524              .with_capability(Capability::inference("llama-7b"))
525              .with_reputation(0.8);
526          
527          assert_eq!(profile.id, "test-agent");
528          assert_eq!(profile.name, Some("Test Agent".into()));
529          assert!(profile.can_infer("llama-7b"));
530          assert!(!profile.can_infer("mistral"));
531          assert!((profile.reputation - 0.8).abs() < 0.01);
532      }
533  
534      #[test]
535      fn test_agent_circle_creation() {
536          let founder = AgentProfile::new("founder")
537              .with_pubkey([1u8; 32])
538              .with_capability(Capability::inference("llama-7b"));
539          
540          let circle = AgentCircle::new("Test Circle", founder)
541              .with_description("A test circle for agents")
542              .with_max_members(10);
543          
544          assert_eq!(circle.name, "Test Circle");
545          assert_eq!(circle.members.len(), 1);
546          assert!(circle.capabilities.contains("inference:llama-7b"));
547      }
548  
549      #[test]
550      fn test_add_member() {
551          let founder = AgentProfile::new("founder")
552              .with_pubkey([1u8; 32])
553              .with_capability(Capability::inference("llama-7b"));
554          
555          let mut circle = AgentCircle::new("Test Circle", founder);
556          
557          let member = AgentProfile::new("member1")
558              .with_pubkey([2u8; 32])
559              .with_capability(Capability::embedding("bge-small", 384));
560          
561          circle.add_member(member).unwrap();
562          
563          assert_eq!(circle.members.len(), 2);
564          assert_eq!(circle.epoch, 2);
565          assert!(circle.capabilities.contains("embedding:bge-small"));
566      }
567  
568      #[test]
569      fn test_find_capable() {
570          let founder = AgentProfile::new("founder")
571              .with_pubkey([1u8; 32])
572              .with_capability(Capability::inference("llama-7b"));
573          
574          let mut circle = AgentCircle::new("Test Circle", founder);
575          
576          let member = AgentProfile::new("member1")
577              .with_pubkey([2u8; 32])
578              .with_capability(Capability::inference("mistral-7b"));
579          
580          circle.add_member(member).unwrap();
581          
582          let llama_capable = circle.find_inference_capable("llama-7b");
583          assert_eq!(llama_capable.len(), 1);
584          assert_eq!(llama_capable[0].id, "founder");
585          
586          let mistral_capable = circle.find_inference_capable("mistral-7b");
587          assert_eq!(mistral_capable.len(), 1);
588          assert_eq!(mistral_capable[0].id, "member1");
589      }
590  
591      #[test]
592      fn test_invite_link() {
593          let founder = AgentProfile::new("founder").with_pubkey([1u8; 32]);
594          let circle = AgentCircle::new("Test", founder);
595          
596          let link = circle.invite_link();
597          assert!(link.starts_with("abzu://circle/"));
598          assert!(link.ends_with("/join"));
599      }
600  
601      #[test]
602      fn test_circle_serialization() {
603          let founder = AgentProfile::new("founder")
604              .with_pubkey([1u8; 32])
605              .with_capability(Capability::inference("test"));
606          
607          let circle = AgentCircle::new("Test", founder);
608          
609          let bytes = circle.to_bytes().unwrap();
610          let restored = AgentCircle::from_bytes(&bytes).unwrap();
611          
612          assert_eq!(restored.name, circle.name);
613          assert_eq!(restored.id, circle.id);
614          assert_eq!(restored.members.len(), 1);
615      }
616  
617      #[test]
618      fn test_invite_token() {
619          let invite = CircleInvite::new([1u8; 32], "Test Circle".into(), [2u8; 32])
620              .expires_in(std::time::Duration::from_secs(3600))
621              .with_max_uses(5);
622          
623          let token = invite.to_token();
624          let restored = CircleInvite::from_token(&token).unwrap();
625          
626          assert_eq!(restored.circle_name, "Test Circle");
627          assert_eq!(restored.max_uses, 5);
628          assert!(restored.is_valid());
629      }
630  }