tmcp 0.4.0

Complete, ergonomic implementation of the Model Context Protocol (MCP)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137

//! Client/server ping integration tests.

#[cfg(test)]
mod tests {
    use std::sync::{Arc, Mutex};

    use async_trait::async_trait;
    use tmcp::{
        ClientCtx, ClientHandler, Result, ServerCtx, ServerHandler,
        schema::*,
        testutils::{
            connected_client_and_server_with_conn, shutdown_client_and_server, test_client_ctx,
        },
    };
    use tokio::sync::mpsc;
    use tracing_subscriber::fmt;

    #[derive(Default, Clone)]
    struct TestClientHandler {
        calls: Arc<Mutex<Vec<String>>>,
    }

    #[async_trait]
    impl ClientHandler for TestClientHandler {
        async fn on_connect(&self, _ctx: &ClientCtx) -> Result<()> {
            self.calls.lock().unwrap().push("on_connect".into());
            Ok(())
        }

        async fn on_shutdown(&self, _ctx: &ClientCtx) -> Result<()> {
            self.calls.lock().unwrap().push("on_shutdown".into());
            Ok(())
        }

        async fn pong(&self, _ctx: &ClientCtx) -> Result<()> {
            self.calls.lock().unwrap().push("ping".into());
            Ok(())
        }

        async fn create_message(
            &self,
            _ctx: &ClientCtx,
            _method: &str,
            _params: CreateMessageParams,
        ) -> Result<CreateMessageResult> {
            self.calls.lock().unwrap().push("create_message".into());
            Ok(CreateMessageResult {
                message: SamplingMessage::assistant_text("Test response"),
                model: "test-model".into(),
                stop_reason: None,
            })
        }

        async fn list_roots(&self, _ctx: &ClientCtx) -> Result<ListRootsResult> {
            self.calls.lock().unwrap().push("list_roots".into());
            Ok(ListRootsResult {
                roots: vec![Root {
                    uri: "test://root".into(),
                    name: Some("Test Root".into()),
                    _meta: None,
                }],
                _meta: None,
            })
        }
    }

    struct TestServerHandler;

    #[async_trait]
    impl ServerHandler for TestServerHandler {
        async fn initialize(
            &self,
            _ctx: &ServerCtx,
            _protocol_version: String,
            _capabilities: ClientCapabilities,
            _client_info: Implementation,
        ) -> Result<InitializeResult> {
            Ok(InitializeResult::new("test-server").with_version("1.0.0"))
        }

        async fn pong(&self, _ctx: &ServerCtx) -> Result<()> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn client_connection_trait_methods() {
        let connection = TestClientHandler::default();

        let (tx, _) = mpsc::unbounded_channel();
        let ctx = test_client_ctx(tx);

        connection.pong(&ctx).await.expect("Ping failed");

        let params = CreateMessageParams::user_message("Hello").with_max_tokens(1000);

        let result = connection
            .create_message(&ctx, "test", params)
            .await
            .expect("Create message failed");
        assert_eq!(result.model, "test-model");

        let roots = connection.list_roots(&ctx).await.unwrap();
        assert_eq!(roots.roots.len(), 1);

        let calls = connection.calls.lock().unwrap();
        assert!(calls.contains(&"ping".to_string()));
        assert!(calls.contains(&"create_message".to_string()));
        assert!(calls.contains(&"list_roots".to_string()));
    }

    #[tokio::test]
    async fn client_server_ping() {
        fmt::try_init().ok();

        let calls = Arc::new(Mutex::new(Vec::new()));

        let (mut client, handle) = connected_client_and_server_with_conn(
            || Box::new(TestServerHandler),
            TestClientHandler {
                calls: calls.clone(),
            },
        )
        .await
        .expect("setup");

        client.init().await.expect("client init");
        client.ping().await.expect("client ping");

        {
            let list = calls.lock().unwrap();
            assert!(list.contains(&"on_connect".to_string()));
        }

        shutdown_client_and_server(client, handle).await;
    }
}