pub mod actions;
pub mod jsonrpc;
use anyhow::Result;
use axum::{
extract::{rejection::JsonRejection, DefaultBodyLimit, Json, State as AxumState},
http::StatusCode,
response::{IntoResponse, Response},
routing::post,
Router,
};
use jsonrpc::{ErrorCode, JsonRpcError, JsonRpcMessage, JsonRpcResponse};
use serde_json::Value;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::mpsc;
use tracing::{debug, error, info, trace, warn};
use crate::console_error;
#[cfg(not(target_arch = "wasm32"))]
use crate::llm::action_helper::call_llm;
#[cfg(not(target_arch = "wasm32"))]
use crate::llm::ollama_client::OllamaClient;
#[cfg(not(target_arch = "wasm32"))]
use crate::logging::emit::Log;
#[cfg(not(target_arch = "wasm32"))]
use crate::protocol::Event;
#[cfg(not(target_arch = "wasm32"))]
use crate::server::connection::ConnectionId;
#[cfg(not(target_arch = "wasm32"))]
use crate::server::McpProtocol;
#[cfg(not(target_arch = "wasm32"))]
use crate::state::app_state::AppState;
#[cfg(not(target_arch = "wasm32"))]
use crate::state::server::{ConnectionStatus, ProtocolConnectionInfo, ServerId};
#[cfg(not(target_arch = "wasm32"))]
use actions::{
MCP_INITIALIZE_EVENT, MCP_PROMPTS_GET_EVENT, MCP_PROMPTS_LIST_EVENT, MCP_RESOURCES_LIST_EVENT,
MCP_RESOURCES_READ_EVENT, MCP_TOOLS_CALL_EVENT, MCP_TOOLS_LIST_EVENT,
};
#[cfg(not(target_arch = "wasm32"))]
use jsonrpc::RequestId;
#[derive(Clone)]
pub struct McpServerState {
pub llm_client: OllamaClient,
pub app_state: Arc<AppState>,
pub status_tx: mpsc::UnboundedSender<String>,
pub server_id: ServerId,
pub protocol: Arc<McpProtocol>,
pub local_addr: SocketAddr,
}
pub const MAX_REQUEST_BODY_BYTES: usize = 2 * 1024 * 1024;
#[cfg(not(target_arch = "wasm32"))]
const BODY_TOO_LARGE_MESSAGE: &str = "request body too large";
#[cfg(not(target_arch = "wasm32"))]
const MAX_TRACE_BYTES: usize = 4096;
#[cfg(not(target_arch = "wasm32"))]
const SUPPORTED_PROTOCOL_VERSIONS: &[&str] = &["2024-11-05", "2025-03-26", "2025-06-18"];
#[cfg(not(target_arch = "wasm32"))]
const MCP_SERVER_BUSY_CODE: i32 = -32000;
#[cfg(not(target_arch = "wasm32"))]
const DEFAULT_PROTOCOL_VERSION: &str = "2024-11-05";
#[cfg(not(target_arch = "wasm32"))]
const FIRST_REQUEST_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
#[cfg(not(target_arch = "wasm32"))]
const IDLE_BETWEEN_REQUESTS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300);
#[cfg(not(target_arch = "wasm32"))]
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
#[cfg(not(target_arch = "wasm32"))]
const CONNECTION_CAP_REFUSAL: &[u8] = b"HTTP/1.1 503 Service Unavailable\r\n\
Content-Type: application/json\r\nContent-Length: 82\r\nRetry-After: 5\r\n\
Connection: close\r\n\r\n\
{\"jsonrpc\":\"2.0\",\"id\":null,\"error\":{\"code\":-32000,\"message\":\"server at capacity\"}}";
#[cfg(not(target_arch = "wasm32"))]
fn mcp_error_from_action(data: &Value) -> JsonRpcError {
let code = data
.get("code")
.and_then(|v| v.as_i64())
.and_then(|n| i32::try_from(n).ok())
.unwrap_or(ErrorCode::InternalError as i32);
JsonRpcError {
code,
message: data
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Internal error")
.to_string(),
data: data.get("data").cloned(),
}
}
#[cfg(not(target_arch = "wasm32"))]
fn llm_failure_error(state: &McpServerState, method: &str, e: anyhow::Error) -> JsonRpcError {
let failure = crate::utils::WireFailure::classify(&e);
let decision = if failure.is_overloaded() {
"fail_closed_llm_error_overloaded"
} else {
"fail_closed_llm_error_unavailable"
};
error!("MCP {} decision={}: {}", method, decision, e);
Log::new(Some(&state.status_tx)).error(format!("MCP {} decision={}: {}", method, decision, e));
if failure.is_overloaded() {
return JsonRpcError {
code: MCP_SERVER_BUSY_CODE,
message: failure.text().to_string(),
data: Some(serde_json::json!({"retryable": true})),
};
}
JsonRpcError {
code: ErrorCode::InternalError as i32,
message: failure.text().to_string(),
data: Some(serde_json::json!({"retryable": false})),
}
}
#[cfg(not(target_arch = "wasm32"))]
fn log_model_reject(state: &McpServerState, method: &str) {
warn!(
"MCP {} decision=model_reject: handler returned mcp_error",
method
);
Log::new(Some(&state.status_tx)).warn(format!(
"MCP {} decision=model_reject (handler returned mcp_error)",
method
));
}
#[cfg(not(target_arch = "wasm32"))]
fn log_model_no_answer(state: &McpServerState, method: &str, fallback: &str) {
warn!(
"MCP {} decision=model_no_answer: replying with {}",
method, fallback
);
Log::new(Some(&state.status_tx)).warn(format!(
"MCP {} decision=model_no_answer (replying with {})",
method, fallback
));
}
#[cfg(not(target_arch = "wasm32"))]
fn recover_request_id(payload: &Value) -> Option<RequestId> {
match payload.get("id")? {
Value::String(s) => Some(RequestId::String(s.clone())),
Value::Number(n) => n.as_i64().map(RequestId::Number),
_ => None,
}
}
pub struct McpServer;
#[cfg(not(target_arch = "wasm32"))]
impl McpServer {
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: ServerId,
) -> Result<SocketAddr> {
let listener = tokio::net::TcpListener::bind(listen_addr).await?;
let local_addr = listener.local_addr()?;
info!("MCP server (JSON-RPC 2.0) listening on {}", local_addr);
Log::new(Some(&status_tx)).info(format!("MCP server listening on {}", local_addr));
let protocol = Arc::new(McpProtocol::new());
let task_registrar = app_state.clone();
let server_state = McpServerState {
llm_client,
app_state,
status_tx: status_tx.clone(),
server_id,
protocol,
local_addr,
};
let app = Router::new()
.route("/", post(handle_jsonrpc))
.layer(DefaultBodyLimit::max(MAX_REQUEST_BODY_BYTES))
.with_state(server_state);
let backend = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
let backend_addr = backend.local_addr()?;
let backend_status = status_tx.clone();
let backend_handle = tokio::spawn(async move {
if let Err(e) = axum::serve(backend, app).await {
console_error!(backend_status, "MCP server error: {}", e);
}
});
let relay_state = task_registrar.clone();
let accept_handle = tokio::spawn(async move {
serve_screened_mcp(listener, backend_addr, status_tx, relay_state, server_id).await;
});
task_registrar
.register_server_task(server_id, backend_handle)
.await;
task_registrar
.register_server_task(server_id, accept_handle)
.await;
Ok(local_addr)
}
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_jsonrpc(
AxumState(state): AxumState<McpServerState>,
payload: Result<Json<Value>, JsonRejection>,
) -> Response {
let payload = match payload {
Ok(Json(payload)) => payload,
Err(rejection) if rejection.status() == StatusCode::PAYLOAD_TOO_LARGE => {
warn!(
"MCP request body over {} bytes decision=fail_closed_body_too_large",
MAX_REQUEST_BODY_BYTES
);
Log::new(Some(&state.status_tx)).warn(format!(
"MCP refused a request body over {} bytes (decision=fail_closed_body_too_large)",
MAX_REQUEST_BODY_BYTES
));
let body = serde_json::json!({
"jsonrpc": "2.0",
"id": null,
"error": {
"code": ErrorCode::InvalidRequest.as_i32(),
"message": BODY_TOO_LARGE_MESSAGE,
},
});
return (StatusCode::PAYLOAD_TOO_LARGE, Json(body)).into_response();
}
Err(rejection) => return rejection.into_response(),
};
trace!(
"MCP received JSON-RPC request: {}",
serde_json::to_string_pretty(&payload).unwrap_or_else(|_| payload.to_string())
);
let mut trace_body = serde_json::to_string(&payload).unwrap_or_default();
if trace_body.len() > MAX_TRACE_BYTES {
let cut = trace_body
.char_indices()
.take_while(|(i, _)| *i <= MAX_TRACE_BYTES)
.last()
.map(|(i, _)| i)
.unwrap_or(0);
trace_body.truncate(cut);
trace_body.push_str("… (truncated)");
}
Log::new(Some(&state.status_tx)).trace(format!("MCP received: {}", trace_body));
let message = match JsonRpcMessage::from_value(payload.clone()) {
Ok(msg) => msg,
Err(e) => {
error!("Failed to parse JSON-RPC message: {:?}", e);
let response = JsonRpcResponse::error(recover_request_id(&payload), e);
return Json(response).into_response();
}
};
match message {
JsonRpcMessage::Request(req) => {
let request_id = req.id.clone();
let method = req.method.clone();
debug!("MCP request: method={}, id={:?}", method, request_id);
let result = match method.as_str() {
"initialize" => handle_initialize(&state, req.params).await,
"ping" => handle_ping(),
"resources/list" => handle_resources_list(&state).await,
"resources/read" => handle_resources_read(&state, req.params).await,
"resources/subscribe" => handle_resources_subscribe(&state, req.params).await,
"resources/unsubscribe" => handle_resources_unsubscribe(&state, req.params).await,
"resources/templates/list" => handle_resources_templates_list(&state).await,
"tools/list" => handle_tools_list(&state).await,
"tools/call" => handle_tools_call(&state, req.params).await,
"prompts/list" => handle_prompts_list(&state).await,
"prompts/get" => handle_prompts_get(&state, req.params).await,
"logging/setLevel" => handle_logging_set_level(&state, req.params).await,
"completion/complete" => handle_completion(&state, req.params).await,
_ => Err(JsonRpcError::new(ErrorCode::MethodNotFound)),
};
let response = match result {
Ok(value) => {
trace!(
"MCP response success: {}",
serde_json::to_string_pretty(&value).unwrap_or_default()
);
JsonRpcResponse::success(request_id, value)
}
Err(e) => {
error!("MCP error: code={}, message={}", e.code, e.message);
Log::new(Some(&state.status_tx)).error(format!("MCP error: {}", e.message));
JsonRpcResponse::error(request_id, e)
}
};
Json(response).into_response()
}
JsonRpcMessage::Notification(notif) => {
let method = notif.method.clone();
debug!("MCP notification: method={}", method);
match method.as_str() {
"notifications/initialized" => {
handle_initialized(&state).await;
}
"notifications/cancelled" => {
handle_cancelled(&state, notif.params).await;
}
"notifications/progress" => {
handle_progress(&state, notif.params).await;
}
_ => {
debug!("Unknown MCP notification: {}", method);
}
}
StatusCode::NO_CONTENT.into_response()
}
}
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_initialize(
state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let connection_id = ConnectionId::new(state.app_state.get_next_unified_id().await);
let result = handle_initialize_inner(state, params, connection_id).await;
state
.app_state
.close_connection_on_server(state.server_id, connection_id)
.await;
result
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_initialize_inner(
state: &McpServerState,
params: Option<Value>,
connection_id: ConnectionId,
) -> Result<Value, JsonRpcError> {
info!("MCP initialize request");
let client_info = params
.as_ref()
.and_then(|p| p.get("clientInfo"))
.map(|c| c.to_string())
.unwrap_or_else(|| "unknown".to_string());
debug!("MCP client: {}", client_info);
let _ = state
.status_tx
.send(format!("→ MCP client initializing: {}", client_info));
state
.app_state
.add_connection_to_server(
state.server_id,
crate::state::ConnectionState {
id: connection_id,
remote_addr: state.local_addr, local_addr: state.local_addr,
bytes_sent: 0,
bytes_received: 0,
packets_sent: 0,
packets_received: 0,
last_activity: crate::utils::clock::Instant::now(),
status: ConnectionStatus::Active,
status_changed_at: crate::utils::clock::Instant::now(),
protocol_info: ProtocolConnectionInfo::empty(),
},
)
.await;
let requested_version = params
.as_ref()
.and_then(|p| p.get("protocolVersion"))
.and_then(|v| v.as_str())
.unwrap_or("unknown")
.to_string();
let event = Event::new(
&MCP_INITIALIZE_EVENT,
serde_json::json!({
"method": "initialize",
"client_info": client_info,
"protocol_version": requested_version,
"capabilities": params.as_ref().and_then(|p| p.get("capabilities")),
}),
);
let protocol = state.protocol.clone();
Log::new(Some(&state.status_tx)).debug("MCP calling LLM for initialize request");
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
Some(connection_id),
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "initialize", e));
}
};
let log = Log::new(Some(&state.status_tx));
for message in &execution_result.messages {
log.info(format!("{}", message));
}
log.debug(format!(
"MCP got {} protocol results",
execution_result.protocol_results.len()
));
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "initialize");
return Err(mcp_error_from_action(data));
}
if name == "mcp_initialize" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "initialize", "the negotiated protocol version");
let agreed_version = if SUPPORTED_PROTOCOL_VERSIONS.contains(&requested_version.as_str()) {
requested_version.as_str()
} else {
DEFAULT_PROTOCOL_VERSION
};
Ok(serde_json::json!({
"protocolVersion": agreed_version,
"capabilities": {
"resources": {},
"tools": {},
"prompts": {}
},
"serverInfo": {
"name": "netget-mcp",
"version": "0.1.0"
}
}))
}
#[cfg(not(target_arch = "wasm32"))]
fn handle_ping() -> Result<Value, JsonRpcError> {
Ok(serde_json::json!({}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_resources_list(state: &McpServerState) -> Result<Value, JsonRpcError> {
Log::new(Some(&state.status_tx)).debug("MCP resources/list request");
let event = Event::new(
&MCP_RESOURCES_LIST_EVENT,
serde_json::json!({
"method": "resources/list",
}),
);
let protocol = state.protocol.clone();
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
None,
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "resources/list", e));
}
};
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "resources/list");
return Err(mcp_error_from_action(data));
}
if name == "mcp_resources_list" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "resources/list", "an empty resources list");
Ok(serde_json::json!({"resources": []}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_resources_read(
state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let uri = params
.as_ref()
.and_then(|p| p.get("uri"))
.and_then(|u| u.as_str())
.ok_or_else(|| JsonRpcError::new(ErrorCode::InvalidParams))?;
Log::new(Some(&state.status_tx)).debug(format!("MCP resources/read: {}", uri));
let event = Event::new(
&MCP_RESOURCES_READ_EVENT,
serde_json::json!({
"method": "resources/read",
"uri": uri,
}),
);
let protocol = state.protocol.clone();
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
None,
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "resources/read", e));
}
};
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "resources/read");
return Err(mcp_error_from_action(data));
}
if name == "mcp_resources_read" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "resources/read", "a JSON-RPC error");
Err(JsonRpcError::custom(
ErrorCode::InternalError,
format!("Resource not found: {}", uri),
))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_resources_subscribe(
_state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let uri = params
.as_ref()
.and_then(|p| p.get("uri"))
.and_then(|u| u.as_str())
.ok_or_else(|| JsonRpcError::new(ErrorCode::InvalidParams))?;
debug!("MCP resources/subscribe: uri={}", uri);
Ok(serde_json::json!({}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_resources_unsubscribe(
_state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let uri = params
.as_ref()
.and_then(|p| p.get("uri"))
.and_then(|u| u.as_str())
.ok_or_else(|| JsonRpcError::new(ErrorCode::InvalidParams))?;
debug!("MCP resources/unsubscribe: uri={}", uri);
Ok(serde_json::json!({}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_resources_templates_list(_state: &McpServerState) -> Result<Value, JsonRpcError> {
debug!("MCP resources/templates/list");
Ok(serde_json::json!({
"resourceTemplates": []
}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_tools_list(state: &McpServerState) -> Result<Value, JsonRpcError> {
Log::new(Some(&state.status_tx)).debug("MCP tools/list request");
let event = Event::new(
&MCP_TOOLS_LIST_EVENT,
serde_json::json!({
"method": "tools/list",
}),
);
let protocol = state.protocol.clone();
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
None,
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "tools/list", e));
}
};
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "tools/list");
return Err(mcp_error_from_action(data));
}
if name == "mcp_tools_list" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "tools/list", "an empty tools list");
Ok(serde_json::json!({"tools": []}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_tools_call(
state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let tool_name = params
.as_ref()
.and_then(|p| p.get("name"))
.and_then(|n| n.as_str())
.ok_or_else(|| JsonRpcError::new(ErrorCode::InvalidParams))?;
let tool_arguments = params.as_ref().and_then(|p| p.get("arguments"));
Log::new(Some(&state.status_tx)).debug(format!("MCP tools/call: {}", tool_name));
let event = Event::new(
&MCP_TOOLS_CALL_EVENT,
serde_json::json!({
"method": "tools/call",
"name": tool_name,
"arguments": tool_arguments,
}),
);
let protocol = state.protocol.clone();
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
None,
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "tools/call", e));
}
};
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "tools/call");
return Err(mcp_error_from_action(data));
}
if name == "mcp_tools_call" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "tools/call", "a JSON-RPC error");
Err(JsonRpcError::custom(
ErrorCode::InternalError,
format!("Tool execution failed: {}", tool_name),
))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_prompts_list(state: &McpServerState) -> Result<Value, JsonRpcError> {
Log::new(Some(&state.status_tx)).debug("MCP prompts/list request");
let event = Event::new(
&MCP_PROMPTS_LIST_EVENT,
serde_json::json!({
"method": "prompts/list",
}),
);
let protocol = state.protocol.clone();
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
None,
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "prompts/list", e));
}
};
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "prompts/list");
return Err(mcp_error_from_action(data));
}
if name == "mcp_prompts_list" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "prompts/list", "an empty prompts list");
Ok(serde_json::json!({"prompts": []}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_prompts_get(
state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let prompt_name = params
.as_ref()
.and_then(|p| p.get("name"))
.and_then(|n| n.as_str())
.ok_or_else(|| JsonRpcError::new(ErrorCode::InvalidParams))?;
let prompt_arguments = params.as_ref().and_then(|p| p.get("arguments"));
Log::new(Some(&state.status_tx)).debug(format!("MCP prompts/get: {}", prompt_name));
let event = Event::new(
&MCP_PROMPTS_GET_EVENT,
serde_json::json!({
"method": "prompts/get",
"name": prompt_name,
"arguments": prompt_arguments,
}),
);
let protocol = state.protocol.clone();
let execution_result = match call_llm(
&state.llm_client,
&state.app_state,
state.server_id,
None,
&event,
protocol.as_ref(),
)
.await
{
Ok(result) => result,
Err(e) => {
return Err(llm_failure_error(state, "prompts/get", e));
}
};
for protocol_result in &execution_result.protocol_results {
use crate::llm::actions::protocol_trait::ActionResult;
if let ActionResult::Custom { name, data } = protocol_result {
if name == "mcp_error" {
log_model_reject(state, "prompts/get");
return Err(mcp_error_from_action(data));
}
if name == "mcp_prompts_get" {
if let Some(response) = data.get("response") {
return Ok(response.clone());
}
}
}
}
log_model_no_answer(state, "prompts/get", "a JSON-RPC error");
Err(JsonRpcError::custom(
ErrorCode::InternalError,
format!("Prompt not found: {}", prompt_name),
))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_logging_set_level(
state: &McpServerState,
params: Option<Value>,
) -> Result<Value, JsonRpcError> {
let level = params
.as_ref()
.and_then(|p| p.get("level"))
.and_then(|l| l.as_str())
.unwrap_or("info");
Log::new(Some(&state.status_tx)).info(format!("MCP log level set to: {}", level));
Ok(serde_json::json!({}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_completion(
_state: &McpServerState,
_params: Option<Value>,
) -> Result<Value, JsonRpcError> {
debug!("MCP completion/complete");
Ok(serde_json::json!({
"completion": {
"values": [],
"total": 0,
"hasMore": false
}
}))
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_initialized(state: &McpServerState) {
Log::new(Some(&state.status_tx)).info("MCP client initialized");
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_cancelled(state: &McpServerState, params: Option<Value>) {
if let Some(req_id) = params.as_ref().and_then(|p| p.get("requestId")) {
Log::new(Some(&state.status_tx)).debug(format!("MCP cancelled: {:?}", req_id));
}
}
#[cfg(not(target_arch = "wasm32"))]
async fn handle_progress(_state: &McpServerState, params: Option<Value>) {
if let Some(progress) = params {
trace!("MCP progress: {}", progress);
}
}
#[cfg(not(target_arch = "wasm32"))]
async fn serve_screened_mcp(
listener: tokio::net::TcpListener,
backend_addr: SocketAddr,
status_tx: mpsc::UnboundedSender<String>,
app_state: std::sync::Arc<crate::state::app_state::AppState>,
server_id: crate::state::ServerId,
) {
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
loop {
let (peer_stream, peer_addr, permit) = match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"MCP",
Some(&status_tx),
)
.await
{
Ok(triple) => triple,
Err(e) => {
console_error!(status_tx, "MCP accept failed, listener stopped: {}", e);
return;
}
};
let status = status_tx.clone();
let task_owner = app_state.clone();
task_owner
.spawn_server_task(server_id, async move {
let _permit = permit;
relay_mcp_connection(peer_stream, peer_addr, backend_addr, status).await;
})
.await;
}
}
#[cfg(not(target_arch = "wasm32"))]
async fn relay_mcp_connection(
peer_stream: tokio::net::TcpStream,
peer_addr: SocketAddr,
backend_addr: SocketAddr,
status_tx: mpsc::UnboundedSender<String>,
) {
use std::sync::atomic::{AtomicBool, Ordering};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let backend_stream = match tokio::net::TcpStream::connect(backend_addr).await {
Ok(stream) => stream,
Err(e) => {
console_error!(status_tx, "MCP could not reach its own backend: {}", e);
return;
}
};
let (mut peer_read, mut peer_write) = tokio::io::split(peer_stream);
let (mut backend_read, mut backend_write) = tokio::io::split(backend_stream);
let awaiting_response = Arc::new(AtomicBool::new(false));
let to_backend_awaiting = Arc::clone(&awaiting_response);
let to_backend_status = status_tx.clone();
let to_backend = tokio::spawn(async move {
let mut buf = vec![0u8; 16 * 1024];
let mut seen_request = false;
loop {
let bound = if seen_request {
IDLE_BETWEEN_REQUESTS_TIMEOUT
} else {
FIRST_REQUEST_READ_TIMEOUT
};
let read = loop {
match tokio::time::timeout(bound, peer_read.read(&mut buf)).await {
Ok(read) => break Some(read),
Err(_) => {
if to_backend_awaiting.load(Ordering::SeqCst) {
continue;
}
Log::new(Some(&to_backend_status)).info(format!(
"MCP peer {} sent nothing for {}s; closing idle connection",
peer_addr,
bound.as_secs()
));
break None;
}
}
};
let n = match read {
Some(Ok(0)) | None => break,
Some(Ok(n)) => n,
Some(Err(_)) => break,
};
seen_request = true;
to_backend_awaiting.store(true, Ordering::SeqCst);
if backend_write.write_all(&buf[..n]).await.is_err() {
break;
}
}
let _ = backend_write.shutdown().await;
});
let to_peer_awaiting = Arc::clone(&awaiting_response);
let to_peer = tokio::spawn(async move {
let mut buf = vec![0u8; 16 * 1024];
loop {
match backend_read.read(&mut buf).await {
Ok(0) | Err(_) => break,
Ok(n) => {
if peer_write.write_all(&buf[..n]).await.is_err() {
break;
}
to_peer_awaiting.store(false, Ordering::SeqCst);
}
}
}
let _ = peer_write.shutdown().await;
});
tokio::select! {
_ = to_backend => {}
_ = to_peer => {}
}
}