Skip to main content

ai_lib_rust/protocol/v2/
capabilities.rs

1//! V2 能力声明系统 — 支持 required/optional 分离和 feature_flags 精细控制
2//!
3//! V2 capability declaration system with structured required/optional separation,
4//! feature flags, and capability-to-module mapping for runtime loading.
5
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9/// Standard capability identifiers aligned with `schemas/v2/capabilities.json`.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
11#[serde(rename_all = "snake_case")]
12pub enum Capability {
13    Text,
14    Streaming,
15    Vision,
16    Audio,
17    Video,
18    Tools,
19    ParallelTools,
20    Agentic,
21    Reasoning,
22    Embeddings,
23    StructuredOutput,
24    Batch,
25    ImageGeneration,
26    ComputerUse,
27    McpClient,
28    McpServer,
29}
30
31impl Capability {
32    /// Map capability to the corresponding Cargo feature flag name.
33    pub fn feature_flag(&self) -> Option<&'static str> {
34        match self {
35            Self::Text | Self::Streaming | Self::Tools | Self::ParallelTools => None, // always loaded
36            Self::Vision => Some("vision"),
37            Self::Audio | Self::Video => Some("multimodal"),
38            Self::Agentic => Some("agentic"),
39            Self::Reasoning => Some("reasoning"),
40            Self::Embeddings => Some("embeddings"),
41            Self::StructuredOutput => Some("structured"),
42            Self::Batch => Some("batch"),
43            Self::ImageGeneration => Some("image_gen"),
44            Self::ComputerUse => Some("computer_use"),
45            Self::McpClient | Self::McpServer => Some("mcp"),
46        }
47    }
48
49    /// Check whether this capability requires a feature flag to be compiled in.
50    pub fn is_feature_gated(&self) -> bool {
51        self.feature_flag().is_some()
52    }
53
54    /// Get the runtime module path this capability maps to.
55    pub fn module_path(&self) -> &'static str {
56        match self {
57            Self::Text => "core",
58            Self::Streaming => "streaming",
59            Self::Vision => "multimodal.vision",
60            Self::Audio => "multimodal.audio",
61            Self::Video => "multimodal.video",
62            Self::Tools => "tools",
63            Self::ParallelTools => "tools.parallel",
64            Self::Agentic => "agentic",
65            Self::Reasoning => "reasoning",
66            Self::Embeddings => "embeddings",
67            Self::StructuredOutput => "structured",
68            Self::Batch => "batch",
69            Self::ImageGeneration => "generation.image",
70            Self::ComputerUse => "computer_use",
71            Self::McpClient => "mcp.client",
72            Self::McpServer => "mcp.server",
73        }
74    }
75}
76
77/// V2 structured capability declaration with required/optional separation.
78#[derive(Debug, Clone, Serialize, Deserialize)]
79#[serde(untagged)]
80pub enum CapabilitiesV2 {
81    /// V2 structured format: `{ required: [...], optional: [...], feature_flags: {...} }`
82    Structured {
83        required: Vec<Capability>,
84        #[serde(default)]
85        optional: Vec<Capability>,
86        #[serde(default)]
87        feature_flags: FeatureFlags,
88    },
89    /// V1 legacy flat format: `{ streaming: true, tools: true, vision: false }`
90    Legacy(LegacyCapabilities),
91}
92
93impl CapabilitiesV2 {
94    /// Get all capabilities (required + optional) as a unified set.
95    pub fn all_capabilities(&self) -> Vec<Capability> {
96        match self {
97            Self::Structured { required, optional, .. } => {
98                let mut all = required.clone();
99                all.extend(optional.iter().cloned());
100                all
101            }
102            Self::Legacy(legacy) => legacy.to_capabilities(),
103        }
104    }
105
106    /// Get only the required capabilities.
107    pub fn required_capabilities(&self) -> Vec<Capability> {
108        match self {
109            Self::Structured { required, .. } => required.clone(),
110            Self::Legacy(legacy) => {
111                // V1 legacy: text is always required, streaming if declared
112                let mut req = vec![Capability::Text];
113                if legacy.streaming {
114                    req.push(Capability::Streaming);
115                }
116                req
117            }
118        }
119    }
120
121    /// Check if a specific capability is declared (required or optional).
122    pub fn has_capability(&self, cap: Capability) -> bool {
123        self.all_capabilities().contains(&cap)
124    }
125
126    /// Get the feature flags.
127    pub fn feature_flags(&self) -> FeatureFlags {
128        match self {
129            Self::Structured { feature_flags, .. } => feature_flags.clone(),
130            Self::Legacy(_) => FeatureFlags::default(),
131        }
132    }
133
134    /// Auto-promote V1 legacy capabilities to V2 structured format.
135    pub fn promote_to_v2(&self) -> Self {
136        match self {
137            Self::Structured { .. } => self.clone(),
138            Self::Legacy(legacy) => {
139                let mut required = vec![Capability::Text];
140                let mut optional = Vec::new();
141
142                if legacy.streaming {
143                    required.push(Capability::Streaming);
144                }
145                if legacy.tools {
146                    optional.push(Capability::Tools);
147                }
148                if legacy.vision {
149                    optional.push(Capability::Vision);
150                }
151                if legacy.agentic {
152                    optional.push(Capability::Agentic);
153                }
154                if legacy.reasoning {
155                    optional.push(Capability::Reasoning);
156                }
157                if legacy.parallel_tools {
158                    optional.push(Capability::ParallelTools);
159                }
160
161                Self::Structured {
162                    required,
163                    optional,
164                    feature_flags: FeatureFlags::default(),
165                }
166            }
167        }
168    }
169}
170
171/// V1 legacy boolean capability flags — backward compatible.
172#[derive(Debug, Clone, Serialize, Deserialize)]
173pub struct LegacyCapabilities {
174    #[serde(default)]
175    pub streaming: bool,
176    #[serde(default)]
177    pub tools: bool,
178    #[serde(default)]
179    pub vision: bool,
180    #[serde(default)]
181    pub agentic: bool,
182    #[serde(default)]
183    pub reasoning: bool,
184    #[serde(default)]
185    pub parallel_tools: bool,
186}
187
188impl LegacyCapabilities {
189    fn to_capabilities(&self) -> Vec<Capability> {
190        let mut caps = vec![Capability::Text];
191        if self.streaming { caps.push(Capability::Streaming); }
192        if self.tools { caps.push(Capability::Tools); }
193        if self.vision { caps.push(Capability::Vision); }
194        if self.agentic { caps.push(Capability::Agentic); }
195        if self.reasoning { caps.push(Capability::Reasoning); }
196        if self.parallel_tools { caps.push(Capability::ParallelTools); }
197        caps
198    }
199}
200
201/// Fine-grained feature toggles within capabilities.
202#[derive(Debug, Clone, Default, Serialize, Deserialize)]
203pub struct FeatureFlags {
204    #[serde(default)]
205    pub structured_output: bool,
206    #[serde(default)]
207    pub parallel_tool_calls: bool,
208    #[serde(default)]
209    pub extended_thinking: bool,
210    #[serde(default)]
211    pub streaming_usage: bool,
212    #[serde(default)]
213    pub system_messages: bool,
214    #[serde(default)]
215    pub image_generation: bool,
216    /// Additional provider-specific flags.
217    #[serde(flatten)]
218    pub extra: HashMap<String, bool>,
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224
225    #[test]
226    fn test_capability_feature_flags() {
227        assert_eq!(Capability::Text.feature_flag(), None);
228        assert_eq!(Capability::McpClient.feature_flag(), Some("mcp"));
229        assert_eq!(Capability::ComputerUse.feature_flag(), Some("computer_use"));
230        assert!(!Capability::Streaming.is_feature_gated());
231        assert!(Capability::Audio.is_feature_gated());
232    }
233
234    #[test]
235    fn test_v2_capabilities_structured() {
236        let json = r#"{
237            "required": ["text", "streaming", "tools"],
238            "optional": ["vision", "mcp_client"],
239            "feature_flags": {"structured_output": true}
240        }"#;
241        let caps: CapabilitiesV2 = serde_json::from_str(json).unwrap();
242        assert!(caps.has_capability(Capability::Text));
243        assert!(caps.has_capability(Capability::McpClient));
244        assert!(!caps.has_capability(Capability::ComputerUse));
245        assert!(caps.feature_flags().structured_output);
246    }
247
248    #[test]
249    fn test_legacy_promotion() {
250        let legacy = LegacyCapabilities {
251            streaming: true,
252            tools: true,
253            vision: true,
254            agentic: false,
255            reasoning: false,
256            parallel_tools: false,
257        };
258        let v1 = CapabilitiesV2::Legacy(legacy);
259        let v2 = v1.promote_to_v2();
260        match &v2 {
261            CapabilitiesV2::Structured { required, optional, .. } => {
262                assert!(required.contains(&Capability::Text));
263                assert!(required.contains(&Capability::Streaming));
264                assert!(optional.contains(&Capability::Tools));
265                assert!(optional.contains(&Capability::Vision));
266            }
267            _ => panic!("Expected Structured"),
268        }
269    }
270}