agentcat-sdk 1.0.0

//! The hooks: what a [`ToolCall`] holds, the order in which the hooks start, that they start before
//! the tool runs and never delay the call, that a panic stays in its task, and what hook mode
//! changes for the listing, the resolution and the mirror. The settling of the hooks and the event
//! body are tested elsewhere.
use std::future::{Future, pending, ready};
use std::sync::{Arc, Mutex, PoisonError};
use std::time::Duration;

use rmcp::model::{
    CallToolRequestParams, CallToolResponse, CallToolResult, InitializeResult, ServerCapabilities,
};
use rmcp::service::{MaybeSendFuture, RequestContext, RoleClient, RunningService};
use rmcp::{ErrorData, RoleServer, ServerHandler};
use serde_json::{Map, Value, json};

use super::{ToolCall, fire};
use crate::client::AgentCat;
use crate::config::{HookError, Options};
use crate::constants::META_PROTOCOL_VERSION_KEY;
use crate::event::CallRecord;
use crate::handles::SessionSource;
use crate::sink::Sink;
use crate::test_support::rmcp::{
    CREDENTIAL_HEADERS, RecordingSink, Sample, assert_credentials_hidden, bare_context, connect,
    credential_parts,
};

type Client = RunningService<RoleClient, rmcp::model::InitializeRequestParams>;

/// The labels that hooks and tools push, in the order they happen.
type Log = Arc<Mutex<Vec<String>>>;

fn push(log: &Log, label: &str) {
    log.lock()
        .unwrap_or_else(PoisonError::into_inner)
        .push(label.to_owned());
}

fn labels(log: &Log) -> Vec<String> {
    log.lock().unwrap_or_else(PoisonError::into_inner).clone()
}

/// A hook that pushes `label` when it is first polled and gives nothing.
fn labelled<T: Send + 'static>(
    log: &Log,
    label: &'static str,
) -> impl Fn(ToolCall) -> std::future::Ready<Result<Option<T>, HookError>> + Send + Sync + 'static {
    let log = Arc::clone(log);
    move |_call| {
        push(&log, label);
        ready(Ok(None))
    }
}

/// Serves `Sample` tracked with `options`; the records go to the returned sink.
async fn serve(options: Options) -> (Client, Arc<RecordingSink>) {
    let sink = Arc::new(RecordingSink::default());
    let agentcat = AgentCat::with_sink("", options, sink.clone());
    (connect(agentcat.track(Sample), "hook-client").await, sink)
}

fn params(name: &str, arguments: Value) -> CallToolRequestParams {
    let Value::Object(arguments) = arguments else {
        panic!("arguments must be an object");
    };
    CallToolRequestParams::new(name.to_owned()).with_arguments(arguments)
}

async fn call(client: &Client, name: &str, arguments: Value) -> CallToolResult {
    client
        .call_tool(params(name, arguments))
        .await
        .unwrap_or_else(|e| panic!("{name}: {e}"))
}

fn only_record(sink: &RecordingSink) -> CallRecord {
    let mut records = sink.take_records();
    assert_eq!(records.len(), 1, "one record: {records:?}");
    records.remove(0)
}

fn texts(result: &CallToolResult) -> Vec<&str> {
    result
        .content
        .iter()
        .filter_map(|block| block.as_text().map(|text| text.text.as_str()))
        .collect()
}

#[tokio::test(flavor = "current_thread")]
async fn hooks_receive_raw_arguments() {
    let seen: Arc<Mutex<Vec<ToolCall>>> = Arc::default();
    let options = Options::default().enable_agent_tracking(true).event_tags({
        let seen = Arc::clone(&seen);
        move |call| {
            seen.lock()
                .unwrap_or_else(PoisonError::into_inner)
                .push(call);
            ready(Ok(None))
        }
    });
    let (client, _sink) = serve(options).await;
    let sent = json!({
        "msg": "hi",
        "session_id": "start",
        "agent_id": "opus|test|a1b2c",
        "context": "Checking the raw arguments."
    });
    call(&client, "echo", sent.clone()).await;

    client.cancel().await.unwrap();
    let seen = seen.lock().unwrap_or_else(PoisonError::into_inner);
    let [tool_call] = seen.as_slice() else {
        panic!("one hook call, got {}", seen.len());
    };
    assert_eq!(tool_call.tool_name(), "echo");
    assert_eq!(
        tool_call.arguments().map(|a| Value::Object(a.clone())),
        Some(sent),
        "injected keys included, in the order sent"
    );
    assert_eq!(tool_call.client_name(), Some("hook-client"));
    assert_eq!(tool_call.client_version(), Some("1.0.0"));
    // rmcp's client puts a progress token in the `_meta` of each request it sends.
    assert!(tool_call.meta().contains_key("progressToken"));
    assert_eq!(tool_call.headers(), None, "no HTTP on an in-memory pipe");
    assert_eq!(tool_call.transport_session_id(), None);
    assert_eq!(tool_call.protocol_version(), None);
    assert!(
        tool_call
            .request_context()
            .extensions
            .get::<crate::CallContext>()
            .is_none(),
        "the snapshot is taken before the CallContext is added"
    );
}

/// The snapshot keeps the credential headers, and its `Debug` output hides them: the values of
/// `authorization`, `proxy-authorization`, `cookie` and `set-cookie` print as `<redacted>` under
/// their names, other headers print as they are, and the request context prints as a summary (its
/// id, whether HTTP parts are present), not with the raw headers of its parts. `Debug` output
/// exists only in Rust.
#[tokio::test(flavor = "current_thread")]
async fn tool_call_debug_hides_credentials() {
    let mut context = bare_context(None);
    context.extensions.insert(credential_parts());
    let call = ToolCall::new(&params("echo", json!({"x": 1})), &context, None);

    let text = format!("{call:?}");
    assert_credentials_hidden(&text);
    assert!(text.contains("<redacted>"), "{text}");
    assert!(text.contains("\"echo\""), "the other fields print: {text}");
    assert!(text.contains("http_parts: true"), "{text}");
    let headers = call.headers().expect("HTTP headers");
    for (name, value) in CREDENTIAL_HEADERS {
        assert_eq!(headers[name], value, "the data keeps {name}");
    }
}

/// The snapshot gives the protocol version the request's `_meta` names and the `mcp-session-id` of
/// a streamable HTTP request, each as sent.
#[tokio::test(flavor = "current_thread")]
async fn tool_call_reads_protocol_version_and_transport_session() {
    let mut context = bare_context(None);
    context
        .meta
        .insert(META_PROTOCOL_VERSION_KEY.to_owned(), json!("2025-11-25"));
    let request = http::Request::builder()
        .method("POST")
        .uri("/mcp")
        .header("host", "localhost:8080")
        .header("mcp-session-id", "transport-session-1")
        .body(())
        .expect("valid parts");
    context.extensions.insert(request.into_parts().0);
    let call = ToolCall::new(&params("echo", json!({})), &context, None);
    assert_eq!(call.protocol_version(), Some("2025-11-25"));
    assert_eq!(call.transport_session_id(), Some("transport-session-1"));
}

#[tokio::test(flavor = "current_thread")]
async fn hooks_fire_in_the_documented_order() {
    let log = Log::default();
    // Set in the reverse of the firing order, so that the test would fail if hooks fired in the
    // order they were set.
    let options = Options::default()
        .event_properties(labelled(&log, "event_properties"))
        .event_tags(labelled(&log, "event_tags"))
        .identify(labelled(&log, "identify"))
        .resolve_session_id(labelled(&log, "resolve_session_id"));
    let (client, _sink) = serve(options).await;
    call(&client, "echo", json!({"msg": "hi"})).await;
    assert_eq!(
        labels(&log),
        [
            "resolve_session_id",
            "identify",
            "event_tags",
            "event_properties"
        ]
    );
    client.cancel().await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn hooks_fire_in_the_documented_order_without_hook_mode() {
    let log = Log::default();
    let options = Options::default()
        .event_properties(labelled(&log, "event_properties"))
        .event_tags(labelled(&log, "event_tags"))
        .identify(labelled(&log, "identify"));
    let (client, _sink) = serve(options).await;
    call(&client, "echo", json!({"msg": "hi"})).await;
    assert_eq!(labels(&log), ["identify", "event_tags", "event_properties"]);
    client.cancel().await.unwrap();
}

/// A tool whose `call_tool` pushes `tool` to the log.
struct Recorder(Log);

impl ServerHandler for Recorder {
    fn call_tool(
        &self,
        _request: CallToolRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> impl Future<Output = Result<CallToolResponse, ErrorData>> + MaybeSendFuture + '_ {
        push(&self.0, "tool");
        ready(Ok(CallToolResult::success(vec![]).into()))
    }

    fn get_info(&self) -> InitializeResult {
        InitializeResult::new(ServerCapabilities::builder().enable_tools().build())
    }
}

#[tokio::test(flavor = "current_thread")]
async fn hooks_fire_before_the_tool_runs() {
    let log = Log::default();
    let options = Options::default()
        .identify(labelled(&log, "identify"))
        .event_tags(labelled(&log, "event_tags"));
    let agentcat = AgentCat::with_sink("", options, Arc::new(RecordingSink::default()));
    let client = connect(agentcat.track(Recorder(Arc::clone(&log))), "c").await;
    call(&client, "anything", json!({})).await;
    assert_eq!(labels(&log), ["identify", "event_tags", "tool"]);
    client.cancel().await.unwrap();
}

/// Each hook set alone also makes its first poll before the tool runs, as every hook is called
/// before the tool.
#[tokio::test(flavor = "current_thread")]
async fn each_hook_alone_fires_before_the_tool_runs() {
    for label in [
        "resolve_session_id",
        "identify",
        "event_tags",
        "event_properties",
    ] {
        let log = Log::default();
        let options = Options::default();
        let options = match label {
            "resolve_session_id" => options.resolve_session_id(labelled(&log, label)),
            "identify" => options.identify(labelled(&log, label)),
            "event_tags" => options.event_tags(labelled(&log, label)),
            _ => options.event_properties(labelled(&log, label)),
        };
        let agentcat = AgentCat::with_sink("", options, Arc::new(RecordingSink::default()));
        let client = connect(agentcat.track(Recorder(Arc::clone(&log))), "c").await;
        call(&client, "anything", json!({})).await;
        assert_eq!(labels(&log), [label, "tool"], "{label}");
        client.cancel().await.unwrap();
    }
}

/// The call returns while all four hooks hang, in hook mode (`resolve_session_id` set), so a
/// hanging `resolve_session_id`, `identify`, `event_tags` or `event_properties` never delays it.
#[tokio::test(flavor = "current_thread")]
async fn hooks_never_delay_the_call() {
    let options = Options::default()
        .resolve_session_id(|_call| pending::<Result<Option<String>, HookError>>())
        .identify(|_call| pending::<Result<Option<crate::UserIdentity>, HookError>>())
        .event_tags(|_call| pending::<Result<Option<Map<String, Value>>, HookError>>())
        .event_properties(|_call| pending::<Result<Option<Map<String, Value>>, HookError>>());
    let (client, sink) = serve(options).await;
    let result = tokio::time::timeout(
        Duration::from_secs(5),
        client.call_tool(params("echo", json!({"msg": "hi"}))),
    )
    .await
    .expect("the call returns while the hooks hang")
    .expect("the call works");
    assert_eq!(texts(&result).last(), Some(&"hi"));

    let record = only_record(&sink);
    assert!(record.resolution.hook_mode);
    let pending = &record.pending;
    let session = pending
        .resolve_session_id
        .as_ref()
        .expect("session started");
    let identify = pending.identify.as_ref().expect("identify started");
    let tags = pending.event_tags.as_ref().expect("tags started");
    let properties = pending
        .event_properties
        .as_ref()
        .expect("properties started");
    assert!(!session.is_finished() && !identify.is_finished());
    assert!(!tags.is_finished() && !properties.is_finished());
    client.cancel().await.unwrap();
}

/// Hanging `resolve_session_id`, `event_tags` and `event_properties` hooks do not delay the tool
/// call: with no `identify` hook, the call returns while those three hooks hang, and the record
/// holds the three unsettled tasks and no `identify` task.
#[tokio::test(flavor = "current_thread")]
async fn hanging_session_tags_and_properties_hooks_do_not_delay_call() {
    let options = Options::default()
        .resolve_session_id(|_call| pending::<Result<Option<String>, HookError>>())
        .event_tags(|_call| pending::<Result<Option<Map<String, Value>>, HookError>>())
        .event_properties(|_call| pending::<Result<Option<Map<String, Value>>, HookError>>());
    let (client, sink) = serve(options).await;
    let result = tokio::time::timeout(
        Duration::from_secs(2),
        client.call_tool(params("echo", json!({"msg": "hi"}))),
    )
    .await
    .expect("the call returns while the three hooks hang")
    .expect("the call works");
    assert!(texts(&result).contains(&"hi"), "{result:?}");

    let record = only_record(&sink);
    let pending = &record.pending;
    assert!(pending.identify.is_none(), "no identify hook, no task");
    let session = pending
        .resolve_session_id
        .as_ref()
        .expect("session started");
    let tags = pending.event_tags.as_ref().expect("tags started");
    let properties = pending
        .event_properties
        .as_ref()
        .expect("properties started");
    assert!(!session.is_finished());
    assert!(!tags.is_finished() && !properties.is_finished());
    client.cancel().await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn hook_panic_does_not_reach_the_call() {
    let options = Options::default()
        .identify(
            |_call| -> std::future::Ready<Result<Option<crate::UserIdentity>, HookError>> {
                panic!("identify blew up")
            },
        )
        .event_tags(|call| async move {
            assert!(call.tool_name().is_empty(), "tags blew up");
            Ok(None::<Map<String, Value>>)
        });
    let (client, sink) = serve(options).await;
    for _ in 0..2 {
        let result = call(&client, "echo", json!({"msg": "hi"})).await;
        assert_eq!(
            texts(&result).last(),
            Some(&"hi"),
            "the server keeps serving"
        );
    }
    for record in sink.take_records() {
        let mut pending = record.pending;
        let identify = pending.identify.take().expect("identify started");
        let tags = pending.event_tags.take().expect("tags started");
        assert!(identify.await.unwrap_err().is_panic());
        assert!(tags.await.unwrap_err().is_panic());
    }
    client.cancel().await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn no_hooks_no_tasks() {
    let (client, sink) = serve(Options::default()).await;
    call(&client, "echo", json!({"msg": "hi"})).await;
    let pending = only_record(&sink).pending;
    assert!(pending.resolve_session_id.is_none());
    assert!(pending.identify.is_none());
    assert!(pending.event_tags.is_none());
    assert!(pending.event_properties.is_none());
    assert!(!pending.any());
    client.cancel().await.unwrap();
}

/// Without a runtime a hook does not start, and the sink says so; no snapshot is built.
#[test]
fn no_runtime_no_start() {
    let options = Options::default()
        .resolve_session_id(|_call| ready(Ok(Some("s".to_owned()))))
        .identify(|_call| ready(Ok(None)))
        .event_tags(|_call| ready(Ok(None)))
        .event_properties(|_call| ready(Ok(None)));
    let sink = RecordingSink::default();
    let pending = fire(&options, true, &sink as &dyn Sink, || {
        panic!("no hook starts, so no snapshot is built")
    });
    assert!(!pending.any());
    assert_eq!(
        sink.logs(),
        [
            "Warning: resolve_session_id hook not started - no tokio runtime",
            "Warning: identify hook not started - no tokio runtime",
            "Warning: event_tags hook not started - no tokio runtime",
            "Warning: event_properties hook not started - no tokio runtime",
        ]
    );
}

fn hook_mode_options() -> Options {
    Options::default()
        .enable_agent_tracking(true)
        .resolve_session_id(|_call| ready(Ok(Some("customer-session".to_owned()))))
}

#[tokio::test(flavor = "current_thread")]
async fn hook_mode_listing_has_no_session_id() {
    let (client, _sink) = serve(hook_mode_options()).await;
    let tools = client.list_all_tools().await.unwrap();
    let echo = tools.iter().find(|tool| tool.name == "echo").unwrap();
    let properties = echo.input_schema["properties"].as_object().unwrap();
    assert!(!properties.contains_key("session_id"));
    assert!(properties.contains_key("agent_id"));
    assert!(properties.contains_key("context"));
    client.cancel().await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn hook_mode_resolution_is_provisional() {
    let (client, sink) = serve(hook_mode_options()).await;
    let result = call(
        &client,
        "echo",
        json!({"msg": "hi", "session_id": "start", "agent_id": "opus|t|abcde"}),
    )
    .await;
    assert_eq!(texts(&result), ["hi"], "no mint-back text");
    let record = only_record(&sink);
    assert_eq!(record.resolution.session_id, "");
    assert_eq!(record.resolution.session_source, SessionSource::Hook);
    assert!(record.resolution.hook_mode);
    assert_eq!(record.tags["agentcat_session_id_source"], "hook");
    client.cancel().await.unwrap();
}

#[tokio::test(flavor = "current_thread")]
async fn hook_mode_mirror_has_only_agent_id() {
    let (client, _sink) = serve(hook_mode_options()).await;
    let result = call(
        &client,
        "structured",
        json!({"msg": "hi", "agent_id": "opus|t|abcde"}),
    )
    .await;
    let structured = result.structured_content.expect("structured content");
    assert_eq!(
        structured["mcp_session"],
        json!({"agent_id": "opus|t|abcde"})
    );
    assert_eq!(structured["msg"], "hi");
    client.cancel().await.unwrap();
}

/// A hook that returns an error gives the call no trace of it, and the error stays in the task for
/// the event pipeline to settle.
#[tokio::test(flavor = "current_thread")]
async fn hook_error_does_not_reach_the_call() {
    let options = Options::default().event_properties(|_call| async {
        Err::<Option<Map<String, Value>>, HookError>("properties failed".into())
    });
    let (client, sink) = serve(options).await;
    let result = call(&client, "echo", json!({"msg": "hi"})).await;
    assert_eq!(texts(&result).last(), Some(&"hi"));
    assert_ne!(result.is_error, Some(true));
    let mut pending = only_record(&sink).pending;
    let settled = pending.event_properties.take().expect("started").await;
    let error = settled.expect("no panic").expect_err("the hook failed");
    assert_eq!(error.to_string(), "properties failed");
    client.cancel().await.unwrap();
}

/// Only `tools/call` runs the hooks. `initialize` and `tools/list` run none and record nothing.
#[tokio::test(flavor = "current_thread")]
async fn listing_fires_no_hooks() {
    let log = Log::default();
    let options = Options::default()
        .resolve_session_id(labelled(&log, "resolve_session_id"))
        .identify(labelled(&log, "identify"))
        .event_tags(labelled(&log, "event_tags"))
        .event_properties(labelled(&log, "event_properties"));
    let (client, sink) = serve(options).await;
    client.list_all_tools().await.unwrap();
    assert_eq!(labels(&log), Vec::<String>::new());
    assert!(sink.take_records().is_empty());
    client.cancel().await.unwrap();
}