Skip to main content

hoist_client/
arm.rs

1//! Azure Resource Manager client for discovering Search and Foundry services
2
3use 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
12/// Azure Resource Manager client for subscription/service discovery
13pub struct ArmClient {
14    http: Client,
15    token: String,
16}
17
18/// Azure subscription
19#[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/// Azure AI Search service
34#[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/// Result of the discovery flow
61#[derive(Debug, Clone)]
62pub struct DiscoveredService {
63    pub name: String,
64    pub subscription_id: String,
65    pub location: String,
66}
67
68/// Azure AI Services account (kind=AIServices)
69#[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    /// Primary endpoint (e.g., "https://name.cognitiveservices.azure.com/")
84    #[serde(default)]
85    pub endpoint: Option<String>,
86}
87
88impl AiServicesAccount {
89    /// Derive the `.services.ai.azure.com` endpoint for the agents API.
90    ///
91    /// Extracts the custom subdomain from the ARM `properties.endpoint`
92    /// (which may differ from the resource name), then constructs the
93    /// AI services endpoint. Falls back to the resource name.
94    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
104/// Extract the subdomain from an Azure endpoint URL.
105///
106/// `"https://my-svc.cognitiveservices.azure.com/"` → `"my-svc"`
107fn 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/// Microsoft Foundry project (sub-resource of AI Services account)
119#[derive(Debug, Clone, Deserialize)]
120pub struct FoundryProject {
121    /// ARM name — may be "accountName/projectName" for sub-resources
122    #[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    /// The project display name (human-friendly, e.g. "proj-default")
140    pub fn display_name(&self) -> &str {
141        if !self.properties.display_name.is_empty() {
142            &self.properties.display_name
143        } else {
144            // Fallback: parse from "account/project" ARM name
145            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/// Azure Storage account
157#[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/// Storage account key
172#[derive(Debug, Clone, Deserialize)]
173struct StorageKey {
174    value: String,
175}
176
177/// Storage account key list response
178#[derive(Debug, Deserialize)]
179struct StorageKeyList {
180    keys: Vec<StorageKey>,
181}
182
183/// ARM list response envelope
184#[derive(Debug, Deserialize)]
185struct ArmListResponse<T> {
186    value: Vec<T>,
187}
188
189impl ArmClient {
190    /// Create a new ARM client using Azure CLI credentials
191    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    /// List subscriptions the user has access to
201    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        // Only return enabled subscriptions
220        Ok(result
221            .value
222            .into_iter()
223            .filter(|s| s.state == "Enabled")
224            .collect())
225    }
226
227    /// List Azure AI Search services in a subscription
228    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    /// Find the resource group of a search service by scanning the subscription.
256    ///
257    /// Returns the resource group name extracted from the service's ARM resource ID.
258    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                // Parse resource group from ARM ID:
268                // /subscriptions/{sub}/resourceGroups/{rg}/providers/...
269                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    /// List Azure AI Services accounts in a subscription (filtered to kind=AIServices)
283    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    /// List Microsoft Foundry projects under a specific AI Services account.
315    ///
316    /// Projects are sub-resources at:
317    /// `Microsoft.CognitiveServices/accounts/{accountName}/projects`
318    ///
319    /// The `account_id` should be the full ARM resource ID of the account,
320    /// from which we extract the resource group.
321    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    /// List storage accounts in a resource group.
355    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    /// Get the primary access key for a storage account.
384    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    /// Build a full connection string for a storage account.
423    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
440/// Parse resource group from an ARM resource ID.
441///
442/// ARM IDs look like: `/subscriptions/{sub}/resourceGroups/{rg}/providers/...`
443fn 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}