rmcp_server_kit/
tool_hooks.rs

1//! Opt-in tool-call instrumentation for `ServerHandler` implementations.
2//!
3//! [`crate::tool_hooks::HookedHandler`] wraps any [`rmcp::ServerHandler`] with:
4//!
5//! - **Before hooks** (async) that observe `(tool_name, arguments, identity,
6//!   role, sub, request_id)` and may [`HookOutcome::Continue`](crate::tool_hooks::HookOutcome::Continue),
7//!   [`HookOutcome::Deny`](crate::tool_hooks::HookOutcome::Deny), or
8//!   [`HookOutcome::Replace`](crate::tool_hooks::HookOutcome::Replace) the call.
9//! - **After hooks** (async) that observe the same context plus a
10//!   [`HookDisposition`](crate::tool_hooks::HookDisposition) describing how the call resolved and the
11//!   approximate result size in bytes.  After-hooks are spawned via
12//!   `tokio::spawn` and never block the response path.
13//! - **Result-size capping**: serialized tool results larger than
14//!   `max_result_bytes` are replaced with a structured error, preventing
15//!   token-expensive or memory-expensive payloads from reaching clients.
16//!   The cap applies both to inner-handler results and to
17//!   [`HookOutcome::Replace`](crate::tool_hooks::HookOutcome::Replace) payloads.
18//!
19//! This is entirely **opt-in** at the application layer - `rmcp_server_kit::serve()`
20//! does not wrap handlers automatically.  Applications that want hooks do:
21//!
22//! ```no_run
23//! use std::sync::Arc;
24//! use rmcp_server_kit::tool_hooks::{HookedHandler, HookOutcome, ToolHooks, with_hooks};
25//!
26//! # #[derive(Clone, Default)]
27//! # struct MyHandler;
28//! # impl rmcp::ServerHandler for MyHandler {}
29//! let handler = MyHandler::default();
30//! let hooks = Arc::new(
31//!     ToolHooks::new()
32//!         .with_max_result_bytes(256 * 1024)
33//!         .with_before(Arc::new(|_ctx| Box::pin(async { HookOutcome::Continue })))
34//!         .with_after(Arc::new(|_ctx, _disp, _bytes| Box::pin(async {}))),
35//! );
36//! let _wrapped = with_hooks(handler, hooks);
37//! ```
38
39use std::{borrow::Cow, fmt, future::Future, io, pin::Pin, sync::Arc};
40
41#[allow(
42    deprecated,
43    reason = "transparent ServerHandler delegation must import legacy logging/subscription parameter types until rmcp removes those methods"
44)]
45use rmcp::{
46    ErrorData, RoleServer, ServerHandler,
47    model::{
48        CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams,
49        CancelledNotificationParam, CompleteRequestParams, CompleteResult, ContentBlock,
50        CustomNotification, CustomRequest, CustomResult, DiscoverResult, GetPromptRequestParams,
51        GetPromptResponse, GetTaskParams, GetTaskResult, InitializeRequestParams, InitializeResult,
52        ListPromptsResult, ListResourceTemplatesResult, ListResourcesResult, ListToolsResult,
53        PaginatedRequestParams, ProgressNotificationParam, ProtocolVersion,
54        ReadResourceRequestParams, ReadResourceResponse, ServerInfo, SetLevelRequestParams,
55        SubscribeRequestParams, SubscriptionFilter, Tool, UnsubscribeRequestParams,
56        UpdateTaskParams,
57    },
58    service::{NotificationContext, RequestContext, SubscriptionContext},
59};
60
61/// Context passed to before/after hooks for a single tool call.
62#[derive(Clone)]
63#[non_exhaustive]
64pub struct ToolCallContext {
65    /// Tool name being invoked.
66    pub tool_name: String,
67    /// JSON arguments as sent by the client (may be `None`).
68    pub arguments: Option<serde_json::Value>,
69    /// Identity name from the authenticated request, if any.
70    pub identity: Option<String>,
71    /// RBAC role associated with the request, if any.
72    pub role: Option<String>,
73    /// OAuth `sub` claim, if present.
74    pub sub: Option<String>,
75    /// Raw JSON-RPC request id rendered as a string, if available.
76    pub request_id: Option<String>,
77}
78
79impl ToolCallContext {
80    /// Construct a [`ToolCallContext`] with the given tool name and all
81    /// optional fields cleared.  Primarily for use in unit tests and
82    /// benchmarks of user-supplied hooks; the runtime path populates
83    /// these fields from the request and task-local RBAC state.
84    #[must_use]
85    pub fn for_tool(tool_name: impl Into<String>) -> Self {
86        Self {
87            tool_name: tool_name.into(),
88            arguments: None,
89            identity: None,
90            role: None,
91            sub: None,
92            request_id: None,
93        }
94    }
95}
96
97impl fmt::Debug for ToolCallContext {
98    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
99        let Self {
100            tool_name,
101            arguments,
102            identity,
103            role,
104            sub,
105            request_id,
106        } = self;
107        let mut debug = f.debug_struct("ToolCallContext");
108        debug.field("tool_name", tool_name);
109        if crate::diagnostics::tool_call_arguments() {
110            debug
111                .field("arguments", arguments)
112                .field("identity", identity)
113                .field("role", role)
114                .field("sub", sub);
115        } else {
116            debug
117                .field("arguments", &"[REDACTED]")
118                .field("identity", &"[REDACTED]")
119                .field("role", &"[REDACTED]")
120                .field("sub", &"[REDACTED]");
121        }
122        debug.field("request_id", request_id).finish()
123    }
124}
125
126/// Outcome returned by a [`BeforeHook`] to control invocation flow.
127///
128/// - [`HookOutcome::Continue`] - proceed with the wrapped handler.
129/// - [`HookOutcome::Deny`] - reject the call with the supplied
130///   [`ErrorData`]; the inner handler is **not** called.
131/// - [`HookOutcome::Replace`] - return the supplied result instead of
132///   invoking the inner handler.  The result is still subject to
133///   `max_result_bytes` capping.
134#[derive(Debug)]
135#[non_exhaustive]
136pub enum HookOutcome {
137    /// Proceed with the wrapped handler.
138    Continue,
139    /// Reject the call.  The error is propagated to the client as-is.
140    Deny(ErrorData),
141    /// Skip the inner handler and return the supplied result instead.
142    Replace(Box<CallToolResult>),
143}
144
145/// How a tool call resolved, passed to the [`AfterHook`].
146#[derive(Debug, Clone, Copy)]
147#[non_exhaustive]
148pub enum HookDisposition {
149    /// The inner handler ran and returned `Ok`.
150    InnerExecuted,
151    /// The inner handler ran and returned `Err`.
152    InnerErrored,
153    /// The before-hook returned [`HookOutcome::Deny`].
154    DeniedBefore,
155    /// The before-hook returned [`HookOutcome::Replace`].
156    ReplacedBefore,
157    /// The result (from inner or replace) exceeded `max_result_bytes`
158    /// and was substituted with a structured error.
159    ResultTooLarge,
160}
161
162/// Async before-hook callback type.
163///
164/// Returns a [`HookOutcome`] controlling whether the inner handler runs.
165/// The borrow of `ToolCallContext` is held for the duration of the
166/// returned future, which avoids forcing implementations to clone the
167/// context for every invocation.
168pub type BeforeHook = Arc<
169    dyn for<'a> Fn(&'a ToolCallContext) -> Pin<Box<dyn Future<Output = HookOutcome> + Send + 'a>>
170        + Send
171        + Sync
172        + 'static,
173>;
174
175/// Async after-hook callback type.
176///
177/// Receives the call context, a [`HookDisposition`] describing how the
178/// call resolved, and the approximate serialized result size in bytes
179/// (`0` for `DeniedBefore` and `InnerErrored`).  Spawned via
180/// `tokio::spawn`, so it must not assume it runs before the response is
181/// flushed.
182pub type AfterHook = Arc<
183    dyn for<'a> Fn(
184            &'a ToolCallContext,
185            HookDisposition,
186            usize,
187        ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
188        + Send
189        + Sync
190        + 'static,
191>;
192
193/// Opt-in hooks applied by [`crate::tool_hooks::HookedHandler`].
194#[allow(clippy::struct_field_names, reason = "before/after read naturally")]
195#[derive(Clone, Default)]
196#[non_exhaustive]
197pub struct ToolHooks {
198    /// Hard cap on serialized `CallToolResult` size in bytes.  When
199    /// exceeded, the result is replaced with an `is_error=true` result
200    /// carrying a `result_too_large` structured error.  `None` disables
201    /// the cap.
202    pub max_result_bytes: Option<usize>,
203    /// Optional before-hook invoked after arg deserialization, before
204    /// the wrapped handler is called.
205    pub before: Option<BeforeHook>,
206    /// Optional after-hook invoked once per call, regardless of how the
207    /// call resolved.  Spawned via `tokio::spawn` and never blocks the
208    /// response path.
209    pub after: Option<AfterHook>,
210}
211
212impl ToolHooks {
213    /// Construct an empty [`ToolHooks`] with no cap and no hooks.
214    ///
215    /// Use the `with_*` builder methods to populate fields; this avoids
216    /// the `#[non_exhaustive]` restriction that prevents struct-literal
217    /// construction from outside the crate.
218    #[must_use]
219    pub fn new() -> Self {
220        Self::default()
221    }
222
223    /// Set the serialized result size cap in bytes.
224    #[must_use]
225    pub fn with_max_result_bytes(mut self, max: usize) -> Self {
226        self.max_result_bytes = Some(max);
227        self
228    }
229
230    /// Set the before-hook.
231    #[must_use]
232    pub fn with_before(mut self, before: BeforeHook) -> Self {
233        self.before = Some(before);
234        self
235    }
236
237    /// Set the after-hook.
238    #[must_use]
239    pub fn with_after(mut self, after: AfterHook) -> Self {
240        self.after = Some(after);
241        self
242    }
243}
244
245const _HOOKED_HANDLER_DOC_ANCHOR: &str = "HookedHandler";
246
247impl fmt::Debug for ToolHooks {
248    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
249        f.debug_struct("ToolHooks")
250            .field("max_result_bytes", &self.max_result_bytes)
251            .field("before", &self.before.as_ref().map(|_| "<fn>"))
252            .field("after", &self.after.as_ref().map(|_| "<fn>"))
253            .finish()
254    }
255}
256
257/// `ServerHandler` wrapper that applies [`ToolHooks`].
258#[derive(Clone)]
259pub struct HookedHandler<H: ServerHandler> {
260    inner: Arc<H>,
261    hooks: Arc<ToolHooks>,
262}
263
264impl<H: ServerHandler> fmt::Debug for HookedHandler<H> {
265    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
266        f.debug_struct("HookedHandler")
267            .field("hooks", &self.hooks)
268            .finish_non_exhaustive()
269    }
270}
271
272/// Construct a [`crate::tool_hooks::HookedHandler`] from an inner handler and hooks.
273///
274/// Returning the wrapped handler is the entire point of this function;
275/// dropping it on the floor would silently disable the supplied hooks.
276#[must_use = "HookedHandler must be wired into a ServerHandler (e.g. via \
277              `serve(..., || hooked)`) to take effect; dropping the returned \
278              value silently disables the supplied hooks"]
279pub fn with_hooks<H: ServerHandler>(inner: H, hooks: Arc<ToolHooks>) -> HookedHandler<H> {
280    HookedHandler {
281        inner: Arc::new(inner),
282        hooks,
283    }
284}
285
286impl<H: ServerHandler> HookedHandler<H> {
287    /// Access the wrapped handler.
288    #[must_use]
289    pub fn inner(&self) -> &H {
290        &self.inner
291    }
292
293    fn build_context(request: &CallToolRequestParams, req_id: Option<String>) -> ToolCallContext {
294        ToolCallContext {
295            tool_name: request.name.to_string(),
296            arguments: request.arguments.clone().map(serde_json::Value::Object),
297            identity: crate::rbac::current_identity(),
298            role: crate::rbac::current_role(),
299            sub: crate::rbac::current_sub(),
300            request_id: req_id,
301        }
302    }
303
304    /// Spawn the after-hook on the current Tokio runtime.  The future
305    /// captures clones of `ctx` and the `Arc<AfterHook>` so it can run
306    /// independently of the request task; panics inside the after-hook
307    /// are caught by Tokio and never poison the response path.
308    ///
309    /// The spawned task is **instrumented** with the request span via
310    /// [`tracing::Instrument`] and re-establishes the per-request RBAC
311    /// task-locals (role, identity, token, sub) via
312    /// [`crate::rbac::with_rbac_scope`]. Without this, after-hooks lose
313    /// their parent span (breaking trace correlation) and observe
314    /// `current_role()` / `current_identity()` as `None`.
315    fn spawn_after(
316        after: Option<&Arc<AfterHookHolder>>,
317        ctx: ToolCallContext,
318        disposition: HookDisposition,
319        size: usize,
320    ) {
321        if let Some(after) = after {
322            use tracing::Instrument;
323
324            let after = Arc::clone(after);
325            // Capture the request span before leaving the request task so
326            // after-hook log lines are correlated with the originating call.
327            let span = tracing::Span::current();
328            // Snapshot RBAC task-locals; defaults are empty strings so the
329            // re-established scope is a no-op when the request had no
330            // authenticated identity (e.g. health checks, anonymous tools).
331            let role = crate::rbac::current_role().unwrap_or_default();
332            let identity = crate::rbac::current_identity().unwrap_or_default();
333            let token = crate::rbac::current_token()
334                .unwrap_or_else(|| secrecy::SecretString::from(String::new()));
335            let sub = crate::rbac::current_sub().unwrap_or_default();
336            tokio::spawn(
337                async move {
338                    crate::rbac::with_rbac_scope(role, identity, token, sub, async move {
339                        let fut = (after.f)(&ctx, disposition, size);
340                        fut.await;
341                    })
342                    .await;
343                }
344                .instrument(span),
345            );
346        }
347    }
348}
349
350/// Internal newtype that owns the [`AfterHook`] so we can `Arc::clone`
351/// the *holder* and let the spawned task borrow `ctx` for the lifetime
352/// of the future without lifetime acrobatics in `tokio::spawn`.
353struct AfterHookHolder {
354    f: AfterHook,
355}
356
357/// Structured error body returned when a result exceeds `max_result_bytes`.
358///
359/// `actual` is `None` when the result could not be serialized, so its true
360/// size is unknown. It is rendered as `"unknown"` rather than a fabricated
361/// number -- operators read `actual_bytes` as a measurement.
362fn too_large_result(limit: usize, actual: Option<usize>, tool: &str) -> CallToolResult {
363    let actual_desc =
364        actual.map_or_else(|| "an unmeasurable number of".to_owned(), |n| n.to_string());
365    let body = serde_json::json!({
366        "error": "result_too_large",
367        "message": format!(
368            "tool '{tool}' result of {actual_desc} bytes exceeds the configured \
369             max_result_bytes={limit}; ask for a narrower query"
370        ),
371        "limit_bytes": limit,
372        "actual_bytes": actual.map_or_else(
373            || serde_json::Value::from("unknown"),
374            serde_json::Value::from,
375        ),
376    });
377    let mut r = CallToolResult::error(vec![ContentBlock::text(body.to_string())]);
378    r.structured_content = None;
379    r
380}
381
382/// Outcome of the `max_result_bytes` policy for a measured -- or
383/// unmeasurable -- result.
384#[derive(Debug, PartialEq, Eq)]
385enum SizeVerdict {
386    /// Within the cap, or no cap configured. Carries the measured size.
387    Pass { size: usize },
388    /// Over the cap, or unmeasurable while a cap is configured.
389    Replace { limit: usize, actual: Option<usize> },
390    /// Unmeasurable and no cap configured: nothing to enforce.
391    PassUnmeasured,
392}
393
394/// Decide what the size cap does, given an optional size-measurement outcome.
395const fn decide_size(size: Option<SizeMeasure>, max: Option<usize>) -> SizeVerdict {
396    match size {
397        Some(SizeMeasure::Exact(size)) => match max {
398            Some(limit) if size > limit => SizeVerdict::Replace {
399                limit,
400                actual: Some(size),
401            },
402            Some(_) | None => SizeVerdict::Pass { size },
403        },
404        Some(SizeMeasure::Exceeded { limit }) => SizeVerdict::Replace {
405            limit,
406            actual: None,
407        },
408        None => match max {
409            Some(limit) => SizeVerdict::Replace {
410                limit,
411                actual: None,
412            },
413            None => SizeVerdict::PassUnmeasured,
414        },
415    }
416}
417
418/// Apply the `max_result_bytes` cap to a result.  Returns the (possibly
419/// replaced) result, the size used for accounting, and whether the cap
420/// fired.
421fn apply_size_cap(
422    result: CallToolResult,
423    max: Option<usize>,
424    tool: &str,
425) -> (CallToolResult, usize, bool) {
426    let size = if max.is_some() {
427        Some(serialized_size(&result, max))
428    } else {
429        None
430    };
431    match decide_size(size, max) {
432        SizeVerdict::Pass { size } => (result, size, false),
433        SizeVerdict::PassUnmeasured => (result, 0, false),
434        SizeVerdict::Replace { limit, actual } => {
435            tracing::warn!(
436                tool = %tool,
437                size_bytes = actual.unwrap_or_default(),
438                size_measured = actual.is_some(),
439                limit_bytes = limit,
440                "tool result exceeds max_result_bytes; replacing with structured error"
441            );
442            let accounted = actual.unwrap_or_else(|| limit.saturating_add(1));
443            (too_large_result(limit, actual, tool), accounted, true)
444        }
445    }
446}
447
448#[allow(
449    deprecated,
450    reason = "transparent ServerHandler delegation must include legacy logging/subscription methods until rmcp removes them"
451)]
452impl<H: ServerHandler> ServerHandler for HookedHandler<H> {
453    async fn ping(&self, context: RequestContext<RoleServer>) -> Result<(), ErrorData> {
454        self.inner.ping(context).await
455    }
456
457    fn get_info(&self) -> ServerInfo {
458        self.inner.get_info()
459    }
460
461    async fn initialize(
462        &self,
463        request: InitializeRequestParams,
464        context: RequestContext<RoleServer>,
465    ) -> Result<InitializeResult, ErrorData> {
466        self.inner.initialize(request, context).await
467    }
468
469    async fn list_tools(
470        &self,
471        request: Option<PaginatedRequestParams>,
472        context: RequestContext<RoleServer>,
473    ) -> Result<ListToolsResult, ErrorData> {
474        self.inner.list_tools(request, context).await
475    }
476
477    async fn complete(
478        &self,
479        request: CompleteRequestParams,
480        context: RequestContext<RoleServer>,
481    ) -> Result<CompleteResult, ErrorData> {
482        self.inner.complete(request, context).await
483    }
484
485    async fn set_level(
486        &self,
487        request: SetLevelRequestParams,
488        context: RequestContext<RoleServer>,
489    ) -> Result<(), ErrorData> {
490        self.inner.set_level(request, context).await
491    }
492
493    fn get_tool(&self, name: &str) -> Option<Tool> {
494        self.inner.get_tool(name)
495    }
496
497    async fn list_prompts(
498        &self,
499        request: Option<PaginatedRequestParams>,
500        context: RequestContext<RoleServer>,
501    ) -> Result<ListPromptsResult, ErrorData> {
502        self.inner.list_prompts(request, context).await
503    }
504
505    async fn get_prompt(
506        &self,
507        request: GetPromptRequestParams,
508        context: RequestContext<RoleServer>,
509    ) -> Result<GetPromptResponse, ErrorData> {
510        self.inner.get_prompt(request, context).await
511    }
512
513    async fn list_resources(
514        &self,
515        request: Option<PaginatedRequestParams>,
516        context: RequestContext<RoleServer>,
517    ) -> Result<ListResourcesResult, ErrorData> {
518        self.inner.list_resources(request, context).await
519    }
520
521    async fn list_resource_templates(
522        &self,
523        request: Option<PaginatedRequestParams>,
524        context: RequestContext<RoleServer>,
525    ) -> Result<ListResourceTemplatesResult, ErrorData> {
526        self.inner.list_resource_templates(request, context).await
527    }
528
529    async fn read_resource(
530        &self,
531        request: ReadResourceRequestParams,
532        context: RequestContext<RoleServer>,
533    ) -> Result<ReadResourceResponse, ErrorData> {
534        self.inner.read_resource(request, context).await
535    }
536
537    // NOT cancel-safe: this awaits consumer-supplied before-hooks and the
538    // consumer's inner handler. After-hooks are dispatched only on the normal
539    // Deny/Replace/Ok/Err paths, so a cancellation between the before-hook and
540    // the response drops the paired after-hook -- an audit hook can therefore
541    // record a started call that is never closed out. Consumers needing
542    // guaranteed pairing should make the after-hook idempotent or run the tool
543    // body detached (see `crate::cancel`).
544    #[allow(
545        clippy::wildcard_enum_match_arm,
546        reason = "CallToolResponse is #[non_exhaustive]; the non-Complete MRTR variants (InputRequired/Task) are passed through unchanged"
547    )]
548    async fn call_tool(
549        &self,
550        request: CallToolRequestParams,
551        context: RequestContext<RoleServer>,
552    ) -> Result<CallToolResponse, ErrorData> {
553        let req_id = Some(format!("{:?}", context.id));
554        let ctx = Self::build_context(&request, req_id);
555        let max = self.hooks.max_result_bytes;
556        let after_holder = self
557            .hooks
558            .after
559            .as_ref()
560            .map(|f| Arc::new(AfterHookHolder { f: Arc::clone(f) }));
561
562        // Before hook: may Continue, Deny, or Replace.
563        if let Some(before) = self.hooks.before.as_ref() {
564            let outcome = before(&ctx).await;
565            match outcome {
566                HookOutcome::Continue => {}
567                HookOutcome::Deny(err) => {
568                    Self::spawn_after(after_holder.as_ref(), ctx, HookDisposition::DeniedBefore, 0);
569                    return Err(err);
570                }
571                HookOutcome::Replace(boxed) => {
572                    let (final_result, size, capped) = apply_size_cap(*boxed, max, &ctx.tool_name);
573                    let disposition = if capped {
574                        HookDisposition::ResultTooLarge
575                    } else {
576                        HookDisposition::ReplacedBefore
577                    };
578                    Self::spawn_after(after_holder.as_ref(), ctx, disposition, size);
579                    return Ok(final_result.into());
580                }
581            }
582        }
583
584        // Inner handler.
585        match self.inner.call_tool(request, context).await {
586            // Completed tool result: subject to the size cap + after hook.
587            Ok(CallToolResponse::Complete(result)) => {
588                let (final_result, size, capped) = apply_size_cap(result, max, &ctx.tool_name);
589                let disposition = if capped {
590                    HookDisposition::ResultTooLarge
591                } else {
592                    HookDisposition::InnerExecuted
593                };
594                Self::spawn_after(after_holder.as_ref(), ctx, disposition, size);
595                Ok(final_result.into())
596            }
597            // MRTR input-required / task responses (rmcp 3.0): no CallToolResult
598            // to size-cap, so pass them through unchanged.
599            Ok(other) => {
600                Self::spawn_after(
601                    after_holder.as_ref(),
602                    ctx,
603                    HookDisposition::InnerExecuted,
604                    0,
605                );
606                Ok(other)
607            }
608            Err(e) => {
609                Self::spawn_after(after_holder.as_ref(), ctx, HookDisposition::InnerErrored, 0);
610                Err(e)
611            }
612        }
613    }
614
615    // rmcp 3.0 added task/subscription/discovery request handlers with defaults;
616    // delegate them to `inner` so wrapping a handler that implements those stays
617    // transparent (otherwise the default would shadow the inner implementation).
618    fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
619        self.inner.supported_protocol_versions()
620    }
621
622    async fn discover(
623        &self,
624        context: RequestContext<RoleServer>,
625    ) -> Result<DiscoverResult, ErrorData> {
626        self.inner.discover(context).await
627    }
628
629    fn accepted_subscription_filter(
630        &self,
631        requested: &SubscriptionFilter,
632    ) -> Option<SubscriptionFilter> {
633        self.inner.accepted_subscription_filter(requested)
634    }
635
636    async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> {
637        self.inner.listen(context).await
638    }
639
640    async fn subscribe(
641        &self,
642        request: SubscribeRequestParams,
643        context: RequestContext<RoleServer>,
644    ) -> Result<(), ErrorData> {
645        self.inner.subscribe(request, context).await
646    }
647
648    async fn unsubscribe(
649        &self,
650        request: UnsubscribeRequestParams,
651        context: RequestContext<RoleServer>,
652    ) -> Result<(), ErrorData> {
653        self.inner.unsubscribe(request, context).await
654    }
655
656    async fn get_task(
657        &self,
658        request: GetTaskParams,
659        context: RequestContext<RoleServer>,
660    ) -> Result<GetTaskResult, ErrorData> {
661        self.inner.get_task(request, context).await
662    }
663
664    async fn update_task(
665        &self,
666        request: UpdateTaskParams,
667        context: RequestContext<RoleServer>,
668    ) -> Result<(), ErrorData> {
669        self.inner.update_task(request, context).await
670    }
671
672    async fn cancel_task(
673        &self,
674        request: CancelTaskParams,
675        context: RequestContext<RoleServer>,
676    ) -> Result<(), ErrorData> {
677        self.inner.cancel_task(request, context).await
678    }
679
680    async fn on_custom_request(
681        &self,
682        request: CustomRequest,
683        context: RequestContext<RoleServer>,
684    ) -> Result<CustomResult, ErrorData> {
685        self.inner.on_custom_request(request, context).await
686    }
687
688    async fn on_cancelled(
689        &self,
690        notification: CancelledNotificationParam,
691        context: NotificationContext<RoleServer>,
692    ) {
693        self.inner.on_cancelled(notification, context).await;
694    }
695
696    async fn on_progress(
697        &self,
698        notification: ProgressNotificationParam,
699        context: NotificationContext<RoleServer>,
700    ) {
701        self.inner.on_progress(notification, context).await;
702    }
703
704    async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
705        self.inner.on_initialized(context).await;
706    }
707
708    async fn on_roots_list_changed(&self, context: NotificationContext<RoleServer>) {
709        self.inner.on_roots_list_changed(context).await;
710    }
711
712    async fn on_custom_notification(
713        &self,
714        notification: CustomNotification,
715        context: NotificationContext<RoleServer>,
716    ) {
717        self.inner
718            .on_custom_notification(notification, context)
719            .await;
720    }
721}
722
723#[derive(Debug, Clone, Copy, PartialEq, Eq)]
724struct SizeLimitExceeded;
725
726impl fmt::Display for SizeLimitExceeded {
727    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
728        f.write_str("serialized result exceeded configured size cap")
729    }
730}
731
732impl std::error::Error for SizeLimitExceeded {}
733
734struct CountingWriter {
735    bytes: usize,
736    limit: Option<usize>,
737}
738
739impl CountingWriter {
740    const fn unbounded() -> Self {
741        Self {
742            bytes: 0,
743            limit: None,
744        }
745    }
746
747    const fn bounded(limit: usize) -> Self {
748        Self {
749            bytes: 0,
750            limit: Some(limit),
751        }
752    }
753}
754
755impl io::Write for CountingWriter {
756    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
757        let next = self.bytes.saturating_add(buf.len());
758        if self.limit.is_some_and(|limit| next > limit) {
759            Err(io::Error::other(SizeLimitExceeded))
760        } else {
761            self.bytes = next;
762            Ok(buf.len())
763        }
764    }
765
766    fn flush(&mut self) -> io::Result<()> {
767        Ok(())
768    }
769}
770
771/// Outcome of measuring serialized result size.
772#[derive(Debug, Clone, Copy, PartialEq, Eq)]
773enum SizeMeasure {
774    /// Exact serialized size in bytes.
775    Exact(usize),
776    /// Serialization crossed the configured size cap and stopped early.
777    Exceeded { limit: usize },
778}
779
780/// Serialized byte length, or a deliberate cap-abort outcome.
781fn serialized_size(result: &CallToolResult, max: Option<usize>) -> SizeMeasure {
782    let mut writer = max.map_or_else(CountingWriter::unbounded, CountingWriter::bounded);
783    match serde_json::to_writer(&mut writer, result) {
784        Ok(()) => SizeMeasure::Exact(writer.bytes),
785        Err(error) if error.io_error_kind() == Some(io::ErrorKind::Other) => {
786            SizeMeasure::Exceeded {
787                limit: max.unwrap_or(writer.bytes),
788            }
789        }
790        Err(_error) => {
791            // `CallToolResult` is made only of infallibly serializable fields
792            // (`String`, `bool`, arrays/maps, and serde_json::Value`). There is
793            // no inhabitable production value that can reach this branch.
794            SizeMeasure::Exact(writer.bytes)
795        }
796    }
797}
798
799#[cfg(test)]
800mod tests {
801    use std::sync::{
802        Arc,
803        atomic::{AtomicUsize, Ordering},
804    };
805
806    #[allow(
807        deprecated,
808        reason = "delegation tests cover legacy logging/subscription methods"
809    )]
810    use rmcp::{
811        ErrorData, RoleServer, ServerHandler,
812        model::{
813            CallToolRequestParams, CallToolResponse, CallToolResult, CancelledNotificationParam,
814            CompleteRequestParams, CompleteResult, CompletionInfo, ContentBlock,
815            CustomNotification, CustomRequest, CustomResult, ProgressNotificationParam, ServerInfo,
816            SetLevelRequestParams, SubscribeRequestParams, UnsubscribeRequestParams,
817        },
818        service::RequestContext,
819    };
820    use serde_json::json;
821    use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream};
822
823    use super::*;
824
825    type DelegationTransport = (
826        DelegationProbe,
827        BufReader<tokio::io::ReadHalf<DuplexStream>>,
828        tokio::io::WriteHalf<DuplexStream>,
829        rmcp::service::RunningService<RoleServer, HookedHandler<DelegationProbe>>,
830    );
831
832    /// Maintenance aid only: rmcp gives every `ServerHandler` method a default,
833    /// so this count cannot guarantee compile-time completeness. It makes future
834    /// upstream method additions visible in review alongside the delegation impl.
835    const HOOKED_HANDLER_DELEGATED_METHODS: &[&str] = &[
836        "ping",
837        "initialize",
838        "supported_protocol_versions",
839        "discover",
840        "complete",
841        "set_level",
842        "get_prompt",
843        "list_prompts",
844        "list_resources",
845        "list_resource_templates",
846        "read_resource",
847        "accepted_subscription_filter",
848        "listen",
849        "subscribe",
850        "unsubscribe",
851        "call_tool",
852        "list_tools",
853        "get_tool",
854        "on_custom_request",
855        "on_cancelled",
856        "on_progress",
857        "on_initialized",
858        "on_roots_list_changed",
859        "on_custom_notification",
860        "get_info",
861        "get_task",
862        "update_task",
863        "cancel_task",
864    ];
865
866    #[derive(Clone, Default)]
867    struct CapturedLogs(Arc<std::sync::Mutex<Vec<u8>>>);
868
869    impl CapturedLogs {
870        fn contents(&self) -> String {
871            let bytes = self.0.lock().map(|guard| guard.clone()).unwrap_or_default();
872            String::from_utf8(bytes).unwrap_or_default()
873        }
874    }
875
876    struct CapturedLogsWriter(Arc<std::sync::Mutex<Vec<u8>>>);
877
878    impl io::Write for CapturedLogsWriter {
879        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
880            if let Ok(mut guard) = self.0.lock() {
881                guard.extend_from_slice(buf);
882            }
883            Ok(buf.len())
884        }
885
886        fn flush(&mut self) -> io::Result<()> {
887            Ok(())
888        }
889    }
890
891    impl<'a> tracing_subscriber::fmt::MakeWriter<'a> for CapturedLogs {
892        type Writer = CapturedLogsWriter;
893
894        fn make_writer(&'a self) -> Self::Writer {
895            CapturedLogsWriter(Arc::clone(&self.0))
896        }
897    }
898
899    /// Minimal in-process `ServerHandler` for tests.
900    #[derive(Clone, Default)]
901    struct TestHandler {
902        /// When Some, `call_tool` returns a body of this many 'x' bytes.
903        body_bytes: Option<usize>,
904    }
905
906    impl ServerHandler for TestHandler {
907        fn get_info(&self) -> ServerInfo {
908            ServerInfo::default()
909        }
910
911        #[allow(
912            clippy::unused_async_trait_impl,
913            reason = "async is mandated by the rmcp ServerHandler trait signature; this test handler does not await"
914        )]
915        async fn call_tool(
916            &self,
917            _request: CallToolRequestParams,
918            _context: RequestContext<RoleServer>,
919        ) -> Result<CallToolResponse, ErrorData> {
920            let body = "x".repeat(self.body_bytes.unwrap_or(4));
921            Ok(CallToolResult::success(vec![ContentBlock::text(body)]).into())
922        }
923    }
924
925    #[derive(Clone, Default)]
926    struct DelegationProbe {
927        seen: Arc<std::sync::Mutex<Vec<&'static str>>>,
928        notify: Arc<tokio::sync::Notify>,
929    }
930
931    impl DelegationProbe {
932        fn record(&self, method: &'static str) {
933            if let Ok(mut seen) = self.seen.lock() {
934                seen.push(method);
935            }
936            self.notify.notify_waiters();
937        }
938
939        fn seen(&self) -> Vec<&'static str> {
940            self.seen
941                .lock()
942                .map(|seen| seen.clone())
943                .unwrap_or_default()
944        }
945
946        async fn wait_for_seen_count(&self, count: usize) {
947            tokio::time::timeout(std::time::Duration::from_secs(1), async {
948                while self.seen().len() < count {
949                    self.notify.notified().await;
950                }
951            })
952            .await
953            .expect("delegated handler methods should be observed");
954        }
955    }
956
957    #[allow(
958        clippy::unused_async_trait_impl,
959        deprecated,
960        reason = "delegation tests cover rmcp async trait methods whose probe implementations return immediately"
961    )]
962    impl ServerHandler for DelegationProbe {
963        fn get_info(&self) -> ServerInfo {
964            ServerInfo::default()
965        }
966
967        async fn ping(&self, _context: RequestContext<RoleServer>) -> Result<(), ErrorData> {
968            self.record("ping");
969            Ok(())
970        }
971
972        async fn complete(
973            &self,
974            _request: CompleteRequestParams,
975            _context: RequestContext<RoleServer>,
976        ) -> Result<CompleteResult, ErrorData> {
977            self.record("complete");
978            let completion = CompletionInfo::with_all_values(vec!["delegated".to_owned()])
979                .expect("single completion is within rmcp max");
980            Ok(CompleteResult::new(completion))
981        }
982
983        async fn set_level(
984            &self,
985            _request: SetLevelRequestParams,
986            _context: RequestContext<RoleServer>,
987        ) -> Result<(), ErrorData> {
988            self.record("set_level");
989            Ok(())
990        }
991
992        async fn subscribe(
993            &self,
994            _request: SubscribeRequestParams,
995            _context: RequestContext<RoleServer>,
996        ) -> Result<(), ErrorData> {
997            self.record("subscribe");
998            Ok(())
999        }
1000
1001        async fn unsubscribe(
1002            &self,
1003            _request: UnsubscribeRequestParams,
1004            _context: RequestContext<RoleServer>,
1005        ) -> Result<(), ErrorData> {
1006            self.record("unsubscribe");
1007            Ok(())
1008        }
1009
1010        async fn call_tool(
1011            &self,
1012            _request: CallToolRequestParams,
1013            _context: RequestContext<RoleServer>,
1014        ) -> Result<CallToolResponse, ErrorData> {
1015            self.record("call_tool");
1016            Ok(CallToolResult::success(vec![ContentBlock::text("inner")]).into())
1017        }
1018
1019        async fn on_custom_request(
1020            &self,
1021            _request: CustomRequest,
1022            _context: RequestContext<RoleServer>,
1023        ) -> Result<CustomResult, ErrorData> {
1024            self.record("on_custom_request");
1025            Ok(CustomResult::new(json!({ "delegated": true })))
1026        }
1027
1028        async fn on_cancelled(
1029            &self,
1030            _notification: CancelledNotificationParam,
1031            _context: NotificationContext<RoleServer>,
1032        ) {
1033            self.record("on_cancelled");
1034        }
1035
1036        async fn on_progress(
1037            &self,
1038            _notification: ProgressNotificationParam,
1039            _context: NotificationContext<RoleServer>,
1040        ) {
1041            self.record("on_progress");
1042        }
1043
1044        async fn on_initialized(&self, _context: NotificationContext<RoleServer>) {
1045            self.record("on_initialized");
1046        }
1047
1048        async fn on_roots_list_changed(&self, _context: NotificationContext<RoleServer>) {
1049            self.record("on_roots_list_changed");
1050        }
1051
1052        async fn on_custom_notification(
1053            &self,
1054            _notification: CustomNotification,
1055            _context: NotificationContext<RoleServer>,
1056        ) {
1057            self.record("on_custom_notification");
1058        }
1059    }
1060
1061    fn delegation_transport(probe: DelegationProbe, hooks: Arc<ToolHooks>) -> DelegationTransport {
1062        let (client, server) = tokio::io::duplex(16 * 1024);
1063        let (client_read, client_write) = tokio::io::split(client);
1064        let service = rmcp::service::serve_directly::<RoleServer, _, _, io::Error, _>(
1065            with_hooks(probe.clone(), hooks),
1066            server,
1067            None,
1068        );
1069        (probe, BufReader::new(client_read), client_write, service)
1070    }
1071
1072    async fn send_json_rpc(
1073        writer: &mut tokio::io::WriteHalf<DuplexStream>,
1074        reader: &mut BufReader<tokio::io::ReadHalf<DuplexStream>>,
1075        request: serde_json::Value,
1076    ) -> serde_json::Value {
1077        writer
1078            .write_all(request.to_string().as_bytes())
1079            .await
1080            .expect("write request");
1081        writer.write_all(b"\n").await.expect("write newline");
1082        writer.flush().await.expect("flush request");
1083
1084        let mut line = String::new();
1085        reader.read_line(&mut line).await.expect("read response");
1086        serde_json::from_str(&line).expect("response is JSON")
1087    }
1088
1089    async fn send_notification(
1090        writer: &mut tokio::io::WriteHalf<DuplexStream>,
1091        notification: serde_json::Value,
1092    ) {
1093        writer
1094            .write_all(notification.to_string().as_bytes())
1095            .await
1096            .expect("write notification");
1097        writer.write_all(b"\n").await.expect("write newline");
1098        writer.flush().await.expect("flush notification");
1099    }
1100
1101    #[test]
1102    fn hooked_handler_delegation_count_is_maintenance_aid() {
1103        assert_eq!(
1104            HOOKED_HANDLER_DELEGATED_METHODS.len(),
1105            28,
1106            "maintenance aid only: update this list and the HookedHandler impl when rmcp adds ServerHandler methods"
1107        );
1108    }
1109
1110    #[tokio::test]
1111    async fn hooked_handler_delegates_ping() {
1112        let (probe, mut reader, mut writer, _service) =
1113            delegation_transport(DelegationProbe::default(), Arc::new(ToolHooks::new()));
1114
1115        let response = send_json_rpc(
1116            &mut writer,
1117            &mut reader,
1118            json!({ "jsonrpc": "2.0", "id": 1, "method": "ping" }),
1119        )
1120        .await;
1121
1122        assert_eq!(response["result"], json!({}));
1123        assert_eq!(probe.seen(), vec!["ping"]);
1124    }
1125
1126    #[tokio::test]
1127    async fn hooked_handler_delegates_notifications() {
1128        let (probe, _reader, mut writer, _service) =
1129            delegation_transport(DelegationProbe::default(), Arc::new(ToolHooks::new()));
1130
1131        send_notification(
1132            &mut writer,
1133            json!({
1134                "jsonrpc": "2.0",
1135                "method": "notifications/cancelled",
1136                "params": { "requestId": 1, "reason": "test" }
1137            }),
1138        )
1139        .await;
1140        send_notification(
1141            &mut writer,
1142            json!({
1143                "jsonrpc": "2.0",
1144                "method": "notifications/progress",
1145                "params": { "progressToken": 1, "progress": 0.5 }
1146            }),
1147        )
1148        .await;
1149        send_notification(
1150            &mut writer,
1151            json!({ "jsonrpc": "2.0", "method": "notifications/initialized" }),
1152        )
1153        .await;
1154        send_notification(
1155            &mut writer,
1156            json!({ "jsonrpc": "2.0", "method": "notifications/roots/list_changed" }),
1157        )
1158        .await;
1159        send_notification(
1160            &mut writer,
1161            json!({ "jsonrpc": "2.0", "method": "notifications/custom/probe" }),
1162        )
1163        .await;
1164
1165        probe.wait_for_seen_count(5).await;
1166        assert_eq!(
1167            probe.seen(),
1168            vec![
1169                "on_cancelled",
1170                "on_progress",
1171                "on_initialized",
1172                "on_roots_list_changed",
1173                "on_custom_notification"
1174            ]
1175        );
1176    }
1177
1178    #[tokio::test]
1179    #[allow(
1180        deprecated,
1181        reason = "set_level is deprecated by rmcp but must delegate"
1182    )]
1183    async fn hooked_handler_delegates_completion_and_level() {
1184        let (probe, mut reader, mut writer, _service) =
1185            delegation_transport(DelegationProbe::default(), Arc::new(ToolHooks::new()));
1186
1187        let completion = send_json_rpc(
1188            &mut writer,
1189            &mut reader,
1190            json!({
1191                "jsonrpc": "2.0",
1192                "id": 1,
1193                "method": "completion/complete",
1194                "params": {
1195                    "ref": { "type": "ref/prompt", "name": "prompt" },
1196                    "argument": { "name": "arg", "value": "de" }
1197                }
1198            }),
1199        )
1200        .await;
1201        let level = send_json_rpc(
1202            &mut writer,
1203            &mut reader,
1204            json!({
1205                "jsonrpc": "2.0",
1206                "id": 2,
1207                "method": "logging/setLevel",
1208                "params": { "level": "debug" }
1209            }),
1210        )
1211        .await;
1212
1213        assert_eq!(
1214            completion["result"]["completion"]["values"],
1215            json!(["delegated"])
1216        );
1217        assert_eq!(level["result"], json!({}));
1218        assert_eq!(probe.seen(), vec!["complete", "set_level"]);
1219    }
1220
1221    #[tokio::test]
1222    #[allow(
1223        deprecated,
1224        reason = "subscribe/unsubscribe are deprecated by rmcp but must delegate"
1225    )]
1226    async fn hooked_handler_delegates_subscriptions() {
1227        let (probe, mut reader, mut writer, _service) =
1228            delegation_transport(DelegationProbe::default(), Arc::new(ToolHooks::new()));
1229
1230        let subscribe = send_json_rpc(
1231            &mut writer,
1232            &mut reader,
1233            json!({
1234                "jsonrpc": "2.0",
1235                "id": 1,
1236                "method": "resources/subscribe",
1237                "params": { "uri": "file:///tmp/a" }
1238            }),
1239        )
1240        .await;
1241        let unsubscribe = send_json_rpc(
1242            &mut writer,
1243            &mut reader,
1244            json!({
1245                "jsonrpc": "2.0",
1246                "id": 2,
1247                "method": "resources/unsubscribe",
1248                "params": { "uri": "file:///tmp/a" }
1249            }),
1250        )
1251        .await;
1252
1253        assert_eq!(subscribe["result"], json!({}));
1254        assert_eq!(unsubscribe["result"], json!({}));
1255        assert_eq!(probe.seen(), vec!["subscribe", "unsubscribe"]);
1256    }
1257
1258    #[tokio::test]
1259    async fn hooked_handler_delegates_custom_request() {
1260        let (probe, mut reader, mut writer, _service) =
1261            delegation_transport(DelegationProbe::default(), Arc::new(ToolHooks::new()));
1262
1263        let response = send_json_rpc(
1264            &mut writer,
1265            &mut reader,
1266            json!({
1267                "jsonrpc": "2.0",
1268                "id": 1,
1269                "method": "requests/custom/probe",
1270                "params": { "x": true }
1271            }),
1272        )
1273        .await;
1274
1275        assert_eq!(response["result"], json!({ "delegated": true }));
1276        assert_eq!(probe.seen(), vec!["on_custom_request"]);
1277    }
1278
1279    #[tokio::test]
1280    async fn hooked_handler_still_applies_hooks_to_call_tool() {
1281        let before_count = Arc::new(AtomicUsize::new(0));
1282        let before_seen = Arc::clone(&before_count);
1283        let before: BeforeHook = Arc::new(move |_ctx| {
1284            let before_seen = Arc::clone(&before_seen);
1285            Box::pin(async move {
1286                before_seen.fetch_add(1, Ordering::Relaxed);
1287                HookOutcome::Continue
1288            })
1289        });
1290        let after_count = Arc::new(AtomicUsize::new(0));
1291        let after_seen = Arc::clone(&after_count);
1292        let after_notify = Arc::new(tokio::sync::Notify::new());
1293        let after_notify_seen = Arc::clone(&after_notify);
1294        let after: AfterHook = Arc::new(move |_ctx, _disp, _size| {
1295            let after_seen = Arc::clone(&after_seen);
1296            let after_notify_seen = Arc::clone(&after_notify_seen);
1297            Box::pin(async move {
1298                after_seen.fetch_add(1, Ordering::Relaxed);
1299                after_notify_seen.notify_waiters();
1300            })
1301        });
1302        let hooks = Arc::new(
1303            ToolHooks::new()
1304                .with_before(before)
1305                .with_after(after)
1306                .with_max_result_bytes(1024),
1307        );
1308        let (probe, mut reader, mut writer, _service) =
1309            delegation_transport(DelegationProbe::default(), hooks);
1310
1311        let response = send_json_rpc(
1312            &mut writer,
1313            &mut reader,
1314            json!({
1315                "jsonrpc": "2.0",
1316                "id": 1,
1317                "method": "tools/call",
1318                "params": { "name": "probe", "arguments": {} }
1319            }),
1320        )
1321        .await;
1322
1323        tokio::time::timeout(std::time::Duration::from_secs(1), async {
1324            while after_count.load(Ordering::Relaxed) == 0 {
1325                after_notify.notified().await;
1326            }
1327        })
1328        .await
1329        .expect("after hook should run");
1330        assert_eq!(response["result"]["content"][0]["text"], "inner");
1331        assert_eq!(probe.seen(), vec!["call_tool"]);
1332        assert_eq!(before_count.load(Ordering::Relaxed), 1);
1333        assert_eq!(after_count.load(Ordering::Relaxed), 1);
1334    }
1335
1336    fn ctx(name: &str) -> ToolCallContext {
1337        ToolCallContext {
1338            tool_name: name.to_owned(),
1339            arguments: None,
1340            identity: None,
1341            role: None,
1342            sub: None,
1343            request_id: None,
1344        }
1345    }
1346
1347    fn sensitive_ctx() -> ToolCallContext {
1348        ToolCallContext {
1349            tool_name: "safe-tool-name".to_owned(),
1350            arguments: Some(serde_json::json!({ "password": "argument-secret" })),
1351            identity: Some("identity-secret".to_owned()),
1352            role: Some("role-secret".to_owned()),
1353            sub: Some("sub-secret".to_owned()),
1354            request_id: Some("request-id-visible".to_owned()),
1355        }
1356    }
1357
1358    #[test]
1359    fn tool_call_context_debug_redacts_sensitive_fields_by_default() {
1360        let _guard = crate::diagnostics::ExposureTestGuard::acquire();
1361        crate::diagnostics::set_diagnostic_exposure(
1362            &crate::diagnostics::DiagnosticExposure::default(),
1363        );
1364
1365        let rendered = format!("{:?}", sensitive_ctx());
1366
1367        assert!(rendered.contains("safe-tool-name"));
1368        assert!(rendered.contains("request-id-visible"));
1369        assert!(rendered.contains("[REDACTED]"));
1370        for secret in [
1371            "argument-secret",
1372            "identity-secret",
1373            "role-secret",
1374            "sub-secret",
1375        ] {
1376            assert!(
1377                !rendered.contains(secret),
1378                "ToolCallContext Debug must not contain {secret}: {rendered}"
1379            );
1380        }
1381    }
1382
1383    #[test]
1384    fn tool_call_context_debug_can_show_sensitive_fields_when_enabled() {
1385        let _guard = crate::diagnostics::ExposureTestGuard::acquire();
1386        crate::diagnostics::set_diagnostic_exposure(&crate::diagnostics::DiagnosticExposure {
1387            tool_call_arguments: true,
1388            ..crate::diagnostics::DiagnosticExposure::default()
1389        });
1390
1391        let rendered = format!("{:?}", sensitive_ctx());
1392
1393        for secret in [
1394            "argument-secret",
1395            "identity-secret",
1396            "role-secret",
1397            "sub-secret",
1398        ] {
1399            assert!(
1400                rendered.contains(secret),
1401                "ToolCallContext Debug must contain {secret} when enabled: {rendered}"
1402            );
1403        }
1404    }
1405
1406    #[tokio::test]
1407    async fn size_cap_replaces_oversized_result() {
1408        let inner = TestHandler {
1409            body_bytes: Some(8_192),
1410        };
1411        let hooks = Arc::new(ToolHooks {
1412            max_result_bytes: Some(256),
1413            before: None,
1414            after: None,
1415        });
1416        let hooked = with_hooks(inner, hooks);
1417
1418        let small = CallToolResult::success(vec![ContentBlock::text("ok".to_owned())]);
1419        assert!(exact_size(&small) < 256);
1420
1421        let big = CallToolResult::success(vec![ContentBlock::text("x".repeat(8_192))]);
1422        let size = exact_size(&big);
1423        assert!(size > 256);
1424
1425        let (replaced, accounted, capped) = apply_size_cap(big, Some(256), "whatever");
1426        assert!(capped);
1427        assert_eq!(accounted, 257);
1428        assert_eq!(replaced.is_error, Some(true));
1429        assert!(matches!(
1430            replaced.content.first(),
1431            Some(rmcp::model::ContentBlock::Text(t)) if t.text.contains("result_too_large")
1432        ));
1433
1434        // Compile-check that HookedHandler instantiates with the test inner.
1435        let _ = hooked;
1436    }
1437
1438    fn exact_size(result: &CallToolResult) -> usize {
1439        match serialized_size(result, None) {
1440            SizeMeasure::Exact(size) => size,
1441            SizeMeasure::Exceeded { limit } => {
1442                panic!("unbounded measurement exceeded impossible limit {limit}");
1443            }
1444        }
1445    }
1446
1447    #[test]
1448    fn serialized_size_under_cap_is_exact() {
1449        let result = CallToolResult::success(vec![ContentBlock::text("ok".to_owned())]);
1450        let exact = serde_json::to_vec(&result).unwrap().len();
1451
1452        let measured = serialized_size(&result, Some(exact));
1453
1454        assert_eq!(measured, SizeMeasure::Exact(exact));
1455    }
1456
1457    #[test]
1458    fn serialized_size_over_cap_stops_with_exceeded() {
1459        let result = CallToolResult::success(vec![ContentBlock::text("x".repeat(8_192))]);
1460
1461        let measured = serialized_size(&result, Some(256));
1462
1463        assert_eq!(measured, SizeMeasure::Exceeded { limit: 256 });
1464    }
1465
1466    #[test]
1467    fn over_cap_replacement_does_not_log_serialization_failure() {
1468        let logs = CapturedLogs::default();
1469        let subscriber = tracing_subscriber::fmt()
1470            .with_max_level(tracing::Level::TRACE)
1471            .with_writer(logs.clone())
1472            .with_ansi(false)
1473            .without_time()
1474            .finish();
1475        let _guard = tracing::subscriber::set_default(subscriber);
1476        let result = CallToolResult::success(vec![ContentBlock::text("x".repeat(8_192))]);
1477
1478        let (_final_result, accounted, capped) = apply_size_cap(result, Some(256), "big_tool");
1479
1480        assert!(capped);
1481        assert_eq!(accounted, 257);
1482        assert!(
1483            logs.contents()
1484                .contains("tool result exceeds max_result_bytes")
1485        );
1486        assert!(
1487            !logs.contents().contains("failed to serialize"),
1488            "cap-abort must not be logged as serialization failure: {}",
1489            logs.contents()
1490        );
1491    }
1492
1493    #[test]
1494    fn disabled_result_cap_skips_measurement() {
1495        let result = CallToolResult::success(vec![ContentBlock::text("x".repeat(8_192))]);
1496
1497        let (_final_result, accounted, capped) = apply_size_cap(result, None, "uncapped_tool");
1498
1499        assert!(!capped);
1500        assert_eq!(accounted, 0);
1501    }
1502
1503    #[tokio::test]
1504    async fn before_hook_deny_builds_error() {
1505        let counter = Arc::new(AtomicUsize::new(0));
1506        let c = Arc::clone(&counter);
1507        let before: BeforeHook = Arc::new(move |ctx_ref| {
1508            let c = Arc::clone(&c);
1509            let name = ctx_ref.tool_name.clone();
1510            Box::pin(async move {
1511                c.fetch_add(1, Ordering::Relaxed);
1512                if name == "forbidden" {
1513                    HookOutcome::Deny(ErrorData::invalid_request("nope", None))
1514                } else {
1515                    HookOutcome::Continue
1516                }
1517            })
1518        });
1519
1520        let hooks = Arc::new(ToolHooks {
1521            max_result_bytes: None,
1522            before: Some(before),
1523            after: None,
1524        });
1525        let hooked = with_hooks(TestHandler::default(), hooks);
1526
1527        let bad_ctx = ctx("forbidden");
1528        let before_fn = hooked.hooks.before.as_ref().unwrap();
1529        let outcome = before_fn(&bad_ctx).await;
1530        assert!(matches!(outcome, HookOutcome::Deny(_)));
1531        assert_eq!(counter.load(Ordering::Relaxed), 1);
1532
1533        let ok_ctx = ctx("allowed");
1534        let outcome2 = before_fn(&ok_ctx).await;
1535        assert!(matches!(outcome2, HookOutcome::Continue));
1536        assert_eq!(counter.load(Ordering::Relaxed), 2);
1537    }
1538
1539    #[test]
1540    fn too_large_result_mentions_limit_and_actual() {
1541        let r = too_large_result(100, Some(500), "my_tool");
1542        let body = serde_json::to_string(&r).unwrap();
1543        assert!(body.contains("result_too_large"));
1544        assert!(body.contains("my_tool"));
1545        assert!(body.contains("100"));
1546        assert!(body.contains("500"));
1547    }
1548
1549    #[test]
1550    fn decide_size_truth_table() {
1551        assert_eq!(
1552            decide_size(Some(SizeMeasure::Exact(10)), Some(100)),
1553            SizeVerdict::Pass { size: 10 }
1554        );
1555        assert_eq!(
1556            decide_size(Some(SizeMeasure::Exact(100)), Some(100)),
1557            SizeVerdict::Pass { size: 100 },
1558            "cap is inclusive: size == limit passes"
1559        );
1560        assert_eq!(
1561            decide_size(Some(SizeMeasure::Exact(101)), Some(100)),
1562            SizeVerdict::Replace {
1563                limit: 100,
1564                actual: Some(101)
1565            }
1566        );
1567        assert_eq!(
1568            decide_size(Some(SizeMeasure::Exact(999)), None),
1569            SizeVerdict::Pass { size: 999 }
1570        );
1571        assert_eq!(
1572            decide_size(None, Some(100)),
1573            SizeVerdict::Replace {
1574                limit: 100,
1575                actual: None
1576            },
1577            "unmeasurable result must fail closed when a cap is configured"
1578        );
1579        assert_eq!(decide_size(None, None), SizeVerdict::PassUnmeasured);
1580        assert_eq!(
1581            decide_size(Some(SizeMeasure::Exceeded { limit: 100 }), Some(100)),
1582            SizeVerdict::Replace {
1583                limit: 100,
1584                actual: None
1585            },
1586            "cap-abort is not an exact measurement"
1587        );
1588    }
1589
1590    #[test]
1591    fn too_large_result_does_not_fabricate_a_size_when_unmeasurable() {
1592        let r = too_large_result(100, None, "my_tool");
1593        let body = serde_json::to_string(&r).unwrap();
1594        assert!(body.contains("result_too_large"));
1595        assert!(body.contains("unknown"));
1596        assert!(
1597            !body.contains("101"),
1598            "the over-limit accounting sentinel must not leak into the client payload"
1599        );
1600    }
1601
1602    #[tokio::test]
1603    async fn replace_outcome_skips_inner_and_returns_payload() {
1604        // Returning Replace from before-hook must yield the supplied
1605        // CallToolResult directly, with no need for the inner handler.
1606        let before: BeforeHook = Arc::new(|_ctx| {
1607            Box::pin(async {
1608                HookOutcome::Replace(Box::new(CallToolResult::success(vec![ContentBlock::text(
1609                    "from-replace".to_owned(),
1610                )])))
1611            })
1612        });
1613        let hooks = Arc::new(ToolHooks {
1614            max_result_bytes: None,
1615            before: Some(before),
1616            after: None,
1617        });
1618        let _hooked = with_hooks(TestHandler::default(), Arc::clone(&hooks));
1619
1620        // Exercise the before-hook closure + apply_size_cap helper directly,
1621        // matching the established test pattern in this module.
1622        let outcome = (hooks.before.as_ref().unwrap())(&ctx("any")).await;
1623        let HookOutcome::Replace(boxed) = outcome else {
1624            panic!("expected HookOutcome::Replace");
1625        };
1626        let (result, size, capped) = apply_size_cap(*boxed, None, "any");
1627        assert!(!capped);
1628        assert_eq!(size, 0);
1629        assert!(!result.is_error.unwrap_or(false));
1630        assert!(matches!(
1631            result.content.first(),
1632            Some(rmcp::model::ContentBlock::Text(t)) if t.text == "from-replace"
1633        ));
1634    }
1635
1636    #[tokio::test]
1637    async fn replace_outcome_subject_to_size_cap() {
1638        // A Replace payload that exceeds max_result_bytes must be rewritten
1639        // to result_too_large just like an inner-handler result would be,
1640        // and the disposition must reflect ResultTooLarge.
1641        let huge = CallToolResult::success(vec![ContentBlock::text("y".repeat(8_192))]);
1642        let huge_size = serde_json::to_vec(&huge).unwrap().len();
1643        assert!(huge_size > 256);
1644
1645        let (final_result, accounted, capped) = apply_size_cap(huge, Some(256), "replaced_tool");
1646        assert!(capped);
1647        assert_eq!(accounted, 257);
1648        assert_eq!(final_result.is_error, Some(true));
1649        assert!(matches!(
1650            final_result.content.first(),
1651            Some(rmcp::model::ContentBlock::Text(t)) if t.text.contains("result_too_large")
1652        ));
1653    }
1654
1655    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1656    async fn after_hook_fires_exactly_once_via_spawn() {
1657        // spawn_after must enqueue the after-hook exactly one time per
1658        // invocation and never block the caller; we wait for the spawned
1659        // task to run by polling the counter with a short timeout.
1660        let counter = Arc::new(AtomicUsize::new(0));
1661        let c = Arc::clone(&counter);
1662        let after: AfterHook = Arc::new(move |_ctx, _disp, _size| {
1663            let c = Arc::clone(&c);
1664            Box::pin(async move {
1665                c.fetch_add(1, Ordering::Relaxed);
1666            })
1667        });
1668        let holder = Arc::new(AfterHookHolder { f: after });
1669
1670        HookedHandler::<TestHandler>::spawn_after(
1671            Some(&holder),
1672            ctx("t"),
1673            HookDisposition::InnerExecuted,
1674            42,
1675        );
1676
1677        // Wait up to 1s for the spawned task to run.
1678        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1);
1679        while counter.load(Ordering::Relaxed) == 0 && std::time::Instant::now() < deadline {
1680            tokio::task::yield_now().await;
1681            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
1682        }
1683        assert_eq!(counter.load(Ordering::Relaxed), 1);
1684    }
1685
1686    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1687    async fn after_hook_panic_is_isolated_from_response_path() {
1688        // A panicking after-hook must not affect the request task.  We
1689        // spawn a panicking after-hook and then verify the current task
1690        // can still complete an unrelated future to completion.
1691        let after: AfterHook = Arc::new(|_ctx, _disp, _size| {
1692            Box::pin(async {
1693                panic!("intentional panic in after-hook");
1694            })
1695        });
1696        let holder = Arc::new(AfterHookHolder { f: after });
1697
1698        HookedHandler::<TestHandler>::spawn_after(
1699            Some(&holder),
1700            ctx("boom"),
1701            HookDisposition::InnerExecuted,
1702            0,
1703        );
1704
1705        // Give Tokio a chance to run + abort the panicking task, then
1706        // confirm we're still alive and the runtime is healthy.
1707        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1708        let still_alive = tokio::spawn(async { 1_u32 + 2 }).await.unwrap();
1709        assert_eq!(still_alive, 3);
1710    }
1711}
rmcp_server_kit/tool_hooks.rs

rmcp_server_kit/
tool_hooks.rs
