use std::{
borrow::Cow,
sync::{Arc, Mutex},
time::Instant,
};
use axum::{
Router,
body::Body,
extract::State,
http::{
Request, StatusCode,
header::{ORIGIN, RETRY_AFTER},
},
middleware::{self, Next},
response::{IntoResponse, Response},
};
use base64::{Engine as _, prelude::BASE64_STANDARD};
use rmcp::{
ErrorData as McpError, RoleServer, ServerHandler,
model::{
CacheScope, CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock,
DiscoverResult, Implementation, ListToolsResult, MetaObject, PaginatedRequestParams,
ProtocolVersion, ResourceContents, ServerCapabilities, ServerInfo, Tool,
},
service::RequestContext,
transport::streamable_http_server::{
StreamableHttpServerConfig, StreamableHttpService, session::never::NeverSessionManager,
},
};
use serde_json::{Value, json};
use tokio::sync::Semaphore;
use tracing::{info, warn};
use crate::{
manifest::Catalog,
upstream::{Upstream, UpstreamResponse},
};
const CACHE_META_KEY: &str = "space.bitview/upstreamCache";
const MAX_REQUEST_BYTES: usize = 64 * 1024;
const MAX_CONCURRENCY: usize = 16;
const REQUESTS_PER_SECOND: u32 = 50;
const CATALOG_TTL_MS: u64 = 3_600_000;
const MAX_UPSTREAM_ERROR_BYTES: usize = 2_048;
static SUPPORTED_PROTOCOL_VERSIONS: [ProtocolVersion; 1] = [ProtocolVersion::V_2026_07_28];
struct AppState {
catalog: Catalog,
upstream: Upstream,
concurrency: Arc<Semaphore>,
rate_limit: Mutex<RateLimit>,
}
struct RateLimit {
window_started: Instant,
used: u32,
}
#[derive(Clone)]
struct BrkMcp {
state: Arc<AppState>,
}
impl AppState {
fn new(api_bases: Vec<String>, catalog: Catalog) -> Self {
let upstream = Upstream::new(api_bases);
Self {
concurrency: Arc::new(Semaphore::new(MAX_CONCURRENCY)),
rate_limit: Mutex::new(RateLimit {
window_started: Instant::now(),
used: 0,
}),
catalog,
upstream,
}
}
fn admit(&self) -> bool {
let mut limit = self.rate_limit.lock().expect("rate limit mutex poisoned");
if limit.window_started.elapsed().as_secs() >= 1 {
limit.window_started = Instant::now();
limit.used = 0;
}
if limit.used >= REQUESTS_PER_SECOND {
return false;
}
limit.used += 1;
true
}
}
pub fn router(api_bases: Vec<String>, catalog: Catalog) -> Router {
let transport_config = StreamableHttpServerConfig::default()
.disable_allowed_hosts()
.with_sse_keep_alive(None)
.with_sse_retry(None)
.with_legacy_session_mode(false)
.with_json_response(true)
.with_max_request_body_bytes(MAX_REQUEST_BYTES)
.with_stateless_protocol_metadata_required(true);
let state = Arc::new(AppState::new(api_bases, catalog));
let handler = BrkMcp {
state: state.clone(),
};
let service = StreamableHttpService::new(
move || Ok(handler.clone()),
Arc::new(NeverSessionManager::default()),
transport_config,
);
Router::new()
.route_service("/", service)
.layer(middleware::from_fn_with_state(state, gateway_guard))
}
impl BrkMcp {
fn server_info(&self) -> ServerInfo {
ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
.with_protocol_version(ProtocolVersion::V_2026_07_28)
.with_server_info(Implementation::new("brk_mcp", env!("CARGO_PKG_VERSION")))
.with_instructions(
"Read-only Bitcoin analytics tools backed by BRK REST GET endpoints. \
Calls use the configured public Cloudflare-cached API and never mutate BRK state. \
Tool descriptions preserve the OpenAPI documentation and identify their underlying \
REST operation. Use available list, search, and info operations to discover \
identifiers before querying dynamic datasets. Binary responses are returned as \
embedded MCP resources.",
)
}
fn tool_error(&self, message: impl Into<String>) -> CallToolResponse {
CallToolResult::error(vec![ContentBlock::text(message)])
.with_meta(Some(server_meta()))
.into()
}
fn render_upstream(
&self,
response: UpstreamResponse,
has_output_schema: bool,
) -> CallToolResponse {
if !(200..300).contains(&response.status) {
let detail = String::from_utf8_lossy(
&response.body[..response.body.len().min(MAX_UPSTREAM_ERROR_BYTES)],
);
let detail = detail.trim();
let message = if detail.is_empty() {
format!("BRK API returned HTTP {}", response.status)
} else {
format!("BRK API returned HTTP {}: {detail}", response.status)
};
return self.tool_error(message);
}
let content_type = response
.content_type
.split(';')
.next()
.unwrap_or("application/octet-stream")
.trim()
.to_ascii_lowercase();
let meta = Some(upstream_meta(&response));
if content_type == "application/json" || content_type.ends_with("+json") {
return match serde_json::from_slice::<Value>(&response.body) {
Ok(value) => CallToolResult::structured(value).with_meta(meta).into(),
Err(_) => self.tool_error("BRK API returned invalid JSON"),
};
}
if content_type.starts_with("text/") {
return match String::from_utf8(response.body) {
Ok(text) => {
let mut result =
CallToolResult::success(vec![ContentBlock::text(text.clone())]);
if has_output_schema {
result.structured_content = Some(Value::String(text));
}
result.with_meta(meta).into()
}
Err(_) => self.tool_error("BRK API returned invalid UTF-8 text"),
};
}
let resource = ResourceContents::blob(BASE64_STANDARD.encode(response.body), response.url)
.with_mime_type(content_type);
CallToolResult::success(vec![ContentBlock::resource(resource)])
.with_meta(meta)
.into()
}
}
impl ServerHandler for BrkMcp {
fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
Cow::Borrowed(&SUPPORTED_PROTOCOL_VERSIONS)
}
fn discover(
&self,
_context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<DiscoverResult, McpError>> + Send + '_ {
std::future::ready(Ok(DiscoverResult::from_server_info(
SUPPORTED_PROTOCOL_VERSIONS.to_vec(),
self.server_info(),
)
.with_ttl_ms(CATALOG_TTL_MS)
.with_cache_scope(CacheScope::Public)))
}
fn list_tools(
&self,
request: Option<PaginatedRequestParams>,
_context: RequestContext<RoleServer>,
) -> impl Future<Output = Result<ListToolsResult, McpError>> + Send + '_ {
let result = if request.and_then(|request| request.cursor).is_some() {
Err(McpError::invalid_params(
"This complete tool catalog does not accept a cursor",
None,
))
} else {
let mut result = ListToolsResult::with_all_items(self.state.catalog.tools().to_vec())
.with_ttl_ms(CATALOG_TTL_MS)
.with_cache_scope(CacheScope::Public);
result.meta = Some(server_meta());
Ok(result)
};
std::future::ready(result)
}
fn get_tool(&self, name: &str) -> Option<Tool> {
self.state
.catalog
.operation(name)
.map(|operation| operation.tool.clone())
}
async fn call_tool(
&self,
request: CallToolRequestParams,
_context: RequestContext<RoleServer>,
) -> Result<CallToolResponse, McpError> {
let name = request.name;
let operation = self
.state
.catalog
.operation(&name)
.ok_or_else(|| McpError::invalid_params("Unknown tool name", None))?;
let arguments = request.arguments.unwrap_or_default();
let arguments = operation
.validate_arguments(&arguments)
.map_err(|error| McpError::invalid_params(error, None))?;
let has_output_schema = operation.tool.output_schema.is_some();
let prepared = self
.state
.upstream
.prepare(operation, arguments)
.map_err(|error| McpError::invalid_params(error, None))?;
let permit = self
.state
.concurrency
.clone()
.try_acquire_owned()
.map_err(|_| McpError::internal_error("MCP proxy is at its concurrency limit", None))?;
let upstream = self.state.upstream.clone();
let response = match tokio::task::spawn_blocking(move || {
let _permit = permit;
upstream.fetch(prepared)
})
.await
{
Ok(Ok(response)) => response,
Ok(Err(error)) => return Ok(self.tool_error(error)),
Err(_) => return Ok(self.tool_error("BRK API request task failed")),
};
if let Some(cache_status) = &response.cache_status {
info!(
tool = name.as_ref(),
cache_status,
cache_age = response.cache_age.as_deref().unwrap_or(""),
"BRK API response"
);
} else {
warn!(
tool = name.as_ref(),
"BRK API response did not include CF-Cache-Status"
);
}
Ok(self.render_upstream(response, has_output_schema))
}
fn get_info(&self) -> ServerInfo {
self.server_info()
}
}
async fn gateway_guard(
State(state): State<Arc<AppState>>,
request: Request<Body>,
next: Next,
) -> Response {
if request.headers().contains_key(ORIGIN) {
return (StatusCode::FORBIDDEN, "Forbidden origin").into_response();
}
if !state.admit() {
let mut response =
(StatusCode::TOO_MANY_REQUESTS, "Request rate limit exceeded").into_response();
response.headers_mut().insert(
RETRY_AFTER,
"1".parse().expect("static retry header is valid"),
);
return response;
}
next.run(request).await
}
fn server_meta() -> MetaObject {
let mut meta = MetaObject::new();
meta.insert(
"io.modelcontextprotocol/serverInfo".to_string(),
json!({
"name": "brk_mcp",
"version": env!("CARGO_PKG_VERSION"),
}),
);
meta
}
fn upstream_meta(response: &UpstreamResponse) -> MetaObject {
let mut meta = server_meta();
meta.insert(
CACHE_META_KEY.to_string(),
json!({
"status": response.cache_status,
"age": response.cache_age,
}),
);
meta
}