grapsus_agent_protocol/v2/
capabilities.rs

1//! Agent capability negotiation for Protocol v2.
2
3use crate::EventType;
4use serde::{Deserialize, Serialize};
5
6/// Agent capabilities declared during handshake.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct AgentCapabilities {
9    pub protocol_version: u32,
10    pub agent_id: String,
11    pub name: String,
12    pub version: String,
13    pub supported_events: Vec<EventType>,
14    #[serde(default)]
15    pub features: AgentFeatures,
16    #[serde(default)]
17    pub limits: AgentLimits,
18    #[serde(default)]
19    pub health: HealthConfig,
20}
21
22impl AgentCapabilities {
23    pub fn new(
24        agent_id: impl Into<String>,
25        name: impl Into<String>,
26        version: impl Into<String>,
27    ) -> Self {
28        Self {
29            protocol_version: super::PROTOCOL_VERSION_2,
30            agent_id: agent_id.into(),
31            name: name.into(),
32            version: version.into(),
33            supported_events: vec![EventType::RequestHeaders],
34            features: AgentFeatures::default(),
35            limits: AgentLimits::default(),
36            health: HealthConfig::default(),
37        }
38    }
39
40    pub fn supports_event(&self, event_type: EventType) -> bool {
41        self.supported_events.contains(&event_type)
42    }
43
44    pub fn with_event(mut self, event_type: EventType) -> Self {
45        if !self.supported_events.contains(&event_type) {
46            self.supported_events.push(event_type);
47        }
48        self
49    }
50
51    pub fn with_features(mut self, features: AgentFeatures) -> Self {
52        self.features = features;
53        self
54    }
55
56    pub fn with_limits(mut self, limits: AgentLimits) -> Self {
57        self.limits = limits;
58        self
59    }
60}
61
62/// Features this agent supports.
63#[derive(Debug, Clone, Default, Serialize, Deserialize)]
64pub struct AgentFeatures {
65    #[serde(default)]
66    pub streaming_body: bool,
67    #[serde(default)]
68    pub websocket: bool,
69    #[serde(default)]
70    pub guardrails: bool,
71    #[serde(default)]
72    pub config_push: bool,
73    #[serde(default)]
74    pub metrics_export: bool,
75    #[serde(default)]
76    pub concurrent_requests: u32,
77    #[serde(default)]
78    pub cancellation: bool,
79    #[serde(default)]
80    pub flow_control: bool,
81    #[serde(default)]
82    pub health_reporting: bool,
83}
84
85impl AgentFeatures {
86    pub fn simple() -> Self {
87        Self::default()
88    }
89    pub fn full() -> Self {
90        Self {
91            streaming_body: true,
92            websocket: true,
93            guardrails: true,
94            config_push: true,
95            metrics_export: true,
96            concurrent_requests: 100,
97            cancellation: true,
98            flow_control: true,
99            health_reporting: true,
100        }
101    }
102}
103
104/// Resource limits.
105#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct AgentLimits {
107    pub max_body_size: usize,
108    pub max_concurrency: u32,
109    pub preferred_chunk_size: usize,
110    pub max_memory: Option<usize>,
111    pub max_processing_time_ms: Option<u64>,
112}
113
114impl Default for AgentLimits {
115    fn default() -> Self {
116        Self {
117            max_body_size: 10 * 1024 * 1024,
118            max_concurrency: 100,
119            preferred_chunk_size: 64 * 1024,
120            max_memory: None,
121            max_processing_time_ms: Some(5000),
122        }
123    }
124}
125
126/// Health check configuration.
127#[derive(Debug, Clone, Serialize, Deserialize)]
128pub struct HealthConfig {
129    pub report_interval_ms: u32,
130    pub include_load_metrics: bool,
131    pub include_resource_metrics: bool,
132}
133
134impl Default for HealthConfig {
135    fn default() -> Self {
136        Self {
137            report_interval_ms: 10_000,
138            include_load_metrics: true,
139            include_resource_metrics: false,
140        }
141    }
142}
143
144/// Handshake request from proxy to agent.
145#[derive(Debug, Clone, Serialize, Deserialize)]
146pub struct HandshakeRequest {
147    pub supported_versions: Vec<u32>,
148    pub proxy_id: String,
149    pub proxy_version: String,
150    pub config: serde_json::Value,
151}
152
153impl HandshakeRequest {
154    pub fn new(proxy_id: impl Into<String>, proxy_version: impl Into<String>) -> Self {
155        Self {
156            supported_versions: vec![super::PROTOCOL_VERSION_2, 1],
157            proxy_id: proxy_id.into(),
158            proxy_version: proxy_version.into(),
159            config: serde_json::Value::Null,
160        }
161    }
162
163    pub fn max_version(&self) -> u32 {
164        self.supported_versions.first().copied().unwrap_or(1)
165    }
166}
167
168/// Handshake response from agent to proxy.
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct HandshakeResponse {
171    pub protocol_version: u32,
172    pub capabilities: AgentCapabilities,
173    pub success: bool,
174    pub error: Option<String>,
175}
176
177impl HandshakeResponse {
178    pub fn success(capabilities: AgentCapabilities) -> Self {
179        Self {
180            protocol_version: capabilities.protocol_version,
181            capabilities,
182            success: true,
183            error: None,
184        }
185    }
186
187    pub fn failure(error: impl Into<String>) -> Self {
188        Self {
189            protocol_version: 0,
190            capabilities: AgentCapabilities::new("", "", ""),
191            success: false,
192            error: Some(error.into()),
193        }
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_capabilities_builder() {
203        let caps = AgentCapabilities::new("test-agent", "Test Agent", "1.0.0")
204            .with_event(EventType::RequestHeaders)
205            .with_features(AgentFeatures::full());
206
207        assert_eq!(caps.agent_id, "test-agent");
208        assert!(caps.supports_event(EventType::RequestHeaders));
209        assert!(caps.features.streaming_body);
210    }
211
212    #[test]
213    fn test_handshake() {
214        let request = HandshakeRequest::new("proxy-1", "0.2.5");
215        assert_eq!(request.max_version(), 2);
216
217        let caps = AgentCapabilities::new("agent-1", "My Agent", "1.0.0");
218        let response = HandshakeResponse::success(caps);
219        assert!(response.success);
220    }
221
222    #[test]
223    fn test_features() {
224        let simple = AgentFeatures::simple();
225        assert!(!simple.streaming_body);
226
227        let full = AgentFeatures::full();
228        assert!(full.streaming_body);
229    }
230}
grapsus_agent_protocol/v2/capabilities.rs

grapsus_agent_protocol/v2/
capabilities.rs
