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    /// Set the Azure resource name.
60    pub fn with_resource(mut self, resource: impl Into<String>) -> Self {
61        self.resource_name = Some(resource.into());
62        self
63    }
64
65    /// Set the base URL (alternative to resource name).
66    pub fn with_base_url(mut self, base_url: impl Into<String>) -> Self {
67        self.base_url = Some(base_url.into());
68        self
69    }
70
71    /// Set the API key for authentication.
72    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
73        self.api_key = Some(key.into());
74        self
75    }
76
77    fn build_messages_url(&self) -> String {
78        if let Some(ref base_url) = self.base_url {
79            let base = base_url.trim_end_matches('/');
80            format!("{}/v1/messages", base)
81        } else if let Some(ref resource) = self.resource_name {
82            format!(
83                "https://{}.services.ai.azure.com/anthropic/v1/messages",
84                resource
85            )
86        } else {
87            unreachable!("validated in from_config")
88        }
89    }
90
91    fn build_request_body(&self, request: &CreateMessageRequest) -> serde_json::Value {
92        build_messages_body(request, None, self.config.thinking_budget)
93    }
94
95    async fn get_auth_header(&self) -> Result<(String, String)> {
96        if let Some(ref api_key) = self.api_key {
97            return Ok(("api-key".into(), api_key.clone()));
98        }
99        let token = self.get_token().await?;
100        Ok(("Authorization".into(), format!("Bearer {}", token)))
101    }
102
103    async fn get_token(&self) -> Result<String> {
104        {
105            let cache = self.token_cache.read().await;
106            if let Some(ref cached) = *cache
107                && !cached.is_expired()
108            {
109                return Ok(cached.token().to_string());
110            }
111        }
112
113        let token_response = self
114            .credential
115            .get_token(&[COGNITIVE_SERVICES_SCOPE], None)
116            .await
117            .map_err(|e| Error::auth(format!("Failed to get Azure token: {}", e)))?;
118
119        let token_str = token_response.token.secret().to_string();
120        let cached = CachedToken::new(token_str.clone(), Duration::from_secs(3600));
121        *self.token_cache.write().await = Some(cached);
122
123        Ok(token_str)
124    }
125
126    async fn execute_request(
127        &self,
128        http: &reqwest::Client,
129        url: &str,
130        body: &serde_json::Value,
131    ) -> Result<reqwest::Response> {
132        let (header_name, header_value) = self.get_auth_header().await?;
133        let headers = vec![
134            (header_name, header_value),
135            ("anthropic-version".into(), ANTHROPIC_VERSION.into()),
136        ];
137        RequestExecutor::post(http, url, body, headers).await
138    }
139}
140
141#[async_trait]
142impl ProviderAdapter for FoundryAdapter {
143    fn config(&self) -> &ProviderConfig {
144        &self.config
145    }
146
147    fn name(&self) -> &'static str {
148        "foundry"
149    }
150
151    async fn build_url(&self, _model: &str, _stream: bool) -> String {
152        self.build_messages_url()
153    }
154
155    async fn transform_request(&self, request: CreateMessageRequest) -> Result<serde_json::Value> {
156        Ok(self.build_request_body(&request))
157    }
158
159    fn transform_response(&self, response: serde_json::Value) -> Result<ApiResponse> {
160        serde_json::from_value(response).map_err(|e| Error::Parse(e.to_string()))
161    }
162
163    async fn send(
164        &self,
165        http: &reqwest::Client,
166        request: CreateMessageRequest,
167    ) -> Result<ApiResponse> {
168        let url = self.build_messages_url();
169        let body = self.build_request_body(&request);
170        let response = self.execute_request(http, &url, &body).await?;
171        let json: serde_json::Value = response.json().await?;
172        self.transform_response(json)
173    }
174
175    async fn send_stream(
176        &self,
177        http: &reqwest::Client,
178        mut request: CreateMessageRequest,
179    ) -> Result<reqwest::Response> {
180        request.stream = Some(true);
181        let url = self.build_messages_url();
182        let body = self.build_request_body(&request);
183        self.execute_request(http, &url, &body).await
184    }
185
186    async fn refresh_credentials(&self) -> Result<()> {
187        *self.token_cache.write().await = None;
188        if self.api_key.is_none() {
189            self.get_token().await?;
190        }
191        Ok(())
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use crate::client::adapter::ModelConfig;
198
199    #[test]
200    fn test_build_url_with_resource() {
201        let url = format!(
202            "https://{}.services.ai.azure.com/anthropic/v1/messages",
203            "my-resource"
204        );
205        assert!(url.contains("services.ai.azure.com"));
206        assert!(url.contains("/anthropic/v1/messages"));
207    }
208
209    #[test]
210    fn test_build_url_with_base_url() {
211        let base_url = "https://custom-endpoint.azure.com/anthropic";
212        let url = format!("{}/v1/messages", base_url.trim_end_matches('/'));
213        assert!(url.contains("custom-endpoint.azure.com"));
214        assert!(url.contains("/v1/messages"));
215    }
216
217    #[test]
218    fn test_model_config() {
219        let config = ModelConfig::foundry();
220        assert!(config.primary.contains("sonnet"));
221    }
222
223    #[test]
224    fn test_anthropic_version() {
225        assert_eq!(super::ANTHROPIC_VERSION, "2023-06-01");
226    }
227}