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, pin::Pin, sync::Arc};
40
41use rmcp::{
42    ErrorData, RoleServer, ServerHandler,
43    model::{
44        CallToolRequestParams, CallToolResponse, CallToolResult, CancelTaskParams, ContentBlock,
45        DiscoverResult, GetPromptRequestParams, GetPromptResponse, GetTaskParams, GetTaskResult,
46        InitializeRequestParams, InitializeResult, ListPromptsResult, ListResourceTemplatesResult,
47        ListResourcesResult, ListToolsResult, PaginatedRequestParams, ProtocolVersion,
48        ReadResourceRequestParams, ReadResourceResponse, ServerInfo, SubscriptionFilter, Tool,
49        UpdateTaskParams,
50    },
51    service::{RequestContext, SubscriptionContext},
52};
53
54/// Context passed to before/after hooks for a single tool call.
55#[derive(Debug, Clone)]
56#[non_exhaustive]
57pub struct ToolCallContext {
58    /// Tool name being invoked.
59    pub tool_name: String,
60    /// JSON arguments as sent by the client (may be `None`).
61    pub arguments: Option<serde_json::Value>,
62    /// Identity name from the authenticated request, if any.
63    pub identity: Option<String>,
64    /// RBAC role associated with the request, if any.
65    pub role: Option<String>,
66    /// OAuth `sub` claim, if present.
67    pub sub: Option<String>,
68    /// Raw JSON-RPC request id rendered as a string, if available.
69    pub request_id: Option<String>,
70}
71
72impl ToolCallContext {
73    /// Construct a [`ToolCallContext`] with the given tool name and all
74    /// optional fields cleared.  Primarily for use in unit tests and
75    /// benchmarks of user-supplied hooks; the runtime path populates
76    /// these fields from the request and task-local RBAC state.
77    #[must_use]
78    pub fn for_tool(tool_name: impl Into<String>) -> Self {
79        Self {
80            tool_name: tool_name.into(),
81            arguments: None,
82            identity: None,
83            role: None,
84            sub: None,
85            request_id: None,
86        }
87    }
88}
89
90/// Outcome returned by a [`BeforeHook`] to control invocation flow.
91///
92/// - [`HookOutcome::Continue`] - proceed with the wrapped handler.
93/// - [`HookOutcome::Deny`] - reject the call with the supplied
94///   [`ErrorData`]; the inner handler is **not** called.
95/// - [`HookOutcome::Replace`] - return the supplied result instead of
96///   invoking the inner handler.  The result is still subject to
97///   `max_result_bytes` capping.
98#[derive(Debug)]
99#[non_exhaustive]
100pub enum HookOutcome {
101    /// Proceed with the wrapped handler.
102    Continue,
103    /// Reject the call.  The error is propagated to the client as-is.
104    Deny(ErrorData),
105    /// Skip the inner handler and return the supplied result instead.
106    Replace(Box<CallToolResult>),
107}
108
109/// How a tool call resolved, passed to the [`AfterHook`].
110#[derive(Debug, Clone, Copy)]
111#[non_exhaustive]
112pub enum HookDisposition {
113    /// The inner handler ran and returned `Ok`.
114    InnerExecuted,
115    /// The inner handler ran and returned `Err`.
116    InnerErrored,
117    /// The before-hook returned [`HookOutcome::Deny`].
118    DeniedBefore,
119    /// The before-hook returned [`HookOutcome::Replace`].
120    ReplacedBefore,
121    /// The result (from inner or replace) exceeded `max_result_bytes`
122    /// and was substituted with a structured error.
123    ResultTooLarge,
124}
125
126/// Async before-hook callback type.
127///
128/// Returns a [`HookOutcome`] controlling whether the inner handler runs.
129/// The borrow of `ToolCallContext` is held for the duration of the
130/// returned future, which avoids forcing implementations to clone the
131/// context for every invocation.
132pub type BeforeHook = Arc<
133    dyn for<'a> Fn(&'a ToolCallContext) -> Pin<Box<dyn Future<Output = HookOutcome> + Send + 'a>>
134        + Send
135        + Sync
136        + 'static,
137>;
138
139/// Async after-hook callback type.
140///
141/// Receives the call context, a [`HookDisposition`] describing how the
142/// call resolved, and the approximate serialized result size in bytes
143/// (`0` for `DeniedBefore` and `InnerErrored`).  Spawned via
144/// `tokio::spawn`, so it must not assume it runs before the response is
145/// flushed.
146pub type AfterHook = Arc<
147    dyn for<'a> Fn(
148            &'a ToolCallContext,
149            HookDisposition,
150            usize,
151        ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>>
152        + Send
153        + Sync
154        + 'static,
155>;
156
157/// Opt-in hooks applied by [`crate::tool_hooks::HookedHandler`].
158#[allow(clippy::struct_field_names, reason = "before/after read naturally")]
159#[derive(Clone, Default)]
160#[non_exhaustive]
161pub struct ToolHooks {
162    /// Hard cap on serialized `CallToolResult` size in bytes.  When
163    /// exceeded, the result is replaced with an `is_error=true` result
164    /// carrying a `result_too_large` structured error.  `None` disables
165    /// the cap.
166    pub max_result_bytes: Option<usize>,
167    /// Optional before-hook invoked after arg deserialization, before
168    /// the wrapped handler is called.
169    pub before: Option<BeforeHook>,
170    /// Optional after-hook invoked once per call, regardless of how the
171    /// call resolved.  Spawned via `tokio::spawn` and never blocks the
172    /// response path.
173    pub after: Option<AfterHook>,
174}
175
176impl ToolHooks {
177    /// Construct an empty [`ToolHooks`] with no cap and no hooks.
178    ///
179    /// Use the `with_*` builder methods to populate fields; this avoids
180    /// the `#[non_exhaustive]` restriction that prevents struct-literal
181    /// construction from outside the crate.
182    #[must_use]
183    pub fn new() -> Self {
184        Self::default()
185    }
186
187    /// Set the serialized result size cap in bytes.
188    #[must_use]
189    pub fn with_max_result_bytes(mut self, max: usize) -> Self {
190        self.max_result_bytes = Some(max);
191        self
192    }
193
194    /// Set the before-hook.
195    #[must_use]
196    pub fn with_before(mut self, before: BeforeHook) -> Self {
197        self.before = Some(before);
198        self
199    }
200
201    /// Set the after-hook.
202    #[must_use]
203    pub fn with_after(mut self, after: AfterHook) -> Self {
204        self.after = Some(after);
205        self
206    }
207}
208
209impl fmt::Debug for ToolHooks {
210    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
211        f.debug_struct("ToolHooks")
212            .field("max_result_bytes", &self.max_result_bytes)
213            .field("before", &self.before.as_ref().map(|_| "<fn>"))
214            .field("after", &self.after.as_ref().map(|_| "<fn>"))
215            .finish()
216    }
217}
218
219/// `ServerHandler` wrapper that applies [`ToolHooks`].
220#[derive(Clone)]
221pub struct HookedHandler<H: ServerHandler> {
222    inner: Arc<H>,
223    hooks: Arc<ToolHooks>,
224}
225
226impl<H: ServerHandler> fmt::Debug for HookedHandler<H> {
227    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
228        f.debug_struct("HookedHandler")
229            .field("hooks", &self.hooks)
230            .finish_non_exhaustive()
231    }
232}
233
234/// Construct a [`crate::tool_hooks::HookedHandler`] from an inner handler and hooks.
235///
236/// Returning the wrapped handler is the entire point of this function;
237/// dropping it on the floor would silently disable the supplied hooks. The
238/// `#[must_use]` attribute would be the natural enforcement here, but adding
239/// it to a public function is a SemVer-minor change per cargo-semver-checks;
240/// it is deferred to the next minor-version bump.
241///
242// NOTE(next-minor): add `#[must_use = "HookedHandler must be wired into a
243// ServerHandler (e.g. via `serve(..., || hooked)`) to take effect; dropping
244// the returned value silently disables the supplied hooks"]` here. Deferred
245// from patch releases on 1.7.x to avoid downstream `-D warnings` churn.
246// Cross-ref: CHANGELOG.md "[Unreleased]" — Quality / lint hygiene section.
247pub fn with_hooks<H: ServerHandler>(inner: H, hooks: Arc<ToolHooks>) -> HookedHandler<H> {
248    HookedHandler {
249        inner: Arc::new(inner),
250        hooks,
251    }
252}
253
254impl<H: ServerHandler> HookedHandler<H> {
255    /// Access the wrapped handler.
256    #[must_use]
257    pub fn inner(&self) -> &H {
258        &self.inner
259    }
260
261    fn build_context(request: &CallToolRequestParams, req_id: Option<String>) -> ToolCallContext {
262        ToolCallContext {
263            tool_name: request.name.to_string(),
264            arguments: request.arguments.clone().map(serde_json::Value::Object),
265            identity: crate::rbac::current_identity(),
266            role: crate::rbac::current_role(),
267            sub: crate::rbac::current_sub(),
268            request_id: req_id,
269        }
270    }
271
272    /// Spawn the after-hook on the current Tokio runtime.  The future
273    /// captures clones of `ctx` and the `Arc<AfterHook>` so it can run
274    /// independently of the request task; panics inside the after-hook
275    /// are caught by Tokio and never poison the response path.
276    ///
277    /// The spawned task is **instrumented** with the request span via
278    /// [`tracing::Instrument`] and re-establishes the per-request RBAC
279    /// task-locals (role, identity, token, sub) via
280    /// [`crate::rbac::with_rbac_scope`]. Without this, after-hooks lose
281    /// their parent span (breaking trace correlation) and observe
282    /// `current_role()` / `current_identity()` as `None`.
283    fn spawn_after(
284        after: Option<&Arc<AfterHookHolder>>,
285        ctx: ToolCallContext,
286        disposition: HookDisposition,
287        size: usize,
288    ) {
289        if let Some(after) = after {
290            use tracing::Instrument;
291
292            let after = Arc::clone(after);
293            // Capture the request span before leaving the request task so
294            // after-hook log lines are correlated with the originating call.
295            let span = tracing::Span::current();
296            // Snapshot RBAC task-locals; defaults are empty strings so the
297            // re-established scope is a no-op when the request had no
298            // authenticated identity (e.g. health checks, anonymous tools).
299            let role = crate::rbac::current_role().unwrap_or_default();
300            let identity = crate::rbac::current_identity().unwrap_or_default();
301            let token = crate::rbac::current_token()
302                .unwrap_or_else(|| secrecy::SecretString::from(String::new()));
303            let sub = crate::rbac::current_sub().unwrap_or_default();
304            tokio::spawn(
305                async move {
306                    crate::rbac::with_rbac_scope(role, identity, token, sub, async move {
307                        let fut = (after.f)(&ctx, disposition, size);
308                        fut.await;
309                    })
310                    .await;
311                }
312                .instrument(span),
313            );
314        }
315    }
316}
317
318/// Internal newtype that owns the [`AfterHook`] so we can `Arc::clone`
319/// the *holder* and let the spawned task borrow `ctx` for the lifetime
320/// of the future without lifetime acrobatics in `tokio::spawn`.
321struct AfterHookHolder {
322    f: AfterHook,
323}
324
325/// Structured error body returned when a result exceeds `max_result_bytes`.
326fn too_large_result(limit: usize, actual: usize, tool: &str) -> CallToolResult {
327    let body = serde_json::json!({
328        "error": "result_too_large",
329        "message": format!(
330            "tool '{tool}' result of {actual} bytes exceeds the configured \
331             max_result_bytes={limit}; ask for a narrower query"
332        ),
333        "limit_bytes": limit,
334        "actual_bytes": actual,
335    });
336    let mut r = CallToolResult::error(vec![ContentBlock::text(body.to_string())]);
337    r.structured_content = None;
338    r
339}
340
341fn serialized_size(result: &CallToolResult) -> usize {
342    serde_json::to_vec(result).map_or(0, |v| v.len())
343}
344
345/// Apply the `max_result_bytes` cap to a result.  Returns the (possibly
346/// replaced) result, the size used for accounting, and whether the cap
347/// fired.
348fn apply_size_cap(
349    result: CallToolResult,
350    max: Option<usize>,
351    tool: &str,
352) -> (CallToolResult, usize, bool) {
353    let size = serialized_size(&result);
354    if let Some(limit) = max
355        && size > limit
356    {
357        tracing::warn!(
358            tool = %tool,
359            size_bytes = size,
360            limit_bytes = limit,
361            "tool result exceeds max_result_bytes; replacing with structured error"
362        );
363        let replaced = too_large_result(limit, size, tool);
364        return (replaced, size, true);
365    }
366    (result, size, false)
367}
368
369impl<H: ServerHandler> ServerHandler for HookedHandler<H> {
370    fn get_info(&self) -> ServerInfo {
371        self.inner.get_info()
372    }
373
374    async fn initialize(
375        &self,
376        request: InitializeRequestParams,
377        context: RequestContext<RoleServer>,
378    ) -> Result<InitializeResult, ErrorData> {
379        self.inner.initialize(request, context).await
380    }
381
382    async fn list_tools(
383        &self,
384        request: Option<PaginatedRequestParams>,
385        context: RequestContext<RoleServer>,
386    ) -> Result<ListToolsResult, ErrorData> {
387        self.inner.list_tools(request, context).await
388    }
389
390    fn get_tool(&self, name: &str) -> Option<Tool> {
391        self.inner.get_tool(name)
392    }
393
394    async fn list_prompts(
395        &self,
396        request: Option<PaginatedRequestParams>,
397        context: RequestContext<RoleServer>,
398    ) -> Result<ListPromptsResult, ErrorData> {
399        self.inner.list_prompts(request, context).await
400    }
401
402    async fn get_prompt(
403        &self,
404        request: GetPromptRequestParams,
405        context: RequestContext<RoleServer>,
406    ) -> Result<GetPromptResponse, ErrorData> {
407        self.inner.get_prompt(request, context).await
408    }
409
410    async fn list_resources(
411        &self,
412        request: Option<PaginatedRequestParams>,
413        context: RequestContext<RoleServer>,
414    ) -> Result<ListResourcesResult, ErrorData> {
415        self.inner.list_resources(request, context).await
416    }
417
418    async fn list_resource_templates(
419        &self,
420        request: Option<PaginatedRequestParams>,
421        context: RequestContext<RoleServer>,
422    ) -> Result<ListResourceTemplatesResult, ErrorData> {
423        self.inner.list_resource_templates(request, context).await
424    }
425
426    async fn read_resource(
427        &self,
428        request: ReadResourceRequestParams,
429        context: RequestContext<RoleServer>,
430    ) -> Result<ReadResourceResponse, ErrorData> {
431        self.inner.read_resource(request, context).await
432    }
433
434    #[allow(
435        clippy::wildcard_enum_match_arm,
436        reason = "CallToolResponse is #[non_exhaustive]; the non-Complete MRTR variants (InputRequired/Task) are passed through unchanged"
437    )]
438    async fn call_tool(
439        &self,
440        request: CallToolRequestParams,
441        context: RequestContext<RoleServer>,
442    ) -> Result<CallToolResponse, ErrorData> {
443        let req_id = Some(format!("{:?}", context.id));
444        let ctx = Self::build_context(&request, req_id);
445        let max = self.hooks.max_result_bytes;
446        let after_holder = self
447            .hooks
448            .after
449            .as_ref()
450            .map(|f| Arc::new(AfterHookHolder { f: Arc::clone(f) }));
451
452        // Before hook: may Continue, Deny, or Replace.
453        if let Some(before) = self.hooks.before.as_ref() {
454            let outcome = before(&ctx).await;
455            match outcome {
456                HookOutcome::Continue => {}
457                HookOutcome::Deny(err) => {
458                    Self::spawn_after(after_holder.as_ref(), ctx, HookDisposition::DeniedBefore, 0);
459                    return Err(err);
460                }
461                HookOutcome::Replace(boxed) => {
462                    let (final_result, size, capped) = apply_size_cap(*boxed, max, &ctx.tool_name);
463                    let disposition = if capped {
464                        HookDisposition::ResultTooLarge
465                    } else {
466                        HookDisposition::ReplacedBefore
467                    };
468                    Self::spawn_after(after_holder.as_ref(), ctx, disposition, size);
469                    return Ok(final_result.into());
470                }
471            }
472        }
473
474        // Inner handler.
475        match self.inner.call_tool(request, context).await {
476            // Completed tool result: subject to the size cap + after hook.
477            Ok(CallToolResponse::Complete(result)) => {
478                let (final_result, size, capped) = apply_size_cap(result, max, &ctx.tool_name);
479                let disposition = if capped {
480                    HookDisposition::ResultTooLarge
481                } else {
482                    HookDisposition::InnerExecuted
483                };
484                Self::spawn_after(after_holder.as_ref(), ctx, disposition, size);
485                Ok(final_result.into())
486            }
487            // MRTR input-required / task responses (rmcp 3.0): no CallToolResult
488            // to size-cap, so pass them through unchanged.
489            Ok(other) => {
490                Self::spawn_after(
491                    after_holder.as_ref(),
492                    ctx,
493                    HookDisposition::InnerExecuted,
494                    0,
495                );
496                Ok(other)
497            }
498            Err(e) => {
499                Self::spawn_after(after_holder.as_ref(), ctx, HookDisposition::InnerErrored, 0);
500                Err(e)
501            }
502        }
503    }
504
505    // rmcp 3.0 added task/subscription/discovery request handlers with defaults;
506    // delegate them to `inner` so wrapping a handler that implements those stays
507    // transparent (otherwise the default would shadow the inner implementation).
508    fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
509        self.inner.supported_protocol_versions()
510    }
511
512    async fn discover(
513        &self,
514        context: RequestContext<RoleServer>,
515    ) -> Result<DiscoverResult, ErrorData> {
516        self.inner.discover(context).await
517    }
518
519    fn accepted_subscription_filter(
520        &self,
521        requested: &SubscriptionFilter,
522    ) -> Option<SubscriptionFilter> {
523        self.inner.accepted_subscription_filter(requested)
524    }
525
526    async fn listen(&self, context: SubscriptionContext) -> Result<(), ErrorData> {
527        self.inner.listen(context).await
528    }
529
530    async fn get_task(
531        &self,
532        request: GetTaskParams,
533        context: RequestContext<RoleServer>,
534    ) -> Result<GetTaskResult, ErrorData> {
535        self.inner.get_task(request, context).await
536    }
537
538    async fn update_task(
539        &self,
540        request: UpdateTaskParams,
541        context: RequestContext<RoleServer>,
542    ) -> Result<(), ErrorData> {
543        self.inner.update_task(request, context).await
544    }
545
546    async fn cancel_task(
547        &self,
548        request: CancelTaskParams,
549        context: RequestContext<RoleServer>,
550    ) -> Result<(), ErrorData> {
551        self.inner.cancel_task(request, context).await
552    }
553}
554
555#[cfg(test)]
556mod tests {
557    use std::sync::{
558        Arc,
559        atomic::{AtomicUsize, Ordering},
560    };
561
562    use rmcp::{
563        ErrorData, RoleServer, ServerHandler,
564        model::{
565            CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, ServerInfo,
566        },
567        service::RequestContext,
568    };
569
570    use super::*;
571
572    /// Minimal in-process `ServerHandler` for tests.
573    #[derive(Clone, Default)]
574    struct TestHandler {
575        /// When Some, `call_tool` returns a body of this many 'x' bytes.
576        body_bytes: Option<usize>,
577    }
578
579    impl ServerHandler for TestHandler {
580        fn get_info(&self) -> ServerInfo {
581            ServerInfo::default()
582        }
583
584        async fn call_tool(
585            &self,
586            _request: CallToolRequestParams,
587            _context: RequestContext<RoleServer>,
588        ) -> Result<CallToolResponse, ErrorData> {
589            let body = "x".repeat(self.body_bytes.unwrap_or(4));
590            Ok(CallToolResult::success(vec![ContentBlock::text(body)]).into())
591        }
592    }
593
594    fn ctx(name: &str) -> ToolCallContext {
595        ToolCallContext {
596            tool_name: name.to_owned(),
597            arguments: None,
598            identity: None,
599            role: None,
600            sub: None,
601            request_id: None,
602        }
603    }
604
605    #[tokio::test]
606    async fn size_cap_replaces_oversized_result() {
607        let inner = TestHandler {
608            body_bytes: Some(8_192),
609        };
610        let hooks = Arc::new(ToolHooks {
611            max_result_bytes: Some(256),
612            before: None,
613            after: None,
614        });
615        let hooked = with_hooks(inner, hooks);
616
617        let small = CallToolResult::success(vec![ContentBlock::text("ok".to_owned())]);
618        assert!(serialized_size(&small) < 256);
619
620        let big = CallToolResult::success(vec![ContentBlock::text("x".repeat(8_192))]);
621        let size = serialized_size(&big);
622        assert!(size > 256);
623
624        let (replaced, accounted, capped) = apply_size_cap(big, Some(256), "whatever");
625        assert!(capped);
626        assert_eq!(accounted, size);
627        assert_eq!(replaced.is_error, Some(true));
628        assert!(matches!(
629            replaced.content.first(),
630            Some(rmcp::model::ContentBlock::Text(t)) if t.text.contains("result_too_large")
631        ));
632
633        // Compile-check that HookedHandler instantiates with the test inner.
634        let _ = hooked;
635    }
636
637    #[tokio::test]
638    async fn before_hook_deny_builds_error() {
639        let counter = Arc::new(AtomicUsize::new(0));
640        let c = Arc::clone(&counter);
641        let before: BeforeHook = Arc::new(move |ctx_ref| {
642            let c = Arc::clone(&c);
643            let name = ctx_ref.tool_name.clone();
644            Box::pin(async move {
645                c.fetch_add(1, Ordering::Relaxed);
646                if name == "forbidden" {
647                    HookOutcome::Deny(ErrorData::invalid_request("nope", None))
648                } else {
649                    HookOutcome::Continue
650                }
651            })
652        });
653
654        let hooks = Arc::new(ToolHooks {
655            max_result_bytes: None,
656            before: Some(before),
657            after: None,
658        });
659        let hooked = with_hooks(TestHandler::default(), hooks);
660
661        let bad_ctx = ctx("forbidden");
662        let before_fn = hooked.hooks.before.as_ref().unwrap();
663        let outcome = before_fn(&bad_ctx).await;
664        assert!(matches!(outcome, HookOutcome::Deny(_)));
665        assert_eq!(counter.load(Ordering::Relaxed), 1);
666
667        let ok_ctx = ctx("allowed");
668        let outcome2 = before_fn(&ok_ctx).await;
669        assert!(matches!(outcome2, HookOutcome::Continue));
670        assert_eq!(counter.load(Ordering::Relaxed), 2);
671    }
672
673    #[test]
674    fn too_large_result_mentions_limit_and_actual() {
675        let r = too_large_result(100, 500, "my_tool");
676        let body = serde_json::to_string(&r).unwrap();
677        assert!(body.contains("result_too_large"));
678        assert!(body.contains("my_tool"));
679        assert!(body.contains("100"));
680        assert!(body.contains("500"));
681    }
682
683    #[tokio::test]
684    async fn replace_outcome_skips_inner_and_returns_payload() {
685        // Returning Replace from before-hook must yield the supplied
686        // CallToolResult directly, with no need for the inner handler.
687        let before: BeforeHook = Arc::new(|_ctx| {
688            Box::pin(async {
689                HookOutcome::Replace(Box::new(CallToolResult::success(vec![ContentBlock::text(
690                    "from-replace".to_owned(),
691                )])))
692            })
693        });
694        let hooks = Arc::new(ToolHooks {
695            max_result_bytes: None,
696            before: Some(before),
697            after: None,
698        });
699        let _hooked = with_hooks(TestHandler::default(), Arc::clone(&hooks));
700
701        // Exercise the before-hook closure + apply_size_cap helper directly,
702        // matching the established test pattern in this module.
703        let outcome = (hooks.before.as_ref().unwrap())(&ctx("any")).await;
704        let HookOutcome::Replace(boxed) = outcome else {
705            panic!("expected HookOutcome::Replace");
706        };
707        let (result, size, capped) = apply_size_cap(*boxed, None, "any");
708        assert!(!capped);
709        assert!(size > 0);
710        assert!(!result.is_error.unwrap_or(false));
711        assert!(matches!(
712            result.content.first(),
713            Some(rmcp::model::ContentBlock::Text(t)) if t.text == "from-replace"
714        ));
715    }
716
717    #[tokio::test]
718    async fn replace_outcome_subject_to_size_cap() {
719        // A Replace payload that exceeds max_result_bytes must be rewritten
720        // to result_too_large just like an inner-handler result would be,
721        // and the disposition must reflect ResultTooLarge.
722        let huge = CallToolResult::success(vec![ContentBlock::text("y".repeat(8_192))]);
723        let huge_size = serialized_size(&huge);
724        assert!(huge_size > 256);
725
726        let (final_result, accounted, capped) = apply_size_cap(huge, Some(256), "replaced_tool");
727        assert!(capped);
728        assert_eq!(accounted, huge_size);
729        assert_eq!(final_result.is_error, Some(true));
730        assert!(matches!(
731            final_result.content.first(),
732            Some(rmcp::model::ContentBlock::Text(t)) if t.text.contains("result_too_large")
733        ));
734    }
735
736    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
737    async fn after_hook_fires_exactly_once_via_spawn() {
738        // spawn_after must enqueue the after-hook exactly one time per
739        // invocation and never block the caller; we wait for the spawned
740        // task to run by polling the counter with a short timeout.
741        let counter = Arc::new(AtomicUsize::new(0));
742        let c = Arc::clone(&counter);
743        let after: AfterHook = Arc::new(move |_ctx, _disp, _size| {
744            let c = Arc::clone(&c);
745            Box::pin(async move {
746                c.fetch_add(1, Ordering::Relaxed);
747            })
748        });
749        let holder = Arc::new(AfterHookHolder { f: after });
750
751        HookedHandler::<TestHandler>::spawn_after(
752            Some(&holder),
753            ctx("t"),
754            HookDisposition::InnerExecuted,
755            42,
756        );
757
758        // Wait up to 1s for the spawned task to run.
759        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(1);
760        while counter.load(Ordering::Relaxed) == 0 && std::time::Instant::now() < deadline {
761            tokio::task::yield_now().await;
762            tokio::time::sleep(std::time::Duration::from_millis(5)).await;
763        }
764        assert_eq!(counter.load(Ordering::Relaxed), 1);
765    }
766
767    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
768    async fn after_hook_panic_is_isolated_from_response_path() {
769        // A panicking after-hook must not affect the request task.  We
770        // spawn a panicking after-hook and then verify the current task
771        // can still complete an unrelated future to completion.
772        let after: AfterHook = Arc::new(|_ctx, _disp, _size| {
773            Box::pin(async {
774                panic!("intentional panic in after-hook");
775            })
776        });
777        let holder = Arc::new(AfterHookHolder { f: after });
778
779        HookedHandler::<TestHandler>::spawn_after(
780            Some(&holder),
781            ctx("boom"),
782            HookDisposition::InnerExecuted,
783            0,
784        );
785
786        // Give Tokio a chance to run + abort the panicking task, then
787        // confirm we're still alive and the runtime is healthy.
788        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
789        let still_alive = tokio::spawn(async { 1_u32 + 2 }).await.unwrap();
790        assert_eq!(still_alive, 3);
791    }
792}
rmcp_server_kit/tool_hooks.rs

rmcp_server_kit/
tool_hooks.rs
