kontext_dev_sdk/

lib.rs

use base64::Engine;
use reqwest::Client;
use reqwest::StatusCode;
use serde::Deserialize;
use serde::Serialize;
use serde::de::DeserializeOwned;
use sha2::Digest;
use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
use tiny_http::Response;
use tiny_http::Server;
use tokio::sync::oneshot;
use tokio::time::timeout;
use url::Url;

pub mod client;
pub mod errors;
mod http_client;
pub mod management;
pub mod mcp;
pub mod oauth;
pub mod orchestrator;
pub mod server;
pub mod verify;

pub use client::ClientState;
pub use client::ConnectSessionResult;
pub use client::IntegrationInfo;
pub use client::KontextClient as SdkKontextClient;
pub use client::KontextClientConfig;
pub use client::ToolResult;
pub use client::create_kontext_client;
pub use errors::*;
pub use management::KontextManagementClient;
pub use management::KontextManagementClientConfig;
pub use management::*;
pub use mcp::KontextMcp;
pub use mcp::KontextMcpConfig;
pub use mcp::RuntimeIntegrationCategory;
pub use mcp::RuntimeIntegrationConnectType;
pub use mcp::RuntimeIntegrationRecord;
pub use oauth::KontextOAuthProvider;
pub use oauth::KontextOAuthProviderConfig;
pub use oauth::ParsedOAuthCallback;
pub use oauth::TokenExchangeConfig;
pub use oauth::exchange_token;
pub use oauth::parse_oauth_callback;
pub use orchestrator::KontextOrchestrator;
pub use orchestrator::KontextOrchestratorConfig;
pub use orchestrator::KontextOrchestratorState;
pub use orchestrator::create_kontext_orchestrator;
pub use server::IntegrationCredential;
pub use server::IntegrationName;
pub use server::IntegrationResolvedCredentials;
pub use server::KnownIntegration;
pub use server::Kontext;
pub use server::KontextOptions;
pub use server::MiddlewareOptions;
pub use verify::JwksClient;
pub use verify::KontextTokenVerifier;
pub use verify::KontextTokenVerifierConfig;
pub use verify::TokenVerificationError;
pub use verify::TokenVerificationErrorCode;
pub use verify::VerifiedTokenClaims;
pub use verify::VerifyResult;

pub use kontext_dev_sdk_core::AccessToken;
pub use kontext_dev_sdk_core::DEFAULT_AUTH_TIMEOUT_SECONDS;
pub use kontext_dev_sdk_core::DEFAULT_RESOURCE;
pub use kontext_dev_sdk_core::DEFAULT_SCOPE;
pub use kontext_dev_sdk_core::DEFAULT_SERVER;
pub use kontext_dev_sdk_core::DEFAULT_SERVER_NAME;
pub use kontext_dev_sdk_core::KontextDevConfig;
pub use kontext_dev_sdk_core::KontextDevCoreError;
pub use kontext_dev_sdk_core::TokenExchangeToken;
pub use kontext_dev_sdk_core::normalize_server_url;
pub use kontext_dev_sdk_core::resolve_authorize_url;
pub use kontext_dev_sdk_core::resolve_connect_session_url;
pub use kontext_dev_sdk_core::resolve_integration_connection_url;
pub use kontext_dev_sdk_core::resolve_integration_oauth_init_url;
pub use kontext_dev_sdk_core::resolve_mcp_url;
pub use kontext_dev_sdk_core::resolve_server_base_url;
pub use kontext_dev_sdk_core::resolve_token_url;

const TOKEN_EXCHANGE_GRANT_TYPE: &str = "urn:ietf:params:oauth:grant-type:token-exchange";
const TOKEN_TYPE_ACCESS_TOKEN: &str = "urn:ietf:params:oauth:token-type:access_token";
const TOKEN_EXPIRY_BUFFER_SECONDS: u64 = 60;

#[derive(Debug, thiserror::Error)]
pub enum KontextDevError {
    #[error(transparent)]
    Core(#[from] kontext_dev_sdk_core::KontextDevCoreError),
    #[error("failed to parse URL `{url}`")]
    InvalidUrl {
        url: String,
        source: url::ParseError,
    },
    #[error("failed to open browser for OAuth authorization")]
    BrowserOpenFailed,
    #[error("OAuth callback timed out after {timeout_seconds}s")]
    OAuthCallbackTimeout { timeout_seconds: i64 },
    #[error("OAuth callback channel was unexpectedly cancelled")]
    OAuthCallbackCancelled,
    #[error("OAuth callback is missing the authorization code")]
    MissingAuthorizationCode,
    #[error("OAuth callback returned an error: {error}")]
    OAuthCallbackError { error: String },
    #[error("OAuth state mismatch")]
    InvalidOAuthState,
    #[error("Kontext-Dev token request failed for {token_url}: {message}")]
    TokenRequest { token_url: String, message: String },
    #[error("Kontext-Dev token exchange failed for resource `{resource}`: {message}")]
    TokenExchange { resource: String, message: String },
    #[error("Kontext-Dev connect session request failed: {message}")]
    ConnectSession { message: String },
    #[error("Kontext-Dev integration OAuth init failed: {message}")]
    IntegrationOAuthInit { message: String },
    #[error("failed to read token cache at `{path}`: {source}")]
    TokenCacheRead {
        path: String,
        source: std::io::Error,
    },
    #[error("failed to write token cache at `{path}`: {source}")]
    TokenCacheWrite {
        path: String,
        source: std::io::Error,
    },
    #[error("failed to deserialize token cache at `{path}`: {source}")]
    TokenCacheDeserialize {
        path: String,
        source: serde_json::Error,
    },
    #[error("failed to serialize token cache: {source}")]
    TokenCacheSerialize { source: serde_json::Error },
    #[error("Kontext-Dev access token is empty")]
    EmptyAccessToken,
    #[error("missing integration UI URL; set `integration_ui_url` in config")]
    MissingIntegrationUiUrl,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct ConnectSession {
    #[serde(rename = "sessionId")]
    pub session_id: String,
    #[serde(rename = "expiresAt")]
    pub expires_at: String,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct IntegrationOAuthInitResponse {
    #[serde(rename = "authorizationUrl")]
    pub authorization_url: String,
    #[serde(default)]
    pub state: Option<String>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct IntegrationConnectionStatus {
    pub connected: bool,
    #[serde(default)]
    pub expires_at: Option<String>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
pub struct KontextAuthSession {
    pub identity_token: AccessToken,
    pub gateway_token: TokenExchangeToken,
    pub browser_auth_performed: bool,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct CachedAccessToken {
    access_token: String,
    token_type: String,
    refresh_token: Option<String>,
    scope: Option<String>,
    expires_at_unix_ms: Option<u64>,
}

#[derive(Clone, Debug, Deserialize, Serialize)]
struct TokenCacheFile {
    client_id: String,
    resource: String,
    identity: CachedAccessToken,
    gateway: CachedAccessToken,
}

impl CachedAccessToken {
    fn from_access_token(token: &AccessToken) -> Result<Self, KontextDevError> {
        if token.access_token.is_empty() {
            return Err(KontextDevError::EmptyAccessToken);
        }

        Ok(Self {
            access_token: token.access_token.clone(),
            token_type: token.token_type.clone(),
            refresh_token: token.refresh_token.clone(),
            scope: token.scope.clone(),
            expires_at_unix_ms: compute_expires_at_unix_ms(token.expires_in, &token.access_token),
        })
    }

    fn from_token_exchange(token: &TokenExchangeToken) -> Result<Self, KontextDevError> {
        if token.access_token.is_empty() {
            return Err(KontextDevError::EmptyAccessToken);
        }

        Ok(Self {
            access_token: token.access_token.clone(),
            token_type: token.token_type.clone(),
            refresh_token: token.refresh_token.clone(),
            scope: token.scope.clone(),
            expires_at_unix_ms: compute_expires_at_unix_ms(token.expires_in, &token.access_token),
        })
    }

    fn is_valid(&self) -> bool {
        match self.expires_at_unix_ms {
            Some(expires_at) => {
                let buffer_ms = TOKEN_EXPIRY_BUFFER_SECONDS * 1000;
                now_unix_ms().saturating_add(buffer_ms) < expires_at
            }
            None => true,
        }
    }

    fn to_access_token(&self) -> AccessToken {
        AccessToken {
            access_token: self.access_token.clone(),
            token_type: self.token_type.clone(),
            expires_in: self
                .expires_at_unix_ms
                .and_then(unix_ms_to_relative_seconds),
            refresh_token: self.refresh_token.clone(),
            scope: self.scope.clone(),
        }
    }

    fn to_token_exchange_token(&self) -> TokenExchangeToken {
        TokenExchangeToken {
            access_token: self.access_token.clone(),
            issued_token_type: TOKEN_TYPE_ACCESS_TOKEN.to_string(),
            token_type: self.token_type.clone(),
            expires_in: self
                .expires_at_unix_ms
                .and_then(unix_ms_to_relative_seconds),
            scope: self.scope.clone(),
            refresh_token: self.refresh_token.clone(),
        }
    }
}

fn now_unix_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_else(|_| Duration::from_secs(0))
        .as_millis() as u64
}

fn unix_ms_to_relative_seconds(unix_ms: u64) -> Option<i64> {
    if unix_ms <= now_unix_ms() {
        return Some(0);
    }

    let delta_ms = unix_ms - now_unix_ms();
    let secs = delta_ms / 1000;
    i64::try_from(secs).ok()
}

fn compute_expires_at_unix_ms(expires_in: Option<i64>, access_token: &str) -> Option<u64> {
    if let Some(expires_in) = expires_in
        && expires_in > 0
    {
        let expires_in_u64 = u64::try_from(expires_in).ok()?;
        return Some(now_unix_ms().saturating_add(expires_in_u64.saturating_mul(1000)));
    }

    decode_jwt_exp(access_token).map(|exp| exp.saturating_mul(1000))
}

fn decode_jwt_exp(token: &str) -> Option<u64> {
    let mut parts = token.split('.');
    let _header = parts.next()?;
    let payload = parts.next()?;

    let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
        .decode(payload)
        .ok()?;
    let value: serde_json::Value = serde_json::from_slice(&decoded).ok()?;
    value.get("exp").and_then(|exp| {
        exp.as_u64()
            .or_else(|| exp.as_i64().and_then(|v| u64::try_from(v).ok()))
    })
}

#[derive(Clone, Debug, Deserialize)]
struct OAuthErrorBody {
    error: Option<String>,
    error_description: Option<String>,
    message: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
struct OAuthAuthorizationServerMetadata {
    authorization_endpoint: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
struct OAuthCallbackPayload {
    code: Option<String>,
    state: Option<String>,
    error: Option<String>,
    error_description: Option<String>,
}

#[derive(Clone, Debug)]
struct PkcePair {
    verifier: String,
    challenge: String,
}

fn generate_pkce_pair() -> PkcePair {
    let mut raw = [0u8; 64];
    rand::RngCore::fill_bytes(&mut rand::rng(), &mut raw);

    let verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw);
    let digest = sha2::Sha256::digest(verifier.as_bytes());
    let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);

    PkcePair {
        verifier,
        challenge,
    }
}

fn generate_state() -> String {
    let mut bytes = [0u8; 16];
    rand::RngCore::fill_bytes(&mut rand::rng(), &mut bytes);
    bytes.iter().map(|b| format!("{b:02x}")).collect()
}

fn normalized_scope(scope: &str) -> Option<&str> {
    let trimmed = scope.trim();
    if trimmed.is_empty() {
        None
    } else {
        Some(trimmed)
    }
}

struct CallbackServerGuard {
    server: Arc<Server>,
}

impl Drop for CallbackServerGuard {
    fn drop(&mut self) {
        self.server.unblock();
    }
}

fn parse_callback(url_path: &str, callback_path: &str) -> Option<OAuthCallbackPayload> {
    let full = format!("http://localhost{url_path}");
    let parsed = Url::parse(&full).ok()?;

    if parsed.path() != callback_path {
        return None;
    }

    let mut payload = OAuthCallbackPayload {
        code: None,
        state: None,
        error: None,
        error_description: None,
    };

    for (key, value) in parsed.query_pairs() {
        match key.as_ref() {
            "code" => payload.code = Some(value.to_string()),
            "state" => payload.state = Some(value.to_string()),
            "error" => payload.error = Some(value.to_string()),
            "error_description" => payload.error_description = Some(value.to_string()),
            _ => {}
        }
    }

    Some(payload)
}

fn spawn_callback_server(
    server: Arc<Server>,
    callback_path: String,
    tx: oneshot::Sender<OAuthCallbackPayload>,
) -> tokio::task::JoinHandle<()> {
    tokio::task::spawn_blocking(move || {
        while let Ok(request) = server.recv() {
            let path = request.url().to_string();
            if let Some(payload) = parse_callback(&path, &callback_path) {
                let response = Response::from_string(
                    "Authentication complete. You can return to your terminal.",
                );
                let _ = request.respond(response);
                let _ = tx.send(payload);
                break;
            }

            let response = Response::from_string("Invalid callback").with_status_code(400);
            let _ = request.respond(response);
        }
    })
}

#[derive(Clone, Debug)]
pub struct KontextDevClient {
    config: KontextDevConfig,
    http: Client,
}

impl KontextDevClient {
    pub fn new(config: KontextDevConfig) -> Self {
        Self {
            config,
            http: Client::new(),
        }
    }

    pub fn config(&self) -> &KontextDevConfig {
        &self.config
    }

    pub fn server_base_url(&self) -> Result<String, KontextDevError> {
        resolve_server_base_url(&self.config).map_err(KontextDevError::from)
    }

    pub fn mcp_url(&self) -> Result<String, KontextDevError> {
        resolve_mcp_url(&self.config).map_err(KontextDevError::from)
    }

    pub fn token_url(&self) -> Result<String, KontextDevError> {
        resolve_token_url(&self.config).map_err(KontextDevError::from)
    }

    pub fn authorize_url(&self) -> Result<String, KontextDevError> {
        resolve_authorize_url(&self.config).map_err(KontextDevError::from)
    }

    pub fn connect_session_url(&self) -> Result<String, KontextDevError> {
        resolve_connect_session_url(&self.config).map_err(KontextDevError::from)
    }

    pub fn clear_token_cache(&self) -> Result<(), KontextDevError> {
        let Some(path) = self.token_cache_path() else {
            return Ok(());
        };

        if !path.exists() {
            return Ok(());
        }

        fs::remove_file(&path).map_err(|source| KontextDevError::TokenCacheWrite {
            path: path.display().to_string(),
            source,
        })
    }

    /// Authenticate a user via browser PKCE, exchange the identity token to a
    /// resource-scoped MCP gateway token, and persist cache for future runs.
    pub async fn authenticate_mcp(&self) -> Result<KontextAuthSession, KontextDevError> {
        let resource = self.config.resource.clone();

        if let Some(cache) = self.read_cache()?
            && cache.client_id == self.config.client_id
            && cache.resource == resource
            && cache.gateway.is_valid()
            && cache.identity.is_valid()
        {
            return Ok(KontextAuthSession {
                identity_token: cache.identity.to_access_token(),
                gateway_token: cache.gateway.to_token_exchange_token(),
                browser_auth_performed: false,
            });
        }

        if let Some(cache) = self.read_cache()?
            && cache.client_id == self.config.client_id
            && cache.resource == resource
        {
            let maybe_refreshed = if cache.identity.is_valid() {
                Some(cache.identity.to_access_token())
            } else if let Some(refresh_token) = &cache.identity.refresh_token {
                Some(self.refresh_identity_token(refresh_token).await?)
            } else {
                None
            };

            if let Some(identity_token) = maybe_refreshed {
                let gateway_token = self
                    .exchange_for_resource(&identity_token.access_token, &resource, None)
                    .await?;
                self.write_cache(&identity_token, &gateway_token)?;
                return Ok(KontextAuthSession {
                    identity_token,
                    gateway_token,
                    browser_auth_performed: false,
                });
            }
        }

        let identity_token = self.authorize_with_browser_pkce().await?;
        let gateway_token = self
            .exchange_for_resource(&identity_token.access_token, &resource, None)
            .await?;

        self.write_cache(&identity_token, &gateway_token)?;

        Ok(KontextAuthSession {
            identity_token,
            gateway_token,
            browser_auth_performed: true,
        })
    }

    /// Request a short-lived connect session and build the hosted integration
    /// connect URL (`.../oauth/connect?session=...`).
    pub async fn create_integration_connect_url(
        &self,
        gateway_access_token: &str,
    ) -> Result<String, KontextDevError> {
        let session = self.create_connect_session(gateway_access_token).await?;
        self.integration_connect_url(&session.session_id)
    }

    pub fn integration_connect_url(&self, session_id: &str) -> Result<String, KontextDevError> {
        if session_id.trim().is_empty() {
            return Err(KontextDevError::ConnectSession {
                message: "connect session id is empty".to_string(),
            });
        }

        let base = if let Some(explicit) = &self.config.integration_ui_url {
            explicit.trim_end_matches('/').to_string()
        } else {
            let server = resolve_server_base_url(&self.config)?;
            if server.contains("api.kontext.dev") {
                "https://app.kontext.dev".to_string()
            } else {
                server
            }
        };

        let mut url = Url::parse(&base).map_err(|source| KontextDevError::InvalidUrl {
            url: base.clone(),
            source,
        })?;
        url.set_path("/oauth/connect");
        url.query_pairs_mut().append_pair("session", session_id);
        Ok(url.to_string())
    }

    /// Opens the integration connect page in the browser.
    pub async fn open_integration_connect_page(
        &self,
        gateway_access_token: &str,
    ) -> Result<String, KontextDevError> {
        let url = self
            .create_integration_connect_url(gateway_access_token)
            .await?;

        if webbrowser::open(&url).is_err() {
            return Err(KontextDevError::BrowserOpenFailed);
        }

        Ok(url)
    }

    pub async fn create_connect_session(
        &self,
        gateway_access_token: &str,
    ) -> Result<ConnectSession, KontextDevError> {
        let url = self.connect_session_url()?;
        let response = self
            .http
            .post(&url)
            .header("Authorization", format!("Bearer {gateway_access_token}"))
            .json(&serde_json::json!({}))
            .send()
            .await
            .map_err(|err| KontextDevError::ConnectSession {
                message: err.to_string(),
            })?;

        if !response.status().is_success() {
            let message = build_error_message(response).await;
            return Err(KontextDevError::ConnectSession { message });
        }

        response
            .json::<ConnectSession>()
            .await
            .map_err(|err| KontextDevError::ConnectSession {
                message: err.to_string(),
            })
    }

    pub async fn initiate_integration_oauth(
        &self,
        gateway_access_token: &str,
        integration_id: &str,
        return_to: Option<&str>,
    ) -> Result<IntegrationOAuthInitResponse, KontextDevError> {
        let url = resolve_integration_oauth_init_url(&self.config, integration_id)?;

        let payload = return_to
            .map(|value| serde_json::json!({ "returnTo": value }))
            .unwrap_or_else(|| serde_json::json!({}));

        let request = self
            .http
            .post(&url)
            .header("Authorization", format!("Bearer {gateway_access_token}"))
            .json(&payload);

        let response =
            request
                .send()
                .await
                .map_err(|err| KontextDevError::IntegrationOAuthInit {
                    message: err.to_string(),
                })?;

        if !response.status().is_success() {
            let message = build_error_message(response).await;
            return Err(KontextDevError::IntegrationOAuthInit { message });
        }

        response
            .json::<IntegrationOAuthInitResponse>()
            .await
            .map_err(|err| KontextDevError::IntegrationOAuthInit {
                message: err.to_string(),
            })
    }

    pub async fn integration_connection_status(
        &self,
        gateway_access_token: &str,
        integration_id: &str,
    ) -> Result<IntegrationConnectionStatus, KontextDevError> {
        let url = resolve_integration_connection_url(&self.config, integration_id)?;
        let response = self
            .http
            .get(url)
            .header("Authorization", format!("Bearer {gateway_access_token}"))
            .send()
            .await
            .map_err(|err| KontextDevError::IntegrationOAuthInit {
                message: err.to_string(),
            })?;

        if !response.status().is_success() {
            let message = build_error_message(response).await;
            return Err(KontextDevError::IntegrationOAuthInit { message });
        }

        response
            .json::<IntegrationConnectionStatus>()
            .await
            .map_err(|err| KontextDevError::IntegrationOAuthInit {
                message: err.to_string(),
            })
    }

    pub async fn wait_for_integration_connection(
        &self,
        gateway_access_token: &str,
        integration_id: &str,
        timeout_ms: u64,
        interval_ms: u64,
    ) -> Result<bool, KontextDevError> {
        let started = now_unix_ms();

        loop {
            let status = self
                .integration_connection_status(gateway_access_token, integration_id)
                .await?;
            if status.connected {
                return Ok(true);
            }

            if now_unix_ms().saturating_sub(started) >= timeout_ms {
                return Ok(false);
            }

            tokio::time::sleep(Duration::from_millis(interval_ms)).await;
        }
    }

    async fn authorize_with_browser_pkce(&self) -> Result<AccessToken, KontextDevError> {
        let auth_url = self.resolve_authorization_url().await?;
        let token_url = self.token_url()?;
        let pkce = generate_pkce_pair();
        let state = generate_state();

        let (callback_url, callback_payload) = self.listen_for_callback().await?;

        let mut url = Url::parse(&auth_url).map_err(|source| KontextDevError::InvalidUrl {
            url: auth_url.clone(),
            source,
        })?;

        {
            let mut query = url.query_pairs_mut();
            query
                .append_pair("client_id", &self.config.client_id)
                .append_pair("response_type", "code")
                .append_pair("redirect_uri", &callback_url)
                .append_pair("state", &state);

            if let Some(scope) = normalized_scope(&self.config.scope) {
                query.append_pair("scope", scope);
            }

            query
                .append_pair("code_challenge_method", "S256")
                .append_pair("code_challenge", &pkce.challenge);
        }

        if webbrowser::open(url.as_str()).is_err() {
            return Err(KontextDevError::BrowserOpenFailed);
        }

        let payload = callback_payload.await?;

        if let Some(error) = payload.error {
            let with_details = payload
                .error_description
                .map(|description| format!("{error}: {description}"))
                .unwrap_or(error);
            return Err(KontextDevError::OAuthCallbackError {
                error: with_details,
            });
        }

        if payload.state.as_deref() != Some(state.as_str()) {
            return Err(KontextDevError::InvalidOAuthState);
        }

        let code = payload
            .code
            .ok_or(KontextDevError::MissingAuthorizationCode)?;

        let mut body = vec![
            ("grant_type", "authorization_code".to_string()),
            ("code", code),
            ("redirect_uri", callback_url),
            ("client_id", self.config.client_id.clone()),
            ("code_verifier", pkce.verifier),
        ];

        if let Some(client_secret) = &self.config.client_secret {
            body.push(("client_secret", client_secret.clone()));
        }

        post_token(&self.http, &token_url, None, &body).await
    }

    async fn resolve_authorization_url(&self) -> Result<String, KontextDevError> {
        if let Some(discovered) = self.discover_authorization_endpoint().await {
            return Ok(discovered);
        }

        let authorize_url = self.authorize_url()?;
        if !self.endpoint_is_missing(&authorize_url).await {
            return Ok(authorize_url);
        }

        let server_base = resolve_server_base_url(&self.config)?;
        Ok(format!("{}/oauth2/auth", server_base.trim_end_matches('/')))
    }

    async fn discover_authorization_endpoint(&self) -> Option<String> {
        let base = resolve_server_base_url(&self.config).ok()?;
        let base = base.trim_end_matches('/');

        let candidates = [
            format!("{base}/.well-known/oauth-authorization-server/mcp"),
            format!("{base}/.well-known/oauth-authorization-server"),
        ];

        for url in candidates {
            let response = match self.http.get(&url).send().await {
                Ok(response) => response,
                Err(_) => continue,
            };

            if !response.status().is_success() {
                continue;
            }

            let metadata = match response.json::<OAuthAuthorizationServerMetadata>().await {
                Ok(metadata) => metadata,
                Err(_) => continue,
            };

            let Some(endpoint) = metadata.authorization_endpoint else {
                continue;
            };

            if Url::parse(&endpoint).is_ok() {
                return Some(endpoint);
            }
        }

        None
    }

    async fn endpoint_is_missing(&self, url: &str) -> bool {
        let probe_client = match reqwest::Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .build()
        {
            Ok(client) => client,
            Err(_) => return false,
        };

        match probe_client.get(url).send().await {
            Ok(response) => response.status() == StatusCode::NOT_FOUND,
            Err(_) => false,
        }
    }

    async fn refresh_identity_token(
        &self,
        refresh_token: &str,
    ) -> Result<AccessToken, KontextDevError> {
        let token_url = self.token_url()?;
        let mut body = vec![
            ("grant_type", "refresh_token".to_string()),
            ("refresh_token", refresh_token.to_string()),
            ("client_id", self.config.client_id.clone()),
        ];

        if let Some(client_secret) = &self.config.client_secret {
            body.push(("client_secret", client_secret.clone()));
        }

        post_token(&self.http, &token_url, None, &body).await
    }

    pub async fn exchange_for_resource(
        &self,
        subject_token: &str,
        resource: &str,
        scope: Option<&str>,
    ) -> Result<TokenExchangeToken, KontextDevError> {
        let token_url = self.token_url()?;

        let mut body = vec![
            ("grant_type", TOKEN_EXCHANGE_GRANT_TYPE.to_string()),
            ("subject_token", subject_token.to_string()),
            ("subject_token_type", TOKEN_TYPE_ACCESS_TOKEN.to_string()),
            ("resource", resource.to_string()),
        ];

        if let Some(scope) = scope {
            body.push(("scope", scope.to_string()));
        }

        let auth_header = self.config.client_secret.as_ref().map(|secret| {
            let raw = format!("{}:{}", self.config.client_id, secret);
            format!(
                "Basic {}",
                base64::engine::general_purpose::STANDARD.encode(raw.as_bytes())
            )
        });

        if self.config.client_secret.is_none() {
            body.push(("client_id", self.config.client_id.clone()));
        }

        let response = post_form_with_optional_auth::<TokenExchangeToken>(
            &self.http,
            &token_url,
            auth_header,
            &body,
        )
        .await
        .map_err(|message| KontextDevError::TokenExchange {
            resource: resource.to_string(),
            message,
        })?;

        if response.access_token.is_empty() {
            return Err(KontextDevError::EmptyAccessToken);
        }

        Ok(response)
    }

    async fn listen_for_callback(
        &self,
    ) -> Result<
        (
            String,
            impl std::future::Future<Output = Result<OAuthCallbackPayload, KontextDevError>>,
        ),
        KontextDevError,
    > {
        let redirect_uri = self.config.redirect_uri.trim().to_string();
        let parsed = Url::parse(&redirect_uri).map_err(|source| KontextDevError::InvalidUrl {
            url: redirect_uri.clone(),
            source,
        })?;

        if parsed.scheme() != "http" {
            return Err(KontextDevError::OAuthCallbackError {
                error: "redirect_uri must use http".to_string(),
            });
        }

        if parsed.query().is_some() || parsed.fragment().is_some() {
            return Err(KontextDevError::OAuthCallbackError {
                error: "redirect_uri must not include query parameters or fragments".to_string(),
            });
        }

        let host = parsed
            .host_str()
            .ok_or_else(|| KontextDevError::OAuthCallbackError {
                error: "redirect_uri host is missing".to_string(),
            })?;

        let port = parsed
            .port()
            .ok_or_else(|| KontextDevError::OAuthCallbackError {
                error: "redirect_uri must include an explicit port".to_string(),
            })?;

        let callback_path = parsed.path().to_string();
        let bind_addr = if host.contains(':') {
            format!("[{host}]:{port}")
        } else {
            format!("{host}:{port}")
        };

        let server = Arc::new(Server::http(&bind_addr).map_err(|err| {
            KontextDevError::OAuthCallbackError {
                error: format!("failed to start callback server at {bind_addr}: {err}"),
            }
        })?);

        let (tx, rx) = oneshot::channel();
        let _join = spawn_callback_server(server.clone(), callback_path, tx);
        let _guard = CallbackServerGuard { server };
        let timeout_seconds = self.config.auth_timeout_seconds.max(1);

        let fut = async move {
            let payload = timeout(Duration::from_secs(timeout_seconds as u64), rx)
                .await
                .map_err(|_| KontextDevError::OAuthCallbackTimeout { timeout_seconds })?
                .map_err(|_| KontextDevError::OAuthCallbackCancelled)?;
            drop(_guard);
            Ok(payload)
        };

        Ok((redirect_uri, fut))
    }

    fn token_cache_path(&self) -> Option<PathBuf> {
        if let Some(explicit) = &self.config.token_cache_path {
            return Some(PathBuf::from(explicit));
        }

        let home = dirs::home_dir()?;
        let mut path = home;
        path.push(".kontext-dev");
        path.push("tokens");

        let sanitized_client_id: String = self
            .config
            .client_id
            .chars()
            .map(|ch| {
                if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
                    ch
                } else {
                    '_'
                }
            })
            .collect();

        path.push(format!("{sanitized_client_id}.json"));
        Some(path)
    }

    fn read_cache(&self) -> Result<Option<TokenCacheFile>, KontextDevError> {
        let Some(path) = self.token_cache_path() else {
            return Ok(None);
        };

        let raw = match fs::read_to_string(&path) {
            Ok(raw) => raw,
            Err(err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
            Err(source) => {
                return Err(KontextDevError::TokenCacheRead {
                    path: path.display().to_string(),
                    source,
                });
            }
        };

        serde_json::from_str(&raw).map(Some).map_err(|source| {
            KontextDevError::TokenCacheDeserialize {
                path: path.display().to_string(),
                source,
            }
        })
    }

    fn write_cache(
        &self,
        identity: &AccessToken,
        gateway: &TokenExchangeToken,
    ) -> Result<(), KontextDevError> {
        let Some(path) = self.token_cache_path() else {
            return Ok(());
        };

        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).map_err(|source| KontextDevError::TokenCacheWrite {
                path: parent.display().to_string(),
                source,
            })?;
        }

        let payload = TokenCacheFile {
            client_id: self.config.client_id.clone(),
            resource: self.config.resource.clone(),
            identity: CachedAccessToken::from_access_token(identity)?,
            gateway: CachedAccessToken::from_token_exchange(gateway)?,
        };

        let serialized = serde_json::to_string_pretty(&payload)
            .map_err(|source| KontextDevError::TokenCacheSerialize { source })?;

        fs::write(&path, serialized).map_err(|source| KontextDevError::TokenCacheWrite {
            path: path.display().to_string(),
            source,
        })
    }
}

/// Legacy helper that uses the client-credentials grant.
///
/// New PKCE-based apps should use `KontextDevClient::authenticate_mcp`.
pub async fn request_access_token(
    config: &KontextDevConfig,
) -> Result<AccessToken, KontextDevError> {
    let token_url = resolve_token_url(config)?;

    let mut body = vec![("grant_type", "client_credentials".to_string())];

    if let Some(scope) = normalized_scope(&config.scope) {
        body.push(("scope", scope.to_string()));
    }

    let auth_header = if let Some(secret) = &config.client_secret {
        let raw = format!("{}:{}", config.client_id, secret);
        Some(format!(
            "Basic {}",
            base64::engine::general_purpose::STANDARD.encode(raw.as_bytes())
        ))
    } else {
        body.push(("client_id", config.client_id.clone()));
        None
    };

    post_token(&Client::new(), &token_url, auth_header.as_deref(), &body).await
}

async fn post_token(
    http: &Client,
    token_url: &str,
    authorization: Option<&str>,
    body: &[(impl AsRef<str>, String)],
) -> Result<AccessToken, KontextDevError> {
    let body_vec: Vec<(&str, String)> = body.iter().map(|(k, v)| (k.as_ref(), v.clone())).collect();

    let response = post_form_with_optional_auth::<AccessToken>(
        http,
        token_url,
        authorization.map(ToString::to_string),
        &body_vec,
    )
    .await
    .map_err(|message| KontextDevError::TokenRequest {
        token_url: token_url.to_string(),
        message,
    })?;

    Ok(response)
}

async fn post_form_with_optional_auth<T>(
    http: &Client,
    url: &str,
    authorization: Option<String>,
    body: &[(&str, String)],
) -> Result<T, String>
where
    T: DeserializeOwned,
{
    let mut request = http
        .post(url)
        .header("Content-Type", "application/x-www-form-urlencoded");

    if let Some(header) = authorization {
        request = request.header("Authorization", header);
    }

    let form = body
        .iter()
        .map(|(k, v)| (k.to_string(), v.to_string()))
        .collect::<Vec<(String, String)>>();

    let response = request
        .form(&form)
        .send()
        .await
        .map_err(|err| err.to_string())?;

    if !response.status().is_success() {
        return Err(build_error_message(response).await);
    }

    response.json::<T>().await.map_err(|err| err.to_string())
}

async fn build_error_message(response: reqwest::Response) -> String {
    let status = response.status();
    let fallback = format!(
        "{} {}",
        status.as_u16(),
        status.canonical_reason().unwrap_or("")
    );

    let body = response.text().await.unwrap_or_default();
    if body.is_empty() {
        return fallback.trim().to_string();
    }

    if let Ok(parsed) = serde_json::from_str::<OAuthErrorBody>(&body) {
        if let Some(description) = parsed.error_description {
            return description;
        }
        if let Some(message) = parsed.message {
            return message;
        }
        if let Some(error) = parsed.error {
            return error;
        }
    }

    format!("{fallback}: {body}")
}

#[cfg(test)]
mod tests {
    use super::*;

    fn config() -> KontextDevConfig {
        KontextDevConfig {
            server: "http://localhost:4000".to_string(),
            client_id: "client_123".to_string(),
            client_secret: None,
            scope: DEFAULT_SCOPE.to_string(),
            server_name: DEFAULT_SERVER_NAME.to_string(),
            resource: DEFAULT_RESOURCE.to_string(),
            integration_ui_url: Some("https://app.kontext.dev".to_string()),
            integration_return_to: None,
            open_connect_page_on_login: true,
            auth_timeout_seconds: DEFAULT_AUTH_TIMEOUT_SECONDS,
            token_cache_path: None,
            redirect_uri: "http://localhost:3333/callback".to_string(),
        }
    }

    #[test]
    fn create_connect_url_uses_hosted_ui() {
        let client = KontextDevClient::new(config());
        let url = client
            .integration_connect_url("session-123")
            .expect("url should be built");
        assert_eq!(
            url,
            "https://app.kontext.dev/oauth/connect?session=session-123"
        );
    }

    #[test]
    fn jwt_exp_decode_reads_exp() {
        let header = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"alg":"none"}"#);
        let payload =
            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(r#"{"exp":4070908800}"#);
        let token = format!("{header}.{payload}.sig");
        assert_eq!(decode_jwt_exp(&token), Some(4_070_908_800));
    }

    #[test]
    fn oauth_metadata_parses_authorization_endpoint() {
        let metadata = serde_json::from_str::<OAuthAuthorizationServerMetadata>(
            r#"{
              "issuer": "https://issuer.example.com",
              "authorization_endpoint": "https://issuer.example.com/oauth2/auth"
            }"#,
        )
        .expect("metadata should parse");

        assert_eq!(
            metadata.authorization_endpoint.as_deref(),
            Some("https://issuer.example.com/oauth2/auth")
        );
    }

    #[test]
    fn normalized_scope_omits_blank_values() {
        assert_eq!(normalized_scope(""), None);
        assert_eq!(normalized_scope("   "), None);
    }

    #[test]
    fn normalized_scope_preserves_non_empty_values() {
        assert_eq!(normalized_scope("mcp:invoke"), Some("mcp:invoke"));
        assert_eq!(
            normalized_scope("  mcp:invoke openid  "),
            Some("mcp:invoke openid")
        );
    }
}