binaryninja/mcp/
tool.rs

1// Copyright 2022-2026 Vector 35 Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Writing and registering MCP tools.
16//!
17//! A tool registered here is offered by every MCP server in the process, alongside the built-in
18//! `bn_*` tools. Implement [`McpTool`] and pass it to [`register_mcp_tool`]. With the `schemars`
19//! feature, implement `TypedMcpTool` instead and pass it to `register_typed_mcp_tool`. The input
20//! schema and argument parsing then come from the arguments type.
21
22use binaryninjacore_sys::*;
23use serde_json::{json, Value};
24use std::ffi::{c_char, c_void, CStr, CString};
25use std::fmt;
26use std::panic::{self, AssertUnwindSafe};
27use std::ptr;
28
29use crate::binary_view::BinaryView;
30use crate::rc::*;
31use crate::string::{BnString, IntoCStr};
32
33pub use binaryninjacore_sys::BNMcpToolScope as McpToolScope;
34
35bitflags::bitflags! {
36    /// Hints a client uses to decide how to present a tool and whether to ask before running it.
37    #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
38    pub struct McpToolAnnotations: u32 {
39        const READ_ONLY = BNMcpToolAnnotation::ReadOnlyHint.0;
40        const DESTRUCTIVE = BNMcpToolAnnotation::DestructiveHint.0;
41        const IDEMPOTENT = BNMcpToolAnnotation::IdempotentHint.0;
42        const OPEN_WORLD = BNMcpToolAnnotation::OpenWorldHint.0;
43    }
44}
45
46/// Everything about a tool except its input schema.
47#[derive(Clone, Debug)]
48pub struct McpToolInfo {
49    pub name: String,
50    pub title: String,
51    pub description: String,
52    pub scope: McpToolScope,
53    pub annotations: McpToolAnnotations,
54    pub output_schema: Option<Value>,
55}
56
57impl McpToolInfo {
58    /// A BinaryView-scoped tool with no title, annotations or output schema.
59    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
60        Self {
61            name: name.into(),
62            title: String::new(),
63            description: description.into(),
64            scope: McpToolScope::BinaryViewScope,
65            annotations: McpToolAnnotations::empty(),
66            output_schema: None,
67        }
68    }
69
70    pub fn with_title(mut self, title: impl Into<String>) -> Self {
71        self.title = title.into();
72        self
73    }
74
75    pub fn with_scope(mut self, scope: McpToolScope) -> Self {
76        self.scope = scope;
77        self
78    }
79
80    pub fn with_annotations(mut self, annotations: McpToolAnnotations) -> Self {
81        self.annotations = annotations;
82        self
83    }
84
85    pub fn with_output_schema(mut self, schema: Value) -> Self {
86        self.output_schema = Some(schema);
87        self
88    }
89
90    /// Completes the definition with a JSON Schema object for the tool's input.
91    pub fn with_input_schema(self, input_schema: Value) -> McpToolDefinition {
92        McpToolDefinition {
93            info: self,
94            input_schema,
95        }
96    }
97}
98
99#[derive(Clone, Debug)]
100pub struct McpToolDefinition {
101    pub info: McpToolInfo,
102    /// A JSON Schema object with `"type": "object"`.
103    pub input_schema: Value,
104}
105
106/// An error result with a machine-readable code.
107#[derive(Clone, Debug, PartialEq)]
108pub struct McpToolError {
109    pub code: String,
110    pub message: String,
111    pub details: Option<Value>,
112}
113
114impl McpToolError {
115    pub fn new(code: impl Into<String>, message: impl Into<String>) -> Self {
116        Self {
117            code: code.into(),
118            message: message.into(),
119            details: None,
120        }
121    }
122
123    pub fn invalid_params(message: impl Into<String>) -> Self {
124        Self::new("invalid_params", message)
125    }
126
127    pub fn with_details(mut self, details: Value) -> Self {
128        self.details = Some(details);
129        self
130    }
131}
132
133impl fmt::Display for McpToolError {
134    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
135        write!(f, "{}: {}", self.code, self.message)
136    }
137}
138
139impl std::error::Error for McpToolError {}
140
141/// Arguments that do not match the type they are parsed into are invalid parameters.
142impl From<serde_json::Error> for McpToolError {
143    fn from(error: serde_json::Error) -> Self {
144        Self::invalid_params(error.to_string())
145    }
146}
147
148/// Tool output is arbitrary text, and a NUL would otherwise panic in `CString::new`.
149fn c_text(text: &str) -> CString {
150    CString::new(text.replace('\0', "\u{FFFD}")).unwrap_or_default()
151}
152
153#[derive(Clone, Debug, Default, PartialEq)]
154pub struct McpToolResult {
155    text: Vec<String>,
156    structured_content: Option<Value>,
157    error: Option<McpToolError>,
158    warnings: Vec<(String, String)>,
159}
160
161impl McpToolResult {
162    pub fn text(text: impl Into<String>) -> Self {
163        Self {
164            text: vec![text.into()],
165            ..Self::default()
166        }
167    }
168
169    /// `content` must be a JSON object. When no text is added, clients that ignore structured
170    /// content see it serialized as text.
171    pub fn structured(content: Value) -> Self {
172        Self {
173            structured_content: Some(content),
174            ..Self::default()
175        }
176    }
177
178    pub fn error(error: McpToolError) -> Self {
179        Self {
180            error: Some(error),
181            ..Self::default()
182        }
183    }
184
185    pub fn with_text(mut self, text: impl Into<String>) -> Self {
186        self.text.push(text.into());
187        self
188    }
189
190    /// An advisory warning about the result, such as analysis that has not finished. Clients see
191    /// it in the structured content's reserved `warnings` member, and after any text.
192    pub fn with_warning(mut self, code: impl Into<String>, message: impl Into<String>) -> Self {
193        self.warnings.push((code.into(), message.into()));
194        self
195    }
196
197    fn apply(&self, result: *mut BNMcpToolResult) {
198        unsafe {
199            if let Some(error) = &self.error {
200                let code = c_text(&error.code);
201                let message = c_text(&error.message);
202                let details = error
203                    .details
204                    .as_ref()
205                    .map(|details| details.to_string().to_cstr());
206                BNSetMcpToolResultError(
207                    result,
208                    code.as_ptr(),
209                    message.as_ptr(),
210                    details
211                        .as_ref()
212                        .map_or(ptr::null(), |details| details.as_ptr()),
213                );
214            } else if let Some(content) = &self.structured_content {
215                let json = content.to_string().to_cstr();
216                if !BNSetMcpToolResultStructuredContent(result, json.as_ptr()) {
217                    let code = "internal_error".to_cstr();
218                    let message =
219                        "The tool produced structured content that is not a JSON object".to_cstr();
220                    BNSetMcpToolResultError(result, code.as_ptr(), message.as_ptr(), ptr::null());
221                    return;
222                }
223            }
224            for text in &self.text {
225                let text = c_text(text);
226                BNAddMcpToolResultText(result, text.as_ptr());
227            }
228            for (code, message) in &self.warnings {
229                let code = c_text(code);
230                let message = c_text(message);
231                BNAddMcpToolResultWarning(result, code.as_ptr(), message.as_ptr());
232            }
233        }
234    }
235}
236
237impl From<McpToolError> for McpToolResult {
238    fn from(error: McpToolError) -> Self {
239        Self::error(error)
240    }
241}
242
243/// One invocation of a tool. The binary view and cancellation state are only available while the
244/// tool is running.
245pub struct McpToolCall {
246    handle: *mut BNMcpToolCall,
247}
248
249impl McpToolCall {
250    pub(crate) unsafe fn ref_from_raw(handle: *mut BNMcpToolCall) -> Ref<Self> {
251        debug_assert!(!handle.is_null());
252        Ref::new(Self { handle })
253    }
254
255    /// The binary view the MCP session targets. Never `None` while a BinaryView-scoped tool runs.
256    pub fn binary_view(&self) -> Option<Ref<BinaryView>> {
257        let view = unsafe { BNGetMcpToolCallBinaryView(self.handle) };
258        (!view.is_null()).then(|| unsafe { BinaryView::ref_from_raw(view) })
259    }
260
261    /// The binary view, or the `no_active_binary_view` error to return with `?` when there is none.
262    pub fn require_binary_view(&self) -> Result<Ref<BinaryView>, McpToolError> {
263        self.binary_view().ok_or_else(|| {
264            McpToolError::new(
265                "no_active_binary_view",
266                "No active Binary Ninja binary view is selected",
267            )
268        })
269    }
270
271    pub fn is_cancelled(&self) -> bool {
272        unsafe { BNIsMcpToolCallCancelled(self.handle) }
273    }
274
275    pub fn report_progress(&self, progress: f64, total: f64, message: &str) {
276        let message = c_text(message);
277        unsafe { BNReportMcpToolCallProgress(self.handle, progress, total, message.as_ptr()) }
278    }
279
280    /// Evaluates an address expression string against the call's binary view, with `here` as the
281    /// value of `$here`.
282    pub fn parse_address(&self, value: &Value, here: u64) -> Result<u64, String> {
283        self.parse(BNParseMcpToolCallAddress, value, here)
284    }
285
286    /// Evaluates an unsigned integer, or an expression string against the call's binary view with
287    /// `here` as the value of `$here`.
288    pub fn parse_integer(&self, value: &Value, here: u64) -> Result<u64, String> {
289        self.parse(BNParseMcpToolCallInteger, value, here)
290    }
291
292    fn parse(
293        &self,
294        parse: unsafe extern "C" fn(
295            *mut BNMcpToolCall,
296            *const c_char,
297            *mut u64,
298            u64,
299            *mut *mut c_char,
300        ) -> bool,
301        value: &Value,
302        here: u64,
303    ) -> Result<u64, String> {
304        let json = value.to_string().to_cstr();
305        let mut result = 0;
306        let mut error: *mut c_char = ptr::null_mut();
307        if unsafe { parse(self.handle, json.as_ptr(), &mut result, here, &mut error) } {
308            Ok(result)
309        } else if error.is_null() {
310            Err(String::new())
311        } else {
312            Err(unsafe { BnString::into_string(error) })
313        }
314    }
315}
316
317unsafe impl Send for McpToolCall {}
318unsafe impl Sync for McpToolCall {}
319
320impl ToOwned for McpToolCall {
321    type Owned = Ref<Self>;
322
323    fn to_owned(&self) -> Self::Owned {
324        unsafe { RefCountable::inc_ref(self) }
325    }
326}
327
328unsafe impl RefCountable for McpToolCall {
329    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
330        Ref::new(Self {
331            handle: BNNewMcpToolCallReference(handle.handle),
332        })
333    }
334
335    unsafe fn dec_ref(handle: &Self) {
336        BNFreeMcpToolCall(handle.handle);
337    }
338}
339
340/// A tool with a hand-written input schema that receives its arguments as JSON.
341pub trait McpTool: 'static + Send + Sync {
342    fn definition(&self) -> McpToolDefinition;
343
344    /// Called on an arbitrary thread. An error produces an error result. Under `panic = "unwind"`
345    /// a panic produces an `internal_error` result, and under `panic = "abort"` it ends the process.
346    fn invoke(&self, call: &McpToolCall, arguments: Value) -> Result<McpToolResult, McpToolError>;
347}
348
349/// Registers a tool for the life of the process. Returns `None` when the definition is invalid or
350/// its name is already registered, or when a string in it contains a NUL.
351pub fn register_mcp_tool<T: McpTool>(tool: T) -> Option<Ref<CoreMcpTool>> {
352    create_with(tool, BNRegisterMcpTool)
353}
354
355pub(super) fn create_with<T: McpTool>(
356    tool: T,
357    create: unsafe extern "C" fn(
358        *const BNMcpToolDefinition,
359        *const BNMcpToolCallbacks,
360    ) -> *mut BNMcpTool,
361) -> Option<Ref<CoreMcpTool>> {
362    struct ToolContext<T> {
363        name: String,
364        tool: T,
365    }
366
367    extern "C" fn cb_invoke<T: McpTool>(
368        ctxt: *mut c_void,
369        call: *mut BNMcpToolCall,
370        arguments: *const c_char,
371        result: *mut BNMcpToolResult,
372    ) {
373        let context = unsafe { &*(ctxt as *const ToolContext<T>) };
374        let outcome = panic::catch_unwind(AssertUnwindSafe(|| {
375            let call = unsafe { McpToolCall::ref_from_raw(BNNewMcpToolCallReference(call)) };
376            let arguments = unsafe { CStr::from_ptr(arguments) }.to_string_lossy();
377            serde_json::from_str::<Value>(&arguments)
378                .map_err(|error| {
379                    McpToolError::invalid_params(format!("Arguments are not valid JSON: {error}"))
380                })
381                .and_then(|arguments| context.tool.invoke(&call, arguments))
382                .unwrap_or_else(McpToolResult::error)
383        }));
384        let tool_result = outcome.unwrap_or_else(|payload| {
385            let message = payload
386                .downcast_ref::<&str>()
387                .copied()
388                .or_else(|| payload.downcast_ref::<String>().map(String::as_str))
389                .unwrap_or("unknown panic");
390            tracing::error!("MCP tool '{}' panicked: {message}", context.name);
391            McpToolError::new("internal_error", format!("The tool panicked: {message}")).into()
392        });
393        tool_result.apply(result);
394    }
395
396    let definition = tool.definition();
397    let name = CString::new(definition.info.name.as_str()).ok()?;
398    let title = CString::new(definition.info.title.as_str()).ok()?;
399    let description = CString::new(definition.info.description.as_str()).ok()?;
400    let input_schema = definition.input_schema.to_string().to_cstr();
401    let output_schema = definition
402        .info
403        .output_schema
404        .as_ref()
405        .map(|schema| schema.to_string().to_cstr());
406    let api_definition = BNMcpToolDefinition {
407        name: name.as_ptr(),
408        title: title.as_ptr(),
409        description: description.as_ptr(),
410        inputSchema: input_schema.as_ptr(),
411        outputSchema: output_schema
412            .as_ref()
413            .map_or(ptr::null(), |schema| schema.as_ptr()),
414        scope: definition.info.scope,
415        annotations: definition.info.annotations.bits(),
416    };
417
418    extern "C" fn cb_free_object<T: McpTool>(ctxt: *mut c_void) {
419        unsafe { drop(Box::from_raw(ctxt as *mut ToolContext<T>)) };
420    }
421
422    let ctxt = Box::into_raw(Box::new(ToolContext {
423        name: definition.info.name.clone(),
424        tool,
425    }));
426    let callbacks = BNMcpToolCallbacks {
427        context: ctxt as *mut c_void,
428        invoke: Some(cb_invoke::<T>),
429        freeObject: Some(cb_free_object::<T>),
430    };
431    let handle = unsafe { create(&api_definition, &callbacks) };
432    if handle.is_null() {
433        unsafe { drop(Box::from_raw(ctxt)) };
434        return None;
435    }
436    Some(unsafe { CoreMcpTool::ref_from_raw(handle) })
437}
438
439/// A tool in the tool registry.
440#[derive(PartialEq, Eq, Hash)]
441pub struct CoreMcpTool {
442    pub(super) handle: *mut BNMcpTool,
443}
444
445impl CoreMcpTool {
446    pub(crate) unsafe fn from_raw(handle: *mut BNMcpTool) -> Self {
447        debug_assert!(!handle.is_null());
448        Self { handle }
449    }
450
451    pub(crate) unsafe fn ref_from_raw(handle: *mut BNMcpTool) -> Ref<Self> {
452        Ref::new(Self::from_raw(handle))
453    }
454
455    /// Every registered tool, sorted by name.
456    pub fn list() -> Array<CoreMcpTool> {
457        let mut count = 0;
458        let tools = unsafe { BNGetMcpToolList(&mut count) };
459        unsafe { Array::new(tools, count, ()) }
460    }
461
462    pub fn from_name(name: &str) -> Option<Ref<Self>> {
463        let name = name.to_cstr();
464        let handle = unsafe { BNGetMcpToolByName(name.as_ptr()) };
465        (!handle.is_null()).then(|| unsafe { Self::ref_from_raw(handle) })
466    }
467
468    pub fn name(&self) -> String {
469        unsafe { BnString::into_string(BNGetMcpToolName(self.handle)) }
470    }
471
472    pub fn title(&self) -> String {
473        unsafe { BnString::into_string(BNGetMcpToolTitle(self.handle)) }
474    }
475
476    pub fn description(&self) -> String {
477        unsafe { BnString::into_string(BNGetMcpToolDescription(self.handle)) }
478    }
479
480    pub fn input_schema(&self) -> Value {
481        let schema = unsafe { BnString::into_string(BNGetMcpToolInputSchema(self.handle)) };
482        serde_json::from_str(&schema).unwrap_or(Value::Null)
483    }
484
485    pub fn output_schema(&self) -> Option<Value> {
486        let schema = unsafe { BnString::into_string(BNGetMcpToolOutputSchema(self.handle)) };
487        serde_json::from_str(&schema).ok()
488    }
489
490    pub fn scope(&self) -> McpToolScope {
491        unsafe { BNGetMcpToolScope(self.handle) }
492    }
493
494    pub fn annotations(&self) -> McpToolAnnotations {
495        McpToolAnnotations::from_bits_retain(unsafe { BNGetMcpToolAnnotations(self.handle) })
496    }
497}
498
499impl fmt::Debug for CoreMcpTool {
500    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
501        f.debug_struct("CoreMcpTool")
502            .field("name", &self.name())
503            .finish()
504    }
505}
506
507unsafe impl Send for CoreMcpTool {}
508unsafe impl Sync for CoreMcpTool {}
509
510impl ToOwned for CoreMcpTool {
511    type Owned = Ref<Self>;
512
513    fn to_owned(&self) -> Self::Owned {
514        unsafe { RefCountable::inc_ref(self) }
515    }
516}
517
518unsafe impl RefCountable for CoreMcpTool {
519    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
520        Self::ref_from_raw(BNNewMcpToolReference(handle.handle))
521    }
522
523    unsafe fn dec_ref(handle: &Self) {
524        BNFreeMcpTool(handle.handle);
525    }
526}
527
528impl CoreArrayProvider for CoreMcpTool {
529    type Raw = *mut BNMcpTool;
530    type Context = ();
531    type Wrapped<'a> = Guard<'a, CoreMcpTool>;
532}
533
534unsafe impl CoreArrayProviderInner for CoreMcpTool {
535    unsafe fn free(raw: *mut Self::Raw, count: usize, _context: &Self::Context) {
536        BNFreeMcpToolList(raw, count);
537    }
538
539    unsafe fn wrap_raw<'a>(raw: &'a Self::Raw, context: &'a Self::Context) -> Self::Wrapped<'a> {
540        Guard::new(CoreMcpTool::from_raw(*raw), context)
541    }
542}
543
544/// An address expression string, evaluated against the call's binary view with
545/// [`McpAddress::resolve`].
546#[derive(Clone, Debug, PartialEq, Eq, serde_derive::Deserialize)]
547#[serde(transparent)]
548pub struct McpAddress(pub String);
549
550impl McpAddress {
551    /// An invalid expression is an `invalid_params` error.
552    pub fn resolve(&self, call: &McpToolCall) -> Result<u64, McpToolError> {
553        call.parse_address(&Value::String(self.0.clone()), 0)
554            .map_err(McpToolError::invalid_params)
555    }
556}
557
558/// A string argument that rejects an empty string.
559#[derive(Clone, Debug, PartialEq, Eq, Hash)]
560pub struct McpNonEmptyString(pub String);
561
562impl<'de> serde::Deserialize<'de> for McpNonEmptyString {
563    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
564        let value = <String as serde::Deserialize>::deserialize(deserializer)?;
565        if value.is_empty() {
566            return Err(serde::de::Error::invalid_value(
567                serde::de::Unexpected::Str(&value),
568                &"a non-empty string",
569            ));
570        }
571        Ok(McpNonEmptyString(value))
572    }
573}
574
575/// An unsigned integer, or an expression string evaluated against the call's binary view with
576/// [`McpIntegerExpression::resolve`].
577#[derive(Clone, Debug, PartialEq, Eq)]
578pub enum McpIntegerExpression {
579    Integer(u64),
580    Expression(String),
581}
582
583impl<'de> serde::Deserialize<'de> for McpIntegerExpression {
584    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
585        struct Visitor;
586
587        impl serde::de::Visitor<'_> for Visitor {
588            type Value = McpIntegerExpression;
589
590            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
591                formatter.write_str("an unsigned integer or an integer expression string")
592            }
593
594            fn visit_u64<E: serde::de::Error>(self, value: u64) -> Result<Self::Value, E> {
595                Ok(McpIntegerExpression::Integer(value))
596            }
597
598            fn visit_i64<E: serde::de::Error>(self, value: i64) -> Result<Self::Value, E> {
599                u64::try_from(value)
600                    .map(McpIntegerExpression::Integer)
601                    .map_err(|_| E::invalid_value(serde::de::Unexpected::Signed(value), &self))
602            }
603
604            fn visit_str<E: serde::de::Error>(self, value: &str) -> Result<Self::Value, E> {
605                Ok(McpIntegerExpression::Expression(value.to_owned()))
606            }
607        }
608
609        deserializer.deserialize_any(Visitor)
610    }
611}
612
613impl McpIntegerExpression {
614    /// An invalid expression is an `invalid_params` error.
615    pub fn resolve(&self, call: &McpToolCall) -> Result<u64, McpToolError> {
616        self.resolve_relative_to(call, 0)
617    }
618
619    /// Resolves an expression with `here` as the value of `$here`, such as a length measured from
620    /// an address argument.
621    pub fn resolve_relative_to(&self, call: &McpToolCall, here: u64) -> Result<u64, McpToolError> {
622        match self {
623            McpIntegerExpression::Integer(value) => Ok(*value),
624            McpIntegerExpression::Expression(expression) => call
625                .parse_integer(&Value::String(expression.clone()), here)
626                .map_err(McpToolError::invalid_params),
627        }
628    }
629}
630
631#[cfg(feature = "schemars")]
632pub(super) mod typed {
633    use super::*;
634    use schemars::{json_schema, JsonSchema, Schema, SchemaGenerator};
635    use serde::Deserialize;
636    use std::borrow::Cow;
637    use std::sync::OnceLock;
638
639    impl JsonSchema for McpAddress {
640        fn schema_name() -> Cow<'static, str> {
641            "McpAddress".into()
642        }
643
644        fn inline_schema() -> bool {
645            true
646        }
647
648        fn json_schema(_: &mut SchemaGenerator) -> Schema {
649            json_schema!({ "type": "string" })
650        }
651    }
652
653    impl JsonSchema for McpNonEmptyString {
654        fn schema_name() -> Cow<'static, str> {
655            "McpNonEmptyString".into()
656        }
657
658        fn inline_schema() -> bool {
659            true
660        }
661
662        fn json_schema(_: &mut SchemaGenerator) -> Schema {
663            json_schema!({ "type": "string", "minLength": 1 })
664        }
665    }
666
667    impl JsonSchema for McpIntegerExpression {
668        fn schema_name() -> Cow<'static, str> {
669            "McpIntegerExpression".into()
670        }
671
672        fn inline_schema() -> bool {
673            true
674        }
675
676        fn json_schema(_: &mut SchemaGenerator) -> Schema {
677            json_schema!({ "type": ["integer", "string"], "minimum": 0 })
678        }
679    }
680
681    /// A tool whose input schema and argument parsing come from its arguments type. Doc comments on
682    /// the arguments type's fields become parameter descriptions, and `Option` fields and
683    /// `#[serde(default)]` make parameters optional, with a null argument treated as absent.
684    /// Undeclared arguments produce an `invalid_params` result, unless the arguments type accepts
685    /// additional properties.
686    ///
687    /// The input schema must be an object, so `Args` is a struct with named fields, or with none for
688    /// a tool that takes no arguments. A unit struct or `()` describes `null` and is rejected.
689    pub trait TypedMcpTool: 'static + Send + Sync {
690        type Args: serde::de::DeserializeOwned + JsonSchema;
691
692        fn info(&self) -> McpToolInfo;
693
694        /// Called on an arbitrary thread with arguments that already match `Args`. A panic is
695        /// handled as it is for [`McpTool::invoke`].
696        fn invoke(
697            &self,
698            call: &McpToolCall,
699            args: Self::Args,
700        ) -> Result<McpToolResult, McpToolError>;
701    }
702
703    struct Parameters {
704        // The only argument names `invoke` accepts, or `None` when it accepts any name.
705        accepted: Option<Vec<String>>,
706        // The parameters whose null arguments are treated as absent.
707        optional: Vec<String>,
708    }
709
710    pub(in crate::mcp) struct TypedAdapter<T> {
711        tool: T,
712        parameters: OnceLock<Parameters>,
713    }
714
715    impl<T: TypedMcpTool> TypedAdapter<T> {
716        pub(in crate::mcp) fn new(tool: T) -> Self {
717            Self {
718                tool,
719                parameters: OnceLock::new(),
720            }
721        }
722
723        fn input_schema() -> Value {
724            let settings = schemars::generate::SchemaSettings::draft2020_12().with(|settings| {
725                settings.inline_subschemas = true;
726            });
727            let mut schema = settings
728                .into_generator()
729                .into_root_schema_for::<T::Args>()
730                .to_value();
731            if let Some(object) = schema.as_object_mut() {
732                object.remove("$schema");
733                object.remove("title");
734                if object.get("type") == Some(&Value::String("object".into()))
735                    && !object.contains_key("additionalProperties")
736                {
737                    object.insert("additionalProperties".into(), Value::Bool(false));
738                }
739            }
740            schema
741        }
742
743        fn parameters(schema: &Value) -> Parameters {
744            let no_properties = serde_json::Map::new();
745            let properties = schema
746                .get("properties")
747                .and_then(Value::as_object)
748                .unwrap_or(&no_properties);
749            let required = schema.get("required").and_then(Value::as_array);
750            let is_required = |name: &String| {
751                required.is_some_and(|required| required.iter().any(|entry| entry == name))
752            };
753            let accepted = (schema.get("additionalProperties") == Some(&Value::Bool(false)))
754                .then(|| properties.keys().cloned().collect());
755            let optional = properties
756                .keys()
757                .filter(|name| !is_required(name))
758                .cloned()
759                .collect();
760            Parameters { accepted, optional }
761        }
762    }
763
764    /// Deserializes the arguments object as `Args`, naming the parameter in each error.
765    struct Arguments<'de>(&'de serde_json::Map<String, Value>);
766
767    impl<'de> serde::Deserializer<'de> for Arguments<'de> {
768        type Error = serde_json::Error;
769
770        fn deserialize_any<V: serde::de::Visitor<'de>>(
771            self,
772            visitor: V,
773        ) -> Result<V::Value, Self::Error> {
774            visitor.visit_map(ArgumentsMap {
775                entries: self.0.iter(),
776                value: None,
777            })
778        }
779
780        serde::forward_to_deserialize_any! {
781            bool i8 i16 i32 i64 i128 u8 u16 u32 u64 u128 f32 f64 char str string bytes byte_buf
782            option unit unit_struct newtype_struct seq tuple tuple_struct map struct enum identifier
783            ignored_any
784        }
785    }
786
787    struct ArgumentsMap<'de> {
788        entries: serde_json::map::Iter<'de>,
789        value: Option<(&'de String, &'de Value)>,
790    }
791
792    impl<'de> serde::de::MapAccess<'de> for ArgumentsMap<'de> {
793        type Error = serde_json::Error;
794
795        fn next_key_seed<K: serde::de::DeserializeSeed<'de>>(
796            &mut self,
797            seed: K,
798        ) -> Result<Option<K::Value>, Self::Error> {
799            use serde::de::IntoDeserializer;
800
801            let Some((name, value)) = self.entries.next() else {
802                return Ok(None);
803            };
804            self.value = Some((name, value));
805            seed.deserialize(name.as_str().into_deserializer())
806                .map(Some)
807        }
808
809        fn next_value_seed<S: serde::de::DeserializeSeed<'de>>(
810            &mut self,
811            seed: S,
812        ) -> Result<S::Value, Self::Error> {
813            let (name, value) = self
814                .value
815                .take()
816                .expect("next_value_seed called before next_key_seed");
817            seed.deserialize(value).map_err(|error| {
818                serde::de::Error::custom(format!("Invalid parameter '{name}': {error}"))
819            })
820        }
821    }
822
823    impl<T: TypedMcpTool> McpTool for TypedAdapter<T> {
824        fn definition(&self) -> McpToolDefinition {
825            self.tool.info().with_input_schema(Self::input_schema())
826        }
827
828        fn invoke(
829            &self,
830            call: &McpToolCall,
831            mut arguments: Value,
832        ) -> Result<McpToolResult, McpToolError> {
833            let parameters = self
834                .parameters
835                .get_or_init(|| Self::parameters(&Self::input_schema()));
836            let Some(arguments) = arguments.as_object_mut() else {
837                return Err(McpToolError::invalid_params("Expected object arguments"));
838            };
839            if let Some(accepted) = &parameters.accepted {
840                if let Some(name) = arguments.keys().find(|name| !accepted.contains(name)) {
841                    return Err(McpToolError::invalid_params(format!(
842                        "Unexpected parameter '{name}'"
843                    )));
844                }
845            }
846            arguments.retain(|name, value| !value.is_null() || !parameters.optional.contains(name));
847            let args = T::Args::deserialize(Arguments(arguments))?;
848            self.tool.invoke(call, args)
849        }
850    }
851
852    /// Registers a [`TypedMcpTool`] for the life of the process. See [`register_mcp_tool`].
853    pub fn register_typed_mcp_tool<T: TypedMcpTool>(tool: T) -> Option<Ref<CoreMcpTool>> {
854        register_mcp_tool(TypedAdapter::new(tool))
855    }
856}
857
858#[cfg(feature = "schemars")]
859pub use typed::{register_typed_mcp_tool, TypedMcpTool};
860
861/// A JSON Schema object for a tool that takes no arguments.
862pub fn empty_input_schema() -> Value {
863    json!({ "type": "object", "additionalProperties": false })
864}
binaryninja/mcp/tool.rs

binaryninja/mcp/
tool.rs
