zag_agent/
capability.rs

1use anyhow::{Result, bail};
2use serde::{Deserialize, Serialize};
3
4/// A feature that can be either natively supported by the provider or implemented by the wrapper.
5#[derive(Debug, Clone, Serialize, Deserialize)]
6pub struct FeatureSupport {
7    pub supported: bool,
8    pub native: bool,
9}
10
11/// Session log support with completeness level.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct SessionLogSupport {
14    pub supported: bool,
15    pub native: bool,
16    /// Completeness level: "full", "partial", or absent when unsupported.
17    #[serde(skip_serializing_if = "Option::is_none")]
18    pub completeness: Option<String>,
19}
20
21/// Size alias mappings for a provider.
22#[derive(Debug, Clone, Serialize, Deserialize)]
23pub struct SizeMappings {
24    pub small: String,
25    pub medium: String,
26    pub large: String,
27}
28
29/// All feature flags for a provider.
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct Features {
32    pub interactive: FeatureSupport,
33    pub non_interactive: FeatureSupport,
34    pub resume: FeatureSupport,
35    pub resume_with_prompt: FeatureSupport,
36    pub session_logs: SessionLogSupport,
37    pub json_output: FeatureSupport,
38    pub stream_json: FeatureSupport,
39    pub json_schema: FeatureSupport,
40    pub input_format: FeatureSupport,
41    pub streaming_input: FeatureSupport,
42    pub worktree: FeatureSupport,
43    pub sandbox: FeatureSupport,
44    pub system_prompt: FeatureSupport,
45    pub auto_approve: FeatureSupport,
46    pub review: FeatureSupport,
47    pub add_dirs: FeatureSupport,
48    pub max_turns: FeatureSupport,
49}
50
51/// Full capability declaration for a provider.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct ProviderCapability {
54    pub provider: String,
55    pub default_model: String,
56    pub available_models: Vec<String>,
57    pub size_mappings: SizeMappings,
58    pub features: Features,
59}
60
61impl FeatureSupport {
62    pub fn native() -> Self {
63        Self {
64            supported: true,
65            native: true,
66        }
67    }
68
69    pub fn wrapper() -> Self {
70        Self {
71            supported: true,
72            native: false,
73        }
74    }
75
76    pub fn unsupported() -> Self {
77        Self {
78            supported: false,
79            native: false,
80        }
81    }
82}
83
84impl SessionLogSupport {
85    pub fn full() -> Self {
86        Self {
87            supported: true,
88            native: true,
89            completeness: Some("full".to_string()),
90        }
91    }
92
93    pub fn partial() -> Self {
94        Self {
95            supported: true,
96            native: true,
97            completeness: Some("partial".to_string()),
98        }
99    }
100
101    pub fn unsupported() -> Self {
102        Self {
103            supported: false,
104            native: false,
105            completeness: None,
106        }
107    }
108}
109
110/// Get capability declarations for a provider.
111pub fn get_capability(provider: &str) -> Result<ProviderCapability> {
112    use crate::agent::{Agent, ModelSize};
113
114    match provider {
115        "claude" => {
116            use crate::providers::claude::{self, Claude};
117            Ok(ProviderCapability {
118                provider: "claude".to_string(),
119                default_model: claude::DEFAULT_MODEL.to_string(),
120                available_models: models_to_vec(claude::AVAILABLE_MODELS),
121                size_mappings: SizeMappings {
122                    small: Claude::model_for_size(ModelSize::Small).to_string(),
123                    medium: Claude::model_for_size(ModelSize::Medium).to_string(),
124                    large: Claude::model_for_size(ModelSize::Large).to_string(),
125                },
126                features: Features {
127                    interactive: FeatureSupport::native(),
128                    non_interactive: FeatureSupport::native(),
129                    resume: FeatureSupport::native(),
130                    resume_with_prompt: FeatureSupport::native(),
131                    session_logs: SessionLogSupport::full(),
132                    json_output: FeatureSupport::native(),
133                    stream_json: FeatureSupport::native(),
134                    json_schema: FeatureSupport::native(),
135                    input_format: FeatureSupport::native(),
136                    streaming_input: FeatureSupport::native(),
137                    worktree: FeatureSupport::wrapper(),
138                    sandbox: FeatureSupport::wrapper(),
139                    system_prompt: FeatureSupport::native(),
140                    auto_approve: FeatureSupport::native(),
141                    review: FeatureSupport::unsupported(),
142                    add_dirs: FeatureSupport::native(),
143                    max_turns: FeatureSupport::native(),
144                },
145            })
146        }
147        "codex" => {
148            use crate::providers::codex::{self, Codex};
149            Ok(ProviderCapability {
150                provider: "codex".to_string(),
151                default_model: codex::DEFAULT_MODEL.to_string(),
152                available_models: models_to_vec(codex::AVAILABLE_MODELS),
153                size_mappings: SizeMappings {
154                    small: Codex::model_for_size(ModelSize::Small).to_string(),
155                    medium: Codex::model_for_size(ModelSize::Medium).to_string(),
156                    large: Codex::model_for_size(ModelSize::Large).to_string(),
157                },
158                features: Features {
159                    interactive: FeatureSupport::native(),
160                    non_interactive: FeatureSupport::native(),
161                    resume: FeatureSupport::native(),
162                    resume_with_prompt: FeatureSupport::native(),
163                    session_logs: SessionLogSupport::partial(),
164                    json_output: FeatureSupport::native(),
165                    stream_json: FeatureSupport::unsupported(),
166                    json_schema: FeatureSupport::wrapper(),
167                    input_format: FeatureSupport::unsupported(),
168                    streaming_input: FeatureSupport::unsupported(),
169                    worktree: FeatureSupport::wrapper(),
170                    sandbox: FeatureSupport::wrapper(),
171                    system_prompt: FeatureSupport::wrapper(),
172                    auto_approve: FeatureSupport::native(),
173                    review: FeatureSupport::native(),
174                    add_dirs: FeatureSupport::native(),
175                    max_turns: FeatureSupport::native(),
176                },
177            })
178        }
179        "gemini" => {
180            use crate::providers::gemini::{self, Gemini};
181            Ok(ProviderCapability {
182                provider: "gemini".to_string(),
183                default_model: gemini::DEFAULT_MODEL.to_string(),
184                available_models: models_to_vec(gemini::AVAILABLE_MODELS),
185                size_mappings: SizeMappings {
186                    small: Gemini::model_for_size(ModelSize::Small).to_string(),
187                    medium: Gemini::model_for_size(ModelSize::Medium).to_string(),
188                    large: Gemini::model_for_size(ModelSize::Large).to_string(),
189                },
190                features: Features {
191                    interactive: FeatureSupport::native(),
192                    non_interactive: FeatureSupport::native(),
193                    resume: FeatureSupport::native(),
194                    resume_with_prompt: FeatureSupport::unsupported(),
195                    session_logs: SessionLogSupport::full(),
196                    json_output: FeatureSupport::wrapper(),
197                    stream_json: FeatureSupport::unsupported(),
198                    json_schema: FeatureSupport::wrapper(),
199                    input_format: FeatureSupport::unsupported(),
200                    streaming_input: FeatureSupport::unsupported(),
201                    worktree: FeatureSupport::wrapper(),
202                    sandbox: FeatureSupport::wrapper(),
203                    system_prompt: FeatureSupport::wrapper(),
204                    auto_approve: FeatureSupport::native(),
205                    review: FeatureSupport::unsupported(),
206                    add_dirs: FeatureSupport::native(),
207                    max_turns: FeatureSupport::native(),
208                },
209            })
210        }
211        "copilot" => {
212            use crate::providers::copilot::{self, Copilot};
213            Ok(ProviderCapability {
214                provider: "copilot".to_string(),
215                default_model: copilot::DEFAULT_MODEL.to_string(),
216                available_models: models_to_vec(copilot::AVAILABLE_MODELS),
217                size_mappings: SizeMappings {
218                    small: Copilot::model_for_size(ModelSize::Small).to_string(),
219                    medium: Copilot::model_for_size(ModelSize::Medium).to_string(),
220                    large: Copilot::model_for_size(ModelSize::Large).to_string(),
221                },
222                features: Features {
223                    interactive: FeatureSupport::native(),
224                    non_interactive: FeatureSupport::native(),
225                    resume: FeatureSupport::native(),
226                    resume_with_prompt: FeatureSupport::unsupported(),
227                    session_logs: SessionLogSupport::full(),
228                    json_output: FeatureSupport::unsupported(),
229                    stream_json: FeatureSupport::unsupported(),
230                    json_schema: FeatureSupport::unsupported(),
231                    input_format: FeatureSupport::unsupported(),
232                    streaming_input: FeatureSupport::unsupported(),
233                    worktree: FeatureSupport::wrapper(),
234                    sandbox: FeatureSupport::wrapper(),
235                    system_prompt: FeatureSupport::wrapper(),
236                    auto_approve: FeatureSupport::native(),
237                    review: FeatureSupport::unsupported(),
238                    add_dirs: FeatureSupport::native(),
239                    max_turns: FeatureSupport::native(),
240                },
241            })
242        }
243        "ollama" => {
244            use crate::providers::ollama;
245            Ok(ProviderCapability {
246                provider: "ollama".to_string(),
247                default_model: ollama::DEFAULT_MODEL.to_string(),
248                available_models: models_to_vec(ollama::AVAILABLE_SIZES),
249                size_mappings: SizeMappings {
250                    small: "2b".to_string(),
251                    medium: "9b".to_string(),
252                    large: "35b".to_string(),
253                },
254                features: Features {
255                    interactive: FeatureSupport::native(),
256                    non_interactive: FeatureSupport::native(),
257                    resume: FeatureSupport::unsupported(),
258                    resume_with_prompt: FeatureSupport::unsupported(),
259                    session_logs: SessionLogSupport::unsupported(),
260                    json_output: FeatureSupport::wrapper(),
261                    stream_json: FeatureSupport::unsupported(),
262                    json_schema: FeatureSupport::wrapper(),
263                    input_format: FeatureSupport::unsupported(),
264                    streaming_input: FeatureSupport::unsupported(),
265                    worktree: FeatureSupport::wrapper(),
266                    sandbox: FeatureSupport::wrapper(),
267                    system_prompt: FeatureSupport::wrapper(),
268                    auto_approve: FeatureSupport::native(),
269                    review: FeatureSupport::unsupported(),
270                    add_dirs: FeatureSupport::unsupported(),
271                    max_turns: FeatureSupport::unsupported(),
272                },
273            })
274        }
275        _ => bail!(
276            "No capabilities defined for provider '{}'. Available: claude, codex, gemini, copilot, ollama",
277            provider
278        ),
279    }
280}
281
282/// Format a capability struct into the requested output format.
283pub fn format_capability(cap: &ProviderCapability, format: &str, pretty: bool) -> Result<String> {
284    match format {
285        "json" => {
286            if pretty {
287                Ok(serde_json::to_string_pretty(cap)?)
288            } else {
289                Ok(serde_json::to_string(cap)?)
290            }
291        }
292        "yaml" => Ok(serde_yaml::to_string(cap)?),
293        "toml" => Ok(toml::to_string_pretty(cap)?),
294        _ => bail!(
295            "Unsupported format '{}'. Available: json, yaml, toml",
296            format
297        ),
298    }
299}
300
301/// Convert a slice of string references into a Vec of owned Strings.
302pub fn models_to_vec(models: &[&str]) -> Vec<String> {
303    models.iter().map(|s| s.to_string()).collect()
304}
305
306#[cfg(test)]
307#[path = "capability_tests.rs"]
308mod tests;
zag_agent/capability.rs

zag_agent/
capability.rs
