aion-server 0.31.0

//! The axum mount for `/assistant/mcp`.
//!
//! One route on the server's existing listener. It does NOT go through the
//! [`HttpCaller`](crate::api::http::auth::HttpCaller) extractor every other
//! route uses, and that is the point: `HttpCaller` resolves a HUMAN — an
//! operator, a token subject, a namespace membership — and this route accepts
//! only the session bearer this server minted for one agent. Running both would
//! mean a human token could reach a session's context; running the session
//! resolution alone means it cannot, whatever headers arrive.

use std::sync::Arc;

use axum::{
    Router,
    body::{Body, Bytes},
    http::{HeaderMap, Method, StatusCode, header::HeaderName},
    response::{IntoResponse, Response},
    routing::any,
};

use crate::ServerState;

use super::caller::{self, AssistantMcpAuthError};
use super::runtime::AssistantMcpRuntime;

/// The assistant-only MCP endpoint path.
pub const ASSISTANT_MCP_PATH: &str = "/assistant/mcp";

/// Build the assistant MCP route over an already-constructed runtime.
pub(crate) fn assistant_mcp_router(runtime: Arc<AssistantMcpRuntime>) -> Router<ServerState> {
    Router::new()
        .route(ASSISTANT_MCP_PATH, any(handle))
        .layer(axum::Extension(runtime))
}

async fn handle(
    axum::Extension(runtime): axum::Extension<Arc<AssistantMcpRuntime>>,
    axum::extract::State(state): axum::extract::State<ServerState>,
    method: Method,
    headers: HeaderMap,
    body: Bytes,
) -> Response {
    let pairs = header_pairs(&headers);
    let caller = match caller::resolve(state.assistant_sessions(), &pairs).await {
        Ok(caller) => caller,
        Err(error) => return unauthorized(&error),
    };
    let request = aion_mcp::HttpRequest::new(
        method.as_str().to_ascii_uppercase(),
        aion_mcp::protocol::headers::HeaderView::new(pairs),
        body.to_vec(),
    );
    into_axum(runtime.server().handle_http(&request, &caller).await)
}

/// The one refusal every un-admitted call gets.
///
/// The variant reaches the LOG and the single shared sentence reaches the
/// caller: an agent holding one session's token must not be able to tell
/// "no such session" from "not your session" from "that session ended", because
/// together those three enumerate other people's conversations.
fn unauthorized(error: &AssistantMcpAuthError) -> Response {
    tracing::warn!(
        target: "aion_server::assistant::mcp",
        %error,
        "a call on the assistant MCP route was not authorized"
    );
    (
        StatusCode::UNAUTHORIZED,
        axum::Json(
            // The authorization-refusal code this deployment already spells,
            // with an `error_type` naming THIS credential — so a client branches
            // on the type without a second wire code being minted for one route.
            aion_proto::WireError::grant_denied(AssistantMcpAuthError::client_message())
                .with_error_type(BEARER_REFUSED_TYPE),
        ),
    )
        .into_response()
}

/// `error_type` for a call that presented no usable session bearer.
pub(crate) const BEARER_REFUSED_TYPE: &str = "AssistantSessionBearerRefused";

/// Header values that are not valid UTF-8 are dropped rather than lossily
/// decoded — the same rule the general route follows, and for the same reason:
/// a header the server cannot read exactly is one it cannot prove agreement
/// with, and here that header might be the credential.
fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
    headers
        .iter()
        .filter_map(|(name, value)| {
            value
                .to_str()
                .ok()
                .map(|value| (name.as_str().to_owned(), value.to_owned()))
        })
        .collect()
}

fn into_axum(response: aion_mcp::HttpResponse) -> Response {
    let status = StatusCode::from_u16(response.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    let mut builder = Response::builder().status(status);
    for (name, value) in response.headers {
        match HeaderName::from_bytes(name.as_bytes()) {
            Ok(name) => builder = builder.header(name, value),
            Err(error) => {
                tracing::error!(
                    target: "aion_server::assistant::mcp",
                    header = %name,
                    "an assistant MCP response header name was not valid and was dropped: {error}"
                );
            }
        }
    }
    builder
        .body(Body::from(response.body))
        .unwrap_or_else(|error| {
            tracing::error!(
                target: "aion_server::assistant::mcp",
                "an assistant MCP response could not be assembled: {error}"
            );
            StatusCode::INTERNAL_SERVER_ERROR.into_response()
        })
}