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/// Streaming input support with mid-turn injection semantics.
22///
23/// Describes what happens when `StreamingSession::send_user_message` is called
24/// while the agent is already producing a response on the current turn.
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct StreamingInputSupport {
27    pub supported: bool,
28    pub native: bool,
29    /// Mid-turn semantics when `send_user_message` is called while the agent
30    /// is already producing a response. One of:
31    /// - `"queue"` — message is buffered and delivered at the next turn boundary
32    ///   (the current turn runs to completion before the new message is processed).
33    /// - `"interrupt"` — message cancels the current turn and starts a new one
34    ///   with the new input.
35    /// - `"between-turns-only"` — calling mid-turn is an error or no-op; callers
36    ///   must wait for the current turn to finish before sending.
37    ///
38    /// Absent when `supported == false`.
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub semantics: Option<String>,
41}
42
43/// Size alias mappings for a provider.
44#[derive(Debug, Clone, Serialize, Deserialize)]
45pub struct SizeMappings {
46    pub small: String,
47    pub medium: String,
48    pub large: String,
49}
50
51/// All feature flags for a provider.
52#[derive(Debug, Clone, Serialize, Deserialize)]
53pub struct Features {
54    pub interactive: FeatureSupport,
55    pub non_interactive: FeatureSupport,
56    pub resume: FeatureSupport,
57    pub resume_with_prompt: FeatureSupport,
58    pub session_logs: SessionLogSupport,
59    pub json_output: FeatureSupport,
60    pub stream_json: FeatureSupport,
61    pub json_schema: FeatureSupport,
62    pub input_format: FeatureSupport,
63    pub streaming_input: StreamingInputSupport,
64    pub worktree: FeatureSupport,
65    pub sandbox: FeatureSupport,
66    pub system_prompt: FeatureSupport,
67    pub auto_approve: FeatureSupport,
68    pub review: FeatureSupport,
69    pub add_dirs: FeatureSupport,
70    pub max_turns: FeatureSupport,
71}
72
73/// Full capability declaration for a provider.
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct ProviderCapability {
76    pub provider: String,
77    pub default_model: String,
78    pub available_models: Vec<String>,
79    pub size_mappings: SizeMappings,
80    pub features: Features,
81}
82
83impl FeatureSupport {
84    pub fn native() -> Self {
85        Self {
86            supported: true,
87            native: true,
88        }
89    }
90
91    pub fn wrapper() -> Self {
92        Self {
93            supported: true,
94            native: false,
95        }
96    }
97
98    pub fn unsupported() -> Self {
99        Self {
100            supported: false,
101            native: false,
102        }
103    }
104}
105
106impl SessionLogSupport {
107    pub fn full() -> Self {
108        Self {
109            supported: true,
110            native: true,
111            completeness: Some("full".to_string()),
112        }
113    }
114
115    pub fn partial() -> Self {
116        Self {
117            supported: true,
118            native: true,
119            completeness: Some("partial".to_string()),
120        }
121    }
122
123    pub fn unsupported() -> Self {
124        Self {
125            supported: false,
126            native: false,
127            completeness: None,
128        }
129    }
130}
131
132impl StreamingInputSupport {
133    /// Mid-turn messages are queued and delivered at the next turn boundary.
134    /// The currently running turn is not interrupted.
135    pub fn queue() -> Self {
136        Self {
137            supported: true,
138            native: true,
139            semantics: Some("queue".to_string()),
140        }
141    }
142
143    /// Mid-turn messages cancel the current turn and start a new one.
144    pub fn interrupt() -> Self {
145        Self {
146            supported: true,
147            native: true,
148            semantics: Some("interrupt".to_string()),
149        }
150    }
151
152    /// Messages may only be sent between turns; mid-turn sends are an error.
153    pub fn between_turns_only() -> Self {
154        Self {
155            supported: true,
156            native: true,
157            semantics: Some("between-turns-only".to_string()),
158        }
159    }
160
161    /// The provider does not support streaming input at all.
162    pub fn unsupported() -> Self {
163        Self {
164            supported: false,
165            native: false,
166            semantics: None,
167        }
168    }
169}
170
171/// Get capability declarations for a provider.
172pub fn get_capability(provider: &str) -> Result<ProviderCapability> {
173    use crate::agent::{Agent, ModelSize};
174
175    match provider {
176        "claude" => {
177            use crate::providers::claude::{self, Claude};
178            Ok(ProviderCapability {
179                provider: "claude".to_string(),
180                default_model: claude::DEFAULT_MODEL.to_string(),
181                available_models: models_to_vec(claude::AVAILABLE_MODELS),
182                size_mappings: SizeMappings {
183                    small: Claude::model_for_size(ModelSize::Small).to_string(),
184                    medium: Claude::model_for_size(ModelSize::Medium).to_string(),
185                    large: Claude::model_for_size(ModelSize::Large).to_string(),
186                },
187                features: Features {
188                    interactive: FeatureSupport::native(),
189                    non_interactive: FeatureSupport::native(),
190                    resume: FeatureSupport::native(),
191                    resume_with_prompt: FeatureSupport::native(),
192                    session_logs: SessionLogSupport::full(),
193                    json_output: FeatureSupport::native(),
194                    stream_json: FeatureSupport::native(),
195                    json_schema: FeatureSupport::native(),
196                    input_format: FeatureSupport::native(),
197                    streaming_input: StreamingInputSupport::queue(),
198                    worktree: FeatureSupport::wrapper(),
199                    sandbox: FeatureSupport::wrapper(),
200                    system_prompt: FeatureSupport::native(),
201                    auto_approve: FeatureSupport::native(),
202                    review: FeatureSupport::unsupported(),
203                    add_dirs: FeatureSupport::native(),
204                    max_turns: FeatureSupport::native(),
205                },
206            })
207        }
208        "codex" => {
209            use crate::providers::codex::{self, Codex};
210            Ok(ProviderCapability {
211                provider: "codex".to_string(),
212                default_model: codex::DEFAULT_MODEL.to_string(),
213                available_models: models_to_vec(codex::AVAILABLE_MODELS),
214                size_mappings: SizeMappings {
215                    small: Codex::model_for_size(ModelSize::Small).to_string(),
216                    medium: Codex::model_for_size(ModelSize::Medium).to_string(),
217                    large: Codex::model_for_size(ModelSize::Large).to_string(),
218                },
219                features: Features {
220                    interactive: FeatureSupport::native(),
221                    non_interactive: FeatureSupport::native(),
222                    resume: FeatureSupport::native(),
223                    resume_with_prompt: FeatureSupport::native(),
224                    session_logs: SessionLogSupport::partial(),
225                    json_output: FeatureSupport::native(),
226                    stream_json: FeatureSupport::unsupported(),
227                    json_schema: FeatureSupport::wrapper(),
228                    input_format: FeatureSupport::unsupported(),
229                    streaming_input: StreamingInputSupport::unsupported(),
230                    worktree: FeatureSupport::wrapper(),
231                    sandbox: FeatureSupport::wrapper(),
232                    system_prompt: FeatureSupport::wrapper(),
233                    auto_approve: FeatureSupport::native(),
234                    review: FeatureSupport::native(),
235                    add_dirs: FeatureSupport::native(),
236                    max_turns: FeatureSupport::native(),
237                },
238            })
239        }
240        "gemini" => {
241            use crate::providers::gemini::{self, Gemini};
242            Ok(ProviderCapability {
243                provider: "gemini".to_string(),
244                default_model: gemini::DEFAULT_MODEL.to_string(),
245                available_models: models_to_vec(gemini::AVAILABLE_MODELS),
246                size_mappings: SizeMappings {
247                    small: Gemini::model_for_size(ModelSize::Small).to_string(),
248                    medium: Gemini::model_for_size(ModelSize::Medium).to_string(),
249                    large: Gemini::model_for_size(ModelSize::Large).to_string(),
250                },
251                features: Features {
252                    interactive: FeatureSupport::native(),
253                    non_interactive: FeatureSupport::native(),
254                    resume: FeatureSupport::native(),
255                    resume_with_prompt: FeatureSupport::unsupported(),
256                    session_logs: SessionLogSupport::full(),
257                    json_output: FeatureSupport::wrapper(),
258                    stream_json: FeatureSupport::unsupported(),
259                    json_schema: FeatureSupport::wrapper(),
260                    input_format: FeatureSupport::unsupported(),
261                    streaming_input: StreamingInputSupport::unsupported(),
262                    worktree: FeatureSupport::wrapper(),
263                    sandbox: FeatureSupport::wrapper(),
264                    system_prompt: FeatureSupport::wrapper(),
265                    auto_approve: FeatureSupport::native(),
266                    review: FeatureSupport::unsupported(),
267                    add_dirs: FeatureSupport::native(),
268                    max_turns: FeatureSupport::native(),
269                },
270            })
271        }
272        "copilot" => {
273            use crate::providers::copilot::{self, Copilot};
274            Ok(ProviderCapability {
275                provider: "copilot".to_string(),
276                default_model: copilot::DEFAULT_MODEL.to_string(),
277                available_models: models_to_vec(copilot::AVAILABLE_MODELS),
278                size_mappings: SizeMappings {
279                    small: Copilot::model_for_size(ModelSize::Small).to_string(),
280                    medium: Copilot::model_for_size(ModelSize::Medium).to_string(),
281                    large: Copilot::model_for_size(ModelSize::Large).to_string(),
282                },
283                features: Features {
284                    interactive: FeatureSupport::native(),
285                    non_interactive: FeatureSupport::native(),
286                    resume: FeatureSupport::native(),
287                    resume_with_prompt: FeatureSupport::unsupported(),
288                    session_logs: SessionLogSupport::full(),
289                    json_output: FeatureSupport::unsupported(),
290                    stream_json: FeatureSupport::unsupported(),
291                    json_schema: FeatureSupport::unsupported(),
292                    input_format: FeatureSupport::unsupported(),
293                    streaming_input: StreamingInputSupport::unsupported(),
294                    worktree: FeatureSupport::wrapper(),
295                    sandbox: FeatureSupport::wrapper(),
296                    system_prompt: FeatureSupport::wrapper(),
297                    auto_approve: FeatureSupport::native(),
298                    review: FeatureSupport::unsupported(),
299                    add_dirs: FeatureSupport::native(),
300                    max_turns: FeatureSupport::native(),
301                },
302            })
303        }
304        "ollama" => {
305            use crate::providers::ollama;
306            Ok(ProviderCapability {
307                provider: "ollama".to_string(),
308                default_model: ollama::DEFAULT_MODEL.to_string(),
309                available_models: models_to_vec(ollama::AVAILABLE_SIZES),
310                size_mappings: SizeMappings {
311                    small: "2b".to_string(),
312                    medium: "9b".to_string(),
313                    large: "35b".to_string(),
314                },
315                features: Features {
316                    interactive: FeatureSupport::native(),
317                    non_interactive: FeatureSupport::native(),
318                    resume: FeatureSupport::unsupported(),
319                    resume_with_prompt: FeatureSupport::unsupported(),
320                    session_logs: SessionLogSupport::unsupported(),
321                    json_output: FeatureSupport::wrapper(),
322                    stream_json: FeatureSupport::unsupported(),
323                    json_schema: FeatureSupport::wrapper(),
324                    input_format: FeatureSupport::unsupported(),
325                    streaming_input: StreamingInputSupport::unsupported(),
326                    worktree: FeatureSupport::wrapper(),
327                    sandbox: FeatureSupport::wrapper(),
328                    system_prompt: FeatureSupport::wrapper(),
329                    auto_approve: FeatureSupport::native(),
330                    review: FeatureSupport::unsupported(),
331                    add_dirs: FeatureSupport::unsupported(),
332                    max_turns: FeatureSupport::unsupported(),
333                },
334            })
335        }
336        _ => bail!(
337            "No capabilities defined for provider '{provider}'. Available: claude, codex, gemini, copilot, ollama"
338        ),
339    }
340}
341
342/// Format a capability struct into the requested output format.
343pub fn format_capability(cap: &ProviderCapability, format: &str, pretty: bool) -> Result<String> {
344    match format {
345        "json" => {
346            if pretty {
347                Ok(serde_json::to_string_pretty(cap)?)
348            } else {
349                Ok(serde_json::to_string(cap)?)
350            }
351        }
352        "yaml" => Ok(serde_yaml::to_string(cap)?),
353        "toml" => Ok(toml::to_string_pretty(cap)?),
354        _ => bail!("Unsupported format '{format}'. Available: json, yaml, toml"),
355    }
356}
357
358/// Canonical list of provider names (excludes "auto" and "mock").
359pub const PROVIDERS: &[&str] = &["claude", "codex", "gemini", "copilot", "ollama"];
360
361/// List all available provider names.
362pub fn list_providers() -> Vec<String> {
363    PROVIDERS.iter().map(|s| s.to_string()).collect()
364}
365
366/// Get capabilities for all providers.
367pub fn get_all_capabilities() -> Vec<ProviderCapability> {
368    PROVIDERS
369        .iter()
370        .filter_map(|p| get_capability(p).ok())
371        .collect()
372}
373
374/// Result of resolving a model alias.
375#[derive(Debug, Clone, Serialize, Deserialize)]
376pub struct ResolvedModel {
377    pub input: String,
378    pub resolved: String,
379    pub is_alias: bool,
380    pub provider: String,
381}
382
383/// Resolve a model name or alias for a given provider.
384///
385/// Size aliases (`small`/`s`, `medium`/`m`/`default`, `large`/`l`/`max`) are
386/// resolved to the provider-specific model. Non-alias names pass through unchanged.
387pub fn resolve_model(provider: &str, model_input: &str) -> Result<ResolvedModel> {
388    use crate::agent::Agent;
389    use crate::providers::{
390        claude::Claude, codex::Codex, copilot::Copilot, gemini::Gemini, ollama::Ollama,
391    };
392
393    let resolved = match provider {
394        "claude" => Claude::resolve_model(model_input),
395        "codex" => Codex::resolve_model(model_input),
396        "gemini" => Gemini::resolve_model(model_input),
397        "copilot" => Copilot::resolve_model(model_input),
398        "ollama" => Ollama::resolve_model(model_input),
399        _ => bail!(
400            "Unknown provider '{}'. Available: {}",
401            provider,
402            PROVIDERS.join(", ")
403        ),
404    };
405
406    Ok(ResolvedModel {
407        input: model_input.to_string(),
408        is_alias: resolved != model_input,
409        resolved,
410        provider: provider.to_string(),
411    })
412}
413
414/// Format a resolved model into the requested output format.
415pub fn format_resolved_model(rm: &ResolvedModel, format: &str, pretty: bool) -> Result<String> {
416    match format {
417        "json" => {
418            if pretty {
419                Ok(serde_json::to_string_pretty(rm)?)
420            } else {
421                Ok(serde_json::to_string(rm)?)
422            }
423        }
424        "yaml" => Ok(serde_yaml::to_string(rm)?),
425        "toml" => Ok(toml::to_string_pretty(rm)?),
426        _ => bail!("Unsupported format '{format}'. Available: json, yaml, toml"),
427    }
428}
429
430/// Format a list of capabilities into the requested output format.
431pub fn format_capabilities(
432    caps: &[ProviderCapability],
433    format: &str,
434    pretty: bool,
435) -> Result<String> {
436    match format {
437        "json" => {
438            if pretty {
439                Ok(serde_json::to_string_pretty(caps)?)
440            } else {
441                Ok(serde_json::to_string(caps)?)
442            }
443        }
444        "yaml" => Ok(serde_yaml::to_string(caps)?),
445        "toml" => {
446            #[derive(Serialize)]
447            struct Wrapper<'a> {
448                providers: &'a [ProviderCapability],
449            }
450            Ok(toml::to_string_pretty(&Wrapper { providers: caps })?)
451        }
452        _ => bail!("Unsupported format '{format}'. Available: json, yaml, toml"),
453    }
454}
455
456/// Format a models listing into the requested output format.
457pub fn format_models(caps: &[ProviderCapability], format: &str, pretty: bool) -> Result<String> {
458    #[derive(Serialize)]
459    struct ModelEntry {
460        provider: String,
461        default_model: String,
462        models: Vec<String>,
463    }
464
465    let entries: Vec<ModelEntry> = caps
466        .iter()
467        .map(|c| ModelEntry {
468            provider: c.provider.clone(),
469            default_model: c.default_model.clone(),
470            models: c.available_models.clone(),
471        })
472        .collect();
473
474    match format {
475        "json" => {
476            if pretty {
477                Ok(serde_json::to_string_pretty(&entries)?)
478            } else {
479                Ok(serde_json::to_string(&entries)?)
480            }
481        }
482        "yaml" => Ok(serde_yaml::to_string(&entries)?),
483        "toml" => bail!("TOML does not support top-level arrays. Use json or yaml"),
484        _ => bail!("Unsupported format '{format}'. Available: json, yaml, toml"),
485    }
486}
487
488/// Convert a slice of string references into a Vec of owned Strings.
489pub fn models_to_vec(models: &[&str]) -> Vec<String> {
490    models.iter().map(|s| s.to_string()).collect()
491}
492
493#[cfg(test)]
494#[path = "capability_tests.rs"]
495mod tests;
zag_agent/capability.rs

zag_agent/
capability.rs
