agentcat_sdk/
hooks.rs

1//! The customer hooks of [`Options`]: the [`ToolCall`] snapshot they receive and [`fire`], which
2//! starts them for one tool call, before the tool runs. The hooks run as `Send` tasks on the
3//! tokio runtime the call runs on, and the call never waits for them: the delivery pipeline
4//! settles the [`PendingHooks`] when it finishes the event. `rmcp/local` is not supported.
5use std::fmt;
6use std::future::Future;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9
10use http::request::Parts;
11use rmcp::RoleServer;
12use rmcp::model::{CallToolRequestParams, JsonObject};
13use rmcp::service::RequestContext;
14use serde_json::{Map, Value};
15use tokio::runtime::Handle;
16use tokio::task::{JoinError, JoinHandle};
17
18use crate::config::{Hook, HookError, Options, UserIdentity};
19use crate::event::{
20    ClientIdentity, project_headers, protocol_version, redact_headers, transport_session_id,
21};
22use crate::sink::Sink;
23
24#[cfg(test)]
25mod tests;
26
27#[cfg(test)]
28thread_local! {
29    /// How many [`ToolCall`] snapshots this thread has built.
30    static SNAPSHOTS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
31}
32
33/// How many [`ToolCall`] snapshots this thread has built. A test that serves and calls a tool on
34/// one current-thread runtime counts the snapshots of its calls with it (the snapshot is built
35/// only when a hook starts, [`fire`]).
36#[cfg(test)]
37pub(crate) fn snapshots_built() -> usize {
38    SNAPSHOTS.get()
39}
40
41/// One tool call as a hook sees it: the request as the client sent it, with the injected
42/// arguments (`session_id`, `agent_id`, `context`) still in it, and what the SDK knows about the
43/// client and the transport. Clones are independent copies.
44///
45/// Its `Debug` output hides credentials: the values of the `authorization`,
46/// `proxy-authorization`, `cookie` and `set-cookie` headers print as `"<redacted>"`, and the
47/// request context prints as its request ID and whether it holds HTTP request parts.
48///
49/// ```
50/// # use agentcat_sdk::{Options, ToolCall};
51/// let options = Options::default().event_tags(|call: ToolCall| async move {
52///     let mut tags = serde_json::Map::new();
53///     tags.insert("tool".into(), call.tool_name().into());
54///     Ok(Some(tags))
55/// });
56/// # let _ = options;
57/// ```
58#[derive(Clone)]
59pub struct ToolCall {
60    tool_name: String,
61    arguments: Option<JsonObject>,
62    meta: JsonObject,
63    headers: Option<Map<String, Value>>,
64    client: Option<ClientIdentity>,
65    protocol_version: Option<String>,
66    transport_session_id: Option<String>,
67    request_context: RequestContext<RoleServer>,
68}
69
70impl ToolCall {
71    /// The snapshot of `request` before the SDK changes it. `client` is the identity of the
72    /// request's client ([`client_identity`](crate::event::client_identity)).
73    pub(crate) fn new(
74        request: &CallToolRequestParams,
75        context: &RequestContext<RoleServer>,
76        client: Option<&ClientIdentity>,
77    ) -> Self {
78        #[cfg(test)]
79        SNAPSHOTS.set(SNAPSHOTS.get() + 1);
80        let meta: &JsonObject = &context.meta;
81        Self {
82            tool_name: request.name.to_string(),
83            arguments: request.arguments.clone(),
84            meta: meta.clone(),
85            headers: context
86                .extensions
87                .get::<Parts>()
88                .map(|parts| project_headers(&parts.headers)),
89            client: client.cloned(),
90            protocol_version: protocol_version(meta),
91            transport_session_id: transport_session_id(context),
92            request_context: context.clone(),
93        }
94    }
95
96    /// The name of the tool.
97    #[must_use]
98    pub fn tool_name(&self) -> &str {
99        &self.tool_name
100    }
101
102    /// The arguments as the client sent them, the `session_id`, `agent_id` and `context` that the
103    /// SDK added to the tool's listing included.
104    #[must_use]
105    pub fn arguments(&self) -> Option<&JsonObject> {
106        self.arguments.as_ref()
107    }
108
109    /// The `_meta` of the request; empty when the client sent none.
110    #[must_use]
111    pub fn meta(&self) -> &JsonObject {
112        &self.meta
113    }
114
115    /// The HTTP headers of the request, every header included, `authorization` and `cookie` too:
116    /// names in lower case, in byte order, after the
117    /// names that are array indexes in numeric order; values trimmed of leading and trailing HTTP
118    /// whitespace; the values of a repeated name joined in arrival order with `", "` (`"; "` for
119    /// `cookie`), except `set-cookie`, which keeps only its last value; a value that is not UTF-8
120    /// read lossily. `None` on a transport without HTTP, such as stdio.
121    #[must_use]
122    pub fn headers(&self) -> Option<&Map<String, Value>> {
123        self.headers.as_ref()
124    }
125
126    /// The name the client gave itself, when it did.
127    #[must_use]
128    pub fn client_name(&self) -> Option<&str> {
129        self.client.as_ref()?.name.as_deref()
130    }
131
132    /// The version the client gave itself, when it did.
133    #[must_use]
134    pub fn client_version(&self) -> Option<&str> {
135        self.client.as_ref()?.version.as_deref()
136    }
137
138    /// The protocol version the request names in its `_meta`, when it does.
139    #[must_use]
140    pub fn protocol_version(&self) -> Option<&str> {
141        self.protocol_version.as_deref()
142    }
143
144    /// The `mcp-session-id` of the streamable HTTP session; `None` on stdio and on a stateless
145    /// HTTP server.
146    #[must_use]
147    pub fn transport_session_id(&self) -> Option<&str> {
148        self.transport_session_id.as_deref()
149    }
150
151    /// The `rmcp` request context of the call, for anything the other getters do not give. The
152    /// SDK has not yet added its [`CallContext`](crate::CallContext) to it.
153    #[must_use]
154    pub fn request_context(&self) -> &RequestContext<RoleServer> {
155        &self.request_context
156    }
157}
158
159impl fmt::Debug for ToolCall {
160    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
161        f.debug_struct("ToolCall")
162            .field("tool_name", &self.tool_name)
163            .field("arguments", &self.arguments)
164            .field("meta", &self.meta)
165            .field("headers", &self.headers.as_ref().map(redact_headers))
166            .field("client", &self.client)
167            .field("protocol_version", &self.protocol_version)
168            .field("transport_session_id", &self.transport_session_id)
169            .field("request_context", &ContextSummary(&self.request_context))
170            .finish()
171    }
172}
173
174/// A request context in `Debug` output: its request ID and whether it holds HTTP request parts,
175/// never the parts themselves, whose headers carry credentials.
176struct ContextSummary<'a>(&'a RequestContext<RoleServer>);
177
178impl fmt::Debug for ContextSummary<'_> {
179    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
180        f.debug_struct("RequestContext")
181            .field("id", &self.0.id)
182            .field("http_parts", &self.0.extensions.get::<Parts>().is_some())
183            .finish_non_exhaustive()
184    }
185}
186
187/// What a hook gives back once it settles.
188type Settled<T> = Result<Option<T>, HookError>;
189
190/// The tasks of the hooks started for one call, `None` for a hook that is not set or did not
191/// start. Dropping it aborts every task that has not finished (a detached task would leak), as
192/// when the queue drops the event (overflow, a shutdown deadline). Take a handle
193/// out with `Option::take` to keep its task.
194#[derive(Debug, Default)]
195pub(crate) struct PendingHooks {
196    /// The `resolve_session_id` task; only in hook mode.
197    pub(crate) resolve_session_id: Option<JoinHandle<Settled<String>>>,
198    /// The `identify` task.
199    pub(crate) identify: Option<JoinHandle<Settled<UserIdentity>>>,
200    /// The `event_tags` task.
201    pub(crate) event_tags: Option<JoinHandle<Settled<Map<String, Value>>>>,
202    /// The `event_properties` task.
203    pub(crate) event_properties: Option<JoinHandle<Settled<Map<String, Value>>>>,
204}
205
206impl PendingHooks {
207    /// Whether any hook started.
208    pub(crate) fn any(&self) -> bool {
209        self.resolve_session_id.is_some()
210            || self.identify.is_some()
211            || self.event_tags.is_some()
212            || self.event_properties.is_some()
213    }
214}
215
216impl Drop for PendingHooks {
217    fn drop(&mut self) {
218        // `abort` does nothing to a task that has finished.
219        if let Some(task) = &self.resolve_session_id {
220            task.abort();
221        }
222        if let Some(task) = &self.identify {
223            task.abort();
224        }
225        if let Some(task) = &self.event_tags {
226            task.abort();
227        }
228        if let Some(task) = &self.event_properties {
229            task.abort();
230        }
231    }
232}
233
234/// A task handle that aborts its task when it is dropped, so a customer hook whose event is
235/// dropped (a cancelled pipeline, a shutdown deadline) does not keep running detached. Awaiting it
236/// awaits the task.
237#[derive(Debug)]
238pub(crate) struct AbortOnDrop<T>(JoinHandle<T>);
239
240impl<T> AbortOnDrop<T> {
241    /// Takes over `task`.
242    pub(crate) fn new(task: JoinHandle<T>) -> Self {
243        Self(task)
244    }
245}
246
247impl<T> Future for AbortOnDrop<T> {
248    type Output = Result<T, JoinError>;
249
250    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
251        Pin::new(&mut self.0).poll(cx)
252    }
253}
254
255impl<T> Drop for AbortOnDrop<T> {
256    fn drop(&mut self) {
257        // `abort` does nothing to a task that has finished.
258        self.0.abort();
259    }
260}
261
262/// Starts the hooks of `options` for one call, in this order: `resolve_session_id` (in hook mode
263/// only), `identify`, `event_tags`, `event_properties`. Each hook runs as a task on the current
264/// tokio runtime, and none is awaited. A panic in a hook stays in its
265/// task. Without a runtime no hook starts, and `sink` gets one log line for each hook that would
266/// have.
267///
268/// `snapshot` builds the [`ToolCall`] the hooks receive. It runs once, and only when at least one
269/// hook starts, so a call without hooks copies nothing.
270pub(crate) fn fire(
271    options: &Options,
272    hook_mode: bool,
273    sink: &dyn Sink,
274    snapshot: impl FnOnce() -> ToolCall,
275) -> PendingHooks {
276    let resolve_session_id = options.resolve_session_id_hook().filter(|_| hook_mode);
277    let identify = options.identify_hook();
278    let event_tags = options.event_tags_hook();
279    let event_properties = options.event_properties_hook();
280    let set = [
281        ("resolve_session_id", resolve_session_id.is_some()),
282        ("identify", identify.is_some()),
283        ("event_tags", event_tags.is_some()),
284        ("event_properties", event_properties.is_some()),
285    ];
286    if !set.iter().any(|&(_, is_set)| is_set) {
287        return PendingHooks::default();
288    }
289    let Ok(handle) = Handle::try_current() else {
290        for (label, _) in set.into_iter().filter(|&(_, is_set)| is_set) {
291            sink.log(format!(
292                "Warning: {label} hook not started - no tokio runtime"
293            ));
294        }
295        return PendingHooks::default();
296    };
297    let call = snapshot();
298    PendingHooks {
299        resolve_session_id: resolve_session_id.map(|hook| spawn(&handle, hook, &call)),
300        identify: identify.map(|hook| spawn(&handle, hook, &call)),
301        event_tags: event_tags.map(|hook| spawn(&handle, hook, &call)),
302        event_properties: event_properties.map(|hook| spawn(&handle, hook, &call)),
303    }
304}
305
306/// Spawns `hook` on `call` on the runtime of `handle`.
307fn spawn<T: Send + 'static>(
308    handle: &Handle,
309    hook: &Hook<T>,
310    call: &ToolCall,
311) -> JoinHandle<Settled<T>> {
312    handle.spawn(hook.call(call.clone()))
313}
agentcat_sdk/hooks.rs

agentcat_sdk/
hooks.rs
