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

rigg_client/
arm.rs
