soma_mcp_client/
oauth.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use futures::future::BoxFuture;
5use serde::Serialize;
6use thiserror::Error;
7
8use crate::config::UpstreamConfig;
9use crate::upstream::http_body_cap::BodyCappedHttpClient;
10
11pub type UpstreamOAuthHttpClient = rmcp::transport::AuthClient<BodyCappedHttpClient>;
12
13#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
14pub struct BeginAuthorization {
15    pub authorization_url: String,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub struct UpstreamOAuthCredentialStatus {
20    pub access_token_expires_at: i64,
21    pub refresh_token_present: bool,
22}
23
24#[derive(Debug, Error, Clone, PartialEq, Eq)]
25#[error("{kind}: {message}")]
26pub struct UpstreamOAuthError {
27    kind: String,
28    message: String,
29}
30
31impl UpstreamOAuthError {
32    #[must_use]
33    pub fn new(kind: impl Into<String>, message: impl Into<String>) -> Self {
34        Self {
35            kind: kind.into(),
36            message: message.into(),
37        }
38    }
39
40    #[must_use]
41    pub fn internal(message: impl Into<String>) -> Self {
42        Self::new("internal_error", message)
43    }
44
45    #[must_use]
46    pub fn kind(&self) -> &str {
47        &self.kind
48    }
49
50    #[must_use]
51    pub fn message(&self) -> &str {
52        &self.message
53    }
54}
55
56pub trait UpstreamOAuthProvider: Send + Sync {
57    fn authenticated_http_client<'a>(
58        &'a self,
59        upstream: &'a UpstreamConfig,
60        subject: &'a str,
61        http_client: BodyCappedHttpClient,
62    ) -> BoxFuture<'a, Result<UpstreamOAuthHttpClient, UpstreamOAuthError>>;
63
64    fn evict_subject(&self, _upstream: &str, _subject: &str) {}
65
66    fn evict_upstream(&self, _upstream: &str) {}
67}
68
69pub trait UpstreamOAuthManager: Send + Sync {
70    fn begin_authorization<'a>(
71        &'a self,
72        subject: &'a str,
73    ) -> BoxFuture<'a, Result<BeginAuthorization, UpstreamOAuthError>>;
74
75    fn complete_authorization<'a>(
76        &'a self,
77        subject: &'a str,
78        code: &'a str,
79        state: &'a str,
80        issuer: Option<&'a str>,
81    ) -> BoxFuture<'a, Result<(), UpstreamOAuthError>>;
82
83    fn credential_status<'a>(
84        &'a self,
85        subject: &'a str,
86    ) -> BoxFuture<'a, Result<Option<UpstreamOAuthCredentialStatus>, UpstreamOAuthError>>;
87
88    fn clear_credentials<'a>(
89        &'a self,
90        subject: &'a str,
91    ) -> BoxFuture<'a, Result<(), UpstreamOAuthError>>;
92
93    fn access_token<'a>(
94        &'a self,
95        subject: &'a str,
96    ) -> BoxFuture<'a, Result<String, UpstreamOAuthError>>;
97}
98
99#[derive(Clone)]
100pub struct UpstreamOAuthRuntime {
101    provider: Arc<dyn UpstreamOAuthProvider>,
102    managers: Arc<BTreeMap<String, Arc<dyn UpstreamOAuthManager>>>,
103}
104
105impl UpstreamOAuthRuntime {
106    #[must_use]
107    pub fn new(
108        provider: Arc<dyn UpstreamOAuthProvider>,
109        managers: BTreeMap<String, Arc<dyn UpstreamOAuthManager>>,
110    ) -> Self {
111        Self {
112            provider,
113            managers: Arc::new(managers),
114        }
115    }
116
117    #[must_use]
118    pub fn provider(&self) -> Arc<dyn UpstreamOAuthProvider> {
119        Arc::clone(&self.provider)
120    }
121
122    #[must_use]
123    pub fn manager(&self, upstream: &str) -> Option<Arc<dyn UpstreamOAuthManager>> {
124        self.managers.get(upstream).cloned()
125    }
126
127    pub fn evict_subject(&self, upstream: &str, subject: &str) {
128        self.provider.evict_subject(upstream, subject);
129    }
130
131    pub fn evict_upstream(&self, upstream: &str) {
132        self.provider.evict_upstream(upstream);
133    }
134}
135
136#[cfg(test)]
137#[path = "oauth_tests.rs"]
138mod tests;
soma_mcp_client/oauth.rs

soma_mcp_client/
oauth.rs
