claude_agent/client/adapter/
foundry.rs1use 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}