rigg_client/

client.rs

//! Azure Search REST API client

use std::time::Duration;

use reqwest::{Client, Method, StatusCode};
use serde_json::Value;
use tracing::{debug, instrument, warn};

use rigg_core::config::SearchServiceConfig;
use rigg_core::resources::ResourceKind;

use crate::auth::{AuthProvider, get_auth_provider};
use crate::error::ClientError;

/// Maximum number of retry attempts for retryable errors
const MAX_RETRIES: u32 = 3;

/// Initial backoff delay in seconds
const INITIAL_BACKOFF_SECS: u64 = 1;

/// Calculate the backoff duration for a given retry attempt.
///
/// For `RateLimited` errors with a `retry_after` value, that value is used directly.
/// For other retryable errors, exponential backoff is applied: 1s, 2s, 4s, etc.
fn retry_delay(error: &ClientError, attempt: u32) -> Duration {
    match error {
        ClientError::RateLimited { retry_after } => Duration::from_secs(*retry_after),
        _ => Duration::from_secs(INITIAL_BACKOFF_SECS * 2u64.pow(attempt)),
    }
}

/// Azure Search API client
pub struct AzureSearchClient {
    http: Client,
    auth: Box<dyn AuthProvider>,
    base_url: String,
    api_version: String,
    preview_api_version: String,
}

/// Indexer creation validates the data source and AI-services connections
/// and starts the first run before responding — routinely slower than a
/// typical management call. 30s produced live false-failure timeouts.
const HTTP_TIMEOUT_SECS: u64 = 120;

impl AzureSearchClient {
    /// Create client from a workspace search connection.
    pub fn from_connection(
        conn: &rigg_core::workspace::SearchConnection,
    ) -> Result<Self, ClientError> {
        let auth = get_auth_provider()?;
        let http = Client::builder()
            .timeout(Duration::from_secs(HTTP_TIMEOUT_SECS))
            .build()?;
        Ok(Self {
            http,
            auth,
            base_url: conn.url(),
            api_version: conn
                .api_version
                .clone()
                .unwrap_or_else(|| rigg_core::registry::SEARCH_STABLE_API_VERSION.to_string()),
            preview_api_version: conn
                .preview_api_version
                .clone()
                .unwrap_or_else(|| rigg_core::registry::SEARCH_PREVIEW_API_VERSION.to_string()),
        })
    }

    /// Create client from a specific search service config (legacy).
    pub fn from_service_config(service: &SearchServiceConfig) -> Result<Self, ClientError> {
        let auth = get_auth_provider()?;
        let http = Client::builder()
            .timeout(Duration::from_secs(HTTP_TIMEOUT_SECS))
            .build()?;

        Ok(Self {
            http,
            auth,
            base_url: service.service_url(),
            api_version: rigg_core::registry::SEARCH_STABLE_API_VERSION.to_string(),
            preview_api_version: service.preview_api_version.clone(),
        })
    }

    /// Create with a custom auth provider and explicit versions (tests).
    pub fn with_auth(
        base_url: String,
        api_version: String,
        preview_api_version: String,
        auth: Box<dyn AuthProvider>,
    ) -> Result<Self, ClientError> {
        let http = Client::builder()
            .timeout(std::time::Duration::from_secs(HTTP_TIMEOUT_SECS))
            .build()?;

        Ok(Self {
            http,
            auth,
            base_url,
            api_version,
            preview_api_version,
        })
    }

    /// Get the API version to use for a resource kind, per the registry
    /// channel: stable kinds use the stable api-version, preview-gated kinds
    /// the preview one. (Pre-0.18 the client always used preview; GA of
    /// agentic retrieval in 2026-04-01 makes stable the default.)
    fn api_version_for(&self, kind: ResourceKind) -> &str {
        match rigg_core::registry::meta(kind).channel {
            rigg_core::registry::Channel::Stable => &self.api_version,
            rigg_core::registry::Channel::Preview => &self.preview_api_version,
        }
    }

    /// Build URL for a resource collection
    fn collection_url(&self, kind: ResourceKind) -> String {
        format!(
            "{}/{}?api-version={}",
            self.base_url,
            kind.api_path(),
            self.api_version_for(kind)
        )
    }

    /// Build URL for a specific resource
    fn resource_url(&self, kind: ResourceKind, name: &str) -> String {
        format!(
            "{}/{}/{}?api-version={}",
            self.base_url,
            kind.api_path(),
            urlencoding::encode(name),
            self.api_version_for(kind)
        )
    }

    /// Execute an HTTP request
    async fn request(
        &self,
        method: Method,
        url: &str,
        body: Option<&Value>,
    ) -> Result<Option<Value>, ClientError> {
        let token = self.auth.get_token()?;

        let mut request = self
            .http
            .request(method.clone(), url)
            .header("Authorization", format!("Bearer {}", token))
            .header("Content-Type", "application/json");

        if let Some(json) = body {
            request = request.json(json);
        } else if method == Method::POST {
            // Azure's front end rejects bodyless POSTs without an explicit
            // Content-Length (HTTP 411) — indexer run/reset take no body.
            request = request.header("Content-Length", "0");
        }

        debug!("Request: {} {}", method, url);
        let response = request.send().await?;
        let status = response.status();

        if status == StatusCode::NO_CONTENT {
            return Ok(None);
        }

        let body = response.text().await?;

        if status.is_success() {
            if body.is_empty() {
                Ok(None)
            } else {
                let value: Value = serde_json::from_str(&body)?;
                Ok(Some(value))
            }
        } else {
            match status {
                StatusCode::NOT_FOUND => Err(ClientError::NotFound {
                    kind: "resource".to_string(),
                    name: url.to_string(),
                }),
                StatusCode::CONFLICT => Err(ClientError::AlreadyExists {
                    kind: "resource".to_string(),
                    name: url.to_string(),
                }),
                StatusCode::TOO_MANY_REQUESTS => {
                    let retry_after = 60; // Default retry time
                    Err(ClientError::RateLimited { retry_after })
                }
                StatusCode::SERVICE_UNAVAILABLE => Err(ClientError::ServiceUnavailable(body)),
                _ => Err(ClientError::from_response_with_url(
                    status.as_u16(),
                    &body,
                    Some(url),
                )),
            }
        }
    }

    /// Execute an HTTP request with retry logic for transient errors.
    ///
    /// Retries up to [`MAX_RETRIES`] times for retryable errors (429 and 503).
    /// Uses exponential backoff (1s, 2s, 4s) for 503 errors and respects the
    /// `retry_after` value for 429 rate-limiting errors.
    async fn request_with_retry(
        &self,
        method: Method,
        url: &str,
        body: Option<&Value>,
    ) -> Result<Option<Value>, ClientError> {
        let mut attempt = 0u32;
        loop {
            match self.request(method.clone(), url, body).await {
                Ok(value) => return Ok(value),
                Err(err) if err.is_retryable() && attempt < MAX_RETRIES => {
                    let delay = retry_delay(&err, attempt);
                    warn!(
                        "Request {} {} failed (attempt {}/{}): {}. Retrying in {:?}",
                        method,
                        url,
                        attempt + 1,
                        MAX_RETRIES + 1,
                        err,
                        delay,
                    );
                    tokio::time::sleep(delay).await;
                    attempt += 1;
                }
                Err(err) => return Err(err),
            }
        }
    }

    /// List all resources of a given kind
    #[instrument(skip(self))]
    pub async fn list(&self, kind: ResourceKind) -> Result<Vec<Value>, ClientError> {
        let url = self.collection_url(kind);
        let response = self.request_with_retry(Method::GET, &url, None).await?;

        match response {
            Some(value) => {
                // Azure returns { "value": [...] }
                let items = value
                    .get("value")
                    .and_then(|v| v.as_array())
                    .cloned()
                    .unwrap_or_default();
                Ok(items)
            }
            None => Ok(Vec::new()),
        }
    }

    /// Get a specific resource
    #[instrument(skip(self))]
    pub async fn get(&self, kind: ResourceKind, name: &str) -> Result<Value, ClientError> {
        let url = self.resource_url(kind, name);
        let response = self.request_with_retry(Method::GET, &url, None).await?;

        response.ok_or_else(|| ClientError::NotFound {
            kind: kind.display_name().to_string(),
            name: name.to_string(),
        })
    }

    /// Create or update a resource
    ///
    /// Returns the response body if the API returns one. Some APIs (especially
    /// preview endpoints like Knowledge Sources) return 204 No Content on
    /// successful update, which yields `Ok(None)`.
    #[instrument(skip(self, definition))]
    pub async fn create_or_update(
        &self,
        kind: ResourceKind,
        name: &str,
        definition: &Value,
    ) -> Result<Option<Value>, ClientError> {
        let url = self.resource_url(kind, name);
        self.request_with_retry(Method::PUT, &url, Some(definition))
            .await
    }

    /// Delete a resource
    #[instrument(skip(self))]
    pub async fn delete(&self, kind: ResourceKind, name: &str) -> Result<(), ClientError> {
        let url = self.resource_url(kind, name);
        self.request_with_retry(Method::DELETE, &url, None).await?;
        Ok(())
    }

    // ------------------------------------------------------------------
    // Runtime operations (rigg az)
    // ------------------------------------------------------------------

    /// Trigger an indexer run now (202 Accepted, no body).
    #[instrument(skip(self))]
    pub async fn indexer_run(&self, name: &str) -> Result<(), ClientError> {
        let url = format!(
            "{}/indexers/{}/run?api-version={}",
            self.base_url,
            urlencoding::encode(name),
            self.api_version_for(ResourceKind::Indexer)
        );
        self.request_with_retry(Method::POST, &url, None).await?;
        Ok(())
    }

    /// Clear an indexer's change-tracking state (the next run reprocesses
    /// every document).
    #[instrument(skip(self))]
    pub async fn indexer_reset(&self, name: &str) -> Result<(), ClientError> {
        let url = format!(
            "{}/indexers/{}/reset?api-version={}",
            self.base_url,
            urlencoding::encode(name),
            self.api_version_for(ResourceKind::Indexer)
        );
        self.request_with_retry(Method::POST, &url, None).await?;
        Ok(())
    }

    /// Execution state and history for an indexer.
    #[instrument(skip(self))]
    pub async fn indexer_status(&self, name: &str) -> Result<Value, ClientError> {
        let url = format!(
            "{}/indexers/{}/status?api-version={}",
            self.base_url,
            urlencoding::encode(name),
            self.api_version_for(ResourceKind::Indexer)
        );
        self.request_with_retry(Method::GET, &url, None)
            .await?
            .ok_or_else(|| ClientError::NotFound {
                kind: "indexer status".to_string(),
                name: name.to_string(),
            })
    }

    /// Document count and storage size for an index.
    #[instrument(skip(self))]
    pub async fn index_stats(&self, name: &str) -> Result<Value, ClientError> {
        let url = format!(
            "{}/indexes/{}/stats?api-version={}",
            self.base_url,
            urlencoding::encode(name),
            self.api_version_for(ResourceKind::Index)
        );
        self.request_with_retry(Method::GET, &url, None)
            .await?
            .ok_or_else(|| ClientError::NotFound {
                kind: "index stats".to_string(),
                name: name.to_string(),
            })
    }

    /// Run a search query against an index's documents.
    #[instrument(skip(self, body))]
    pub async fn search_docs(&self, index: &str, body: &Value) -> Result<Value, ClientError> {
        let url = format!(
            "{}/indexes/{}/docs/search?api-version={}",
            self.base_url,
            urlencoding::encode(index),
            self.api_version_for(ResourceKind::Index)
        );
        self.request_with_retry(Method::POST, &url, Some(body))
            .await?
            .ok_or_else(|| ClientError::Api {
                status: 0,
                message: "search returned no body".to_string(),
            })
    }

    /// Agentic retrieval against a knowledge base (the `retrieve` action).
    /// A 206 Partial response (some knowledge sources errored) still carries
    /// a usable body and is returned as success.
    #[instrument(skip(self, body))]
    pub async fn kb_retrieve(&self, kb: &str, body: &Value) -> Result<Value, ClientError> {
        let url = format!(
            "{}/knowledgebases('{}')/retrieve?api-version={}",
            self.base_url,
            urlencoding::encode(kb),
            self.api_version_for(ResourceKind::KnowledgeBase)
        );
        self.request_with_retry(Method::POST, &url, Some(body))
            .await?
            .ok_or_else(|| ClientError::Api {
                status: 0,
                message: "retrieve returned no body".to_string(),
            })
    }

    /// Check if a resource exists
    pub async fn exists(&self, kind: ResourceKind, name: &str) -> Result<bool, ClientError> {
        match self.get(kind, name).await {
            Ok(_) => Ok(true),
            Err(ClientError::NotFound { .. }) => Ok(false),
            Err(e) => Err(e),
        }
    }

    /// Get the authentication method being used
    pub fn auth_method(&self) -> &'static str {
        self.auth.method_name()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::auth::{AuthError, AuthProvider};

    struct FakeAuth;
    impl AuthProvider for FakeAuth {
        fn get_token(&self) -> Result<String, AuthError> {
            Ok("fake-token".to_string())
        }
        fn method_name(&self) -> &'static str {
            "Fake"
        }
    }

    fn make_client() -> AzureSearchClient {
        AzureSearchClient::with_auth(
            "https://test-svc.search.windows.net".to_string(),
            "2026-04-01".to_string(),
            "2026-05-01-preview".to_string(),
            Box::new(FakeAuth),
        )
        .unwrap()
    }

    #[test]
    fn stable_kinds_use_stable_version() {
        let client = make_client();
        let url = client.collection_url(ResourceKind::Index);
        assert_eq!(
            url,
            "https://test-svc.search.windows.net/indexes?api-version=2026-04-01"
        );
        // agentic retrieval is GA in 2026-04-01
        let url = client.collection_url(ResourceKind::KnowledgeBase);
        assert_eq!(
            url,
            "https://test-svc.search.windows.net/knowledgeBases?api-version=2026-04-01"
        );
        let url = client.resource_url(ResourceKind::KnowledgeSource, "ks");
        assert_eq!(
            url,
            "https://test-svc.search.windows.net/knowledgeSources/ks?api-version=2026-04-01"
        );
    }

    #[test]
    fn resource_url_percent_encodes_name() {
        let client = make_client();
        let url = client.resource_url(ResourceKind::Index, "my index");
        assert!(url.contains("/indexes/my%20index?"));
    }

    #[test]
    fn search_kinds_route_via_registry_channel() {
        let client = make_client();
        for kind in ResourceKind::search_kinds() {
            let url = client.collection_url(kind);
            let expected = match rigg_core::registry::meta(kind).channel {
                rigg_core::registry::Channel::Stable => "2026-04-01",
                rigg_core::registry::Channel::Preview => "2026-05-01-preview",
            };
            assert!(
                url.contains(expected),
                "{kind:?} should use {expected}, got: {url}"
            );
        }
    }

    #[test]
    fn test_retry_delay_exponential_backoff_attempt_0() {
        let err = ClientError::ServiceUnavailable("down".to_string());
        let delay = retry_delay(&err, 0);
        assert_eq!(delay, Duration::from_secs(1));
    }

    #[test]
    fn test_retry_delay_exponential_backoff_attempt_1() {
        let err = ClientError::ServiceUnavailable("down".to_string());
        let delay = retry_delay(&err, 1);
        assert_eq!(delay, Duration::from_secs(2));
    }

    #[test]
    fn test_retry_delay_exponential_backoff_attempt_2() {
        let err = ClientError::ServiceUnavailable("down".to_string());
        let delay = retry_delay(&err, 2);
        assert_eq!(delay, Duration::from_secs(4));
    }

    #[test]
    fn test_retry_delay_rate_limited_uses_retry_after() {
        let err = ClientError::RateLimited { retry_after: 30 };
        // retry_after should be used regardless of attempt number
        assert_eq!(retry_delay(&err, 0), Duration::from_secs(30));
        assert_eq!(retry_delay(&err, 1), Duration::from_secs(30));
        assert_eq!(retry_delay(&err, 2), Duration::from_secs(30));
    }

    #[test]
    fn test_retry_delay_rate_limited_default_retry_after() {
        let err = ClientError::RateLimited { retry_after: 60 };
        let delay = retry_delay(&err, 0);
        assert_eq!(delay, Duration::from_secs(60));
    }

    #[test]
    fn test_retry_constants() {
        assert_eq!(MAX_RETRIES, 3);
        assert_eq!(INITIAL_BACKOFF_SECS, 1);
    }

    #[test]
    fn test_retry_delay_backoff_sequence() {
        let err = ClientError::ServiceUnavailable("temporarily unavailable".to_string());
        let delays: Vec<Duration> = (0..MAX_RETRIES).map(|i| retry_delay(&err, i)).collect();
        assert_eq!(
            delays,
            vec![
                Duration::from_secs(1),
                Duration::from_secs(2),
                Duration::from_secs(4),
            ]
        );
    }

    #[test]
    fn test_non_retryable_error_still_computes_delay() {
        // retry_delay computes a delay regardless; the caller decides whether to retry.
        // This verifies the function doesn't panic on non-retryable errors.
        let err = ClientError::Api {
            status: 400,
            message: "bad request".to_string(),
        };
        let delay = retry_delay(&err, 0);
        assert_eq!(delay, Duration::from_secs(1));
    }
}