Skip to main content

wowlab_sentinel/mcp/
mod.rs

1pub(crate) mod schema;
2mod tools;
3
4crate::db_client!(McpDb, "mcp", max = 8, statement_timeout_secs = 30);
5
6tokio::task_local! {
7    pub static MCP_USER_ID: Option<uuid::Uuid>;
8}
9
10pub(crate) fn current_user_id() -> Option<uuid::Uuid> {
11    MCP_USER_ID.try_with(|id| *id).ok().flatten()
12}
13
14pub(crate) async fn capture_user_middleware(
15    request: axum::extract::Request,
16    next: axum::middleware::Next,
17) -> axum::response::Response {
18    let user_id = request
19        .extensions()
20        .get::<crate::http::auth::AuthenticatedUser>()
21        .map(|u| u.user_id);
22
23    MCP_USER_ID.scope(user_id, next.run(request)).await
24}
25
26/// Define an MCP tool: route builder fn plus inventory metadata.
27macro_rules! mcp_tool {
28    ($fn_name:ident, $order:expr, $desc:literal, $param:ty,
29     tips = [$($tip:literal),* $(,)?],
30     |$srv:ident, $p:ident| async $body:block
31    ) => {
32        paste::paste! {
33            pub(super) fn [< build_ $fn_name _route >]() -> rmcp::handler::server::router::tool::ToolRoute<$crate::mcp::McpHandler> {
34                use rmcp::{handler::server::tool::schema_for_type, handler::server::wrapper::Parameters};
35
36                let tool = rmcp::model::Tool::new(
37                    stringify!($fn_name),
38                    $desc,
39                    schema_for_type::<Parameters<$param>>(),
40                );
41
42                rmcp::handler::server::router::tool::ToolRoute::new_dyn(tool, |ctx| {
43                    Box::pin(async move {
44                        let $srv: &$crate::mcp::McpHandler = ctx.service;
45                        let $p: $param = serde_json::from_value(
46                            serde_json::Value::Object(ctx.arguments.unwrap_or_default()),
47                        ).map_err(|e| rmcp::ErrorData::invalid_params(
48                            format!("invalid parameters: {e}"), None,
49                        ))?;
50
51                        $crate::mcp::with_metrics(
52                            $crate::telemetry::McpTool::[< $fn_name:camel >],
53                            || async $body,
54                        ).await
55                    })
56                })
57            }
58
59            inventory::submit! {
60                $crate::mcp::schema::McpToolMeta {
61                    name: stringify!($fn_name),
62                    summary: $desc,
63                    order: $order,
64                    tips: &[$($tip),*],
65                }
66            }
67        }
68    };
69}
70
71pub(crate) use mcp_tool;
72
73/// Register an MCP-queryable table via `inventory`.
74macro_rules! mcp_table {
75    (meta, $cat:literal, $desc:literal, [ $($tokens:tt)* ]) => {
76        $crate::mcp::mcp_table_parse!(@cols
77            $crate::mcp::schema::McpTableId::Metadata,
78            $cat, $desc, [], [], $($tokens)*);
79    };
80    ($table:path, $cat:literal, $desc:literal, [ $($tokens:tt)* ]) => {
81        $crate::mcp::mcp_table_parse!(@cols
82            $crate::mcp::schema::McpTableId::GameData($table),
83            $cat, $desc, [], [], $($tokens)*);
84    };
85}
86
87pub(crate) use mcp_table;
88
89/// Internal tt-munching helper for `mcp_table!`.
90macro_rules! mcp_table_parse {
91    (@cols $identity:expr, $cat:literal, $desc:literal,
92     [$($cols:tt)*], [$($schemas:tt)*],
93     $col:literal : Json($($ty:tt)+), $($rest:tt)*) => {
94        $crate::mcp::mcp_table_parse!(@cols $identity, $cat, $desc,
95            [$($cols)* { $col, Json },],
96            [$($schemas)* { $identity, $col, $($ty)+ },],
97            $($rest)*);
98    };
99    (@cols $identity:expr, $cat:literal, $desc:literal,
100     [$($cols:tt)*], [$($schemas:tt)*],
101     $col:literal : $typ:ident, $($rest:tt)*) => {
102        $crate::mcp::mcp_table_parse!(@cols $identity, $cat, $desc,
103            [$($cols)* { $col, $typ },],
104            [$($schemas)*],
105            $($rest)*);
106    };
107    (@cols $identity:expr, $cat:literal, $desc:literal,
108     [$($cols:tt)*], [$($schemas:tt)*],
109     $col:literal : Json($($ty:tt)+)) => {
110        $crate::mcp::mcp_table_parse!(@emit $identity, $cat, $desc,
111            [$($cols)* { $col, Json },],
112            [$($schemas)* { $identity, $col, $($ty)+ },]);
113    };
114    (@cols $identity:expr, $cat:literal, $desc:literal,
115     [$($cols:tt)*], [$($schemas:tt)*],
116     $col:literal : $typ:ident) => {
117        $crate::mcp::mcp_table_parse!(@emit $identity, $cat, $desc,
118            [$($cols)* { $col, $typ },],
119            [$($schemas)*]);
120    };
121    (@cols $identity:expr, $cat:literal, $desc:literal,
122     [$($cols:tt)*], [$($schemas:tt)*], ) => {
123        $crate::mcp::mcp_table_parse!(@emit $identity, $cat, $desc,
124            [$($cols)*], [$($schemas)*]);
125    };
126
127    (@emit $identity:expr, $cat:literal, $desc:literal,
128     [$({ $col:literal, $typ:ident },)*],
129     [$({ $table:expr, $jcol:literal, $($jty:tt)+ },)*]) => {
130        inventory::submit! {
131            $crate::mcp::schema::Table {
132                identity: $identity,
133                category: $cat,
134                description: $desc,
135                columns: &[
136                    $( $crate::mcp::schema::Column {
137                        name: $col,
138                        typ: $crate::mcp::schema::ColType::$typ,
139                    } ),*
140                ],
141            }
142        }
143        $(
144            inventory::submit! {
145                $crate::mcp::schema::JsonColumnSchema {
146                    table: $table,
147                    column: $jcol,
148                    schema_fn: || {
149                        serde_json::to_value(
150                            schemars::schema_for!($($jty)+)
151                        )
152                            .unwrap_or_default()
153                    },
154                }
155            }
156        )*
157    };
158}
159
160use std::sync::Arc;
161
162pub(crate) use mcp_table_parse;
163use rmcp::{
164    ErrorData, ServerHandler,
165    handler::server::router::tool::ToolRouter,
166    model::{CallToolResult, Content, Implementation, ServerCapabilities, ServerInfo},
167    tool_handler,
168    transport::streamable_http_server::{
169        StreamableHttpServerConfig, session::local::LocalSessionManager,
170        tower::StreamableHttpService,
171    },
172};
173use sqlx::PgPool;
174use wowlab_common::{markdown::Doc, time::Instant};
175
176use crate::{
177    telemetry::{self, McpErrorType, McpTool},
178    utils::meta,
179};
180
181pub(crate) async fn with_metrics<F, Fut>(tool: McpTool, f: F) -> Result<CallToolResult, ErrorData>
182where
183    F: FnOnce() -> Fut,
184    Fut: Future<Output = Result<CallToolResult, ErrorData>>,
185{
186    let start = Instant::now();
187
188    telemetry::record_mcp_request(tool);
189    let result = f().await;
190
191    if result.is_ok() {
192        telemetry::record_mcp_duration(tool, start.elapsed().as_secs_f64());
193    }
194
195    result
196}
197
198fn mcp_error(tool: McpTool, error: &impl std::fmt::Display) -> ErrorData {
199    let message = error.to_string();
200
201    telemetry::record_mcp_error(tool, classify_error(&message));
202
203    ErrorData::invalid_params(message, None)
204}
205
206fn classify_error(msg: &str) -> McpErrorType {
207    const PATTERNS: &[(&[&str], McpErrorType)] = &[
208        (&["timeout"], McpErrorType::Timeout),
209        (
210            &[
211                "Unknown table",
212                "not filterable",
213                "not sortable",
214                "Batch requires",
215                "Batch supports",
216            ],
217            McpErrorType::InvalidParams,
218        ),
219    ];
220
221    PATTERNS
222        .iter()
223        .find(|(keywords, _)| keywords.iter().any(|k| msg.contains(k)))
224        .map_or(McpErrorType::Database, |(_, err)| *err)
225}
226
227fn json_result(v: impl serde::Serialize) -> CallToolResult {
228    let text = serde_json::to_string(&v).unwrap_or_else(|_| "{}".into());
229
230    CallToolResult::success(vec![Content::text(text)])
231}
232
233#[derive(Debug, Default, schemars::JsonSchema, serde::Deserialize)]
234#[schemars(description = "Fetch project documentation.")]
235pub(crate) struct GetDocsInput {
236    #[serde(default)]
237    #[schemars(description = "Return full content instead of index. Default: false.")]
238    pub full: Option<bool>,
239}
240
241#[derive(Debug, Default, schemars::JsonSchema, serde::Deserialize)]
242#[schemars(description = "No parameters.")]
243#[expect(
244    clippy::empty_structs_with_brackets,
245    reason = "MCP input must remain a JSON object rather than a unit value"
246)]
247pub(crate) struct EmptyInput {}
248
249#[derive(Clone)]
250pub(crate) struct McpHandler {
251    pub db: PgPool,
252    pub mcp: crate::state::McpStateHandle,
253    tool_router: ToolRouter<Self>,
254}
255
256impl McpHandler {
257    pub(crate) fn new(db: PgPool, mcp: crate::state::McpStateHandle) -> Self {
258        let router = ToolRouter::new()
259            .with_route(tools::build_get_docs_route())
260            .with_route(tools::build_list_tables_route())
261            .with_route(tools::build_get_schema_route())
262            .with_route(tools::build_query_route())
263            .with_route(tools::build_query_batch_route())
264            .with_route(tools::build_query_count_route())
265            .with_route(tools::build_resolve_effects_route())
266            .with_route(tools::build_decode_loadout_route())
267            .with_route(tools::build_get_rotation_route())
268            .with_route(tools::build_list_rotations_route())
269            .with_route(tools::build_list_spell_labels_route())
270            .with_route(tools::build_parse_simc_route());
271
272        Self {
273            db,
274            mcp,
275            tool_router: router,
276        }
277    }
278}
279
280impl std::fmt::Debug for McpHandler {
281    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
282        f.debug_struct("McpHandler")
283            .field("db", &"<PgPool>")
284            .field("mcp", &self.mcp)
285            .finish_non_exhaustive()
286    }
287}
288
289#[tool_handler(router = self.tool_router)]
290impl ServerHandler for McpHandler {
291    fn get_info(&self) -> ServerInfo {
292        ServerInfo::new(ServerCapabilities::builder().enable_tools().build())
293            .with_server_info(
294                Implementation::new("wowlab", env!("CARGO_PKG_VERSION"))
295                    .with_title("WoW Lab MCP Server")
296                    .with_website_url(meta::WEBSITE),
297            )
298            .with_instructions(mcp_instructions(&self.mcp.schema))
299    }
300}
301
302pub(crate) fn create_service(
303    db: PgPool,
304    mcp: crate::state::McpStateHandle,
305) -> StreamableHttpService<McpHandler, LocalSessionManager> {
306    StreamableHttpService::new(
307        move || Ok(McpHandler::new(db.clone(), Arc::clone(&mcp))),
308        Arc::new(LocalSessionManager::default()),
309        StreamableHttpServerConfig::default()
310            .with_stateful_mode(false)
311            .disable_allowed_hosts(),
312    )
313}
314
315fn mcp_instructions(catalog: &schema::SchemaCatalog) -> String {
316    let doc = schema::build_workflow_doc(
317        Doc::new()
318            .line("WoW game data query API.")
319            .blank()
320            .h2("Workflow"),
321    )
322    .blank()
323    .h2("Tables");
324
325    let doc = catalog.build_tables_doc(doc).blank().h2("Tips");
326
327    schema::build_tips_doc(doc).build()
328}