mcp: small changes

crates/brk_interface/src/index.rs

@@ -12,29 +12,53 @@ use serde::{Deserialize, de::Error};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, JsonSchema)]
 pub enum Index {
+    #[schemars(description = "Date/day index")]
     DateIndex,
+    #[schemars(description = "Decade index")]
     DecadeIndex,
+    #[schemars(description = "Difficulty epoch index (equivalent to ~2 weeks)")]
     DifficultyEpoch,
+    #[schemars(description = "Empty output index")]
     EmptyOutputIndex,
+    #[schemars(description = "Halving epoch index (equivalent to ~4 years)")]
     HalvingEpoch,
+    #[schemars(description = "Height/block index")]
     Height,
+    #[schemars(description = "Transaction input index (based on total)")]
     InputIndex,
+    #[schemars(description = "Month index")]
     MonthIndex,
+    #[schemars(description = "Op return index")]
     OpReturnIndex,
+    #[schemars(description = "Transaction output index (based on total)")]
     OutputIndex,
+    #[schemars(description = "Index of P2A address")]
     P2AAddressIndex,
+    #[schemars(description = "Index of P2MS output")]
     P2MSOutputIndex,
+    #[schemars(description = "Index of P2PK (33 bytes) address")]
     P2PK33AddressIndex,
+    #[schemars(description = "Index of P2PK (65 bytes) address")]
     P2PK65AddressIndex,
+    #[schemars(description = "Index of P2PKH address")]
     P2PKHAddressIndex,
+    #[schemars(description = "Index of P2SH address")]
     P2SHAddressIndex,
+    #[schemars(description = "Index of P2TR address")]
     P2TRAddressIndex,
+    #[schemars(description = "Index of P2WPKH address")]
     P2WPKHAddressIndex,
+    #[schemars(description = "Index of P2WSH address")]
     P2WSHAddressIndex,
+    #[schemars(description = "Quarter index")]
     QuarterIndex,
+    #[schemars(description = "Transaction index")]
     TxIndex,
+    #[schemars(description = "Unknown output index")]
     UnknownOutputIndex,
+    #[schemars(description = "Week index")]
     WeekIndex,
+    #[schemars(description = "Year index")]
     YearIndex,
 }
 

crates/brk_interface/src/params.rs

@@ -12,11 +12,11 @@ use crate::{
 #[derive(Debug, Deserialize, JsonSchema)]
 pub struct Params {
     #[serde(alias = "i")]
-    #[schemars(description = "Index of the values requested")]
+    #[schemars(description = "Index of requested vecs")]
     pub index: Index,
 
     #[serde(alias = "v")]
-    #[schemars(description = "Ids of the requested vecs")]
+    #[schemars(description = "Ids of requested vecs")]
     pub ids: MaybeIds,
 
     #[serde(flatten)]
@@ -45,18 +45,24 @@ impl From<((Index, String), ParamsOpt)> for Params {
 pub struct ParamsOpt {
     #[serde(default, alias = "f", deserialize_with = "de_unquote_i64")]
     /// Inclusive starting index, if negative will be from the end
+    #[schemars(description = "Inclusive starting index, if negative will be from the end")]
     from: Option<i64>,
 
     #[serde(default, alias = "t", deserialize_with = "de_unquote_i64")]
     /// Exclusive ending index, if negative will be from the end, overrides 'count'
+    #[schemars(
+        description = "Exclusive ending index, if negative will be from the end, overrides 'count'"
+    )]
     to: Option<i64>,
 
     #[serde(default, alias = "c", deserialize_with = "de_unquote_usize")]
-    /// Number of values
+    /// Number of values requested
+    #[schemars(description = "Number of values requested")]
     count: Option<usize>,
 
-    /// Format of the output
     #[serde(default)]
+    /// Format of the output
+    #[schemars(description = "Format of the output")]
     format: Option<Format>,
 }
 

crates/brk_mcp/src/lib.rs

@@ -6,12 +6,10 @@
 use brk_interface::{IdParam, Interface, PaginatedIndexParam, PaginationParam, Params};
 use brk_rmcp::{
     Error as McpError, RoleServer, ServerHandler,
-    model::{
-        CallToolResult, Content, Implementation, InitializeRequestParam, InitializeResult,
-        ProtocolVersion, ServerCapabilities, ServerInfo,
-    },
+    handler::server::{router::tool::ToolRouter, tool::Parameters},
+    model::*,
     service::RequestContext,
-    tool,
+    tool, tool_handler, tool_router,
 };
 use log::info;
 
@@ -20,14 +18,18 @@ pub mod route;
 #[derive(Clone)]
 pub struct MCP {
     interface: &'static Interface<'static>,
+    tool_router: ToolRouter<MCP>,
 }
 
 const VERSION: &str = env!("CARGO_PKG_VERSION");
 
-#[tool(tool_box)]
+#[tool_router]
 impl MCP {
     pub fn new(interface: &'static Interface<'static>) -> Self {
-        Self { interface }
+        Self {
+            interface,
+            tool_router: Self::tool_router(),
+        }
     }
 
     #[tool(description = "
@@ -54,8 +56,8 @@ Get the count of all existing vec ids.
 Get the count of all existing vecs.
 Equals to the sum of supported Indexes of each vec id.
 ")]
-    async fn get_variant_count(&self) -> Result<CallToolResult, McpError> {
-        info!("mcp: get_variant_count");
+    async fn get_vec_count(&self) -> Result<CallToolResult, McpError> {
+        info!("mcp: get_vec_count");
         Ok(CallToolResult::success(vec![
             Content::json(self.interface.get_vec_count()).unwrap(),
         ]))
@@ -88,7 +90,7 @@ If the `page` param is omitted, it will default to the first page.
 ")]
     async fn get_vecids(
         &self,
-        #[tool(aggr)] pagination: PaginationParam,
+        Parameters(pagination): Parameters<PaginationParam>,
     ) -> Result<CallToolResult, McpError> {
         info!("mcp: get_vecids");
         Ok(CallToolResult::success(vec![
@@ -103,7 +105,7 @@ If the `page` param is omitted, it will default to the first page.
 ")]
     async fn get_index_to_vecids(
         &self,
-        #[tool(aggr)] paginated_index: PaginatedIndexParam,
+        Parameters(paginated_index): Parameters<PaginatedIndexParam>,
     ) -> Result<CallToolResult, McpError> {
         info!("mcp: get_index_to_vecids");
         Ok(CallToolResult::success(vec![
@@ -117,7 +119,7 @@ The list will be empty if the vec id isn't correct.
 ")]
     async fn get_vecid_to_indexes(
         &self,
-        #[tool(aggr)] param: IdParam,
+        Parameters(param): Parameters<IdParam>,
     ) -> Result<CallToolResult, McpError> {
         info!("mcp: get_vecid_to_indexes");
         Ok(CallToolResult::success(vec![
@@ -127,14 +129,12 @@ The list will be empty if the vec id isn't correct.
 
     #[tool(description = "
 Get one or multiple vecs depending on given parameters.
-If you'd like to request multiple vec ids, simply separate them with a ','.
-To get the last value set `-1` to the `from` parameter.
 The response's format will depend on the given parameters, it will be:
 - A value: If requested only one vec and the given range returns one value (for example: `from=-1`)
 - A list: If requested only one vec and the given range returns multiple values (for example: `from=-1000&count=100` or `from=-444&to=-333`)
 - A matrix: When multiple vecs are requested, even if they each return one value.
 ")]
-    fn get_vecs(&self, #[tool(aggr)] params: Params) -> Result<CallToolResult, McpError> {
+    fn get_vecs(&self, Parameters(params): Parameters<Params>) -> Result<CallToolResult, McpError> {
         info!("mcp: get_vecs");
         Ok(CallToolResult::success(vec![
             Content::json(self.interface.search_and_format(params).unwrap()).unwrap(),
@@ -152,7 +152,7 @@ Get the running version of the Bitcoin Research Kit.
     }
 }
 
-#[tool(tool_box)]
+#[tool_handler]
 impl ServerHandler for MCP {
     fn get_info(&self) -> ServerInfo {
         ServerInfo {
@@ -181,13 +181,8 @@ An 'Index' (or indexes) is the timeframe of a dataset.
     async fn initialize(
         &self,
         _request: InitializeRequestParam,
-        context: RequestContext<RoleServer>,
+        _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())
     }
 }