agentcat_sdk/
tracked.rs

1//! [`Tracked<H>`]: a `ServerHandler` that wraps the customer's handler and forwards every method
2//! of rmcp's trait to it, in the order and with the signatures of rmcp 3.0's own wrapper impls
3//! (`impl_server_handler_for_wrapper!`). `negotiate_initialize` (rmcp 3.3) is not forwarded:
4//! only rmcp's default `initialize` calls it, and `initialize` is forwarded. The four methods the
5//! SDK changes (`initialize`, `list_tools`, `get_tool`, `call_tool`) forward untouched when
6//! `AGENTCAT_ENABLED` turned the handle off; `call_tool` then still adds a
7//! [`CallContext`](crate::CallContext) without session or agent.
8
9use std::any::type_name;
10use std::borrow::Cow;
11use std::fmt;
12use std::sync::{Mutex, PoisonError};
13use std::time::SystemTime;
14
15use rmcp::model::{
16    CallToolRequestParams, CallToolResponse, CancelTaskParams, CancelledNotificationParam,
17    CompleteRequestParams, CompleteResult, CustomNotification, CustomRequest, CustomResult,
18    DiscoverResult, GetPromptRequestParams, GetPromptResponse, GetTaskParams, GetTaskResult,
19    Implementation, InitializeRequestParams, InitializeResult, ListPromptsResult,
20    ListResourceTemplatesResult, ListResourcesResult, ListToolsResult, PaginatedRequestParams,
21    ProgressNotificationParam, ProtocolVersion, ReadResourceRequestParams, ReadResourceResponse,
22    SubscribeRequestParams, SubscriptionFilter, Tool, UnsubscribeRequestParams, UpdateTaskParams,
23};
24use rmcp::service::{MaybeSendFuture, NotificationContext, RequestContext, SubscriptionContext};
25use rmcp::{ErrorData, RoleServer, ServerHandler};
26
27use crate::call::run::{run_call, server_info};
28use crate::client::AgentCat;
29use crate::constants::GET_MORE_TOOLS_NAME;
30use crate::custom::{CustomEvent, publish_server};
31use crate::error::Error;
32use crate::event::ClientIdentity;
33use crate::inject::{
34    InjectConfig, ToolEntry, build_injected_list, get_more_tools_tool, inject_tool,
35};
36
37#[cfg(test)]
38mod tests;
39
40/// A server handler wrapped by [`AgentCat::track`] or [`track`](crate::track). It implements
41/// rmcp's `ServerHandler`, so it is served like the handler it wraps. When `AGENTCAT_ENABLED`
42/// turned the handle off, it forwards every method to the wrapped handler untouched, and a tool
43/// still gets a [`CallContext`](crate::CallContext), with no session and no agent.
44pub struct Tracked<H> {
45    inner: H,
46    agentcat: AgentCat,
47    /// The client that `initialize` named on this instance. A clone starts with none.
48    identity: Mutex<Option<Implementation>>,
49}
50
51impl<H: Clone> Clone for Tracked<H> {
52    /// Clones the handler and the handle; the clone has seen no `initialize`, so it holds no
53    /// client identity.
54    fn clone(&self) -> Self {
55        Self::new(self.inner.clone(), self.agentcat.clone())
56    }
57}
58
59/// The warning logged for a listing with no tools.
60const EMPTY_LISTING_WARNING: &str = "Warning: No tools found in the original list. This is likely due to the tools not being registered before `track`.";
61
62const _: () = {
63    const fn send_sync_static<T: Send + Sync + 'static>() {}
64    const fn tracked<H: Send + Sync + 'static>() {
65        send_sync_static::<Tracked<H>>();
66    }
67    tracked::<()>();
68};
69
70impl<H> Tracked<H> {
71    /// Wraps `inner` for `agentcat`.
72    pub(crate) fn new(inner: H, agentcat: AgentCat) -> Self {
73        Self {
74            inner,
75            agentcat,
76            identity: Mutex::new(None),
77        }
78    }
79
80    /// The wrapped handler.
81    pub fn inner(&self) -> &H {
82        &self.inner
83    }
84
85    /// The wrapped handler, mutably.
86    pub fn inner_mut(&mut self) -> &mut H {
87        &mut self.inner
88    }
89
90    /// Unwraps the handler.
91    pub fn into_inner(self) -> H {
92        self.inner
93    }
94
95    /// The handle this wrapper reports to.
96    pub fn agentcat(&self) -> &AgentCat {
97        &self.agentcat
98    }
99
100    /// The handler's tool `name` as the listing shows it, with the entry the pipeline gives it
101    /// (its `declares_session` is the foreign status of `session_id`). The pipeline runs on
102    /// a copy of the registry's state and nothing is written back, so a lookup never uses up a
103    /// log line that only a listing writes. The lock is held to copy the state only.
104    pub(crate) fn lookup_tool(&self, name: &str) -> Option<(Tool, ToolEntry)>
105    where
106        H: ServerHandler,
107    {
108        let tool = self.inner.get_tool(name)?;
109        let config = InjectConfig::from_options(&self.agentcat.inner().options);
110        let mut state = self
111            .agentcat
112            .inner()
113            .with_registry::<H, _>(|registry| registry.state_mut().clone());
114        Some(inject_tool(&config, tool, &mut state, &mut Vec::new()))
115    }
116
117    /// The client that `initialize` named on this instance, if it has run.
118    pub(crate) fn client_identity(&self) -> Option<Implementation> {
119        self.identity
120            .lock()
121            .unwrap_or_else(PoisonError::into_inner)
122            .clone()
123    }
124
125    /// Publishes a custom event with the identity of this instance. It only queues the event, so
126    /// it returns at once. The client is the one this instance's `initialize` named (none before
127    /// it, and none on a clone, which has seen no `initialize`); the server is what the handler's
128    /// `get_info` gives now, and a panic there is logged and leaves the event without server
129    /// fields.
130    ///
131    /// The event goes to the project of the handle, not to `project_id`, which only has to be
132    /// non-empty. Nothing is queued while tracing is off, or by a handle that `AGENTCAT_ENABLED`
133    /// turned off. The `redact_event` hook of the options runs on the event, on the runtime of the
134    /// caller; the string hook does not. The event has no session unless the [`CustomEvent`] names
135    /// one: it takes none from anywhere else. Inside a tool, pass
136    /// [`CallContext::session_id`](crate::CallContext::session_id) in the [`CustomEvent`] to
137    /// attribute the event to the call. A `duration` of zero or none is the time since the event
138    /// was built.
139    ///
140    /// ```no_run
141    /// # fn run<H: rmcp::ServerHandler>(
142    /// #     tracked: agentcat_sdk::Tracked<H>,
143    /// # ) -> Result<(), agentcat_sdk::Error> {
144    /// use agentcat_sdk::CustomEvent;
145    ///
146    /// tracked.publish_custom_event("proj_123", CustomEvent::new().resource_name("startup"))?;
147    /// # Ok(())
148    /// # }
149    /// ```
150    ///
151    /// # Errors
152    ///
153    /// [`Error::ProjectIdRequired`] when `project_id` is empty, except on a handle that
154    /// `AGENTCAT_ENABLED` turned off, which returns `Ok(())` and does nothing, not even call the
155    /// handler's `get_info`.
156    pub fn publish_custom_event(&self, project_id: &str, event: CustomEvent) -> Result<(), Error>
157    where
158        H: ServerHandler,
159    {
160        if !self.agentcat.inner().enabled {
161            return Ok(());
162        }
163        if project_id.is_empty() {
164            return Err(Error::ProjectIdRequired);
165        }
166        let created = SystemTime::now();
167        let client = self.client_identity().as_ref().map(ClientIdentity::from);
168        let server = server_info(self.agentcat.inner(), &self.inner);
169        publish_server(
170            &self.agentcat,
171            project_id,
172            event,
173            client.as_ref(),
174            server.as_ref(),
175            created,
176            SystemTime::now(),
177        )
178    }
179}
180
181impl<H> fmt::Debug for Tracked<H> {
182    /// Names the wrapped handler's type, so `H` needs no `Debug` impl.
183    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
184        f.debug_struct("Tracked")
185            .field("inner", &type_name::<H>())
186            .field("agentcat", &self.agentcat)
187            .finish_non_exhaustive()
188    }
189}
190
191impl<H: ServerHandler> ServerHandler for Tracked<H> {
192    fn ping(
193        &self,
194        context: RequestContext<RoleServer>,
195    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
196        self.inner.ping(context)
197    }
198
199    #[expect(
200        clippy::manual_async_fn,
201        reason = "the signature of rmcp's trait method: a returned future with MaybeSendFuture"
202    )]
203    fn initialize(
204        &self,
205        request: InitializeRequestParams,
206        context: RequestContext<RoleServer>,
207    ) -> impl Future<Output = Result<InitializeResult, ErrorData>> + MaybeSendFuture + '_ {
208        async move {
209            if !self.agentcat.inner().enabled {
210                return self.inner.initialize(request, context).await;
211            }
212            let client = request.client_info.clone();
213            let result = self.inner.initialize(request, context).await?;
214            *self.identity.lock().unwrap_or_else(PoisonError::into_inner) = Some(client);
215            Ok(result)
216        }
217    }
218
219    fn supported_protocol_versions(&self) -> Cow<'static, [ProtocolVersion]> {
220        self.inner.supported_protocol_versions()
221    }
222
223    fn discover(
224        &self,
225        context: RequestContext<RoleServer>,
226    ) -> impl Future<Output = Result<DiscoverResult, ErrorData>> + MaybeSendFuture + '_ {
227        self.inner.discover(context)
228    }
229
230    fn complete(
231        &self,
232        request: CompleteRequestParams,
233        context: RequestContext<RoleServer>,
234    ) -> impl Future<Output = Result<CompleteResult, ErrorData>> + MaybeSendFuture + '_ {
235        self.inner.complete(request, context)
236    }
237
238    #[expect(
239        deprecated,
240        reason = "SetLevelRequestParams is SEP-2577-deprecated on every rmcp 3.x; the method is still forwarded"
241    )]
242    fn set_level(
243        &self,
244        request: rmcp::model::SetLevelRequestParams,
245        context: RequestContext<RoleServer>,
246    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
247        self.inner.set_level(request, context)
248    }
249
250    fn get_prompt(
251        &self,
252        request: GetPromptRequestParams,
253        context: RequestContext<RoleServer>,
254    ) -> impl Future<Output = Result<GetPromptResponse, ErrorData>> + MaybeSendFuture + '_ {
255        self.inner.get_prompt(request, context)
256    }
257
258    fn list_prompts(
259        &self,
260        request: Option<PaginatedRequestParams>,
261        context: RequestContext<RoleServer>,
262    ) -> impl Future<Output = Result<ListPromptsResult, ErrorData>> + MaybeSendFuture + '_ {
263        self.inner.list_prompts(request, context)
264    }
265
266    fn list_resources(
267        &self,
268        request: Option<PaginatedRequestParams>,
269        context: RequestContext<RoleServer>,
270    ) -> impl Future<Output = Result<ListResourcesResult, ErrorData>> + MaybeSendFuture + '_ {
271        self.inner.list_resources(request, context)
272    }
273
274    fn list_resource_templates(
275        &self,
276        request: Option<PaginatedRequestParams>,
277        context: RequestContext<RoleServer>,
278    ) -> impl Future<Output = Result<ListResourceTemplatesResult, ErrorData>> + MaybeSendFuture + '_
279    {
280        self.inner.list_resource_templates(request, context)
281    }
282
283    fn read_resource(
284        &self,
285        request: ReadResourceRequestParams,
286        context: RequestContext<RoleServer>,
287    ) -> impl Future<Output = Result<ReadResourceResponse, ErrorData>> + MaybeSendFuture + '_ {
288        self.inner.read_resource(request, context)
289    }
290
291    fn accepted_subscription_filter(
292        &self,
293        requested: &SubscriptionFilter,
294    ) -> Option<SubscriptionFilter> {
295        self.inner.accepted_subscription_filter(requested)
296    }
297
298    fn listen(
299        &self,
300        context: SubscriptionContext,
301    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
302        self.inner.listen(context)
303    }
304
305    #[expect(
306        deprecated,
307        reason = "resources/subscribe is legacy-only and deprecated on every rmcp 3.x; it is still forwarded"
308    )]
309    fn subscribe(
310        &self,
311        request: SubscribeRequestParams,
312        context: RequestContext<RoleServer>,
313    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
314        self.inner.subscribe(request, context)
315    }
316
317    #[expect(
318        deprecated,
319        reason = "resources/unsubscribe is legacy-only and deprecated on every rmcp 3.x; it is still forwarded"
320    )]
321    fn unsubscribe(
322        &self,
323        request: UnsubscribeRequestParams,
324        context: RequestContext<RoleServer>,
325    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
326        self.inner.unsubscribe(request, context)
327    }
328
329    /// The tool-call pipeline around the handler's `call_tool`: interception, stripping,
330    /// mint-back and the record of the call; on a handle that is off, the handler's `call_tool`
331    /// with a [`CallContext`](crate::CallContext) without session or agent.
332    fn call_tool(
333        &self,
334        request: CallToolRequestParams,
335        context: RequestContext<RoleServer>,
336    ) -> impl Future<Output = Result<CallToolResponse, ErrorData>> + MaybeSendFuture + '_ {
337        run_call(self, request, context)
338    }
339
340    #[expect(
341        clippy::manual_async_fn,
342        reason = "the signature of rmcp's trait method: a returned future with MaybeSendFuture"
343    )]
344    fn list_tools(
345        &self,
346        request: Option<PaginatedRequestParams>,
347        context: RequestContext<RoleServer>,
348    ) -> impl Future<Output = Result<ListToolsResult, ErrorData>> + MaybeSendFuture + '_ {
349        async move {
350            let shared = self.agentcat.inner();
351            if !shared.enabled {
352                return self.inner.list_tools(request, context).await;
353            }
354            let mut result = self.inner.list_tools(request, context).await?;
355            if result.tools.is_empty() {
356                shared.sink.log(EMPTY_LISTING_WARNING.to_owned());
357                return Ok(result);
358            }
359            let config = InjectConfig::from_options(&shared.options);
360            let tools = std::mem::take(&mut result.tools);
361            let list = shared.with_registry::<H, _>(|registry| {
362                let list = build_injected_list(&config, tools, registry.state_mut());
363                registry.record(&list);
364                list
365            });
366            for line in list.logs {
367                shared.sink.log(line);
368            }
369            result.tools = list.tools;
370            Ok(result)
371        }
372    }
373
374    /// The injected descriptor of `get_more_tools`, or the handler's tool through the pipeline,
375    /// so the call path sees what the listing showed. The pipeline's log lines are dropped: the
376    /// listing logs them once. On a handle that is off, the handler's tool.
377    fn get_tool(&self, name: &str) -> Option<Tool> {
378        let shared = self.agentcat.inner();
379        if !shared.enabled {
380            return self.inner.get_tool(name);
381        }
382        if shared.options.report_missing() && name == GET_MORE_TOOLS_NAME {
383            return Some(get_more_tools_tool());
384        }
385        self.lookup_tool(name).map(|(tool, _entry)| tool)
386    }
387
388    fn on_custom_request(
389        &self,
390        request: CustomRequest,
391        context: RequestContext<RoleServer>,
392    ) -> impl Future<Output = Result<CustomResult, ErrorData>> + MaybeSendFuture + '_ {
393        self.inner.on_custom_request(request, context)
394    }
395
396    fn on_cancelled(
397        &self,
398        notification: CancelledNotificationParam,
399        context: NotificationContext<RoleServer>,
400    ) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
401        self.inner.on_cancelled(notification, context)
402    }
403
404    fn on_progress(
405        &self,
406        notification: ProgressNotificationParam,
407        context: NotificationContext<RoleServer>,
408    ) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
409        self.inner.on_progress(notification, context)
410    }
411
412    fn on_initialized(
413        &self,
414        context: NotificationContext<RoleServer>,
415    ) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
416        self.inner.on_initialized(context)
417    }
418
419    fn on_roots_list_changed(
420        &self,
421        context: NotificationContext<RoleServer>,
422    ) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
423        self.inner.on_roots_list_changed(context)
424    }
425
426    fn on_custom_notification(
427        &self,
428        notification: CustomNotification,
429        context: NotificationContext<RoleServer>,
430    ) -> impl Future<Output = ()> + MaybeSendFuture + '_ {
431        self.inner.on_custom_notification(notification, context)
432    }
433
434    fn get_info(&self) -> InitializeResult {
435        self.inner.get_info()
436    }
437
438    fn get_task(
439        &self,
440        request: GetTaskParams,
441        context: RequestContext<RoleServer>,
442    ) -> impl Future<Output = Result<GetTaskResult, ErrorData>> + MaybeSendFuture + '_ {
443        self.inner.get_task(request, context)
444    }
445
446    fn update_task(
447        &self,
448        request: UpdateTaskParams,
449        context: RequestContext<RoleServer>,
450    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
451        self.inner.update_task(request, context)
452    }
453
454    fn cancel_task(
455        &self,
456        request: CancelTaskParams,
457        context: RequestContext<RoleServer>,
458    ) -> impl Future<Output = Result<(), ErrorData>> + MaybeSendFuture + '_ {
459        self.inner.cancel_task(request, context)
460    }
461}
agentcat_sdk/tracked.rs

agentcat_sdk/
tracked.rs
