a2a_rs/domain/core/
agent.rs

1use std::collections::HashMap;
2
3// Re-export generated types so downstream code gets them from `domain::core::agent`
4pub use crate::domain::generated::{
5    APIKeySecurityScheme, AgentCapabilities, AgentCard, AgentCardSignature, AgentExtension,
6    AgentInterface, AgentProvider, AgentSkill, AuthenticationInfo, AuthorizationCodeOAuthFlow,
7    ClientCredentialsOAuthFlow, DeviceCodeOAuthFlow, HTTPAuthSecurityScheme, ImplicitOAuthFlow,
8    MutualTlsSecurityScheme, OAuth2SecurityScheme, OAuthFlows, OpenIdConnectSecurityScheme,
9    PasswordOAuthFlow, SecurityRequirement, SecurityScheme, StringList, o_auth_flows,
10    security_scheme,
11};
12
13pub type PushNotificationAuthenticationInfo = AuthenticationInfo;
14
15/// JSON-RPC 2.0 over HTTP — the spec-default protocol binding, served by
16/// [`jsonrpc_router`](crate::adapter::transport::jsonrpc_router).
17pub const PROTOCOL_BINDING_JSONRPC: &str = "JSONRPC";
18
19/// ConnectRPC — the in-tree first-class streaming binding, served by
20/// [`HttpServer`](crate::adapter::HttpServer) over a
21/// [`ConnectRpcAdapter`](crate::adapter::ConnectRpcAdapter).
22pub const PROTOCOL_BINDING_CONNECTRPC: &str = "CONNECTRPC";
23
24/// Plain REST/JSON over HTTP, served by
25/// [`rest_router`](crate::adapter::transport::rest_router).
26pub const PROTOCOL_BINDING_HTTP_JSON: &str = "HTTP+JSON";
27
28impl AgentSkill {
29    /// Create a new skill with the minimum required fields
30    pub fn new(id: String, name: String, description: String, tags: Vec<String>) -> Self {
31        Self {
32            id,
33            name,
34            description,
35            tags,
36            ..Default::default()
37        }
38    }
39
40    /// Add examples to the skill
41    pub fn with_examples(mut self, examples: Vec<String>) -> Self {
42        self.examples = examples;
43        self
44    }
45
46    /// Add input modes to the skill
47    pub fn with_input_modes(mut self, input_modes: Vec<String>) -> Self {
48        self.input_modes = input_modes;
49        self
50    }
51
52    /// Add output modes to the skill
53    pub fn with_output_modes(mut self, output_modes: Vec<String>) -> Self {
54        self.output_modes = output_modes;
55        self
56    }
57
58    /// Add security requirements to the skill
59    pub fn with_security(mut self, security: Vec<HashMap<String, Vec<String>>>) -> Self {
60        self.security_requirements = security
61            .into_iter()
62            .map(|req| {
63                let schemes = req
64                    .into_iter()
65                    .map(|(k, v)| {
66                        (
67                            k,
68                            StringList {
69                                list: v,
70                                ..Default::default()
71                            },
72                        )
73                    })
74                    .collect();
75                SecurityRequirement {
76                    schemes,
77                    ..Default::default()
78                }
79            })
80            .collect();
81        self
82    }
83
84    /// Create a comprehensive skill with all details in one call
85    #[allow(clippy::too_many_arguments)]
86    pub fn comprehensive(
87        id: String,
88        name: String,
89        description: String,
90        tags: Vec<String>,
91        examples: Option<Vec<String>>,
92        input_modes: Option<Vec<String>>,
93        output_modes: Option<Vec<String>>,
94        security: Option<Vec<HashMap<String, Vec<String>>>>,
95    ) -> Self {
96        let mut skill = Self::new(id, name, description, tags);
97        if let Some(ex) = examples {
98            skill = skill.with_examples(ex);
99        }
100        if let Some(im) = input_modes {
101            skill = skill.with_input_modes(im);
102        }
103        if let Some(om) = output_modes {
104            skill = skill.with_output_modes(om);
105        }
106        if let Some(sec) = security {
107            skill = skill.with_security(sec);
108        }
109        skill
110    }
111}
112
113impl SecurityScheme {
114    pub fn api_key(name: String, location: String, description: Option<String>) -> Self {
115        Self {
116            scheme: Some(security_scheme::Scheme::ApiKeySecurityScheme(Box::new(
117                APIKeySecurityScheme {
118                    name,
119                    location,
120                    description: description.unwrap_or_default(),
121                    ..Default::default()
122                },
123            ))),
124            ..Default::default()
125        }
126    }
127
128    pub fn http(
129        scheme_name: String,
130        bearer_format: Option<String>,
131        description: Option<String>,
132    ) -> Self {
133        Self {
134            scheme: Some(security_scheme::Scheme::HttpAuthSecurityScheme(Box::new(
135                HTTPAuthSecurityScheme {
136                    scheme: scheme_name,
137                    bearer_format: bearer_format.unwrap_or_default(),
138                    description: description.unwrap_or_default(),
139                    ..Default::default()
140                },
141            ))),
142            ..Default::default()
143        }
144    }
145
146    pub fn oauth2(
147        flows: OAuthFlows,
148        description: Option<String>,
149        oauth2_metadata_url: Option<String>,
150    ) -> Self {
151        Self {
152            scheme: Some(security_scheme::Scheme::Oauth2SecurityScheme(Box::new(
153                OAuth2SecurityScheme {
154                    flows: ::buffa::MessageField::some(flows),
155                    description: description.unwrap_or_default(),
156                    oauth2_metadata_url: oauth2_metadata_url.unwrap_or_default(),
157                    ..Default::default()
158                },
159            ))),
160            ..Default::default()
161        }
162    }
163
164    pub fn open_id_connect(open_id_connect_url: String, description: Option<String>) -> Self {
165        Self {
166            scheme: Some(security_scheme::Scheme::OpenIdConnectSecurityScheme(
167                Box::new(OpenIdConnectSecurityScheme {
168                    open_id_connect_url,
169                    description: description.unwrap_or_default(),
170                    ..Default::default()
171                }),
172            )),
173            ..Default::default()
174        }
175    }
176
177    pub fn mutual_tls(description: Option<String>) -> Self {
178        Self {
179            scheme: Some(security_scheme::Scheme::MtlsSecurityScheme(Box::new(
180                MutualTlsSecurityScheme {
181                    description: description.unwrap_or_default(),
182                    ..Default::default()
183                },
184            ))),
185            ..Default::default()
186        }
187    }
188}
189
190impl OAuthFlows {
191    pub fn authorization_code(flow: AuthorizationCodeOAuthFlow) -> Self {
192        Self {
193            flow: Some(o_auth_flows::Flow::AuthorizationCode(Box::new(flow))),
194            ..Default::default()
195        }
196    }
197
198    pub fn client_credentials(flow: ClientCredentialsOAuthFlow) -> Self {
199        Self {
200            flow: Some(o_auth_flows::Flow::ClientCredentials(Box::new(flow))),
201            ..Default::default()
202        }
203    }
204
205    pub fn device_code(flow: DeviceCodeOAuthFlow) -> Self {
206        Self {
207            flow: Some(o_auth_flows::Flow::DeviceCode(Box::new(flow))),
208            ..Default::default()
209        }
210    }
211}
212
213impl AgentCapabilities {
214    pub fn streaming(&self) -> bool {
215        self.streaming.unwrap_or(false)
216    }
217
218    pub fn push_notifications(&self) -> bool {
219        self.push_notifications.unwrap_or(false)
220    }
221
222    pub fn extended_agent_card(&self) -> bool {
223        self.extended_agent_card.unwrap_or(false)
224    }
225}
226
227impl AgentCardSignature {
228    pub fn new(
229        protected: String,
230        signature: String,
231        header: Option<::buffa_types::google::protobuf::Struct>,
232    ) -> Self {
233        Self {
234            protected,
235            signature,
236            header: header.into(),
237            ..Default::default()
238        }
239    }
240}
241
242impl AgentCard {
243    pub fn builder() -> AgentCardBuilder {
244        AgentCardBuilder::new()
245    }
246
247    pub fn url(&self) -> &str {
248        self.supported_interfaces
249            .first()
250            .map(|i| i.url.as_str())
251            .unwrap_or("")
252    }
253
254    pub fn protocol_version(&self) -> &str {
255        self.supported_interfaces
256            .first()
257            .map(|i| i.protocol_version.as_str())
258            .unwrap_or("1.0")
259    }
260
261    pub fn preferred_transport(&self) -> &str {
262        self.supported_interfaces
263            .first()
264            .map(|i| i.protocol_binding.as_str())
265            .unwrap_or(PROTOCOL_BINDING_JSONRPC)
266    }
267
268    pub fn supports_extended_agent_card(&self) -> bool {
269        self.capabilities.extended_agent_card.unwrap_or(false)
270    }
271}
272
273pub struct AgentCardBuilder {
274    name: String,
275    description: String,
276    url: String,
277    provider: Option<AgentProvider>,
278    version: String,
279    protocol_version: Option<String>,
280    preferred_transport: Option<String>,
281    supported_interfaces: Vec<AgentInterface>,
282    icon_url: Option<String>,
283    documentation_url: Option<String>,
284    capabilities: Option<AgentCapabilities>,
285    security_schemes: HashMap<String, SecurityScheme>,
286    security_requirements: Vec<SecurityRequirement>,
287    default_input_modes: Vec<String>,
288    default_output_modes: Vec<String>,
289    skills: Vec<AgentSkill>,
290    signatures: Vec<AgentCardSignature>,
291}
292
293impl Default for AgentCardBuilder {
294    fn default() -> Self {
295        Self::new()
296    }
297}
298
299impl AgentCardBuilder {
300    pub fn new() -> Self {
301        Self {
302            name: String::new(),
303            description: String::new(),
304            url: String::new(),
305            provider: None,
306            version: String::new(),
307            protocol_version: None,
308            preferred_transport: None,
309            supported_interfaces: Vec::new(),
310            icon_url: None,
311            documentation_url: None,
312            capabilities: None,
313            security_schemes: HashMap::new(),
314            security_requirements: Vec::new(),
315            default_input_modes: vec!["text".to_string()],
316            default_output_modes: vec!["text".to_string()],
317            skills: Vec::new(),
318            signatures: Vec::new(),
319        }
320    }
321
322    pub fn name(mut self, name: String) -> Self {
323        self.name = name;
324        self
325    }
326
327    pub fn description(mut self, description: String) -> Self {
328        self.description = description;
329        self
330    }
331
332    pub fn url(mut self, url: String) -> Self {
333        self.url = url;
334        self
335    }
336
337    pub fn provider(mut self, provider: AgentProvider) -> Self {
338        self.provider = Some(provider);
339        self
340    }
341
342    pub fn version(mut self, version: String) -> Self {
343        self.version = version;
344        self
345    }
346
347    pub fn protocol_version(mut self, protocol_version: String) -> Self {
348        self.protocol_version = Some(protocol_version);
349        self
350    }
351
352    pub fn preferred_transport(mut self, preferred_transport: String) -> Self {
353        self.preferred_transport = Some(preferred_transport);
354        self
355    }
356
357    pub fn additional_interfaces(mut self, interfaces: Vec<AgentInterface>) -> Self {
358        self.supported_interfaces.extend(interfaces);
359        self
360    }
361
362    pub fn icon_url(mut self, icon_url: String) -> Self {
363        self.icon_url = Some(icon_url);
364        self
365    }
366
367    pub fn documentation_url(mut self, documentation_url: String) -> Self {
368        self.documentation_url = Some(documentation_url);
369        self
370    }
371
372    pub fn capabilities(mut self, capabilities: AgentCapabilities) -> Self {
373        self.capabilities = Some(capabilities);
374        self
375    }
376
377    pub fn security_schemes(mut self, security_schemes: HashMap<String, SecurityScheme>) -> Self {
378        self.security_schemes = security_schemes;
379        self
380    }
381
382    pub fn security(mut self, security: Vec<HashMap<String, Vec<String>>>) -> Self {
383        self.security_requirements = security
384            .into_iter()
385            .map(|req| {
386                let schemes = req
387                    .into_iter()
388                    .map(|(k, v)| {
389                        (
390                            k,
391                            StringList {
392                                list: v,
393                                ..Default::default()
394                            },
395                        )
396                    })
397                    .collect();
398                SecurityRequirement {
399                    schemes,
400                    ..Default::default()
401                }
402            })
403            .collect();
404        self
405    }
406
407    pub fn default_input_modes(mut self, default_input_modes: Vec<String>) -> Self {
408        self.default_input_modes = default_input_modes;
409        self
410    }
411
412    pub fn default_output_modes(mut self, default_output_modes: Vec<String>) -> Self {
413        self.default_output_modes = default_output_modes;
414        self
415    }
416
417    pub fn skills(mut self, skills: Vec<AgentSkill>) -> Self {
418        self.skills = skills;
419        self
420    }
421
422    pub fn signatures(mut self, signatures: Vec<AgentCardSignature>) -> Self {
423        self.signatures = signatures;
424        self
425    }
426
427    pub fn supports_extended_agent_card(mut self, val: bool) -> Self {
428        let caps = self
429            .capabilities
430            .get_or_insert_with(AgentCapabilities::default);
431        caps.extended_agent_card = Some(val);
432        self
433    }
434
435    pub fn build(self) -> AgentCard {
436        let mut supported_interfaces = self.supported_interfaces;
437        // Make sure the primary interface exists and is first
438        if !self.url.is_empty() {
439            let primary = AgentInterface {
440                url: self.url,
441                protocol_binding: self
442                    .preferred_transport
443                    .unwrap_or_else(|| PROTOCOL_BINDING_JSONRPC.to_string()),
444                protocol_version: self.protocol_version.unwrap_or_else(|| "1.0".to_string()),
445                ..Default::default()
446            };
447            supported_interfaces.insert(0, primary);
448        }
449
450        AgentCard {
451            name: self.name,
452            description: self.description,
453            supported_interfaces,
454            provider: self.provider.into(),
455            version: self.version,
456            documentation_url: self.documentation_url,
457            capabilities: self.capabilities.unwrap_or_default().into(),
458            security_schemes: self.security_schemes,
459            security_requirements: self.security_requirements,
460            default_input_modes: self.default_input_modes,
461            default_output_modes: self.default_output_modes,
462            skills: self.skills,
463            signatures: self.signatures,
464            icon_url: self.icon_url,
465            ..Default::default()
466        }
467    }
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473
474    #[test]
475    fn test_security_scheme_api_key_serialization() {
476        let scheme = SecurityScheme::api_key(
477            "X-API-Key".to_string(),
478            "header".to_string(),
479            Some("API Key authentication".to_string()),
480        );
481
482        let json_value = serde_json::to_value(&scheme).expect("Failed to serialize SecurityScheme");
483        // Verify output matches protobuf JSON mappings
484        assert_eq!(json_value["apiKeySecurityScheme"]["location"], "header");
485        assert_eq!(json_value["apiKeySecurityScheme"]["name"], "X-API-Key");
486    }
487
488    #[test]
489    fn test_security_scheme_http_serialization() {
490        let scheme = SecurityScheme::http(
491            "bearer".to_string(),
492            Some("JWT".to_string()),
493            Some("Bearer token authentication".to_string()),
494        );
495
496        let json_value = serde_json::to_value(&scheme).expect("Failed to serialize SecurityScheme");
497        assert_eq!(json_value["httpAuthSecurityScheme"]["scheme"], "bearer");
498        assert_eq!(json_value["httpAuthSecurityScheme"]["bearerFormat"], "JWT");
499    }
500
501    #[test]
502    fn test_security_scheme_mtls_serialization() {
503        let scheme = SecurityScheme::mutual_tls(Some("Mutual TLS authentication".to_string()));
504
505        let json_value = serde_json::to_value(&scheme).expect("Failed to serialize SecurityScheme");
506        assert_eq!(
507            json_value["mtlsSecurityScheme"]["description"],
508            "Mutual TLS authentication"
509        );
510    }
511
512    #[test]
513    fn test_security_scheme_oauth2_with_metadata() {
514        let flows = OAuthFlows::authorization_code(AuthorizationCodeOAuthFlow {
515            authorization_url: "https://example.com/oauth/authorize".to_string(),
516            token_url: "https://example.com/oauth/token".to_string(),
517            refresh_url: String::new(),
518            scopes: HashMap::new(),
519            ..Default::default()
520        });
521
522        let scheme = SecurityScheme::oauth2(
523            flows,
524            Some("OAuth2 authentication".to_string()),
525            Some("https://example.com/.well-known/oauth-authorization-server".to_string()),
526        );
527
528        let json_value = serde_json::to_value(&scheme).expect("Failed to serialize SecurityScheme");
529        assert_eq!(
530            json_value["oauth2SecurityScheme"]["oauth2MetadataUrl"],
531            "https://example.com/.well-known/oauth-authorization-server"
532        );
533    }
534}
a2a_rs/domain/core/agent.rs

a2a_rs/domain/core/
agent.rs
