aion_server/assistant/mcp/
route.rs1use 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
26pub const ASSISTANT_MCP_PATH: &str = "/assistant/mcp";
28
29pub(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
56fn 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 aion_proto::WireError::grant_denied(AssistantMcpAuthError::client_message())
75 .with_error_type(BEARER_REFUSED_TYPE),
76 ),
77 )
78 .into_response()
79}
80
81pub(crate) const BEARER_REFUSED_TYPE: &str = "AssistantSessionBearerRefused";
83
84fn 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}