claude_agent/client/adapter/
config.rs

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

claude_agent/client/adapter/
config.rs
