aether-mcp-utils 0.7.14

MCP client and server utilities for the Aether AI agent framework
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

use aether_auth::SigV4Signer;
use futures::stream::BoxStream;
use reqwest::header::{HeaderName, HeaderValue};
use rmcp::model::ClientJsonRpcMessage;
use rmcp::transport::auth::AuthError;
use rmcp::transport::streamable_http_client::{
    SseError, StreamableHttpClient, StreamableHttpError, StreamableHttpPostResponse,
};
use sse_stream::Sse;
use std::collections::HashMap;
use std::sync::Arc;

#[derive(Clone)]
pub(super) struct SigV4HttpClient {
    http: reqwest::Client,
    signer: SigV4Signer,
}

impl SigV4HttpClient {
    pub(super) fn new(signer: SigV4Signer) -> Self {
        Self { http: reqwest::Client::default(), signer }
    }

    async fn sign_message(
        &self,
        uri: &str,
        message: &ClientJsonRpcMessage,
        headers: HashMap<HeaderName, HeaderValue>,
    ) -> Result<HashMap<HeaderName, HeaderValue>, StreamableHttpError<reqwest::Error>> {
        self.sign("POST", uri, &serde_json::to_vec(message)?, headers).await
    }

    async fn sign(
        &self,
        method: &str,
        uri: &str,
        body: &[u8],
        mut headers: HashMap<HeaderName, HeaderValue>,
    ) -> Result<HashMap<HeaderName, HeaderValue>, StreamableHttpError<reqwest::Error>> {
        let signed = self
            .signer
            .headers(method, uri, body)
            .await
            .map_err(|error| StreamableHttpError::Auth(AuthError::AuthorizationFailed(error.to_string())))?;

        headers.extend(signed.into_iter().filter_map(|(name, value)| Some((name?, value))));
        Ok(headers)
    }
}

impl StreamableHttpClient for SigV4HttpClient {
    type Error = reqwest::Error;

    async fn post_message(
        &self,
        uri: Arc<str>,
        message: ClientJsonRpcMessage,
        session_id: Option<Arc<str>>,
        auth_header: Option<String>,
        custom_headers: HashMap<HeaderName, HeaderValue>,
    ) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
        let headers = self.sign_message(&uri, &message, custom_headers).await?;
        self.http.post_message(uri, message, session_id, auth_header, headers).await
    }

    async fn post_message_with_max_sse_event_size(
        &self,
        uri: Arc<str>,
        message: ClientJsonRpcMessage,
        session_id: Option<Arc<str>>,
        auth_header: Option<String>,
        custom_headers: HashMap<HeaderName, HeaderValue>,
        max_sse_event_size: usize,
    ) -> Result<StreamableHttpPostResponse, StreamableHttpError<Self::Error>> {
        let headers = self.sign_message(&uri, &message, custom_headers).await?;
        self.http
            .post_message_with_max_sse_event_size(uri, message, session_id, auth_header, headers, max_sse_event_size)
            .await
    }

    async fn delete_session(
        &self,
        uri: Arc<str>,
        session_id: Arc<str>,
        auth_header: Option<String>,
        custom_headers: HashMap<HeaderName, HeaderValue>,
    ) -> Result<(), StreamableHttpError<Self::Error>> {
        let headers = self.sign("DELETE", &uri, &[], custom_headers).await?;
        self.http.delete_session(uri, session_id, auth_header, headers).await
    }

    async fn get_stream(
        &self,
        uri: Arc<str>,
        session_id: Option<Arc<str>>,
        last_event_id: Option<String>,
        auth_header: Option<String>,
        custom_headers: HashMap<HeaderName, HeaderValue>,
    ) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
        let headers = self.sign("GET", &uri, &[], custom_headers).await?;
        self.http.get_stream(uri, session_id, last_event_id, auth_header, headers).await
    }

    async fn get_stream_with_max_sse_event_size(
        &self,
        uri: Arc<str>,
        session_id: Option<Arc<str>>,
        last_event_id: Option<String>,
        auth_header: Option<String>,
        custom_headers: HashMap<HeaderName, HeaderValue>,
        max_sse_event_size: usize,
    ) -> Result<BoxStream<'static, Result<Sse, SseError>>, StreamableHttpError<Self::Error>> {
        let headers = self.sign("GET", &uri, &[], custom_headers).await?;
        self.http
            .get_stream_with_max_sse_event_size(
                uri,
                session_id,
                last_event_id,
                auth_header,
                headers,
                max_sse_event_size,
            )
            .await
    }
}