Skip to main content

aion_server/assistant/mcp/
route.rs

1//! The axum mount for `/assistant/mcp`.
2//!
3//! One route on the server's existing listener. It does NOT go through the
4//! [`HttpCaller`](crate::api::http::auth::HttpCaller) extractor every other
5//! route uses, and that is the point: `HttpCaller` resolves a HUMAN — an
6//! operator, a token subject, a namespace membership — and this route accepts
7//! only the session bearer this server minted for one agent. Running both would
8//! mean a human token could reach a session's context; running the session
9//! resolution alone means it cannot, whatever headers arrive.
10
11use std::sync::Arc;
12
13use axum::{
14    Router,
15    body::{Body, Bytes},
16    http::{HeaderMap, Method, StatusCode, header::HeaderName},
17    response::{IntoResponse, Response},
18    routing::any,
19};
20
21use crate::ServerState;
22
23use super::caller::{self, AssistantMcpAuthError};
24use super::runtime::AssistantMcpRuntime;
25
26/// The assistant-only MCP endpoint path.
27pub const ASSISTANT_MCP_PATH: &str = "/assistant/mcp";
28
29/// Build the assistant MCP route over an already-constructed runtime.
30pub(crate) fn assistant_mcp_router(runtime: Arc<AssistantMcpRuntime>) -> Router<ServerState> {
31    Router::new()
32        .route(ASSISTANT_MCP_PATH, any(handle))
33        .layer(axum::Extension(runtime))
34}
35
36async fn handle(
37    axum::Extension(runtime): axum::Extension<Arc<AssistantMcpRuntime>>,
38    axum::extract::State(state): axum::extract::State<ServerState>,
39    method: Method,
40    headers: HeaderMap,
41    body: Bytes,
42) -> Response {
43    let pairs = header_pairs(&headers);
44    let caller = match caller::resolve(state.assistant_sessions(), &pairs).await {
45        Ok(caller) => caller,
46        Err(error) => return unauthorized(&error),
47    };
48    let request = aion_mcp::HttpRequest::new(
49        method.as_str().to_ascii_uppercase(),
50        aion_mcp::protocol::headers::HeaderView::new(pairs),
51        body.to_vec(),
52    );
53    into_axum(runtime.server().handle_http(&request, &caller).await)
54}
55
56/// The one refusal every un-admitted call gets.
57///
58/// The variant reaches the LOG and the single shared sentence reaches the
59/// caller: an agent holding one session's token must not be able to tell
60/// "no such session" from "not your session" from "that session ended", because
61/// together those three enumerate other people's conversations.
62fn unauthorized(error: &AssistantMcpAuthError) -> Response {
63    tracing::warn!(
64        target: "aion_server::assistant::mcp",
65        %error,
66        "a call on the assistant MCP route was not authorized"
67    );
68    (
69        StatusCode::UNAUTHORIZED,
70        axum::Json(
71            // The authorization-refusal code this deployment already spells,
72            // with an `error_type` naming THIS credential — so a client branches
73            // on the type without a second wire code being minted for one route.
74            aion_proto::WireError::grant_denied(AssistantMcpAuthError::client_message())
75                .with_error_type(BEARER_REFUSED_TYPE),
76        ),
77    )
78        .into_response()
79}
80
81/// `error_type` for a call that presented no usable session bearer.
82pub(crate) const BEARER_REFUSED_TYPE: &str = "AssistantSessionBearerRefused";
83
84/// Header values that are not valid UTF-8 are dropped rather than lossily
85/// decoded — the same rule the general route follows, and for the same reason:
86/// a header the server cannot read exactly is one it cannot prove agreement
87/// with, and here that header might be the credential.
88fn header_pairs(headers: &HeaderMap) -> Vec<(String, String)> {
89    headers
90        .iter()
91        .filter_map(|(name, value)| {
92            value
93                .to_str()
94                .ok()
95                .map(|value| (name.as_str().to_owned(), value.to_owned()))
96        })
97        .collect()
98}
99
100fn into_axum(response: aion_mcp::HttpResponse) -> Response {
101    let status = StatusCode::from_u16(response.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
102    let mut builder = Response::builder().status(status);
103    for (name, value) in response.headers {
104        match HeaderName::from_bytes(name.as_bytes()) {
105            Ok(name) => builder = builder.header(name, value),
106            Err(error) => {
107                tracing::error!(
108                    target: "aion_server::assistant::mcp",
109                    header = %name,
110                    "an assistant MCP response header name was not valid and was dropped: {error}"
111                );
112            }
113        }
114    }
115    builder
116        .body(Body::from(response.body))
117        .unwrap_or_else(|error| {
118            tracing::error!(
119                target: "aion_server::assistant::mcp",
120                "an assistant MCP response could not be assembled: {error}"
121            );
122            StatusCode::INTERNAL_SERVER_ERROR.into_response()
123        })
124}