1use std::time::Duration;
7
8use reqwest::{Client, Method, StatusCode};
9use serde_json::{Map, Value};
10use tracing::{debug, instrument, warn};
11
12use hoist_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}
31
32impl FoundryClient {
33 pub fn new(config: &FoundryServiceConfig) -> Result<Self, ClientError> {
35 let auth = get_auth_provider_for(hoist_core::ServiceDomain::Foundry)?;
36 let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
37
38 Ok(Self {
39 http,
40 auth,
41 base_url: config.service_url(),
42 project: config.project.clone(),
43 api_version: config.api_version.clone(),
44 })
45 }
46
47 pub fn with_auth(
49 base_url: String,
50 project: String,
51 api_version: String,
52 auth: Box<dyn AuthProvider>,
53 ) -> Result<Self, ClientError> {
54 let http = Client::builder().timeout(Duration::from_secs(30)).build()?;
55
56 Ok(Self {
57 http,
58 auth,
59 base_url,
60 project,
61 api_version,
62 })
63 }
64
65 fn agents_url(&self) -> String {
67 format!(
68 "{}/api/projects/{}/agents?api-version={}",
69 self.base_url, self.project, self.api_version
70 )
71 }
72
73 fn agent_url(&self, id: &str) -> String {
75 format!(
76 "{}/api/projects/{}/agents/{}?api-version={}",
77 self.base_url,
78 self.project,
79 urlencoding::encode(id),
80 self.api_version
81 )
82 }
83
84 fn agent_versions_url(&self, name: &str) -> String {
86 format!(
87 "{}/api/projects/{}/agents/{}/versions?api-version={}",
88 self.base_url,
89 self.project,
90 urlencoding::encode(name),
91 self.api_version
92 )
93 }
94
95 async fn request(
97 &self,
98 method: Method,
99 url: &str,
100 body: Option<&Value>,
101 ) -> Result<Option<Value>, ClientError> {
102 let token = self.auth.get_token()?;
103
104 let mut request = self
105 .http
106 .request(method.clone(), url)
107 .header("Authorization", format!("Bearer {}", token))
108 .header("Content-Type", "application/json");
109
110 if let Some(json) = body {
111 request = request.json(json);
112 }
113
114 debug!("Request: {} {}", method, url);
115 let response = request.send().await?;
116 let status = response.status();
117
118 if status == StatusCode::NO_CONTENT {
119 return Ok(None);
120 }
121
122 let body = response.text().await?;
123
124 if status.is_success() {
125 if body.is_empty() {
126 Ok(None)
127 } else {
128 let value: Value = serde_json::from_str(&body)?;
129 Ok(Some(value))
130 }
131 } else {
132 match status {
133 StatusCode::NOT_FOUND => Err(ClientError::NotFound {
134 kind: "agent".to_string(),
135 name: url.to_string(),
136 }),
137 StatusCode::TOO_MANY_REQUESTS => {
138 let retry_after = 60;
139 Err(ClientError::RateLimited { retry_after })
140 }
141 StatusCode::SERVICE_UNAVAILABLE => Err(ClientError::ServiceUnavailable(body)),
142 _ => Err(ClientError::from_response_with_url(
143 status.as_u16(),
144 &body,
145 Some(url),
146 )),
147 }
148 }
149 }
150
151 async fn request_with_retry(
153 &self,
154 method: Method,
155 url: &str,
156 body: Option<&Value>,
157 ) -> Result<Option<Value>, ClientError> {
158 let mut attempt = 0u32;
159 loop {
160 match self.request(method.clone(), url, body).await {
161 Ok(value) => return Ok(value),
162 Err(err) if err.is_retryable() && attempt < MAX_RETRIES => {
163 let delay = match &err {
164 ClientError::RateLimited { retry_after } => {
165 Duration::from_secs(*retry_after)
166 }
167 _ => Duration::from_secs(INITIAL_BACKOFF_SECS * 2u64.pow(attempt)),
168 };
169 warn!(
170 "Request {} {} failed (attempt {}/{}): {}. Retrying in {:?}",
171 method,
172 url,
173 attempt + 1,
174 MAX_RETRIES + 1,
175 err,
176 delay,
177 );
178 tokio::time::sleep(delay).await;
179 attempt += 1;
180 }
181 Err(err) => return Err(err),
182 }
183 }
184 }
185
186 #[instrument(skip(self))]
188 pub async fn list_agents(&self) -> Result<Vec<Value>, ClientError> {
189 let url = self.agents_url();
190 let response = self.request_with_retry(Method::GET, &url, None).await?;
191
192 match response {
193 Some(value) => {
194 let items = value
195 .get("data")
196 .and_then(|v| v.as_array())
197 .cloned()
198 .unwrap_or_default();
199 Ok(items.iter().map(flatten_agent_response).collect())
201 }
202 None => Ok(Vec::new()),
203 }
204 }
205
206 #[instrument(skip(self))]
208 pub async fn get_agent(&self, id: &str) -> Result<Value, ClientError> {
209 let url = self.agent_url(id);
210 let response = self.request_with_retry(Method::GET, &url, None).await?;
211
212 let raw = response.ok_or_else(|| ClientError::NotFound {
213 kind: "Agent".to_string(),
214 name: id.to_string(),
215 })?;
216 Ok(flatten_agent_response(&raw))
217 }
218
219 #[instrument(skip(self, definition))]
224 pub async fn create_agent(&self, definition: &Value) -> Result<Value, ClientError> {
225 let name = definition
226 .get("name")
227 .and_then(|n| n.as_str())
228 .ok_or_else(|| ClientError::Api {
229 status: 400,
230 message: "Agent definition missing 'name' field".to_string(),
231 })?;
232 let payload = wrap_agent_payload(definition);
233 let url = self.agent_versions_url(name);
234 let response = self
235 .request_with_retry(Method::POST, &url, Some(&payload))
236 .await?;
237
238 let raw = response.ok_or_else(|| ClientError::Api {
239 status: 500,
240 message: "No response body from agent creation".to_string(),
241 })?;
242 Ok(flatten_agent_response(&raw))
243 }
244
245 #[instrument(skip(self, definition))]
250 pub async fn update_agent(&self, id: &str, definition: &Value) -> Result<Value, ClientError> {
251 let payload = wrap_agent_payload(definition);
252 let url = self.agent_versions_url(id);
253 let response = self
254 .request_with_retry(Method::POST, &url, Some(&payload))
255 .await?;
256
257 let raw = response.ok_or_else(|| ClientError::Api {
258 status: 500,
259 message: "No response body from agent update".to_string(),
260 })?;
261 Ok(flatten_agent_response(&raw))
262 }
263
264 #[instrument(skip(self))]
266 pub async fn delete_agent(&self, id: &str) -> Result<(), ClientError> {
267 let url = self.agent_url(id);
268 self.request_with_retry(Method::DELETE, &url, None).await?;
269 Ok(())
270 }
271
272 pub fn auth_method(&self) -> &'static str {
274 self.auth.method_name()
275 }
276}
277
278fn wrap_agent_payload(flat: &Value) -> Value {
295 let obj = match flat.as_object() {
296 Some(o) => o,
297 None => return flat.clone(),
298 };
299
300 const VERSION_LEVEL_FIELDS: &[&str] = &["metadata", "description"];
302
303 const EXCLUDED_FIELDS: &[&str] = &["id", "name", "version", "created_at", "object"];
305
306 let mut wrapper = Map::new();
307 let mut definition = Map::new();
308
309 for (key, value) in obj {
310 if EXCLUDED_FIELDS.contains(&key.as_str()) {
311 continue;
312 } else if VERSION_LEVEL_FIELDS.contains(&key.as_str()) {
313 wrapper.insert(key.clone(), value.clone());
314 } else {
315 definition.insert(key.clone(), value.clone());
316 }
317 }
318
319 if !definition.contains_key("kind") {
321 definition.insert("kind".to_string(), Value::String("prompt".to_string()));
322 }
323
324 wrapper.insert("definition".to_string(), Value::Object(definition));
325 Value::Object(wrapper)
326}
327
328fn flatten_agent_response(agent: &Value) -> Value {
354 let obj = match agent.as_object() {
355 Some(o) => o,
356 None => return agent.clone(),
357 };
358
359 let mut flat = Map::new();
360
361 if let Some(id) = obj.get("id") {
363 flat.insert("id".to_string(), id.clone());
364 }
365 if let Some(name) = obj.get("name") {
366 flat.insert("name".to_string(), name.clone());
367 }
368
369 if let Some(latest) = obj
371 .get("versions")
372 .and_then(|v| v.get("latest"))
373 .and_then(|l| l.as_object())
374 {
375 if let Some(metadata) = latest.get("metadata") {
377 flat.insert("metadata".to_string(), metadata.clone());
378 }
379 if let Some(description) = latest.get("description") {
380 flat.insert("description".to_string(), description.clone());
381 }
382 if let Some(version) = latest.get("version") {
383 flat.insert("version".to_string(), version.clone());
384 }
385 if let Some(created_at) = latest.get("created_at") {
386 flat.insert("created_at".to_string(), created_at.clone());
387 }
388
389 if let Some(definition) = latest.get("definition").and_then(|d| d.as_object()) {
391 for (key, value) in definition {
392 flat.insert(key.clone(), value.clone());
393 }
394 }
395 }
396
397 flat.entry("tools".to_string())
399 .or_insert_with(|| Value::Array(Vec::new()));
400 flat.entry("tool_resources".to_string())
401 .or_insert_with(|| Value::Object(Map::new()));
402
403 Value::Object(flat)
404}
405
406#[cfg(test)]
407mod tests {
408 use super::*;
409 use crate::auth::{AuthError, AuthProvider};
410 use serde_json::json;
411
412 struct FakeAuth;
413 impl AuthProvider for FakeAuth {
414 fn get_token(&self) -> Result<String, AuthError> {
415 Ok("fake-token".to_string())
416 }
417 fn method_name(&self) -> &'static str {
418 "Fake"
419 }
420 }
421
422 fn make_client() -> FoundryClient {
423 FoundryClient::with_auth(
424 "https://my-ai-svc.services.ai.azure.com".to_string(),
425 "my-project".to_string(),
426 "2025-05-15-preview".to_string(),
427 Box::new(FakeAuth),
428 )
429 .unwrap()
430 }
431
432 #[test]
433 fn test_agents_url() {
434 let client = make_client();
435 let url = client.agents_url();
436 assert_eq!(
437 url,
438 "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents?api-version=2025-05-15-preview"
439 );
440 }
441
442 #[test]
443 fn test_agent_url() {
444 let client = make_client();
445 let url = client.agent_url("Regulus");
446 assert_eq!(
447 url,
448 "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/Regulus?api-version=2025-05-15-preview"
449 );
450 }
451
452 #[test]
453 fn test_agent_versions_url() {
454 let client = make_client();
455 let url = client.agent_versions_url("KITT");
456 assert_eq!(
457 url,
458 "https://my-ai-svc.services.ai.azure.com/api/projects/my-project/agents/KITT/versions?api-version=2025-05-15-preview"
459 );
460 }
461
462 #[test]
463 fn test_auth_method() {
464 let client = make_client();
465 assert_eq!(client.auth_method(), "Fake");
466 }
467
468 #[test]
469 fn test_wrap_agent_payload() {
470 let flat = json!({
471 "id": "KITT",
472 "name": "KITT",
473 "model": "gpt-5.2-chat",
474 "kind": "prompt",
475 "instructions": "You are KITT.",
476 "tools": [{"type": "code_interpreter"}],
477 "metadata": {"logo": "kitt.svg"},
478 "description": "A smart car",
479 "version": "3",
480 "created_at": 1234567890
481 });
482
483 let wrapped = wrap_agent_payload(&flat);
484 let obj = wrapped.as_object().unwrap();
485
486 assert!(obj.contains_key("definition"));
488 assert!(obj.contains_key("metadata"));
489 assert!(obj.contains_key("description"));
490
491 assert!(!obj.contains_key("id"));
493 assert!(!obj.contains_key("name"));
494 assert!(!obj.contains_key("version"));
495 assert!(!obj.contains_key("created_at"));
496
497 let def = obj.get("definition").unwrap().as_object().unwrap();
499 assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
500 assert_eq!(def.get("kind").unwrap(), "prompt");
501 assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
502 assert!(def.get("tools").unwrap().as_array().unwrap().len() == 1);
503
504 assert!(!def.contains_key("id"));
506 assert!(!def.contains_key("name"));
507 assert!(!def.contains_key("metadata"));
508 }
509
510 #[test]
511 fn test_wrap_agent_payload_adds_default_kind() {
512 let flat = json!({
513 "name": "simple",
514 "model": "gpt-4o",
515 "instructions": "Be helpful."
516 });
517
518 let wrapped = wrap_agent_payload(&flat);
519 let def = wrapped.get("definition").unwrap().as_object().unwrap();
520 assert_eq!(def.get("kind").unwrap(), "prompt");
521 }
522
523 #[test]
524 fn test_flatten_then_wrap_roundtrip() {
525 let api_response = json!({
526 "object": "agent",
527 "id": "KITT",
528 "name": "KITT",
529 "versions": {
530 "latest": {
531 "metadata": {"logo": "kitt.svg"},
532 "version": "3",
533 "description": "Smart car",
534 "created_at": 1234567890,
535 "definition": {
536 "kind": "prompt",
537 "model": "gpt-5.2-chat",
538 "instructions": "You are KITT.",
539 "tools": [{"type": "code_interpreter"}]
540 }
541 }
542 }
543 });
544
545 let flat = flatten_agent_response(&api_response);
546 let wrapped = wrap_agent_payload(&flat);
547
548 let def = wrapped.get("definition").unwrap().as_object().unwrap();
550 assert_eq!(def.get("model").unwrap(), "gpt-5.2-chat");
551 assert_eq!(def.get("instructions").unwrap(), "You are KITT.");
552 assert_eq!(def.get("kind").unwrap(), "prompt");
553 }
554
555 #[test]
556 fn test_flatten_agent_response_full() {
557 let api_response = json!({
558 "object": "agent",
559 "id": "Regulus",
560 "name": "Regulus",
561 "versions": {
562 "latest": {
563 "metadata": {
564 "logo": "Avatar_Default.svg",
565 "description": "",
566 "modified_at": "1769974547"
567 },
568 "object": "agent.version",
569 "id": "Regulus:5",
570 "name": "Regulus",
571 "version": "5",
572 "description": "",
573 "created_at": 1769974549,
574 "definition": {
575 "kind": "prompt",
576 "model": "gpt-5.2-chat",
577 "instructions": "You are Regulus.",
578 "tools": [
579 {"type": "mcp", "server_label": "kb_test"}
580 ]
581 }
582 }
583 }
584 });
585
586 let flat = flatten_agent_response(&api_response);
587 let obj = flat.as_object().unwrap();
588
589 assert_eq!(obj.get("id").unwrap(), "Regulus");
590 assert_eq!(obj.get("name").unwrap(), "Regulus");
591 assert_eq!(obj.get("model").unwrap(), "gpt-5.2-chat");
592 assert_eq!(obj.get("kind").unwrap(), "prompt");
593 assert_eq!(obj.get("instructions").unwrap(), "You are Regulus.");
594 assert_eq!(obj.get("version").unwrap(), "5");
595 assert_eq!(obj.get("description").unwrap(), "");
596 assert!(obj.get("metadata").is_some());
597 assert!(obj.get("tools").unwrap().as_array().unwrap().len() == 1);
598
599 assert!(!obj.contains_key("versions"));
601 assert!(!obj.contains_key("object"));
602 }
603
604 #[test]
605 fn test_flatten_agent_response_minimal() {
606 let api_response = json!({
607 "object": "agent",
608 "id": "simple",
609 "name": "simple"
610 });
611
612 let flat = flatten_agent_response(&api_response);
613 let obj = flat.as_object().unwrap();
614
615 assert_eq!(obj.get("id").unwrap(), "simple");
616 assert_eq!(obj.get("name").unwrap(), "simple");
617 assert!(!obj.contains_key("model"));
618 assert_eq!(obj.get("tools").unwrap(), &json!([]));
620 assert_eq!(obj.get("tool_resources").unwrap(), &json!({}));
621 }
622
623 #[test]
624 fn test_flatten_agent_response_non_object() {
625 let flat = flatten_agent_response(&json!("not an object"));
626 assert_eq!(flat, json!("not an object"));
627 }
628
629 #[test]
630 fn test_wrap_flatten_roundtrip_preserves_tool_permissions() {
631 let flat = json!({
632 "name": "test-agent",
633 "kind": "prompt",
634 "model": "gpt-4o",
635 "tools": [{
636 "type": "mcp",
637 "server_label": "kb_test",
638 "require_approval": "never",
639 "allowed_tools": ["tool_a", "tool_b"]
640 }]
641 });
642
643 let wrapped = wrap_agent_payload(&flat);
645 let def_tools = wrapped["definition"]["tools"].as_array().unwrap();
646 assert_eq!(def_tools[0]["require_approval"], "never");
647 assert_eq!(def_tools[0]["allowed_tools"][0], "tool_a");
648 assert_eq!(def_tools[0]["allowed_tools"][1], "tool_b");
649
650 let api_response = json!({
652 "object": "agent",
653 "id": "test-agent",
654 "name": "test-agent",
655 "versions": {
656 "latest": {
657 "version": "1",
658 "definition": {
659 "kind": "prompt",
660 "model": "gpt-4o",
661 "tools": [{
662 "type": "mcp",
663 "server_label": "kb_test",
664 "require_approval": "never",
665 "allowed_tools": ["tool_a", "tool_b"]
666 }]
667 }
668 }
669 }
670 });
671
672 let flattened = flatten_agent_response(&api_response);
673 let tools = flattened["tools"].as_array().unwrap();
674 assert_eq!(tools[0]["require_approval"], "never");
675 let allowed = tools[0]["allowed_tools"].as_array().unwrap();
676 assert_eq!(allowed.len(), 2);
677 assert_eq!(allowed[0], "tool_a");
678 assert_eq!(allowed[1], "tool_b");
679 }
680
681 #[test]
682 fn test_flatten_preserves_require_approval_object_form() {
683 let api_response = json!({
684 "object": "agent",
685 "id": "granular-agent",
686 "name": "granular-agent",
687 "versions": {
688 "latest": {
689 "version": "1",
690 "definition": {
691 "kind": "prompt",
692 "model": "gpt-4o",
693 "tools": [{
694 "type": "mcp",
695 "server_label": "kb_test",
696 "require_approval": {
697 "never": {"tool_names": ["safe_tool"]},
698 "always": {"tool_names": ["dangerous_tool"]}
699 }
700 }]
701 }
702 }
703 }
704 });
705
706 let flat = flatten_agent_response(&api_response);
707 let ra = &flat["tools"][0]["require_approval"];
708 assert_eq!(ra["never"]["tool_names"][0], "safe_tool");
709 assert_eq!(ra["always"]["tool_names"][0], "dangerous_tool");
710
711 let re_wrapped = wrap_agent_payload(&flat);
713 let def_ra = &re_wrapped["definition"]["tools"][0]["require_approval"];
714 assert_eq!(def_ra["never"]["tool_names"][0], "safe_tool");
715 assert_eq!(def_ra["always"]["tool_names"][0], "dangerous_tool");
716 }
717}