1use 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
17const MAX_RETRIES: u32 = 3;
19
20const INITIAL_BACKOFF_SECS: u64 = 1;
22
23pub struct FoundryClient {
25 http: Client,
26 auth: Box<dyn AuthProvider>,
27 base_url: String,
28 project: String,
29 api_version: String,
30 features: Vec<String>,
33}
34
35impl FoundryClient {
36 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 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 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 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 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 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 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 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 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 #[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 Ok(items.iter().map(flatten_agent_response).collect())
239 }
240 None => Ok(Vec::new()),
241 }
242 }
243
244 #[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 #[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 #[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 #[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 pub fn auth_method(&self) -> &'static str {
312 self.auth.method_name()
313 }
314}
315
316fn 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 const VERSION_LEVEL_FIELDS: &[&str] = &["metadata", "description"];
340
341 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 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
366fn 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 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 let top_level = obj
410 .get("definition")
411 .is_some()
412 .then_some(obj)
413 .map(|o| o as &Map<String, Value>);
414
415 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 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 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 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 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 assert!(obj.contains_key("definition"));
556 assert!(obj.contains_key("metadata"));
557 assert!(obj.contains_key("description"));
558
559 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 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 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 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 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 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 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 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 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}