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
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    // Version-create responses put the definition at the top level
408    // (`{name, version, definition: {...}}`) instead of `versions.latest`.
409    let top_level = obj
410        .get("definition")
411        .is_some()
412        .then_some(obj)
413        .map(|o| o as &Map<String, Value>);
414
415    // Extract from versions.latest (GET/list shape) or the top level (POST shape)
416    if let Some(latest) = obj
417        .get("versions")
418        .and_then(|v| v.get("latest"))
419        .and_then(|l| l.as_object())
420        .or(top_level)
421    {
422        // Version-level fields
423        if let Some(metadata) = latest.get("metadata") {
424            flat.insert("metadata".to_string(), metadata.clone());
425        }
426        if let Some(description) = latest.get("description") {
427            flat.insert("description".to_string(), description.clone());
428        }
429        if let Some(version) = latest.get("version") {
430            flat.insert("version".to_string(), version.clone());
431        }
432        if let Some(created_at) = latest.get("created_at") {
433            flat.insert("created_at".to_string(), created_at.clone());
434        }
435
436        // Definition-level fields (model, instructions, tools, kind, etc.)
437        if let Some(definition) = latest.get("definition").and_then(|d| d.as_object()) {
438            for (key, value) in definition {
439                flat.insert(key.clone(), value.clone());
440            }
441        }
442    }
443
444    // Ensure tools and tool_resources always present (API may omit when empty)
445    flat.entry("tools".to_string())
446        .or_insert_with(|| Value::Array(Vec::new()));
447    flat.entry("tool_resources".to_string())
448        .or_insert_with(|| Value::Object(Map::new()));
449
450    Value::Object(flat)
451}
452
453#[cfg(test)]
454mod tests {
455    use super::*;
456    use crate::auth::{AuthError, AuthProvider};
457    use serde_json::json;
458
459    struct FakeAuth;
460    impl AuthProvider for FakeAuth {
461        fn get_token(&self) -> Result<String, AuthError> {
462            Ok("fake-token".to_string())
463        }
464        fn method_name(&self) -> &'static str {
465            "Fake"
466        }
467    }
468
469    fn make_client() -> FoundryClient {
470        FoundryClient::with_auth(
471            "https://my-ai-svc.services.ai.azure.com".to_string(),
472            "my-project".to_string(),
473            "v1".to_string(),
474            Box::new(FakeAuth),
475        )
476        .unwrap()
477    }
478
479    #[test]
480    fn test_agents_url() {
481        let client = make_client();
482        let url = client.agents_url();
483        assert_eq!(
484            url,
485            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents?api-version=v1"
486        );
487    }
488
489    #[test]
490    fn test_agent_url() {
491        let client = make_client();
492        let url = client.agent_url("Regulus");
493        assert_eq!(
494            url,
495            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/Regulus?api-version=v1"
496        );
497    }
498
499    #[test]
500    fn test_agent_versions_url() {
501        let client = make_client();
502        let url = client.agent_versions_url("KITT");
503        assert_eq!(
504            url,
505            "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/KITT/versions?api-version=v1"
506        );
507    }
508
509    #[test]
510    fn test_auth_method() {
511        let client = make_client();
512        assert_eq!(client.auth_method(), "Fake");
513    }
514
515    #[test]
516    fn flatten_handles_version_create_response_shape() {
517        // POST /agents/{n}/versions returns the definition at the top level
518        let response = json!({
519            "object": "agent.version",
520            "name": "helper",
521            "version": "1",
522            "definition": {
523                "kind": "prompt",
524                "model": "gpt-4.1-mini",
525                "instructions": "Be nice.",
526                "tools": [{"type": "mcp"}]
527            }
528        });
529        let flat = flatten_agent_response(&response);
530        assert_eq!(flat["name"], "helper");
531        assert_eq!(flat["model"], "gpt-4.1-mini");
532        assert_eq!(flat["instructions"], "Be nice.");
533        assert_eq!(flat["tools"][0]["type"], "mcp");
534    }
535
536    #[test]
537    fn test_wrap_agent_payload() {
538        let flat = json!({
539            "id": "KITT",
540            "name": "KITT",
541            "model": "gpt-5.2-chat",
542            "kind": "prompt",
543            "instructions": "You are KITT.",
544            "tools": [{"type": "code_interpreter"}],
545            "metadata": {"logo": "kitt.svg"},
546            "description": "A smart car",
547            "version": "3",
548            "created_at": 1234567890
549        });
550
551        let wrapped = wrap_agent_payload(&flat);
552        let obj = wrapped.as_object().unwrap();
553
554        // Top level: metadata, description, definition
555        assert!(obj.contains_key("definition"));
556        assert!(obj.contains_key("metadata"));
557        assert!(obj.contains_key("description"));
558
559        // Excluded from payload
560        assert!(!obj.contains_key("id"));
561        assert!(!obj.contains_key("name"));
562        assert!(!obj.contains_key("version"));
563        assert!(!obj.contains_key("created_at"));
564
565        // Definition should contain model, instructions, tools, kind
566        let def = obj.get("definition").unwrap().as_object().unwrap();
567        assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
568        assert_eq!(def.get("kind").unwrap(), "prompt");
569        assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
570        assert!(def.get("tools").unwrap().as_array().unwrap().len() == 1);
571
572        // Definition should NOT contain excluded or version-level fields
573        assert!(!def.contains_key("id"));
574        assert!(!def.contains_key("name"));
575        assert!(!def.contains_key("metadata"));
576    }
577
578    #[test]
579    fn test_wrap_agent_payload_adds_default_kind() {
580        let flat = json!({
581            "name": "simple",
582            "model": "gpt-4o",
583            "instructions": "Be helpful."
584        });
585
586        let wrapped = wrap_agent_payload(&flat);
587        let def = wrapped.get("definition").unwrap().as_object().unwrap();
588        assert_eq!(def.get("kind").unwrap(), "prompt");
589    }
590
591    #[test]
592    fn test_flatten_then_wrap_roundtrip() {
593        let api_response = json!({
594            "object": "agent",
595            "id": "KITT",
596            "name": "KITT",
597            "versions": {
598                "latest": {
599                    "metadata": {"logo": "kitt.svg"},
600                    "version": "3",
601                    "description": "Smart car",
602                    "created_at": 1234567890,
603                    "definition": {
604                        "kind": "prompt",
605                        "model": "gpt-5.2-chat",
606                        "instructions": "You are KITT.",
607                        "tools": [{"type": "code_interpreter"}]
608                    }
609                }
610            }
611        });
612
613        let flat = flatten_agent_response(&api_response);
614        let wrapped = wrap_agent_payload(&flat);
615
616        // The wrapped payload should have a definition with the same content
617        let def = wrapped.get("definition").unwrap().as_object().unwrap();
618        assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
619        assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
620        assert_eq!(def.get("kind").unwrap(), "prompt");
621    }
622
623    #[test]
624    fn test_flatten_agent_response_full() {
625        let api_response = json!({
626            "object": "agent",
627            "id": "Regulus",
628            "name": "Regulus",
629            "versions": {
630                "latest": {
631                    "metadata": {
632                        "logo": "Avatar_Default.svg",
633                        "description": "",
634                        "modified_at": "1769974547"
635                    },
636                    "object": "agent.version",
637                    "id": "Regulus:5",
638                    "name": "Regulus",
639                    "version": "5",
640                    "description": "",
641                    "created_at": 1769974549,
642                    "definition": {
643                        "kind": "prompt",
644                        "model": "gpt-5.2-chat",
645                        "instructions": "You are Regulus.",
646                        "tools": [
647                            {"type": "mcp", "server_label": "kb_test"}
648                        ]
649                    }
650                }
651            }
652        });
653
654        let flat = flatten_agent_response(&api_response);
655        let obj = flat.as_object().unwrap();
656
657        assert_eq!(obj.get("id").unwrap(), "Regulus");
658        assert_eq!(obj.get("name").unwrap(), "Regulus");
659        assert_eq!(obj.get("model").unwrap(), "gpt-5.2-chat");
660        assert_eq!(obj.get("kind").unwrap(), "prompt");
661        assert_eq!(obj.get("instructions").unwrap(), "You are Regulus.");
662        assert_eq!(obj.get("version").unwrap(), "5");
663        assert_eq!(obj.get("description").unwrap(), "");
664        assert!(obj.get("metadata").is_some());
665        assert!(obj.get("tools").unwrap().as_array().unwrap().len() == 1);
666
667        // Should NOT have the nested versions structure
668        assert!(!obj.contains_key("versions"));
669        assert!(!obj.contains_key("object"));
670    }
671
672    #[test]
673    fn test_flatten_agent_response_minimal() {
674        let api_response = json!({
675            "object": "agent",
676            "id": "simple",
677            "name": "simple"
678        });
679
680        let flat = flatten_agent_response(&api_response);
681        let obj = flat.as_object().unwrap();
682
683        assert_eq!(obj.get("id").unwrap(), "simple");
684        assert_eq!(obj.get("name").unwrap(), "simple");
685        assert!(!obj.contains_key("model"));
686        // tools and tool_resources always present with defaults
687        assert_eq!(obj.get("tools").unwrap(), &json!([]));
688        assert_eq!(obj.get("tool_resources").unwrap(), &json!({}));
689    }
690
691    #[test]
692    fn test_flatten_agent_response_non_object() {
693        let flat = flatten_agent_response(&json!("not an object"));
694        assert_eq!(flat, json!("not an object"));
695    }
696
697    #[test]
698    fn test_wrap_flatten_roundtrip_preserves_tool_permissions() {
699        let flat = json!({
700            "name": "test-agent",
701            "kind": "prompt",
702            "model": "gpt-4o",
703            "tools": [{
704                "type": "mcp",
705                "server_label": "kb_test",
706                "require_approval": "never",
707                "allowed_tools": ["tool_a", "tool_b"]
708            }]
709        });
710
711        // Wrap for API submission
712        let wrapped = wrap_agent_payload(&flat);
713        let def_tools = wrapped["definition"]["tools"].as_array().unwrap();
714        assert_eq!(def_tools[0]["require_approval"], "never");
715        assert_eq!(def_tools[0]["allowed_tools"][0], "tool_a");
716        assert_eq!(def_tools[0]["allowed_tools"][1], "tool_b");
717
718        // Simulate API response containing the same tools
719        let api_response = json!({
720            "object": "agent",
721            "id": "test-agent",
722            "name": "test-agent",
723            "versions": {
724                "latest": {
725                    "version": "1",
726                    "definition": {
727                        "kind": "prompt",
728                        "model": "gpt-4o",
729                        "tools": [{
730                            "type": "mcp",
731                            "server_label": "kb_test",
732                            "require_approval": "never",
733                            "allowed_tools": ["tool_a", "tool_b"]
734                        }]
735                    }
736                }
737            }
738        });
739
740        let flattened = flatten_agent_response(&api_response);
741        let tools = flattened["tools"].as_array().unwrap();
742        assert_eq!(tools[0]["require_approval"], "never");
743        let allowed = tools[0]["allowed_tools"].as_array().unwrap();
744        assert_eq!(allowed.len(), 2);
745        assert_eq!(allowed[0], "tool_a");
746        assert_eq!(allowed[1], "tool_b");
747    }
748
749    #[test]
750    fn test_flatten_preserves_require_approval_object_form() {
751        let api_response = json!({
752            "object": "agent",
753            "id": "granular-agent",
754            "name": "granular-agent",
755            "versions": {
756                "latest": {
757                    "version": "1",
758                    "definition": {
759                        "kind": "prompt",
760                        "model": "gpt-4o",
761                        "tools": [{
762                            "type": "mcp",
763                            "server_label": "kb_test",
764                            "require_approval": {
765                                "never": {"tool_names": ["safe_tool"]},
766                                "always": {"tool_names": ["dangerous_tool"]}
767                            }
768                        }]
769                    }
770                }
771            }
772        });
773
774        let flat = flatten_agent_response(&api_response);
775        let ra = &flat["tools"][0]["require_approval"];
776        assert_eq!(ra["never"]["tool_names"][0], "safe_tool");
777        assert_eq!(ra["always"]["tool_names"][0], "dangerous_tool");
778
779        // Round-trip: flatten → wrap → verify definition still has structured form
780        let re_wrapped = wrap_agent_payload(&flat);
781        let def_ra = &re_wrapped["definition"]["tools"][0]["require_approval"];
782        assert_eq!(def_ra["never"]["tool_names"][0], "safe_tool");
783        assert_eq!(def_ra["always"]["tool_names"][0], "dangerous_tool");
784    }
785}