tirea_agentos/composition/
delegation.rs

1use super::agent_definition::AgentDefinition;
2use crate::composition::{AgentRegistry, InMemoryAgentRegistry};
3use std::collections::HashMap;
4use std::sync::{Arc, RwLock};
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct AgentDescriptor {
8    pub id: String,
9    pub name: String,
10    pub description: String,
11}
12
13impl AgentDescriptor {
14    #[must_use]
15    pub fn new(id: impl Into<String>) -> Self {
16        let id = id.into();
17        Self {
18            name: id.clone(),
19            id,
20            description: String::new(),
21        }
22    }
23
24    #[must_use]
25    pub fn with_name(mut self, name: impl Into<String>) -> Self {
26        self.name = name.into();
27        self
28    }
29
30    #[must_use]
31    pub fn with_description(mut self, description: impl Into<String>) -> Self {
32        self.description = description.into();
33        self
34    }
35
36    #[must_use]
37    pub fn into_local(self) -> ResolvedAgent {
38        ResolvedAgent {
39            descriptor: self,
40            binding: AgentBinding::Local,
41        }
42    }
43
44    #[must_use]
45    pub fn into_a2a(self, binding: A2aAgentBinding) -> ResolvedAgent {
46        ResolvedAgent {
47            descriptor: self,
48            binding: AgentBinding::A2a(binding),
49        }
50    }
51}
52
53#[derive(Debug, Clone)]
54pub enum AgentDefinitionSpec {
55    Local(Box<AgentDefinition>),
56    Remote(RemoteAgentDefinition),
57}
58
59impl AgentDefinitionSpec {
60    #[must_use]
61    pub fn local(definition: AgentDefinition) -> Self {
62        Self::Local(Box::new(definition))
63    }
64
65    #[must_use]
66    pub fn local_with_id(agent_id: impl Into<String>, mut definition: AgentDefinition) -> Self {
67        definition.id = agent_id.into();
68        Self::local(definition)
69    }
70
71    #[must_use]
72    pub fn a2a(descriptor: AgentDescriptor, binding: A2aAgentBinding) -> Self {
73        Self::Remote(RemoteAgentDefinition::a2a(descriptor, binding))
74    }
75
76    #[must_use]
77    pub fn a2a_with_id(agent_id: impl Into<String>, binding: A2aAgentBinding) -> Self {
78        Self::a2a(AgentDescriptor::new(agent_id), binding)
79    }
80
81    #[must_use]
82    pub fn id(&self) -> &str {
83        match self {
84            Self::Local(definition) => &definition.id,
85            Self::Remote(definition) => definition.id(),
86        }
87    }
88
89    #[must_use]
90    pub fn descriptor(&self) -> AgentDescriptor {
91        match self {
92            Self::Local(definition) => definition.descriptor(),
93            Self::Remote(definition) => definition.descriptor(),
94        }
95    }
96}
97
98#[derive(Debug, Clone)]
99pub struct RemoteAgentDefinition {
100    pub descriptor: AgentDescriptor,
101    pub binding: RemoteAgentBinding,
102}
103
104impl RemoteAgentDefinition {
105    #[must_use]
106    pub fn a2a(descriptor: AgentDescriptor, binding: A2aAgentBinding) -> Self {
107        Self {
108            descriptor,
109            binding: RemoteAgentBinding::A2a(binding),
110        }
111    }
112
113    #[must_use]
114    pub fn id(&self) -> &str {
115        &self.descriptor.id
116    }
117
118    #[must_use]
119    pub fn descriptor(&self) -> AgentDescriptor {
120        self.descriptor.clone()
121    }
122
123    #[must_use]
124    pub fn into_resolved_agent(self) -> ResolvedAgent {
125        match self.binding {
126            RemoteAgentBinding::A2a(binding) => self.descriptor.into_a2a(binding),
127        }
128    }
129}
130
131#[derive(Clone)]
132pub struct ResolvedAgent {
133    pub descriptor: AgentDescriptor,
134    pub binding: AgentBinding,
135}
136
137impl ResolvedAgent {
138    #[must_use]
139    pub fn local(id: impl Into<String>) -> Self {
140        AgentDescriptor::new(id).into_local()
141    }
142
143    #[must_use]
144    pub fn a2a(
145        id: impl Into<String>,
146        base_url: impl Into<String>,
147        remote_agent_id: impl Into<String>,
148    ) -> Self {
149        AgentDescriptor::new(id).into_a2a(A2aAgentBinding::new(base_url, remote_agent_id))
150    }
151
152    #[must_use]
153    pub fn with_name(mut self, name: impl Into<String>) -> Self {
154        self.descriptor.name = name.into();
155        self
156    }
157
158    #[must_use]
159    pub fn with_description(mut self, description: impl Into<String>) -> Self {
160        self.descriptor.description = description.into();
161        self
162    }
163
164    #[must_use]
165    pub fn kind_tag(&self) -> &'static str {
166        self.binding.kind_tag()
167    }
168
169    #[must_use]
170    pub fn is_resumable(&self) -> bool {
171        self.binding.is_resumable()
172    }
173}
174
175impl std::fmt::Debug for ResolvedAgent {
176    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
177        f.debug_struct("ResolvedAgent")
178            .field("descriptor", &self.descriptor)
179            .field("binding", &self.binding)
180            .finish()
181    }
182}
183
184#[derive(Debug, Clone)]
185pub enum RemoteAgentBinding {
186    A2a(A2aAgentBinding),
187}
188
189#[derive(Debug, Clone)]
190pub enum AgentBinding {
191    Local,
192    A2a(A2aAgentBinding),
193}
194
195impl AgentBinding {
196    #[must_use]
197    pub fn kind_tag(&self) -> &'static str {
198        match self {
199            Self::Local => "local",
200            Self::A2a(_) => "remote_a2a",
201        }
202    }
203
204    #[must_use]
205    pub fn is_resumable(&self) -> bool {
206        matches!(self, Self::Local)
207    }
208}
209
210#[derive(Clone)]
211pub struct A2aAgentBinding {
212    pub base_url: String,
213    pub remote_agent_id: String,
214    pub auth: Option<RemoteSecurityConfig>,
215    pub poll_interval_ms: u64,
216}
217
218impl A2aAgentBinding {
219    #[must_use]
220    pub fn new(base_url: impl Into<String>, remote_agent_id: impl Into<String>) -> Self {
221        Self {
222            base_url: base_url.into(),
223            remote_agent_id: remote_agent_id.into(),
224            auth: None,
225            poll_interval_ms: 500,
226        }
227    }
228
229    #[must_use]
230    pub fn with_auth(mut self, auth: RemoteSecurityConfig) -> Self {
231        self.auth = Some(auth);
232        self
233    }
234
235    #[must_use]
236    pub fn with_poll_interval_ms(mut self, poll_interval_ms: u64) -> Self {
237        self.poll_interval_ms = poll_interval_ms.max(50);
238        self
239    }
240}
241
242impl std::fmt::Debug for A2aAgentBinding {
243    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
244        f.debug_struct("A2aAgentBinding")
245            .field("base_url", &self.base_url)
246            .field("remote_agent_id", &self.remote_agent_id)
247            .field("auth", &self.auth)
248            .field("poll_interval_ms", &self.poll_interval_ms)
249            .finish()
250    }
251}
252
253#[derive(Clone)]
254pub enum RemoteSecurityConfig {
255    BearerToken(String),
256    Header { name: String, value: String },
257}
258
259impl std::fmt::Debug for RemoteSecurityConfig {
260    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
261        match self {
262            Self::BearerToken(_) => f.write_str("BearerToken([redacted])"),
263            Self::Header { name, .. } => f
264                .debug_struct("Header")
265                .field("name", name)
266                .field("value", &"[redacted]")
267                .finish(),
268        }
269    }
270}
271
272#[derive(Debug, thiserror::Error)]
273pub enum AgentCatalogError {
274    #[error("agent id already registered: {0}")]
275    AgentIdConflict(String),
276}
277
278pub trait AgentCatalog: Send + Sync {
279    fn len(&self) -> usize;
280
281    fn is_empty(&self) -> bool {
282        self.len() == 0
283    }
284
285    fn get(&self, id: &str) -> Option<ResolvedAgent>;
286
287    fn descriptor(&self, id: &str) -> Option<AgentDescriptor> {
288        self.get(id).map(|agent| agent.descriptor)
289    }
290
291    fn ids(&self) -> Vec<String>;
292
293    fn snapshot(&self) -> HashMap<String, ResolvedAgent>;
294
295    fn descriptors(&self) -> HashMap<String, AgentDescriptor> {
296        self.snapshot()
297            .into_iter()
298            .map(|(id, agent)| (id, agent.descriptor))
299            .collect()
300    }
301}
302
303#[derive(Debug, Clone, Default)]
304pub struct InMemoryAgentCatalog {
305    agents: HashMap<String, ResolvedAgent>,
306}
307
308impl InMemoryAgentCatalog {
309    #[must_use]
310    pub fn new() -> Self {
311        Self::default()
312    }
313
314    pub fn register(
315        &mut self,
316        agent_id: impl Into<String>,
317        mut agent: ResolvedAgent,
318    ) -> Result<(), AgentCatalogError> {
319        let agent_id = agent_id.into();
320        if self.agents.contains_key(&agent_id) {
321            return Err(AgentCatalogError::AgentIdConflict(agent_id));
322        }
323        agent.descriptor.id = agent_id.clone();
324        if agent.descriptor.name.trim().is_empty() {
325            agent.descriptor.name = agent_id.clone();
326        }
327        self.agents.insert(agent_id, agent);
328        Ok(())
329    }
330
331    pub fn upsert(&mut self, agent_id: impl Into<String>, mut agent: ResolvedAgent) {
332        let agent_id = agent_id.into();
333        agent.descriptor.id = agent_id.clone();
334        if agent.descriptor.name.trim().is_empty() {
335            agent.descriptor.name = agent_id.clone();
336        }
337        self.agents.insert(agent_id, agent);
338    }
339
340    pub fn extend_upsert(&mut self, agents: HashMap<String, ResolvedAgent>) {
341        for (id, agent) in agents {
342            self.upsert(id, agent);
343        }
344    }
345
346    pub fn extend_catalog(&mut self, other: &dyn AgentCatalog) -> Result<(), AgentCatalogError> {
347        for (id, agent) in other.snapshot() {
348            self.register(id, agent)?;
349        }
350        Ok(())
351    }
352}
353
354impl AgentCatalog for InMemoryAgentCatalog {
355    fn len(&self) -> usize {
356        self.agents.len()
357    }
358
359    fn get(&self, id: &str) -> Option<ResolvedAgent> {
360        self.agents.get(id).cloned()
361    }
362
363    fn ids(&self) -> Vec<String> {
364        let mut ids: Vec<String> = self.agents.keys().cloned().collect();
365        ids.sort();
366        ids
367    }
368
369    fn snapshot(&self) -> HashMap<String, ResolvedAgent> {
370        self.agents.clone()
371    }
372}
373
374#[derive(Clone)]
375pub struct HostedAgentCatalog {
376    agents: Arc<dyn AgentRegistry>,
377}
378
379impl HostedAgentCatalog {
380    #[must_use]
381    pub fn new(agents: Arc<dyn AgentRegistry>) -> Self {
382        Self { agents }
383    }
384}
385
386impl std::fmt::Debug for HostedAgentCatalog {
387    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
388        f.debug_struct("HostedAgentCatalog")
389            .field("len", &self.agents.len())
390            .finish()
391    }
392}
393
394impl AgentCatalog for HostedAgentCatalog {
395    fn len(&self) -> usize {
396        self.agents.len()
397    }
398
399    fn get(&self, id: &str) -> Option<ResolvedAgent> {
400        self.agents
401            .get(id)
402            .map(|definition| definition.descriptor().into_local())
403    }
404
405    fn ids(&self) -> Vec<String> {
406        self.agents.ids()
407    }
408
409    fn snapshot(&self) -> HashMap<String, ResolvedAgent> {
410        self.agents
411            .snapshot()
412            .into_iter()
413            .map(|(id, definition)| (id, definition.descriptor().into_local()))
414            .collect()
415    }
416}
417
418impl AgentCatalog for InMemoryAgentRegistry {
419    fn len(&self) -> usize {
420        AgentRegistry::len(self)
421    }
422
423    fn get(&self, id: &str) -> Option<ResolvedAgent> {
424        AgentRegistry::get(self, id).map(|definition| definition.descriptor().into_local())
425    }
426
427    fn ids(&self) -> Vec<String> {
428        AgentRegistry::ids(self)
429    }
430
431    fn snapshot(&self) -> HashMap<String, ResolvedAgent> {
432        AgentRegistry::snapshot(self)
433            .into_iter()
434            .map(|(id, definition)| (id, definition.descriptor().into_local()))
435            .collect()
436    }
437}
438
439#[derive(Clone, Default)]
440pub struct CompositeAgentCatalog {
441    catalogs: Vec<Arc<dyn AgentCatalog>>,
442    cached_snapshot: Arc<RwLock<HashMap<String, ResolvedAgent>>>,
443}
444
445impl CompositeAgentCatalog {
446    pub fn try_new(
447        catalogs: impl IntoIterator<Item = Arc<dyn AgentCatalog>>,
448    ) -> Result<Self, AgentCatalogError> {
449        let catalogs: Vec<Arc<dyn AgentCatalog>> = catalogs.into_iter().collect();
450        let merged = Self::merge_snapshots(&catalogs)?;
451        Ok(Self {
452            catalogs,
453            cached_snapshot: Arc::new(RwLock::new(merged)),
454        })
455    }
456
457    fn merge_snapshots(
458        catalogs: &[Arc<dyn AgentCatalog>],
459    ) -> Result<HashMap<String, ResolvedAgent>, AgentCatalogError> {
460        let mut merged = InMemoryAgentCatalog::new();
461        for catalog in catalogs {
462            merged.extend_catalog(catalog.as_ref())?;
463        }
464        Ok(merged.snapshot())
465    }
466
467    fn refresh_snapshot(&self) -> Result<HashMap<String, ResolvedAgent>, AgentCatalogError> {
468        Self::merge_snapshots(&self.catalogs)
469    }
470
471    fn read_cached_snapshot(&self) -> HashMap<String, ResolvedAgent> {
472        match self.cached_snapshot.read() {
473            Ok(guard) => guard.clone(),
474            Err(poisoned) => poisoned.into_inner().clone(),
475        }
476    }
477
478    fn write_cached_snapshot(&self, snapshot: HashMap<String, ResolvedAgent>) {
479        match self.cached_snapshot.write() {
480            Ok(mut guard) => *guard = snapshot,
481            Err(poisoned) => *poisoned.into_inner() = snapshot,
482        };
483    }
484}
485
486impl std::fmt::Debug for CompositeAgentCatalog {
487    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
488        let snapshot = match self.cached_snapshot.read() {
489            Ok(guard) => guard,
490            Err(poisoned) => poisoned.into_inner(),
491        };
492        f.debug_struct("CompositeAgentCatalog")
493            .field("catalogs", &self.catalogs.len())
494            .field("len", &snapshot.len())
495            .finish()
496    }
497}
498
499impl AgentCatalog for CompositeAgentCatalog {
500    fn len(&self) -> usize {
501        self.snapshot().len()
502    }
503
504    fn get(&self, id: &str) -> Option<ResolvedAgent> {
505        self.snapshot().get(id).cloned()
506    }
507
508    fn ids(&self) -> Vec<String> {
509        let mut ids: Vec<String> = self.snapshot().keys().cloned().collect();
510        ids.sort();
511        ids
512    }
513
514    fn snapshot(&self) -> HashMap<String, ResolvedAgent> {
515        match self.refresh_snapshot() {
516            Ok(snapshot) => {
517                self.write_cached_snapshot(snapshot.clone());
518                snapshot
519            }
520            Err(_) => self.read_cached_snapshot(),
521        }
522    }
523}
524
525#[cfg(test)]
526mod tests {
527    use super::*;
528    use crate::composition::InMemoryAgentRegistry;
529
530    #[test]
531    fn hosted_agent_catalog_projects_agent_registry_without_mutating_it() {
532        let mut agents = InMemoryAgentRegistry::new();
533        agents.upsert(
534            "worker",
535            crate::composition::AgentDefinition::new("mock")
536                .with_name("Worker")
537                .with_description("Local worker"),
538        );
539        let catalog = HostedAgentCatalog::new(Arc::new(agents));
540        let agent = catalog.get("worker").expect("agent should exist");
541        assert_eq!(agent.descriptor.id, "worker");
542        assert_eq!(agent.descriptor.name, "Worker");
543        assert_eq!(agent.descriptor.description, "Local worker");
544        assert_eq!(agent.kind_tag(), "local");
545        assert!(agent.is_resumable());
546    }
547
548    #[test]
549    fn composite_agent_catalog_rejects_duplicate_ids() {
550        let mut left = InMemoryAgentCatalog::new();
551        left.upsert("worker", ResolvedAgent::local("worker"));
552        let mut right = InMemoryAgentCatalog::new();
553        right.upsert(
554            "worker",
555            ResolvedAgent::a2a("worker", "https://example.test/v1/a2a", "remote-worker"),
556        );
557        let result = CompositeAgentCatalog::try_new(vec![
558            Arc::new(left) as Arc<dyn AgentCatalog>,
559            Arc::new(right) as Arc<dyn AgentCatalog>,
560        ]);
561        assert!(matches!(result, Err(AgentCatalogError::AgentIdConflict(id)) if id == "worker"));
562    }
563
564    #[test]
565    fn agent_definition_spec_keeps_public_descriptor_shape_for_local_and_remote() {
566        let local = AgentDefinitionSpec::local_with_id(
567            "local-worker",
568            crate::composition::AgentDefinition::new("mock")
569                .with_name("Local Worker")
570                .with_description("Hosted locally"),
571        );
572        assert_eq!(local.id(), "local-worker");
573        assert_eq!(local.descriptor().name, "Local Worker");
574
575        let remote = AgentDefinitionSpec::a2a(
576            AgentDescriptor::new("remote-worker")
577                .with_name("Remote Worker")
578                .with_description("Delegated over A2A"),
579            A2aAgentBinding::new("https://example.test/v1/a2a", "remote-worker"),
580        );
581        assert_eq!(remote.id(), "remote-worker");
582        assert_eq!(remote.descriptor().description, "Delegated over A2A");
583    }
584
585    #[test]
586    fn agent_definition_spec_convenience_constructors_normalize_common_ids() {
587        let local = AgentDefinitionSpec::local_with_id(
588            "worker",
589            crate::composition::AgentDefinition::new("mock"),
590        );
591        assert_eq!(local.id(), "worker");
592
593        let remote = AgentDefinitionSpec::a2a_with_id(
594            "researcher",
595            A2aAgentBinding::new("https://example.test/v1/a2a", "remote-researcher"),
596        );
597        assert_eq!(remote.id(), "researcher");
598        assert_eq!(remote.descriptor().name, "researcher");
599    }
600
601    #[test]
602    fn remote_security_debug_redacts_secret_material() {
603        let auth = RemoteSecurityConfig::BearerToken("secret".to_string());
604        let rendered = format!("{auth:?}");
605        assert!(!rendered.contains("secret"));
606        assert!(rendered.contains("redacted"));
607    }
608}