Skip to main content

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