server: mcp + global: refactor

crates/brk_server/src/api/interface/bridge.rs

@@ -1,6 +1,6 @@
 use std::{fs, io, path::Path};
 
-use brk_query::{Index, Query};
+use brk_interface::{Index, Interface};
 
 use crate::{VERSION, Website};
 
@@ -11,7 +11,7 @@ pub trait Bridge {
     fn generate_bridge_file(&self, website: Website, websites_path: &Path) -> io::Result<()>;
 }
 
-impl Bridge for Query<'static> {
+impl Bridge for Interface<'static> {
     fn generate_bridge_file(&self, website: Website, websites_path: &Path) -> io::Result<()> {
         if website.is_none() {
             return Ok(());
@@ -68,8 +68,7 @@ export const VERSION = \"v{}\";
 
         contents += "  return {\n";
 
-        self.vec_trees
-            .id_to_index_to_vec
+        self.id_to_index_to_vec()
             .iter()
             .for_each(|(id, index_to_vec)| {
                 let indexes = index_to_vec

crates/brk_server/src/api/interface/mod.rs

@@ -1,10 +1,11 @@
 use axum::{
     Json,
-    extract::{Query as AxumQuery, State},
+    extract::{Query, State},
     http::{HeaderMap, StatusCode},
     response::{IntoResponse, Response},
 };
-use brk_query::{Format, Index, Output, Params};
+use brk_core::DateIndex;
+use brk_interface::{Format, Output, Params};
 use brk_vec::{CollectableVec, StoredVec};
 use color_eyre::eyre::eyre;
 
@@ -20,7 +21,7 @@ const MAX_WEIGHT: usize = 320_000;
 
 pub async fn handler(
     headers: HeaderMap,
-    query: AxumQuery<Params>,
+    query: Query<Params>,
     State(app_state): State<AppState>,
 ) -> Response {
     match req_to_response_res(headers, query, app_state) {
@@ -36,33 +37,24 @@ pub async fn handler(
 
 fn req_to_response_res(
     headers: HeaderMap,
-    AxumQuery(Params {
-        index,
-        values,
-        rest,
-    }): AxumQuery<Params>,
-    AppState { query, .. }: AppState,
+    Query(params): Query<Params>,
+    AppState { interface, .. }: AppState,
 ) -> color_eyre::Result<Response> {
-    let index = Index::try_from(index.as_str())?;
-
-    let vecs = query.search(
-        index,
-        &values.iter().map(|v| v.as_str()).collect::<Vec<_>>(),
-    );
+    let vecs = interface.search(&params);
 
     if vecs.is_empty() {
         return Ok(Json(vec![] as Vec<usize>).into_response());
     }
 
-    let from = rest.from();
-    let to = rest.to();
-    let format = rest.format();
+    let from = params.from();
+    let to = params.to();
+    let format = params.format();
 
     let weight = vecs
         .iter()
         .map(|(_, v)| {
             let len = v.len();
-            let count = StoredVec::<usize, usize>::range_count(from, to, len);
+            let count = StoredVec::<DateIndex, usize>::range_count(from, to, len);
             count * v.value_type_to_size_of()
         })
         .sum::<usize>();
@@ -91,15 +83,15 @@ fn req_to_response_res(
         }
     }
 
-    let output = query.format(vecs, from, to, format)?;
+    let output = interface.format(vecs, &params.rest)?;
 
     let mut response = match output {
         Output::CSV(s) => s.into_response(),
         Output::TSV(s) => s.into_response(),
         Output::Json(v) => match v {
-            brk_query::Value::Single(v) => Json(v).into_response(),
-            brk_query::Value::List(v) => Json(v).into_response(),
-            brk_query::Value::Matrix(v) => Json(v).into_response(),
+            brk_interface::Value::Single(v) => Json(v).into_response(),
+            brk_interface::Value::List(v) => Json(v).into_response(),
+            brk_interface::Value::Matrix(v) => Json(v).into_response(),
         },
         Output::MD(s) => s.into_response(),
     };

crates/brk_server/src/api/mod.rs

@@ -1,5 +1,3 @@
-use std::collections::BTreeMap;
-
 use axum::{
     Json, Router,
     extract::{Path, Query, State},
@@ -7,131 +5,132 @@ use axum::{
     response::{IntoResponse, Redirect, Response},
     routing::get,
 };
-use brk_query::{Params, ParamsOpt};
+use brk_interface::{Index, Pagination, Params, ParamsOpt};
 
 use super::AppState;
 
 mod explorer;
-mod query;
+mod interface;
 
-pub use query::Bridge;
+pub use interface::Bridge;
 
 pub trait ApiRoutes {
     fn add_api_routes(self) -> Self;
 }
 
+const TO_SEPARATOR: &str = "-to-";
+
 impl ApiRoutes for Router<AppState> {
     fn add_api_routes(self) -> Self {
-        self.route("/api/query", get(query::handler))
-            .route("/api/vecs/id-count", get(id_count_handler))
-            .route("/api/vecs/index-count", get(index_count_handler))
-            .route("/api/vecs/variant-count", get(variant_count_handler))
-            .route("/api/vecs/ids", get(ids_handler))
-            .route("/api/vecs/indexes", get(indexes_handler))
-            .route("/api/vecs/variants", get(variants_handler))
-            .route("/api/vecs/id-to-indexes", get(id_to_indexes_handler))
-            .route("/api/vecs/index-to-ids", get(index_to_ids_handler))
-            .route("/api/{variant}", get(variant_handler))
-            .route(
-                "/api",
-                get(|| async {
-                    Redirect::temporary(
-                        "https://github.com/bitcoinresearchkit/brk/tree/main/crates/brk_server#api",
-                    )
-                }),
-            )
+        self.route(
+            "/api/vecs/index-count",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(app_state.interface.get_index_count()).into_response()
+            }),
+        )
+        .route(
+            "/api/vecs/id-count",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(app_state.interface.get_vecid_count()).into_response()
+            }),
+        )
+        .route(
+            "/api/vecs/vec-count",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(app_state.interface.get_vec_count()).into_response()
+            }),
+        )
+        .route(
+            "/api/vecs/indexes",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(app_state.interface.get_indexes()).into_response()
+            }),
+        )
+        .route(
+            "/api/vecs/accepted-indexes",
+            get(async |State(app_state): State<AppState>| -> Response {
+                Json(app_state.interface.get_accepted_indexes()).into_response()
+            }),
+        )
+        .route(
+            "/api/vecs/ids",
+            get(
+                async |State(app_state): State<AppState>,
+                       Query(pagination): Query<Pagination>|
+                       -> Response {
+                    Json(app_state.interface.get_vecids(pagination)).into_response()
+                },
+            ),
+        )
+        .route(
+            "/api/vecs/indexes-to-ids",
+            get(
+                async |State(app_state): State<AppState>,
+                       Query(pagination): Query<Pagination>|
+                       -> Response {
+                    Json(app_state.interface.get_indexes_to_vecids(pagination)).into_response()
+                },
+            ),
+        )
+        .route(
+            "/api/vecs/ids-to-indexes",
+            get(
+                async |State(app_state): State<AppState>,
+                       Query(pagination): Query<Pagination>|
+                       -> Response {
+                    Json(app_state.interface.get_vecids_to_indexes(pagination)).into_response()
+                },
+            ),
+        )
+        // .route("/api/vecs/variants", get(variants_handler))
+        .route("/api/vecs/query", get(interface::handler))
+        .route(
+            "/api/vecs/{variant}",
+            get(
+                async |headers: HeaderMap,
+                       Path(variant): Path<String>,
+                       Query(params_opt): Query<ParamsOpt>,
+                       state: State<AppState>|
+                       -> Response {
+                    let variant = variant.replace("_", "-");
+                    let mut split = variant.split(TO_SEPARATOR);
+                    let params = Params::from((
+                        (
+                            Index::try_from(split.next().unwrap()).unwrap(),
+                            split.collect::<Vec<_>>().join(TO_SEPARATOR),
+                        ),
+                        params_opt,
+                    ));
+                    interface::handler(headers, Query(params), state).await
+                },
+            ),
+        )
+        .route(
+            "/api",
+            get(|| async {
+                Redirect::temporary(
+                    "https://github.com/bitcoinresearchkit/brk/tree/main/crates/brk_server#api",
+                )
+            }),
+        )
     }
 }
 
-pub async fn ids_handler(State(app_state): State<AppState>) -> Response {
-    Json(
-        app_state
-            .query
-            .vec_trees
-            .id_to_index_to_vec
-            .keys()
-            .collect::<Vec<_>>(),
-    )
-    .into_response()
-}
-
-pub async fn variant_count_handler(State(app_state): State<AppState>) -> Response {
-    Json(
-        app_state
-            .query
-            .vec_trees
-            .index_to_id_to_vec
-            .values()
-            .map(|tree| tree.len())
-            .sum::<usize>(),
-    )
-    .into_response()
-}
-
-pub async fn id_count_handler(State(app_state): State<AppState>) -> Response {
-    Json(app_state.query.vec_trees.id_to_index_to_vec.keys().count()).into_response()
-}
-
-pub async fn index_count_handler(State(app_state): State<AppState>) -> Response {
-    Json(app_state.query.vec_trees.index_to_id_to_vec.keys().count()).into_response()
-}
-
-pub async fn indexes_handler(State(app_state): State<AppState>) -> Response {
-    Json(
-        app_state
-            .query
-            .vec_trees
-            .index_to_id_to_vec
-            .keys()
-            .map(|i| (i.to_string().to_lowercase(), i.possible_values()))
-            .collect::<BTreeMap<_, _>>(),
-    )
-    .into_response()
-}
-
-pub async fn variants_handler(State(app_state): State<AppState>) -> Response {
-    Json(
-        app_state
-            .query
-            .vec_trees
-            .index_to_id_to_vec
-            .iter()
-            .flat_map(|(index, id_to_vec)| {
-                let index_ser = index.serialize_long();
-                id_to_vec
-                    .keys()
-                    .map(|id| format!("{}-to-{}", index_ser, id))
-                    .collect::<Vec<_>>()
-            })
-            .collect::<Vec<_>>(),
-    )
-    .into_response()
-}
-
-pub async fn id_to_indexes_handler(State(app_state): State<AppState>) -> Response {
-    Json(app_state.query.vec_trees.serialize_id_to_index_to_vec()).into_response()
-}
-
-pub async fn index_to_ids_handler(State(app_state): State<AppState>) -> Response {
-    Json(app_state.query.vec_trees.serialize_index_to_id_to_vec()).into_response()
-}
-
-const TO_SEPARATOR: &str = "-to-";
-
-pub async fn variant_handler(
-    headers: HeaderMap,
-    Path(variant): Path<String>,
-    Query(params_opt): Query<ParamsOpt>,
-    state: State<AppState>,
-) -> Response {
-    let variant = variant.replace("_", "-");
-    let mut split = variant.split(TO_SEPARATOR);
-    let params = Params::from((
-        (
-            split.next().unwrap().to_string(),
-            split.collect::<Vec<_>>().join(TO_SEPARATOR),
-        ),
-        params_opt,
-    ));
-    query::handler(headers, Query(params), state).await
-}
+// pub async fn variants_handler(State(app_state): State<AppState>) -> Response {
+//     Json(
+//         app_state
+//             .query
+//             .vec_trees
+//             .index_to_id_to_vec
+//             .iter()
+//             .flat_map(|(index, id_to_vec)| {
+//                 let index_ser = index.serialize_long();
+//                 id_to_vec
+//                     .keys()
+//                     .map(|id| format!("{}-to-{}", index_ser, id))
+//                     .collect::<Vec<_>>()
+//             })
+//             .collect::<Vec<_>>(),
+//     )
+//     .into_response()
+// }

crates/brk_server/src/lib.rs

@@ -23,7 +23,7 @@ use brk_bundler::bundle;
 use brk_computer::Computer;
 use brk_core::dot_brk_path;
 use brk_indexer::Indexer;
-use brk_query::Query;
+use brk_interface::Interface;
 use color_eyre::owo_colors::OwoColorize;
 use files::FilesRoutes;
 use log::{error, info};
@@ -32,16 +32,16 @@ use tower_http::{compression::CompressionLayer, trace::TraceLayer};
 
 mod api;
 mod files;
+mod mcp;
 mod traits;
 
 pub use files::Website;
+use mcp::*;
 use tracing::Span;
 
 #[derive(Clone)]
 pub struct AppState {
-    // indexer: &'static Indexer,
-    // computer: &'static Computer,
-    query: &'static Query<'static>,
+    interface: &'static Interface<'static>,
     website: Website,
     websites_path: Option<PathBuf>,
 }
@@ -67,7 +67,7 @@ impl Server {
     pub fn new(indexer: Indexer, computer: Computer, website: Website) -> color_eyre::Result<Self> {
         let indexer = Box::leak(Box::new(indexer));
         let computer = Box::leak(Box::new(computer));
-        let query = Box::leak(Box::new(Query::build(indexer, computer)));
+        let interface = Box::leak(Box::new(Interface::build(indexer, computer)));
 
         let websites_path = if website.is_some() {
             let websites_dev_path = Path::new(DEV_PATH).join(WEBSITES);
@@ -99,7 +99,7 @@ impl Server {
                 downloaded_websites_path
             };
 
-            query.generate_bridge_file(website, websites_path.as_path())?;
+            interface.generate_bridge_file(website, websites_path.as_path())?;
 
             Some(websites_path)
         } else {
@@ -107,7 +107,7 @@ impl Server {
         };
 
         Ok(Self(AppState {
-            query,
+            interface,
             website,
             websites_path,
         }))
@@ -162,6 +162,7 @@ impl Server {
         let router = Router::new()
             .add_api_routes()
             .add_website_routes(state.website)
+            .add_mcp_routes(state.interface)
             .route("/version", get(Json(VERSION)))
             .with_state(state)
             .layer(compression_layer)

crates/brk_server/src/mcp/api.rs

@@ -0,0 +1,153 @@
+use brk_interface::{Interface, Pagination, Params};
+use rmcp::{
+    Error as McpError, RoleServer, ServerHandler,
+    model::{
+        CallToolResult, Content, Implementation, InitializeRequestParam, InitializeResult,
+        ProtocolVersion, ServerCapabilities, ServerInfo,
+    },
+    service::RequestContext,
+    tool,
+};
+
+#[derive(Clone)]
+pub struct API {
+    interface: &'static Interface<'static>,
+}
+
+const VERSION: &str = env!("CARGO_PKG_VERSION");
+
+#[tool(tool_box)]
+impl API {
+    pub fn new(interface: &'static Interface<'static>) -> Self {
+        Self { interface }
+    }
+
+    #[tool(description = "
+Get the count of all existing indexes
+")]
+    async fn get_index_count(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_index_count()).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get the count of all existing vec ids
+")]
+    async fn get_vecid_count(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_vecid_count()).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get the count of all existing vecs
+")]
+    async fn get_variant_count(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_vec_count()).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get the list of all existing indexes
+")]
+    async fn get_indexes(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_indexes()).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get an object which has all existing indexes as keys and a list of their accepted variants as values
+")]
+    async fn get_accepted_indexes(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_accepted_indexes()).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get the list of all existing vec ids
+")]
+    async fn get_vecids(
+        &self,
+        #[tool(aggr)] pagination: Pagination,
+    ) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_vecids(pagination)).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get an object which has all existing indexes as keys and a list of ids of vecs which support that index as values
+")]
+    async fn get_indexes_to_vecids(
+        &self,
+        #[tool(aggr)] pagination: Pagination,
+    ) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_indexes_to_vecids(pagination)).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get an object which has all existing vec ids as keys and a list of indexes supported by that vec id as values
+")]
+    async fn get_vecids_to_indexes(
+        &self,
+        #[tool(aggr)] pagination: Pagination,
+    ) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.get_vecids_to_indexes(pagination)).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "Get one or multiple vecs depending on given parameters")]
+    fn get_vecs(&self, #[tool(aggr)] params: Params) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![
+            Content::json(self.interface.search_and_format(params).unwrap()).unwrap(),
+        ]))
+    }
+
+    #[tool(description = "
+Get the running version of the Bitcoin Research Kit
+")]
+    async fn get_version(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![Content::text(format!(
+            "v{VERSION}"
+        ))]))
+    }
+}
+
+#[tool(tool_box)]
+impl ServerHandler for API {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo {
+            protocol_version: ProtocolVersion::V_2025_03_26,
+            capabilities: ServerCapabilities::builder().enable_tools().build(),
+            server_info: Implementation::from_build_env(),
+            instructions: Some(
+                "
+This server provides an interface to communicate with a running instance of the Bitcoin Research Kit (brk).
+Multiple tools are at your disposal including ones to fetch all sorts of Bitcoin on-chain data.
+Arrays are also called Vectors (or Vecs).
+"
+                .to_string(),
+            ),
+        }
+    }
+
+    async fn initialize(
+        &self,
+        _request: InitializeRequestParam,
+        context: RequestContext<RoleServer>,
+    ) -> Result<InitializeResult, McpError> {
+        if let Some(http_request_part) = context.extensions.get::<axum::http::request::Parts>() {
+            let initialize_headers = &http_request_part.headers;
+            let initialize_uri = &http_request_part.uri;
+            tracing::info!(?initialize_headers, %initialize_uri, "initialize from http server");
+        }
+        Ok(self.get_info())
+    }
+}

crates/brk_server/src/mcp/mod.rs

@@ -0,0 +1,32 @@
+use axum::Router;
+use brk_interface::Interface;
+use rmcp::transport::{
+    StreamableHttpServerConfig,
+    streamable_http_server::{StreamableHttpService, session::local::LocalSessionManager},
+};
+
+mod api;
+use api::*;
+
+use crate::AppState;
+
+pub trait MCPRoutes {
+    fn add_mcp_routes(self, interface: &'static Interface<'static>) -> Self;
+}
+
+impl MCPRoutes for Router<AppState> {
+    fn add_mcp_routes(self, interface: &'static Interface<'static>) -> Self {
+        let config = StreamableHttpServerConfig {
+            // stateful_mode: false,
+            ..Default::default()
+        };
+
+        let service = StreamableHttpService::new(
+            move || Ok(API::new(interface)),
+            LocalSessionManager::default().into(),
+            config,
+        );
+
+        self.nest_service("/mcp", service)
+    }
+}