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

rigg_client/
foundry.rs
