rigg_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/// Azure OpenAI model deployment
184#[derive(Debug, Clone, Deserialize)]
185pub struct ModelDeployment {
186    pub name: String,
187    #[serde(default)]
188    pub properties: ModelDeploymentProperties,
189    #[serde(default)]
190    pub sku: ModelDeploymentSku,
191}
192
193#[derive(Debug, Clone, Default, Deserialize)]
194pub struct ModelDeploymentProperties {
195    #[serde(default)]
196    pub model: ModelDeploymentModel,
197}
198
199#[derive(Debug, Clone, Default, Deserialize)]
200pub struct ModelDeploymentModel {
201    #[serde(default)]
202    pub name: String,
203    #[serde(default)]
204    pub version: String,
205}
206
207#[derive(Debug, Clone, Default, Deserialize)]
208pub struct ModelDeploymentSku {
209    #[serde(default)]
210    pub name: String,
211    #[serde(default)]
212    pub capacity: u32,
213}
214
215impl std::fmt::Display for ModelDeployment {
216    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
217        write!(
218            f,
219            "{} ({}, {})",
220            self.name, self.properties.model.name, self.sku.name
221        )
222    }
223}
224
225/// ARM list response envelope
226#[derive(Debug, Deserialize)]
227struct ArmListResponse<T> {
228    value: Vec<T>,
229}
230
231impl ArmClient {
232    /// Create a new ARM client using Azure CLI credentials
233    pub fn new() -> Result<Self, ClientError> {
234        let token = AzCliAuth::get_arm_token()?;
235        let http = Client::builder()
236            .timeout(std::time::Duration::from_secs(30))
237            .build()?;
238
239        Ok(Self { http, token })
240    }
241
242    /// The ARM bearer token this client authenticated with.
243    pub fn token(&self) -> &str {
244        &self.token
245    }
246
247    /// List subscriptions the user has access to
248    pub async fn list_subscriptions(&self) -> Result<Vec<Subscription>, ClientError> {
249        let url = format!("{}/subscriptions?api-version=2022-12-01", ARM_BASE_URL);
250        debug!("Listing subscriptions: {}", url);
251
252        let response = self
253            .http
254            .get(&url)
255            .header("Authorization", format!("Bearer {}", self.token))
256            .send()
257            .await?;
258
259        let status = response.status();
260        if !status.is_success() {
261            let body = response.text().await?;
262            return Err(ClientError::from_response(status.as_u16(), &body));
263        }
264
265        let result: ArmListResponse<Subscription> = response.json().await?;
266        // Only return enabled subscriptions
267        Ok(result
268            .value
269            .into_iter()
270            .filter(|s| s.state == "Enabled")
271            .collect())
272    }
273
274    /// List Azure AI Search services in a subscription
275    pub async fn list_search_services(
276        &self,
277        subscription_id: &str,
278    ) -> Result<Vec<SearchService>, ClientError> {
279        let url = format!(
280            "{}/subscriptions/{}/providers/Microsoft.Search/searchServices?api-version=2023-11-01",
281            ARM_BASE_URL, subscription_id
282        );
283        debug!("Listing search services: {}", url);
284
285        let response = self
286            .http
287            .get(&url)
288            .header("Authorization", format!("Bearer {}", self.token))
289            .send()
290            .await?;
291
292        let status = response.status();
293        if !status.is_success() {
294            let body = response.text().await?;
295            return Err(ClientError::from_response(status.as_u16(), &body));
296        }
297
298        let result: ArmListResponse<SearchService> = response.json().await?;
299        Ok(result.value)
300    }
301
302    /// Find the resource group of a search service by scanning the subscription.
303    ///
304    /// Returns the resource group name extracted from the service's ARM resource ID.
305    pub async fn find_resource_group(
306        &self,
307        subscription_id: &str,
308        service_name: &str,
309    ) -> Result<String, ClientError> {
310        let services = self.list_search_services(subscription_id).await?;
311
312        for svc in &services {
313            if svc.name.eq_ignore_ascii_case(service_name) {
314                // Parse resource group from ARM ID:
315                // /subscriptions/{sub}/resourceGroups/{rg}/providers/...
316                return parse_resource_group(&svc.id).ok_or_else(|| ClientError::Api {
317                    status: 0,
318                    message: format!("Could not parse resource group from ARM ID: {}", svc.id),
319                });
320            }
321        }
322
323        Err(ClientError::NotFound {
324            kind: "Search service".to_string(),
325            name: service_name.to_string(),
326        })
327    }
328
329    /// List Azure AI Services accounts in a subscription (filtered to kind=AIServices)
330    pub async fn list_ai_services_accounts(
331        &self,
332        subscription_id: &str,
333    ) -> Result<Vec<AiServicesAccount>, ClientError> {
334        let url = format!(
335            "{}/subscriptions/{}/providers/Microsoft.CognitiveServices/accounts?api-version=2024-10-01",
336            ARM_BASE_URL, subscription_id
337        );
338        debug!("Listing AI Services accounts: {}", url);
339
340        let response = self
341            .http
342            .get(&url)
343            .header("Authorization", format!("Bearer {}", self.token))
344            .send()
345            .await?;
346
347        let status = response.status();
348        if !status.is_success() {
349            let body = response.text().await?;
350            return Err(ClientError::from_response(status.as_u16(), &body));
351        }
352
353        let result: ArmListResponse<AiServicesAccount> = response.json().await?;
354        Ok(result
355            .value
356            .into_iter()
357            .filter(|a| a.kind.eq_ignore_ascii_case("AIServices"))
358            .collect())
359    }
360
361    /// List Microsoft Foundry projects under a specific AI Services account.
362    ///
363    /// Projects are sub-resources at:
364    /// `Microsoft.CognitiveServices/accounts/{accountName}/projects`
365    ///
366    /// The `account_id` should be the full ARM resource ID of the account,
367    /// from which we extract the resource group.
368    pub async fn list_foundry_projects(
369        &self,
370        account: &AiServicesAccount,
371        subscription_id: &str,
372    ) -> Result<Vec<FoundryProject>, ClientError> {
373        let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
374            status: 0,
375            message: format!("Could not parse resource group from ARM ID: {}", account.id),
376        })?;
377
378        let url = format!(
379            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/projects?api-version=2025-06-01",
380            ARM_BASE_URL, subscription_id, resource_group, account.name
381        );
382        debug!("Listing Foundry projects: {}", url);
383
384        let response = self
385            .http
386            .get(&url)
387            .header("Authorization", format!("Bearer {}", self.token))
388            .send()
389            .await?;
390
391        let status = response.status();
392        if !status.is_success() {
393            let body = response.text().await?;
394            return Err(ClientError::from_response(status.as_u16(), &body));
395        }
396
397        let result: ArmListResponse<FoundryProject> = response.json().await?;
398        Ok(result.value)
399    }
400
401    /// List storage accounts in a resource group.
402    pub async fn list_storage_accounts(
403        &self,
404        subscription_id: &str,
405        resource_group: &str,
406    ) -> Result<Vec<StorageAccount>, ClientError> {
407        let url = format!(
408            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts?api-version=2023-05-01",
409            ARM_BASE_URL, subscription_id, resource_group
410        );
411        debug!("Listing storage accounts: {}", url);
412
413        let response = self
414            .http
415            .get(&url)
416            .header("Authorization", format!("Bearer {}", self.token))
417            .send()
418            .await?;
419
420        let status = response.status();
421        if !status.is_success() {
422            let body = response.text().await?;
423            return Err(ClientError::from_response(status.as_u16(), &body));
424        }
425
426        let result: ArmListResponse<StorageAccount> = response.json().await?;
427        Ok(result.value)
428    }
429
430    /// Get the primary access key for a storage account.
431    pub async fn get_storage_account_key(
432        &self,
433        subscription_id: &str,
434        resource_group: &str,
435        account_name: &str,
436    ) -> Result<String, ClientError> {
437        let url = format!(
438            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.Storage/storageAccounts/{}/listKeys?api-version=2023-05-01",
439            ARM_BASE_URL, subscription_id, resource_group, account_name
440        );
441        debug!("Getting storage account keys: {}", url);
442
443        let response = self
444            .http
445            .post(&url)
446            .header("Authorization", format!("Bearer {}", self.token))
447            .header("Content-Length", "0")
448            .send()
449            .await?;
450
451        let status = response.status();
452        if !status.is_success() {
453            let body = response.text().await?;
454            return Err(ClientError::from_response(status.as_u16(), &body));
455        }
456
457        let key_list: StorageKeyList = response.json().await?;
458        key_list
459            .keys
460            .into_iter()
461            .next()
462            .map(|k| k.value)
463            .ok_or_else(|| ClientError::Api {
464                status: 0,
465                message: "No keys found for storage account".to_string(),
466            })
467    }
468
469    /// List model deployments for an AI Services account.
470    pub async fn list_model_deployments(
471        &self,
472        account: &AiServicesAccount,
473        subscription_id: &str,
474    ) -> Result<Vec<ModelDeployment>, ClientError> {
475        let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
476            status: 0,
477            message: format!("Could not parse resource group from ARM ID: {}", account.id),
478        })?;
479
480        let url = format!(
481            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/deployments?api-version=2024-10-01",
482            ARM_BASE_URL, subscription_id, resource_group, account.name
483        );
484        debug!("Listing model deployments: {}", url);
485
486        let response = self
487            .http
488            .get(&url)
489            .header("Authorization", format!("Bearer {}", self.token))
490            .send()
491            .await?;
492
493        let status = response.status();
494        if !status.is_success() {
495            let body = response.text().await?;
496            return Err(ClientError::from_response(status.as_u16(), &body));
497        }
498
499        let result: ArmListResponse<ModelDeployment> = response.json().await?;
500        Ok(result.value)
501    }
502
503    /// Create a model deployment on an AI Services account.
504    pub async fn create_model_deployment(
505        &self,
506        account: &AiServicesAccount,
507        subscription_id: &str,
508        deployment_name: &str,
509        model_name: &str,
510        model_version: &str,
511    ) -> Result<(), ClientError> {
512        let resource_group = parse_resource_group(&account.id).ok_or_else(|| ClientError::Api {
513            status: 0,
514            message: format!("Could not parse resource group from ARM ID: {}", account.id),
515        })?;
516
517        let url = format!(
518            "{}/subscriptions/{}/resourceGroups/{}/providers/Microsoft.CognitiveServices/accounts/{}/deployments/{}?api-version=2024-10-01",
519            ARM_BASE_URL, subscription_id, resource_group, account.name, deployment_name
520        );
521        debug!("Creating model deployment: {}", url);
522
523        let body = serde_json::json!({
524            "sku": {
525                "name": "GlobalStandard",
526                "capacity": 1
527            },
528            "properties": {
529                "model": {
530                    "format": "OpenAI",
531                    "name": model_name,
532                    "version": model_version
533                }
534            }
535        });
536
537        let response = self
538            .http
539            .put(&url)
540            .header("Authorization", format!("Bearer {}", self.token))
541            .json(&body)
542            .send()
543            .await?;
544
545        let status = response.status();
546        if !status.is_success() {
547            let body = response.text().await?;
548            return Err(ClientError::from_response(status.as_u16(), &body));
549        }
550
551        Ok(())
552    }
553
554    /// Build a full connection string for a storage account.
555    pub async fn get_storage_connection_string(
556        &self,
557        subscription_id: &str,
558        resource_group: &str,
559        account_name: &str,
560    ) -> Result<String, ClientError> {
561        let key = self
562            .get_storage_account_key(subscription_id, resource_group, account_name)
563            .await?;
564
565        Ok(format!(
566            "DefaultEndpointsProtocol=https;AccountName={};AccountKey={};EndpointSuffix=core.windows.net",
567            account_name, key
568        ))
569    }
570}
571
572/// Parse resource group from an ARM resource ID.
573///
574/// ARM IDs look like: `/subscriptions/{sub}/resourceGroups/{rg}/providers/...`
575fn parse_resource_group(arm_id: &str) -> Option<String> {
576    let parts: Vec<&str> = arm_id.split('/').collect();
577    for (i, part) in parts.iter().enumerate() {
578        if part.eq_ignore_ascii_case("resourceGroups")
579            || part.eq_ignore_ascii_case("resourcegroups")
580        {
581            return parts.get(i + 1).map(|s| s.to_string());
582        }
583    }
584    None
585}
586
587#[cfg(test)]
588mod tests {
589    use super::*;
590
591    #[test]
592    fn test_parse_resource_group() {
593        let id = "/subscriptions/abc-123/resourceGroups/my-rg/providers/Microsoft.Search/searchServices/my-svc";
594        assert_eq!(parse_resource_group(id), Some("my-rg".to_string()));
595    }
596
597    #[test]
598    fn test_parse_resource_group_case_insensitive() {
599        let id = "/subscriptions/abc/resourcegroups/MyRG/providers/Something";
600        assert_eq!(parse_resource_group(id), Some("MyRG".to_string()));
601    }
602
603    #[test]
604    fn test_parse_resource_group_missing() {
605        let id = "/subscriptions/abc/providers/Something";
606        assert_eq!(parse_resource_group(id), None);
607    }
608
609    #[test]
610    fn test_ai_services_account_display() {
611        let account = AiServicesAccount {
612            name: "my-ai-service".to_string(),
613            location: "eastus".to_string(),
614            kind: "AIServices".to_string(),
615            id: String::new(),
616            properties: AiServicesAccountProperties::default(),
617        };
618        assert_eq!(format!("{}", account), "my-ai-service (eastus)");
619    }
620
621    #[test]
622    fn test_agents_endpoint_from_arm_endpoint() {
623        let account = AiServicesAccount {
624            name: "irma-prod-foundry".to_string(),
625            location: "swedencentral".to_string(),
626            kind: "AIServices".to_string(),
627            id: String::new(),
628            properties: AiServicesAccountProperties {
629                endpoint: Some("https://custom-subdomain.cognitiveservices.azure.com/".to_string()),
630            },
631        };
632        assert_eq!(
633            account.agents_endpoint(),
634            "https://custom-subdomain.services.ai.azure.com"
635        );
636    }
637
638    #[test]
639    fn test_agents_endpoint_fallback_to_name() {
640        let account = AiServicesAccount {
641            name: "irma-prod-foundry".to_string(),
642            location: "swedencentral".to_string(),
643            kind: "AIServices".to_string(),
644            id: String::new(),
645            properties: AiServicesAccountProperties::default(),
646        };
647        assert_eq!(
648            account.agents_endpoint(),
649            "https://irma-prod-foundry.services.ai.azure.com"
650        );
651    }
652
653    #[test]
654    fn test_extract_subdomain() {
655        assert_eq!(
656            extract_subdomain("https://my-svc.cognitiveservices.azure.com/"),
657            Some("my-svc")
658        );
659        assert_eq!(
660            extract_subdomain("https://custom.services.ai.azure.com"),
661            Some("custom")
662        );
663        assert_eq!(extract_subdomain("not-a-url"), None);
664    }
665
666    #[test]
667    fn test_foundry_project_display_with_display_name() {
668        let project = FoundryProject {
669            name: "my-account/my-project".to_string(),
670            location: "westus2".to_string(),
671            id: String::new(),
672            properties: FoundryProjectProperties {
673                display_name: "my-project".to_string(),
674            },
675        };
676        assert_eq!(format!("{}", project), "my-project (westus2)");
677        assert_eq!(project.display_name(), "my-project");
678    }
679
680    #[test]
681    fn test_model_deployment_display() {
682        let deployment = ModelDeployment {
683            name: "gpt-4o-mini".to_string(),
684            properties: ModelDeploymentProperties {
685                model: ModelDeploymentModel {
686                    name: "gpt-4o-mini".to_string(),
687                    version: "2024-07-18".to_string(),
688                },
689            },
690            sku: ModelDeploymentSku {
691                name: "GlobalStandard".to_string(),
692                capacity: 1,
693            },
694        };
695        assert_eq!(
696            format!("{}", deployment),
697            "gpt-4o-mini (gpt-4o-mini, GlobalStandard)"
698        );
699    }
700
701    #[test]
702    fn test_foundry_project_display_name_fallback() {
703        let project = FoundryProject {
704            name: "my-account/proj-default".to_string(),
705            location: "swedencentral".to_string(),
706            id: String::new(),
707            properties: FoundryProjectProperties::default(),
708        };
709        assert_eq!(project.display_name(), "proj-default");
710        assert_eq!(format!("{}", project), "proj-default (swedencentral)");
711    }
712}
rigg_client/arm.rs

rigg_client/
arm.rs
