Skip to main content

claude_agent/client/adapter/
foundry.rs

1//! Azure AI Foundry adapter with API key and Entra ID authentication.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use async_trait::async_trait;
7use azure_core::credentials::TokenCredential;
8use azure_identity::DeveloperToolsCredential;
9
10use super::base::RequestExecutor;
11use super::config::ProviderConfig;
12use super::request::build_messages_body;
13use super::token_cache::{CachedToken, TokenCache, new_token_cache};
14use super::traits::ProviderAdapter;
15use crate::client::messages::CreateMessageRequest;
16use crate::config::FoundryConfig;
17use crate::types::ApiResponse;
18use crate::{Error, Result};
19
20const ANTHROPIC_VERSION: &str = "2023-06-01";
21const COGNITIVE_SERVICES_SCOPE: &str = "https://cognitiveservices.azure.com/.default";
22
23#[derive(Debug)]
24pub struct FoundryAdapter {
25    config: ProviderConfig,
26    resource_name: Option<String>,
27    base_url: Option<String>,
28    credential: Arc<DeveloperToolsCredential>,
29    api_key: Option<String>,
30    token_cache: TokenCache,
31}
32
33impl FoundryAdapter {
34    pub async fn from_env(config: ProviderConfig) -> Result<Self> {
35        let foundry_config = FoundryConfig::from_env();
36        Self::from_config(config, foundry_config).await
37    }
38
39    pub async fn from_config(config: ProviderConfig, foundry: FoundryConfig) -> Result<Self> {
40        if foundry.resource.is_none() && foundry.base_url.is_none() {
41            return Err(Error::auth(
42                "Either ANTHROPIC_FOUNDRY_RESOURCE or ANTHROPIC_FOUNDRY_BASE_URL must be set",
43            ));
44        }
45
46        let credential = DeveloperToolsCredential::new(None)
47            .map_err(|e| Error::auth(format!("Failed to create Azure credential: {}", e)))?;
48
49        Ok(Self {
50            config,
51            resource_name: foundry.resource,
52            base_url: foundry.base_url,
53            credential,
54            api_key: foundry.api_key,
55            token_cache: new_token_cache(),
56        })
57    }
58
59    pub fn resource(mut self, resource: impl Into<String>) -> Self {
60        self.resource_name = Some(resource.into());
61        self
62    }
63
64    pub fn base_url(mut self, base_url: impl Into<String>) -> Self {
65        self.base_url = Some(base_url.into());
66        self
67    }
68
69    pub fn api_key(mut self, key: impl Into<String>) -> Self {
70        self.api_key = Some(key.into());
71        self
72    }
73
74    fn build_messages_url(&self) -> String {
75        if let Some(ref base_url) = self.base_url {
76            let base = base_url.trim_end_matches('/');
77            format!("{}/v1/messages", base)
78        } else if let Some(ref resource) = self.resource_name {
79            format!(
80                "https://{}.services.ai.azure.com/anthropic/v1/messages",
81                resource
82            )
83        } else {
84            unreachable!(
85                "FoundryAdapter requires base_url or resource_name, enforced by from_config"
86            )
87        }
88    }
89
90    fn build_request_body(&self, request: &CreateMessageRequest) -> serde_json::Value {
91        build_messages_body(request, None, self.config.thinking_budget)
92    }
93
94    async fn get_auth_header(&self) -> Result<(String, String)> {
95        if let Some(ref api_key) = self.api_key {
96            return Ok(("api-key".into(), api_key.clone()));
97        }
98        let token = self.get_token().await?;
99        Ok(("Authorization".into(), format!("Bearer {}", token)))
100    }
101
102    async fn get_token(&self) -> Result<String> {
103        {
104            let cache = self.token_cache.read().await;
105            if let Some(ref cached) = *cache
106                && !cached.is_expired()
107            {
108                return Ok(cached.token().to_string());
109            }
110        }
111
112        let token_response: azure_core::credentials::AccessToken = self
113            .credential
114            .get_token(&[COGNITIVE_SERVICES_SCOPE], None)
115            .await
116            .map_err(|e| Error::auth(format!("Failed to get Azure token: {}", e)))?;
117
118        let token_str = token_response.token.secret().to_string();
119        let cached = CachedToken::new(token_str.clone(), Duration::from_secs(3600));
120        *self.token_cache.write().await = Some(cached);
121
122        Ok(token_str)
123    }
124
125    async fn execute_request(
126        &self,
127        http: &reqwest::Client,
128        url: &str,
129        body: &serde_json::Value,
130    ) -> Result<reqwest::Response> {
131        let (header_name, header_value) = self.get_auth_header().await?;
132        let headers = vec![
133            (header_name, header_value),
134            ("anthropic-version".into(), ANTHROPIC_VERSION.into()),
135        ];
136        RequestExecutor::post(http, url, body, headers).await
137    }
138}
139
140#[async_trait]
141impl ProviderAdapter for FoundryAdapter {
142    fn config(&self) -> &ProviderConfig {
143        &self.config
144    }
145
146    fn name(&self) -> &'static str {
147        "foundry"
148    }
149
150    async fn build_url(&self, _model: &str, _stream: bool) -> String {
151        self.build_messages_url()
152    }
153
154    async fn transform_request(&self, request: CreateMessageRequest) -> Result<serde_json::Value> {
155        Ok(self.build_request_body(&request))
156    }
157
158    async fn send(
159        &self,
160        http: &reqwest::Client,
161        request: CreateMessageRequest,
162    ) -> Result<ApiResponse> {
163        let url = self.build_messages_url();
164        let body = self.build_request_body(&request);
165        let response = self.execute_request(http, &url, &body).await?;
166        let json: serde_json::Value = response.json().await?;
167        self.transform_response(json)
168    }
169
170    async fn send_stream(
171        &self,
172        http: &reqwest::Client,
173        mut request: CreateMessageRequest,
174    ) -> Result<reqwest::Response> {
175        request.stream = Some(true);
176        let url = self.build_messages_url();
177        let body = self.build_request_body(&request);
178        self.execute_request(http, &url, &body).await
179    }
180
181    async fn refresh_credentials(&self) -> Result<()> {
182        *self.token_cache.write().await = None;
183        if self.api_key.is_none() {
184            self.get_token().await?;
185        }
186        Ok(())
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use crate::client::adapter::ModelConfig;
193
194    #[test]
195    fn test_build_url_with_resource() {
196        let url = format!(
197            "https://{}.services.ai.azure.com/anthropic/v1/messages",
198            "my-resource"
199        );
200        assert!(url.contains("services.ai.azure.com"));
201        assert!(url.contains("/anthropic/v1/messages"));
202    }
203
204    #[test]
205    fn test_build_url_with_base_url() {
206        let base_url = "https://custom-endpoint.azure.com/anthropic";
207        let url = format!("{}/v1/messages", base_url.trim_end_matches('/'));
208        assert!(url.contains("custom-endpoint.azure.com"));
209        assert!(url.contains("/v1/messages"));
210    }
211
212    #[test]
213    fn test_model_config() {
214        let config = ModelConfig::foundry();
215        assert!(config.primary.contains("sonnet"));
216    }
217
218    #[test]
219    fn test_anthropic_version() {
220        assert_eq!(super::ANTHROPIC_VERSION, "2023-06-01");
221    }
222}