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.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 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 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 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 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 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 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 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 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 #[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 Ok(items.iter().map(flatten_agent_response).collect())
235 }
236 None => Ok(Vec::new()),
237 }
238 }
239
240 #[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 #[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 #[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 #[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 pub fn auth_method(&self) -> &'static str {
308 self.auth.method_name()
309 }
310}
311
312fn 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 const VERSION_LEVEL_FIELDS: &[&str] = &["metadata", "description"];
336
337 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 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
362fn 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 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 let top_level = obj
406 .get("definition")
407 .is_some()
408 .then_some(obj)
409 .map(|o| o as &Map<String, Value>);
410
411 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 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 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 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 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 assert!(obj.contains_key("definition"));
552 assert!(obj.contains_key("metadata"));
553 assert!(obj.contains_key("description"));
554
555 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 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 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 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 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 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 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 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 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}