magi-code 0.104.0

use super::{
    ACCEPT_POST, HEADER_MCP_PROTOCOL_VERSION, HEADER_MCP_SESSION_ID,
    sanitize_reqwest_error_with_timeout,
};
use crate::mcp::{McpError, McpResult, oauth::TokenProvider};
use reqwest::{
    StatusCode,
    blocking::{Client, Response},
    header::{ACCEPT, AUTHORIZATION, CONTENT_TYPE, HeaderMap},
};
use serde_json::Value;
use std::{
    sync::{Arc, Mutex},
    time::Duration,
};

pub(super) struct HttpPostContext {
    pub(super) client: Client,
    pub(super) url: String,
    pub(super) display_url: String,
    pub(super) server_name: Option<String>,
    pub(super) headers: HeaderMap,
    pub(super) token_provider: Option<TokenProvider>,
    pub(super) session_id: Arc<Mutex<Option<String>>>,
    pub(super) protocol_version: Arc<Mutex<Option<String>>>,
    pub(super) timeout: Duration,
}

impl HttpPostContext {
    pub(super) fn send_post_json(
        &self,
        body: &Value,
        force_refresh: bool,
        timeout: Duration,
    ) -> McpResult<Response> {
        self.request_with_session_headers(self.client.post(&self.url), force_refresh)?
            .header(CONTENT_TYPE, "application/json")
            .header(ACCEPT, ACCEPT_POST)
            .timeout(timeout)
            .json(body)
            .send()
            .map_err(|error| sanitize_reqwest_error_with_timeout(error, timeout))
    }

    fn request_with_session_headers(
        &self,
        builder: reqwest::blocking::RequestBuilder,
        force_refresh: bool,
    ) -> McpResult<reqwest::blocking::RequestBuilder> {
        let mut builder = builder.headers(self.headers.clone());
        if let Some(provider) = &self.token_provider {
            let token = if force_refresh {
                provider.force_refresh_access_token()?
            } else {
                provider.access_token()?
            };
            builder = builder.header(AUTHORIZATION, format!("Bearer {token}"));
        }
        if let Some(session_id) = self.session_id.lock().ok().and_then(|id| id.clone()) {
            builder = builder.header(HEADER_MCP_SESSION_ID, session_id);
        }
        if let Some(protocol_version) = self.protocol_version.lock().ok().and_then(|v| v.clone()) {
            builder = builder.header(HEADER_MCP_PROTOCOL_VERSION, protocol_version);
        }
        Ok(builder)
    }

    pub(super) fn send_reply(&self, body: &Value) -> McpResult<()> {
        let timeout = self.timeout.min(Duration::from_secs(1));
        let mut response = self.send_post_json(body, false, timeout)?;
        if response.status() == StatusCode::UNAUTHORIZED && self.token_provider.is_some() {
            response = self.send_post_json(body, true, timeout)?;
        }
        match response.status() {
            StatusCode::ACCEPTED | StatusCode::NO_CONTENT => Ok(()),
            _ => Err(McpError::Transport(
                "MCP HTTP server rejected a client reply".to_string(),
            )),
        }
    }
}