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}
161
162impl BetaFeature {
163    pub fn header_value(&self) -> &'static str {
164        match self {
165            Self::InterleavedThinking => "interleaved-thinking-2025-05-14",
166            Self::ContextManagement => "context-management-2025-06-27",
167            Self::StructuredOutputs => "structured-outputs-2025-11-13",
168            Self::PromptCaching => "prompt-caching-2024-07-31",
169            Self::MaxTokens128k => "max-tokens-3-5-sonnet-2024-07-15",
170            Self::CodeExecution => "code-execution-2025-01-24",
171            Self::Mcp => "mcp-2025-04-08",
172            Self::WebSearch => "web-search-2025-03-05",
173            Self::WebFetch => "web-fetch-2025-09-10",
174            Self::OAuth => "oauth-2025-04-20",
175            Self::FilesApi => "files-api-2025-04-14",
176            Self::Effort => "effort-2025-11-24",
177            Self::Context1M => "context-1m-2025-08-07",
178        }
179    }
180
181    fn from_header(value: &str) -> Option<Self> {
182        match value {
183            "interleaved-thinking-2025-05-14" => Some(Self::InterleavedThinking),
184            "context-management-2025-06-27" => Some(Self::ContextManagement),
185            "structured-outputs-2025-11-13" => Some(Self::StructuredOutputs),
186            "prompt-caching-2024-07-31" => Some(Self::PromptCaching),
187            "max-tokens-3-5-sonnet-2024-07-15" => Some(Self::MaxTokens128k),
188            "code-execution-2025-01-24" => Some(Self::CodeExecution),
189            "mcp-2025-04-08" => Some(Self::Mcp),
190            "web-search-2025-03-05" => Some(Self::WebSearch),
191            "web-fetch-2025-09-10" => Some(Self::WebFetch),
192            "oauth-2025-04-20" => Some(Self::OAuth),
193            "files-api-2025-04-14" => Some(Self::FilesApi),
194            "effort-2025-11-24" => Some(Self::Effort),
195            "context-1m-2025-08-07" => Some(Self::Context1M),
196            _ => None,
197        }
198    }
199
200    pub fn all() -> &'static [BetaFeature] {
201        &[
202            Self::InterleavedThinking,
203            Self::ContextManagement,
204            Self::StructuredOutputs,
205            Self::PromptCaching,
206            Self::MaxTokens128k,
207            Self::CodeExecution,
208            Self::Mcp,
209            Self::WebSearch,
210            Self::WebFetch,
211            Self::OAuth,
212            Self::FilesApi,
213            Self::Effort,
214            Self::Context1M,
215        ]
216    }
217}
218
219#[derive(Clone, Debug, Default)]
220pub struct BetaConfig {
221    features: HashSet<BetaFeature>,
222    custom: Vec<String>,
223}
224
225impl BetaConfig {
226    pub fn new() -> Self {
227        Self::default()
228    }
229
230    pub fn all() -> Self {
231        Self {
232            features: BetaFeature::all().iter().copied().collect(),
233            custom: Vec::new(),
234        }
235    }
236
237    pub fn with(mut self, feature: BetaFeature) -> Self {
238        self.features.insert(feature);
239        self
240    }
241
242    pub fn with_custom(mut self, flag: impl Into<String>) -> Self {
243        self.custom.push(flag.into());
244        self
245    }
246
247    pub fn add(&mut self, feature: BetaFeature) {
248        self.features.insert(feature);
249    }
250
251    pub fn add_custom(&mut self, flag: impl Into<String>) {
252        self.custom.push(flag.into());
253    }
254
255    pub fn from_env() -> Self {
256        let mut config = Self::new();
257
258        if let Ok(flags) = env::var("ANTHROPIC_BETA_FLAGS") {
259            for flag in flags.split(',').map(str::trim).filter(|s| !s.is_empty()) {
260                if let Some(feature) = BetaFeature::from_header(flag) {
261                    config.features.insert(feature);
262                } else {
263                    config.custom.push(flag.to_string());
264                }
265            }
266        }
267
268        config
269    }
270
271    pub fn header_value(&self) -> Option<String> {
272        let mut flags: Vec<&str> = self.features.iter().map(|f| f.header_value()).collect();
273        flags.sort();
274
275        for custom in &self.custom {
276            if !flags.contains(&custom.as_str()) {
277                flags.push(custom);
278            }
279        }
280
281        if flags.is_empty() {
282            None
283        } else {
284            Some(flags.join(","))
285        }
286    }
287
288    pub fn is_empty(&self) -> bool {
289        self.features.is_empty() && self.custom.is_empty()
290    }
291
292    pub fn has(&self, feature: BetaFeature) -> bool {
293        self.features.contains(&feature)
294    }
295}
296
297#[derive(Clone, Debug)]
298pub struct ProviderConfig {
299    pub models: ModelConfig,
300    pub max_tokens: u32,
301    pub thinking_budget: Option<u32>,
302    pub enable_caching: bool,
303    pub api_version: String,
304    pub beta: BetaConfig,
305    pub extra_headers: HashMap<String, String>,
306}
307
308impl ProviderConfig {
309    pub fn new(models: ModelConfig) -> Self {
310        Self {
311            models,
312            max_tokens: DEFAULT_MAX_TOKENS,
313            thinking_budget: None,
314            enable_caching: !env::var("DISABLE_PROMPT_CACHING")
315                .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
316                .unwrap_or(false),
317            api_version: "2023-06-01".into(),
318            beta: BetaConfig::from_env(),
319            extra_headers: HashMap::new(),
320        }
321    }
322
323    pub fn with_max_tokens(mut self, tokens: u32) -> Self {
324        self.max_tokens = tokens;
325        if tokens > DEFAULT_MAX_TOKENS {
326            self.beta.add(BetaFeature::MaxTokens128k);
327        }
328        self
329    }
330
331    pub fn with_thinking(mut self, budget: u32) -> Self {
332        self.thinking_budget = Some(budget.max(MIN_THINKING_BUDGET));
333        self.beta.add(BetaFeature::InterleavedThinking);
334        self
335    }
336
337    pub fn disable_caching(mut self) -> Self {
338        self.enable_caching = false;
339        self
340    }
341
342    pub fn with_api_version(mut self, version: impl Into<String>) -> Self {
343        self.api_version = version.into();
344        self
345    }
346
347    pub fn with_beta(mut self, feature: BetaFeature) -> Self {
348        self.beta.add(feature);
349        self
350    }
351
352    pub fn with_beta_config(mut self, config: BetaConfig) -> Self {
353        self.beta = config;
354        self
355    }
356
357    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
358        self.extra_headers.insert(key.into(), value.into());
359        self
360    }
361
362    pub fn requires_128k_beta(&self) -> bool {
363        self.max_tokens > DEFAULT_MAX_TOKENS
364    }
365}
366
367impl Default for ProviderConfig {
368    fn default() -> Self {
369        Self::new(ModelConfig::default())
370    }
371}
372
373#[cfg(test)]
374mod tests {
375    use super::*;
376
377    #[test]
378    fn test_model_config_get() {
379        let config = ModelConfig::anthropic();
380        assert!(config.get(ModelType::Primary).contains("sonnet"));
381        assert!(config.get(ModelType::Small).contains("haiku"));
382        assert!(config.get(ModelType::Reasoning).contains("opus"));
383    }
384
385    #[test]
386    fn test_provider_config_default_max_tokens() {
387        let config = ProviderConfig::default();
388        assert_eq!(config.max_tokens, DEFAULT_MAX_TOKENS);
389        assert!(!config.requires_128k_beta());
390    }
391
392    #[test]
393    fn test_provider_config_builder() {
394        let config = ProviderConfig::new(ModelConfig::anthropic())
395            .with_max_tokens(16384)
396            .with_thinking(10000)
397            .disable_caching();
398
399        assert_eq!(config.max_tokens, 16384);
400        assert_eq!(config.thinking_budget, Some(10000));
401        assert!(!config.enable_caching);
402        assert!(config.requires_128k_beta());
403        assert!(config.beta.has(BetaFeature::MaxTokens128k));
404        assert!(config.beta.has(BetaFeature::InterleavedThinking));
405    }
406
407    #[test]
408    fn test_provider_config_auto_128k_beta() {
409        let config = ProviderConfig::default().with_max_tokens(DEFAULT_MAX_TOKENS);
410        assert!(!config.beta.has(BetaFeature::MaxTokens128k));
411
412        let config = ProviderConfig::default().with_max_tokens(DEFAULT_MAX_TOKENS + 1);
413        assert!(config.beta.has(BetaFeature::MaxTokens128k));
414    }
415
416    #[test]
417    fn test_provider_config_thinking_auto_beta() {
418        let config = ProviderConfig::default().with_thinking(5000);
419        assert!(config.beta.has(BetaFeature::InterleavedThinking));
420        assert_eq!(config.thinking_budget, Some(5000));
421    }
422
423    #[test]
424    fn test_provider_config_thinking_min_budget() {
425        let config = ProviderConfig::default().with_thinking(500);
426        assert_eq!(config.thinking_budget, Some(MIN_THINKING_BUDGET));
427    }
428
429    #[test]
430    fn test_beta_feature_header() {
431        assert_eq!(
432            BetaFeature::InterleavedThinking.header_value(),
433            "interleaved-thinking-2025-05-14"
434        );
435        assert_eq!(
436            BetaFeature::MaxTokens128k.header_value(),
437            "max-tokens-3-5-sonnet-2024-07-15"
438        );
439    }
440
441    #[test]
442    fn test_beta_config_with_features() {
443        let config = BetaConfig::new()
444            .with(BetaFeature::InterleavedThinking)
445            .with(BetaFeature::ContextManagement);
446
447        assert!(config.has(BetaFeature::InterleavedThinking));
448        assert!(config.has(BetaFeature::ContextManagement));
449        assert!(!config.has(BetaFeature::MaxTokens128k));
450
451        let header = config.header_value().unwrap();
452        assert!(header.contains("interleaved-thinking"));
453        assert!(header.contains("context-management"));
454    }
455
456    #[test]
457    fn test_beta_config_custom() {
458        let config = BetaConfig::new()
459            .with(BetaFeature::InterleavedThinking)
460            .with_custom("new-feature-2026-01-01");
461
462        let header = config.header_value().unwrap();
463        assert!(header.contains("interleaved-thinking"));
464        assert!(header.contains("new-feature-2026-01-01"));
465    }
466
467    #[test]
468    fn test_beta_config_all() {
469        let config = BetaConfig::all();
470        assert!(config.has(BetaFeature::InterleavedThinking));
471        assert!(config.has(BetaFeature::ContextManagement));
472        assert!(config.has(BetaFeature::MaxTokens128k));
473    }
474
475    #[test]
476    fn test_provider_config_beta() {
477        let config = ProviderConfig::default()
478            .with_beta(BetaFeature::InterleavedThinking)
479            .with_beta_config(
480                BetaConfig::new()
481                    .with(BetaFeature::InterleavedThinking)
482                    .with_custom("experimental-feature"),
483            );
484
485        assert!(config.beta.has(BetaFeature::InterleavedThinking));
486        let header = config.beta.header_value().unwrap();
487        assert!(header.contains("experimental-feature"));
488    }
489
490    #[test]
491    fn test_beta_config_empty() {
492        let config = BetaConfig::new();
493        assert!(config.is_empty());
494        assert!(config.header_value().is_none());
495    }
496
497    #[test]
498    fn test_provider_config_extra_headers() {
499        let config = ProviderConfig::default()
500            .with_header("x-custom", "value")
501            .with_header("x-another", "test");
502
503        assert_eq!(config.extra_headers.get("x-custom"), Some(&"value".into()));
504        assert_eq!(config.extra_headers.get("x-another"), Some(&"test".into()));
505    }
506}
claude_agent/client/adapter/config.rs

claude_agent/client/adapter/
config.rs
