Skip to main content

hoist_client/
foundry.rs

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