rigg_client/
foundry.rs

1//! Microsoft Foundry REST API client
2//!
3//! Manages Foundry agents via the project-scoped `/agents` API
4//! (v1 stable data plane at `https://{account}.services.ai.azure.com/api/projects/{project}`).
5
6use std::time::Duration;
7
8use reqwest::{Client, Method, StatusCode};
9use serde_json::{Map, Value};
10use tracing::{debug, instrument, warn};
11
12use rigg_core::config::FoundryServiceConfig;
13
14use crate::auth::{AuthProvider, get_auth_provider_for};
15use crate::error::ClientError;
16
17/// Maximum number of retry attempts for retryable errors
18const MAX_RETRIES: u32 = 3;
19
20/// Initial backoff delay in seconds
21const INITIAL_BACKOFF_SECS: u64 = 1;
22
23/// Microsoft Foundry API client
24pub struct FoundryClient {
25    http: Client,
26    auth: Box<dyn AuthProvider>,
27    base_url: String,
28    project: String,
29    api_version: String,
30    /// Optional `Foundry-Features` header values gating preview features
31    /// (e.g. `HostedAgents=V1Preview`).
32    features: Vec<String>,
33}
34
35impl FoundryClient {
36    /// Create a client from a workspace foundry connection (v1 data plane).
37    pub fn from_connection(
38        conn: &rigg_core::workspace::FoundryConnection,
39    ) -> Result<Self, ClientError> {
40        let auth = get_auth_provider_for(rigg_core::ServiceDomain::Foundry)?;
41        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
42        Ok(Self {
43            http,
44            auth,
45            base_url: conn
46                .endpoint
47                .clone()
48                .map(|e| e.trim_end_matches('/').to_string())
49                .unwrap_or_else(|| format!("https://{}.services.ai.azure.com", conn.account)),
50            project: conn.project.clone(),
51            api_version: conn
52                .api_version
53                .clone()
54                .unwrap_or_else(|| rigg_core::registry::FOUNDRY_API_VERSION.to_string()),
55            features: Vec::new(),
56        })
57    }
58
59    /// Create a new Foundry client from service configuration (legacy).
60    pub fn new(config: &FoundryServiceConfig) -> Result<Self, ClientError> {
61        let auth = get_auth_provider_for(rigg_core::ServiceDomain::Foundry)?;
62        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
63
64        Ok(Self {
65            http,
66            auth,
67            base_url: config.service_url(),
68            project: config.project.clone(),
69            api_version: rigg_core::registry::FOUNDRY_API_VERSION.to_string(),
70            features: Vec::new(),
71        })
72    }
73
74    /// Create with a custom auth provider (for testing)
75    pub fn with_auth(
76        base_url: String,
77        project: String,
78        api_version: String,
79        auth: Box<dyn AuthProvider>,
80    ) -> Result<Self, ClientError> {
81        let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
82
83        Ok(Self {
84            http,
85            auth,
86            base_url,
87            project,
88            api_version,
89            features: Vec::new(),
90        })
91    }
92
93    /// Enable preview feature gates sent via the `Foundry-Features` header.
94    pub fn with_features(mut self, features: &[&str]) -> Self {
95        self.features = features.iter().map(|s| s.to_string()).collect();
96        self
97    }
98
99    /// Build URL for the agents collection
100    fn agents_url(&self) -> String {
101        format!(
102            "{}/api/projects/{}/agents?api-version={}",
103            self.base_url, self.project, self.api_version
104        )
105    }
106
107    /// Build URL for a specific agent
108    fn agent_url(&self, id: &str) -> String {
109        format!(
110            "{}/api/projects/{}/agents/{}?api-version={}",
111            self.base_url,
112            self.project,
113            urlencoding::encode(id),
114            self.api_version
115        )
116    }
117
118    /// Build URL for creating/updating agent versions
119    fn agent_versions_url(&self, name: &str) -> String {
120        format!(
121            "{}/api/projects/{}/agents/{}/versions?api-version={}",
122            self.base_url,
123            self.project,
124            urlencoding::encode(name),
125            self.api_version
126        )
127    }
128
129    /// Execute an HTTP request
130    async fn request(
131        &self,
132        method: Method,
133        url: &str,
134        body: Option<&Value>,
135    ) -> Result<Option<Value>, ClientError> {
136        let token = self.auth.get_token()?;
137
138        let mut request = self
139            .http
140            .request(method.clone(), url)
141            .header("Authorization", format!("Bearer {}", token))
142            .header("Content-Type", "application/json");
143
144        if !self.features.is_empty() {
145            request = request.header("Foundry-Features", self.features.join(","));
146        }
147
148        if let Some(json) = body {
149            request = request.json(json);
150        }
151
152        debug!("Request: {} {}", method, url);
153        let response = request.send().await?;
154        let status = response.status();
155
156        if status == StatusCode::NO_CONTENT {
157            return Ok(None);
158        }
159
160        let body = response.text().await?;
161
162        if status.is_success() {
163            if body.is_empty() {
164                Ok(None)
165            } else {
166                let value: Value = serde_json::from_str(&body)?;
167                Ok(Some(value))
168            }
169        } else {
170            match status {
171                StatusCode::NOT_FOUND => Err(ClientError::NotFound {
172                    kind: "agent".to_string(),
173                    name: url.to_string(),
174                }),
175                StatusCode::TOO_MANY_REQUESTS => {
176                    let retry_after = 60;
177                    Err(ClientError::RateLimited { retry_after })
178                }
179                StatusCode::SERVICE_UNAVAILABLE => Err(ClientError::ServiceUnavailable(body)),
180                _ => Err(ClientError::from_response_with_url(
181                    status.as_u16(),
182                    &body,
183                    Some(url),
184                )),
185            }
186        }
187    }
188
189    /// Execute an HTTP request with retry logic
190    async fn request_with_retry(
191        &self,
192        method: Method,
193        url: &str,
194        body: Option<&Value>,
195    ) -> Result<Option<Value>, ClientError> {
196        let mut attempt = 0u32;
197        loop {
198            match self.request(method.clone(), url, body).await {
199                Ok(value) => return Ok(value),
200                Err(err) if err.is_retryable() && attempt < MAX_RETRIES => {
201                    let delay = match &err {
202                        ClientError::RateLimited { retry_after } => {
203                            Duration::from_secs(*retry_after)
204                        }
205                        _ => Duration::from_secs(INITIAL_BACKOFF_SECS * 2u64.pow(attempt)),
206                    };
207                    warn!(
208                        "Request {} {} failed (attempt {}/{}): {}. Retrying in {:?}",
209                        method,
210                        url,
211                        attempt + 1,
212                        MAX_RETRIES + 1,
213                        err,
214                        delay,
215                    );
216                    tokio::time::sleep(delay).await;
217                    attempt += 1;
218                }
219                Err(err) => return Err(err),
220            }
221        }
222    }
223
224    /// List all agents in the project
225    #[instrument(skip(self))]
226    pub async fn list_agents(&self) -> Result<Vec<Value>, ClientError> {
227        let url = self.agents_url();
228        let response = self.request_with_retry(Method::GET, &url, None).await?;
229
230        match response {
231            Some(value) => {
232                let items = value
233                    .get("data")
234                    .and_then(|v| v.as_array())
235                    .cloned()
236                    .unwrap_or_default();
237                // Flatten versioned response into flat agent objects
238                Ok(items.iter().map(flatten_agent_response).collect())
239            }
240            None => Ok(Vec::new()),
241        }
242    }
243
244    /// Get a specific agent by ID
245    #[instrument(skip(self))]
246    pub async fn get_agent(&self, id: &str) -> Result<Value, ClientError> {
247        let url = self.agent_url(id);
248        let response = self.request_with_retry(Method::GET, &url, None).await?;
249
250        let raw = response.ok_or_else(|| ClientError::NotFound {
251            kind: "Agent".to_string(),
252            name: id.to_string(),
253        })?;
254        Ok(flatten_agent_response(&raw))
255    }
256
257    /// Create a new agent (creates first version)
258    ///
259    /// Takes a flat agent definition and wraps it in the API format
260    /// before posting to `/agents/{name}/versions`.
261    #[instrument(skip(self, definition))]
262    pub async fn create_agent(&self, definition: &Value) -> Result<Value, ClientError> {
263        let name = definition
264            .get("name")
265            .and_then(|n| n.as_str())
266            .ok_or_else(|| ClientError::Api {
267                status: 400,
268                message: "Agent definition missing 'name' field".to_string(),
269            })?;
270        let payload = wrap_agent_payload(definition);
271        let url = self.agent_versions_url(name);
272        let response = self
273            .request_with_retry(Method::POST, &url, Some(&payload))
274            .await?;
275
276        let raw = response.ok_or_else(|| ClientError::Api {
277            status: 500,
278            message: "No response body from agent creation".to_string(),
279        })?;
280        Ok(flatten_agent_response(&raw))
281    }
282
283    /// Update an existing agent (creates new version)
284    ///
285    /// Takes a flat agent definition and wraps it in the API format
286    /// before posting to `/agents/{name}/versions`.
287    #[instrument(skip(self, definition))]
288    pub async fn update_agent(&self, id: &str, definition: &Value) -> Result<Value, ClientError> {
289        let payload = wrap_agent_payload(definition);
290        let url = self.agent_versions_url(id);
291        let response = self
292            .request_with_retry(Method::POST, &url, Some(&payload))
293            .await?;
294
295        let raw = response.ok_or_else(|| ClientError::Api {
296            status: 500,
297            message: "No response body from agent update".to_string(),
298        })?;
299        Ok(flatten_agent_response(&raw))
300    }
301
302    /// Delete an agent
303    #[instrument(skip(self))]
304    pub async fn delete_agent(&self, id: &str) -> Result<(), ClientError> {
305        let url = self.agent_url(id);
306        self.request_with_retry(Method::DELETE, &url, None).await?;
307        Ok(())
308    }
309
310    /// Get the authentication method being used
311    pub fn auth_method(&self) -> &'static str {
312        self.auth.method_name()
313    }
314}
315
316/// Wrap a flat agent definition into the API request format.
317///
318/// Converts from flat: `{ "name", "model", "instructions", "tools", ... }`
319/// To API format:
320/// ```json
321/// {
322///   "metadata": {...},
323///   "description": "...",
324///   "definition": {
325///     "kind": "prompt",
326///     "model": "...",
327///     "instructions": "...",
328///     "tools": [...]
329///   }
330/// }
331/// ```
332fn wrap_agent_payload(flat: &Value) -> Value {
333    let obj = match flat.as_object() {
334        Some(o) => o,
335        None => return flat.clone(),
336    };
337
338    // Fields that go at the version level (outside definition)
339    const VERSION_LEVEL_FIELDS: &[&str] = &["metadata", "description"];
340
341    // Fields that are response-only and should not be sent
342    const EXCLUDED_FIELDS: &[&str] = &["id", "name", "version", "created_at", "object"];
343
344    let mut wrapper = Map::new();
345    let mut definition = Map::new();
346
347    for (key, value) in obj {
348        if EXCLUDED_FIELDS.contains(&key.as_str()) {
349            continue;
350        } else if VERSION_LEVEL_FIELDS.contains(&key.as_str()) {
351            wrapper.insert(key.clone(), value.clone());
352        } else {
353            definition.insert(key.clone(), value.clone());
354        }
355    }
356
357    // Ensure kind is set (default to "prompt")
358    if !definition.contains_key("kind") {
359        definition.insert("kind".to_string(), Value::String("prompt".to_string()));
360    }
361
362    wrapper.insert("definition".to_string(), Value::Object(definition));
363    Value::Object(wrapper)
364}
365
366/// Flatten a new Foundry agents API response into a flat structure
367/// compatible with the agent decomposition pipeline.
368///
369/// The new Foundry API returns a versioned structure:
370/// ```json
371/// {
372///   "object": "agent",
373///   "id": "MyAgent",
374///   "name": "MyAgent",
375///   "versions": {
376///     "latest": {
377///       "metadata": {...},
378///       "version": "5",
379///       "definition": {
380///         "kind": "prompt",
381///         "model": "gpt-5.2-chat",
382///         "instructions": "...",
383///         "tools": [...]
384///       }
385///     }
386///   }
387/// }
388/// ```
389///
390/// This flattens to: `{ "id", "name", "model", "instructions", "tools", ... }`
391fn flatten_agent_response(agent: &Value) -> Value {
392    let obj = match agent.as_object() {
393        Some(o) => o,
394        None => return agent.clone(),
395    };
396
397    let mut flat = Map::new();
398
399    // Top-level fields
400    if let Some(id) = obj.get("id") {
401        flat.insert("id".to_string(), id.clone());
402    }
403    if let Some(name) = obj.get("name") {
404        flat.insert("name".to_string(), name.clone());
405    }
406
407    // Extract from versions.latest
408    if let Some(latest) = obj
409        .get("versions")
410        .and_then(|v| v.get("latest"))
411        .and_then(|l| l.as_object())
412    {
413        // Version-level fields
414        if let Some(metadata) = latest.get("metadata") {
415            flat.insert("metadata".to_string(), metadata.clone());
416        }
417        if let Some(description) = latest.get("description") {
418            flat.insert("description".to_string(), description.clone());
419        }
420        if let Some(version) = latest.get("version") {
421            flat.insert("version".to_string(), version.clone());
422        }
423        if let Some(created_at) = latest.get("created_at") {
424            flat.insert("created_at".to_string(), created_at.clone());
425        }
426
427        // Definition-level fields (model, instructions, tools, kind, etc.)
428        if let Some(definition) = latest.get("definition").and_then(|d| d.as_object()) {
429            for (key, value) in definition {
430                flat.insert(key.clone(), value.clone());
431            }
432        }
433    }
434
435    // Ensure tools and tool_resources always present (API may omit when empty)
436    flat.entry("tools".to_string())
437        .or_insert_with(|| Value::Array(Vec::new()));
438    flat.entry("tool_resources".to_string())
439        .or_insert_with(|| Value::Object(Map::new()));
440
441    Value::Object(flat)
442}
443
444#[cfg(test)]
445mod tests {
446    use super::*;
447    use crate::auth::{AuthError, AuthProvider};
448    use serde_json::json;
449
450    struct FakeAuth;
451    impl AuthProvider for FakeAuth {
452        fn get_token(&self) -> Result<String, AuthError> {
453            Ok("fake-token".to_string())
454        }
455        fn method_name(&self) -> &'static str {
456            "Fake"
457        }
458    }
459
460    fn make_client() -> FoundryClient {
461        FoundryClient::with_auth(
462            "https://my-ai-svc.services.ai.azure.com".to_string(),
463            "my-project".to_string(),
464            "v1".to_string(),
465            Box::new(FakeAuth),
466        )
467        .unwrap()
468    }
469
470    #[test]
471    fn test_agents_url() {
472        let client = make_client();
473        let url = client.agents_url();
474        assert_eq!(
475            url,
476            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents?api-version=v1"
477        );
478    }
479
480    #[test]
481    fn test_agent_url() {
482        let client = make_client();
483        let url = client.agent_url("Regulus");
484        assert_eq!(
485            url,
486            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/Regulus?api-version=v1"
487        );
488    }
489
490    #[test]
491    fn test_agent_versions_url() {
492        let client = make_client();
493        let url = client.agent_versions_url("KITT");
494        assert_eq!(
495            url,
496            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/KITT/versions?api-version=v1"
497        );
498    }
499
500    #[test]
501    fn test_auth_method() {
502        let client = make_client();
503        assert_eq!(client.auth_method(), "Fake");
504    }
505
506    #[test]
507    fn test_wrap_agent_payload() {
508        let flat = json!({
509            "id": "KITT",
510            "name": "KITT",
511            "model": "gpt-5.2-chat",
512            "kind": "prompt",
513            "instructions": "You are KITT.",
514            "tools": [{"type": "code_interpreter"}],
515            "metadata": {"logo": "kitt.svg"},
516            "description": "A smart car",
517            "version": "3",
518            "created_at": 1234567890
519        });
520
521        let wrapped = wrap_agent_payload(&flat);
522        let obj = wrapped.as_object().unwrap();
523
524        // Top level: metadata, description, definition
525        assert!(obj.contains_key("definition"));
526        assert!(obj.contains_key("metadata"));
527        assert!(obj.contains_key("description"));
528
529        // Excluded from payload
530        assert!(!obj.contains_key("id"));
531        assert!(!obj.contains_key("name"));
532        assert!(!obj.contains_key("version"));
533        assert!(!obj.contains_key("created_at"));
534
535        // Definition should contain model, instructions, tools, kind
536        let def = obj.get("definition").unwrap().as_object().unwrap();
537        assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
538        assert_eq!(def.get("kind").unwrap(), "prompt");
539        assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
540        assert!(def.get("tools").unwrap().as_array().unwrap().len() == 1);
541
542        // Definition should NOT contain excluded or version-level fields
543        assert!(!def.contains_key("id"));
544        assert!(!def.contains_key("name"));
545        assert!(!def.contains_key("metadata"));
546    }
547
548    #[test]
549    fn test_wrap_agent_payload_adds_default_kind() {
550        let flat = json!({
551            "name": "simple",
552            "model": "gpt-4o",
553            "instructions": "Be helpful."
554        });
555
556        let wrapped = wrap_agent_payload(&flat);
557        let def = wrapped.get("definition").unwrap().as_object().unwrap();
558        assert_eq!(def.get("kind").unwrap(), "prompt");
559    }
560
561    #[test]
562    fn test_flatten_then_wrap_roundtrip() {
563        let api_response = json!({
564            "object": "agent",
565            "id": "KITT",
566            "name": "KITT",
567            "versions": {
568                "latest": {
569                    "metadata": {"logo": "kitt.svg"},
570                    "version": "3",
571                    "description": "Smart car",
572                    "created_at": 1234567890,
573                    "definition": {
574                        "kind": "prompt",
575                        "model": "gpt-5.2-chat",
576                        "instructions": "You are KITT.",
577                        "tools": [{"type": "code_interpreter"}]
578                    }
579                }
580            }
581        });
582
583        let flat = flatten_agent_response(&api_response);
584        let wrapped = wrap_agent_payload(&flat);
585
586        // The wrapped payload should have a definition with the same content
587        let def = wrapped.get("definition").unwrap().as_object().unwrap();
588        assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
589        assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
590        assert_eq!(def.get("kind").unwrap(), "prompt");
591    }
592
593    #[test]
594    fn test_flatten_agent_response_full() {
595        let api_response = json!({
596            "object": "agent",
597            "id": "Regulus",
598            "name": "Regulus",
599            "versions": {
600                "latest": {
601                    "metadata": {
602                        "logo": "Avatar_Default.svg",
603                        "description": "",
604                        "modified_at": "1769974547"
605                    },
606                    "object": "agent.version",
607                    "id": "Regulus:5",
608                    "name": "Regulus",
609                    "version": "5",
610                    "description": "",
611                    "created_at": 1769974549,
612                    "definition": {
613                        "kind": "prompt",
614                        "model": "gpt-5.2-chat",
615                        "instructions": "You are Regulus.",
616                        "tools": [
617                            {"type": "mcp", "server_label": "kb_test"}
618                        ]
619                    }
620                }
621            }
622        });
623
624        let flat = flatten_agent_response(&api_response);
625        let obj = flat.as_object().unwrap();
626
627        assert_eq!(obj.get("id").unwrap(), "Regulus");
628        assert_eq!(obj.get("name").unwrap(), "Regulus");
629        assert_eq!(obj.get("model").unwrap(), "gpt-5.2-chat");
630        assert_eq!(obj.get("kind").unwrap(), "prompt");
631        assert_eq!(obj.get("instructions").unwrap(), "You are Regulus.");
632        assert_eq!(obj.get("version").unwrap(), "5");
633        assert_eq!(obj.get("description").unwrap(), "");
634        assert!(obj.get("metadata").is_some());
635        assert!(obj.get("tools").unwrap().as_array().unwrap().len() == 1);
636
637        // Should NOT have the nested versions structure
638        assert!(!obj.contains_key("versions"));
639        assert!(!obj.contains_key("object"));
640    }
641
642    #[test]
643    fn test_flatten_agent_response_minimal() {
644        let api_response = json!({
645            "object": "agent",
646            "id": "simple",
647            "name": "simple"
648        });
649
650        let flat = flatten_agent_response(&api_response);
651        let obj = flat.as_object().unwrap();
652
653        assert_eq!(obj.get("id").unwrap(), "simple");
654        assert_eq!(obj.get("name").unwrap(), "simple");
655        assert!(!obj.contains_key("model"));
656        // tools and tool_resources always present with defaults
657        assert_eq!(obj.get("tools").unwrap(), &json!([]));
658        assert_eq!(obj.get("tool_resources").unwrap(), &json!({}));
659    }
660
661    #[test]
662    fn test_flatten_agent_response_non_object() {
663        let flat = flatten_agent_response(&json!("not an object"));
664        assert_eq!(flat, json!("not an object"));
665    }
666
667    #[test]
668    fn test_wrap_flatten_roundtrip_preserves_tool_permissions() {
669        let flat = json!({
670            "name": "test-agent",
671            "kind": "prompt",
672            "model": "gpt-4o",
673            "tools": [{
674                "type": "mcp",
675                "server_label": "kb_test",
676                "require_approval": "never",
677                "allowed_tools": ["tool_a", "tool_b"]
678            }]
679        });
680
681        // Wrap for API submission
682        let wrapped = wrap_agent_payload(&flat);
683        let def_tools = wrapped["definition"]["tools"].as_array().unwrap();
684        assert_eq!(def_tools[0]["require_approval"], "never");
685        assert_eq!(def_tools[0]["allowed_tools"][0], "tool_a");
686        assert_eq!(def_tools[0]["allowed_tools"][1], "tool_b");
687
688        // Simulate API response containing the same tools
689        let api_response = json!({
690            "object": "agent",
691            "id": "test-agent",
692            "name": "test-agent",
693            "versions": {
694                "latest": {
695                    "version": "1",
696                    "definition": {
697                        "kind": "prompt",
698                        "model": "gpt-4o",
699                        "tools": [{
700                            "type": "mcp",
701                            "server_label": "kb_test",
702                            "require_approval": "never",
703                            "allowed_tools": ["tool_a", "tool_b"]
704                        }]
705                    }
706                }
707            }
708        });
709
710        let flattened = flatten_agent_response(&api_response);
711        let tools = flattened["tools"].as_array().unwrap();
712        assert_eq!(tools[0]["require_approval"], "never");
713        let allowed = tools[0]["allowed_tools"].as_array().unwrap();
714        assert_eq!(allowed.len(), 2);
715        assert_eq!(allowed[0], "tool_a");
716        assert_eq!(allowed[1], "tool_b");
717    }
718
719    #[test]
720    fn test_flatten_preserves_require_approval_object_form() {
721        let api_response = json!({
722            "object": "agent",
723            "id": "granular-agent",
724            "name": "granular-agent",
725            "versions": {
726                "latest": {
727                    "version": "1",
728                    "definition": {
729                        "kind": "prompt",
730                        "model": "gpt-4o",
731                        "tools": [{
732                            "type": "mcp",
733                            "server_label": "kb_test",
734                            "require_approval": {
735                                "never": {"tool_names": ["safe_tool"]},
736                                "always": {"tool_names": ["dangerous_tool"]}
737                            }
738                        }]
739                    }
740                }
741            }
742        });
743
744        let flat = flatten_agent_response(&api_response);
745        let ra = &flat["tools"][0]["require_approval"];
746        assert_eq!(ra["never"]["tool_names"][0], "safe_tool");
747        assert_eq!(ra["always"]["tool_names"][0], "dangerous_tool");
748
749        // Round-trip: flatten → wrap → verify definition still has structured form
750        let re_wrapped = wrap_agent_payload(&flat);
751        let def_ra = &re_wrapped["definition"]["tools"][0]["require_approval"];
752        assert_eq!(def_ra["never"]["tool_names"][0], "safe_tool");
753        assert_eq!(def_ra["always"]["tool_names"][0], "dangerous_tool");
754    }
755}
rigg_client/foundry.rs

rigg_client/
foundry.rs
