claude_agent/client/adapter/
config.rs

1//! Provider and model configuration.
2
3use std::collections::{HashMap, HashSet};
4use std::env;
5
6// Anthropic API models
7pub const DEFAULT_MODEL: &str = "claude-sonnet-4-5-20250929";
8pub const DEFAULT_SMALL_MODEL: &str = "claude-haiku-4-5-20251001";
9pub const DEFAULT_REASONING_MODEL: &str = "claude-opus-4-5-20251101";
10pub const FRONTIER_MODEL: &str = DEFAULT_REASONING_MODEL;
11
12// AWS Bedrock models (using global endpoint prefix for maximum availability)
13#[cfg(feature = "aws")]
14pub const BEDROCK_MODEL: &str = "global.anthropic.claude-sonnet-4-5-20250929-v1:0";
15#[cfg(feature = "aws")]
16pub const BEDROCK_SMALL_MODEL: &str = "global.anthropic.claude-haiku-4-5-20251001-v1:0";
17#[cfg(feature = "aws")]
18pub const BEDROCK_REASONING_MODEL: &str = "global.anthropic.claude-opus-4-5-20251101-v1:0";
19
20// GCP Vertex AI models
21#[cfg(feature = "gcp")]
22pub const VERTEX_MODEL: &str = "claude-sonnet-4-5@20250929";
23#[cfg(feature = "gcp")]
24pub const VERTEX_SMALL_MODEL: &str = "claude-haiku-4-5@20251001";
25#[cfg(feature = "gcp")]
26pub const VERTEX_REASONING_MODEL: &str = "claude-opus-4-5@20251101";
27
28// Azure Foundry models
29#[cfg(feature = "azure")]
30pub const FOUNDRY_MODEL: &str = "claude-sonnet-4-5";
31#[cfg(feature = "azure")]
32pub const FOUNDRY_SMALL_MODEL: &str = "claude-haiku-4-5";
33#[cfg(feature = "azure")]
34pub const FOUNDRY_REASONING_MODEL: &str = "claude-opus-4-5";
35
36#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
37#[serde(rename_all = "lowercase")]
38pub enum ModelType {
39    #[default]
40    Primary,
41    Small,
42    Reasoning,
43}
44
45#[derive(Clone, Debug)]
46pub struct ModelConfig {
47    pub primary: String,
48    pub small: String,
49    pub reasoning: Option<String>,
50}
51
52impl ModelConfig {
53    pub fn new(primary: impl Into<String>, small: impl Into<String>) -> Self {
54        Self {
55            primary: primary.into(),
56            small: small.into(),
57            reasoning: None,
58        }
59    }
60
61    pub fn anthropic() -> Self {
62        Self::from_env_with_defaults(DEFAULT_MODEL, DEFAULT_SMALL_MODEL, DEFAULT_REASONING_MODEL)
63    }
64
65    fn from_env_with_defaults(
66        default_primary: &str,
67        default_small: &str,
68        default_reasoning: &str,
69    ) -> Self {
70        Self {
71            primary: env::var("ANTHROPIC_MODEL").unwrap_or_else(|_| default_primary.into()),
72            small: env::var("ANTHROPIC_SMALL_FAST_MODEL").unwrap_or_else(|_| default_small.into()),
73            reasoning: Some(
74                env::var("ANTHROPIC_REASONING_MODEL").unwrap_or_else(|_| default_reasoning.into()),
75            ),
76        }
77    }
78
79    #[cfg(feature = "aws")]
80    pub fn bedrock() -> Self {
81        Self::from_env_with_defaults(BEDROCK_MODEL, BEDROCK_SMALL_MODEL, BEDROCK_REASONING_MODEL)
82    }
83
84    #[cfg(feature = "gcp")]
85    pub fn vertex() -> Self {
86        Self::from_env_with_defaults(VERTEX_MODEL, VERTEX_SMALL_MODEL, VERTEX_REASONING_MODEL)
87    }
88
89    #[cfg(feature = "azure")]
90    pub fn foundry() -> Self {
91        Self::from_env_with_defaults(FOUNDRY_MODEL, FOUNDRY_SMALL_MODEL, FOUNDRY_REASONING_MODEL)
92    }
93
94    pub fn with_primary(mut self, model: impl Into<String>) -> Self {
95        self.primary = model.into();
96        self
97    }
98
99    pub fn with_small(mut self, model: impl Into<String>) -> Self {
100        self.small = model.into();
101        self
102    }
103
104    pub fn with_reasoning(mut self, model: impl Into<String>) -> Self {
105        self.reasoning = Some(model.into());
106        self
107    }
108
109    pub fn get(&self, model_type: ModelType) -> &str {
110        match model_type {
111            ModelType::Primary => &self.primary,
112            ModelType::Small => &self.small,
113            ModelType::Reasoning => self.reasoning.as_deref().unwrap_or(&self.primary),
114        }
115    }
116
117    pub fn resolve_alias<'a>(&'a self, alias: &'a str) -> &'a str {
118        match alias {
119            "sonnet" => &self.primary,
120            "haiku" => &self.small,
121            "opus" => self.reasoning.as_deref().unwrap_or(&self.primary),
122            other => other,
123        }
124    }
125
126    pub fn model_type_from_alias(alias: &str) -> Option<ModelType> {
127        match alias {
128            "sonnet" => Some(ModelType::Primary),
129            "haiku" => Some(ModelType::Small),
130            "opus" => Some(ModelType::Reasoning),
131            _ => None,
132        }
133    }
134}
135
136impl Default for ModelConfig {
137    fn default() -> Self {
138        Self::anthropic()
139    }
140}
141
142#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
143pub enum BetaFeature {
144    InterleavedThinking,
145    ContextManagement,
146    StructuredOutputs,
147    PromptCaching,
148    MaxTokens128k,
149    CodeExecution,
150    Mcp,
151    WebSearch,
152    WebFetch,
153    OAuth,
154    FilesApi,
155    Effort,
156    /// 1M token context window (for Sonnet 4 and 4.5 on Bedrock/Vertex).
157    Context1M,
158}
159
160impl BetaFeature {
161    pub fn header_value(&self) -> &'static str {
162        match self {
163            Self::InterleavedThinking => "interleaved-thinking-2025-05-14",
164            Self::ContextManagement => "context-management-2025-06-27",
165            Self::StructuredOutputs => "structured-outputs-2025-11-13",
166            Self::PromptCaching => "prompt-caching-2024-07-31",
167            Self::MaxTokens128k => "max-tokens-3-5-sonnet-2024-07-15",
168            Self::CodeExecution => "code-execution-2025-01-24",
169            Self::Mcp => "mcp-2025-04-08",
170            Self::WebSearch => "web-search-2025-03-05",
171            Self::WebFetch => "web-fetch-2025-09-10",
172            Self::OAuth => "oauth-2025-04-20",
173            Self::FilesApi => "files-api-2025-04-14",
174            Self::Effort => "effort-2025-11-24",
175            Self::Context1M => "context-1m-2025-08-07",
176        }
177    }
178
179    fn from_header(value: &str) -> Option<Self> {
180        match value {
181            "interleaved-thinking-2025-05-14" => Some(Self::InterleavedThinking),
182            "context-management-2025-06-27" => Some(Self::ContextManagement),
183            "structured-outputs-2025-11-13" => Some(Self::StructuredOutputs),
184            "prompt-caching-2024-07-31" => Some(Self::PromptCaching),
185            "max-tokens-3-5-sonnet-2024-07-15" => Some(Self::MaxTokens128k),
186            "code-execution-2025-01-24" => Some(Self::CodeExecution),
187            "mcp-2025-04-08" => Some(Self::Mcp),
188            "web-search-2025-03-05" => Some(Self::WebSearch),
189            "web-fetch-2025-09-10" => Some(Self::WebFetch),
190            "oauth-2025-04-20" => Some(Self::OAuth),
191            "files-api-2025-04-14" => Some(Self::FilesApi),
192            "effort-2025-11-24" => Some(Self::Effort),
193            "context-1m-2025-08-07" => Some(Self::Context1M),
194            _ => None,
195        }
196    }
197
198    pub fn all() -> &'static [BetaFeature] {
199        &[
200            Self::InterleavedThinking,
201            Self::ContextManagement,
202            Self::StructuredOutputs,
203            Self::PromptCaching,
204            Self::MaxTokens128k,
205            Self::CodeExecution,
206            Self::Mcp,
207            Self::WebSearch,
208            Self::WebFetch,
209            Self::OAuth,
210            Self::FilesApi,
211            Self::Effort,
212            Self::Context1M,
213        ]
214    }
215}
216
217#[derive(Clone, Debug, Default)]
218pub struct BetaConfig {
219    features: HashSet<BetaFeature>,
220    custom: Vec<String>,
221}
222
223impl BetaConfig {
224    pub fn new() -> Self {
225        Self::default()
226    }
227
228    pub fn all() -> Self {
229        Self {
230            features: BetaFeature::all().iter().copied().collect(),
231            custom: Vec::new(),
232        }
233    }
234
235    pub fn with(mut self, feature: BetaFeature) -> Self {
236        self.features.insert(feature);
237        self
238    }
239
240    pub fn with_custom(mut self, flag: impl Into<String>) -> Self {
241        self.custom.push(flag.into());
242        self
243    }
244
245    pub fn add(&mut self, feature: BetaFeature) {
246        self.features.insert(feature);
247    }
248
249    pub fn add_custom(&mut self, flag: impl Into<String>) {
250        self.custom.push(flag.into());
251    }
252
253    pub fn from_env() -> Self {
254        let mut config = Self::new();
255
256        if let Ok(flags) = env::var("ANTHROPIC_BETA_FLAGS") {
257            for flag in flags.split(',').map(str::trim).filter(|s| !s.is_empty()) {
258                if let Some(feature) = BetaFeature::from_header(flag) {
259                    config.features.insert(feature);
260                } else {
261                    config.custom.push(flag.to_string());
262                }
263            }
264        }
265
266        config
267    }
268
269    pub fn header_value(&self) -> Option<String> {
270        let mut flags: Vec<&str> = self.features.iter().map(|f| f.header_value()).collect();
271        flags.sort();
272
273        for custom in &self.custom {
274            if !flags.contains(&custom.as_str()) {
275                flags.push(custom);
276            }
277        }
278
279        if flags.is_empty() {
280            None
281        } else {
282            Some(flags.join(","))
283        }
284    }
285
286    pub fn is_empty(&self) -> bool {
287        self.features.is_empty() && self.custom.is_empty()
288    }
289
290    pub fn has(&self, feature: BetaFeature) -> bool {
291        self.features.contains(&feature)
292    }
293}
294
295#[derive(Clone, Debug)]
296pub struct ProviderConfig {
297    pub models: ModelConfig,
298    pub max_tokens: u32,
299    pub thinking_budget: Option<u32>,
300    pub enable_caching: bool,
301    pub api_version: String,
302    pub beta: BetaConfig,
303    pub extra_headers: HashMap<String, String>,
304}
305
306impl ProviderConfig {
307    pub fn new(models: ModelConfig) -> Self {
308        Self {
309            models,
310            max_tokens: 8192,
311            thinking_budget: None,
312            enable_caching: !env::var("DISABLE_PROMPT_CACHING")
313                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
314                .unwrap_or(false),
315            api_version: "2023-06-01".into(),
316            beta: BetaConfig::from_env(),
317            extra_headers: HashMap::new(),
318        }
319    }
320
321    pub fn with_max_tokens(mut self, tokens: u32) -> Self {
322        self.max_tokens = tokens;
323        self
324    }
325
326    pub fn with_thinking(mut self, budget: u32) -> Self {
327        self.thinking_budget = Some(budget);
328        self
329    }
330
331    pub fn disable_caching(mut self) -> Self {
332        self.enable_caching = false;
333        self
334    }
335
336    pub fn with_api_version(mut self, version: impl Into<String>) -> Self {
337        self.api_version = version.into();
338        self
339    }
340
341    pub fn with_beta(mut self, feature: BetaFeature) -> Self {
342        self.beta.add(feature);
343        self
344    }
345
346    pub fn with_beta_config(mut self, config: BetaConfig) -> Self {
347        self.beta = config;
348        self
349    }
350
351    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
352        self.extra_headers.insert(key.into(), value.into());
353        self
354    }
355}
356
357impl Default for ProviderConfig {
358    fn default() -> Self {
359        Self::new(ModelConfig::default())
360    }
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366
367    #[test]
368    fn test_model_config_get() {
369        let config = ModelConfig::anthropic();
370        assert!(config.get(ModelType::Primary).contains("sonnet"));
371        assert!(config.get(ModelType::Small).contains("haiku"));
372        assert!(config.get(ModelType::Reasoning).contains("opus"));
373    }
374
375    #[test]
376    fn test_provider_config_builder() {
377        let config = ProviderConfig::new(ModelConfig::anthropic())
378            .with_max_tokens(16384)
379            .with_thinking(10000)
380            .disable_caching();
381
382        assert_eq!(config.max_tokens, 16384);
383        assert_eq!(config.thinking_budget, Some(10000));
384        assert!(!config.enable_caching);
385    }
386
387    #[test]
388    fn test_beta_feature_header() {
389        assert_eq!(
390            BetaFeature::InterleavedThinking.header_value(),
391            "interleaved-thinking-2025-05-14"
392        );
393    }
394
395    #[test]
396    fn test_beta_config_with_features() {
397        let config = BetaConfig::new()
398            .with(BetaFeature::InterleavedThinking)
399            .with(BetaFeature::ContextManagement);
400
401        assert!(config.has(BetaFeature::InterleavedThinking));
402        assert!(config.has(BetaFeature::ContextManagement));
403        assert!(!config.has(BetaFeature::MaxTokens128k));
404
405        let header = config.header_value().unwrap();
406        assert!(header.contains("interleaved-thinking"));
407        assert!(header.contains("context-management"));
408    }
409
410    #[test]
411    fn test_beta_config_custom() {
412        let config = BetaConfig::new()
413            .with(BetaFeature::InterleavedThinking)
414            .with_custom("new-feature-2026-01-01");
415
416        let header = config.header_value().unwrap();
417        assert!(header.contains("interleaved-thinking"));
418        assert!(header.contains("new-feature-2026-01-01"));
419    }
420
421    #[test]
422    fn test_beta_config_all() {
423        let config = BetaConfig::all();
424        assert!(config.has(BetaFeature::InterleavedThinking));
425        assert!(config.has(BetaFeature::ContextManagement));
426        assert!(config.has(BetaFeature::MaxTokens128k));
427    }
428
429    #[test]
430    fn test_provider_config_beta() {
431        let config = ProviderConfig::default()
432            .with_beta(BetaFeature::InterleavedThinking)
433            .with_beta_config(
434                BetaConfig::new()
435                    .with(BetaFeature::InterleavedThinking)
436                    .with_custom("experimental-feature"),
437            );
438
439        assert!(config.beta.has(BetaFeature::InterleavedThinking));
440        let header = config.beta.header_value().unwrap();
441        assert!(header.contains("experimental-feature"));
442    }
443
444    #[test]
445    fn test_beta_config_empty() {
446        let config = BetaConfig::new();
447        assert!(config.is_empty());
448        assert!(config.header_value().is_none());
449    }
450
451    #[test]
452    fn test_provider_config_extra_headers() {
453        let config = ProviderConfig::default()
454            .with_header("x-custom", "value")
455            .with_header("x-another", "test");
456
457        assert_eq!(config.extra_headers.get("x-custom"), Some(&"value".into()));
458        assert_eq!(config.extra_headers.get("x-another"), Some(&"test".into()));
459    }
460}
claude_agent/client/adapter/config.rs

claude_agent/client/adapter/
config.rs
