kontext_dev_sdk/

mcp.rs

use std::sync::Arc;

use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use tokio::sync::Mutex;

use crate::KontextAuthSession;
use crate::KontextDevClient;
use crate::KontextDevConfig;
use crate::KontextDevError;

pub const DEFAULT_SERVER: &str = "https://api.kontext.dev";
const MCP_SESSION_HEADER: &str = "Mcp-Session-Id";
const META_SEARCH_TOOLS: &str = "SEARCH_TOOLS";
const META_EXECUTE_TOOL: &str = "EXECUTE_TOOL";
const DEFAULT_MCP_PROTOCOL_VERSION: &str = "2025-06-18";
const STREAMABLE_HTTP_ACCEPT: &str = "application/json, text/event-stream";
const STREAM_CONTENT_TYPE: &str = "text/event-stream";

pub fn normalize_kontext_server_url(server: &str) -> String {
    let mut url = server.trim_end_matches('/').to_string();
    if let Some(stripped) = url.strip_suffix("/api/v1") {
        url = stripped.to_string();
    }
    if let Some(stripped) = url.strip_suffix("/mcp") {
        url = stripped.to_string();
    }
    url.trim_end_matches('/').to_string()
}

#[derive(Clone, Debug)]
pub struct KontextMcpConfig {
    pub client_session_id: String,
    pub client_id: String,
    pub redirect_uri: String,
    pub url: Option<String>,
    pub server: Option<String>,
    pub client_secret: Option<String>,
    pub scope: Option<String>,
    pub resource: Option<String>,
    pub session_key: Option<String>,
    pub integration_ui_url: Option<String>,
    pub integration_return_to: Option<String>,
    pub auth_timeout_seconds: Option<i64>,
    pub open_connect_page_on_login: Option<bool>,
    pub token_cache_path: Option<String>,
}

impl Default for KontextMcpConfig {
    fn default() -> Self {
        Self {
            client_session_id: String::new(),
            client_id: String::new(),
            redirect_uri: "http://localhost:3333/callback".to_string(),
            url: None,
            server: Some(DEFAULT_SERVER.to_string()),
            client_secret: None,
            scope: None,
            resource: None,
            session_key: None,
            integration_ui_url: None,
            integration_return_to: None,
            auth_timeout_seconds: None,
            open_connect_page_on_login: None,
            token_cache_path: None,
        }
    }
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum RuntimeIntegrationCategory {
    GatewayRemoteMcp,
    InternalMcpCredentials,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum RuntimeIntegrationConnectType {
    Oauth,
    UserToken,
    Credentials,
    None,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeIntegrationRecord {
    pub id: String,
    pub name: String,
    pub url: String,
    pub category: RuntimeIntegrationCategory,
    pub connect_type: RuntimeIntegrationConnectType,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub auth_mode: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub credential_schema: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub requires_oauth: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub connection: Option<RuntimeIntegrationConnection>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeIntegrationConnection {
    pub connected: bool,
    pub status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub expires_at: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub display_name: Option<String>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct KontextTool {
    pub id: String,
    pub name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub input_schema: Option<serde_json::Value>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub server: Option<KontextToolServer>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct KontextToolServer {
    pub id: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

#[derive(Clone, Debug, Default)]
struct McpSessionState {
    session_id: Option<String>,
    access_token: Option<String>,
}

#[derive(Clone, Debug)]
pub struct KontextMcp {
    config: KontextMcpConfig,
    client: KontextDevClient,
    http: reqwest::Client,
    session: Arc<Mutex<McpSessionState>>,
}

impl KontextMcp {
    pub fn new(config: KontextMcpConfig) -> Self {
        let server =
            normalize_kontext_server_url(config.server.as_deref().unwrap_or(DEFAULT_SERVER));
        let sdk_config = KontextDevConfig {
            server,
            client_id: config.client_id.clone(),
            client_secret: config.client_secret.clone(),
            scope: config.scope.clone().unwrap_or_default(),
            server_name: "kontext-dev".to_string(),
            resource: config
                .resource
                .clone()
                .unwrap_or_else(|| "mcp-gateway".to_string()),
            integration_ui_url: config.integration_ui_url.clone(),
            integration_return_to: config.integration_return_to.clone(),
            open_connect_page_on_login: config.open_connect_page_on_login.unwrap_or(true),
            auth_timeout_seconds: config.auth_timeout_seconds.unwrap_or(300),
            token_cache_path: config.token_cache_path.clone(),
            redirect_uri: config.redirect_uri.clone(),
        };

        Self {
            config,
            client: KontextDevClient::new(sdk_config),
            http: reqwest::Client::new(),
            session: Arc::new(Mutex::new(McpSessionState::default())),
        }
    }

    pub fn client(&self) -> &KontextDevClient {
        &self.client
    }

    pub async fn authenticate_mcp(&self) -> Result<KontextAuthSession, KontextDevError> {
        self.client.authenticate_mcp().await
    }

    pub fn mcp_url(&self) -> Result<String, KontextDevError> {
        if let Some(url) = &self.config.url {
            return Ok(url.clone());
        }
        self.client.mcp_url()
    }

    pub async fn clear_cached_session(&self) {
        self.invalidate_session().await;
    }

    pub async fn list_integrations(
        &self,
    ) -> Result<Vec<RuntimeIntegrationRecord>, KontextDevError> {
        let session = self.authenticate_mcp().await?;
        let base = self.client.server_base_url()?;
        let response = self
            .http
            .get(format!("{}/mcp/integrations", base.trim_end_matches('/')))
            .bearer_auth(session.gateway_token.access_token)
            .send()
            .await
            .map_err(|err| KontextDevError::ConnectSession {
                message: err.to_string(),
            })?;

        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(KontextDevError::ConnectSession {
                message: format!("{status}: {body}"),
            });
        }

        #[derive(Deserialize)]
        struct IntegrationsResponse {
            items: Vec<RuntimeIntegrationRecord>,
        }

        let payload = response
            .json::<IntegrationsResponse>()
            .await
            .map_err(|err| KontextDevError::ConnectSession {
                message: err.to_string(),
            })?;

        Ok(payload.items)
    }

    pub async fn list_tools(&self) -> Result<Vec<KontextTool>, KontextDevError> {
        let session = self.authenticate_mcp().await?;
        self.list_tools_with_access_token(&session.gateway_token.access_token)
            .await
    }

    pub async fn list_tools_with_access_token(
        &self,
        access_token: &str,
    ) -> Result<Vec<KontextTool>, KontextDevError> {
        let result = self
            .json_rpc_with_session(
                access_token,
                "tools/list",
                json!({}),
                Some("list-tools"),
                true,
            )
            .await?;
        parse_tools_list_result(&result)
    }

    pub async fn call_tool(
        &self,
        tool_id: &str,
        args: Option<serde_json::Map<String, serde_json::Value>>,
    ) -> Result<serde_json::Value, KontextDevError> {
        let session = self.authenticate_mcp().await?;
        self.call_tool_with_access_token(&session.gateway_token.access_token, tool_id, args)
            .await
    }

    pub async fn call_tool_with_access_token(
        &self,
        access_token: &str,
        tool_id: &str,
        args: Option<serde_json::Map<String, serde_json::Value>>,
    ) -> Result<serde_json::Value, KontextDevError> {
        self.json_rpc_with_session(
            access_token,
            "tools/call",
            json!({ "name": tool_id, "arguments": args.unwrap_or_default() }),
            Some("call-tool"),
            true,
        )
        .await
    }

    async fn json_rpc_with_session(
        &self,
        access_token: &str,
        method: &str,
        params: Value,
        id: Option<&str>,
        allow_session_retry: bool,
    ) -> Result<Value, KontextDevError> {
        let max_attempts = if allow_session_retry { 2 } else { 1 };
        for attempt in 0..max_attempts {
            let session_id = self.ensure_mcp_session(access_token).await?;

            let response = self
                .http
                .post(self.mcp_url()?)
                .bearer_auth(access_token)
                .header(reqwest::header::ACCEPT, STREAMABLE_HTTP_ACCEPT)
                .header(MCP_SESSION_HEADER, &session_id)
                .json(&json!({
                    "jsonrpc": "2.0",
                    "id": id.unwrap_or("1"),
                    "method": method,
                    "params": params,
                }))
                .send()
                .await
                .map_err(|err| KontextDevError::ConnectSession {
                    message: err.to_string(),
                })?;

            if !response.status().is_success() {
                let status = response.status();
                let body = response.text().await.unwrap_or_default();
                let retryable =
                    attempt + 1 < max_attempts && is_invalid_session_response_body(body.as_str());
                if retryable {
                    self.invalidate_session().await;
                    continue;
                }
                return Err(KontextDevError::ConnectSession {
                    message: format!("{status}: {body}"),
                });
            }

            let payload = parse_json_or_streamable_response(response).await?;

            if let Some(error) = payload.get("error") {
                let message = extract_jsonrpc_error_message(error);
                let retryable =
                    attempt + 1 < max_attempts && is_invalid_session_jsonrpc_error(error);
                if retryable {
                    self.invalidate_session().await;
                    continue;
                }
                return Err(KontextDevError::ConnectSession { message });
            }

            return Ok(payload.get("result").cloned().unwrap_or(Value::Null));
        }

        Err(KontextDevError::ConnectSession {
            message: "MCP request failed after session retry".to_string(),
        })
    }

    async fn ensure_mcp_session(&self, access_token: &str) -> Result<String, KontextDevError> {
        {
            let guard = self.session.lock().await;
            if guard.access_token.as_deref() == Some(access_token)
                && let Some(session_id) = guard.session_id.clone()
            {
                return Ok(session_id);
            }
        }

        let initialize_response = self
            .http
            .post(self.mcp_url()?)
            .bearer_auth(access_token)
            .header(reqwest::header::ACCEPT, STREAMABLE_HTTP_ACCEPT)
            .json(&json!({
                "jsonrpc": "2.0",
                "id": "initialize",
                "method": "initialize",
                "params": {
                    "protocolVersion": DEFAULT_MCP_PROTOCOL_VERSION,
                    "capabilities": {
                        "tools": {}
                    },
                    "clientInfo": {
                        "name": "kontext-dev-sdk-rs",
                        "version": env!("CARGO_PKG_VERSION"),
                        "sessionId": self.config.client_session_id
                    }
                }
            }))
            .send()
            .await
            .map_err(|err| KontextDevError::ConnectSession {
                message: err.to_string(),
            })?;

        if !initialize_response.status().is_success() {
            let status = initialize_response.status();
            let body = initialize_response.text().await.unwrap_or_default();
            return Err(KontextDevError::ConnectSession {
                message: format!("{status}: {body}"),
            });
        }

        let session_header = initialize_response
            .headers()
            .get(MCP_SESSION_HEADER)
            .or_else(|| initialize_response.headers().get("mcp-session-id"))
            .and_then(|value| value.to_str().ok())
            .map(|value| value.trim().to_string());

        let initialize_payload = parse_json_or_streamable_response(initialize_response).await?;

        if let Some(error) = initialize_payload.get("error") {
            return Err(KontextDevError::ConnectSession {
                message: extract_jsonrpc_error_message(error),
            });
        }

        let session_id = session_header
            .or_else(|| {
                initialize_payload
                    .get("result")
                    .and_then(|result| result.get("sessionId"))
                    .and_then(|value| value.as_str())
                    .map(|value| value.to_string())
            })
            .or_else(|| {
                initialize_payload
                    .get("result")
                    .and_then(|result| result.get("session_id"))
                    .and_then(|value| value.as_str())
                    .map(|value| value.to_string())
            })
            .ok_or_else(|| KontextDevError::ConnectSession {
                message: "MCP initialize did not return a session id".to_string(),
            })?;

        // Best-effort initialized notification. Many servers accept requests
        // without it, but send it to follow the Streamable HTTP MCP handshake.
        let _ = self
            .http
            .post(self.mcp_url()?)
            .bearer_auth(access_token)
            .header(reqwest::header::ACCEPT, STREAMABLE_HTTP_ACCEPT)
            .header(MCP_SESSION_HEADER, &session_id)
            .json(&json!({
                "jsonrpc": "2.0",
                "method": "notifications/initialized",
                "params": {}
            }))
            .send()
            .await;

        {
            let mut guard = self.session.lock().await;
            guard.session_id = Some(session_id.clone());
            guard.access_token = Some(access_token.to_string());
        }

        Ok(session_id)
    }

    async fn invalidate_session(&self) {
        let mut guard = self.session.lock().await;
        guard.session_id = None;
        guard.access_token = None;
    }
}

fn extract_jsonrpc_error_message(error: &Value) -> String {
    error
        .get("message")
        .and_then(|value| value.as_str())
        .map(ToString::to_string)
        .or_else(|| {
            error
                .get("error_description")
                .and_then(|value| value.as_str())
                .map(ToString::to_string)
        })
        .unwrap_or_else(|| error.to_string())
}

fn is_invalid_session_error(message: &str) -> bool {
    let lower = message.to_ascii_lowercase();
    lower.contains("no valid session id")
        || lower.contains("no valid session-id")
        || lower.contains("invalid session")
}

fn is_session_not_found_error(message: &str) -> bool {
    let lower = message.to_ascii_lowercase();
    lower.contains("session") && lower.contains("not found")
}

fn is_invalid_session_jsonrpc_error(error: &Value) -> bool {
    let message = extract_jsonrpc_error_message(error);
    if is_invalid_session_error(message.as_str()) {
        return true;
    }

    let code = error.get("code").and_then(Value::as_i64);
    code == Some(-32000) && is_session_not_found_error(message.as_str())
}

fn is_invalid_session_response_body(body: &str) -> bool {
    if is_invalid_session_error(body) || is_session_not_found_error(body) {
        return true;
    }

    if let Ok(payload) = serde_json::from_str::<Value>(body)
        && let Some(error) = payload.get("error")
    {
        return is_invalid_session_jsonrpc_error(error);
    }

    false
}

async fn parse_json_or_streamable_response(
    response: reqwest::Response,
) -> Result<Value, KontextDevError> {
    let content_type = response
        .headers()
        .get(reqwest::header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .map(|value| value.to_ascii_lowercase())
        .unwrap_or_default();
    let body = response
        .text()
        .await
        .map_err(|err| KontextDevError::ConnectSession {
            message: err.to_string(),
        })?;

    parse_json_or_streamable_body(&body, &content_type)
        .map_err(|message| KontextDevError::ConnectSession { message })
}

fn parse_json_or_streamable_body(body: &str, content_type: &str) -> Result<Value, String> {
    let parse_json = || serde_json::from_str::<Value>(body).map_err(|err| err.to_string());
    let parse_sse = || parse_sse_last_json_event(body);

    if content_type.contains(STREAM_CONTENT_TYPE) {
        return parse_sse().ok_or_else(|| {
            "failed to parse streamable MCP response as SSE JSON events".to_string()
        });
    }

    parse_json().or_else(|json_err| {
        parse_sse().ok_or_else(|| format!("failed to decode response body: {json_err}"))
    })
}

fn parse_sse_last_json_event(body: &str) -> Option<Value> {
    let mut current_data = Vec::<String>::new();
    let mut last_json = None;

    let flush_data = |current_data: &mut Vec<String>, last_json: &mut Option<Value>| {
        if current_data.is_empty() {
            return;
        }
        let data = current_data.join("\n");
        current_data.clear();
        let trimmed = data.trim();
        if trimmed.is_empty() || trimmed == "[DONE]" {
            return;
        }
        if let Ok(value) = serde_json::from_str::<Value>(trimmed) {
            *last_json = Some(value);
        }
    };

    for line in body.lines() {
        let line = line.trim_end_matches('\r');
        if line.is_empty() {
            flush_data(&mut current_data, &mut last_json);
            continue;
        }
        if let Some(data) = line.strip_prefix("data:") {
            current_data.push(data.trim_start().to_string());
            continue;
        }
        if let Ok(value) = serde_json::from_str::<Value>(line) {
            last_json = Some(value);
        }
    }
    flush_data(&mut current_data, &mut last_json);

    last_json
}

pub(crate) fn has_meta_gateway_tools(tools: &[KontextTool]) -> bool {
    let mut has_search = false;
    let mut has_execute = false;
    for tool in tools {
        if tool.name == META_SEARCH_TOOLS {
            has_search = true;
        } else if tool.name == META_EXECUTE_TOOL {
            has_execute = true;
        }
    }
    has_search && has_execute
}

pub(crate) fn extract_json_resource_text(result: &Value) -> Option<String> {
    let content = result.get("content")?.as_array()?;
    for item in content {
        if item.get("type").and_then(Value::as_str) != Some("resource") {
            continue;
        }
        let Some(resource) = item.get("resource") else {
            continue;
        };
        if resource.get("mimeType").and_then(Value::as_str) != Some("application/json") {
            continue;
        }
        if let Some(text) = resource.get("text").and_then(Value::as_str) {
            return Some(text.to_string());
        }
    }
    None
}

pub(crate) fn extract_text_content(result: &Value) -> String {
    let Some(content) = result.get("content").and_then(Value::as_array) else {
        return result.to_string();
    };

    let mut text_items = Vec::new();
    for item in content {
        if item.get("type").and_then(Value::as_str) == Some("text")
            && let Some(text) = item.get("text").and_then(Value::as_str)
        {
            text_items.push(text.to_string());
        }
    }
    if !text_items.is_empty() {
        return text_items.join("\n");
    }

    let mut resource_items = Vec::new();
    for item in content {
        if item.get("type").and_then(Value::as_str) != Some("resource") {
            continue;
        }
        let Some(resource_text) = item
            .get("resource")
            .and_then(|resource| resource.get("text"))
            .and_then(Value::as_str)
        else {
            continue;
        };

        let parsed = serde_json::from_str::<Value>(resource_text)
            .ok()
            .map(|value| extract_text_content(&value))
            .unwrap_or_else(|| resource_text.to_string());
        resource_items.push(parsed);
    }

    if !resource_items.is_empty() {
        return resource_items.join("\n");
    }

    content
        .iter()
        .map(Value::to_string)
        .collect::<Vec<_>>()
        .join("\n")
}

#[derive(Clone, Debug)]
pub(crate) struct GatewayToolsPayload {
    pub tools: Vec<KontextTool>,
    pub errors: Vec<GatewayToolError>,
    pub elicitations: Vec<GatewayElicitation>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct GatewayToolError {
    pub server_id: String,
    #[serde(default)]
    pub server_name: Option<String>,
    #[serde(default)]
    pub reason: Option<String>,
}

#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "camelCase")]
pub struct GatewayElicitation {
    pub url: String,
    #[serde(default)]
    pub message: Option<String>,
    #[serde(default)]
    pub integration_id: Option<String>,
    #[serde(default)]
    pub integration_name: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct GatewayToolSummary {
    id: String,
    name: String,
    #[serde(default)]
    description: Option<String>,
    #[serde(default)]
    input_schema: Option<Value>,
    #[serde(default)]
    server: Option<GatewayToolServer>,
}

#[derive(Clone, Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct GatewayToolServer {
    #[serde(default)]
    id: Option<String>,
    #[serde(default)]
    name: Option<String>,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct RawTool {
    name: String,
    #[serde(default)]
    description: Option<String>,
    #[serde(default)]
    input_schema: Option<serde_json::Value>,
}

fn parse_tools_list_result(result: &Value) -> Result<Vec<KontextTool>, KontextDevError> {
    let tools = result
        .get("tools")
        .and_then(|value| value.as_array())
        .cloned()
        .unwrap_or_default();

    tools
        .into_iter()
        .map(|tool| {
            let raw: RawTool =
                serde_json::from_value(tool).map_err(|err| KontextDevError::ConnectSession {
                    message: format!("invalid tool payload: {err}"),
                })?;

            Ok(KontextTool {
                id: raw.name.clone(),
                name: raw.name,
                description: raw.description,
                input_schema: raw.input_schema,
                server: None,
            })
        })
        .collect()
}

pub(crate) fn parse_gateway_tools_payload(
    raw: &Value,
) -> Result<GatewayToolsPayload, KontextDevError> {
    let json_text =
        extract_json_resource_text(raw).ok_or_else(|| KontextDevError::ConnectSession {
            message: "SEARCH_TOOLS did not return JSON resource content".to_string(),
        })?;

    let parsed = serde_json::from_str::<Value>(&json_text).map_err(|err| {
        KontextDevError::ConnectSession {
            message: format!("SEARCH_TOOLS returned invalid JSON: {err}"),
        }
    })?;

    if let Some(items) = parsed.as_array() {
        let tools = items
            .iter()
            .cloned()
            .map(serde_json::from_value::<GatewayToolSummary>)
            .collect::<Result<Vec<_>, _>>()
            .map_err(|err| KontextDevError::ConnectSession {
                message: format!("SEARCH_TOOLS returned invalid tool entry: {err}"),
            })?
            .into_iter()
            .map(to_kontext_gateway_tool)
            .collect();
        return Ok(GatewayToolsPayload {
            tools,
            errors: Vec::new(),
            elicitations: Vec::new(),
        });
    }

    let Some(obj) = parsed.as_object() else {
        return Err(KontextDevError::ConnectSession {
            message: "SEARCH_TOOLS response was not a JSON array or object".to_string(),
        });
    };

    let tools = obj
        .get("items")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .map(serde_json::from_value::<GatewayToolSummary>)
        .collect::<Result<Vec<_>, _>>()
        .map_err(|err| KontextDevError::ConnectSession {
            message: format!("SEARCH_TOOLS items contained invalid tool data: {err}"),
        })?
        .into_iter()
        .map(to_kontext_gateway_tool)
        .collect::<Vec<_>>();

    let errors = obj
        .get("errors")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .filter_map(|value| serde_json::from_value::<GatewayToolError>(value).ok())
        .collect::<Vec<_>>();

    let elicitations = obj
        .get("elicitations")
        .and_then(Value::as_array)
        .cloned()
        .unwrap_or_default()
        .into_iter()
        .filter_map(|value| serde_json::from_value::<GatewayElicitation>(value).ok())
        .collect::<Vec<_>>();

    Ok(GatewayToolsPayload {
        tools,
        errors,
        elicitations,
    })
}

fn to_kontext_gateway_tool(summary: GatewayToolSummary) -> KontextTool {
    let server = summary.server.and_then(|server| {
        server.id.map(|id| KontextToolServer {
            id,
            name: server.name,
        })
    });

    KontextTool {
        id: summary.id,
        name: summary.name,
        description: summary.description,
        input_schema: summary.input_schema,
        server,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Arc;
    use std::sync::atomic::AtomicUsize;
    use std::sync::atomic::Ordering;
    use wiremock::Mock;
    use wiremock::MockServer;
    use wiremock::ResponseTemplate;
    use wiremock::matchers::method;
    use wiremock::matchers::path;

    #[test]
    fn parse_json_or_streamable_body_parses_json_payload() {
        let parsed = parse_json_or_streamable_body(
            r#"{"jsonrpc":"2.0","result":{"ok":true}}"#,
            "application/json",
        )
        .expect("json should parse");
        assert_eq!(parsed["result"]["ok"], Value::Bool(true));
    }

    #[test]
    fn parse_json_or_streamable_body_parses_sse_payload() {
        let parsed = parse_json_or_streamable_body(
            "event: message\ndata: {\"jsonrpc\":\"2.0\",\"result\":{\"sessionId\":\"abc\"}}\n\n",
            "text/event-stream",
        )
        .expect("sse should parse");
        assert_eq!(
            parsed["result"]["sessionId"],
            Value::String("abc".to_string())
        );
    }

    #[test]
    fn parse_json_or_streamable_body_falls_back_to_sse_when_content_type_is_json() {
        let parsed = parse_json_or_streamable_body(
            "data: {\"jsonrpc\":\"2.0\",\"result\":{\"tools\":[]}}\n\n",
            "application/json",
        )
        .expect("sse fallback should parse");
        assert_eq!(parsed["result"]["tools"], Value::Array(Vec::new()));
    }

    #[test]
    fn raw_tool_parses_input_schema_from_camel_case_key() {
        let parsed: RawTool = serde_json::from_value(serde_json::json!({
            "name": "SEARCH_TOOLS",
            "description": "Search available tools",
            "inputSchema": { "type": "object", "properties": { "limit": { "type": "number" } } }
        }))
        .expect("raw tool should deserialize");

        assert_eq!(parsed.name, "SEARCH_TOOLS");
        assert_eq!(
            parsed
                .input_schema
                .as_ref()
                .and_then(|value| value.get("type"))
                .and_then(Value::as_str),
            Some("object")
        );
    }

    #[test]
    fn extract_json_resource_text_skips_resource_items_without_resource_payload() {
        let payload = serde_json::json!({
            "content": [
                { "type": "resource" },
                {
                    "type": "resource",
                    "resource": {
                        "mimeType": "application/json",
                        "text": "{\"ok\":true}"
                    }
                }
            ]
        });

        assert_eq!(
            extract_json_resource_text(&payload),
            Some("{\"ok\":true}".to_string())
        );
    }

    #[test]
    fn runtime_integration_record_parses_user_token_connect_type() {
        let parsed: RuntimeIntegrationRecord = serde_json::from_value(serde_json::json!({
            "id": "convex-int",
            "name": "Convex",
            "url": "https://convex.example.com/mcp",
            "category": "gateway_remote_mcp",
            "connectType": "user_token"
        }))
        .expect("record should deserialize");

        assert_eq!(
            parsed.connect_type,
            RuntimeIntegrationConnectType::UserToken
        );
    }

    #[test]
    fn runtime_integration_record_rejects_unknown_connect_type() {
        let err = serde_json::from_value::<RuntimeIntegrationRecord>(serde_json::json!({
            "id": "convex-int",
            "name": "Convex",
            "url": "https://convex.example.com/mcp",
            "category": "gateway_remote_mcp",
            "connectType": "api_key"
        }))
        .expect_err("record should reject unknown connect type");

        assert!(err.to_string().contains("unknown variant"));
    }

    #[derive(Clone, Copy, Debug)]
    enum SessionFailureKind {
        HttpNotFound,
        JsonRpcNotFound,
    }

    fn create_test_mcp(server: &MockServer) -> KontextMcp {
        KontextMcp::new(KontextMcpConfig {
            client_session_id: "client-session".to_string(),
            client_id: "client-id".to_string(),
            redirect_uri: "http://localhost:3333/callback".to_string(),
            url: Some(format!("{}/mcp", server.uri())),
            server: Some(server.uri()),
            client_secret: None,
            scope: None,
            resource: None,
            session_key: None,
            integration_ui_url: None,
            integration_return_to: None,
            auth_timeout_seconds: None,
            open_connect_page_on_login: None,
            token_cache_path: None,
        })
    }

    async fn mount_retrying_tools_list_server(
        server: &MockServer,
        failure_kind: SessionFailureKind,
        recover_on_retry: bool,
    ) -> (Arc<AtomicUsize>, Arc<AtomicUsize>) {
        let initialize_calls = Arc::new(AtomicUsize::new(0));
        let tools_list_calls = Arc::new(AtomicUsize::new(0));

        let initialize_calls_for_mock = Arc::clone(&initialize_calls);
        let tools_list_calls_for_mock = Arc::clone(&tools_list_calls);
        Mock::given(method("POST"))
            .and(path("/mcp"))
            .respond_with(move |request: &wiremock::Request| {
                let payload: Value = serde_json::from_slice(&request.body)
                    .expect("MCP requests should be valid JSON payloads");
                let method = payload
                    .get("method")
                    .and_then(Value::as_str)
                    .expect("MCP requests should include method");

                match method {
                    "initialize" => {
                        let initialize_call =
                            initialize_calls_for_mock.fetch_add(1, Ordering::SeqCst);
                        let session_id = if initialize_call == 0 {
                            "stale-session"
                        } else {
                            "fresh-session"
                        };

                        ResponseTemplate::new(200)
                            .append_header("Mcp-Session-Id", session_id)
                            .set_body_json(json!({
                                "jsonrpc": "2.0",
                                "id": "initialize",
                                "result": {
                                    "sessionId": session_id
                                }
                            }))
                    }
                    "notifications/initialized" => {
                        ResponseTemplate::new(200).set_body_json(json!({
                            "jsonrpc": "2.0",
                            "result": {}
                        }))
                    }
                    "tools/list" => {
                        let tools_list_call =
                            tools_list_calls_for_mock.fetch_add(1, Ordering::SeqCst);
                        if tools_list_call == 0 || !recover_on_retry {
                            return match failure_kind {
                                SessionFailureKind::HttpNotFound => ResponseTemplate::new(400)
                                    .set_body_string(
                                        "Request rejected: Session stale-session not found",
                                    ),
                                SessionFailureKind::JsonRpcNotFound => ResponseTemplate::new(200)
                                    .set_body_json(json!({
                                        "jsonrpc": "2.0",
                                        "id": "list-tools",
                                        "error": {
                                            "code": -32000,
                                            "message": "Session stale-session not found"
                                        }
                                    })),
                            };
                        }

                        ResponseTemplate::new(200).set_body_json(json!({
                            "jsonrpc": "2.0",
                            "id": "list-tools",
                            "result": {
                                "tools": [{
                                    "name": "github.search",
                                    "description": "Search GitHub",
                                    "inputSchema": {
                                        "type": "object"
                                    }
                                }]
                            }
                        }))
                    }
                    _ => ResponseTemplate::new(500),
                }
            })
            .mount(server)
            .await;

        (initialize_calls, tools_list_calls)
    }

    #[tokio::test]
    async fn list_tools_recovers_from_http_session_not_found() {
        let server = MockServer::start().await;
        let (initialize_calls, tools_list_calls) =
            mount_retrying_tools_list_server(&server, SessionFailureKind::HttpNotFound, true).await;

        let mcp = create_test_mcp(&server);
        let tools = mcp
            .list_tools_with_access_token("access-token")
            .await
            .expect("HTTP session-not-found should recover");

        assert_eq!(
            tools,
            vec![KontextTool {
                id: "github.search".to_string(),
                name: "github.search".to_string(),
                description: Some("Search GitHub".to_string()),
                input_schema: Some(json!({
                    "type": "object"
                })),
                server: None,
            }]
        );
        assert_eq!(initialize_calls.load(Ordering::SeqCst), 2);
        assert_eq!(tools_list_calls.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn list_tools_recovers_from_jsonrpc_session_not_found() {
        let server = MockServer::start().await;
        let (initialize_calls, tools_list_calls) =
            mount_retrying_tools_list_server(&server, SessionFailureKind::JsonRpcNotFound, true)
                .await;

        let mcp = create_test_mcp(&server);
        let tools = mcp
            .list_tools_with_access_token("access-token")
            .await
            .expect("JSON-RPC session-not-found should recover");

        assert_eq!(tools.len(), 1);
        assert_eq!(initialize_calls.load(Ordering::SeqCst), 2);
        assert_eq!(tools_list_calls.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn list_tools_stale_session_retry_happens_once() {
        let server = MockServer::start().await;
        let (initialize_calls, tools_list_calls) =
            mount_retrying_tools_list_server(&server, SessionFailureKind::HttpNotFound, false)
                .await;

        let mcp = create_test_mcp(&server);
        let err = mcp
            .list_tools_with_access_token("access-token")
            .await
            .expect_err("recovery should fail when stale session persists");

        assert!(err.to_string().contains("Session stale-session not found"));
        assert_eq!(initialize_calls.load(Ordering::SeqCst), 2);
        assert_eq!(tools_list_calls.load(Ordering::SeqCst), 2);
    }
}