adk_agent/team/
discovery.rs

1use std::cmp::Reverse;
2use std::collections::HashMap;
3use std::sync::Arc;
4
5use adk_core::Agent;
6use async_trait::async_trait;
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10use super::{CompiledTeam, ResolvedTeamMember, TeamError, TeamMemberSpec, TeamSpec};
11
12/// Health advertised by a registry candidate at resolution time.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
14#[serde(rename_all = "camelCase")]
15pub enum TeamAgentHealth {
16    /// Candidate is accepting work normally.
17    #[default]
18    Healthy,
19    /// Candidate is available with reduced capacity or functionality.
20    Degraded,
21    /// Candidate must not receive new work.
22    Unavailable,
23}
24
25/// Serializable candidate metadata returned by a team agent registry.
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
27#[serde(rename_all = "camelCase", deny_unknown_fields)]
28pub struct TeamAgentDescriptor {
29    /// Registry-specific immutable binding identifier.
30    pub binding: String,
31    /// Agent's advertised name.
32    pub name: String,
33    /// Human-readable capability summary.
34    #[serde(default, skip_serializing_if = "String::is_empty")]
35    pub description: String,
36    /// Machine-comparable capability identifiers.
37    #[serde(default, skip_serializing_if = "Vec::is_empty")]
38    pub capabilities: Vec<String>,
39    /// Higher-priority candidates win deterministic selection.
40    #[serde(default)]
41    pub priority: i32,
42    /// Provider or semantic version of this immutable binding.
43    #[serde(default, skip_serializing_if = "Option::is_none")]
44    pub version: Option<String>,
45    /// Immutable content or configuration digest.
46    #[serde(default, skip_serializing_if = "Option::is_none")]
47    pub digest: Option<String>,
48    /// Deployment-defined trust labels such as `internal` or `pii-approved`.
49    #[serde(default, skip_serializing_if = "Vec::is_empty")]
50    pub trust_labels: Vec<String>,
51    /// Current registry health.
52    #[serde(default)]
53    pub health: TeamAgentHealth,
54    /// Unix expiry timestamp in milliseconds for ephemeral registrations.
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub expires_at_ms: Option<u64>,
57}
58
59/// Discovers and resolves local or remote agents for portable team members.
60#[async_trait]
61pub trait TeamAgentRegistry: Send + Sync {
62    /// Lists candidates visible to the current deployment and caller.
63    async fn candidates(&self) -> adk_core::Result<Vec<TeamAgentDescriptor>>;
64
65    /// Resolves one immutable binding identifier to an executable agent.
66    async fn resolve(&self, binding: &str) -> adk_core::Result<Arc<dyn Agent>>;
67
68    /// Authorizes a candidate for one portable member after metadata filtering.
69    async fn authorize(
70        &self,
71        _member: &TeamMemberSpec,
72        _candidate: &TeamAgentDescriptor,
73    ) -> adk_core::Result<bool> {
74        Ok(true)
75    }
76}
77
78/// Deterministic in-process registry useful for generated projects and tests.
79#[derive(Default)]
80pub struct StaticTeamAgentRegistry {
81    entries: HashMap<String, (TeamAgentDescriptor, Arc<dyn Agent>)>,
82}
83
84impl StaticTeamAgentRegistry {
85    /// Creates an empty registry.
86    pub fn new() -> Self {
87        Self::default()
88    }
89
90    /// Registers one immutable descriptor and executable binding.
91    pub fn register(
92        mut self,
93        descriptor: TeamAgentDescriptor,
94        agent: Arc<dyn Agent>,
95    ) -> std::result::Result<Self, TeamError> {
96        if self.entries.contains_key(&descriptor.binding) {
97            return Err(TeamError::Registry(format!(
98                "duplicate team registry binding '{}'",
99                descriptor.binding
100            )));
101        }
102        self.entries.insert(descriptor.binding.clone(), (descriptor, agent));
103        Ok(self)
104    }
105}
106
107#[async_trait]
108impl TeamAgentRegistry for StaticTeamAgentRegistry {
109    async fn candidates(&self) -> adk_core::Result<Vec<TeamAgentDescriptor>> {
110        Ok(self.entries.values().map(|(descriptor, _)| descriptor.clone()).collect())
111    }
112
113    async fn resolve(&self, binding: &str) -> adk_core::Result<Arc<dyn Agent>> {
114        self.entries
115            .get(binding)
116            .map(|(_, agent)| agent.clone())
117            .ok_or_else(|| adk_core::AdkError::agent(format!("unknown team binding '{binding}'")))
118    }
119}
120
121impl TeamSpec {
122    /// Resolves missing concrete members by capability and freezes the roster.
123    ///
124    /// Selection is reproducible: exact advertised-name matches win, followed
125    /// by descending registry priority and then lexical binding identifier.
126    pub async fn compile_with_registry(
127        &self,
128        registry: Arc<dyn TeamAgentRegistry>,
129    ) -> std::result::Result<CompiledTeam, TeamError> {
130        self.validate()?;
131        #[cfg(not(feature = "team-tools"))]
132        if self.relationships.iter().any(|edge| edge.kind == super::RelationshipKind::Delegate) {
133            return Err(TeamError::DelegationFeatureDisabled);
134        }
135        let candidates =
136            registry.candidates().await.map_err(|error| TeamError::Registry(error.to_string()))?;
137        let mut bindings = HashMap::new();
138        let mut frozen = Vec::with_capacity(self.members.len());
139        for member in &self.members {
140            let mut eligible: Vec<&TeamAgentDescriptor> = candidates
141                .iter()
142                .filter(|candidate| {
143                    let now_ms = std::time::SystemTime::now()
144                        .duration_since(std::time::UNIX_EPOCH)
145                        .unwrap_or_default()
146                        .as_millis() as u64;
147                    let has_capabilities = member.required_capabilities.iter().all(|required| {
148                        candidate.capabilities.iter().any(|available| available == required)
149                    });
150                    let health_allowed = candidate.health != TeamAgentHealth::Unavailable
151                        && (!member.registry.require_healthy
152                            || candidate.health == TeamAgentHealth::Healthy);
153                    let trust_allowed = member.registry.trust_labels.is_empty()
154                        || member.registry.trust_labels.iter().any(|required| {
155                            candidate.trust_labels.iter().any(|label| label == required)
156                        });
157                    has_capabilities
158                        && health_allowed
159                        && trust_allowed
160                        && candidate.expires_at_ms.is_none_or(|expiry| expiry > now_ms)
161                        && member
162                            .registry
163                            .version
164                            .as_ref()
165                            .is_none_or(|version| candidate.version.as_ref() == Some(version))
166                        && member
167                            .registry
168                            .digest
169                            .as_ref()
170                            .is_none_or(|digest| candidate.digest.as_ref() == Some(digest))
171                        && (!member.required_capabilities.is_empty()
172                            || candidate.name == member.name)
173                })
174                .collect();
175            eligible.sort_by_key(|candidate| {
176                (
177                    Reverse(candidate.name == member.name),
178                    Reverse(candidate.priority),
179                    candidate.binding.as_str(),
180                )
181            });
182            let mut selected = None;
183            for candidate in eligible {
184                if registry
185                    .authorize(member, candidate)
186                    .await
187                    .map_err(|error| TeamError::Registry(error.to_string()))?
188                {
189                    selected = Some(candidate);
190                    break;
191                }
192            }
193            let selected = selected.ok_or_else(|| TeamError::NoRegistryCandidate {
194                member: member.name.clone(),
195                capabilities: member.required_capabilities.clone(),
196            })?;
197            let agent = registry
198                .resolve(&selected.binding)
199                .await
200                .map_err(|error| TeamError::Registry(error.to_string()))?;
201            bindings.insert(member.name.clone(), agent);
202            frozen.push(ResolvedTeamMember {
203                member: member.name.clone(),
204                binding: selected.binding.clone(),
205                capabilities: selected.capabilities.clone(),
206                version: selected.version.clone(),
207                digest: selected.digest.clone(),
208                trust_labels: selected.trust_labels.clone(),
209            });
210        }
211        self.compile_registry(bindings, frozen, Vec::new())
212    }
213}
214
215#[cfg(test)]
216mod tests {
217    use super::*;
218    use crate::TeamMemberSpec;
219    use adk_core::{EventStream, InvocationContext, Result};
220
221    struct NamedAgent(String);
222
223    #[async_trait]
224    impl Agent for NamedAgent {
225        fn name(&self) -> &str {
226            &self.0
227        }
228
229        fn description(&self) -> &str {
230            "registry test agent"
231        }
232
233        fn sub_agents(&self) -> &[Arc<dyn Agent>] {
234            &[]
235        }
236
237        async fn run(&self, _ctx: Arc<dyn InvocationContext>) -> Result<EventStream> {
238            Ok(Box::pin(futures::stream::empty()))
239        }
240    }
241
242    fn descriptor(
243        binding: &str,
244        name: &str,
245        capabilities: &[&str],
246        priority: i32,
247    ) -> TeamAgentDescriptor {
248        TeamAgentDescriptor {
249            binding: binding.to_string(),
250            name: name.to_string(),
251            description: String::new(),
252            capabilities: capabilities.iter().map(ToString::to_string).collect(),
253            priority,
254            version: None,
255            digest: None,
256            trust_labels: Vec::new(),
257            health: TeamAgentHealth::Healthy,
258            expires_at_ms: None,
259        }
260    }
261
262    #[tokio::test]
263    async fn resolves_capabilities_deterministically_and_freezes_roster() {
264        let spec = TeamSpec {
265            name: "discovered".to_string(),
266            description: String::new(),
267            coordinator: "planner".to_string(),
268            members: vec![TeamMemberSpec::new("planner").with_capabilities(["plan"])],
269            relationships: Vec::new(),
270            policy: super::super::TeamPolicy::default(),
271        };
272        let registry = StaticTeamAgentRegistry::new()
273            .register(
274                descriptor("remote-low", "planner", &["plan"], 1),
275                Arc::new(NamedAgent("remote_low".to_string())),
276            )
277            .unwrap()
278            .register(
279                descriptor("remote-high", "alternate", &["plan", "web"], 9),
280                Arc::new(NamedAgent("remote_high".to_string())),
281            )
282            .unwrap();
283
284        let team = spec.compile_with_registry(Arc::new(registry)).await.unwrap();
285        team.runtime.check_budget("registry-test").unwrap();
286        let snapshot = team.runtime.snapshot("registry-test").unwrap();
287        assert_eq!(snapshot.roster[0].member, "planner");
288        assert_eq!(snapshot.roster[0].binding, "remote-low");
289        assert_eq!(team.sub_agents()[0].name(), "planner");
290    }
291
292    #[tokio::test]
293    async fn rejects_registry_without_required_capabilities() {
294        let spec = TeamSpec {
295            name: "discovered".to_string(),
296            description: String::new(),
297            coordinator: "planner".to_string(),
298            members: vec![TeamMemberSpec::new("planner").with_capabilities(["plan", "secure"])],
299            relationships: Vec::new(),
300            policy: super::super::TeamPolicy::default(),
301        };
302        let registry = StaticTeamAgentRegistry::new()
303            .register(
304                descriptor("candidate", "planner", &["plan"], 1),
305                Arc::new(NamedAgent("candidate".to_string())),
306            )
307            .unwrap();
308        assert!(matches!(
309            spec.compile_with_registry(Arc::new(registry)).await,
310            Err(TeamError::NoRegistryCandidate { .. })
311        ));
312    }
313
314    #[tokio::test]
315    async fn enforces_registry_health_trust_version_and_digest() {
316        let requirement = super::super::TeamRegistryRequirement {
317            version: Some("2.1.0".to_string()),
318            digest: Some("sha256:trusted".to_string()),
319            trust_labels: vec!["internal".to_string()],
320            require_healthy: true,
321        };
322        let spec = TeamSpec {
323            name: "governed".to_string(),
324            description: String::new(),
325            coordinator: "planner".to_string(),
326            members: vec![TeamMemberSpec::new("planner").with_registry_requirement(requirement)],
327            relationships: Vec::new(),
328            policy: super::super::TeamPolicy::default(),
329        };
330        let mut rejected = descriptor("rejected", "planner", &[], 10);
331        rejected.health = TeamAgentHealth::Degraded;
332        rejected.version = Some("2.1.0".to_string());
333        rejected.digest = Some("sha256:trusted".to_string());
334        rejected.trust_labels.push("internal".to_string());
335        let mut selected = descriptor("selected", "planner", &[], 1);
336        selected.version = Some("2.1.0".to_string());
337        selected.digest = Some("sha256:trusted".to_string());
338        selected.trust_labels.push("internal".to_string());
339        let registry = StaticTeamAgentRegistry::new()
340            .register(rejected, Arc::new(NamedAgent("rejected".to_string())))
341            .unwrap()
342            .register(selected, Arc::new(NamedAgent("selected".to_string())))
343            .unwrap();
344        let team = spec.compile_with_registry(Arc::new(registry)).await.unwrap();
345        team.runtime.check_budget("governed-run").unwrap();
346        let roster = team.runtime.snapshot("governed-run").unwrap().roster;
347        assert_eq!(roster[0].binding, "selected");
348        assert_eq!(roster[0].version.as_deref(), Some("2.1.0"));
349        assert_eq!(roster[0].trust_labels, ["internal"]);
350    }
351}
adk_agent/team/discovery.rs

adk_agent/team/
discovery.rs
