1use reqwest::Client;
4use serde::Deserialize;
5use tracing::debug;
6
7use crate::auth::AzCliAuth;
8use crate::error::ClientError;
9
10const ARM_BASE_URL: &str = "https://management.azure.com";
11
12pub struct ArmClient {
14 http: Client,
15 token: String,
16}
17
18#[derive(Debug, Clone, Deserialize)]
20#[serde(rename_all = "camelCase")]
21pub struct Subscription {
22 pub subscription_id: String,
23 pub display_name: String,
24 pub state: String,
25}
26
27impl std::fmt::Display for Subscription {
28 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
29 write!(f, "{} ({})", self.display_name, self.subscription_id)
30 }
31}
32
33#[derive(Debug, Clone, Deserialize)]
35pub struct SearchService {
36 pub name: String,
37 pub location: String,
38 pub sku: SearchServiceSku,
39 #[serde(default)]
40 pub id: String,
41}
42
43#[derive(Debug, Clone, Deserialize)]
44pub struct SearchServiceSku {
45 pub name: String,
46}
47
48impl std::fmt::Display for SearchService {
49 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50 write!(
51 f,
52 "{} ({}, {})",
53 self.name,
54 self.location,
55 self.sku.name.to_uppercase()
56 )
57 }
58}
59
60#[derive(Debug, Clone)]
62pub struct DiscoveredService {
63 pub name: String,
64 pub subscription_id: String,
65 pub location: String,
66}
67
68#[derive(Debug, Clone, Deserialize)]
70pub struct AiServicesAccount {
71 pub name: String,
72 pub location: String,
73 #[serde(default)]
74 pub kind: String,
75 #[serde(default)]
76 pub id: String,
77 #[serde(default)]
78 pub properties: AiServicesAccountProperties,
79}
80
81#[derive(Debug, Clone, Default, Deserialize)]
82pub struct AiServicesAccountProperties {
83 #[serde(default)]
85 pub endpoint: Option<String>,
86}
87
88impl AiServicesAccount {
89 pub fn agents_endpoint(&self) -> String {
95 if let Some(ref endpoint) = self.properties.endpoint {
96 if let Some(subdomain) = extract_subdomain(endpoint) {
97 return format!("https://{}.services.ai.azure.com", subdomain);
98 }
99 }
100 format!("https://{}.services.ai.azure.com", self.name)
101 }
102}
103
104fn extract_subdomain(endpoint: &str) -> Option<&str> {
108 let host = endpoint.strip_prefix("https://")?.split('/').next()?;
109 host.split('.').next()
110}
111
112impl std::fmt::Display for AiServicesAccount {
113 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114 write!(f, "{} ({})", self.name, self.location)
115 }
116}
117
118#[derive(Debug, Clone, Deserialize)]
120pub struct FoundryProject {
121 #[serde(default)]
123 name: String,
124 pub location: String,
125 #[serde(default)]
126 pub id: String,
127 #[serde(default)]
128 pub properties: FoundryProjectProperties,
129}
130
131#[derive(Debug, Clone, Default, Deserialize)]
132#[serde(rename_all = "camelCase")]
133pub struct FoundryProjectProperties {
134 #[serde(default)]
135 pub display_name: String,
136}
137
138impl FoundryProject {
139 pub fn display_name(&self) -> &str {
141 if !self.properties.display_name.is_empty() {
142 &self.properties.display_name
143 } else {
144 self.name.rsplit('/').next().unwrap_or(&self.name)
146 }
147 }
148}
149
150impl std::fmt::Display for FoundryProject {
151 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
152 write!(f, "{} ({})", self.display_name(), self.location)
153 }
154}
155
156#[derive(Debug, Clone, Deserialize)]
158pub struct StorageAccount {
159 pub name: String,
160 pub location: String,
161 #[serde(default)]
162 pub id: String,
163}
164
165impl std::fmt::Display for StorageAccount {
166 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
167 write!(f, "{} ({})", self.name, self.location)
168 }
169}
170
171#[derive(Debug, Clone, Deserialize)]
173struct StorageKey {
174 value: String,
175}
176
177#[derive(Debug, Deserialize)]
179struct StorageKeyList {
180 keys: Vec<StorageKey>,
181}
182
183#[derive(Debug, Deserialize)]
185struct ArmListResponse<T> {
186 value: Vec<T>,
187}
188
189impl ArmClient {
190 pub fn new() -> Result<Self, ClientError> {
192 let token = AzCliAuth::get_arm_token()?;
193 let http = Client::builder()
194 .timeout(std::time::Duration::from_secs(30))
195 .build()?;
196
197 Ok(Self { http, token })
198 }
199
200 pub async fn list_subscriptions(&self) -> Result<Vec<Subscription>, ClientError> {
202 let url = format!("{}/subscriptions?api-version=2022-12-01", ARM_BASE_URL);
203 debug!("Listing subscriptions: {}", url);
204
205 let response = self
206 .http
207 .get(&url)
208 .header("Authorization", format!("Bearer {}", self.token))
209 .send()
210 .await?;
211
212 let status = response.status();
213 if !status.is_success() {
214 let body = response.text().await?;
215 return Err(ClientError::from_response(status.as_u16(), &body));
216 }
217
218 let result: ArmListResponse<Subscription> = response.json().await?;
219 Ok(result
221 .value
222 .into_iter()
223 .filter(|s| s.state == "Enabled")
224 .collect())
225 }
226
227 pub async fn list_search_services(
229 &self,
230 subscription_id: &str,
231 ) -> Result<Vec<SearchService>, ClientError> {
232 let url = format!(
233 "{}/subscriptions/{}/providers/Microsoft.Search/searchServices?api-version=2023-11-01",
234 ARM_BASE_URL, subscription_id
235 );
236 debug!("Listing search services: {}", url);
237
238 let response = self
239 .http
240 .get(&url)
241 .header("Authorization", format!("Bearer {}", self.token))
242 .send()
243 .await?;
244
245 let status = response.status();
246 if !status.is_success() {
247 let body = response.text().await?;
248 return Err(ClientError::from_response(status.as_u16(), &body));
249 }
250
251 let result: ArmListResponse<SearchService> = response.json().await?;
252 Ok(result.value)
253 }
254
255 pub async fn find_resource_group(
259 &self,
260 subscription_id: &str,
261 service_name: &str,
262 ) -> Result<String, ClientError> {
263 let services = self.list_search_services(subscription_id).await?;
264
265 for svc in &services {
266 if svc.name.eq_ignore_ascii_case(service_name) {
267 return parse_resource_group(&svc.id).ok_or_else(|| ClientError::Api {
270 status: 0,
271 message: format!("Could not parse resource group from ARM ID: {}", svc.id),
272 });
273 }
274 }
275
276 Err(ClientError::NotFound {
277 kind: "Search service".to_string(),
278 name: service_name.to_string(),
279 })
280 }
281
282 pub async fn list_ai_services_accounts(
284 &self,
285 subscription_id: &str,
286 ) -> Result<Vec<AiServicesAccount>, ClientError> {
287 let url = format!(
288 "{}/subscriptions/{}/providers/Microsoft.CognitiveServices/accounts?api-version=2024-10-01",
289 ARM_BASE_URL, subscription_id
290 );
291 debug!("Listing AI Services accounts: {}", url);
292
293 let response = self
294 .http
295 .get(&url)
296 .header("Authorization", format!("Bearer {}", self.token))
297 .send()
298 .await?;
299
300 let status = response.status();
301 if !status.is_success() {
302 let body = response.text().await?;
303 return Err(ClientError::from_response(status.as_u16(), &body));
304 }
305
306 let result: ArmListResponse<AiServicesAccount> = response.json().await?;
307 Ok(result
308 .value
309 .into_iter()
310 .filter(|a| a.kind.eq_ignore_ascii_case("AIServices"))
311 .collect())
312 }
313
314 pub async fn list_foundry_projects(
322 &self,
323 account: &AiServicesAccount,
324 subscription_id: &str,
325 ) -> Result<Vec<FoundryProject>, ClientError> {
326 let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
327 status: 0,
328 message: format!("Could not parse resource group from ARM ID: {}", account.id),
329 })?;
330
331 let url = format!(
332 "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/projects?api-version=2025-06-01",
333 ARM_BASE_URL, subscription_id, resource_group, account.name
334 );
335 debug!("Listing Foundry projects: {}", url);
336
337 let response = self
338 .http
339 .get(&url)
340 .header("Authorization", format!("Bearer {}", self.token))
341 .send()
342 .await?;
343
344 let status = response.status();
345 if !status.is_success() {
346 let body = response.text().await?;
347 return Err(ClientError::from_response(status.as_u16(), &body));
348 }
349
350 let result: ArmListResponse<FoundryProject> = response.json().await?;
351 Ok(result.value)
352 }
353
354 pub async fn list_storage_accounts(
356 &self,
357 subscription_id: &str,
358 resource_group: &str,
359 ) -> Result<Vec<StorageAccount>, ClientError> {
360 let url = format!(
361 "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts?api-version=2023-05-01",
362 ARM_BASE_URL, subscription_id, resource_group
363 );
364 debug!("Listing storage accounts: {}", url);
365
366 let response = self
367 .http
368 .get(&url)
369 .header("Authorization", format!("Bearer {}", self.token))
370 .send()
371 .await?;
372
373 let status = response.status();
374 if !status.is_success() {
375 let body = response.text().await?;
376 return Err(ClientError::from_response(status.as_u16(), &body));
377 }
378
379 let result: ArmListResponse<StorageAccount> = response.json().await?;
380 Ok(result.value)
381 }
382
383 pub async fn get_storage_account_key(
385 &self,
386 subscription_id: &str,
387 resource_group: &str,
388 account_name: &str,
389 ) -> Result<String, ClientError> {
390 let url = format!(
391 "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts/{}/listKeys?api-version=2023-05-01",
392 ARM_BASE_URL, subscription_id, resource_group, account_name
393 );
394 debug!("Getting storage account keys: {}", url);
395
396 let response = self
397 .http
398 .post(&url)
399 .header("Authorization", format!("Bearer {}", self.token))
400 .header("Content-Length", "0")
401 .send()
402 .await?;
403
404 let status = response.status();
405 if !status.is_success() {
406 let body = response.text().await?;
407 return Err(ClientError::from_response(status.as_u16(), &body));
408 }
409
410 let key_list: StorageKeyList = response.json().await?;
411 key_list
412 .keys
413 .into_iter()
414 .next()
415 .map(|k| k.value)
416 .ok_or_else(|| ClientError::Api {
417 status: 0,
418 message: "No keys found for storage account".to_string(),
419 })
420 }
421
422 pub async fn get_storage_connection_string(
424 &self,
425 subscription_id: &str,
426 resource_group: &str,
427 account_name: &str,
428 ) -> Result<String, ClientError> {
429 let key = self
430 .get_storage_account_key(subscription_id, resource_group, account_name)
431 .await?;
432
433 Ok(format!(
434 "DefaultEndpointsProtocol=https;AccountName={};AccountKey={};EndpointSuffix=core.windows.net",
435 account_name, key
436 ))
437 }
438}
439
440fn parse_resource_group(arm_id: &str) -> Option<String> {
444 let parts: Vec<&str> = arm_id.split('/').collect();
445 for (i, part) in parts.iter().enumerate() {
446 if part.eq_ignore_ascii_case("resourceGroups")
447 || part.eq_ignore_ascii_case("resourcegroups")
448 {
449 return parts.get(i + 1).map(|s| s.to_string());
450 }
451 }
452 None
453}
454
455#[cfg(test)]
456mod tests {
457 use super::*;
458
459 #[test]
460 fn test_parse_resource_group() {
461 let id = "/subscriptions/abc-123/resourceGroups/my-rg/providers/Microsoft.Search/searchServices/my-svc";
462 assert_eq!(parse_resource_group(id), Some("my-rg".to_string()));
463 }
464
465 #[test]
466 fn test_parse_resource_group_case_insensitive() {
467 let id = "/subscriptions/abc/resourcegroups/MyRG/providers/Something";
468 assert_eq!(parse_resource_group(id), Some("MyRG".to_string()));
469 }
470
471 #[test]
472 fn test_parse_resource_group_missing() {
473 let id = "/subscriptions/abc/providers/Something";
474 assert_eq!(parse_resource_group(id), None);
475 }
476
477 #[test]
478 fn test_ai_services_account_display() {
479 let account = AiServicesAccount {
480 name: "my-ai-service".to_string(),
481 location: "eastus".to_string(),
482 kind: "AIServices".to_string(),
483 id: String::new(),
484 properties: AiServicesAccountProperties::default(),
485 };
486 assert_eq!(format!("{}", account), "my-ai-service (eastus)");
487 }
488
489 #[test]
490 fn test_agents_endpoint_from_arm_endpoint() {
491 let account = AiServicesAccount {
492 name: "irma-prod-foundry".to_string(),
493 location: "swedencentral".to_string(),
494 kind: "AIServices".to_string(),
495 id: String::new(),
496 properties: AiServicesAccountProperties {
497 endpoint: Some("https://custom-subdomain.cognitiveservices.azure.com/".to_string()),
498 },
499 };
500 assert_eq!(
501 account.agents_endpoint(),
502 "https://custom-subdomain.services.ai.azure.com"
503 );
504 }
505
506 #[test]
507 fn test_agents_endpoint_fallback_to_name() {
508 let account = AiServicesAccount {
509 name: "irma-prod-foundry".to_string(),
510 location: "swedencentral".to_string(),
511 kind: "AIServices".to_string(),
512 id: String::new(),
513 properties: AiServicesAccountProperties::default(),
514 };
515 assert_eq!(
516 account.agents_endpoint(),
517 "https://irma-prod-foundry.services.ai.azure.com"
518 );
519 }
520
521 #[test]
522 fn test_extract_subdomain() {
523 assert_eq!(
524 extract_subdomain("https://my-svc.cognitiveservices.azure.com/"),
525 Some("my-svc")
526 );
527 assert_eq!(
528 extract_subdomain("https://custom.services.ai.azure.com"),
529 Some("custom")
530 );
531 assert_eq!(extract_subdomain("not-a-url"), None);
532 }
533
534 #[test]
535 fn test_foundry_project_display_with_display_name() {
536 let project = FoundryProject {
537 name: "my-account/my-project".to_string(),
538 location: "westus2".to_string(),
539 id: String::new(),
540 properties: FoundryProjectProperties {
541 display_name: "my-project".to_string(),
542 },
543 };
544 assert_eq!(format!("{}", project), "my-project (westus2)");
545 assert_eq!(project.display_name(), "my-project");
546 }
547
548 #[test]
549 fn test_foundry_project_display_name_fallback() {
550 let project = FoundryProject {
551 name: "my-account/proj-default".to_string(),
552 location: "swedencentral".to_string(),
553 id: String::new(),
554 properties: FoundryProjectProperties::default(),
555 };
556 assert_eq!(project.display_name(), "proj-default");
557 assert_eq!(format!("{}", project), "proj-default (swedencentral)");
558 }
559}