Rust: rmcp

cargo add agentcat-sdk

[dependencies]
agentcat-sdk = { version = "1", default-features = false, features = ["native-tls"] }

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // One line: wrap the server with AgentCat.
    let server = agentcat_sdk::track(Calculator, "proj_...", Options::default());
    let running = server.serve(rmcp::transport::stdio()).await?;
    running.waiting().await?;
    Ok(())
}

use agentcat_sdk::{AgentCat, Options};

// AGENTCAT_PROJECT_ID=proj_123 in the environment.
let agentcat = AgentCat::from_env(Options::default());

use agentcat_sdk::{HookError, Options, ToolCall, UserIdentity};

let options = Options::default()
    .enable_agent_tracking(true)
    .identify(|call: ToolCall| async move {
        // Read your own data about the caller, for example an HTTP header.
        let user = call
            .headers()
            .and_then(|headers| headers.get("x-user-id"))
            .and_then(|value| value.as_str())
            .map(UserIdentity::new);
        Ok::<_, HookError>(user)
    });

use agentcat_sdk::{AgentCat, Options};
use rmcp::handler::server::wrapper::Parameters;
use rmcp::model::{Implementation, InitializeResult, ServerCapabilities};
use rmcp::transport::stdio;
use rmcp::{ServerHandler, ServiceExt, tool, tool_handler, tool_router};
use schemars::JsonSchema;
use serde::Deserialize;

// ...

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Reads AGENTCAT_PROJECT_ID, and AGENTCAT_API_URL when set. Without a project ID, the handle
    // builds events but sends none.
    let server = AgentCat::from_env(Options::default()).track(Demo);

    // Serve the wrapper where you served `Demo`.
    let running = server.serve(stdio()).await?;
    running.waiting().await?;
    Ok(())
}

use std::sync::Arc;
use std::time::Duration;

use agentcat_sdk::{AgentCat, Options};
use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
use rmcp::transport::streamable_http_server::{StreamableHttpServerConfig, StreamableHttpService};

// ...

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // One handle for the whole process. It reads AGENTCAT_PROJECT_ID, and AGENTCAT_API_URL when set.
    let agentcat = AgentCat::from_env(Options::default());

    // rmcp calls this closure once for each new MCP session.
    let handle = agentcat.clone();
    let service = StreamableHttpService::new(
        move || Ok(handle.track(Demo)),
        Arc::new(LocalSessionManager::default()),
        StreamableHttpServerConfig::default(),
    );
    let router = axum::Router::new().route_service("/mcp", service);

    // ...
    tokio::select! {
        served = axum::serve(listener, router) => served?,
        () = stop_signal() => {}
    }

    // Deliver the events that are still queued before the process ends.
    let delivered = agentcat.shutdown(Duration::from_secs(5)).await;
    // ...
    Ok(())
}

// One handle for the whole process. It reads AGENTCAT_PROJECT_ID, and AGENTCAT_API_URL when set.
let agentcat = AgentCat::from_env(Options::default());

// rmcp calls this closure once for each request.
let handle = agentcat.clone();
let service = StreamableHttpService::new(
    move || Ok(handle.track(Demo)),
    Arc::new(LocalSessionManager::default()),
    // No sessions, and JSON responses instead of SSE.
    StreamableHttpServerConfig::default()
        .with_legacy_session_mode(false)
        .with_json_response(true),
);

/// Resolves on SIGTERM (Unix) or Ctrl-C.
async fn stop_signal() {
    #[cfg(unix)]
    {
        use tokio::signal::unix::{SignalKind, signal};
        if let Ok(mut terminate) = signal(SignalKind::terminate()) {
            tokio::select! {
                _ = terminate.recv() => {}
                _ = tokio::signal::ctrl_c() => {}
            }
            return;
        }
    }
    let _ = tokio::signal::ctrl_c().await;
}

/// Sends the text back, and publishes one custom event for each call. The SDK passes the
/// `CallContext` of the call to the tool as an `Extension`.
#[tool(description = "Sends the text back")]
async fn echo(
    &self,
    Parameters(args): Parameters<EchoArgs>,
    Extension(call): Extension<CallContext>,
) -> String {
    let mut event = CustomEvent::new()
        .resource_name("echo-called")
        .message("the echo tool ran");
    // Attach the event to the session of the call.
    if let Some(session) = call.session_id() {
        event = event.session_id(session);
    }
    // Queues the event and returns at once. An event that cannot be queued is not an error
    // for the tool.
    let _ = call.publish_custom_event(&self.project_id, event);
    args.text
}

#[tool(description = "Always fails")]
async fn fail(&self) -> Result<String, String> {
    let error = std::io::Error::new(std::io::ErrorKind::NotFound, "the record does not exist");
    capture_error(&error);
    Err(error.to_string())
}

use agentcat_sdk::{AgentCat, CustomEvent, Options};

let agentcat = AgentCat::new("proj_123", Options::default());
let event = CustomEvent::new().resource_name("cache-warmed").message("40 keys");
agentcat.publish_custom_event("ses_1", "proj_123", event)?;

tracked.publish_custom_event("proj_123", CustomEvent::new().resource_name("startup"))?;

[dev-dependencies]
agentcat-sdk = { version = "1", features = ["test-util"] }
rmcp = { version = "3", features = ["client"] }
tokio = { version = "1", features = ["io-util", "macros", "rt"] }

use agentcat_sdk::Options;
use agentcat_sdk::testing::Fake;
use rmcp::ServiceExt;
use rmcp::model::{
    CallToolRequestParams, ClientCapabilities, Implementation, InitializeRequestParams,
};

let fake = Fake::new();
let server = fake.agentcat("proj_test", Options::default()).track(Calculator);

// The server on one end of an in-memory stream, an rmcp client on the other.
let (server_end, client_end) = tokio::io::duplex(64 * 1024);
tokio::spawn(async move {
    if let Ok(running) = server.serve(server_end).await {
        let _ = running.waiting().await;
    }
});
let client = InitializeRequestParams::new(
    ClientCapabilities::default(),
    Implementation::new("test", "1.0.0"),
)
.serve(client_end)
.await?;

let mut arguments = rmcp::model::JsonObject::new();
arguments.insert("a".into(), 1.into());
arguments.insert("b".into(), 2.into());
client
    .call_tool(CallToolRequestParams::new("add").with_arguments(arguments))
    .await?;

fake.assert_tracked("add").await;

use agentcat_sdk::{Event, HookError};

// The closure that you pass to `Options::redact_event`.
let hook = |mut event: Event| async move {
    event.response = None;
    Ok::<_, HookError>(Some(event))
};

let mut event = Event::default();
event.response = Some(serde_json::json!({"token": "fake-secret-42"}));
let redacted = hook(event).await?;
assert_eq!(redacted.and_then(|event| event.response), None);

use std::time::Duration;

// ... serve until the server stops ...
let delivered = agentcat.shutdown(Duration::from_secs(5)).await;
if !delivered {
    // Some events did not go out within 5 s.
}